feat(ops): add B02 TSELS and TPRELU support

.github/workflows/ci.yml

@@ -1074,7 +1074,7 @@ jobs:
         # against the runtime's bundled pto-isa, which still clones from the
         # stale PTO-ISA/pto-isa mirror (lacks pto::Coalesce) until the runtime
         # submodule is bumped past simpler #806. Re-enable once that lands.
-        run: pytest tests/st/runtime/ops/test_assemble.py tests/st/runtime/ops/test_gather.py tests/st/runtime/ops/test_mscatter.py tests/st/runtime/ops/test_random.py tests/st/runtime/framework_and_models/test_qwen3_decode_scope3_mixed.py tests/st/runtime/control_flow/test_dyn_orch_shape.py::TestDynOrchShapeOperations::test_dyn_orch_paged_attention -v --platform=a5sim --forked -k "not TestMscatter"
+        run: pytest tests/st/runtime/ops/test_assemble.py tests/st/runtime/ops/test_gather.py tests/st/runtime/ops/test_mscatter.py tests/st/runtime/ops/test_prelu.py tests/st/runtime/ops/test_random.py tests/st/runtime/ops/test_sels.py tests/st/runtime/framework_and_models/test_qwen3_decode_scope3_mixed.py tests/st/runtime/control_flow/test_dyn_orch_shape.py::TestDynOrchShapeOperations::test_dyn_orch_paged_attention -v --platform=a5sim --forked -k "not TestMscatter"
 
       - name: Test A5 cross-core system tests (simulator)
         run: pytest tests/st/runtime/cross_core/test_cross_core.py -v --forked --platform=a5sim

docs/en/dev/passes/31-memory_reuse.md

