diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d63b1842c4..05186818bb 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/docs/en/dev/passes/31-memory_reuse.md b/docs/en/dev/passes/31-memory_reuse.md index 188d4b7726..4790bb15fd 100644 --- a/docs/en/dev/passes/31-memory_reuse.md +++ b/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) diff --git a/docs/en/dev/ptoas-op-status.md b/docs/en/dev/ptoas-op-status.md index d1ffa63420..cd700cdb6d 100644 --- a/docs/en/dev/ptoas-op-status.md +++ b/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 | diff --git a/docs/en/user/02-operation_reference.md b/docs/en/user/02-operation_reference.md index 5679baa350..506fad07f1 100644 --- a/docs/en/user/02-operation_reference.md +++ b/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.*`) diff --git a/docs/zh/dev/passes/31-memory_reuse.md b/docs/zh/dev/passes/31-memory_reuse.md index 66601c4f0c..51b0fb968c 100644 --- a/docs/zh/dev/passes/31-memory_reuse.md +++ b/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 间接到达的禁止操作数也被遵守（物理缓冲区解析） diff --git a/docs/zh/dev/ptoas-op-status.md b/docs/zh/dev/ptoas-op-status.md index f54ef12af9..b36ad27fd4 100644 --- a/docs/zh/dev/ptoas-op-status.md +++ b/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 复验 | diff --git a/docs/zh/user/02-operation_reference.md b/docs/zh/user/02-operation_reference.md index 1315b1b85d..ee057400df 100644 --- a/docs/zh/user/02-operation_reference.md +++ b/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.*`） diff --git a/python/pypto/debug/torch_codegen.py b/python/pypto/debug/torch_codegen.py index 7d53ff9dca..cc8b7d1fcf 100644 --- a/python/pypto/debug/torch_codegen.py +++ b/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 diff --git a/python/pypto/ir/op/tile_ops.py b/python/pypto/ir/op/tile_ops.py index bdd42c4a51..1a7db6dda1 100644 --- a/python/pypto/ir/op/tile_ops.py +++ b/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: diff --git a/python/pypto/ir/utils.py b/python/pypto/ir/utils.py index 44a6151237..8ac5ab8ea6 100644 --- a/python/pypto/ir/utils.py +++ b/python/pypto/ir/utils.py @@ -292,6 +292,7 @@ def _normalize_scalar_operand( *, fallback_int_dtype: DataType = DataType.INT32, fallback_float_dtype: DataType = DataType.FP32, + retype_constants: bool = False, ) -> _ir.Expr: """Normalize an untyped scalar constant to the paired tile/tensor element dtype. @@ -302,10 +303,13 @@ def _normalize_scalar_operand( treated as "dtype not yet decided" and re-stamped to the ``operand`` element dtype, alongside raw Python literals which carry no dtype at all. - Any constant that already carries a real dtype is left untouched -- an explicit - ``pl.const(42, pl.INT32)`` is a deliberate user annotation, not a placeholder. - A float literal paired with an integer operand keeps ``fallback_float_dtype`` - so existing promotion semantics (``int32_tensor * 2.5 -> fp32``) are preserved. + Any constant that already carries a real dtype is normally left untouched -- an + explicit ``pl.const(42, pl.INT32)`` is a deliberate user annotation, not a + placeholder. Operators whose instruction contract requires immediate constants + to match the paired operand can opt into retyping all constants. + Unless ``retype_constants`` is enabled, a float literal paired with an integer + operand keeps ``fallback_float_dtype`` so existing promotion semantics + (``int32_tensor * 2.5 -> fp32``) are preserved. Args: operand: The tile/tensor the scalar is paired with. @@ -313,6 +317,7 @@ def _normalize_scalar_operand( span: Span for any constant created here. fallback_int_dtype: Int dtype used when ``operand`` is not statically typed. fallback_float_dtype: Float dtype used when ``operand`` is not statically typed. + retype_constants: Restamp typed integer/float constants to the operand dtype. Returns: An expression whose dtype matches the operand element dtype where the rule @@ -323,16 +328,25 @@ def _normalize_scalar_operand( ``_check_not_index_scalar``) -- convert it with ``pl.cast``. """ target = _elem_dtype(operand) - value = _placeholder_value(scalar, target) + if retype_constants and isinstance(scalar, (_ir.ConstInt, _ir.ConstFloat)): + value = scalar.value + else: + value = _placeholder_value(scalar, target) if value is None: assert isinstance(scalar, _ir.Expr) # _placeholder_value returns None only for exprs return scalar # already-typed expr, kept as-is # Unknown operand type, or a float constant on an integer operand: fall back # to the literal-kind default so promotion behaviour is unchanged. - if target is None or (isinstance(value, float) and target.is_int()): + if target is None or (not retype_constants and isinstance(value, float) and target.is_int()): target = fallback_float_dtype if isinstance(value, float) else fallback_int_dtype + if retype_constants and target.is_int() and isinstance(value, float) and not value.is_integer(): + raise ValueError( + f"Cannot retype non-integral floating-point constant {value} to integer dtype " + f"{target}; use an integral value or an explicit cast" + ) + if target.is_float() or target.is_int(): return _const_at_dtype(value, target, span) diff --git a/python/pypto/language/op/tile_ops.py b/python/pypto/language/op/tile_ops.py index dd115d8de9..446130086d 100644 --- a/python/pypto/language/op/tile_ops.py +++ b/python/pypto/language/op/tile_ops.py @@ -2433,22 +2433,23 @@ def sel(mask: Tile, lhs: Tile, rhs: Tile, tmp: Tile) -> Tile: return Tile(expr=call_expr) -def sels(lhs: Tile, rhs: Tile, select_mode: int | float | Expr | Scalar) -> Tile: - """Select between two tiles based on a scalar mode. +def sels(mask: Tile, src: Tile, tmp: Tile, scalar: int | float | Expr | Scalar) -> Tile: + """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 - rhs: Source tile 1 - select_mode: Scalar select mode + mask: Predicate mask tile; encoding is target-defined + src: Source tile, selected where mask is true + tmp: Scratch tile required by TSELS + scalar: Scalar value, selected where mask is false Returns: Tile wrapping the sels operation """ - select_mode_expr = select_mode.unwrap() if isinstance(select_mode, Scalar) else select_mode - call_expr = _ir_ops.sels(lhs.unwrap(), rhs.unwrap(), select_mode_expr) + scalar_expr = scalar.unwrap() if isinstance(scalar, Scalar) else scalar + call_expr = _ir_ops.sels(mask.unwrap(), src.unwrap(), tmp.unwrap(), scalar_expr) return Tile(expr=call_expr) diff --git a/src/backend/common/pto_ops_elementwise.cpp b/src/backend/common/pto_ops_elementwise.cpp index 591526b571..a5fea24faa 100644 --- a/src/backend/common/pto_ops_elementwise.cpp +++ b/src/backend/common/pto_ops_elementwise.cpp @@ -15,6 +15,7 @@ */ #include +#include #include #include #include @@ -24,14 +25,17 @@ #include #include "pypto/backend/common/backend.h" +#include "pypto/backend/common/backend_handler.h" #include "pypto/codegen/codegen_base.h" #include "pypto/codegen/pto/pto_codegen.h" #include "pypto/core/logging.h" #include "pypto/ir/expr.h" #include "pypto/ir/kind_traits.h" +#include "pypto/ir/memref.h" #include "pypto/ir/scalar_expr.h" #include "pypto/ir/tile_view_semantics.h" #include "pypto/ir/type.h" +#include "pypto/ir/type_inference.h" #include "src/backend/common/pto_ops_internal.h" namespace pypto { @@ -79,6 +83,7 @@ static bool RequiresRowMajorLayout(std::string_view op_name) { "tile.sqrt", "tile.recip", "tile.not", + "tile.prelu", "tile.relu", // Tile x Scalar ops "tile.adds", @@ -88,6 +93,7 @@ static bool RequiresRowMajorLayout(std::string_view op_name) { "tile.fmods", "tile.maximums", "tile.lrelu", + "tile.sels", // Ternary scalar ops (Tile x Scalar x Tile) "tile.addsc", "tile.subsc", @@ -348,6 +354,132 @@ static std::string MakePrintCodegenPTO(const std::string& pto_op_name, const Cal return ""; } +static std::string MakeSelsCodegenPTO(const CallPtr& op, codegen::CodegenBase& codegen_base) { + auto& codegen = AsPto(codegen_base); + CheckArity(op, "pto.tsels", 4); + auto mask_type = As(op->args_[0]->GetType()); + auto src_type = As(op->args_[1]->GetType()); + auto tmp_type = As(op->args_[2]->GetType()); + INTERNAL_CHECK(mask_type && src_type && tmp_type); + const auto* handler = codegen.GetBackendHandler(); + const bool is_a5 = handler->GetPtoTargetArch() == "a5"; + const auto dtype = src_type->dtype_; + const bool supported_on_a2a3 = dtype == DataType::INT16 || dtype == DataType::UINT16 || + dtype == DataType::INT32 || dtype == DataType::UINT32 || + dtype == DataType::FP16 || dtype == DataType::FP32; + CHECK_SPAN(supported_on_a2a3 || is_a5, op->span_) + << "tile.sels with integer src dtype " << src_type->dtype_.ToString() + << " is only supported on the 'a5' backend; A2/A3 supports 16/32-bit integers, FP16, and FP32"; + + auto dst_var = codegen.GetCurrentResultVar(); + INTERNAL_CHECK_SPAN(dst_var, op->span_) << "Internal error: tile.sels requires an assignment target"; + auto dst_type = As(dst_var->GetType()); + INTERNAL_CHECK_SPAN(dst_type, op->span_) << "Internal error: tile.sels result must be a TileType"; + + std::vector>> operands = { + {"mask", mask_type}, {"src", src_type}, {"tmp", tmp_type}, {"dst", dst_type}}; + std::vector> regions; + regions.reserve(operands.size()); + for (const auto& [name, type] : operands) { + INTERNAL_CHECK_SPAN(type->memref_.has_value(), op->span_) + << "Internal error: tile.sels " << name << " must carry a MemRef before PTO codegen"; + regions.emplace_back(name, *type->memref_); + } + CHECK_SPAN(!ir::MemRef::MayAlias(regions[0].second, regions[3].second), op->span_) + << "tile.sels requires mask and dst to use non-overlapping memory regions"; + if (!is_a5) { + for (const size_t other : {size_t{0}, size_t{1}}) { + CHECK_SPAN(!ir::MemRef::MayAlias(regions[2].second, regions[other].second), op->span_) + << "tile.sels on A2/A3 requires tmp not to overlap mask or src, but tmp overlaps " + << regions[other].first; + } + } + return MakeNaryCodegenPTO("pto.tsels", 4, op, codegen_base); +} + +static std::string MakePreluCodegenPTO(const CallPtr& op, codegen::CodegenBase& codegen_base) { + auto& codegen = AsPto(codegen_base); + CheckArity(op, "pto.tprelu", 3); + auto src_type = As(op->args_[0]->GetType()); + auto slope_type = As(op->args_[1]->GetType()); + auto tmp_type = As(op->args_[2]->GetType()); + INTERNAL_CHECK(src_type && slope_type && tmp_type); + + auto dst_var = codegen.GetCurrentResultVar(); + INTERNAL_CHECK_SPAN(dst_var, op->span_) << "Internal error: tile.prelu requires an assignment target"; + auto dst_type = As(dst_var->GetType()); + INTERNAL_CHECK_SPAN(dst_type, op->span_) << "Internal error: tile.prelu result must be a TileType"; + + if (codegen.GetBackendHandler()->GetPtoTargetArch() == "a5") { + INTERNAL_CHECK_SPAN(src_type->memref_.has_value(), op->span_) + << "Internal error: tile.prelu src must carry a MemRef before PTO codegen"; + INTERNAL_CHECK_SPAN(slope_type->memref_.has_value(), op->span_) + << "Internal error: tile.prelu slope must carry a MemRef before PTO codegen"; + INTERNAL_CHECK_SPAN(tmp_type->memref_.has_value(), op->span_) + << "Internal error: tile.prelu tmp must carry a MemRef before PTO codegen"; + INTERNAL_CHECK_SPAN(dst_type->memref_.has_value(), op->span_) + << "Internal error: tile.prelu dst must carry a MemRef before PTO codegen"; + CHECK_SPAN(!ir::MemRef::MayAlias(*src_type->memref_, *dst_type->memref_), op->span_) + << "tile.prelu on A5 requires dst not to overlap src"; + CHECK_SPAN(!ir::MemRef::MayAlias(*slope_type->memref_, *dst_type->memref_), op->span_) + << "tile.prelu on A5 requires dst not to overlap slope"; + EmitInsOuts(codegen, "pto.tprelu", + {{codegen.GetExprAsCode(op->args_[0]), codegen.GetExprTypeAnnotation(op->args_[0])}, + {codegen.GetExprAsCode(op->args_[1]), codegen.GetExprTypeAnnotation(op->args_[1])}, + {codegen.GetExprAsCode(op->args_[2]), codegen.GetExprTypeAnnotation(op->args_[2])}}); + return ""; + } + + CHECK_SPAN(tmp_type->dtype_ == DataType::UINT8, op->args_[2]->span_) + << "tile.prelu on A2/A3 requires UINT8 tmp scratch, but got " << tmp_type->dtype_.ToString(); + const auto src_valid_shape = ir::GetValidShape(src_type); + const auto tmp_valid_shape = ir::GetValidShape(tmp_type); + const auto required_rows = ir::MakeAdd( + src_valid_shape[0], std::make_shared(1, DataType::INDEX, op->span_), op->span_); + ir::ExprPtr required_cols; + if (auto const_cols = As(src_valid_shape[1])) { + required_cols = + std::make_shared((const_cols->value_ + 7) / 8, DataType::INDEX, const_cols->span_); + } else { + required_cols = ir::MakeFloorDiv( + ir::MakeAdd(src_valid_shape[1], + std::make_shared(7, DataType::INDEX, src_valid_shape[1]->span_), + src_valid_shape[1]->span_), + std::make_shared(8, DataType::INDEX, src_valid_shape[1]->span_), + src_valid_shape[1]->span_); + } + CHECK_SPAN(ir::ProveValidExtentLessEqual(required_rows, tmp_type->shape_[0]) == ir::ProofResult::kTrue, + op->args_[2]->span_) + << "tile.prelu on A2/A3 requires UINT8 tmp physical rows >= src valid rows + 1"; + CHECK_SPAN(ir::ProveValidExtentLessEqual(required_cols, tmp_valid_shape[1]) == ir::ProofResult::kTrue, + op->args_[2]->span_) + << "tile.prelu on A2/A3 requires UINT8 tmp valid columns >= ceil(src valid columns / 8)"; + + std::vector>> operands = { + {"src", src_type}, {"slope", slope_type}, {"tmp", tmp_type}, {"dst", dst_type}}; + std::vector> regions; + regions.reserve(operands.size()); + for (const auto& [name, type] : operands) { + INTERNAL_CHECK_SPAN(type->memref_.has_value(), op->span_) + << "Internal error: tile.prelu " << name << " must carry a MemRef before PTO codegen"; + regions.emplace_back(name, *type->memref_); + } + for (size_t i = 0; i < regions.size(); ++i) { + for (size_t j = i + 1; j < regions.size(); ++j) { + CHECK_SPAN(!ir::MemRef::MayAlias(regions[i].second, regions[j].second), op->span_) + << "tile.prelu on A2/A3 requires src, slope, tmp, and dst to use pairwise non-overlapping " + "memory regions, but " + << regions[i].first << " overlaps " << regions[j].first; + } + } + + EmitInsOuts(codegen, "pto.tprelu", + {{codegen.GetExprAsCode(op->args_[0]), codegen.GetExprTypeAnnotation(op->args_[0])}, + {codegen.GetExprAsCode(op->args_[1]), codegen.GetExprTypeAnnotation(op->args_[1])}, + {codegen.GetExprAsCode(op->args_[2]), codegen.GetExprTypeAnnotation(op->args_[2])}}); + return ""; +} + struct SimpleOpEntry { const char* op_name; const char* pto_op_name; @@ -378,7 +510,6 @@ static const SimpleOpEntry kSimpleOps[] = { // Tile x Tile comparison/selection operations {"tile.maximum", "pto.tmax", 2}, {"tile.minimum", "pto.tmin", 2}, - {"tile.prelu", "pto.tprelu", 2}, // Unary operations {"tile.abs", "pto.tabs", 1}, {"tile.exp", "pto.texp", 1}, @@ -491,6 +622,24 @@ void RegisterElementwiseOps(Backend& backend, const std::unordered_set 0) return; diff --git a/src/ir/op/tile_ops/elementwise.cpp b/src/ir/op/tile_ops/elementwise.cpp index 1c67b6bbbe..f20db73ff2 100644 --- a/src/ir/op/tile_ops/elementwise.cpp +++ b/src/ir/op/tile_ops/elementwise.cpp @@ -43,6 +43,8 @@ namespace pypto { namespace ir { +constexpr int64_t kPackedPredicateBitsPerByte = 8; + static ExprPtr MakeIndexConst(int64_t value, const Span& span = Span::unknown()) { return std::make_shared(value, DataType::INDEX, span); } @@ -73,24 +75,41 @@ static bool IsTSubsDataType(DataType dtype) { dtype == DataType::FP16 || dtype == DataType::FP32 || dtype == DataType::BF16; } +static bool IsTSelsDataType(DataType dtype) { + return dtype == DataType::INT8 || dtype == DataType::UINT8 || dtype == DataType::INT16 || + dtype == DataType::UINT16 || dtype == DataType::INT32 || dtype == DataType::UINT32 || + dtype == DataType::FP16 || dtype == DataType::FP32; +} + +static DataType GetTSelsScalarDataType(DataType src_dtype) { + if (src_dtype == DataType::UINT8) return DataType::INT8; + if (src_dtype == DataType::UINT16) return DataType::INT16; + if (src_dtype == DataType::UINT32) return DataType::INT32; + return src_dtype; +} + +static bool IsTSelsMaskDataType(DataType dtype) { + return dtype == DataType::INT8 || dtype == DataType::UINT8 || dtype == DataType::INT16 || + dtype == DataType::UINT16 || dtype == DataType::INT32 || dtype == DataType::UINT32; +} + static std::shared_ptr MakePackedPredicateTileType( const std::vector& logical_shape, const std::shared_ptr& source_tile_type) { INTERNAL_CHECK(!logical_shape.empty()) << "tile.cmp/tile.cmps require a non-empty tile shape for packed predicate mask inference"; - constexpr int64_t kA2A3PredicateBitsPerByte = 8; constexpr int64_t kA2A3PredicateColAlignment = 32; const size_t col_axis = logical_shape.size() - 1; std::vector mask_shape = logical_shape; mask_shape[col_axis] = MakeRoundUpIndex( - MakeCeilDivIndex(logical_shape[col_axis], kA2A3PredicateBitsPerByte), kA2A3PredicateColAlignment); + MakeCeilDivIndex(logical_shape[col_axis], kPackedPredicateBitsPerByte), kA2A3PredicateColAlignment); auto logical_valid_shape = GetValidShape(source_tile_type); TileView tile_view; tile_view.valid_shape = logical_valid_shape; tile_view.valid_shape[col_axis] = - MakeCeilDivIndex(logical_valid_shape[col_axis], kA2A3PredicateBitsPerByte); + MakeCeilDivIndex(logical_valid_shape[col_axis], kPackedPredicateBitsPerByte); InheritTileViewLayout(tile_view, source_tile_type); return std::make_shared(mask_shape, DataType::UINT8, std::nullopt, tile_view); } @@ -842,6 +861,70 @@ REGISTER_OP("tile.xors") return DeduceTileOpXorScalarType(args, kwargs, "tile.xors"); }); +// Type deduction for tile.prelu (Src x Slope x Tmp -> Tile). +// TPRELU requires src, slope, and dst to share their dtype, physical shape, and +// valid region. Target-specific tmp and alias rules are checked during codegen. +TypePtr DeduceTilePreluType(const std::vector& args, + const std::vector>& kwargs, + const std::string& op_name) { + CHECK(args.size() == 3) << "The operator " << op_name << " requires exactly 3 arguments, but got " + << args.size(); + + auto src_type = As(args[0]->GetType()); + auto slope_type = As(args[1]->GetType()); + auto tmp_type = As(args[2]->GetType()); + CHECK_SPAN(src_type, args[0]->span_) + << "The operator " << op_name << " requires src to be a TileType, but got " + << args[0]->GetType()->TypeName(); + CHECK_SPAN(slope_type, args[1]->span_) + << "The operator " << op_name << " requires slope to be a TileType, but got " + << args[1]->GetType()->TypeName(); + CHECK_SPAN(tmp_type, args[2]->span_) + << "The operator " << op_name << " requires tmp to be a TileType, but got " + << args[2]->GetType()->TypeName(); + + CHECK_SPAN(src_type->dtype_ == DataType::FP16 || src_type->dtype_ == DataType::FP32, args[0]->span_) + << "The operator " << op_name << " requires src dtype in {FP16, FP32}, but got " + << src_type->dtype_.ToString(); + CHECK_SPAN(slope_type->dtype_ == src_type->dtype_, args[1]->span_) + << "The operator " << op_name << " requires slope dtype to match src dtype, but got " + << slope_type->dtype_.ToString() << " and " << src_type->dtype_.ToString(); + CHECK_SPAN(src_type->shape_.size() == 2, args[0]->span_) + << "The operator " << op_name << " requires a rank-2 src tile, but got rank " + << src_type->shape_.size(); + CHECK_SPAN(slope_type->shape_.size() == src_type->shape_.size(), args[1]->span_) + << "The operator " << op_name << " requires slope and src to have the same rank, but got " + << slope_type->shape_.size() << " and " << src_type->shape_.size(); + for (size_t i = 0; i < src_type->shape_.size(); ++i) { + CHECK_SPAN(DimensionsEqual(src_type->shape_[i], slope_type->shape_[i]), args[1]->span_) + << "The operator " << op_name + << " requires slope and src to have the same physical shape, but dimension " << i + << " differs; got slope shape " << FormatShape(slope_type->shape_) << " and src shape " + << FormatShape(src_type->shape_); + } + + const auto src_valid_shape = GetValidShape(src_type); + const auto slope_valid_shape = GetValidShape(slope_type); + CHECK_SPAN(slope_valid_shape.size() == src_valid_shape.size(), args[1]->span_) + << "The operator " << op_name << " requires slope and src to have the same valid_shape rank"; + for (size_t i = 0; i < src_valid_shape.size(); ++i) { + CHECK_SPAN(ProveValidExtentEqual(src_valid_shape[i], slope_valid_shape[i]) == ProofResult::kTrue, + args[1]->span_) + << "The operator " << op_name << " requires slope and src to have the same valid_shape, but " + << "dimension " << i << " differs; got slope valid_shape " << FormatShape(slope_valid_shape) + << " and src valid_shape " << FormatShape(src_valid_shape); + } + + CHECK_SPAN(tmp_type->shape_.size() == 2, args[2]->span_) + << "The operator " << op_name << " requires a rank-2 tmp tile, but got rank " + << tmp_type->shape_.size(); + + TileView tile_view; + tile_view.valid_shape = src_valid_shape; + InheritTileViewLayout(tile_view, src_type); + return std::make_shared(src_type->shape_, src_type->dtype_, std::nullopt, tile_view); +} + REGISTER_OP("tile.prelu") .set_op_category("TileOp") .set_description("Element-wise parametric ReLU of a tile with slope tile and temporary buffer") @@ -852,9 +935,10 @@ REGISTER_OP("tile.prelu") .set_input_memory(1, MemorySpace::Vec) .set_input_memory(2, MemorySpace::Vec) .set_output_memory(MemorySpace::Vec) + .not_inplace_safe() .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { - return DeduceTileOpTernaryType(args, kwargs, "tile.prelu"); + return DeduceTilePreluType(args, kwargs, "tile.prelu"); }); REGISTER_OP("tile.addc") @@ -995,55 +1079,94 @@ REGISTER_OP("tile.sel") return DeduceTileSelType(args, kwargs, "tile.sel"); }); -// Type deduction for tile.sels (Tile x Tile x Scalar -> Tile) -TypePtr DeduceTileSelScalarType(const std::vector& args, - const std::vector>& kwargs, - const std::string& op_name) { - CHECK(args.size() == 3) << "The operator " << op_name << " requires exactly 3 arguments, but got " +// Type deduction for tile.sels (Mask x Src x Tmp x Scalar -> Tile). +// dst[i,j] = mask[i,j] ? src[i,j] : scalar; the result mirrors src. +TypePtr DeduceTileSelsType(const std::vector& args, + const std::vector>& kwargs, + const std::string& op_name) { + CHECK(args.size() == 4) << "The operator " << op_name << " requires exactly 4 arguments, but got " << args.size(); - auto tile_type1 = As(args[0]->GetType()); - auto tile_type2 = As(args[1]->GetType()); - CHECK(tile_type1) << "The operator " << op_name - << " requires first argument (lhs) to be a TileType, but got " - << args[0]->GetType()->TypeName(); - CHECK(tile_type2) << "The operator " << op_name - << " requires second argument (rhs) to be a TileType, but got " - << args[1]->GetType()->TypeName(); - - CHECK(As(args[2]->GetType())) - << "The operator " << op_name << " requires third argument (select_mode) to be a ScalarType, but got " + auto mask_type = As(args[0]->GetType()); + auto src_type = As(args[1]->GetType()); + auto tmp_type = As(args[2]->GetType()); + auto scalar_type = As(args[3]->GetType()); + CHECK_SPAN(mask_type, args[0]->span_) + << "The operator " << op_name << " requires mask to be a TileType, but got " + << args[0]->GetType()->TypeName(); + CHECK_SPAN(src_type, args[1]->span_) + << "The operator " << op_name << " requires src to be a TileType, but got " + << args[1]->GetType()->TypeName(); + CHECK_SPAN(tmp_type, args[2]->span_) + << "The operator " << op_name << " requires tmp to be a TileType, but got " << args[2]->GetType()->TypeName(); + CHECK_SPAN(scalar_type, args[3]->span_) + << "The operator " << op_name << " requires scalar to be a ScalarType, but got " + << args[3]->GetType()->TypeName(); - auto result_dtype = PromoteDataTypes(tile_type1->dtype_, tile_type2->dtype_); - CHECK(result_dtype) << "The operator " << op_name << " requires compatible data types, but got " - << tile_type1->dtype_.ToString() << " and " << tile_type2->dtype_.ToString(); - - auto broadcast_result = BroadcastShapes(tile_type1->shape_, tile_type2->shape_); - CHECK(broadcast_result.success) << "The operator " << op_name << " requires compatible shapes, but got " - << FormatShape(tile_type1->shape_) << " and " - << FormatShape(tile_type2->shape_); + CHECK_SPAN(mask_type->shape_.size() == 2, args[0]->span_) + << "The operator " << op_name << " requires a rank-2 mask tile, but got rank " + << mask_type->shape_.size(); + CHECK_SPAN(IsTSelsMaskDataType(mask_type->dtype_), args[0]->span_) + << "The operator " << op_name << " requires an 8-, 16-, or 32-bit integer mask, but got " + << mask_type->dtype_.ToString(); + CHECK_SPAN(src_type->shape_.size() == 2, args[1]->span_) + << "The operator " << op_name << " requires a rank-2 src tile, but got rank " + << src_type->shape_.size(); + CHECK_SPAN(IsTSelsDataType(src_type->dtype_), args[1]->span_) + << "The operator " << op_name + << " requires src dtype in {INT8, UINT8, INT16, UINT16, INT32, UINT32, FP16, FP32}, but got " + << src_type->dtype_.ToString(); + CHECK_SPAN(tmp_type->shape_.size() == 2, args[2]->span_) + << "The operator " << op_name << " requires a rank-2 tmp tile, but got rank " + << tmp_type->shape_.size(); + const DataType expected_scalar_dtype = GetTSelsScalarDataType(src_type->dtype_); + CHECK_SPAN(scalar_type->dtype_ == expected_scalar_dtype, args[3]->span_) + << "The operator " << op_name << " requires scalar dtype " << expected_scalar_dtype.ToString() + << " for src dtype " << src_type->dtype_.ToString() << ", but got " << scalar_type->dtype_.ToString(); + + const auto mask_valid_shape = GetValidShape(mask_type); + const auto src_valid_shape = GetValidShape(src_type); + CHECK_SPAN(ProveValidExtentLessEqual(src_valid_shape[0], mask_valid_shape[0]) == ProofResult::kTrue, + args[0]->span_) + << "The operator " << op_name + << " requires mask carrier rows to cover src valid rows, but got mask valid_shape " + << FormatShape(mask_valid_shape) << " and src valid_shape " << FormatShape(src_valid_shape); + const auto required_mask_bytes = MakeCeilDivIndex(src_valid_shape[1], kPackedPredicateBitsPerByte); + const auto mask_row_bytes = MakeMul( + mask_valid_shape[1], MakeIndexConst(static_cast(mask_type->dtype_.GetByte()), args[0]->span_), + args[0]->span_); + CHECK_SPAN(ProveValidExtentLessEqual(required_mask_bytes, mask_row_bytes) == ProofResult::kTrue, + args[0]->span_) + << "The operator " << op_name + << " requires each mask carrier row to hold at least ceil(src valid columns / 8) packed bytes, " + "but got mask valid_shape " + << FormatShape(mask_valid_shape) << " with dtype " << mask_type->dtype_.ToString() + << " and src valid_shape " << FormatShape(src_valid_shape); - // TODO(YunjiQin): assumes both src tiles have the same valid_shape; may need refinement - // for cases where lhs and rhs have different valid_shape values (e.g. after broadcasting). TileView tile_view; - tile_view.valid_shape = GetValidShape(tile_type1); - InheritTileViewLayout(tile_view, tile_type1); - return std::make_shared(broadcast_result.shape, *result_dtype, std::nullopt, tile_view); + tile_view.valid_shape = src_valid_shape; + InheritTileViewLayout(tile_view, src_type); + return std::make_shared(src_type->shape_, src_type->dtype_, std::nullopt, tile_view); } REGISTER_OP("tile.sels") .set_op_category("TileOp") - .set_description("Select between two tiles based on a scalar mode. Maps to the TSELS hardware intrinsic.") - .add_argument("lhs", "Source tile 0 (TileType)") - .add_argument("rhs", "Source tile 1 (TileType)") - .add_argument("select_mode", "Scalar select mode (ScalarType)") + .set_description( + "Per-element selection between a source tile and a scalar using a predicate mask tile. " + "dst[i,j] = mask[i,j] ? src[i,j] : scalar. Maps to the TSELS hardware intrinsic.") + .add_argument("mask", "Predicate mask tile; encoding is target-defined (TileType)") + .add_argument("src", "Source tile, selected where mask is true (TileType)") + .add_argument("tmp", "Scratch tile required by TSELS (TileType)") + .add_argument("scalar", "Scalar value, selected where mask is false (ScalarType)") .set_input_memory(0, MemorySpace::Vec) .set_input_memory(1, MemorySpace::Vec) + .set_input_memory(2, MemorySpace::Vec) .set_output_memory(MemorySpace::Vec) + .forbid_output_alias(0) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { - return DeduceTileSelScalarType(args, kwargs, "tile.sels"); + return DeduceTileSelsType(args, kwargs, "tile.sels"); }); // Type deduction for tile.cmp and tile.cmps (comparison operations) diff --git a/src/ir/transforms/memory_reuse_pass.cpp b/src/ir/transforms/memory_reuse_pass.cpp index 607293c7d0..3573f9e6e5 100644 --- a/src/ir/transforms/memory_reuse_pass.cpp +++ b/src/ir/transforms/memory_reuse_pass.cpp @@ -309,6 +309,14 @@ inline bool SubtreeWritesBase(const StmtPtr& stmt, const Var* target_base) { return c.bases.count(target_base) > 0; } +bool IsA5Target() { + if (!backend::BackendConfig::IsConfigured()) return false; + const auto* ctx = PassContext::Current(); + return ctx != nullptr && ctx->GetBackendHandler()->GetPtoTargetArch() == "a5"; +} + +bool IsA5Prelu(const CallPtr& call) { return IsOp(call, "tile.prelu") && IsA5Target(); } + /// Plans top-down retypes. Produces (old Var -> new Type) map. class TopDownRetargeter { public: @@ -624,7 +632,10 @@ class TopDownRetargeter { return false; } if (FixedOutputMemoryConflicts(entry, *call, target_memory)) return false; - if (!entry.IsInplaceSafe() && CallReadsBase(*call, target->base_.get())) return false; + if (!entry.IsInplaceSafe()) { + const size_t read_arg_count = IsA5Prelu(call) ? 2 : call->args_.size(); + if (CallReadsBase(*call, target->base_.get(), read_arg_count)) return false; + } // Unconstrained: check liveness, then plan retype. (Skipped for if-phi // branch coalescing, where branch exclusivity is a stronger guarantee.) @@ -636,9 +647,11 @@ class TopDownRetargeter { /// True if any argument of the call is a TileType Var whose MemRef base /// is `target_base`. Used to detect would-be in-place execution before /// we retype the output onto the same buffer. - static bool CallReadsBase(const Call& call, const Var* target_base) { + static bool CallReadsBase(const Call& call, const Var* target_base, size_t arg_count) { SubtreeReadBaseCollector c; - for (const auto& arg : call.args_) c.VisitExpr(arg); + for (size_t i = 0; i < std::min(arg_count, call.args_.size()); ++i) { + c.VisitExpr(call.args_[i]); + } return c.bases.count(target_base) > 0; } @@ -1623,8 +1636,11 @@ class ForbidAliasCollector : public IRVisitor { } }; if (!entry.IsInplaceSafe()) { - // src != dst required: the output must not alias any input operand. - for (size_t i = 0; i < call->args_.size(); ++i) forbid_arg(i); + // Non-in-place ops forbid output aliasing active inputs. A5 TPRELU + // retains tmp (arg 2) in the ABI but does not read it, so only src + // and slope remain active. + const size_t forbidden_arg_count = IsA5Prelu(call) ? 2 : call->args_.size(); + for (size_t i = 0; i < forbidden_arg_count; ++i) forbid_arg(i); } else { for (size_t i : entry.ForbidOutputAliasArgs()) forbid_arg(i); } diff --git a/tests/st/runtime/ops/test_activation_ops.py b/tests/st/runtime/ops/test_activation_ops.py index 328a048fe6..c06f552ee2 100644 --- a/tests/st/runtime/ops/test_activation_ops.py +++ b/tests/st/runtime/ops/test_activation_ops.py @@ -18,9 +18,8 @@ the slope arg. Scope is a2a3 only (``@pytest.mark.platforms("a2a3")``); a5 coverage is a -separate PR. - -(prelu is omitted: its 3-arg DSL form mismatches codegen pto.tprelu — KNOWN_ISSUES.) +separate PR. PReLU has its own same-name test module because it also exercises +the slope-tile and scratch-buffer contracts. """ from typing import Any diff --git a/tests/st/runtime/ops/test_prelu.py b/tests/st/runtime/ops/test_prelu.py new file mode 100644 index 0000000000..a23191e0bd --- /dev/null +++ b/tests/st/runtime/ops/test_prelu.py @@ -0,0 +1,245 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""Same-name hardware tests for the three-input ``pto.tprelu`` chain.""" + +from typing import Any + +import pypto.language as pl +import pytest +import torch +from harness.core.harness import ONBOARD_PLATFORMS, DataType, PTOTestCase, TensorSpec +from pypto.runtime.runner import RunConfig + +_PL_DT = { + DataType.INT8: pl.INT8, + DataType.UINT8: pl.UINT8, + DataType.INT32: pl.INT32, + DataType.FP16: pl.FP16, + DataType.FP32: pl.FP32, +} +_TORCH_DT = {DataType.FP16: torch.float16, DataType.FP32: torch.float32} + + +def _source(m: int, n: int, dtype: DataType) -> torch.Tensor: + values = torch.arange(m * n, dtype=torch.float32).reshape(m, n).remainder(23) - 11 + return values.to(_TORCH_DT[dtype]) + + +def _slope(m: int, n: int, dtype: DataType) -> torch.Tensor: + choices = torch.tensor([-1.0, 0.0, 0.25, 1.5], dtype=torch.float32) + indices = torch.arange(m * n).reshape(m, n).remainder(len(choices)) + return choices[indices].to(_TORCH_DT[dtype]) + + +class TilePreluTestCase(PTOTestCase): + """Execute TPRELU with portable UINT8 scratch.""" + + __test__ = False + + def __init__( + self, + *, + m: int = 16, + n: int = 64, + valid_shape: tuple[int, int] | None = None, + dtype: DataType = DataType.FP32, + tmp_dtype: DataType = DataType.UINT8, + oversized_tmp: bool = False, + a5_placeholder_tmp: bool = False, + platform: str | None = None, + ): + config = RunConfig(rtol=2e-3, atol=2e-3) if dtype == DataType.FP16 else None + super().__init__(config, platform=platform) + self._m = m + self._n = n + self._valid_shape = valid_shape + self._dtype = dtype + self._tmp_dtype = tmp_dtype + self._oversized_tmp = oversized_tmp + self._a5_placeholder_tmp = a5_placeholder_tmp + + def get_name(self) -> str: + valid = self._valid_shape or (self._m, self._n) + if self._a5_placeholder_tmp: + tmp = "a5_placeholder" + elif self._oversized_tmp: + tmp = "oversized" + else: + tmp = "minimum" + tmp_dtype = DataType.INT32 if self._a5_placeholder_tmp else self._tmp_dtype + return ( + f"tile_prelu_{self._dtype.value}_{self._m}x{self._n}_v{valid[0]}x{valid[1]}_" + f"{tmp}_{tmp_dtype.value}_tmp" + ) + + def define_tensors(self) -> list[TensorSpec]: + valid_rows, valid_cols = self._valid_shape or (self._m, self._n) + if self._a5_placeholder_tmp: + tmp_shape = [1, 8] + tmp_dtype = DataType.INT32 + else: + tmp_shape = [valid_rows + (8 if self._oversized_tmp else 1), 64 if self._oversized_tmp else 32] + tmp_dtype = self._tmp_dtype + return [ + TensorSpec( + "src", + [self._m, self._n], + self._dtype, + init_value=lambda: _source(self._m, self._n, self._dtype), + ), + TensorSpec( + "slope", + [self._m, self._n], + self._dtype, + init_value=lambda: _slope(self._m, self._n, self._dtype), + ), + TensorSpec( + "tmp", + tmp_shape, + tmp_dtype, + init_value=torch.zeros, + ), + TensorSpec( + "out", + [self._m, self._n], + self._dtype, + is_output=True, + init_value=torch.zeros, + ), + ] + + def get_program(self) -> Any: + m, n = self._m, self._n + valid_shape = list(self._valid_shape or (m, n)) + valid_rows, valid_cols = valid_shape + dtype = _PL_DT[self._dtype] + if self._a5_placeholder_tmp: + tmp_rows, tmp_cols = 1, 8 + tmp_valid_shape = [1, 1] + tmp_dtype = pl.INT32 + else: + tmp_rows = valid_rows + (8 if self._oversized_tmp else 1) + tmp_cols = 64 if self._oversized_tmp else 32 + tmp_valid_shape = [valid_rows, (valid_cols + 7) // 8] + tmp_dtype = _PL_DT[self._tmp_dtype] + + @pl.program + class PreluProgram: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[m, n], dtype], + slope: pl.Tensor[[m, n], dtype], + tmp_in: pl.Tensor[[tmp_rows, tmp_cols], tmp_dtype], + out: pl.InOut[pl.Tensor[[m, n], dtype]], + ) -> pl.Tensor[[m, n], dtype]: + src_tile: pl.Tile[[m, n], dtype] = pl.load(src, [0, 0], [m, n], valid_shape=valid_shape) + slope_tile: pl.Tile[[m, n], dtype] = pl.load(slope, [0, 0], [m, n], valid_shape=valid_shape) + tmp: pl.Tile[[tmp_rows, tmp_cols], tmp_dtype] = pl.load( + tmp_in, + [0, 0], + [tmp_rows, tmp_cols], + valid_shape=tmp_valid_shape, + ) + result: pl.Tile[[m, n], dtype] = pl.tile.prelu(src_tile, slope_tile, tmp) + out = pl.store(result, [0, 0], out) + return out + + @pl.function(type=pl.FunctionType.Orchestration) + def orchestrator( + self, + src: pl.Tensor[[m, n], dtype], + slope: pl.Tensor[[m, n], dtype], + tmp_in: pl.Tensor[[tmp_rows, tmp_cols], tmp_dtype], + out: pl.InOut[pl.Tensor[[m, n], dtype]], + ) -> pl.Tensor[[m, n], dtype]: + out = self.kernel(src, slope, tmp_in, out) + return out + + return PreluProgram + + def compute_expected(self, tensors: dict[str, torch.Tensor], params=None) -> None: + src = tensors["src"] + slope = tensors["slope"] + valid_rows, valid_cols = self._valid_shape or (self._m, self._n) + expected = torch.zeros_like(tensors["out"]) + valid_src = src[:valid_rows, :valid_cols] + valid_slope = slope[:valid_rows, :valid_cols] + expected[:valid_rows, :valid_cols] = torch.where( + valid_src > 0, + valid_src, + valid_src * valid_slope, + ) + tensors["out"][:] = expected + + +class TestTilePrelu: + """TPRELU dtype, scratch, and valid-shape branches on hardware.""" + + @pytest.mark.parametrize("platform", ONBOARD_PLATFORMS) + @pytest.mark.parametrize( + "dtype", + [ + pytest.param(DataType.FP16, id="fp16"), + pytest.param(DataType.FP32, id="fp32"), + ], + ) + def test_dtypes(self, test_runner, platform, dtype): + result = test_runner.run(TilePreluTestCase(dtype=dtype, platform=platform)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.parametrize("platform", ONBOARD_PLATFORMS) + @pytest.mark.parametrize( + "valid_shape", + [ + pytest.param(None, id="full"), + pytest.param((9, 64), id="row-tail"), + pytest.param((16, 37), id="col-tail"), + pytest.param((9, 37), id="row-col-tail"), + ], + ) + def test_valid_shape(self, test_runner, platform, valid_shape): + result = test_runner.run(TilePreluTestCase(valid_shape=valid_shape, platform=platform)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.parametrize("platform", ONBOARD_PLATFORMS) + def test_oversized_tmp(self, test_runner, platform): + result = test_runner.run(TilePreluTestCase(oversized_tmp=True, platform=platform)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.platforms("a2a3") + @pytest.mark.parametrize("platform", [pytest.param("a2a3", id="a2a3")]) + def test_a2a3_uint8_tmp(self, test_runner, platform): + result = test_runner.run(TilePreluTestCase(tmp_dtype=DataType.UINT8, platform=platform)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.platforms("a5") + @pytest.mark.parametrize("platform", [pytest.param("a5", id="a5")]) + def test_a5_unused_placeholder_tmp(self, test_runner, platform): + result = test_runner.run(TilePreluTestCase(a5_placeholder_tmp=True, platform=platform)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.parametrize("platform", ONBOARD_PLATFORMS) + @pytest.mark.parametrize( + "physical_shape", + [ + pytest.param((1, 256), id="one-row-wide"), + pytest.param((64, 16), id="tall-narrow"), + ], + ) + def test_boundary_physical_shapes(self, test_runner, platform, physical_shape): + result = test_runner.run( + TilePreluTestCase(m=physical_shape[0], n=physical_shape[1], platform=platform) + ) + assert result.passed, f"Test failed: {result.error}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/st/runtime/ops/test_sels.py b/tests/st/runtime/ops/test_sels.py new file mode 100644 index 0000000000..db6777150a --- /dev/null +++ b/tests/st/runtime/ops/test_sels.py @@ -0,0 +1,440 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""Same-name hardware tests for ``pto.tsels``. + +The mask is produced by ``tile.cmps`` and feeds the canonical four-input +``tile.sels(mask, src, tmp, scalar)`` chain. Coverage includes every comparison +mode, target-supported scalar types, and row/column/combined valid tails. +""" + +from typing import Any + +import pypto.language as pl +import pytest +import torch +from harness.core.harness import ONBOARD_PLATFORMS, DataType, PTOTestCase, TensorSpec + +_PL_DT = { + DataType.INT8: pl.INT8, + DataType.UINT8: pl.UINT8, + DataType.INT16: pl.INT16, + DataType.UINT16: pl.UINT16, + DataType.INT32: pl.INT32, + DataType.UINT32: pl.UINT32, + DataType.FP16: pl.FP16, + DataType.FP32: pl.FP32, +} +_TORCH_DT = { + DataType.INT8: torch.int8, + DataType.UINT8: torch.uint8, + DataType.INT16: torch.int16, + DataType.UINT16: torch.int16, + DataType.INT32: torch.int32, + DataType.UINT32: torch.int32, + DataType.FP16: torch.float16, + DataType.FP32: torch.float32, +} +_CMP = { + 0: torch.eq, + 1: torch.ne, + 2: torch.lt, + 3: torch.le, + 4: torch.gt, + 5: torch.ge, +} +_A5_ONBOARD_PLATFORMS = [pytest.param("a5", id="a5")] +_A2A3_ONBOARD_PLATFORMS = [pytest.param("a2a3", id="a2a3")] +_UNSIGNED_DTYPES = {DataType.UINT8, DataType.UINT16, DataType.UINT32} +_INTEGER_BYTES = { + DataType.INT8: 1, + DataType.UINT8: 1, + DataType.INT16: 2, + DataType.UINT16: 2, + DataType.INT32: 4, + DataType.UINT32: 4, +} +_DTYPE_BYTES = { + **_INTEGER_BYTES, + DataType.FP16: 2, + DataType.FP32: 4, +} +_ALTERNATING_MASK_VALUE = { + DataType.INT8: -86, + DataType.UINT8: 0xAA, + DataType.INT16: -21846, + DataType.UINT16: -21846, + DataType.INT32: -1431655766, + DataType.UINT32: -1431655766, +} + + +def _source(m: int, n: int, dtype: DataType) -> torch.Tensor: + values = torch.arange(m * n, dtype=torch.int64).reshape(m, n).remainder(17) + if dtype in _UNSIGNED_DTYPES: + values += 1 << (_DTYPE_BYTES[dtype] * 8 - 1) + else: + values = values - 8 + return values.to(_TORCH_DT[dtype]) + + +class TileSelsTestCase(PTOTestCase): + """Execute one canonical TSELS branch on hardware.""" + + __test__ = False + + def __init__( + self, + *, + m: int = 16, + n: int = 64, + valid_shape: tuple[int, int] | None = None, + dtype: DataType = DataType.FP32, + cmp_type: int = 4, + threshold: int | float = 0, + scalar: int | float = -3, + tmp_dtype: DataType | None = None, + minimal_tmp: bool = False, + platform: str | None = None, + ): + super().__init__(platform=platform) + self._m = m + self._n = n + self._valid_shape = valid_shape + self._dtype = dtype + self._cmp_type = cmp_type + self._threshold = threshold + self._scalar = scalar + self._tmp_dtype = tmp_dtype or (dtype if platform == "a2a3" else DataType.UINT8) + self._minimal_tmp = minimal_tmp + + def get_name(self) -> str: + valid = self._valid_shape or (self._m, self._n) + tmp_suffix = f"_{self._tmp_dtype.value}_{'min' if self._minimal_tmp else 'default'}_tmp" + return ( + f"tile_sels_{self._dtype.value}_{self._m}x{self._n}_v{valid[0]}x{valid[1]}_" + f"cmp{self._cmp_type}{tmp_suffix}" + ) + + def define_tensors(self) -> list[TensorSpec]: + return [ + TensorSpec( + "src", + [self._m, self._n], + self._dtype, + init_value=lambda: _source(self._m, self._n, self._dtype), + ), + TensorSpec( + "out", + [self._m, self._n], + self._dtype, + is_output=True, + init_value=torch.zeros, + ), + ] + + def get_program(self) -> Any: + m, n = self._m, self._n + valid_shape = list(self._valid_shape or (m, n)) + dtype = _PL_DT[self._dtype] + cmp_type = self._cmp_type + threshold = self._threshold + scalar = self._scalar + mask_cols = ((n + 7) // 8 + 31) // 32 * 32 + tmp_dtype = _PL_DT[self._tmp_dtype] + aligned_minimum_cols = 32 // _DTYPE_BYTES[self._tmp_dtype] + tmp_rows, tmp_cols = (1, aligned_minimum_cols) if self._minimal_tmp else (1, 32) + + @pl.program + class SelsProgram: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[m, n], dtype], + out: pl.InOut[pl.Tensor[[m, n], dtype]], + ) -> pl.Tensor[[m, n], dtype]: + src_tile: pl.Tile[[m, n], dtype] = pl.load(src, [0, 0], [m, n], valid_shape=valid_shape) + mask: pl.Tile[[m, mask_cols], pl.UINT8] = pl.tile.cmps(src_tile, threshold, cmp_type=cmp_type) + tmp: pl.Tile[[tmp_rows, tmp_cols], tmp_dtype] = pl.tile.create( + [tmp_rows, tmp_cols], + dtype=tmp_dtype, + ) + result: pl.Tile[[m, n], dtype] = pl.tile.sels(mask, src_tile, tmp, scalar) + out = pl.store(result, [0, 0], out) + return out + + @pl.function(type=pl.FunctionType.Orchestration) + def orchestrator( + self, + src: pl.Tensor[[m, n], dtype], + out: pl.InOut[pl.Tensor[[m, n], dtype]], + ) -> pl.Tensor[[m, n], dtype]: + out = self.kernel(src, out) + return out + + return SelsProgram + + def compute_expected(self, tensors: dict[str, torch.Tensor], params=None) -> None: + src = tensors["src"] + valid_rows, valid_cols = self._valid_shape or (self._m, self._n) + expected = torch.zeros_like(tensors["out"]) + valid_src = src[:valid_rows, :valid_cols] + mask = _CMP[self._cmp_type](valid_src, self._threshold) + expected[:valid_rows, :valid_cols] = torch.where( + mask, + valid_src, + torch.as_tensor(self._scalar, dtype=valid_src.dtype), + ) + tensors["out"][:] = expected + + +class TileSelsMaskCarrierTestCase(PTOTestCase): + """Execute TSELS with an explicitly loaded 8/16/32-bit mask carrier.""" + + __test__ = False + + def __init__( + self, + mask_dtype: DataType, + platform: str, + *, + src_dtype: DataType = DataType.FP32, + scalar: int | float = -3.0, + ): + super().__init__(platform=platform) + self._mask_dtype = mask_dtype + self._platform = platform + self._src_dtype = src_dtype + self._scalar = scalar + + def get_name(self) -> str: + return f"tile_sels_{self._src_dtype.value}_mask_{self._mask_dtype.value}" + + def define_tensors(self) -> list[TensorSpec]: + mask_cols = 32 // _INTEGER_BYTES[self._mask_dtype] + mask_value = _ALTERNATING_MASK_VALUE[self._mask_dtype] + return [ + TensorSpec( + "src", + [2, 16], + self._src_dtype, + init_value=lambda: _source(2, 16, self._src_dtype), + ), + TensorSpec( + "mask", + [2, mask_cols], + self._mask_dtype, + init_value=lambda: torch.full( + (2, mask_cols), + mask_value, + dtype=_TORCH_DT[self._mask_dtype], + ), + ), + TensorSpec("out", [2, 16], self._src_dtype, is_output=True, init_value=torch.zeros), + ] + + def get_program(self) -> Any: + mask_dtype = _PL_DT[self._mask_dtype] + mask_cols = 32 // _INTEGER_BYTES[self._mask_dtype] + src_dtype = _PL_DT[self._src_dtype] + scalar = self._scalar + tmp_dtype = src_dtype if self._platform == "a2a3" else pl.UINT8 + tmp_cols = 32 // (_DTYPE_BYTES[self._src_dtype] if self._platform == "a2a3" else 1) + + @pl.program + class SelsMaskProgram: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[2, 16], src_dtype], + mask_in: pl.Tensor[[2, mask_cols], mask_dtype], + out: pl.InOut[pl.Tensor[[2, 16], src_dtype]], + ) -> pl.Tensor[[2, 16], src_dtype]: + src_tile: pl.Tile[[2, 16], src_dtype] = pl.load(src, [0, 0], [2, 16]) + mask: pl.Tile[[2, mask_cols], mask_dtype] = pl.load( + mask_in, + [0, 0], + [2, mask_cols], + ) + tmp: pl.Tile[[1, tmp_cols], tmp_dtype] = pl.tile.create([1, tmp_cols], dtype=tmp_dtype) + result: pl.Tile[[2, 16], src_dtype] = pl.tile.sels(mask, src_tile, tmp, scalar) + return pl.store(result, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def orchestrator( + self, + src: pl.Tensor[[2, 16], src_dtype], + mask_in: pl.Tensor[[2, mask_cols], mask_dtype], + out: pl.InOut[pl.Tensor[[2, 16], src_dtype]], + ) -> pl.Tensor[[2, 16], src_dtype]: + return self.kernel(src, mask_in, out) + + return SelsMaskProgram + + def compute_expected(self, tensors: dict[str, torch.Tensor], params=None) -> None: + scalar = self._scalar + if self._src_dtype in _UNSIGNED_DTYPES and isinstance(scalar, int): + bits = _DTYPE_BYTES[self._src_dtype] * 8 + if scalar >= 1 << (bits - 1): + scalar -= 1 << bits + expected = torch.full_like(tensors["out"], scalar) + expected[:, 1::2] = tensors["src"][:, 1::2] + tensors["out"][:] = expected + + +class TestTileSels: + """TSELS semantic branches on every onboard platform.""" + + @pytest.mark.parametrize("platform", ONBOARD_PLATFORMS) + @pytest.mark.parametrize("cmp_type", range(6), ids=("eq", "ne", "lt", "le", "gt", "ge")) + def test_comparison_modes(self, test_runner, platform, cmp_type): + result = test_runner.run(TileSelsTestCase(cmp_type=cmp_type, platform=platform)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.parametrize("platform", ONBOARD_PLATFORMS) + @pytest.mark.parametrize( + "dtype,scalar", + [ + pytest.param(DataType.FP16, -0.5, id="fp16"), + pytest.param(DataType.FP32, 1.25, id="fp32"), + ], + ) + def test_scalar_dtypes(self, test_runner, platform, dtype, scalar): + result = test_runner.run(TileSelsTestCase(dtype=dtype, scalar=scalar, platform=platform)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.parametrize("platform", _A2A3_ONBOARD_PLATFORMS) + @pytest.mark.parametrize( + "dtype,scalar", + [ + pytest.param(DataType.INT16, -7, id="int16"), + pytest.param(DataType.UINT16, 0x8007, id="uint16-high-bit"), + pytest.param(DataType.INT32, -11, id="int32"), + pytest.param(DataType.UINT32, 0x8000000B, id="uint32-high-bit"), + ], + ) + def test_a2a3_integer_scalar_dtypes(self, test_runner, platform, dtype, scalar): + result = test_runner.run( + TileSelsMaskCarrierTestCase( + DataType.UINT8, + platform, + src_dtype=dtype, + scalar=scalar, + ) + ) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.parametrize("platform", _A5_ONBOARD_PLATFORMS) + @pytest.mark.parametrize( + "dtype,scalar", + [ + pytest.param(DataType.INT8, -2, id="int8"), + pytest.param(DataType.UINT8, 0x82, id="uint8-high-bit"), + pytest.param(DataType.INT16, -7, id="int16"), + pytest.param(DataType.UINT16, 0x8007, id="uint16-high-bit"), + pytest.param(DataType.INT32, 11, id="int32"), + pytest.param(DataType.UINT32, 0x8000000B, id="uint32-high-bit"), + ], + ) + def test_a5_integer_scalar_dtypes(self, test_runner, platform, dtype, scalar): + result = test_runner.run( + TileSelsMaskCarrierTestCase( + DataType.UINT8, + platform, + src_dtype=dtype, + scalar=scalar, + ) + ) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.parametrize("platform", ONBOARD_PLATFORMS) + @pytest.mark.parametrize( + "mask_dtype", + [ + pytest.param(DataType.INT8, id="int8"), + pytest.param(DataType.UINT8, id="uint8"), + pytest.param(DataType.INT16, id="int16"), + pytest.param(DataType.UINT16, id="uint16"), + pytest.param(DataType.INT32, id="int32"), + pytest.param(DataType.UINT32, id="uint32"), + ], + ) + def test_mask_carrier_dtypes(self, test_runner, platform, mask_dtype): + result = test_runner.run(TileSelsMaskCarrierTestCase(mask_dtype, platform)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.parametrize( + "platform,src_dtype,mask_dtype", + [ + pytest.param("a2a3", DataType.INT16, DataType.INT16, id="a2a3-i16-mask-i16"), + pytest.param("a2a3", DataType.UINT32, DataType.UINT32, id="a2a3-u32-mask-u32"), + pytest.param("a5", DataType.INT8, DataType.INT16, id="a5-i8-mask-i16"), + pytest.param("a5", DataType.UINT16, DataType.UINT32, id="a5-u16-mask-u32"), + ], + ) + def test_source_and_mask_width_interactions(self, test_runner, platform, src_dtype, mask_dtype): + result = test_runner.run(TileSelsMaskCarrierTestCase(mask_dtype, platform, src_dtype=src_dtype)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.platforms("a2a3") + @pytest.mark.parametrize("platform", _A2A3_ONBOARD_PLATFORMS) + def test_a2a3_minimum_typed_tmp(self, test_runner, platform): + result = test_runner.run(TileSelsTestCase(minimal_tmp=True, platform=platform)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.platforms("a5") + @pytest.mark.parametrize("platform", _A5_ONBOARD_PLATFORMS) + def test_a5_unrestricted_tmp_placeholder(self, test_runner, platform): + result = test_runner.run( + TileSelsTestCase( + tmp_dtype=DataType.INT32, + minimal_tmp=True, + platform=platform, + ) + ) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.parametrize("platform", ONBOARD_PLATFORMS) + @pytest.mark.parametrize( + "valid_shape", + [ + pytest.param(None, id="full"), + pytest.param((9, 64), id="row-tail"), + pytest.param((16, 37), id="col-tail"), + pytest.param((9, 37), id="row-col-tail"), + ], + ) + def test_valid_shape(self, test_runner, platform, valid_shape): + result = test_runner.run(TileSelsTestCase(valid_shape=valid_shape, platform=platform)) + assert result.passed, f"Test failed: {result.error}" + + @pytest.mark.parametrize("platform", ONBOARD_PLATFORMS) + @pytest.mark.parametrize( + "physical_shape,valid_shape", + [ + pytest.param((1, 64), None, id="one-row"), + pytest.param((64, 16), None, id="tall-narrow"), + pytest.param((2, 256), None, id="packed-32-byte-boundary"), + pytest.param((2, 264), (2, 257), id="packed-33-byte-boundary"), + ], + ) + def test_boundary_physical_shapes(self, test_runner, platform, physical_shape, valid_shape): + result = test_runner.run( + TileSelsTestCase( + m=physical_shape[0], + n=physical_shape[1], + valid_shape=valid_shape, + platform=platform, + ) + ) + assert result.passed, f"Test failed: {result.error}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/ut/codegen/test_pto_codegen_ops.py b/tests/ut/codegen/test_pto_codegen_ops.py index fdf7114374..064d1fdf38 100644 --- a/tests/ut/codegen/test_pto_codegen_ops.py +++ b/tests/ut/codegen/test_pto_codegen_ops.py @@ -15,6 +15,7 @@ and verifies the generated orchestration code. """ +import re import warnings import pypto.language as pl @@ -738,6 +739,402 @@ def kernel( assert self._ins_operand_count(three_operand_line) == 3 +class TestB02SelectionAndPreluCodegen: + """Exact PTOAS forms for TSELS and TPRELU.""" + + def _generate_mlir(self, program_cls, backend_type=BackendType.Ascend910B) -> str: + backend.reset_for_testing() + backend.set_backend_type(backend_type) + + optimized = PassManager.get_strategy(OptimizationStrategy.Default).run_passes(program_cls) + funcs = list(optimized.functions.values()) + assert funcs, "Program has no functions" + single = ir.Program([funcs[0]], funcs[0].name, optimized.span) + return codegen.PTOCodegen().generate(single) + + @staticmethod + def _op_line(mlir: str, op_name: str) -> str: + line = next((line for line in mlir.splitlines() if op_name in line), "") + assert line, f"{op_name} not found in MLIR:\n{mlir}" + return line + + @staticmethod + def _ins_outs_ssas(line: str) -> tuple[list[str], list[str]]: + match = re.fullmatch(r"\s*pto\.\w+\s+ins\((.*?)\)\s+outs\((.*?)\)\s*", line) + assert match, f"expected exact PTO ins(...)/outs(...) form, got: {line}" + ins = re.findall(r"%[\w.$-]+", match.group(1).split(":", 1)[0]) + outs = re.findall(r"%[\w.$-]+", match.group(2).split(":", 1)[0]) + return ins, outs + + @staticmethod + def _assert_named_ssas(ssas: list[str], expected_names: list[str]) -> None: + assert len(ssas) == len(expected_names) + for ssa, expected_name in zip(ssas, expected_names, strict=True): + assert expected_name in ssa, f"expected {expected_name!r} SSA at this position, got {ssa!r}" + + @staticmethod + def _alloc_addr_for_named_ssa(mlir: str, expected_name: str) -> str: + lines = [ + line + for line in mlir.splitlines() + if "= pto.alloc_tile" in line and expected_name in line.split("=", 1)[0] + ] + assert len(lines) == 1, f"expected one alloc_tile for {expected_name!r}, got {lines}:\n{mlir}" + match = re.search(r"addr = (%[\w.$-]+)", lines[0]) + assert match, f"expected a baked PyPTO address in: {lines[0]}" + return match.group(1) + + def test_tsels_emits_mask_src_tmp_and_typed_scalar(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[16, 16], pl.INT32], + out: pl.Tensor[[16, 16], pl.INT32], + ) -> pl.Tensor[[16, 16], pl.INT32]: + src_tile: pl.Tile[[16, 16], pl.INT32] = pl.load(src, [0, 0], [16, 16]) + mask: pl.Tile[[16, 32], pl.UINT8] = pl.tile.cmps(src_tile, 0, cmp_type=4) + tmp: pl.Tile[[1, 32], pl.UINT8] = pl.tile.create([1, 32], dtype=pl.UINT8) + result: pl.Tile[[16, 16], pl.INT32] = pl.tile.sels(mask, src_tile, tmp, -3) + return pl.store(result, [0, 0], out) + + for backend_type in (BackendType.Ascend910B, BackendType.Ascend950): + mlir = self._generate_mlir(Prog, backend_type) + line = self._op_line(mlir, "pto.tsels") + ins, outs = self._ins_outs_ssas(line) + scalar = re.search(r"^\s*(%[\w.$-]+)\s*=\s*arith\.constant -3 : i32\s*$", mlir, re.MULTILINE) + assert scalar + self._assert_named_ssas(ins[:3], ["mask", "src_tile", "tmp"]) + assert ins[3:] == [scalar.group(1)] + self._assert_named_ssas(outs, ["result"]) + assert "i32" in line + + def test_tsels_unsigned_src_emits_signed_bit_compatible_scalar(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[16, 16], pl.UINT32], + out: pl.Tensor[[16, 16], pl.UINT32], + ) -> pl.Tensor[[16, 16], pl.UINT32]: + src_tile: pl.Tile[[16, 16], pl.UINT32] = pl.load(src, [0, 0], [16, 16]) + mask: pl.Tile[[16, 32], pl.UINT8] = pl.tile.cmps(src_tile, 0, cmp_type=4) + tmp: pl.Tile[[1, 32], pl.UINT8] = pl.tile.create([1, 32], dtype=pl.UINT8) + result: pl.Tile[[16, 16], pl.UINT32] = pl.tile.sels(mask, src_tile, tmp, 0x8000000B) + return pl.store(result, [0, 0], out) + + mlir = self._generate_mlir(Prog) + line = self._op_line(mlir, "pto.tsels") + ins, outs = self._ins_outs_ssas(line) + scalar = re.search( + r"^\s*(%[\w.$-]+)\s*=\s*arith\.constant -2147483637 : i32\s*$", + mlir, + re.MULTILINE, + ) + assert scalar + self._assert_named_ssas(ins[:3], ["mask", "src_tile", "tmp"]) + assert ins[3:] == [scalar.group(1)] + self._assert_named_ssas(outs, ["result"]) + assert "ui32" in line + assert "i32" in line + + def test_tsels_rejects_int8_on_a2a3(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[16, 16], pl.INT8], + out: pl.Tensor[[16, 16], pl.INT8], + ) -> pl.Tensor[[16, 16], pl.INT8]: + src_tile: pl.Tile[[16, 16], pl.INT8] = pl.load(src, [0, 0], [16, 16]) + mask: pl.Tile[[16, 32], pl.UINT8] = pl.tile.cmps(src_tile, 0, cmp_type=4) + tmp: pl.Tile[[1, 1], pl.UINT8] = pl.tile.create([1, 1], dtype=pl.UINT8) + result: pl.Tile[[16, 16], pl.INT8] = pl.tile.sels(mask, src_tile, tmp, -3) + return pl.store(result, [0, 0], out) + + with pytest.raises(ValueError, match="only supported on the 'a5' backend"): + self._generate_mlir(Prog) + + def test_tsels_tmp_may_alias_src_only_on_a5(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + mask_in: pl.Tensor[[2, 16], pl.INT32], + src: pl.Tensor[[2, 16], pl.INT32], + out: pl.Tensor[[2, 16], pl.INT32], + ) -> pl.Tensor[[2, 16], pl.INT32]: + mask: pl.Tile[[2, 16], pl.INT32] = pl.load(mask_in, [0, 0], [2, 16]) + src_tile: pl.Tile[[2, 16], pl.INT32] = pl.load(src, [0, 0], [2, 16]) + result: pl.Tile[[2, 16], pl.INT32] = pl.tile.sels(mask, src_tile, src_tile, -3) + return pl.store(result, [0, 0], out) + + with pytest.raises(ValueError, match="tmp overlaps src"): + self._generate_mlir(Prog, BackendType.Ascend910B) + + line = self._op_line(self._generate_mlir(Prog, BackendType.Ascend950), "pto.tsels") + ins, outs = self._ins_outs_ssas(line) + assert ins[1] == ins[2] + self._assert_named_ssas(outs, ["result"]) + + def test_tsels_tmp_may_alias_mask_only_on_a5(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + mask_in: pl.Tensor[[2, 16], pl.INT32], + src: pl.Tensor[[2, 16], pl.INT32], + out: pl.Tensor[[2, 16], pl.INT32], + ) -> pl.Tensor[[2, 16], pl.INT32]: + mask: pl.Tile[[2, 16], pl.INT32] = pl.load(mask_in, [0, 0], [2, 16]) + src_tile: pl.Tile[[2, 16], pl.INT32] = pl.load(src, [0, 0], [2, 16]) + result: pl.Tile[[2, 16], pl.INT32] = pl.tile.sels(mask, src_tile, mask, -3) + return pl.store(result, [0, 0], out) + + with pytest.raises(ValueError, match="tmp overlaps mask"): + self._generate_mlir(Prog, BackendType.Ascend910B) + + line = self._op_line(self._generate_mlir(Prog, BackendType.Ascend950), "pto.tsels") + ins, outs = self._ins_outs_ssas(line) + assert ins[0] == ins[2] + self._assert_named_ssas(outs, ["result"]) + + def test_tsels_a2a3_rejects_overlapping_tmp_view(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + mask_in: pl.Tensor[[2, 16], pl.INT32], + src: pl.Tensor[[2, 24], pl.INT32], + out: pl.Tensor[[2, 16], pl.INT32], + ) -> pl.Tensor[[2, 16], pl.INT32]: + mask: pl.Tile[[2, 16], pl.INT32] = pl.load(mask_in, [0, 0], [2, 16]) + base: pl.Tile[[2, 24], pl.INT32] = pl.load(src, [0, 0], [2, 24]) + src_view: pl.Tile[[2, 16], pl.INT32] = pl.tile.slice(base, [2, 16], [0, 0]) + tmp_view: pl.Tile[[2, 16], pl.INT32] = pl.tile.slice(base, [2, 16], [0, 8]) + result: pl.Tile[[2, 16], pl.INT32] = pl.tile.sels(mask, src_view, tmp_view, -3) + return pl.store(result, [0, 0], out) + + with pytest.raises(ValueError, match="tmp overlaps src"): + self._generate_mlir(Prog, BackendType.Ascend910B) + + assert "pto.tsels" in self._generate_mlir(Prog, BackendType.Ascend950) + + def test_tsels_tmp_may_alias_result_on_a2a3(self): + """A2/A3 consumes tmp through set_cmpmask before the first dst write.""" + + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + mask_in: pl.Tensor[[2, 16], pl.INT32], + src_in: pl.Tensor[[2, 16], pl.INT32], + tmp_in: pl.Tensor[[2, 16], pl.INT32], + out: pl.Tensor[[2, 16], pl.INT32], + ) -> pl.Tensor[[2, 16], pl.INT32]: + mask: pl.Tile[[2, 16], pl.INT32] = pl.load(mask_in, [0, 0], [2, 16]) + src: pl.Tile[[2, 16], pl.INT32] = pl.load(src_in, [0, 0], [2, 16]) + tmp: pl.Tile[[2, 16], pl.INT32] = pl.load(tmp_in, [0, 0], [2, 16]) + result: pl.Tile[[2, 16], pl.INT32] = pl.tile.sels(mask, src, tmp, -3) + keep_src_live: pl.Tile[[2, 16], pl.INT32] = pl.tile.add(src, result) + return pl.store(keep_src_live, [0, 0], out) + + mlir = self._generate_mlir(Prog, BackendType.Ascend910B) + line = self._op_line(mlir, "pto.tsels") + ins, outs = self._ins_outs_ssas(line) + self._assert_named_ssas(ins[:3], ["mask", "src", "tmp"]) + self._assert_named_ssas(outs, ["result"]) + tmp_addr = self._alloc_addr_for_named_ssa(mlir, "tmp") + result_addr = self._alloc_addr_for_named_ssa(mlir, "result") + assert tmp_addr == result_addr + + def test_tprelu_emits_target_specific_exact_operands(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[16, 16], pl.FP32], + slope: pl.Tensor[[16, 16], pl.FP32], + out: pl.Tensor[[16, 16], pl.FP32], + ) -> pl.Tensor[[16, 16], pl.FP32]: + src_tile: pl.Tile[[16, 16], pl.FP32] = pl.load(src, [0, 0], [16, 16]) + slope_tile: pl.Tile[[16, 16], pl.FP32] = pl.load(slope, [0, 0], [16, 16]) + tmp: pl.Tile[[17, 32], pl.UINT8] = pl.tile.create([17, 32], dtype=pl.UINT8) + result: pl.Tile[[16, 16], pl.FP32] = pl.tile.prelu(src_tile, slope_tile, tmp) + return pl.store(result, [0, 0], out) + + a3_line = self._op_line(self._generate_mlir(Prog, BackendType.Ascend910B), "pto.tprelu") + a3_ins, a3_outs = self._ins_outs_ssas(a3_line) + self._assert_named_ssas(a3_ins, ["src_tile", "slope_tile", "tmp"]) + self._assert_named_ssas(a3_outs, ["result"]) + + a5_line = self._op_line(self._generate_mlir(Prog, BackendType.Ascend950), "pto.tprelu") + a5_ins, a5_outs = self._ins_outs_ssas(a5_line) + self._assert_named_ssas(a5_ins, ["src_tile", "slope_tile", "tmp"]) + self._assert_named_ssas(a5_outs, ["result"]) + + def test_tprelu_signed_scratch_is_a5_only(self): + """A2/A3 requires UINT8 scratch even though the pinned verifier accepts signed i8.""" + + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[16, 16], pl.FP32], + slope: pl.Tensor[[16, 16], pl.FP32], + out: pl.Tensor[[16, 16], pl.FP32], + ) -> pl.Tensor[[16, 16], pl.FP32]: + src_tile: pl.Tile[[16, 16], pl.FP32] = pl.load(src, [0, 0], [16, 16]) + slope_tile: pl.Tile[[16, 16], pl.FP32] = pl.load(slope, [0, 0], [16, 16]) + tmp: pl.Tile[[17, 32], pl.INT8] = pl.tile.create([17, 32], dtype=pl.INT8) + result: pl.Tile[[16, 16], pl.FP32] = pl.tile.prelu(src_tile, slope_tile, tmp) + return pl.store(result, [0, 0], out) + + with pytest.raises(ValueError, match="A2/A3 requires UINT8 tmp scratch"): + self._generate_mlir(Prog, BackendType.Ascend910B) + line = self._op_line(self._generate_mlir(Prog, BackendType.Ascend950), "pto.tprelu") + ins, outs = self._ins_outs_ssas(line) + self._assert_named_ssas(ins, ["src_tile", "slope_tile", "tmp"]) + self._assert_named_ssas(outs, ["result"]) + + def test_tprelu_a3_rejects_overlapping_views_but_a5_accepts_them(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[16, 24], pl.FP32], + out: pl.Tensor[[16, 16], pl.FP32], + ) -> pl.Tensor[[16, 16], pl.FP32]: + base: pl.Tile[[16, 24], pl.FP32] = pl.load(src, [0, 0], [16, 24]) + src_view: pl.Tile[[16, 16], pl.FP32] = pl.tile.slice(base, [16, 16], [0, 0]) + slope_view: pl.Tile[[16, 16], pl.FP32] = pl.tile.slice(base, [16, 16], [0, 8]) + tmp: pl.Tile[[17, 32], pl.UINT8] = pl.tile.create([17, 32], dtype=pl.UINT8) + result: pl.Tile[[16, 16], pl.FP32] = pl.tile.prelu(src_view, slope_view, tmp) + return pl.store(result, [0, 0], out) + + with pytest.raises(ValueError, match="src overlaps slope"): + self._generate_mlir(Prog, BackendType.Ascend910B) + + mlir = self._generate_mlir(Prog, BackendType.Ascend950) + line = self._op_line(mlir, "pto.tprelu") + ins, outs = self._ins_outs_ssas(line) + subview_results = [ + subview_line.split("=", 1)[0].strip() + for subview_line in mlir.splitlines() + if "= pto.subview " in subview_line + ] + assert ins[:2] == subview_results[:2] + self._assert_named_ssas(ins[2:], ["tmp"]) + self._assert_named_ssas(outs, ["result"]) + + def test_tprelu_undersized_tmp_is_a3_only_validation(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[16, 16], pl.FP32], + slope: pl.Tensor[[16, 16], pl.FP32], + out: pl.Tensor[[16, 16], pl.FP32], + ) -> pl.Tensor[[16, 16], pl.FP32]: + src_tile: pl.Tile[[16, 16], pl.FP32] = pl.load(src, [0, 0], [16, 16]) + slope_tile: pl.Tile[[16, 16], pl.FP32] = pl.load(slope, [0, 0], [16, 16]) + tmp: pl.Tile[[1, 1], pl.UINT8] = pl.tile.create([1, 1], dtype=pl.UINT8) + result: pl.Tile[[16, 16], pl.FP32] = pl.tile.prelu(src_tile, slope_tile, tmp) + return pl.store(result, [0, 0], out) + + with pytest.raises(ValueError, match="physical rows"): + self._generate_mlir(Prog, BackendType.Ascend910B) + + line = self._op_line(self._generate_mlir(Prog, BackendType.Ascend950), "pto.tprelu") + ins, outs = self._ins_outs_ssas(line) + self._assert_named_ssas(ins, ["src_tile", "slope_tile", "tmp"]) + self._assert_named_ssas(outs, ["result"]) + + def test_tprelu_a3_rejects_tmp_with_insufficient_valid_columns(self): + """Exercise the packed-column bound independently of the row bound.""" + + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[16, 16], pl.FP32], + slope: pl.Tensor[[16, 16], pl.FP32], + tmp_in: pl.Tensor[[17, 32], pl.UINT8], + out: pl.Tensor[[16, 16], pl.FP32], + ) -> pl.Tensor[[16, 16], pl.FP32]: + src_tile: pl.Tile[[16, 16], pl.FP32] = pl.load(src, [0, 0], [16, 16]) + slope_tile: pl.Tile[[16, 16], pl.FP32] = pl.load(slope, [0, 0], [16, 16]) + tmp: pl.Tile[[17, 32], pl.UINT8] = pl.load(tmp_in, [0, 0], [17, 32], valid_shape=[17, 1]) + result: pl.Tile[[16, 16], pl.FP32] = pl.tile.prelu(src_tile, slope_tile, tmp) + return pl.store(result, [0, 0], out) + + with pytest.raises(ValueError, match="valid columns"): + self._generate_mlir(Prog, BackendType.Ascend910B) + + line = self._op_line(self._generate_mlir(Prog, BackendType.Ascend950), "pto.tprelu") + ins, outs = self._ins_outs_ssas(line) + self._assert_named_ssas(ins, ["src_tile", "slope_tile", "tmp"]) + self._assert_named_ssas(outs, ["result"]) + + def test_tprelu_a2a3_rejects_unproven_dynamic_tmp_rows(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[16, 16], pl.FP32], + slope: pl.Tensor[[16, 16], pl.FP32], + tmp_in: pl.Tensor[[17, 32], pl.UINT8], + out: pl.Tensor[[16, 16], pl.FP32], + rows: pl.Scalar[pl.INDEX], + ) -> pl.Tensor[[16, 16], pl.FP32]: + src_tile: pl.Tile[[16, 16], pl.FP32] = pl.load(src, [0, 0], [16, 16], valid_shape=[rows, 16]) + slope_tile: pl.Tile[[16, 16], pl.FP32] = pl.load( + slope, [0, 0], [16, 16], valid_shape=[rows, 16] + ) + tmp: pl.Tile[[17, 32], pl.UINT8] = pl.load(tmp_in, [0, 0], [17, 32]) + result: pl.Tile[[16, 16], pl.FP32] = pl.tile.prelu(src_tile, slope_tile, tmp) + return pl.store(result, [0, 0], out) + + with pytest.raises(ValueError, match="physical rows"): + self._generate_mlir(Prog, BackendType.Ascend910B) + + def test_tprelu_a2a3_rejects_unproven_dynamic_tmp_columns(self): + @pl.program + class Prog: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[16, 16], pl.FP32], + slope: pl.Tensor[[16, 16], pl.FP32], + tmp_in: pl.Tensor[[17, 32], pl.UINT8], + out: pl.Tensor[[16, 16], pl.FP32], + cols: pl.Scalar[pl.INDEX], + ) -> pl.Tensor[[16, 16], pl.FP32]: + src_tile: pl.Tile[[16, 16], pl.FP32] = pl.load(src, [0, 0], [16, 16], valid_shape=[16, cols]) + slope_tile: pl.Tile[[16, 16], pl.FP32] = pl.load( + slope, [0, 0], [16, 16], valid_shape=[16, cols] + ) + tmp: pl.Tile[[17, 32], pl.UINT8] = pl.load(tmp_in, [0, 0], [17, 32], valid_shape=[17, 2]) + result: pl.Tile[[16, 16], pl.FP32] = pl.tile.prelu(src_tile, slope_tile, tmp) + return pl.store(result, [0, 0], out) + + with pytest.raises(ValueError, match="valid columns"): + self._generate_mlir(Prog, BackendType.Ascend910B) + + class TestTileReadWriteOffsetCodegen: """Tests verifying tile.read/write multi-dimensional indices generate correct flat offsets.""" diff --git a/tests/ut/debug/test_torch_codegen.py b/tests/ut/debug/test_torch_codegen.py index 5cf972b4d5..8dc855d11e 100644 --- a/tests/ut/debug/test_torch_codegen.py +++ b/tests/ut/debug/test_torch_codegen.py @@ -408,6 +408,33 @@ def test_bitwise_not_reference(): assert "torch.bitwise_not(a)" in torch_codegen(func) +def test_tile_sels_and_prelu(): + """Selection and PReLU debug codegen must ignore scratch operands.""" + mask = _tile_var("mask", [16, 32], DataType.UINT8) + src = _tile_var("src", [16, 16]) + slope = _tile_var("slope", [16, 16]) + sels_tmp = _tile_var("sels_tmp", [1, 32], DataType.UINT8) + prelu_tmp = _tile_var("prelu_tmp", [17, 32], DataType.UINT8) + sels_out = _tile_var("sels_out", [16, 16]) + prelu_out = _tile_var("prelu_out", [16, 16]) + + sels_call = _op_call("tile.sels", [mask, src, sels_tmp, _float(-1.0)]) + prelu_call = _op_call("tile.prelu", [src, slope, prelu_tmp]) + body = ir.SeqStmts( + [ + ir.AssignStmt(sels_out, sels_call, _span()), + ir.AssignStmt(prelu_out, prelu_call, _span()), + ], + _span(), + ) + func = _simple_function("f", [mask, src, slope, sels_tmp, prelu_tmp], body) + + code = torch_codegen(func) + + assert "torch.where(mask, src, -1.0)" in code + assert "torch.where(src > 0, src, src * slope)" in code + + def test_tile_matmul_acc(): """tile.matmul_acc should emit (acc + torch.matmul(lhs, rhs)).""" acc = _tile_var("acc", [64, 64]) diff --git a/tests/ut/ir/operators/test_tile_ops.py b/tests/ut/ir/operators/test_tile_ops.py index a00d3c87be..d800406307 100644 --- a/tests/ut/ir/operators/test_tile_ops.py +++ b/tests/ut/ir/operators/test_tile_ops.py @@ -3877,8 +3877,8 @@ def main( slope: pl.Tile[[16, 16], pl.FP32] = pl.tile.create( [16, 16], dtype=pl.FP32, target_memory=pl.MemorySpace.Vec ) - tmp: pl.Tile[[16, 16], pl.FP32] = pl.tile.create( - [16, 16], dtype=pl.FP32, target_memory=pl.MemorySpace.Vec + tmp: pl.Tile[[17, 32], pl.UINT8] = pl.tile.create( + [17, 32], dtype=pl.UINT8, target_memory=pl.MemorySpace.Vec ) tile_c: pl.Tile[[16, 16], pl.FP32] = pl.prelu(tile_x, slope, tmp) result: pl.Tensor[[128, 128], pl.FP32] = pl.store(tile_c, [0, 0], output) @@ -3886,6 +3886,91 @@ def main( ir_str = str(Program) assert "tile.prelu" in ir_str + reparsed = pl.parse_program(ir_str) + ir.assert_structural_equal(Program, reparsed) + + def test_tile_prelu_preserves_valid_shape(self): + """TPRELU result mirrors the source physical and valid shapes.""" + src = _partial_tile([16, 16], [8, 12], name="src") + slope = _partial_tile([16, 16], [8, 12], name="slope") + span = ir.Span.unknown() + tmp = ir.Var("tmp", ir.TileType([9, 32], DataType.UINT8), span) + + result = tile.prelu(src, slope, tmp).type + + assert isinstance(result, ir.TileType) + assert [dim.value for dim in result.shape if isinstance(dim, ir.ConstInt)] == [16, 16] + assert _valid_of(result) == [8, 12] + + def test_tile_prelu_defers_target_specific_tmp_validation(self): + """IR deduction accepts a small UINT8 placeholder; A2/A3 validates it in codegen.""" + span = ir.Span.unknown() + src = ir.Var("src", ir.TileType([16, 16], DataType.FP32), span) + slope = ir.Var("slope", ir.TileType([16, 16], DataType.FP32), span) + tmp = ir.Var("tmp", ir.TileType([1, 1], DataType.UINT8), span) + + result = tile.prelu(src, slope, tmp).type + + assert isinstance(result, ir.TileType) + assert result.dtype == DataType.FP32 + + def test_tile_prelu_defers_alias_validation_to_target_codegen(self): + """Expression identity is not an alias proof, and A5 permits overlapping operands.""" + span = ir.Span.unknown() + src = ir.Var("src", ir.TileType([16, 16], DataType.FP32), span) + tmp = ir.Var("tmp", ir.TileType([17, 32], DataType.UINT8), span) + + result = tile.prelu(src, src, tmp) + + assert isinstance(result.type, ir.TileType) + + @pytest.mark.parametrize( + "slope_type,error", + [ + (ir.TileType([8, 16], DataType.FP32), "physical shape"), + (ir.TileType([16, 16], DataType.FP16), "slope dtype"), + ( + ir.TileType([16, 16], DataType.FP32, tile_view=ir.TileView(valid_shape=[8, 16])), + "valid_shape", + ), + ], + ) + def test_tile_prelu_rejects_incompatible_slope(self, slope_type, error): + """TPRELU rejects slope contracts that PTOAS cannot assemble.""" + span = ir.Span.unknown() + src = ir.Var("src", ir.TileType([16, 16], DataType.FP32), span) + slope = ir.Var("slope", slope_type, span) + tmp = ir.Var("tmp", ir.TileType([17, 32], DataType.UINT8), span) + + with pytest.raises(ValueError, match=error): + tile.prelu(src, slope, tmp) + + def test_tile_prelu_rejects_non_rank2_tmp(self): + """The target-independent ABI still requires a rank-2 tile placeholder.""" + span = ir.Span.unknown() + src = ir.Var("src", ir.TileType([16, 16], DataType.FP32), span) + slope = ir.Var("slope", ir.TileType([16, 16], DataType.FP32), span) + tmp = ir.Var("tmp", ir.TileType([16], DataType.UINT8), span) + + with pytest.raises(ValueError, match="rank-2 tmp"): + tile.prelu(src, slope, tmp) + + @pytest.mark.parametrize( + ("src_type", "tmp_type", "error"), + [ + (ir.TileType([16, 16], DataType.INT32), ir.TileType([17, 32], DataType.UINT8), "src dtype"), + (ir.TileType([256], DataType.FP32), ir.TileType([17, 32], DataType.UINT8), "rank-2 src"), + ], + ) + def test_tile_prelu_rejects_invalid_src_contract(self, src_type, tmp_type, error): + """TPRELU rejects unsupported source dtypes and ranks.""" + span = ir.Span.unknown() + src = ir.Var("src", src_type, span) + slope = ir.Var("slope", src_type, span) + tmp = ir.Var("tmp", tmp_type, span) + + with pytest.raises(ValueError, match=error): + tile.prelu(src, slope, tmp) def test_tile_not(self): """Test tile.not operator - element-wise bitwise NOT of a tile (int16/uint16 only).""" @@ -4014,7 +4099,7 @@ def main( assert "tile.lrelu" in ir_str def test_tile_sels(self): - """Test tile.sels operator - select between two tiles via integer scalar mode.""" + """Test tile.sels operator - select between a tile and scalar via mask.""" @pl.program class Program: @@ -4022,17 +4107,194 @@ class Program: def main( self, a: pl.Tensor[[128, 128], pl.FP32], - b: pl.Tensor[[128, 128], pl.FP32], output: pl.Tensor[[128, 128], pl.FP32], ) -> pl.Tensor[[128, 128], pl.FP32]: tile_a: pl.Tile[[32, 32], pl.FP32] = pl.load(a, [0, 0], [32, 32]) - tile_b: pl.Tile[[32, 32], pl.FP32] = pl.load(b, [0, 0], [32, 32]) - tile_out: pl.Tile[[32, 32], pl.FP32] = pl.sels(tile_a, tile_b, 1) + mask: pl.Tile[[32, 32], pl.UINT8] = pl.cmps(tile_a, 0.0, cmp_type=4) + tmp: pl.Tile[[1, 32], pl.UINT8] = pl.tile.create([1, 32], dtype=pl.UINT8) + tile_out: pl.Tile[[32, 32], pl.FP32] = pl.sels(mask, tile_a, tmp, -1.0) result: pl.Tensor[[128, 128], pl.FP32] = pl.store(tile_out, [0, 0], output) return result ir_str = str(Program) assert "tile.sels" in ir_str + reparsed = pl.parse_program(ir_str) + ir.assert_structural_equal(Program, reparsed) + + def test_tile_sels_preserves_src_type_and_valid_shape(self): + """TSELS result mirrors src rather than the packed mask or tmp.""" + span = ir.Span.unknown() + mask = ir.Var( + "mask", + ir.TileType([16, 32], DataType.UINT8, tile_view=ir.TileView(valid_shape=[8, 2])), + span, + ) + src = _partial_tile([16, 16], [8, 12], name="src") + tmp = ir.Var("tmp", ir.TileType([1, 32], DataType.UINT8), span) + + result = tile.sels(mask, src, tmp, -2.5).type + + assert isinstance(result, ir.TileType) + assert result.dtype == DataType.FP32 + assert [dim.value for dim in result.shape if isinstance(dim, ir.ConstInt)] == [16, 16] + assert _valid_of(result) == [8, 12] + + def test_tile_sels_retypes_constant_to_src_dtype(self): + """A parser-produced constant adopts the selected source dtype.""" + span = ir.Span.unknown() + mask = ir.Var("mask", ir.TileType([16, 32], DataType.UINT8), span) + src = ir.Var("src", ir.TileType([16, 16], DataType.FP16), span) + tmp = ir.Var("tmp", ir.TileType([1, 32], DataType.UINT8), span) + scalar = ir.ConstFloat(-1.0, DataType.FP32, span) + + call = tile.sels(mask, src, tmp, scalar) + + assert _operand_dtype(call.args[3]) == DataType.FP16 + + def test_tile_sels_rejects_fractional_constant_for_integer_src(self): + """Retyping a scalar must not silently truncate a fractional value.""" + span = ir.Span.unknown() + mask = ir.Var("mask", ir.TileType([16, 32], DataType.UINT8), span) + src = ir.Var("src", ir.TileType([16, 16], DataType.INT32), span) + tmp = ir.Var("tmp", ir.TileType([1, 32], DataType.UINT8), span) + + with pytest.raises(ValueError, match="non-integral"): + tile.sels(mask, src, tmp, -1.5) + + def test_tile_sels_rejects_scalar_dtype_mismatch(self): + """A non-constant scalar expression must match the selected source dtype.""" + span = ir.Span.unknown() + mask = ir.Var("mask", ir.TileType([16, 32], DataType.UINT8), span) + src = ir.Var("src", ir.TileType([16, 16], DataType.FP16), span) + tmp = ir.Var("tmp", ir.TileType([1, 32], DataType.UINT8), span) + scalar = ir.Var("scalar", ir.ScalarType(DataType.FP32), span) + + with pytest.raises(ValueError, match="scalar dtype"): + tile.sels(mask, src, tmp, scalar) + + @pytest.mark.parametrize( + "mask_type,error", + [ + (ir.TileType([16, 32], DataType.FP32), "integer mask"), + (ir.TileType([32], DataType.UINT8), "rank-2 mask"), + ], + ) + def test_tile_sels_rejects_invalid_mask(self, mask_type, error): + """TSELS requires a rank-2 packed integer predicate tile.""" + span = ir.Span.unknown() + mask = ir.Var("mask", mask_type, span) + src = ir.Var("src", ir.TileType([16, 16], DataType.FP32), span) + tmp = ir.Var("tmp", ir.TileType([1, 32], DataType.UINT8), span) + + with pytest.raises(ValueError, match=error): + tile.sels(mask, src, tmp, -1.0) + + @pytest.mark.parametrize( + "mask_type,error", + [ + ( + ir.TileType([7, 64], DataType.UINT8), + "mask carrier rows", + ), + ( + ir.TileType([8, 32], DataType.UINT8), + "each mask carrier row", + ), + ], + ) + def test_tile_sels_rejects_mask_too_small_for_src_valid_shape(self, mask_type, error): + """A packed mask must cover every valid source row and column bit.""" + span = ir.Span.unknown() + mask = ir.Var("mask", mask_type, span) + src = ir.Var("src", ir.TileType([8, 257], DataType.FP32), span) + tmp = ir.Var("tmp", ir.TileType([1, 32], DataType.UINT8), span) + + with pytest.raises(ValueError, match=error): + tile.sels(mask, src, tmp, -1.0) + + def test_tile_sels_accepts_provable_dynamic_mask_coverage(self): + """Shared symbolic rows and the exact packed-byte expression are provably safe.""" + span = ir.Span.unknown() + valid_rows = ir.Var("valid_rows", ir.ScalarType(DataType.INDEX), span) + valid_cols = ir.Var("valid_cols", ir.ScalarType(DataType.INDEX), span) + packed_cols = (valid_cols + 7) // 8 + mask = ir.Var( + "mask", + ir.TileType( + [16, 64], + DataType.UINT8, + tile_view=ir.TileView(valid_shape=[valid_rows, packed_cols]), + ), + span, + ) + src = ir.Var( + "src", + ir.TileType( + [16, 512], + DataType.FP32, + tile_view=ir.TileView(valid_shape=[valid_rows, valid_cols]), + ), + span, + ) + tmp = ir.Var("tmp", ir.TileType([1, 32], DataType.UINT8), span) + + result = tile.sels(mask, src, tmp, -1.0) + + assert isinstance(result.type, ir.TileType) + + @pytest.mark.parametrize( + ("mask_dtype", "physical_cols", "valid_cols", "accepted"), + [ + (DataType.INT16, 32, 16, False), + (DataType.INT16, 32, 17, True), + (DataType.UINT16, 32, 16, False), + (DataType.UINT16, 32, 17, True), + (DataType.INT32, 16, 8, False), + (DataType.INT32, 16, 9, True), + (DataType.UINT32, 16, 8, False), + (DataType.UINT32, 16, 9, True), + ], + ) + def test_tile_sels_packed_mask_capacity_respects_carrier_width( + self, mask_dtype, physical_cols, valid_cols, accepted + ): + """Packed-mask capacity is measured in bytes for every integer carrier.""" + span = ir.Span.unknown() + mask = ir.Var( + "mask", + ir.TileType( + [2, physical_cols], + mask_dtype, + tile_view=ir.TileView(valid_shape=[2, valid_cols]), + ), + span, + ) + src = ir.Var("src", ir.TileType([2, 257], DataType.FP32), span) + tmp = ir.Var("tmp", ir.TileType([1, 32], DataType.UINT8), span) + + if accepted: + assert isinstance(tile.sels(mask, src, tmp, -1.0).type, ir.TileType) + else: + with pytest.raises(ValueError, match="each mask carrier row"): + tile.sels(mask, src, tmp, -1.0) + + @pytest.mark.parametrize( + ("src_type", "tmp_type", "error"), + [ + (ir.TileType([16, 16], DataType.BF16), ir.TileType([1, 32], DataType.UINT8), "src dtype"), + (ir.TileType([256], DataType.FP32), ir.TileType([1, 32], DataType.UINT8), "rank-2 src"), + (ir.TileType([16, 16], DataType.FP32), ir.TileType([32], DataType.UINT8), "rank-2 tmp"), + ], + ) + def test_tile_sels_rejects_invalid_src_and_tmp_contract(self, src_type, tmp_type, error): + """TSELS rejects unsupported source dtypes and non-2D operands.""" + span = ir.Span.unknown() + mask = ir.Var("mask", ir.TileType([16, 32], DataType.UINT8), span) + src = ir.Var("src", src_type, span) + tmp = ir.Var("tmp", tmp_type, span) + + with pytest.raises(ValueError, match=error): + tile.sels(mask, src, tmp, -1.0) def test_tile_sel(self): """Test tile.sel operator - per-element selection between two tiles via mask tile.""" @@ -4246,13 +4508,44 @@ def test_lrelu_slope_stays_fp32(self): call = tile.lrelu(ir.Var("t", ir.TileType([32, 32], DataType.FP32), ir.Span.unknown()), 1) assert _operand_dtype(call.args[1]) == DataType.FP32 - def test_sels_mode_stays_int32(self): - """tile.sels keeps its select-mode flag at INT32 and never index.""" - span = ir.Span.unknown() - lhs = ir.Var("a", ir.TileType([32, 32], DataType.FP32), span) - rhs = ir.Var("b", ir.TileType([32, 32], DataType.FP32), span) - call = tile.sels(lhs, rhs, 1) - assert _operand_dtype(call.args[2]) == DataType.INT32 + @pytest.mark.parametrize( + "dtype,scalar,expected_dtype,expected_value", + [ + (DataType.INT8, -2, DataType.INT8, -2), + (DataType.UINT8, 0x82, DataType.INT8, -126), + (DataType.INT16, -3, DataType.INT16, -3), + (DataType.UINT16, 0x8007, DataType.INT16, -32761), + (DataType.INT32, 7, DataType.INT32, 7), + (DataType.UINT32, 0x8000000B, DataType.INT32, -2147483637), + (DataType.FP16, -0.5, DataType.FP16, -0.5), + (DataType.FP32, 1.25, DataType.FP32, 1.25), + ], + ) + def test_sels_scalar_adopts_ptoas_dtype(self, dtype, scalar, expected_dtype, expected_value): + """tile.sels uses signed bit-compatible scalars for unsigned sources.""" + span = ir.Span.unknown() + mask = ir.Var("mask", ir.TileType([32, 32], DataType.UINT8), span) + src = ir.Var("src", ir.TileType([32, 32], dtype), span) + tmp = ir.Var("tmp", ir.TileType([1, 32], DataType.UINT8), span) + call = tile.sels(mask, src, tmp, scalar) + scalar_arg = call.args[3] + assert isinstance(scalar_arg, (ir.ConstInt, ir.ConstFloat)) + assert _operand_dtype(scalar_arg) == expected_dtype + assert scalar_arg.value == expected_value + + def test_sels_unsigned_src_accepts_only_signed_same_width_scalar_expr(self): + """PTOAS scalar operands are signed even when the selected tile is unsigned.""" + span = ir.Span.unknown() + mask = ir.Var("mask", ir.TileType([32, 32], DataType.UINT8), span) + src = ir.Var("src", ir.TileType([32, 32], DataType.UINT16), span) + tmp = ir.Var("tmp", ir.TileType([1, 32], DataType.UINT8), span) + + call = tile.sels(mask, src, tmp, ir.Var("signed_scalar", ir.ScalarType(DataType.INT16), span)) + assert isinstance(call.type, ir.TileType) + assert call.type.dtype == DataType.UINT16 + + with pytest.raises(ValueError, match="requires scalar dtype int16 for src dtype uint16"): + tile.sels(mask, src, tmp, ir.Var("unsigned_scalar", ir.ScalarType(DataType.UINT16), span)) class TestTileLoadOp: diff --git a/tests/ut/ir/transforms/test_memory_reuse.py b/tests/ut/ir/transforms/test_memory_reuse.py index 93776597df..e9e415a787 100644 --- a/tests/ut/ir/transforms/test_memory_reuse.py +++ b/tests/ut/ir/transforms/test_memory_reuse.py @@ -3890,6 +3890,113 @@ def main( f"tile.sel output must not alias its tmp buffer, but both bind to {bases['dst']}" ) + @pytest.mark.parametrize("backend_type", [BackendType.Ascend910B, BackendType.Ascend950]) + def test_sels_output_may_reuse_dead_tmp(self, backend_type): + """TSELS consumes tmp before dst writes on A2/A3; A5 leaves tmp unread.""" + + @pl.program + class Before: + @pl.function + def main( + self, + a: pl.Tensor[[16, 16], pl.FP32], + b: pl.Tensor[[16, 16], pl.FP32], + tmp_in: pl.Tensor[[16, 16], pl.FP32], + out: pl.Out[pl.Tensor[[16, 16], pl.FP32]], + ) -> pl.Tensor[[16, 16], pl.FP32]: + t0: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.load(a, [0, 0], [16, 16]) + dead: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.add(t0, t0) + src: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.load(b, [0, 0], [16, 16]) + mask: pl.Tile[[16, 32], pl.UINT8, pl.MemorySpace.Vec] = pl.cmps(dead, 0.0, cmp_type=4) + tmp: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.load(tmp_in, [0, 0], [16, 16]) + dst: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.sels(mask, src, tmp, -1.0) + keep_src_live: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.add(src, dst) + res: pl.Tensor[[16, 16], pl.FP32] = pl.store(keep_src_live, [0, 0], out) + return res + + backend.reset_for_testing() + backend.set_backend_type(backend_type) + try: + After = _run_pipeline(Before) + finally: + backend.reset_for_testing() + + bases = _collect_tile_memref_bases(After) + for name in ("dst", "src", "mask", "tmp"): + assert name in bases, f"Expected {name} in After IR; got bases: {bases}" + assert bases["dst"] == bases["tmp"] + assert bases["dst"] != bases["src"] + assert bases["dst"] != bases["mask"] + + def test_prelu_output_does_not_alias_any_input(self): + """A2/A3 TPRELU reads src, slope, and tmp while writing dst.""" + + @pl.program + class Before: + @pl.function + def main( + self, + src_in: pl.Tensor[[16, 16], pl.FP32], + slope_in: pl.Tensor[[16, 16], pl.FP32], + tmp_in: pl.Tensor[[17, 32], pl.UINT8], + out: pl.Out[pl.Tensor[[16, 16], pl.FP32]], + ) -> pl.Tensor[[16, 16], pl.FP32]: + src: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.load(src_in, [0, 0], [16, 16]) + slope: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.load(slope_in, [0, 0], [16, 16]) + tmp: pl.Tile[[17, 32], pl.UINT8, pl.MemorySpace.Vec] = pl.load(tmp_in, [0, 0], [17, 32]) + dst: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.prelu(src, slope, tmp) + res: pl.Tensor[[16, 16], pl.FP32] = pl.store(dst, [0, 0], out) + return res + + backend.reset_for_testing() + backend.set_backend_type(BackendType.Ascend910B) + try: + After = _run_pipeline(Before) + finally: + backend.reset_for_testing() + bases = _collect_tile_memref_bases(After) + for name in ("dst", "src", "slope", "tmp"): + assert name in bases, f"Expected {name} in After IR; got bases: {bases}" + assert bases["dst"] != bases["src"] + assert bases["dst"] != bases["slope"] + assert bases["dst"] != bases["tmp"] + + def test_a5_prelu_output_may_reuse_dead_tmp(self): + """A5 retains unread TPRELU tmp, so dst may reuse it while src/slope stay live.""" + + @pl.program + class Before: + @pl.function + def main( + self, + src_in: pl.Tensor[[16, 16], pl.FP32], + slope_in: pl.Tensor[[16, 16], pl.FP32], + tmp_in: pl.Tensor[[16, 16], pl.FP32], + out: pl.Out[pl.Tensor[[16, 16], pl.FP32]], + ) -> pl.Tensor[[16, 16], pl.FP32]: + src: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.load(src_in, [0, 0], [16, 16]) + slope: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.load(slope_in, [0, 0], [16, 16]) + tmp: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.load(tmp_in, [0, 0], [16, 16]) + dst: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.prelu(src, slope, tmp) + live_inputs: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.add(src, slope) + result: pl.Tile[[16, 16], pl.FP32, pl.MemorySpace.Vec] = pl.add(dst, live_inputs) + res: pl.Tensor[[16, 16], pl.FP32] = pl.store(result, [0, 0], out) + return res + + backend.reset_for_testing() + backend.set_backend_type(BackendType.Ascend950) + try: + After = _run_pipeline(Before) + finally: + backend.reset_for_testing() + + bases = _collect_tile_memref_bases(After) + for name in ("dst", "src", "slope", "tmp"): + assert name in bases, f"Expected {name} in After IR; got bases: {bases}" + assert bases["dst"] == bases["tmp"] + assert bases["dst"] != bases["src"] + assert bases["dst"] != bases["slope"] + def test_row_sum_output_does_not_alias_input_or_tmp(self): """A row reduction output must not share a buffer with its input or tmp.