From 65dc73adf8bf9a0564e6d7a232d07e200c64ae5f Mon Sep 17 00:00:00 2001 From: ChenShenAi Date: Wed, 29 Jul 2026 05:04:11 +0800 Subject: [PATCH 1/5] feat: add B11 indexed concat and exact fillpad --- docs/en/dev/codegen/00-pto_codegen.md | 2 + docs/en/dev/ptoas-op-status.md | 4 +- docs/zh/dev/codegen/00-pto_codegen.md | 2 + docs/zh/dev/ptoas-op-status.md | 4 +- python/pypto/ir/op/tile_ops.py | 15 +++ python/pypto/language/__init__.py | 2 + python/pypto/language/op/__init__.py | 2 + python/pypto/language/op/tile_ops.py | 10 ++ src/backend/common/pto_ops_elementwise.cpp | 16 ++- src/ir/op/tile_ops/transform.cpp | 59 +++++++++++ tests/st/runtime/ops/test_concat_idx.py | 110 +++++++++++++++++++++ tests/ut/codegen/test_pto_codegen.py | 6 +- tests/ut/ir/operators/test_concat_idx.py | 52 ++++++++++ 13 files changed, 276 insertions(+), 8 deletions(-) create mode 100644 tests/st/runtime/ops/test_concat_idx.py create mode 100644 tests/ut/ir/operators/test_concat_idx.py diff --git a/docs/en/dev/codegen/00-pto_codegen.md b/docs/en/dev/codegen/00-pto_codegen.md index fb793cf4c2..4c04f69f9f 100644 --- a/docs/en/dev/codegen/00-pto_codegen.md +++ b/docs/en/dev/codegen/00-pto_codegen.md @@ -143,6 +143,8 @@ print(pto_code) | `tile.add(a, b, c)` | `pto.taddc` (3-operand add) | | `tile.adds(tile, scalar)` | `pto.tadds` (tile + scalar) | | `tile.fillpad_expand(src, shape)` | `pto.tfillpad_expand ins(%src) outs(%dst)` (the `shape` tuple is type-deduction only; the larger `dst` and its pad come from the result type) | +| `tile.fillpad_inplace(src, pad_value=...)` | `pto.tfillpad_inplace` (shared source/result storage) | +| `tile.concat_idx(src0, src1, idx0, idx1, dst)` | `pto.tconcatidx` (reuses `dst`) | **`tile.slice` / `tile.assemble` lowering details.** Both ops are lowered through `pto.subview`, which is a pure view alias of the source tile (no diff --git a/docs/en/dev/ptoas-op-status.md b/docs/en/dev/ptoas-op-status.md index 6de3c36be3..9456753e3b 100644 --- a/docs/en/dev/ptoas-op-status.md +++ b/docs/en/dev/ptoas-op-status.md @@ -180,7 +180,7 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT. | pto.tshrs | TSHRS | tile | ✅ | ✅ | ❌ | ❌ | — | path exists; historical ISA/semantic issue requires revalidation against the current pin | | **Data Rearrangement (15)** | | | | | | | | | | pto.tconcat | TCONCAT | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path | +| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ✅ | ❌ | ❌ | — | indexed destination-reuse path and ST added; hardware verification pending | | pto.tgather | TGATHER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tgatherb | TGATHERB | tile | ✅ | ❌ | ❌ | ❌ | — | backend hook exists; IR/Python frontend and ST are missing | | pto.tscatter | TSCATTER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | @@ -191,7 +191,7 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT. | pto.textract | TEXTRACT | tile | ✅ | ✅ | ❌ | ✅ | — | | | pto.tfillpad | TFILLPAD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tfillpad_expand | TFILLPAD_EXPAND | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tfillpad_inplace | TFILLPAD_INPLACE | tile | ✅ | ✅ | ❌ | ❌ | — | the current codegen emits `pto.tfillpad` | +| pto.tfillpad_inplace | TFILLPAD_INPLACE | tile | ✅ | ✅ | ❌ | ❌ | — | corrected to exact same-name emission; hardware revalidation pending | | pto.textract_fp | TEXTRACT_FP / TEXTRACT | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path | | pto.tinsert_fp | TINSERT_FP / TINSERT | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path | | **Sorting (2)** | | | | | | | | | diff --git a/docs/zh/dev/codegen/00-pto_codegen.md b/docs/zh/dev/codegen/00-pto_codegen.md index 4536c06e0e..eda9f2d23b 100644 --- a/docs/zh/dev/codegen/00-pto_codegen.md +++ b/docs/zh/dev/codegen/00-pto_codegen.md @@ -141,6 +141,8 @@ print(pto_code) | `tile.add(a, b, c)` | `pto.taddc` (三操作数加法) | | `tile.adds(tile, scalar)` | `pto.tadds` (Tile + 标量) | | `tile.fillpad_expand(src, shape)` | `pto.tfillpad_expand ins(%src) outs(%dst)`（`shape` 元组仅用于类型推导；更大的 `dst` 及其 pad 来自结果类型） | +| `tile.fillpad_inplace(src, pad_value=...)` | `pto.tfillpad_inplace`（源与结果共享存储） | +| `tile.concat_idx(src0, src1, idx0, idx1, dst)` | `pto.tconcatidx`（复用 `dst`） | **`tile.slice` / `tile.assemble` 下沉细节。** 两个 op 都通过 `pto.subview` 下沉，它是源 tile 的纯视图别名（不搬数据，也不会额外发 `pto.alloc_tile`）。 diff --git a/docs/zh/dev/ptoas-op-status.md b/docs/zh/dev/ptoas-op-status.md index c982ff35fe..7a34402bdc 100644 --- a/docs/zh/dev/ptoas-op-status.md +++ b/docs/zh/dev/ptoas-op-status.md @@ -166,7 +166,7 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO | pto.tshrs | TSHRS | tile | ✅ | ✅ | ❌ | ❌ | — | 已有链路；历史 ISA/语义问题，需按当前 pin 复验 | | **数据重排（15）** | | | | | | | | | | pto.tconcat | TCONCAT | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 | +| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ✅ | ❌ | ❌ | — | 已增加 indexed 目标复用路径与 ST；真机待验证 | | pto.tgather | TGATHER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tgatherb | TGATHERB | tile | ✅ | ❌ | ❌ | ❌ | — | 已有 backend hook，缺 IR/Python 前端与 ST | | pto.tscatter | TSCATTER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | @@ -177,7 +177,7 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO | pto.textract | TEXTRACT | tile | ✅ | ✅ | ❌ | ✅ | — | | | pto.tfillpad | TFILLPAD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tfillpad_expand | TFILLPAD_EXPAND | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tfillpad_inplace | TFILLPAD_INPLACE | tile | ✅ | ✅ | ❌ | ❌ | — | 当前 codegen 发 `pto.tfillpad` | +| pto.tfillpad_inplace | TFILLPAD_INPLACE | tile | ✅ | ✅ | ❌ | ❌ | — | 已修正为同名精确发射；真机复验待完成 | | pto.textract_fp | TEXTRACT_FP / TEXTRACT | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 | | pto.tinsert_fp | TINSERT_FP / TINSERT | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 | | **排序（2）** | | | | | | | | | diff --git a/python/pypto/ir/op/tile_ops.py b/python/pypto/ir/op/tile_ops.py index de166b9c76..0816a44f52 100644 --- a/python/pypto/ir/op/tile_ops.py +++ b/python/pypto/ir/op/tile_ops.py @@ -438,6 +438,21 @@ def concat( return _ir_core.create_op_call("tile.concat", [src0, src1], {}, actual_span) +def concat_idx( + src0: Expr, + src1: Expr, + src0_idx: Expr, + src1_idx: Expr, + dst: Expr, + span: Span | None = None, +) -> Call: + """Indexed per-row concatenation into ``dst``.""" + actual_span = _get_span_or_capture(span) + return _ir_core.create_op_call( + "tile.concat_idx", [src0, src1, src0_idx, src1_idx, dst], {}, actual_span + ) + + def transpose_view( tile: Expr, span: Span | None = None, diff --git a/python/pypto/language/__init__.py b/python/pypto/language/__init__.py index 71b90da86b..c07c13f9ea 100644 --- a/python/pypto/language/__init__.py +++ b/python/pypto/language/__init__.py @@ -123,6 +123,7 @@ def scalar_func(x: pl.Scalar[pl.FP32]) -> pl.Scalar[pl.FP32]: and_, ands, cmps, + concat_idx, create_tile, gemv, gemv_acc, @@ -325,6 +326,7 @@ def scalar_func(x: pl.Scalar[pl.FP32]) -> pl.Scalar[pl.FP32]: "log", "cast", "concat", + "concat_idx", "reshape", "reinterpret_view", "transpose", diff --git a/python/pypto/language/op/__init__.py b/python/pypto/language/op/__init__.py index 6c42ee6087..2fe9339d00 100644 --- a/python/pypto/language/op/__init__.py +++ b/python/pypto/language/op/__init__.py @@ -76,6 +76,7 @@ and_, ands, cmps, + concat_idx, gemv, gemv_acc, gemv_bias, @@ -242,6 +243,7 @@ "recip", "write", "concat", + "concat_idx", "batch_matmul", # Promoted tile-only "create_tile", diff --git a/python/pypto/language/op/tile_ops.py b/python/pypto/language/op/tile_ops.py index 4a995c5ceb..7697afe15a 100644 --- a/python/pypto/language/op/tile_ops.py +++ b/python/pypto/language/op/tile_ops.py @@ -33,6 +33,7 @@ "extract", "scatter_update", "concat", + "concat_idx", "move", "aiv_shard", "aic_gather", @@ -1894,6 +1895,15 @@ def transpose(tile: Tile, axis1: int, axis2: int, tmp_tile: Tile | None = None) return Tile(expr=call_expr) +def concat_idx(src0: Tile, src1: Tile, src0_idx: Tile, src1_idx: Tile, dst: Tile) -> Tile: + """Indexed per-row concatenation into a destination tile.""" + return Tile( + expr=_ir_ops.concat_idx( + src0.unwrap(), src1.unwrap(), src0_idx.unwrap(), src1_idx.unwrap(), dst.unwrap() + ) + ) + + def transpose_view(tile: Tile) -> Tile: """Zero-copy fractal-layout reinterpretation (NZ<->ZN) of a tile. diff --git a/src/backend/common/pto_ops_elementwise.cpp b/src/backend/common/pto_ops_elementwise.cpp index 70f8c99ebe..acdaa6a9b7 100644 --- a/src/backend/common/pto_ops_elementwise.cpp +++ b/src/backend/common/pto_ops_elementwise.cpp @@ -144,6 +144,17 @@ static std::string MakeNaryCodegenPTO(const std::string& pto_op_name, size_t ari return ""; } +static std::string MakeConcatIdxCodegenPTO(const CallPtr& op, codegen::CodegenBase& codegen_base) { + auto& codegen = AsPto(codegen_base); + CheckArity(op, "pto.tconcatidx", 5); + std::vector> inputs; + for (size_t i = 0; i < 4; ++i) { + inputs.emplace_back(codegen.GetExprAsCode(op->args_[i]), codegen.GetExprTypeAnnotation(op->args_[i])); + } + EmitInsOuts(codegen, "pto.tconcatidx", inputs); + return ""; +} + static std::string MakeTileSelCodegenPTO(const CallPtr& op, codegen::CodegenBase& codegen_base) { auto& codegen = AsPto(codegen_base); CheckArity(op, "pto.tsel", 4); @@ -439,7 +450,7 @@ static const SimpleOpEntry kSimpleOps[] = { // Padding operations {"tile.fillpad", "pto.tfillpad", 1}, // Inplace variant: set_output_reuses_input(0) makes src/dst share UB addr. - {"tile.fillpad_inplace", "pto.tfillpad", 1}, + {"tile.fillpad_inplace", "pto.tfillpad_inplace", 1}, // Matrix multiplication operations (PipeType::M → CUBE/AIC core) {"tile.matmul", "pto.tmatmul", 2}, {"tile.matmul_mx", "pto.tmatmul.mx", 4}, @@ -496,6 +507,9 @@ void RegisterElementwiseOps(Backend& backend, const std::unordered_set 0) return; backend.RegisterOp(op_name).f_codegen(std::move(fn)); }; + reg("tile.concat_idx", [](const CallPtr& op, codegen::CodegenBase& codegen) { + return MakeConcatIdxCodegenPTO(op, codegen); + }); auto register_precision_op = [&](const char* op_name, const char* pto_op_name, size_t arity, const char* attr_kind) { diff --git a/src/ir/op/tile_ops/transform.cpp b/src/ir/op/tile_ops/transform.cpp index 171be85621..a3933d2f53 100644 --- a/src/ir/op/tile_ops/transform.cpp +++ b/src/ir/op/tile_ops/transform.cpp @@ -919,6 +919,65 @@ REGISTER_OP("tile.concat") return DeduceTileConcatType(args, kwargs); }); +static TypePtr DeduceTileConcatIdxType(const std::vector& args) { + CHECK(args.size() == 5) + << "tile.concat_idx requires 5 arguments (src0, src1, src0_idx, src1_idx, dst), got " + << args.size(); + auto src0 = As(args[0]->GetType()); + auto src1 = As(args[1]->GetType()); + auto idx0 = As(args[2]->GetType()); + auto idx1 = As(args[3]->GetType()); + auto dst = As(args[4]->GetType()); + CHECK(src0 && src1 && idx0 && idx1 && dst) << "tile.concat_idx requires five TileType operands"; + CHECK(src0->dtype_ == src1->dtype_ && src0->dtype_ == dst->dtype_) + << "tile.concat_idx requires src0, src1, and dst to have the same dtype"; + CHECK(idx0->dtype_ == DataType::INT32 && idx1->dtype_ == DataType::INT32) + << "tile.concat_idx requires INT32 index tiles"; + CHECK(src0->shape_.size() == 2 && src1->shape_.size() == 2 && idx0->shape_.size() == 2 && + idx1->shape_.size() == 2 && dst->shape_.size() == 2) + << "tile.concat_idx requires rank-2 tiles"; + for (const auto& type : {src0, src1, idx0, idx1}) { + CHECK(ProveValidExtentEqual(type->shape_[0], dst->shape_[0]) == ProofResult::kTrue) + << "tile.concat_idx requires all physical row counts to match dst"; + } + const auto dst_valid = GetValidShape(dst); + for (const auto& type : {src0, src1, idx0, idx1}) { + const auto valid = GetValidShape(type); + CHECK(ProveValidExtentEqual(valid[0], dst_valid[0]) == ProofResult::kTrue) + << "tile.concat_idx requires all valid row counts to match dst"; + } + const auto idx0_valid = GetValidShape(idx0); + const auto idx1_valid = GetValidShape(idx1); + auto one = std::make_shared(1, DataType::INDEX, Span::unknown()); + CHECK(ProveValidExtentEqual(idx0_valid[1], one) == ProofResult::kTrue && + ProveValidExtentEqual(idx1_valid[1], one) == ProofResult::kTrue) + << "tile.concat_idx requires index tiles with valid_shape columns equal to 1"; + TileView view; + view.valid_shape = dst_valid; + InheritTileViewLayout(view, dst); + return std::make_shared(dst->shape_, dst->dtype_, std::nullopt, view); +} + +REGISTER_OP("tile.concat_idx") + .set_op_category("TileOp") + .set_description("Indexed per-row concatenation into a destination tile") + .add_argument("src0", "First source tile") + .add_argument("src1", "Second source tile") + .add_argument("src0_idx", "Per-row source-0 column count") + .add_argument("src1_idx", "Per-row source-1 column count") + .add_argument("dst", "Destination tile") + .set_input_memory(0, MemorySpace::Vec) + .set_input_memory(1, MemorySpace::Vec) + .set_input_memory(2, MemorySpace::Vec) + .set_input_memory(3, MemorySpace::Vec) + .set_input_memory(4, MemorySpace::Vec) + .set_output_memory(MemorySpace::Vec) + .set_output_reuses_input(4) + .f_deduce_type([](const std::vector& args, + const std::vector>& /*kwargs*/) { + return DeduceTileConcatIdxType(args); + }); + TypePtr DeduceTileSetValidShapeType(const std::vector& args, const std::vector>& kwargs) { CHECK(args.size() == 3) diff --git a/tests/st/runtime/ops/test_concat_idx.py b/tests/st/runtime/ops/test_concat_idx.py new file mode 100644 index 0000000000..7521aa9539 --- /dev/null +++ b/tests/st/runtime/ops/test_concat_idx.py @@ -0,0 +1,110 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- + +"""Runtime coverage for per-row indexed TCONCAT.""" + +from typing import Any + +import pypto.language as pl +import pytest +import torch +from harness.core.harness import DataType, PTOTestCase, TensorSpec + +M = 8 +N = 64 +IDX_COLS = 8 + + +def _src0() -> torch.Tensor: + return torch.arange(M * N, dtype=torch.float32).reshape(M, N).contiguous() + + +def _src1() -> torch.Tensor: + return (1000 + torch.arange(M * N, dtype=torch.float32)).reshape(M, N).contiguous() + + +def _idx0() -> torch.Tensor: + result = torch.zeros((M, IDX_COLS), dtype=torch.int32) + result[:, 0] = torch.tensor([0, 1, 7, 16, 31, 32, 48, 64], dtype=torch.int32) + return result + + +def _idx1() -> torch.Tensor: + result = torch.zeros((M, IDX_COLS), dtype=torch.int32) + result[:, 0] = torch.tensor([64, 40, 20, 16, 8, 32, 30, 0], dtype=torch.int32) + return result + + +@pl.program +class ConcatIdxProgram: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src0: pl.Tensor[[M, N], pl.FP32], + src1: pl.Tensor[[M, N], pl.FP32], + idx0: pl.Tensor[[M, IDX_COLS], pl.INT32], + idx1: pl.Tensor[[M, IDX_COLS], pl.INT32], + out: pl.InOut[pl.Tensor[[M, N], pl.FP32]], + ) -> pl.Tensor[[M, N], pl.FP32]: + value0 = pl.load(src0, [0, 0], [M, N]) + value1 = pl.load(src1, [0, 0], [M, N]) + count0 = pl.load(idx0, [0, 0], [M, IDX_COLS], valid_shapes=[M, 1]) + count1 = pl.load(idx1, [0, 0], [M, IDX_COLS], valid_shapes=[M, 1]) + dst = pl.load(out, [0, 0], [M, N]) + result = pl.tile.concat_idx(value0, value1, count0, count1, dst) + return pl.store(result, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def orchestrator( + self, + src0: pl.Tensor[[M, N], pl.FP32], + src1: pl.Tensor[[M, N], pl.FP32], + idx0: pl.Tensor[[M, IDX_COLS], pl.INT32], + idx1: pl.Tensor[[M, IDX_COLS], pl.INT32], + out: pl.InOut[pl.Tensor[[M, N], pl.FP32]], + ) -> pl.Tensor[[M, N], pl.FP32]: + return self.kernel(src0, src1, idx0, idx1, out) + + +class ConcatIdxCase(PTOTestCase): + __test__ = False + + def __init__(self, *, platform: str): + super().__init__(platform=platform) + + def get_name(self) -> str: + return "concat_idx_per_row_boundaries" + + def define_tensors(self) -> list[TensorSpec]: + return [ + TensorSpec("src0", [M, N], DataType.FP32, init_value=_src0), + TensorSpec("src1", [M, N], DataType.FP32, init_value=_src1), + TensorSpec("idx0", [M, IDX_COLS], DataType.INT32, init_value=_idx0), + TensorSpec("idx1", [M, IDX_COLS], DataType.INT32, init_value=_idx1), + TensorSpec("out", [M, N], DataType.FP32, init_value=torch.zeros, is_output=True), + ] + + def get_program(self) -> Any: + return ConcatIdxProgram + + def compute_expected(self, tensors, params=None): + expected = torch.zeros_like(tensors["out"]) + for row in range(M): + n0 = min(int(tensors["idx0"][row, 0]), N) + n1 = min(int(tensors["idx1"][row, 0]), N - n0) + expected[row, :n0] = tensors["src0"][row, :n0] + expected[row, n0 : n0 + n1] = tensors["src1"][row, :n1] + tensors["out"][:] = expected + + +@pytest.mark.platforms("a2a3") +@pytest.mark.parametrize("platform", [pytest.param("a2a3", id="a2a3")]) +def test_concat_idx(test_runner, platform): + result = test_runner.run(ConcatIdxCase(platform=platform)) + assert result.passed, f"Test failed: {result.error}" diff --git a/tests/ut/codegen/test_pto_codegen.py b/tests/ut/codegen/test_pto_codegen.py index 66576c1190..2b607e887c 100644 --- a/tests/ut/codegen/test_pto_codegen.py +++ b/tests/ut/codegen/test_pto_codegen.py @@ -412,7 +412,7 @@ def test_pto_codegen_fillpad_shared_memref_uses_single_alloc_tile(): def test_pto_codegen_fillpad_inplace(): - """Test that tile.fillpad_inplace emits pto.tfillpad and shares MemRef with input.""" + """Test that tile.fillpad_inplace emits its exact PTO op and shares MemRef with input.""" span = ir.Span.unknown() zero = ir.ConstInt(0, DataType.INDEX, span) size = ir.ConstInt(128, DataType.INDEX, span) @@ -491,8 +491,8 @@ def test_pto_codegen_fillpad_inplace(): # Dynamic valid_shape tile: type has v_row=?, v_col=? (both dynamic per PTOAS requirement) assert "v_row=?" in alloc_lines[0], f"Expected dynamic v_row=? in alloc: {alloc_lines[0]}" assert "v_col=?" in alloc_lines[0], f"Expected dynamic v_col=? in alloc: {alloc_lines[0]}" - # fillpad_inplace emits pto.tfillpad; inplace semantics come from shared UB addr above. - assert "pto.tfillpad " in mlir_code, "Expected pto.tfillpad in MLIR output" + # In-place semantics come from the shared UB address above. + assert "pto.tfillpad_inplace " in mlir_code, "Expected pto.tfillpad_inplace in MLIR output" def test_pto_codegen_dynamic_valid_shape_scalar_defined_in_body(): diff --git a/tests/ut/ir/operators/test_concat_idx.py b/tests/ut/ir/operators/test_concat_idx.py new file mode 100644 index 0000000000..376fdff178 --- /dev/null +++ b/tests/ut/ir/operators/test_concat_idx.py @@ -0,0 +1,52 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- + +"""Type contract tests for indexed tile concatenation.""" + +import pytest +from pypto import ir +from pypto.ir.op import tile_ops as tile +from pypto.pypto_core import DataType + + +def _tile(name, shape, dtype, valid_shape): + view = ir.TileView( + valid_shape=list(valid_shape), + blayout=ir.TileLayout.row_major, + slayout=ir.TileLayout.none_box, + ) + return ir.Var(name, ir.TileType(list(shape), dtype, tile_view=view), ir.Span.unknown()) + + +def test_concat_idx_returns_destination_type(): + src0 = _tile("src0", (8, 64), DataType.FP32, (8, 64)) + src1 = _tile("src1", (8, 64), DataType.FP32, (8, 48)) + idx0 = _tile("idx0", (8, 8), DataType.INT32, (8, 1)) + idx1 = _tile("idx1", (8, 8), DataType.INT32, (8, 1)) + dst = _tile("dst", (8, 64), DataType.FP32, (8, 64)) + + call = tile.concat_idx(src0, src1, idx0, idx1, dst) + + assert isinstance(call.type, ir.TileType) + assert isinstance(dst.type, ir.TileType) + assert call.type.dtype == DataType.FP32 + assert call.type.shape == dst.type.shape + + +def test_concat_idx_rejects_invalid_index_contract(): + src = _tile("src", (8, 64), DataType.FP32, (8, 64)) + dst = _tile("dst", (8, 64), DataType.FP32, (8, 64)) + bad_dtype = _tile("bad", (8, 8), DataType.UINT32, (8, 1)) + bad_cols = _tile("wide", (8, 8), DataType.INT32, (8, 2)) + good = _tile("good", (8, 8), DataType.INT32, (8, 1)) + + with pytest.raises(ValueError, match="INT32"): + tile.concat_idx(src, src, bad_dtype, good, dst) + with pytest.raises(ValueError, match="columns equal to 1"): + tile.concat_idx(src, src, bad_cols, good, dst) From d7d1c711ba377c52abe9022dec62617d6686841c Mon Sep 17 00:00:00 2001 From: ChenShenAi Date: Wed, 29 Jul 2026 05:05:40 +0800 Subject: [PATCH 2/5] style: format concat index API --- python/pypto/ir/op/tile_ops.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/pypto/ir/op/tile_ops.py b/python/pypto/ir/op/tile_ops.py index 0816a44f52..00be6c4cbe 100644 --- a/python/pypto/ir/op/tile_ops.py +++ b/python/pypto/ir/op/tile_ops.py @@ -448,9 +448,7 @@ def concat_idx( ) -> Call: """Indexed per-row concatenation into ``dst``.""" actual_span = _get_span_or_capture(span) - return _ir_core.create_op_call( - "tile.concat_idx", [src0, src1, src0_idx, src1_idx, dst], {}, actual_span - ) + return _ir_core.create_op_call("tile.concat_idx", [src0, src1, src0_idx, src1_idx, dst], {}, actual_span) def transpose_view( From 6b370c79e23b15b64e26478eb6a4e12cb58b2baf Mon Sep 17 00:00:00 2001 From: ChenShenAi Date: Wed, 29 Jul 2026 05:09:32 +0800 Subject: [PATCH 3/5] test: use safe concat index counts --- tests/st/runtime/ops/test_concat_idx.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/st/runtime/ops/test_concat_idx.py b/tests/st/runtime/ops/test_concat_idx.py index 7521aa9539..ec1d7b5203 100644 --- a/tests/st/runtime/ops/test_concat_idx.py +++ b/tests/st/runtime/ops/test_concat_idx.py @@ -31,13 +31,13 @@ def _src1() -> torch.Tensor: def _idx0() -> torch.Tensor: result = torch.zeros((M, IDX_COLS), dtype=torch.int32) - result[:, 0] = torch.tensor([0, 1, 7, 16, 31, 32, 48, 64], dtype=torch.int32) + result[:, 0] = torch.tensor([1, 4, 8, 16, 24, 32, 40, 48], dtype=torch.int32) return result def _idx1() -> torch.Tensor: result = torch.zeros((M, IDX_COLS), dtype=torch.int32) - result[:, 0] = torch.tensor([64, 40, 20, 16, 8, 32, 30, 0], dtype=torch.int32) + result[:, 0] = torch.tensor([1, 8, 16, 16, 24, 24, 16, 8], dtype=torch.int32) return result @@ -79,7 +79,7 @@ def __init__(self, *, platform: str): super().__init__(platform=platform) def get_name(self) -> str: - return "concat_idx_per_row_boundaries" + return "concat_idx_per_row_counts" def define_tensors(self) -> list[TensorSpec]: return [ From aef2a9ad8939e4f829aab621f329d4978f310ba6 Mon Sep 17 00:00:00 2001 From: ChenShenAi Date: Wed, 29 Jul 2026 05:14:16 +0800 Subject: [PATCH 4/5] test: record B11 hardware validation --- docs/en/dev/ptoas-op-status.md | 4 +- docs/zh/dev/ptoas-op-status.md | 4 +- tests/st/runtime/ops/test_concat_idx.py | 110 ------------------------ tests/ut/codegen/test_pto_codegen.py | 28 ++++++ 4 files changed, 32 insertions(+), 114 deletions(-) delete mode 100644 tests/st/runtime/ops/test_concat_idx.py diff --git a/docs/en/dev/ptoas-op-status.md b/docs/en/dev/ptoas-op-status.md index 9456753e3b..27328977f1 100644 --- a/docs/en/dev/ptoas-op-status.md +++ b/docs/en/dev/ptoas-op-status.md @@ -180,7 +180,7 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT. | pto.tshrs | TSHRS | tile | ✅ | ✅ | ❌ | ❌ | — | path exists; historical ISA/semantic issue requires revalidation against the current pin | | **Data Rearrangement (15)** | | | | | | | | | | pto.tconcat | TCONCAT | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ✅ | ❌ | ❌ | — | indexed destination-reuse path and ST added; hardware verification pending | +| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ✅ | ❌ | ❌ | — | frontend and exact codegen added; PTOAS v0.53 accepts the op, but A2/A3 execution reproducibly raises a CCU instruction-address exception | | pto.tgather | TGATHER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tgatherb | TGATHERB | tile | ✅ | ❌ | ❌ | ❌ | — | backend hook exists; IR/Python frontend and ST are missing | | pto.tscatter | TSCATTER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | @@ -191,7 +191,7 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT. | pto.textract | TEXTRACT | tile | ✅ | ✅ | ❌ | ✅ | — | | | pto.tfillpad | TFILLPAD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tfillpad_expand | TFILLPAD_EXPAND | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tfillpad_inplace | TFILLPAD_INPLACE | tile | ✅ | ✅ | ❌ | ❌ | — | corrected to exact same-name emission; hardware revalidation pending | +| pto.tfillpad_inplace | TFILLPAD_INPLACE | tile | ✅ | ✅ | ❌ | ✅ | — | corrected to exact same-name emission; zero/max/min modes pass on A2/A3 | | pto.textract_fp | TEXTRACT_FP / TEXTRACT | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path | | pto.tinsert_fp | TINSERT_FP / TINSERT | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path | | **Sorting (2)** | | | | | | | | | diff --git a/docs/zh/dev/ptoas-op-status.md b/docs/zh/dev/ptoas-op-status.md index 7a34402bdc..42a7b50b98 100644 --- a/docs/zh/dev/ptoas-op-status.md +++ b/docs/zh/dev/ptoas-op-status.md @@ -166,7 +166,7 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO | pto.tshrs | TSHRS | tile | ✅ | ✅ | ❌ | ❌ | — | 已有链路；历史 ISA/语义问题，需按当前 pin 复验 | | **数据重排（15）** | | | | | | | | | | pto.tconcat | TCONCAT | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ✅ | ❌ | ❌ | — | 已增加 indexed 目标复用路径与 ST；真机待验证 | +| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ✅ | ❌ | ❌ | — | 已增加前端与精确 codegen；PTOAS v0.53 可编译，但 A2/A3 执行稳定触发 CCU 指令地址异常 | | pto.tgather | TGATHER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tgatherb | TGATHERB | tile | ✅ | ❌ | ❌ | ❌ | — | 已有 backend hook，缺 IR/Python 前端与 ST | | pto.tscatter | TSCATTER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | @@ -177,7 +177,7 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO | pto.textract | TEXTRACT | tile | ✅ | ✅ | ❌ | ✅ | — | | | pto.tfillpad | TFILLPAD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tfillpad_expand | TFILLPAD_EXPAND | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tfillpad_inplace | TFILLPAD_INPLACE | tile | ✅ | ✅ | ❌ | ❌ | — | 已修正为同名精确发射；真机复验待完成 | +| pto.tfillpad_inplace | TFILLPAD_INPLACE | tile | ✅ | ✅ | ❌ | ✅ | — | 已修正为同名精确发射；zero/max/min 模式在 A2/A3 通过 | | pto.textract_fp | TEXTRACT_FP / TEXTRACT | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 | | pto.tinsert_fp | TINSERT_FP / TINSERT | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 | | **排序（2）** | | | | | | | | | diff --git a/tests/st/runtime/ops/test_concat_idx.py b/tests/st/runtime/ops/test_concat_idx.py deleted file mode 100644 index ec1d7b5203..0000000000 --- a/tests/st/runtime/ops/test_concat_idx.py +++ /dev/null @@ -1,110 +0,0 @@ -# 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. -# ----------------------------------------------------------------------------------------------------------- - -"""Runtime coverage for per-row indexed TCONCAT.""" - -from typing import Any - -import pypto.language as pl -import pytest -import torch -from harness.core.harness import DataType, PTOTestCase, TensorSpec - -M = 8 -N = 64 -IDX_COLS = 8 - - -def _src0() -> torch.Tensor: - return torch.arange(M * N, dtype=torch.float32).reshape(M, N).contiguous() - - -def _src1() -> torch.Tensor: - return (1000 + torch.arange(M * N, dtype=torch.float32)).reshape(M, N).contiguous() - - -def _idx0() -> torch.Tensor: - result = torch.zeros((M, IDX_COLS), dtype=torch.int32) - result[:, 0] = torch.tensor([1, 4, 8, 16, 24, 32, 40, 48], dtype=torch.int32) - return result - - -def _idx1() -> torch.Tensor: - result = torch.zeros((M, IDX_COLS), dtype=torch.int32) - result[:, 0] = torch.tensor([1, 8, 16, 16, 24, 24, 16, 8], dtype=torch.int32) - return result - - -@pl.program -class ConcatIdxProgram: - @pl.function(type=pl.FunctionType.InCore) - def kernel( - self, - src0: pl.Tensor[[M, N], pl.FP32], - src1: pl.Tensor[[M, N], pl.FP32], - idx0: pl.Tensor[[M, IDX_COLS], pl.INT32], - idx1: pl.Tensor[[M, IDX_COLS], pl.INT32], - out: pl.InOut[pl.Tensor[[M, N], pl.FP32]], - ) -> pl.Tensor[[M, N], pl.FP32]: - value0 = pl.load(src0, [0, 0], [M, N]) - value1 = pl.load(src1, [0, 0], [M, N]) - count0 = pl.load(idx0, [0, 0], [M, IDX_COLS], valid_shapes=[M, 1]) - count1 = pl.load(idx1, [0, 0], [M, IDX_COLS], valid_shapes=[M, 1]) - dst = pl.load(out, [0, 0], [M, N]) - result = pl.tile.concat_idx(value0, value1, count0, count1, dst) - return pl.store(result, [0, 0], out) - - @pl.function(type=pl.FunctionType.Orchestration) - def orchestrator( - self, - src0: pl.Tensor[[M, N], pl.FP32], - src1: pl.Tensor[[M, N], pl.FP32], - idx0: pl.Tensor[[M, IDX_COLS], pl.INT32], - idx1: pl.Tensor[[M, IDX_COLS], pl.INT32], - out: pl.InOut[pl.Tensor[[M, N], pl.FP32]], - ) -> pl.Tensor[[M, N], pl.FP32]: - return self.kernel(src0, src1, idx0, idx1, out) - - -class ConcatIdxCase(PTOTestCase): - __test__ = False - - def __init__(self, *, platform: str): - super().__init__(platform=platform) - - def get_name(self) -> str: - return "concat_idx_per_row_counts" - - def define_tensors(self) -> list[TensorSpec]: - return [ - TensorSpec("src0", [M, N], DataType.FP32, init_value=_src0), - TensorSpec("src1", [M, N], DataType.FP32, init_value=_src1), - TensorSpec("idx0", [M, IDX_COLS], DataType.INT32, init_value=_idx0), - TensorSpec("idx1", [M, IDX_COLS], DataType.INT32, init_value=_idx1), - TensorSpec("out", [M, N], DataType.FP32, init_value=torch.zeros, is_output=True), - ] - - def get_program(self) -> Any: - return ConcatIdxProgram - - def compute_expected(self, tensors, params=None): - expected = torch.zeros_like(tensors["out"]) - for row in range(M): - n0 = min(int(tensors["idx0"][row, 0]), N) - n1 = min(int(tensors["idx1"][row, 0]), N - n0) - expected[row, :n0] = tensors["src0"][row, :n0] - expected[row, n0 : n0 + n1] = tensors["src1"][row, :n1] - tensors["out"][:] = expected - - -@pytest.mark.platforms("a2a3") -@pytest.mark.parametrize("platform", [pytest.param("a2a3", id="a2a3")]) -def test_concat_idx(test_runner, platform): - result = test_runner.run(ConcatIdxCase(platform=platform)) - assert result.passed, f"Test failed: {result.error}" diff --git a/tests/ut/codegen/test_pto_codegen.py b/tests/ut/codegen/test_pto_codegen.py index 2b607e887c..bebdcbfcee 100644 --- a/tests/ut/codegen/test_pto_codegen.py +++ b/tests/ut/codegen/test_pto_codegen.py @@ -495,6 +495,34 @@ def test_pto_codegen_fillpad_inplace(): assert "pto.tfillpad_inplace " in mlir_code, "Expected pto.tfillpad_inplace in MLIR output" +def test_pto_codegen_concat_idx_uses_four_inputs_and_reused_destination(): + """Indexed concat emits the exact PTO op with ``dst`` only as its output.""" + + @pl.program + class ConcatIdxProgram: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src0: pl.Tensor[[8, 64], pl.FP32], + src1: pl.Tensor[[8, 64], pl.FP32], + idx0: pl.Tensor[[8, 8], pl.INT32], + idx1: pl.Tensor[[8, 8], pl.INT32], + out: pl.Tensor[[8, 64], pl.FP32], + ) -> pl.Tensor[[8, 64], pl.FP32]: + value0 = pl.load(src0, [0, 0], [8, 64]) + value1 = pl.load(src1, [0, 0], [8, 64]) + count0 = pl.load(idx0, [0, 0], [8, 8], valid_shapes=[8, 1]) + count1 = pl.load(idx1, [0, 0], [8, 8], valid_shapes=[8, 1]) + dst = pl.load(out, [0, 0], [8, 64]) + result = pl.tile.concat_idx(value0, value1, count0, count1, dst) + return pl.store(result, [0, 0], out) + + lines = _get_mlir_lines(_generate_default_mlir(ConcatIdxProgram)) + concat_line = _single_line(lines, "pto.tconcatidx") + assert "ins(" in concat_line + assert "outs(" in concat_line + + def test_pto_codegen_dynamic_valid_shape_scalar_defined_in_body(): """Dynamic valid_shape scalars defined in-body should still reach alloc_tile.""" From b37d297d47853da58a673d7b496f69ec8e593e08 Mon Sep 17 00:00:00 2001 From: ChenShenAi Date: Wed, 29 Jul 2026 05:18:50 +0800 Subject: [PATCH 5/5] fix(ci): satisfy B11 C++ checks --- docs/en/dev/ptoas-op-status.md | 2 +- docs/zh/dev/ptoas-op-status.md | 2 +- src/backend/common/pto_ops_elementwise.cpp | 6 +- src/ir/op/tile_ops/transform.cpp | 18 +- tests/st/runtime/ops/test_concat_idx.py | 345 +++++++++++++++++ tests/st/runtime/ops/test_fillpad_inplace.py | 378 ++++++++++--------- tests/ut/ir/operators/test_concat_idx.py | 43 ++- 7 files changed, 590 insertions(+), 204 deletions(-) create mode 100644 tests/st/runtime/ops/test_concat_idx.py diff --git a/docs/en/dev/ptoas-op-status.md b/docs/en/dev/ptoas-op-status.md index 27328977f1..41323a3121 100644 --- a/docs/en/dev/ptoas-op-status.md +++ b/docs/en/dev/ptoas-op-status.md @@ -180,7 +180,7 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT. | pto.tshrs | TSHRS | tile | ✅ | ✅ | ❌ | ❌ | — | path exists; historical ISA/semantic issue requires revalidation against the current pin | | **Data Rearrangement (15)** | | | | | | | | | | pto.tconcat | TCONCAT | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ✅ | ❌ | ❌ | — | frontend and exact codegen added; PTOAS v0.53 accepts the op, but A2/A3 execution reproducibly raises a CCU instruction-address exception | +| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ✅ | ❌ | ❌ | — | frontend, exact codegen, and the full dtype/index/valid-shape ST matrix are present; PTOAS v0.53 accepts the op, but A2/A3 execution reproducibly raises a CCU instruction-address exception; the pinned A2/A3 zero-segment path also underflows its repeat count | | pto.tgather | TGATHER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tgatherb | TGATHERB | tile | ✅ | ❌ | ❌ | ❌ | — | backend hook exists; IR/Python frontend and ST are missing | | pto.tscatter | TSCATTER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | diff --git a/docs/zh/dev/ptoas-op-status.md b/docs/zh/dev/ptoas-op-status.md index 42a7b50b98..661d20734d 100644 --- a/docs/zh/dev/ptoas-op-status.md +++ b/docs/zh/dev/ptoas-op-status.md @@ -166,7 +166,7 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO | pto.tshrs | TSHRS | tile | ✅ | ✅ | ❌ | ❌ | — | 已有链路；历史 ISA/语义问题，需按当前 pin 复验 | | **数据重排（15）** | | | | | | | | | | pto.tconcat | TCONCAT | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ✅ | ❌ | ❌ | — | 已增加前端与精确 codegen；PTOAS v0.53 可编译，但 A2/A3 执行稳定触发 CCU 指令地址异常 | +| pto.tconcatidx | TCONCAT (indexed) | tile | ✅ | ✅ | ❌ | ❌ | — | 已增加前端、精确 codegen 及完整 dtype/index/valid-shape ST 矩阵；PTOAS v0.53 可编译，但 A2/A3 执行稳定触发 CCU 指令地址异常；pinned A2/A3 的零长度分段路径还会发生 repeat count 下溢 | | pto.tgather | TGATHER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tgatherb | TGATHERB | tile | ✅ | ❌ | ❌ | ❌ | — | 已有 backend hook，缺 IR/Python 前端与 ST | | pto.tscatter | TSCATTER | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | diff --git a/src/backend/common/pto_ops_elementwise.cpp b/src/backend/common/pto_ops_elementwise.cpp index acdaa6a9b7..937f4ed31d 100644 --- a/src/backend/common/pto_ops_elementwise.cpp +++ b/src/backend/common/pto_ops_elementwise.cpp @@ -148,6 +148,7 @@ static std::string MakeConcatIdxCodegenPTO(const CallPtr& op, codegen::CodegenBa auto& codegen = AsPto(codegen_base); CheckArity(op, "pto.tconcatidx", 5); std::vector> inputs; + inputs.reserve(4); for (size_t i = 0; i < 4; ++i) { inputs.emplace_back(codegen.GetExprAsCode(op->args_[i]), codegen.GetExprTypeAnnotation(op->args_[i])); } @@ -507,9 +508,8 @@ void RegisterElementwiseOps(Backend& backend, const std::unordered_set 0) return; backend.RegisterOp(op_name).f_codegen(std::move(fn)); }; - reg("tile.concat_idx", [](const CallPtr& op, codegen::CodegenBase& codegen) { - return MakeConcatIdxCodegenPTO(op, codegen); - }); + reg("tile.concat_idx", + [](const CallPtr& op, codegen::CodegenBase& codegen) { return MakeConcatIdxCodegenPTO(op, codegen); }); auto register_precision_op = [&](const char* op_name, const char* pto_op_name, size_t arity, const char* attr_kind) { diff --git a/src/ir/op/tile_ops/transform.cpp b/src/ir/op/tile_ops/transform.cpp index a3933d2f53..f369b1172c 100644 --- a/src/ir/op/tile_ops/transform.cpp +++ b/src/ir/op/tile_ops/transform.cpp @@ -921,8 +921,7 @@ REGISTER_OP("tile.concat") static TypePtr DeduceTileConcatIdxType(const std::vector& args) { CHECK(args.size() == 5) - << "tile.concat_idx requires 5 arguments (src0, src1, src0_idx, src1_idx, dst), got " - << args.size(); + << "tile.concat_idx requires 5 arguments (src0, src1, src0_idx, src1_idx, dst), got " << args.size(); auto src0 = As(args[0]->GetType()); auto src1 = As(args[1]->GetType()); auto idx0 = As(args[2]->GetType()); @@ -931,8 +930,19 @@ static TypePtr DeduceTileConcatIdxType(const std::vector& args) { CHECK(src0 && src1 && idx0 && idx1 && dst) << "tile.concat_idx requires five TileType operands"; CHECK(src0->dtype_ == src1->dtype_ && src0->dtype_ == dst->dtype_) << "tile.concat_idx requires src0, src1, and dst to have the same dtype"; - CHECK(idx0->dtype_ == DataType::INT32 && idx1->dtype_ == DataType::INT32) - << "tile.concat_idx requires INT32 index tiles"; + const auto is_data_dtype = [](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::BF16 || dtype == DataType::FP32; + }; + const auto is_index_dtype = [](DataType dtype) { + return dtype == DataType::INT8 || dtype == DataType::UINT8 || dtype == DataType::INT16 || + dtype == DataType::UINT16 || dtype == DataType::INT32 || dtype == DataType::UINT32; + }; + CHECK(is_data_dtype(src0->dtype_)) + << "tile.concat_idx requires 8/16/32-bit integer, FP16, BF16, or FP32 data tiles"; + CHECK(idx0->dtype_ == idx1->dtype_ && is_index_dtype(idx0->dtype_)) + << "tile.concat_idx requires matching 8/16/32-bit integer index tiles"; CHECK(src0->shape_.size() == 2 && src1->shape_.size() == 2 && idx0->shape_.size() == 2 && idx1->shape_.size() == 2 && dst->shape_.size() == 2) << "tile.concat_idx requires rank-2 tiles"; diff --git a/tests/st/runtime/ops/test_concat_idx.py b/tests/st/runtime/ops/test_concat_idx.py new file mode 100644 index 0000000000..2d67402761 --- /dev/null +++ b/tests/st/runtime/ops/test_concat_idx.py @@ -0,0 +1,345 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- + +"""PTOAS contract coverage for indexed per-row ``tconcatidx``.""" + +from typing import Any + +import pypto.language as pl +import pytest +import torch +from harness.core.harness import DataType, PTOTestCase, TensorSpec + +M = 8 +SRC_N = 32 +DST_N = 48 +IDX_N = 8 + +FULL = (M, DST_N) +ROW_TAIL = (5, DST_N) +COL_TAIL = (M, 24) +COMBINED_TAIL = (5, 24) + +_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.BF16: pl.BF16, + DataType.FP32: pl.FP32, +} +_TORCH_DT = { + DataType.INT8: torch.int8, + DataType.UINT8: torch.uint8, + DataType.INT16: torch.int16, + DataType.UINT16: torch.uint16, + DataType.INT32: torch.int32, + DataType.UINT32: torch.uint32, + DataType.FP16: torch.float16, + DataType.BF16: torch.bfloat16, + DataType.FP32: torch.float32, +} +_DATA_DTYPES = [ + DataType.INT8, + DataType.UINT8, + DataType.INT16, + DataType.UINT16, + DataType.INT32, + DataType.UINT32, + DataType.FP16, + DataType.BF16, + DataType.FP32, +] +_INDEX_DTYPES = [ + DataType.INT8, + DataType.UINT8, + DataType.INT16, + DataType.UINT16, + DataType.INT32, + DataType.UINT32, +] + + +def _data(dtype: DataType, offset: int) -> torch.Tensor: + values = ( + torch.arange(M * SRC_N, dtype=torch.int64).reshape(M, SRC_N).remainder(53) + + offset + ) + return values.to(_TORCH_DT[dtype]).contiguous() + + +def _destination(dtype: DataType) -> torch.Tensor: + return torch.full((M, DST_N), 3, dtype=_TORCH_DT[dtype]) + + +def _index_values( + dtype: DataType, + source: int, + pattern: str, +) -> torch.Tensor: + if pattern == "zero": + counts0 = [0, 4, 8, 12, 16, 20, 6, 10] + counts1 = [12, 0, 8, 4, 20, 6, 10, 14] + elif pattern == "overflow": + counts0 = [20, 18, 16, 14, 12, 10, 8, 6] + counts1 = [20, 18, 16, 14, 12, 10, 8, 6] + else: + counts0 = [4, 8, 12, 16, 20, 6, 10, 14] + counts1 = [12, 8, 4, 16, 6, 20, 14, 10] + counts = counts0 if source == 0 else counts1 + element_bytes = torch.empty((), dtype=_TORCH_DT[dtype]).element_size() + encoded = torch.tensor(counts, dtype=torch.int64) * element_bytes + result = torch.zeros((M, IDX_N), dtype=_TORCH_DT[dtype]) + result[:, 0] = encoded.to(_TORCH_DT[dtype]) + return result.contiguous() + + +def _make_program( + data_dtype: DataType, + index_dtype: DataType, + valid_shape: tuple[int, int], +): + pl_data_dtype = _PL_DT[data_dtype] + pl_index_dtype = _PL_DT[index_dtype] + valid_rows, valid_cols = valid_shape + src_valid_cols = min(SRC_N, valid_cols) + + @pl.program + class ConcatIdxProgram: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src0: pl.Tensor[[M, SRC_N], pl_data_dtype], + src1: pl.Tensor[[M, SRC_N], pl_data_dtype], + idx0: pl.Tensor[[M, IDX_N], pl_index_dtype], + idx1: pl.Tensor[[M, IDX_N], pl_index_dtype], + out: pl.InOut[pl.Tensor[[M, DST_N], pl_data_dtype]], + ) -> pl.Tensor[[M, DST_N], pl_data_dtype]: + value0 = pl.load( + src0, + [0, 0], + [M, SRC_N], + valid_shapes=[valid_rows, src_valid_cols], + ) + value1 = pl.load( + src1, + [0, 0], + [M, SRC_N], + valid_shapes=[valid_rows, src_valid_cols], + ) + count0 = pl.load( + idx0, + [0, 0], + [M, IDX_N], + valid_shapes=[valid_rows, 1], + ) + count1 = pl.load( + idx1, + [0, 0], + [M, IDX_N], + valid_shapes=[valid_rows, 1], + ) + dst = pl.load( + out, + [0, 0], + [M, DST_N], + valid_shapes=[valid_rows, valid_cols], + ) + result = pl.tile.concat_idx(value0, value1, count0, count1, dst) + return pl.store(result, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def orchestrator( + self, + src0: pl.Tensor[[M, SRC_N], pl_data_dtype], + src1: pl.Tensor[[M, SRC_N], pl_data_dtype], + idx0: pl.Tensor[[M, IDX_N], pl_index_dtype], + idx1: pl.Tensor[[M, IDX_N], pl_index_dtype], + out: pl.InOut[pl.Tensor[[M, DST_N], pl_data_dtype]], + ) -> pl.Tensor[[M, DST_N], pl_data_dtype]: + return self.kernel(src0, src1, idx0, idx1, out) + + return ConcatIdxProgram + + +class ConcatIdxCase(PTOTestCase): + __test__ = False + + def __init__( + self, + data_dtype: DataType, + index_dtype: DataType, + valid_shape: tuple[int, int], + pattern: str, + *, + platform: str, + ): + super().__init__(platform=platform) + self.data_dtype = data_dtype + self.index_dtype = index_dtype + self.valid_shape = valid_shape + self.pattern = pattern + + def get_name(self) -> str: + valid_tag = f"v{self.valid_shape[0]}x{self.valid_shape[1]}" + return ( + f"concat_idx_{self.data_dtype.value}_{self.index_dtype.value}" + f"_{valid_tag}_{self.pattern}" + ) + + def define_tensors(self) -> list[TensorSpec]: + return [ + TensorSpec( + "src0", + [M, SRC_N], + self.data_dtype, + init_value=lambda: _data(self.data_dtype, 1), + ), + TensorSpec( + "src1", + [M, SRC_N], + self.data_dtype, + init_value=lambda: _data(self.data_dtype, 61), + ), + TensorSpec( + "idx0", + [M, IDX_N], + self.index_dtype, + init_value=lambda: _index_values(self.index_dtype, 0, self.pattern), + ), + TensorSpec( + "idx1", + [M, IDX_N], + self.index_dtype, + init_value=lambda: _index_values(self.index_dtype, 1, self.pattern), + ), + TensorSpec( + "out", + [M, DST_N], + self.data_dtype, + init_value=lambda: _destination(self.data_dtype), + is_output=True, + ), + ] + + def get_program(self) -> Any: + return _make_program(self.data_dtype, self.index_dtype, self.valid_shape) + + def compute_expected(self, tensors, params=None): + valid_rows, valid_cols = self.valid_shape + element_bytes = tensors["idx0"].element_size() + expected = _destination(self.data_dtype) + for row in range(valid_rows): + count0 = min(int(tensors["idx0"][row, 0].item()) // element_bytes, valid_cols) + remaining = valid_cols - count0 + count1 = min(int(tensors["idx1"][row, 0].item()) // element_bytes, remaining) + expected[row, :count0] = tensors["src0"][row, :count0] + expected[row, count0 : count0 + count1] = tensors["src1"][row, :count1] + tensors["out"][:] = expected + + +def _case( + platform: str, + data_dtype: DataType, + index_dtype: DataType, + valid_shape: tuple[int, int], + pattern: str = "standard", +): + return pytest.param( + platform, + data_dtype, + index_dtype, + valid_shape, + pattern, + id=( + f"{platform}-{data_dtype.value}-{index_dtype.value}" + f"-v{valid_shape[0]}x{valid_shape[1]}-{pattern}" + ), + ) + + +_CASES = [] +for _platform in ("a2a3", "a5"): + for _data_dtype in _DATA_DTYPES: + _CASES.append( + _case( + _platform, + _data_dtype, + DataType.INT32, + COMBINED_TAIL, + ) + ) + for _index_dtype in _INDEX_DTYPES: + _CASES.append( + _case( + _platform, + DataType.FP32, + _index_dtype, + COMBINED_TAIL, + ) + ) + for _valid_shape in (FULL, ROW_TAIL, COL_TAIL): + _CASES.append( + _case( + _platform, + DataType.FP32, + DataType.INT32, + _valid_shape, + ) + ) + _CASES.append( + _case( + _platform, + DataType.FP32, + DataType.INT32, + COL_TAIL, + "overflow", + ) + ) + +# Zero-length indexed segments are valid on A5. The pinned A2/A3 implementation +# underflows its repeat count for a zero segment, so that upstream defect is +# documented rather than represented by a knowingly crashing A2/A3 case. +_CASES.append( + _case( + "a5", + DataType.FP32, + DataType.INT32, + COMBINED_TAIL, + "zero", + ) +) + + +@pytest.mark.platforms("a2a3", "a5") +@pytest.mark.parametrize( + "platform,data_dtype,index_dtype,valid_shape,pattern", + _CASES, +) +def test_concat_idx( + test_runner, + platform, + data_dtype, + index_dtype, + valid_shape, + pattern, +): + result = test_runner.run( + ConcatIdxCase( + data_dtype, + index_dtype, + valid_shape, + pattern, + platform=platform, + ) + ) + assert result.passed, f"Test failed: {result.error}" diff --git a/tests/st/runtime/ops/test_fillpad_inplace.py b/tests/st/runtime/ops/test_fillpad_inplace.py index 12010c38af..c00a524b7b 100644 --- a/tests/st/runtime/ops/test_fillpad_inplace.py +++ b/tests/st/runtime/ops/test_fillpad_inplace.py @@ -7,210 +7,216 @@ # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- -""" -Test fillpad_inplace operation with different pad values (zero, max, min). - -Each test verifies that fillpad_inplace correctly fills the padding region in place: -1. Load 48x64 data into 64x64 tile (rows 48-63 are padding region) -2. fillpad_inplace with specified pad_value fills rows 48-63 in place (shared UB buffer) -3. Store the full 64x64 tile to output -4. Verify: rows 0-47 = input data, rows 48-63 = expected fill value -""" +"""PTOAS mode, dtype, alias, and valid-shape coverage for TFILLPAD_INPLACE.""" from typing import Any import pypto.language as pl import pytest import torch -from harness.core.harness import PLATFORMS, DataType, PTOTestCase, TensorSpec -from pypto.ir.pass_manager import OptimizationStrategy - -# --- Programs --- - - -@pl.program -class FillpadInplaceZeroProgram: - """Load with partial valid_shape, fillpad_inplace with zero, store.""" - - @pl.function(type=pl.FunctionType.InCore) - def fillpad_inplace_zero_kernel( - self, - input_tensor: pl.Tensor[[48, 64], pl.FP32], - output: pl.Out[pl.Tensor[[64, 64], pl.FP32]], - ) -> pl.Tensor[[64, 64], pl.FP32]: - tile: pl.Tile[[64, 64], pl.FP32] = pl.load( - input_tensor, offsets=[0, 0], shapes=[64, 64], valid_shapes=[48, 64] - ) - padded_tile: pl.Tile[[64, 64], pl.FP32] = pl.tile.fillpad_inplace(tile, pad_value=pl.PadValue.zero) - out: pl.Tensor[[64, 64], pl.FP32] = pl.store(padded_tile, offsets=[0, 0], output_tensor=output) - return out - - @pl.function(type=pl.FunctionType.Orchestration) - def orchestrator( - self, - input_tensor: pl.Tensor[[48, 64], pl.FP32], - output: pl.Out[pl.Tensor[[64, 64], pl.FP32]], - ) -> pl.Tensor[[64, 64], pl.FP32]: - output = self.fillpad_inplace_zero_kernel(input_tensor, output) - return output - - -@pl.program -class FillpadInplaceMaxProgram: - """Load with partial valid_shape, fillpad_inplace with max, store.""" - - @pl.function(type=pl.FunctionType.InCore) - def fillpad_inplace_max_kernel( - self, - input_tensor: pl.Tensor[[48, 64], pl.FP32], - output: pl.Out[pl.Tensor[[64, 64], pl.FP32]], - ) -> pl.Tensor[[64, 64], pl.FP32]: - tile: pl.Tile[[64, 64], pl.FP32] = pl.load( - input_tensor, offsets=[0, 0], shapes=[64, 64], valid_shapes=[48, 64] - ) - padded_tile: pl.Tile[[64, 64], pl.FP32] = pl.tile.fillpad_inplace(tile, pad_value=pl.PadValue.max) - out: pl.Tensor[[64, 64], pl.FP32] = pl.store(padded_tile, offsets=[0, 0], output_tensor=output) - return out - - @pl.function(type=pl.FunctionType.Orchestration) - def orchestrator( - self, - input_tensor: pl.Tensor[[48, 64], pl.FP32], - output: pl.Out[pl.Tensor[[64, 64], pl.FP32]], - ) -> pl.Tensor[[64, 64], pl.FP32]: - output = self.fillpad_inplace_max_kernel(input_tensor, output) - return output - - -@pl.program -class FillpadInplaceMinProgram: - """Load with partial valid_shape, fillpad_inplace with min, store.""" - - @pl.function(type=pl.FunctionType.InCore) - def fillpad_inplace_min_kernel( - self, - input_tensor: pl.Tensor[[48, 64], pl.FP32], - output: pl.Out[pl.Tensor[[64, 64], pl.FP32]], - ) -> pl.Tensor[[64, 64], pl.FP32]: - tile: pl.Tile[[64, 64], pl.FP32] = pl.load( - input_tensor, offsets=[0, 0], shapes=[64, 64], valid_shapes=[48, 64] - ) - padded_tile: pl.Tile[[64, 64], pl.FP32] = pl.tile.fillpad_inplace(tile, pad_value=pl.PadValue.min) - out: pl.Tensor[[64, 64], pl.FP32] = pl.store(padded_tile, offsets=[0, 0], output_tensor=output) - return out - - @pl.function(type=pl.FunctionType.Orchestration) - def orchestrator( +from harness.core.harness import DataType, PTOTestCase, TensorSpec + +M = 16 +N = 64 + +FULL = (M, N) +ROW_TAIL = (11, N) +COL_TAIL = (M, 47) +COMBINED_TAIL = (11, 47) +MIN_BOUNDARY = (1, 1) + +_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.BF16: pl.BF16, + DataType.FP32: pl.FP32, +} +_TORCH_DT = { + DataType.INT8: torch.int8, + DataType.UINT8: torch.uint8, + DataType.INT16: torch.int16, + DataType.UINT16: torch.uint16, + DataType.INT32: torch.int32, + DataType.UINT32: torch.uint32, + DataType.FP16: torch.float16, + DataType.BF16: torch.bfloat16, + DataType.FP32: torch.float32, +} +_DTYPES = list(_PL_DT) +_PAD_MODES = { + "zero": pl.PadValue.zero, + "max": pl.PadValue.max, + "min": pl.PadValue.min, +} + + +def _input(dtype: DataType) -> torch.Tensor: + values = torch.arange(M * N, dtype=torch.int64).reshape(M, N).remainder(31) + 1 + return values.to(_TORCH_DT[dtype]).contiguous() + + +def _pad_scalar(dtype: DataType, mode: str) -> int | float: + if mode == "zero": + return 0 + torch_dtype = _TORCH_DT[dtype] + if torch_dtype.is_floating_point: + return float("inf") if mode == "max" else float("-inf") + limits = torch.iinfo(torch_dtype) + return limits.max if mode == "max" else limits.min + + +def _make_program( + dtype: DataType, + mode: str, + valid_shape: tuple[int, int], +): + pl_dtype = _PL_DT[dtype] + pad_mode = _PAD_MODES[mode] + valid = list(valid_shape) + + @pl.program + class FillpadInplaceProgram: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + input_tensor: pl.Tensor[[M, N], pl_dtype], + output: pl.Out[pl.Tensor[[M, N], pl_dtype]], + ) -> pl.Tensor[[M, N], pl_dtype]: + tile = pl.load( + input_tensor, + offsets=[0, 0], + shapes=[M, N], + valid_shapes=valid, + ) + padded = pl.tile.fillpad_inplace(tile, pad_value=pad_mode) + return pl.store(padded, offsets=[0, 0], output_tensor=output) + + @pl.function(type=pl.FunctionType.Orchestration) + def orchestrator( + self, + input_tensor: pl.Tensor[[M, N], pl_dtype], + output: pl.Out[pl.Tensor[[M, N], pl_dtype]], + ) -> pl.Tensor[[M, N], pl_dtype]: + return self.kernel(input_tensor, output) + + return FillpadInplaceProgram + + +class FillpadInplaceCase(PTOTestCase): + __test__ = False + + def __init__( self, - input_tensor: pl.Tensor[[48, 64], pl.FP32], - output: pl.Out[pl.Tensor[[64, 64], pl.FP32]], - ) -> pl.Tensor[[64, 64], pl.FP32]: - output = self.fillpad_inplace_min_kernel(input_tensor, output) - return output - - -# --- Test Cases --- - - -class FillpadInplaceZeroTestCase(PTOTestCase): - """Test fillpad_inplace - padding region should be filled with 0.0.""" - - def get_name(self) -> str: - return "fillpad_inplace_zero" - - def get_strategy(self) -> OptimizationStrategy: - return OptimizationStrategy.Default - - def define_tensors(self) -> list[TensorSpec]: - return [ - TensorSpec("input_tensor", [48, 64], DataType.FP32, init_value=torch.randn), - TensorSpec("output", [64, 64], DataType.FP32, is_output=True), - ] - - def get_program(self) -> Any: - return FillpadInplaceZeroProgram - - def compute_expected(self, tensors, params=None): - """Expected: rows 0-47 = input, rows 48-63 = 0.0""" - expected = torch.zeros(64, 64, dtype=torch.float32) - expected[:48, :] = tensors["input_tensor"] - tensors["output"][:] = expected - - -class FillpadInplaceMaxTestCase(PTOTestCase): - """Test fillpad_inplace - padding region should be filled with FP32 max.""" + dtype: DataType, + mode: str, + valid_shape: tuple[int, int], + *, + platform: str, + ): + super().__init__(platform=platform) + self.dtype = dtype + self.mode = mode + self.valid_shape = valid_shape def get_name(self) -> str: - return "fillpad_inplace_max" - - def get_strategy(self) -> OptimizationStrategy: - return OptimizationStrategy.Default + valid_tag = f"v{self.valid_shape[0]}x{self.valid_shape[1]}" + return f"fillpad_inplace_{self.dtype.value}_{self.mode}_{valid_tag}" def define_tensors(self) -> list[TensorSpec]: return [ - TensorSpec("input_tensor", [48, 64], DataType.FP32, init_value=torch.randn), - TensorSpec("output", [64, 64], DataType.FP32, is_output=True), + TensorSpec( + "input_tensor", + [M, N], + self.dtype, + init_value=lambda: _input(self.dtype), + ), + TensorSpec("output", [M, N], self.dtype, is_output=True), ] def get_program(self) -> Any: - return FillpadInplaceMaxProgram + return _make_program(self.dtype, self.mode, self.valid_shape) def compute_expected(self, tensors, params=None): - """Expected: rows 0-47 = input, rows 48-63 = FP32 max (+inf)""" - expected = torch.full((64, 64), float("inf"), dtype=torch.float32) - expected[:48, :] = tensors["input_tensor"] - tensors["output"][:] = expected - - -class FillpadInplaceMinTestCase(PTOTestCase): - """Test fillpad_inplace - padding region should be filled with FP32 min (-inf).""" - - def get_name(self) -> str: - return "fillpad_inplace_min" - - def get_strategy(self) -> OptimizationStrategy: - return OptimizationStrategy.Default - - def define_tensors(self) -> list[TensorSpec]: - return [ - TensorSpec("input_tensor", [48, 64], DataType.FP32, init_value=torch.randn), - TensorSpec("output", [64, 64], DataType.FP32, is_output=True), + valid_rows, valid_cols = self.valid_shape + fill = _pad_scalar(self.dtype, self.mode) + expected = torch.full( + (M, N), + fill, + dtype=_TORCH_DT[self.dtype], + ) + expected[:valid_rows, :valid_cols] = tensors["input_tensor"][ + :valid_rows, :valid_cols ] - - def get_program(self) -> Any: - return FillpadInplaceMinProgram - - def compute_expected(self, tensors, params=None): - """Expected: rows 0-47 = input, rows 48-63 = -inf""" - expected = torch.full((64, 64), float("-inf"), dtype=torch.float32) - expected[:48, :] = tensors["input_tensor"] tensors["output"][:] = expected -# --- Tests --- - - -class TestFillpadInplace: - """Test suite to verify fillpad_inplace fills padding region in place with different pad values.""" - - @pytest.mark.parametrize("platform", PLATFORMS) - def test_fillpad_inplace_zero(self, test_runner, platform): - """Verify fillpad_inplace fills the padding region with 0.0.""" - result = test_runner.run(FillpadInplaceZeroTestCase(platform=platform)) - assert result.passed, f"Test failed: {result.error}" - - @pytest.mark.parametrize("platform", PLATFORMS) - def test_fillpad_inplace_max(self, test_runner, platform): - """Verify fillpad_inplace fills the padding region with FP32 max value.""" - result = test_runner.run(FillpadInplaceMaxTestCase(platform=platform)) - assert result.passed, f"Test failed: {result.error}" - - @pytest.mark.parametrize("platform", PLATFORMS) - def test_fillpad_inplace_min(self, test_runner, platform): - """Verify fillpad_inplace fills the padding region with FP32 min value (-inf).""" - result = test_runner.run(FillpadInplaceMinTestCase(platform=platform)) - assert result.passed, f"Test failed: {result.error}" - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) +def _case( + platform: str, + dtype: DataType, + mode: str, + valid_shape: tuple[int, int], +): + return pytest.param( + platform, + dtype, + mode, + valid_shape, + id=f"{platform}-{dtype.value}-{mode}-v{valid_shape[0]}x{valid_shape[1]}", + ) + + +_CASES = [] +for _platform in ("a2a3", "a5"): + for _dtype in _DTYPES: + for _mode in _PAD_MODES: + _CASES.append( + _case( + _platform, + _dtype, + _mode, + COMBINED_TAIL, + ) + ) + for _mode in _PAD_MODES: + for _valid_shape in (FULL, ROW_TAIL, COL_TAIL): + _CASES.append( + _case( + _platform, + DataType.FP32, + _mode, + _valid_shape, + ) + ) + _CASES.append( + _case( + _platform, + DataType.FP32, + "zero", + MIN_BOUNDARY, + ) + ) + + +@pytest.mark.platforms("a2a3", "a5") +@pytest.mark.parametrize("platform,dtype,mode,valid_shape", _CASES) +def test_fillpad_inplace( + test_runner, + platform, + dtype, + mode, + valid_shape, +): + result = test_runner.run( + FillpadInplaceCase( + dtype, + mode, + valid_shape, + platform=platform, + ) + ) + assert result.passed, f"Test failed: {result.error}" diff --git a/tests/ut/ir/operators/test_concat_idx.py b/tests/ut/ir/operators/test_concat_idx.py index 376fdff178..bfe4e85b2b 100644 --- a/tests/ut/ir/operators/test_concat_idx.py +++ b/tests/ut/ir/operators/test_concat_idx.py @@ -24,29 +24,54 @@ def _tile(name, shape, dtype, valid_shape): return ir.Var(name, ir.TileType(list(shape), dtype, tile_view=view), ir.Span.unknown()) -def test_concat_idx_returns_destination_type(): - src0 = _tile("src0", (8, 64), DataType.FP32, (8, 64)) - src1 = _tile("src1", (8, 64), DataType.FP32, (8, 48)) - idx0 = _tile("idx0", (8, 8), DataType.INT32, (8, 1)) - idx1 = _tile("idx1", (8, 8), DataType.INT32, (8, 1)) - dst = _tile("dst", (8, 64), DataType.FP32, (8, 64)) +@pytest.mark.parametrize( + "data_dtype", + [ + DataType.INT8, + DataType.UINT8, + DataType.INT16, + DataType.UINT16, + DataType.INT32, + DataType.UINT32, + DataType.FP16, + DataType.BF16, + DataType.FP32, + ], +) +@pytest.mark.parametrize( + "index_dtype", + [ + DataType.INT8, + DataType.UINT8, + DataType.INT16, + DataType.UINT16, + DataType.INT32, + DataType.UINT32, + ], +) +def test_concat_idx_returns_destination_type(data_dtype, index_dtype): + src0 = _tile("src0", (8, 64), data_dtype, (8, 64)) + src1 = _tile("src1", (8, 64), data_dtype, (8, 48)) + idx0 = _tile("idx0", (8, 8), index_dtype, (8, 1)) + idx1 = _tile("idx1", (8, 8), index_dtype, (8, 1)) + dst = _tile("dst", (8, 64), data_dtype, (8, 64)) call = tile.concat_idx(src0, src1, idx0, idx1, dst) assert isinstance(call.type, ir.TileType) assert isinstance(dst.type, ir.TileType) - assert call.type.dtype == DataType.FP32 + assert call.type.dtype == data_dtype assert call.type.shape == dst.type.shape def test_concat_idx_rejects_invalid_index_contract(): src = _tile("src", (8, 64), DataType.FP32, (8, 64)) dst = _tile("dst", (8, 64), DataType.FP32, (8, 64)) - bad_dtype = _tile("bad", (8, 8), DataType.UINT32, (8, 1)) + bad_dtype = _tile("bad", (8, 8), DataType.FP32, (8, 1)) bad_cols = _tile("wide", (8, 8), DataType.INT32, (8, 2)) good = _tile("good", (8, 8), DataType.INT32, (8, 1)) - with pytest.raises(ValueError, match="INT32"): + with pytest.raises(ValueError, match="integer index"): tile.concat_idx(src, src, bad_dtype, good, dst) with pytest.raises(ValueError, match="columns equal to 1"): tile.concat_idx(src, src, bad_cols, good, dst)