@@ -73,6 +73,8 @@ program_optimized = reuse_pass(program)
   | `tile.fmod`, `tile.fmods` | `not_inplace_safe` | `TFMOD`/`TFMODS` compute `a - trunc(a/b)*b` by overwriting `dst = a/b` first, then re-reading the original `src0` (`a`) for the final subtraction; when `dst == src0` that subtraction sees the already-clobbered quotient and yields `0` for every element |
   | `tile.transpose` | `not_inplace_safe` | `pto.ttrans` is not in-place safe: the a2a3 unaligned scalar path writes `dst` directly from `src` (no tmp staging), so `dst == src` corrupts the data mid-write. The output always gets a fresh buffer (also enforced in InitMemRef, which never inherits the input's buffer for it). |
   | `tile.sel` | `forbid_output_alias(0)` (mask), `(3)` (tmp) | `TSEL` reads the predicate mask + tmp scratch while writing `dst` |
+  | `tile.sels` | target-aware | `TSELS` keeps `dst` disjoint from the predicate mask and may reuse `src` or `tmp`; A2/A3 writes the scalar into `tmp` and loads it with `set_cmpmask` before writing `dst`, so `tmp` may alias `dst` but must remain disjoint from mask/src; A5 retains an unread `tmp` ABI operand that may alias any operand |
+  | `tile.prelu` | target-aware | A2/A3 is `not_inplace_safe` because `TPRELU` reads `src`, `slope`, and `tmp` while writing `dst`; A5 retains the ABI-required `tmp` operand but does not read it, so `dst` may reuse `tmp` but not the active `src`/`slope` inputs |
   | `tile.{row,col}_expand{,_mul,_add,_sub,_div}` | `forbid_output_alias(1)` (broadcast vector) | the row/col vector (arg 1) is re-read for **every** output row/col, so an output aliasing it is overwritten after the first row/col |
   | `tile.cast` (widening only) | output ≠ input buffer (conditional, in `ForbidAliasCollector`) | wider output's write cursor outruns the read cursor (see above) |
 
@@ -245,6 +247,8 @@ passes.def("memory_reuse", &pass::MemoryReuse, "Memory reuse optimization");
 - Tests the no-alias guard (`TestForbidOutputAlias` + `TestInplaceOps`), one case per constraint above:
   - `tile.recip` / `tile.rsqrt` / `tile.row_sum` — output must not alias input (`not_inplace_safe`)
   - `tile.sel` — output must not alias the mask / tmp (`forbid_output_alias`)
+  - `tile.sels` — output never aliases mask; both A2/A3 and A5 permit tmp/output alias, while A2/A3 backend validation still rejects tmp overlap with mask/src
+  - `tile.prelu` — A2/A3 output must not alias any input; A5 output may alias only the unused `tmp`
   - `tile.col_expand_mul` — output must not alias the broadcast vector
   - widening `tile.cast` — output must not alias the (narrower) input
   - a forbidden operand reached through a VIEW is still honored (physical-buffer resolution)

docs/en/dev/ptoas-op-status.md

@@ -100,7 +100,7 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT.
 | pto.tpartargmax | TPARTARGMAX | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path |
 | pto.tpartargmin | TPARTARGMIN | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path |
 | pto.tpartmul | TPARTMUL | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
-| pto.tprelu | TPRELU | tile | ✅ | ✅ | ❌ | ❌ | — | path exists; historical ISA/semantic issue requires revalidation against the current pin |
+| pto.tprelu | TPRELU | tile | ✅ | ✅ | ❌ | ✅ | — | canonical 3-input path; verified on A2/A3 hardware, A5 hardware verification pending |
 | pto.tadds | TADDS | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tsubs | TSUBS | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | verified on A2/A3 hardware; A5 hardware verification pending |
 | pto.tmuls | TMULS | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
@@ -165,7 +165,7 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT.
 | pto.tcmp | TCMP | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tcmps | TCMPS | tile | ✅ | ✅ | ❌ | ✅ | — |  |
 | pto.tsel | TSEL | tile | ✅ | ✅ | ❌ | ✅ | — |  |
-| pto.tsels | TSELS | tile | ✅ | ✅ | ❌ | ❌ | — | frontend/codegen path exists; same-name ST is missing |
+| pto.tsels | TSELS | tile | ✅ | ✅ | ❌ | ✅ | — | canonical 4-input path; verified on A2/A3 hardware, A5 hardware verification pending |
 | **Bitwise Operations (11)** |  |  |  |  |  |  |  |  |
 | pto.tand | TAND | tile+tensor | ✅ | ✅ | ✅ | ❌ | — | path exists; historical ISA/semantic issue requires revalidation against the current pin |
 | pto.tor | TOR | tile+tensor | ✅ | ✅ | ✅ | ❌ | — | path exists; historical ISA/semantic issue requires revalidation against the current pin |

docs/en/user/02-operation_reference.md

@@ -235,7 +235,7 @@ scratch tile to materialize numeric results on A2/A3.
 | `cmp` | `(lhs: Tile, rhs: Tile, cmp_type: int = 0) -> Tile` | Compare two tiles |
 | `cmps` | `(lhs: Tile, rhs: int \| float \| Scalar, cmp_type: int = 0) -> Tile` | Compare tile with scalar |
 | `sel` | `(mask: Tile, lhs: Tile, rhs: Tile, tmp: Tile) -> Tile` | Select: `lhs if mask else rhs`; `tmp` is TSEL scratch |
-| `sels` | `(lhs: Tile, rhs: Tile, select_mode: int \| float \| Scalar) -> Tile` | Select by scalar mode |
+| `sels` | `(mask: Tile, src: Tile, tmp: Tile, scalar: int \| float \| Scalar) -> Tile` | Select `src` where `mask` is true, otherwise `scalar`; `mask` must have enough valid rows and packed bytes per row to cover `src`; A2/A3 supports signed/unsigned 16/32-bit integers plus FP16/FP32, requires `tmp` not to overlap mask/src, and permits `tmp` to alias the result; A5 also supports signed/unsigned 8-bit integers and retains an unread `tmp` ABI operand that may alias any operand |
 
 ## Bitwise (`pl.tensor.*`)
 
@@ -303,7 +303,7 @@ normalize it.
 | ---- | --------- | ----------- |
 | `relu` | `(tile: Tile) -> Tile` | ReLU: `max(0, x)` |
 | `lrelu` | `(tile: Tile, slope: int \| float \| Scalar) -> Tile` | Leaky ReLU with scalar slope |
-| `prelu` | `(tile: Tile, slope: Tile, tmp: Tile) -> Tile` | Parametric ReLU (requires tmp) |
+| `prelu` | `(tile: Tile, slope: Tile, tmp: Tile) -> Tile` | FP16/FP32 parametric ReLU; A2/A3 requires pairwise non-overlapping `tile`/`slope`/`tmp`/result regions and UINT8 packed-mask scratch; A5 retains the ABI-required but unread `tmp`, permits `tile`/`slope` overlap and `tmp`/result aliasing, but keeps the result disjoint from the active inputs |
 
 ## Shape Operations (`pl.tile.*`)
 

docs/zh/dev/passes/31-memory_reuse.md

@@ -73,6 +73,8 @@ program_optimized = reuse_pass(program)
   | `tile.fmod`、`tile.fmods` | `not_inplace_safe` | `TFMOD`/`TFMODS` 按 `a - trunc(a/b)*b` 计算，先用 `dst = a/b` 覆盖输出，再重新读取原始 `src0`（`a`）做最后的减法；当 `dst == src0` 时该减法读到的是已被覆盖的商，导致每个元素都算成 `0` |
   | `tile.transpose` | `not_inplace_safe` | `pto.ttrans` 非 in-place 安全：a2a3 非对齐标量路径直接从 `src` 写 `dst`（不经 tmp 暂存），`dst == src` 会边写边读损坏数据。输出始终分配新 buffer（InitMemRef 也不会为其继承输入的 buffer）。 |
   | `tile.sel` | `forbid_output_alias(0)`（mask）、`(3)`（tmp） | `TSEL` 在写 `dst` 时读取 mask + tmp scratch |
+  | `tile.sels` | 感知 target | `TSELS` 始终要求 `dst` 与 predicate mask 分离，并允许复用 `src` 或 `tmp`；A2/A3 会先将 scalar 写入 `tmp`，再通过 `set_cmpmask` 读取它，之后才写 `dst`，因此 `tmp` 可以 alias `dst`，但不得与 mask/src 重叠；A5 保留但不读取 ABI 中的 `tmp`，允许其 alias 任一操作数 |
+  | `tile.prelu` | 感知 target | A2/A3 的 `TPRELU` 在写 `dst` 时读取 `src`、`slope` 与 `tmp`，因此是 `not_inplace_safe`；A5 保留 ABI 要求的 `tmp` 操作数但不读取它，所以 `dst` 可复用 `tmp`，但不可复用仍参与运算的 `src`/`slope` |
   | `tile.{row,col}_expand{,_mul,_add,_sub,_div}` | `forbid_output_alias(1)`（广播向量） | 行/列向量（arg 1）会被**每个**输出行/列重读,输出若 alias 它则在第一行/列后被覆盖 |
   | `tile.cast`（仅升精度） | 输出 ≠ 输入缓冲区（条件式,在 `ForbidAliasCollector`） | 更宽的输出写指针超前于读指针（见上） |
 
@@ -237,6 +239,8 @@ passes.def("memory_reuse", &pass::MemoryReuse, "Memory reuse optimization");
 - 测试 no-alias 守护（`TestForbidOutputAlias` + `TestInplaceOps`），上表每条约束一个用例：
   - `tile.recip` / `tile.rsqrt` / `tile.row_sum` —— 输出不得 alias 输入（`not_inplace_safe`）
   - `tile.sel` —— 输出不得 alias mask / tmp（`forbid_output_alias`）
+  - `tile.sels` —— 输出始终不得 alias mask；A2/A3 与 A5 均允许 tmp/输出 alias，但 A2/A3 backend 仍会拒绝 tmp 与 mask/src 重叠
+  - `tile.prelu` —— A2/A3 输出不得 alias 任一输入；A5 输出仅可 alias 未使用的 `tmp`
   - `tile.col_expand_mul` —— 输出不得 alias 广播向量
   - 升精度 `tile.cast` —— 输出不得 alias（更窄的）输入
   - 经 VIEW 间接到达的禁止操作数也被遵守（物理缓冲区解析）

docs/zh/dev/ptoas-op-status.md

@@ -86,7 +86,7 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO
 | pto.tpartargmax | TPARTARGMAX | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 |
 | pto.tpartargmin | TPARTARGMIN | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 |
 | pto.tpartmul | TPARTMUL | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
-| pto.tprelu | TPRELU | tile | ✅ | ✅ | ❌ | ❌ | — | 已有链路；历史 ISA/语义问题，需按当前 pin 复验 |
+| pto.tprelu | TPRELU | tile | ✅ | ✅ | ❌ | ✅ | — | 已补齐规范 3 输入链路；A2/A3 真机已验证，A5 真机待验证 |
 | pto.tadds | TADDS | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tsubs | TSUBS | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | A2/A3 真机已验证；A5 真机待验证 |
 | pto.tmuls | TMULS | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
@@ -151,7 +151,7 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO
 | pto.tcmp | TCMP | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tcmps | TCMPS | tile | ✅ | ✅ | ❌ | ✅ | — |  |
 | pto.tsel | TSEL | tile | ✅ | ✅ | ❌ | ✅ | — |  |
-| pto.tsels | TSELS | tile | ✅ | ✅ | ❌ | ❌ | — | 前端/codegen 已有，缺同名 ST |
+| pto.tsels | TSELS | tile | ✅ | ✅ | ❌ | ✅ | — | 已补齐规范 4 输入链路；A2/A3 真机已验证，A5 真机待验证 |
 | **位运算（11）** |  |  |  |  |  |  |  |  |
 | pto.tand | TAND | tile+tensor | ✅ | ✅ | ✅ | ❌ | — | 已有链路；历史 ISA/语义问题，需按当前 pin 复验 |
 | pto.tor | TOR | tile+tensor | ✅ | ✅ | ✅ | ❌ | — | 已有链路；历史 ISA/语义问题，需按当前 pin 复验 |

docs/zh/user/02-operation_reference.md

@@ -233,7 +233,7 @@ packed predicate mask；A2/A3 上如需得到数值结果，请配合 `sel` 和
 | `cmp` | `(lhs: Tile, rhs: Tile, cmp_type: int = 0) -> Tile` | 比较两个 tile |
 | `cmps` | `(lhs: Tile, rhs: int \| float \| Scalar, cmp_type: int = 0) -> Tile` | tile 与标量比较 |
 | `sel` | `(mask: Tile, lhs: Tile, rhs: Tile, tmp: Tile) -> Tile` | 选择：`mask 为真取 lhs，否则取 rhs`；`tmp` 是 TSEL scratch |
-| `sels` | `(lhs: Tile, rhs: Tile, select_mode: int \| float \| Scalar) -> Tile` | 按标量模式选择 |
+| `sels` | `(mask: Tile, src: Tile, tmp: Tile, scalar: int \| float \| Scalar) -> Tile` | `mask` 为真时选择 `src`，否则选择 `scalar`；`mask` 的有效行数和每行 packed 字节数必须足以覆盖 `src`；A2/A3 支持有符号/无符号 16/32 位整数及 FP16/FP32，要求 `tmp` 不与 mask/src 重叠，但允许 `tmp` alias 结果；A5 还支持有符号/无符号 8 位整数，并保留但不读取 ABI 中的 `tmp`，允许其 alias 任一操作数 |
 
 ## 位运算（`pl.tensor.*`）
 
@@ -297,7 +297,7 @@ kernel 可以在后端修复后自动受益，无需改动前端。
 | ---- | ---- | ---- |
 | `relu` | `(tile: Tile) -> Tile` | ReLU：`max(0, x)` |
 | `lrelu` | `(tile: Tile, slope: int \| float \| Scalar) -> Tile` | 带标量斜率的 Leaky ReLU |
-| `prelu` | `(tile: Tile, slope: Tile, tmp: Tile) -> Tile` | 参数化 ReLU（需要 tmp） |
+| `prelu` | `(tile: Tile, slope: Tile, tmp: Tile) -> Tile` | FP16/FP32 参数化 ReLU；A2/A3 要求 `tile`、`slope`、`tmp`、结果的内存区间两两不重叠，并需要 UINT8 packed-mask 临时空间；A5 保留 ABI 要求但不读取的 `tmp`，允许 `tile`/`slope` 重叠及 `tmp`/结果 alias，但结果仍须与有效输入分离 |
 
 ## 形状操作（`pl.tile.*`）
 

python/pypto/debug/torch_codegen.py

@@ -1027,7 +1027,7 @@ def _register_ops() -> None:  # noqa: PLR0915
 
     # tile selection
     m["tile.sel"] = lambda a, _kw: f"torch.where({a[0]}, {a[1]}, {a[2]})"
-    m["tile.sels"] = lambda a, _kw: f"torch.where({a[0]}, {a[1]}, {a[2]})"
+    m["tile.sels"] = lambda a, _kw: f"torch.where({a[0]}, {a[1]}, {a[3]})"
     m["tile.lrelu"] = lambda a, _kw: f"torch.where({a[0]} > 0, {a[0]}, {a[0]} * {a[1]})"
 
     # tile ternary add/sub with carry

python/pypto/ir/op/tile_ops.py

@@ -77,6 +77,28 @@ def _create_tile_binary_call(
     return _ir_core.create_op_call(tile_op_name, [lhs, rhs_expr], {}, span)
 
 
+def _normalize_sels_scalar_operand(src: Expr, scalar: int | float | Expr, span: Span) -> Expr:
+    """Normalize TSELS scalar constants to the PTOAS-compatible element dtype."""
+    scalar_expr = _normalize_scalar_operand(src, scalar, span, retype_constants=True)
+    src_type = src.type
+    if not isinstance(src_type, _ir_core.TileType) or not isinstance(scalar_expr, ConstInt):
+        return scalar_expr
+
+    signed_dtype_and_bits = {
+        DataType.UINT8: (DataType.INT8, 8),
+        DataType.UINT16: (DataType.INT16, 16),
+        DataType.UINT32: (DataType.INT32, 32),
+    }.get(src_type.dtype)
+    if signed_dtype_and_bits is None:
+        return scalar_expr
+
+    signed_dtype, bits = signed_dtype_and_bits
+    value = scalar_expr.value
+    if value >= 1 << (bits - 1):
+        value -= 1 << bits
+    return ConstInt(value, signed_dtype, span)
+
+
 # ============================================================================
 # Memory Operations
 # ============================================================================
@@ -1306,25 +1328,33 @@ def sel(mask: Expr, lhs: Expr, rhs: Expr, tmp: Expr, span: Span | None = None) -
     return _ir_core.create_op_call("tile.sel", [mask, lhs, rhs, tmp], {}, actual_span)
 
 
-def sels(lhs: Expr, rhs: Expr, select_mode: int | float | Expr, span: Span | None = None) -> Call:
-    """Select between two tiles based on a scalar mode.
+def sels(
+    mask: Expr,
+    src: Expr,
+    tmp: Expr,
+    scalar: int | float | Expr,
+    span: Span | None = None,
+) -> Call:
+    """Per-element selection between a source tile and a scalar.
 
-    Maps to the TSELS hardware intrinsic. The interpretation of select_mode values
-    is target-dependent and enforced by codegen.
+    For each element (i, j): dst[i,j] = src[i,j] if mask[i,j] is true,
+    else scalar. Maps to the TSELS hardware intrinsic.
 
     Args:
-        lhs: Source tile 0 (TileType)
-        rhs: Source tile 1 (TileType)
-        select_mode: Scalar select mode
+        mask: Predicate mask tile (TileType); encoding is target-defined
+        src: Source tile, selected where mask is true (TileType)
+        tmp: Scratch tile required by TSELS (TileType)
+        scalar: Scalar value, selected where mask is false. For an unsigned
+            integer src, constants use the same-width signed PTOAS scalar type
+            while preserving their bit pattern.
         span: Optional source span for debugging (auto-captured if not provided)
 
     Returns:
-        Call expression for tile select
+        Call expression for per-element tile/scalar selection
     """
     actual_span = _get_span_or_capture(span)
-    # select_mode is a mode flag interpreted by codegen, not a tile element value.
-    select_mode_expr = _normalize_const_to_dtype(select_mode, DataType.INT32, actual_span)
-    return _ir_core.create_op_call("tile.sels", [lhs, rhs, select_mode_expr], {}, actual_span)
+    scalar_expr = _normalize_sels_scalar_operand(src, scalar, actual_span)
+    return _ir_core.create_op_call("tile.sels", [mask, src, tmp, scalar_expr], {}, actual_span)
 
 
 def muls(lhs: Expr, rhs: int | float | Expr, span: Span | None = None) -> Call: