diff --git a/docs/en/dev/codegen/00-pto_codegen.md b/docs/en/dev/codegen/00-pto_codegen.md index fb793cf4c2..22dc135339 100644 --- a/docs/en/dev/codegen/00-pto_codegen.md +++ b/docs/en/dev/codegen/00-pto_codegen.md @@ -142,6 +142,9 @@ print(pto_code) | `tile.mul(lhs, rhs)` | `pto.tmul` | | `tile.add(a, b, c)` | `pto.taddc` (3-operand add) | | `tile.adds(tile, scalar)` | `pto.tadds` (tile + scalar) | +| `tile.part_argmax(src0, src1, src0_idx, src1_idx)` | `pto.tpartargmax` (selected value and paired index outputs) | +| `tile.part_argmin(src0, src1, src0_idx, src1_idx)` | `pto.tpartargmin` (selected value and paired index outputs) | +| `tile.histogram(src, idx, byte=...)` | `pto.thistogram` (A5 cumulative 256-bin histogram) | | `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.slice` / `tile.assemble` lowering details.** Both ops are lowered diff --git a/docs/en/dev/ptoas-op-status.md b/docs/en/dev/ptoas-op-status.md index 6de3c36be3..1ff3f7ad25 100644 --- a/docs/en/dev/ptoas-op-status.md +++ b/docs/en/dev/ptoas-op-status.md @@ -97,8 +97,8 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT. | pto.tpartadd | TPARTADD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tpartmax | TPARTMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tpartmin | TPARTMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| 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.tpartargmax | TPARTARGMAX | tile | ✅ | ✅ | ✅ | ✅ | — | verified on A2/A3 hardware | +| pto.tpartargmin | TPARTARGMIN | tile | ✅ | ✅ | ✅ | ✅ | — | verified on A2/A3 hardware | | pto.tpartmul | TPARTMUL | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tprelu | TPRELU | tile | ✅ | ✅ | ❌ | ❌ | — | path exists; historical ISA/semantic issue requires revalidation against the current pin | | pto.tadds | TADDS | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | @@ -136,7 +136,7 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT. | pto.trowmin | TROWMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | verified on A2/A3 hardware; A5 hardware verification pending | | pto.trowargmin | TROWARGMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.trowprod | TROWPROD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.thistogram | THISTOGRAM | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path | +| pto.thistogram | THISTOGRAM | tile | ✅ | ✅ | ✅ | ❌ | — | A5-only; same-name A5sim PTOAS compilation coverage added, A5 execution pending | | pto.tcolsum | TCOLSUM | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tcolmax | TCOLMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tcolargmax | TCOLARGMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | @@ -273,6 +273,6 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT. | **Source Compatibility / Manual Mode (1)** | | | | | | | | | | pto.tassign | TASSIGN | internal | ✅ | — | — | — | — | inactive backend hook; no standalone ST | -**Stats**: 204 public/compatibility PTOAS ops; 113 have a pypto tile frontend and 75 have a tensor frontend; -110 have same-name ST coverage (106 regular STs and 4 distributed STs); 62 lack same-name ST coverage -(52 regular and 10 distributed); within these 204, another 32 ops are not suitable for standalone STs. +**Stats**: 204 public/compatibility PTOAS ops; 116 have a pypto tile frontend and 75 have a tensor frontend; +112 have same-name ST coverage (108 regular STs and 4 distributed STs); 60 lack same-name ST coverage +(50 regular and 10 distributed); within these 204, another 32 ops are not suitable for standalone STs. diff --git a/docs/zh/dev/codegen/00-pto_codegen.md b/docs/zh/dev/codegen/00-pto_codegen.md index 4536c06e0e..b7a1397598 100644 --- a/docs/zh/dev/codegen/00-pto_codegen.md +++ b/docs/zh/dev/codegen/00-pto_codegen.md @@ -140,6 +140,9 @@ print(pto_code) | `tile.mul(lhs, rhs)` | `pto.tmul` | | `tile.add(a, b, c)` | `pto.taddc` (三操作数加法) | | `tile.adds(tile, scalar)` | `pto.tadds` (Tile + 标量) | +| `tile.part_argmax(src0, src1, src0_idx, src1_idx)` | `pto.tpartargmax`（输出选中的值及配对索引） | +| `tile.part_argmin(src0, src1, src0_idx, src1_idx)` | `pto.tpartargmin`（输出选中的值及配对索引） | +| `tile.histogram(src, idx, byte=...)` | `pto.thistogram`（A5 累计 256-bin 直方图） | | `tile.fillpad_expand(src, shape)` | `pto.tfillpad_expand ins(%src) outs(%dst)`（`shape` 元组仅用于类型推导；更大的 `dst` 及其 pad 来自结果类型） | **`tile.slice` / `tile.assemble` 下沉细节。** 两个 op 都通过 `pto.subview` diff --git a/docs/zh/dev/ptoas-op-status.md b/docs/zh/dev/ptoas-op-status.md index c982ff35fe..fc968b3a29 100644 --- a/docs/zh/dev/ptoas-op-status.md +++ b/docs/zh/dev/ptoas-op-status.md @@ -83,8 +83,8 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO | pto.tpartadd | TPARTADD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tpartmax | TPARTMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tpartmin | TPARTMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.tpartargmax | TPARTARGMAX | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 | -| pto.tpartargmin | TPARTARGMIN | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 | +| pto.tpartargmax | TPARTARGMAX | tile | ✅ | ✅ | ✅ | ✅ | — | A2/A3 真机已验证 | +| pto.tpartargmin | TPARTARGMIN | tile | ✅ | ✅ | ✅ | ✅ | — | A2/A3 真机已验证 | | pto.tpartmul | TPARTMUL | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tprelu | TPRELU | tile | ✅ | ✅ | ❌ | ❌ | — | 已有链路；历史 ISA/语义问题，需按当前 pin 复验 | | pto.tadds | TADDS | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | @@ -122,7 +122,7 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO | pto.trowmin | TROWMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | A2/A3 真机已验证；A5 真机待验证 | | pto.trowargmin | TROWARGMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.trowprod | TROWPROD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | -| pto.thistogram | THISTOGRAM | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 | +| pto.thistogram | THISTOGRAM | tile | ✅ | ✅ | ✅ | ❌ | — | 仅 A5；已增加同名 A5sim PTOAS 编译覆盖，A5 执行待验证 | | pto.tcolsum | TCOLSUM | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tcolmax | TCOLMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | | pto.tcolargmax | TCOLARGMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | | @@ -259,6 +259,6 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO | **源码兼容 / 手动模式（1）** | | | | | | | | | | pto.tassign | TASSIGN | internal | ✅ | — | — | — | — | 失活 backend hook，不独立建 ST | -**统计**：共 204 个 PTOAS 公共/兼容 op；pypto tile 前端 113 个，tensor 前端 75 个； -同名 ST 覆盖 110 个（普通 ST 106，distributed ST 4）；无同名 ST 62 个 -（普通 52，distributed 10）；这 204 个中另有 32 个 op 不适合独立 ST。 +**统计**：共 204 个 PTOAS 公共/兼容 op；pypto tile 前端 116 个，tensor 前端 75 个； +同名 ST 覆盖 112 个（普通 ST 108，distributed ST 4）；无同名 ST 60 个 +（普通 50，distributed 10）；这 204 个中另有 32 个 op 不适合独立 ST。 diff --git a/python/pypto/ir/op/tile_ops.py b/python/pypto/ir/op/tile_ops.py index de166b9c76..02a2159447 100644 --- a/python/pypto/ir/op/tile_ops.py +++ b/python/pypto/ir/op/tile_ops.py @@ -935,6 +935,38 @@ def part_min(src0: Expr, src1: Expr, span: Span | None = None) -> Call: return _ir_core.create_op_call("tile.part_min", [src0, src1], {}, actual_span) +def part_argmax( + src0: Expr, + src1: Expr, + src0_idx: Expr, + src1_idx: Expr, + span: Span | None = None, +) -> Call: + """Partial element-wise maximum returning selected values and indices.""" + actual_span = _get_span_or_capture(span) + return _ir_core.create_op_call("tile.part_argmax", [src0, src1, src0_idx, src1_idx], {}, actual_span) + + +def part_argmin( + src0: Expr, + src1: Expr, + src0_idx: Expr, + src1_idx: Expr, + span: Span | None = None, +) -> Call: + """Partial element-wise minimum returning selected values and indices.""" + actual_span = _get_span_or_capture(span) + return _ir_core.create_op_call("tile.part_argmin", [src0, src1, src0_idx, src1_idx], {}, actual_span) + + +def histogram(src: Expr, idx: Expr, byte: int = 1, span: Span | None = None) -> Call: + """A5 per-row cumulative histogram for one byte of a UINT16/UINT32 tile.""" + if not isinstance(byte, int) or isinstance(byte, bool) or not 0 <= byte <= 3: + raise ValueError(f"histogram byte must be an integer in [0, 3], got {byte!r}") + actual_span = _get_span_or_capture(span) + return _ir_core.create_op_call("tile.histogram", [src, idx], {"byte": byte}, actual_span) + + def fmod(lhs: Expr, rhs: Expr, span: Span | None = None) -> Call: """Element-wise floating-point remainder of two tiles. diff --git a/python/pypto/language/__init__.py b/python/pypto/language/__init__.py index 71b90da86b..9a70ddd6ff 100644 --- a/python/pypto/language/__init__.py +++ b/python/pypto/language/__init__.py @@ -127,6 +127,7 @@ def scalar_func(x: pl.Scalar[pl.FP32]) -> pl.Scalar[pl.FP32]: gemv, gemv_acc, gemv_bias, + histogram, load, lrelu, matmul_bias, @@ -138,6 +139,8 @@ def scalar_func(x: pl.Scalar[pl.FP32]) -> pl.Scalar[pl.FP32]: not_, or_, ors, + part_argmax, + part_argmin, prelu, relu, rem, @@ -320,6 +323,9 @@ def scalar_func(x: pl.Scalar[pl.FP32]) -> pl.Scalar[pl.FP32]: "part_mul", "part_max", "part_min", + "part_argmax", + "part_argmin", + "histogram", "maximum", "exp", "log", diff --git a/python/pypto/language/op/__init__.py b/python/pypto/language/op/__init__.py index 6c42ee6087..a3f19df5cb 100644 --- a/python/pypto/language/op/__init__.py +++ b/python/pypto/language/op/__init__.py @@ -79,6 +79,7 @@ gemv, gemv_acc, gemv_bias, + histogram, load, lrelu, matmul_bias, @@ -91,6 +92,8 @@ not_, or_, ors, + part_argmax, + part_argmin, prelu, relu, rem, @@ -193,6 +196,9 @@ "part_mul", "part_max", "part_min", + "part_argmax", + "part_argmin", + "histogram", "maximum", "min", "max", diff --git a/python/pypto/language/op/tile_ops.py b/python/pypto/language/op/tile_ops.py index 4a995c5ceb..aef0172c0e 100644 --- a/python/pypto/language/op/tile_ops.py +++ b/python/pypto/language/op/tile_ops.py @@ -120,6 +120,9 @@ "part_mul", "part_max", "part_min", + "part_argmax", + "part_argmin", + "histogram", "fmod", "fmods", "and_", @@ -2039,6 +2042,31 @@ def part_min(src0: Tile, src1: Tile) -> Tile: return Tile(expr=call_expr) +def part_argmax(src0: Tile, src1: Tile, src0_idx: Tile, src1_idx: Tile) -> tuple[Tile, Tile]: + """Partial element-wise maximum returning ``(value, index)`` tiles.""" + call_expr = _ir_ops.part_argmax(src0.unwrap(), src1.unwrap(), src0_idx.unwrap(), src1_idx.unwrap()) + span = call_expr.span + return ( + Tile(expr=_ir_core.TupleGetItemExpr(call_expr, 0, span)), + Tile(expr=_ir_core.TupleGetItemExpr(call_expr, 1, span)), + ) + + +def part_argmin(src0: Tile, src1: Tile, src0_idx: Tile, src1_idx: Tile) -> tuple[Tile, Tile]: + """Partial element-wise minimum returning ``(value, index)`` tiles.""" + call_expr = _ir_ops.part_argmin(src0.unwrap(), src1.unwrap(), src0_idx.unwrap(), src1_idx.unwrap()) + span = call_expr.span + return ( + Tile(expr=_ir_core.TupleGetItemExpr(call_expr, 0, span)), + Tile(expr=_ir_core.TupleGetItemExpr(call_expr, 1, span)), + ) + + +def histogram(src: Tile, idx: Tile, byte: int = 1) -> Tile: + """A5 per-row cumulative histogram for a selected source byte.""" + return Tile(expr=_ir_ops.histogram(src.unwrap(), idx.unwrap(), byte=byte)) + + def fmod(lhs: Tile, rhs: Tile) -> Tile: """Element-wise floating-point remainder of two tiles. diff --git a/src/backend/common/pto_ops_elementwise.cpp b/src/backend/common/pto_ops_elementwise.cpp index 70f8c99ebe..18ce0bc163 100644 --- a/src/backend/common/pto_ops_elementwise.cpp +++ b/src/backend/common/pto_ops_elementwise.cpp @@ -14,7 +14,10 @@ * @brief PTO codegen registration for elementwise / compute tile ops. */ +#include +#include #include +#include #include #include #include @@ -31,6 +34,7 @@ #include "pypto/ir/kind_traits.h" #include "pypto/ir/scalar_expr.h" #include "pypto/ir/tile_view_semantics.h" +#include "pypto/ir/transforms/utils/memref_utils.h" #include "pypto/ir/type.h" #include "src/backend/common/pto_ops_internal.h" @@ -348,6 +352,88 @@ static std::string MakePrintCodegenPTO(const std::string& pto_op_name, const Cal return ""; } +static std::string MakePartArgCodegenPTO(const std::string& pto_op_name, const CallPtr& op, + codegen::CodegenBase& codegen_base) { + auto& codegen = AsPto(codegen_base); + CHECK(op->args_.size() == 4) << op->op_->name_ + << " requires 4 arguments (src0, src1, src0_idx, src1_idx), but got " + << op->args_.size(); + + auto tuple_var = codegen.GetCurrentResultVar(); + INTERNAL_CHECK_SPAN(tuple_var, op->span_) + << "Internal error: " << op->op_->name_ << " codegen requires current_result_var"; + auto element_vars = codegen.ResolveTupleResultElements(tuple_var, 2); + INTERNAL_CHECK_SPAN(element_vars[0] && element_vars[1], op->span_) + << "Internal error: " << op->op_->name_ << " requires two TupleGetItemExpr consumers"; + + std::array, 2> result_types; + for (size_t i = 0; i < result_types.size(); ++i) { + result_types[i] = ir::GetTileTypeWithMemRef(element_vars[i]->GetType()); + INTERNAL_CHECK_SPAN(result_types[i], element_vars[i]->span_) + << "Internal error: " << op->op_->name_ << " result " << i + << " must have TileType with MemRef set by InitMemRef"; + codegen.EmitAllocTileForVar(element_vars[i], result_types[i]); + } + + std::array inputs; + std::array input_types; + for (size_t i = 0; i < inputs.size(); ++i) { + inputs[i] = codegen.GetExprAsCode(op->args_[i]); + input_types[i] = codegen.GetExprTypeAnnotation(op->args_[i]); + INTERNAL_CHECK_SPAN(!inputs[i].empty(), op->args_[i]->span_) + << "Internal error: " << op->op_->name_ << " input " << i << " requires an SSA value"; + } + const std::string dst = codegen.GetVarName(element_vars[0]); + const std::string dst_idx = codegen.GetVarName(element_vars[1]); + const std::string dst_type = codegen.GetTileBufTypeStringFromTileType(result_types[0]); + const std::string dst_idx_type = codegen.GetTileBufTypeStringFromTileType(result_types[1]); + INTERNAL_CHECK_SPAN(!dst.empty() && !dst_idx.empty(), op->span_) + << "Internal error: " << op->op_->name_ << " requires value and index outputs"; + + std::ostringstream oss; + oss << pto_op_name << " ins(" << inputs[0] << ", " << inputs[1] << ", " << inputs[2] << ", " << inputs[3]; + if (std::all_of(input_types.begin(), input_types.end(), + [](const std::string& type) { return !type.empty(); })) { + oss << " : " << input_types[0] << ", " << input_types[1] << ", " << input_types[2] << ", " + << input_types[3]; + } + oss << ") outs(" << dst << ", " << dst_idx; + if (!dst_type.empty() && !dst_idx_type.empty()) { + oss << " : " << dst_type << ", " << dst_idx_type; + } + oss << ")"; + codegen.Emit(oss.str()); + return ""; +} + +static std::string MakeHistogramCodegenPTO(const CallPtr& op, codegen::CodegenBase& codegen_base) { + auto& codegen = AsPto(codegen_base); + CHECK(op->args_.size() == 2) << "tile.histogram requires 2 arguments (src, idx), but got " + << op->args_.size(); + const int byte = op->GetKwarg("byte", 1); + std::string src = codegen.GetExprAsCode(op->args_[0]); + std::string idx = codegen.GetExprAsCode(op->args_[1]); + std::string src_type = codegen.GetExprTypeAnnotation(op->args_[0]); + std::string idx_type = codegen.GetExprTypeAnnotation(op->args_[1]); + std::string dst = codegen.GetCurrentResultTarget(); + std::string dst_type = codegen.GetCurrentResultTileBufTypeString(); + INTERNAL_CHECK_SPAN(!src.empty() && !idx.empty() && !dst.empty(), op->span_) + << "Internal error: tile.histogram requires input and output SSA values"; + + std::ostringstream oss; + oss << "pto.thistogram ins(" << src << ", " << idx; + if (!src_type.empty() && !idx_type.empty()) { + oss << " : " << src_type << ", " << idx_type; + } + oss << ") outs(" << dst; + if (!dst_type.empty()) { + oss << " : " << dst_type; + } + oss << ") {byte = " << byte << " : i32}"; + codegen.Emit(oss.str()); + return ""; +} + struct SimpleOpEntry { const char* op_name; const char* pto_op_name; @@ -497,6 +583,22 @@ void RegisterElementwiseOps(Backend& backend, const std::unordered_set 0) return; + auto entry = backend.RegisterOp(op_name); + entry.f_codegen([pto_op = std::string(pto_op_name)](const CallPtr& op, codegen::CodegenBase& codegen) { + return MakePartArgCodegenPTO(pto_op, op, codegen); + }); + for (size_t i = 0; i < 4; ++i) { + entry.set_input_layout(i, ir::TileLayout::row_major); + } + entry.set_output_layout(ir::TileLayout::row_major); + }; + register_part_arg("tile.part_argmax", "pto.tpartargmax"); + register_part_arg("tile.part_argmin", "pto.tpartargmin"); + reg("tile.histogram", + [](const CallPtr& op, codegen::CodegenBase& codegen) { return MakeHistogramCodegenPTO(op, codegen); }); + auto register_precision_op = [&](const char* op_name, const char* pto_op_name, size_t arity, const char* attr_kind) { if (exclude_ops.count(op_name) > 0) return; diff --git a/src/ir/op/tile_ops/elementwise.cpp b/src/ir/op/tile_ops/elementwise.cpp index 8a5a292b3c..a05e36c5a3 100644 --- a/src/ir/op/tile_ops/elementwise.cpp +++ b/src/ir/op/tile_ops/elementwise.cpp @@ -21,6 +21,7 @@ */ #include +#include #include #include #include @@ -37,6 +38,7 @@ #include "pypto/ir/op_registry.h" #include "pypto/ir/scalar_expr.h" #include "pypto/ir/span.h" +#include "pypto/ir/tile_view_semantics.h" #include "pypto/ir/type.h" #include "pypto/ir/type_inference.h" @@ -449,6 +451,179 @@ REGISTER_OP("tile.part_min") return DeduceTileOpElementwiseBinaryType(args, kwargs, "tile.part_min"); }); +static TypePtr DeduceTilePartArgType(const std::vector& args, const std::string& op_name) { + CHECK(args.size() == 4) << "The operator " << op_name + << " requires 4 arguments (src0, src1, src0_idx, src1_idx), but got " + << args.size(); + + std::array, 4> types; + for (size_t i = 0; i < args.size(); ++i) { + types[i] = As(args[i]->GetType()); + CHECK(types[i]) << "The operator " << op_name << " requires argument " << i + << " to be a TileType, but got " << args[i]->GetType()->TypeName(); + CHECK(types[i]->shape_.size() == 2) << "The operator " << op_name << " requires rank-2 tile operands"; + } + + CHECK(types[0]->dtype_ == DataType::FP16 || types[0]->dtype_ == DataType::FP32) + << "The operator " << op_name << " requires value dtype in {FP16, FP32}, but got " + << types[0]->dtype_.ToString(); + CHECK(types[1]->dtype_ == types[0]->dtype_) + << "The operator " << op_name << " requires src0 and src1 to have the same value dtype"; + CHECK(types[2]->dtype_ == DataType::INT32 || types[2]->dtype_ == DataType::UINT32) + << "The operator " << op_name << " requires index dtype in {INT32, UINT32}, but got " + << types[2]->dtype_.ToString(); + CHECK(types[3]->dtype_ == types[2]->dtype_) + << "The operator " << op_name << " requires src0_idx and src1_idx to have the same index dtype"; + + for (size_t operand = 1; operand < types.size(); ++operand) { + for (size_t dim = 0; dim < types[0]->shape_.size(); ++dim) { + CHECK(ProveValidExtentEqual(types[0]->shape_[dim], types[operand]->shape_[dim]) == ProofResult::kTrue) + << "The operator " << op_name << " requires all value and index operands to have the same shape"; + } + } + + const auto src0_valid = GetValidShape(types[0]); + const auto src1_valid = GetValidShape(types[1]); + const auto src0_idx_valid = GetValidShape(types[2]); + const auto src1_idx_valid = GetValidShape(types[3]); + for (size_t dim = 0; dim < src0_valid.size(); ++dim) { + CHECK(ProveValidExtentEqual(src0_valid[dim], src0_idx_valid[dim]) == ProofResult::kTrue) + << "The operator " << op_name << " requires src0 and src0_idx to have the same valid_shape"; + CHECK(ProveValidExtentEqual(src1_valid[dim], src1_idx_valid[dim]) == ProofResult::kTrue) + << "The operator " << op_name << " requires src1 and src1_idx to have the same valid_shape"; + } + + bool src0_dominates = true; + bool src1_dominates = true; + for (size_t dim = 0; dim < src0_valid.size(); ++dim) { + src0_dominates &= ProveValidExtentLessEqual(src1_valid[dim], src0_valid[dim]) == ProofResult::kTrue; + src1_dominates &= ProveValidExtentLessEqual(src0_valid[dim], src1_valid[dim]) == ProofResult::kTrue; + } + CHECK(src0_dominates || src1_dominates) + << "The operator " << op_name + << " requires one source valid_shape to contain the other source valid_shape"; + const bool use_src0 = src0_dominates; + const auto& result_valid = use_src0 ? src0_valid : src1_valid; + + TileView value_view; + value_view.valid_shape = result_valid; + InheritTileViewLayout(value_view, use_src0 ? types[0] : types[1]); + auto value_type = std::make_shared(types[0]->shape_, types[0]->dtype_, std::nullopt, value_view); + + TileView index_view; + index_view.valid_shape = result_valid; + InheritTileViewLayout(index_view, use_src0 ? types[2] : types[3]); + auto index_type = std::make_shared(types[2]->shape_, types[2]->dtype_, std::nullopt, index_view); + return std::make_shared(std::vector{value_type, index_type}); +} + +REGISTER_OP("tile.part_argmax") + .set_op_category("TileOp") + .set_description("Partial element-wise maximum returning selected values and indices") + .add_argument("src0", "First value tile (FP16/FP32)") + .add_argument("src1", "Second value tile (same dtype/shape as src0)") + .add_argument("src0_idx", "Indices paired with src0 (INT32/UINT32)") + .add_argument("src1_idx", "Indices paired with src1 (same dtype/shape as src0_idx)") + .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_output_memory(MemorySpace::Vec) + .not_inplace_safe() + .f_deduce_type([](const std::vector& args, + const std::vector>& /*kwargs*/) { + return DeduceTilePartArgType(args, "tile.part_argmax"); + }); + +REGISTER_OP("tile.part_argmin") + .set_op_category("TileOp") + .set_description("Partial element-wise minimum returning selected values and indices") + .add_argument("src0", "First value tile (FP16/FP32)") + .add_argument("src1", "Second value tile (same dtype/shape as src0)") + .add_argument("src0_idx", "Indices paired with src0 (INT32/UINT32)") + .add_argument("src1_idx", "Indices paired with src1 (same dtype/shape as src0_idx)") + .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_output_memory(MemorySpace::Vec) + .not_inplace_safe() + .f_deduce_type([](const std::vector& args, + const std::vector>& /*kwargs*/) { + return DeduceTilePartArgType(args, "tile.part_argmin"); + }); + +static TypePtr DeduceTileHistogramType(const std::vector& args, + const std::vector>& kwargs) { + CHECK(args.size() == 2) << "The operator tile.histogram requires 2 arguments (src, idx), but got " + << args.size(); + auto src = As(args[0]->GetType()); + auto idx = As(args[1]->GetType()); + CHECK(src && idx) << "The operator tile.histogram requires src and idx to be TileType operands"; + CHECK(src->shape_.size() == 2 && idx->shape_.size() == 2) + << "The operator tile.histogram requires rank-2 src and idx tiles"; + CHECK(src->dtype_ == DataType::UINT16 || src->dtype_ == DataType::UINT32) + << "The operator tile.histogram requires src dtype in {UINT16, UINT32}, but got " + << src->dtype_.ToString(); + CHECK(idx->dtype_ == DataType::UINT8) + << "The operator tile.histogram requires UINT8 idx, but got " << idx->dtype_.ToString(); + + const int byte = GetKwargOr(kwargs, "byte", 1); + CHECK(byte >= 0 && byte <= 3) << "The operator tile.histogram requires byte in [0, 3], but got " << byte; + + const TileView src_view = tile_view_semantics::GetEffectiveTileView(*src); + const TileView idx_view = tile_view_semantics::GetEffectiveTileView(*idx); + CHECK(src_view.blayout == TileLayout::row_major && src_view.slayout == TileLayout::none_box) + << "The operator tile.histogram requires row_major + none_box src layout"; + + const auto src_valid = GetValidShape(src); + const auto idx_valid = GetValidShape(idx); + if (src->dtype_ == DataType::UINT16) { + CHECK(byte <= 1) << "The operator tile.histogram only supports byte 0 or 1 for UINT16 src"; + CHECK(idx_view.blayout == TileLayout::col_major && idx_view.slayout == TileLayout::none_box) + << "The operator tile.histogram requires col_major + none_box idx layout for UINT16 src"; + CHECK(ProveValidExtentEqual(src->shape_[0], idx->shape_[0]) == ProofResult::kTrue && + ProveValidExtentEqual(src_valid[0], idx_valid[0]) == ProofResult::kTrue) + << "The operator tile.histogram requires idx rows/valid rows to match UINT16 src"; + const auto idx_cols = GetConstantDimension(idx->shape_[1]); + CHECK(idx_cols && *idx_cols == 1) + << "The operator tile.histogram requires UINT16 idx to have exactly one column"; + } else { + CHECK(idx_view.blayout == TileLayout::row_major && idx_view.slayout == TileLayout::none_box) + << "The operator tile.histogram requires row_major + none_box idx layout for UINT32 src"; + const int expected_rows = byte == 0 ? 3 : byte == 1 ? 2 : 1; + const auto idx_rows = GetConstantDimension(idx->shape_[0]); + CHECK(idx_rows && *idx_rows == expected_rows) + << "The operator tile.histogram requires UINT32 idx rows to equal " << expected_rows << " for byte " + << byte; + CHECK(ProveValidExtentEqual(src->shape_[1], idx->shape_[1]) == ProofResult::kTrue) + << "The operator tile.histogram requires UINT32 idx columns to match src columns"; + } + + auto bins = std::make_shared(256, DataType::INDEX, Span::unknown()); + TileView result_view; + result_view.blayout = TileLayout::row_major; + result_view.slayout = TileLayout::none_box; + result_view.valid_shape = {src_valid[0], bins}; + return std::make_shared(std::vector{src->shape_[0], bins}, DataType::UINT32, + std::nullopt, result_view); +} + +REGISTER_OP("tile.histogram") + .set_op_category("TileOp") + .set_description("A5 per-row cumulative 256-bin histogram for one selected source byte") + .add_argument("src", "UINT16/UINT32 source tile") + .add_argument("idx", "UINT8 byte-filter index tile") + .set_attr("byte") + .set_input_memory(0, MemorySpace::Vec) + .set_input_memory(1, MemorySpace::Vec) + .set_output_memory(MemorySpace::Vec) + .not_inplace_safe() + .f_deduce_type([](const std::vector& args, + const std::vector>& kwargs) { + return DeduceTileHistogramType(args, kwargs); + }); + REGISTER_OP("tile.fmod") .set_op_category("TileOp") .set_description("Element-wise floating-point remainder of two tiles with broadcasting") diff --git a/src/ir/transforms/memory_reuse_pass.cpp b/src/ir/transforms/memory_reuse_pass.cpp index 2f50fe2c53..c5e395fd33 100644 --- a/src/ir/transforms/memory_reuse_pass.cpp +++ b/src/ir/transforms/memory_reuse_pass.cpp @@ -1564,37 +1564,12 @@ class ForbidAliasCollector : public IRVisitor { void VisitStmt_(const AssignStmtPtr& op) override { if (auto call = As(op->value_); call && call->op_) { - const auto& reg = OpRegistry::GetInstance(); - if (reg.IsRegistered(call->op_->name_)) { - const auto& entry = reg.GetEntry(call->op_->name_); - auto rep_it = member_to_rep_.find(op->var_.get()); - const Var* out_key = rep_it != member_to_rep_.end() ? rep_it->second : op->var_.get(); - auto forbid_arg = [&](size_t i) { - if (i < call->args_.size()) { - if (auto v = AsVarLike(call->args_[i])) forbidden_[out_key].push_back(v); - } - }; - 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); - } else { - for (size_t i : entry.ForbidOutputAliasArgs()) forbid_arg(i); - } - // A dtype-widening cast (output element wider than its input) cannot run - // in place: element i is read at i*in_bytes but written at i*out_bytes, - // so with out_bytes > in_bytes the write cursor outruns the read cursor - // and clobbers input elements not yet converted -> corrupt results. - // Narrowing / same-width casts are in-place-safe and keep the cross-dtype - // reuse the removed gate enables, so forbid only the widening direction. - if (IsOp(call, "tile.cast") && !call->args_.empty()) { - auto out_t = As(op->var_->GetType()); - auto in_t = As(call->args_[0]->GetType()); - if (out_t && in_t && out_t->dtype_.GetBit() > in_t->dtype_.GetBit()) forbid_arg(0); - } - // tile.transpose is registered not_inplace_safe(), so its output is - // already forbidden from aliasing any input above (pto.ttrans writes - // dst directly from src on the scalar path — dst == src corrupts). - } + if (As(op->var_->GetType())) tuple_calls_[op->var_.get()] = call; + RecordForCall(op->var_, call); + } else if (auto tuple_get = As(op->value_)) { + auto tuple_var = AsVarLike(tuple_get->tuple_); + auto call_it = tuple_var ? tuple_calls_.find(tuple_var.get()) : tuple_calls_.end(); + if (call_it != tuple_calls_.end()) RecordForCall(op->var_, call_it->second); } IRVisitor::VisitStmt_(op); } @@ -1602,8 +1577,46 @@ class ForbidAliasCollector : public IRVisitor { ForbidAliasMap Take() { return std::move(forbidden_); } private: + void RecordForCall(const VarPtr& output, const CallPtr& call) { + if (!output || !call || !call->op_) return; + if (!As(output->GetType())) return; + const auto& reg = OpRegistry::GetInstance(); + if (!reg.IsRegistered(call->op_->name_)) return; + + const auto& entry = reg.GetEntry(call->op_->name_); + auto rep_it = member_to_rep_.find(output.get()); + const Var* out_key = rep_it != member_to_rep_.end() ? rep_it->second : output.get(); + auto forbid_arg = [&](size_t i) { + if (i < call->args_.size()) { + if (auto v = AsVarLike(call->args_[i])) forbidden_[out_key].push_back(v); + } + }; + if (!entry.IsInplaceSafe()) { + // src != dst required: every scalar or tuple-element output must not alias + // any input operand read by the call. + for (size_t i = 0; i < call->args_.size(); ++i) forbid_arg(i); + } else { + for (size_t i : entry.ForbidOutputAliasArgs()) forbid_arg(i); + } + // A dtype-widening cast (output element wider than its input) cannot run + // in place: element i is read at i*in_bytes but written at i*out_bytes, + // so with out_bytes > in_bytes the write cursor outruns the read cursor + // and clobbers input elements not yet converted -> corrupt results. + // Narrowing / same-width casts are in-place-safe and keep the cross-dtype + // reuse the removed gate enables, so forbid only the widening direction. + if (IsOp(call, "tile.cast") && !call->args_.empty()) { + auto out_t = As(output->GetType()); + auto in_t = As(call->args_[0]->GetType()); + if (out_t && in_t && out_t->dtype_.GetBit() > in_t->dtype_.GetBit()) forbid_arg(0); + } + // tile.transpose is registered not_inplace_safe(), so its output is + // already forbidden from aliasing any input above (pto.ttrans writes + // dst directly from src on the scalar path — dst == src corrupts). + } + ForbidAliasMap forbidden_; std::map member_to_rep_; ///< sharing-group member -> representative + std::map tuple_calls_; ///< tuple result Var -> defining call }; /// True only for Ascend910B AIV split-mode functions, which need the load + diff --git a/src/ir/transforms/utils/dead_code_elimination.cpp b/src/ir/transforms/utils/dead_code_elimination.cpp index 5abe78daca..87dc64daff 100644 --- a/src/ir/transforms/utils/dead_code_elimination.cpp +++ b/src/ir/transforms/utils/dead_code_elimination.cpp @@ -347,8 +347,24 @@ std::vector EliminateDeadCodeCore(const std::vector& stmts, /// Predicate for the default `EliminateDeadCode`: any AssignStmt that is not /// a known side-effect op is a removal candidate. +bool IsMultiTileTupleElement(const AssignStmtPtr& assign) { + auto tuple_get = As(assign ? assign->value_ : ExprPtr{}); + if (!tuple_get) return false; + auto tuple_var = AsVarLike(tuple_get->tuple_); + auto tuple_type = As(tuple_var ? tuple_var->GetType() : TypePtr{}); + if (!tuple_type || tuple_type->types_.size() < 2) return false; + for (const auto& element_type : tuple_type->types_) { + if (!As(element_type)) return false; + } + return true; +} + bool IsRemovableForDefaultDce(const StmtPtr& stmt) { - return std::dynamic_pointer_cast(stmt) != nullptr && !IsSideEffectOp(stmt); + auto assign = std::dynamic_pointer_cast(stmt); + if (!assign || IsSideEffectOp(stmt)) return false; + // PTO instructions with multiple tile outputs must materialize every `outs` + // operand even when a later pass only consumes one result. + return !IsMultiTileTupleElement(assign); } /// Walk an expression tree and report whether any Call or Submit appears. diff --git a/tests/st/runtime/ops/test_histogram.py b/tests/st/runtime/ops/test_histogram.py new file mode 100644 index 0000000000..cca371cc50 --- /dev/null +++ b/tests/st/runtime/ops/test_histogram.py @@ -0,0 +1,222 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- + +"""A5 coverage for the cumulative ``thistogram`` instruction.""" + +from typing import Any + +import pypto.language as pl +import pytest +import torch +from harness.core.harness import DataType, PTOTestCase, TensorSpec + +M = 32 +N = 32 + + +def _src16() -> torch.Tensor: + rows = torch.arange(M, dtype=torch.int32).reshape(M, 1) + cols = torch.arange(N, dtype=torch.int32).reshape(1, N) + return ((rows << 8) | ((cols * 17 + rows) & 0xFF)).to(torch.uint16).contiguous() + + +def _idx16() -> torch.Tensor: + return torch.arange(M, dtype=torch.uint8).reshape(1, M).contiguous() + + +def _src32() -> torch.Tensor: + rows = torch.arange(M, dtype=torch.int64).reshape(M, 1) + cols = torch.arange(N, dtype=torch.int64).reshape(1, N) + low = (cols * 13 + rows * 7 + 5) & 0xFF + high = torch.where(cols.remainder(3) == 0, 0x99, 0x12) + return ((high << 24) | (0x34 << 16) | (0x56 << 8) | low).to(torch.uint32).contiguous() + + +def _idx32(rows: int) -> torch.Tensor: + values = torch.tensor([0x12, 0x34, 0x56], dtype=torch.uint8).reshape(3, 1) + return values[:rows].expand(rows, N).contiguous() + + +def _histogram16(byte: int, valid_shape: tuple[int, int]): + valid_rows, valid_cols = valid_shape + + @pl.program + class Histogram16: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[M, N], pl.UINT16], + idx: pl.Tensor[[1, M], pl.UINT8], + out: pl.Out[pl.Tensor[[M, 256], pl.UINT32]], + ) -> pl.Tensor[[M, 256], pl.UINT32]: + src_tile = pl.load(src, [0, 0], [M, N], valid_shapes=[valid_rows, valid_cols]) + idx_row = pl.load(idx, [0, 0], [1, M], valid_shapes=[1, valid_rows]) + idx_col = pl.tile.reshape(idx_row, [M, 1]) + result = pl.tile.histogram(src_tile, idx_col, byte=byte) + return pl.store(result, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def orchestrator( + self, + src: pl.Tensor[[M, N], pl.UINT16], + idx: pl.Tensor[[1, M], pl.UINT8], + out: pl.Out[pl.Tensor[[M, 256], pl.UINT32]], + ) -> pl.Tensor[[M, 256], pl.UINT32]: + return self.kernel(src, idx, out) + + return Histogram16 + + +def _histogram32(byte: int, idx_rows: int, valid_shape: tuple[int, int]): + valid_rows, valid_cols = valid_shape + + @pl.program + class Histogram32: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src: pl.Tensor[[M, N], pl.UINT32], + idx: pl.Tensor[[idx_rows, N], pl.UINT8], + out: pl.Out[pl.Tensor[[M, 256], pl.UINT32]], + ) -> pl.Tensor[[M, 256], pl.UINT32]: + src_tile = pl.load(src, [0, 0], [M, N], valid_shapes=[valid_rows, valid_cols]) + idx_tile = pl.load( + idx, + [0, 0], + [idx_rows, N], + valid_shapes=[idx_rows, valid_cols], + ) + result = pl.tile.histogram(src_tile, idx_tile, byte=byte) + return pl.store(result, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def orchestrator( + self, + src: pl.Tensor[[M, N], pl.UINT32], + idx: pl.Tensor[[idx_rows, N], pl.UINT8], + out: pl.Out[pl.Tensor[[M, 256], pl.UINT32]], + ) -> pl.Tensor[[M, 256], pl.UINT32]: + return self.kernel(src, idx, out) + + return Histogram32 + + +def _cumulative(values: torch.Tensor) -> torch.Tensor: + counts = torch.bincount(values.to(torch.int64), minlength=256) + return torch.cumsum(counts, dim=0).to(torch.uint32) + + +class HistogramTestCase(PTOTestCase): + __test__ = False + + def __init__( + self, + dtype: DataType, + byte: int, + valid_shape: tuple[int, int], + *, + platform=None, + config=None, + ): + super().__init__(config, platform=platform) + self._dtype = dtype + self._byte = byte + self._valid_shape = valid_shape + + def get_name(self) -> str: + dtype_name = "uint16" if self._dtype == DataType.UINT16 else "uint32" + valid_tag = f"v{self._valid_shape[0]}x{self._valid_shape[1]}" + return f"histogram_{dtype_name}_byte{self._byte}_{valid_tag}" + + def define_tensors(self) -> list[TensorSpec]: + if self._dtype == DataType.UINT16: + return [ + TensorSpec("src", [M, N], DataType.UINT16, init_value=_src16), + TensorSpec("idx", [1, M], DataType.UINT8, init_value=_idx16), + TensorSpec("out", [M, 256], DataType.UINT32, is_output=True), + ] + rows = 3 if self._byte == 0 else 2 if self._byte == 1 else 1 + return [ + TensorSpec("src", [M, N], DataType.UINT32, init_value=_src32), + TensorSpec("idx", [rows, N], DataType.UINT8, init_value=lambda: _idx32(rows)), + TensorSpec("out", [M, 256], DataType.UINT32, is_output=True), + ] + + def get_program(self) -> Any: + if self._dtype == DataType.UINT16: + return _histogram16(self._byte, self._valid_shape) + rows = 3 if self._byte == 0 else 2 if self._byte == 1 else 1 + return _histogram32(self._byte, rows, self._valid_shape) + + def compute_expected(self, tensors, params=None): + valid_rows, valid_cols = self._valid_shape + src = tensors["src"].to(torch.int64) + tensors["out"].zero_() + if self._dtype == DataType.UINT16: + for row in range(valid_rows): + row_src = src[row, :valid_cols] + values = (row_src >> (8 * self._byte)) & 0xFF + if self._byte == 0: + values = values[((row_src >> 8) & 0xFF) == row] + tensors["out"][row] = _cumulative(values) + return + + for row in range(valid_rows): + row_src = src[row, :valid_cols] + values = (row_src >> (8 * self._byte)) & 0xFF + if self._byte < 3: + selected = torch.ones(valid_cols, dtype=torch.bool) + for filter_byte in range(self._byte + 1, 4): + idx_row = 3 - filter_byte + selected &= ((row_src >> (8 * filter_byte)) & 0xFF) == tensors["idx"][ + idx_row, 0 + ].to(torch.int64) + values = values[selected] + tensors["out"][row] = _cumulative(values) + + +_FULL = (M, N) +_ROW_TAIL = (11, N) +_COL_TAIL = (M, 23) +_COMBINED_TAIL = (11, 23) +_CASES = [ + *[ + pytest.param(dtype, byte, valid_shape, id=f"{dtype.value}-byte{byte}-{shape_id}") + for dtype, bytes_ in ( + (DataType.UINT16, (0, 1)), + (DataType.UINT32, (0, 1, 2, 3)), + ) + for byte in bytes_ + for valid_shape, shape_id in ( + (_FULL, "full"), + (_COMBINED_TAIL, "combined-tail"), + ) + ], + *[ + pytest.param(dtype, 0, valid_shape, id=f"{dtype.value}-byte0-{shape_id}") + for dtype in (DataType.UINT16, DataType.UINT32) + for valid_shape, shape_id in ( + (_ROW_TAIL, "row-tail"), + (_COL_TAIL, "col-tail"), + ) + ], +] + + +@pytest.mark.platforms("a5", "a5sim") +@pytest.mark.parametrize( + "platform", + [pytest.param("a5", id="a5"), pytest.param("a5sim", id="a5sim")], +) +@pytest.mark.parametrize("dtype,byte,valid_shape", _CASES) +def test_histogram(test_runner, platform, dtype, byte, valid_shape): + result = test_runner.run( + HistogramTestCase(dtype, byte, valid_shape, platform=platform) + ) + assert result.passed, f"Test failed: {result.error}" diff --git a/tests/st/runtime/ops/test_partial_arg.py b/tests/st/runtime/ops/test_partial_arg.py new file mode 100644 index 0000000000..d4c4b38016 --- /dev/null +++ b/tests/st/runtime/ops/test_partial_arg.py @@ -0,0 +1,313 @@ +# 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 ``tpartargmax`` and ``tpartargmin``. + +The source valid regions cover equal, source-0-dominant, and +source-1-dominant cases. Within their overlap, source 0 wins ties on the +pinned A2/A3 and A5 implementations; outside it, the only valid source is +copied together with its paired index. +""" + +from typing import Any + +import pypto.language as pl +import pytest +import torch +from harness.core.harness import DataType, PTOTestCase, TensorSpec + +M = 16 +N = 16 + +_PL_VALUE_DTYPES = { + DataType.FP16: pl.FP16, + DataType.FP32: pl.FP32, +} +_PL_INDEX_DTYPES = { + DataType.INT32: pl.INT32, + DataType.UINT32: pl.UINT32, +} + + +def _torch_value_dtype(dtype: DataType) -> torch.dtype: + return torch.float16 if dtype == DataType.FP16 else torch.float32 + + +def _torch_index_dtype(dtype: DataType) -> torch.dtype: + return torch.int32 if dtype == DataType.INT32 else torch.uint32 + + +def _src0(dtype: DataType) -> torch.Tensor: + return ( + (torch.arange(M * N, dtype=torch.float32).reshape(M, N).remainder(11) - 5) + .to(_torch_value_dtype(dtype)) + .contiguous() + ) + + +def _src1(dtype: DataType) -> torch.Tensor: + values = (torch.arange(M * N, dtype=torch.float32).reshape(M, N).remainder(7) - 3).to( + _torch_value_dtype(dtype) + ) + values[:, ::5] = _src0(dtype)[:, ::5] + return values.contiguous() + + +def _idx0(dtype: DataType) -> torch.Tensor: + return ( + torch.arange(M * N, dtype=torch.int64) + .reshape(M, N) + .to(_torch_index_dtype(dtype)) + .contiguous() + ) + + +def _idx1(dtype: DataType) -> torch.Tensor: + return ( + (1000 + torch.arange(M * N, dtype=torch.int64)) + .reshape(M, N) + .to(_torch_index_dtype(dtype)) + .contiguous() + ) + + +def _part_argmax( + value_dtype: DataType, + index_dtype: DataType, + src0_valid: tuple[int, int], + src1_valid: tuple[int, int], +): + pl_value_dtype = _PL_VALUE_DTYPES[value_dtype] + pl_index_dtype = _PL_INDEX_DTYPES[index_dtype] + src0_valid_shape = list(src0_valid) + src1_valid_shape = list(src1_valid) + + @pl.program + class PartArgMax: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src0: pl.Tensor[[M, N], pl_value_dtype], + src1: pl.Tensor[[M, N], pl_value_dtype], + idx0: pl.Tensor[[M, N], pl_index_dtype], + idx1: pl.Tensor[[M, N], pl_index_dtype], + value_out: pl.Out[pl.Tensor[[M, N], pl_value_dtype]], + index_out: pl.Out[pl.Tensor[[M, N], pl_index_dtype]], + ) -> tuple[pl.Tensor[[M, N], pl_value_dtype], pl.Tensor[[M, N], pl_index_dtype]]: + value0 = pl.load(src0, [0, 0], [M, N], valid_shapes=src0_valid_shape) + value1 = pl.load(src1, [0, 0], [M, N], valid_shapes=src1_valid_shape) + index0 = pl.load(idx0, [0, 0], [M, N], valid_shapes=src0_valid_shape) + index1 = pl.load(idx1, [0, 0], [M, N], valid_shapes=src1_valid_shape) + value, index = pl.tile.part_argmax(value0, value1, index0, index1) + value_out = pl.store(value, [0, 0], value_out) + index_out = pl.store(index, [0, 0], index_out) + return value_out, index_out + + @pl.function(type=pl.FunctionType.Orchestration) + def orchestrator( + self, + src0: pl.Tensor[[M, N], pl_value_dtype], + src1: pl.Tensor[[M, N], pl_value_dtype], + idx0: pl.Tensor[[M, N], pl_index_dtype], + idx1: pl.Tensor[[M, N], pl_index_dtype], + value_out: pl.Out[pl.Tensor[[M, N], pl_value_dtype]], + index_out: pl.Out[pl.Tensor[[M, N], pl_index_dtype]], + ) -> tuple[pl.Tensor[[M, N], pl_value_dtype], pl.Tensor[[M, N], pl_index_dtype]]: + return self.kernel(src0, src1, idx0, idx1, value_out, index_out) + + return PartArgMax + + +def _part_argmin( + value_dtype: DataType, + index_dtype: DataType, + src0_valid: tuple[int, int], + src1_valid: tuple[int, int], +): + pl_value_dtype = _PL_VALUE_DTYPES[value_dtype] + pl_index_dtype = _PL_INDEX_DTYPES[index_dtype] + src0_valid_shape = list(src0_valid) + src1_valid_shape = list(src1_valid) + + @pl.program + class PartArgMin: + @pl.function(type=pl.FunctionType.InCore) + def kernel( + self, + src0: pl.Tensor[[M, N], pl_value_dtype], + src1: pl.Tensor[[M, N], pl_value_dtype], + idx0: pl.Tensor[[M, N], pl_index_dtype], + idx1: pl.Tensor[[M, N], pl_index_dtype], + value_out: pl.Out[pl.Tensor[[M, N], pl_value_dtype]], + index_out: pl.Out[pl.Tensor[[M, N], pl_index_dtype]], + ) -> tuple[pl.Tensor[[M, N], pl_value_dtype], pl.Tensor[[M, N], pl_index_dtype]]: + value0 = pl.load(src0, [0, 0], [M, N], valid_shapes=src0_valid_shape) + value1 = pl.load(src1, [0, 0], [M, N], valid_shapes=src1_valid_shape) + index0 = pl.load(idx0, [0, 0], [M, N], valid_shapes=src0_valid_shape) + index1 = pl.load(idx1, [0, 0], [M, N], valid_shapes=src1_valid_shape) + value, index = pl.tile.part_argmin(value0, value1, index0, index1) + value_out = pl.store(value, [0, 0], value_out) + index_out = pl.store(index, [0, 0], index_out) + return value_out, index_out + + @pl.function(type=pl.FunctionType.Orchestration) + def orchestrator( + self, + src0: pl.Tensor[[M, N], pl_value_dtype], + src1: pl.Tensor[[M, N], pl_value_dtype], + idx0: pl.Tensor[[M, N], pl_index_dtype], + idx1: pl.Tensor[[M, N], pl_index_dtype], + value_out: pl.Out[pl.Tensor[[M, N], pl_value_dtype]], + index_out: pl.Out[pl.Tensor[[M, N], pl_index_dtype]], + ) -> tuple[pl.Tensor[[M, N], pl_value_dtype], pl.Tensor[[M, N], pl_index_dtype]]: + return self.kernel(src0, src1, idx0, idx1, value_out, index_out) + + return PartArgMin + + +class PartialArgTestCase(PTOTestCase): + __test__ = False + + def __init__( + self, + op_name: str, + value_dtype: DataType, + index_dtype: DataType, + src0_valid: tuple[int, int], + src1_valid: tuple[int, int], + *, + platform=None, + config=None, + ): + super().__init__(config, platform=platform) + self._op_name = op_name + self._value_dtype = value_dtype + self._index_dtype = index_dtype + self._src0_valid = src0_valid + self._src1_valid = src1_valid + + def get_name(self) -> str: + src0_tag = f"{self._src0_valid[0]}x{self._src0_valid[1]}" + src1_tag = f"{self._src1_valid[0]}x{self._src1_valid[1]}" + return ( + f"{self._op_name}_{self._value_dtype.value}_{self._index_dtype.value}" + f"_s0-{src0_tag}_s1-{src1_tag}" + ) + + def define_tensors(self) -> list[TensorSpec]: + return [ + TensorSpec("src0", [M, N], self._value_dtype, init_value=lambda: _src0(self._value_dtype)), + TensorSpec("src1", [M, N], self._value_dtype, init_value=lambda: _src1(self._value_dtype)), + TensorSpec("idx0", [M, N], self._index_dtype, init_value=lambda: _idx0(self._index_dtype)), + TensorSpec("idx1", [M, N], self._index_dtype, init_value=lambda: _idx1(self._index_dtype)), + TensorSpec("value_out", [M, N], self._value_dtype, is_output=True), + TensorSpec("index_out", [M, N], self._index_dtype, is_output=True), + ] + + def get_program(self) -> Any: + factory = _part_argmax if self._op_name == "part_argmax" else _part_argmin + return factory( + self._value_dtype, + self._index_dtype, + self._src0_valid, + self._src1_valid, + ) + + def compute_expected(self, tensors, params=None): + src0 = tensors["src0"] + src1 = tensors["src1"] + rows = torch.arange(M).reshape(M, 1) + cols = torch.arange(N).reshape(1, N) + valid0 = (rows < self._src0_valid[0]) & (cols < self._src0_valid[1]) + valid1 = (rows < self._src1_valid[0]) & (cols < self._src1_valid[1]) + preferred0 = src0 >= src1 if self._op_name == "part_argmax" else src0 <= src1 + choose0 = valid0 & (~valid1 | preferred0) + + value = torch.zeros_like(src0) + index = torch.zeros_like(tensors["idx0"]) + value[valid1] = src1[valid1] + index[valid1] = tensors["idx1"][valid1] + value[choose0] = src0[choose0] + index[choose0] = tensors["idx0"][choose0] + tensors["value_out"][:] = value + tensors["index_out"][:] = index + + +_DTYPE_PAIRS = [ + (DataType.FP16, DataType.INT32), + (DataType.FP16, DataType.UINT32), + (DataType.FP32, DataType.INT32), + (DataType.FP32, DataType.UINT32), +] +_VALID_SCENARIOS = [ + ((M, N), (M, N), "full"), + ((11, N), (11, N), "row-tail"), + ((M, 11), (M, 11), "col-tail"), + ((11, 13), (11, 13), "combined-tail"), + ((M, N), (11, 13), "src0-dominant"), + ((11, 13), (M, N), "src1-dominant"), +] +_CASES = [ + *[ + pytest.param( + op_name, + value_dtype, + index_dtype, + (11, 13), + (11, 13), + id=f"{op_name}-{value_dtype.value}-{index_dtype.value}-combined", + ) + for op_name in ("part_argmax", "part_argmin") + for value_dtype, index_dtype in _DTYPE_PAIRS + ], + *[ + pytest.param( + op_name, + DataType.FP32, + DataType.INT32, + src0_valid, + src1_valid, + id=f"{op_name}-fp32-int32-{scenario}", + ) + for op_name in ("part_argmax", "part_argmin") + for src0_valid, src1_valid, scenario in _VALID_SCENARIOS + if scenario != "combined-tail" + ], +] + + +@pytest.mark.platforms("a2a3", "a5") +@pytest.mark.parametrize( + "platform", + [pytest.param("a2a3", id="a2a3"), pytest.param("a5", id="a5")], +) +@pytest.mark.parametrize( + "op_name,value_dtype,index_dtype,src0_valid,src1_valid", + _CASES, +) +def test_partial_arg( + test_runner, + platform, + op_name, + value_dtype, + index_dtype, + src0_valid, + src1_valid, +): + result = test_runner.run( + PartialArgTestCase( + op_name, + value_dtype, + index_dtype, + src0_valid, + src1_valid, + platform=platform, + ) + ) + assert result.passed, f"Test failed: {result.error}" diff --git a/tests/ut/ir/operators/test_partial_arg_histogram.py b/tests/ut/ir/operators/test_partial_arg_histogram.py new file mode 100644 index 0000000000..fe27ef1ddb --- /dev/null +++ b/tests/ut/ir/operators/test_partial_arg_histogram.py @@ -0,0 +1,127 @@ +# 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 contracts for TPARTARGMAX/MIN and THISTOGRAM.""" + +import pytest +from pypto import ir +from pypto.ir.op import tile_ops as tile +from pypto.pypto_core import DataType + + +def _view(valid_shape, layout=ir.TileLayout.row_major): + return ir.TileView( + valid_shape=valid_shape, + blayout=layout, + slayout=ir.TileLayout.none_box, + ) + + +def _tile(name, dtype, valid_shape, layout=ir.TileLayout.row_major, shape=(8, 16)): + return ir.Var( + name, + ir.TileType(list(shape), dtype, tile_view=_view(list(valid_shape), layout)), + ir.Span.unknown(), + ) + + +def _const_values(shape): + values = [] + for dim in shape: + assert isinstance(dim, ir.ConstInt) + values.append(dim.value) + return values + + +@pytest.mark.parametrize("op", [tile.part_argmax, tile.part_argmin]) +@pytest.mark.parametrize("value_dtype", [DataType.FP16, DataType.FP32]) +@pytest.mark.parametrize("index_dtype", [DataType.INT32, DataType.UINT32]) +def test_partial_arg_contract_returns_value_and_index_tiles(op, value_dtype, index_dtype): + src0 = _tile("src0", value_dtype, (8, 16)) + src1 = _tile("src1", value_dtype, (7, 13)) + idx0 = _tile("idx0", index_dtype, (8, 16)) + idx1 = _tile("idx1", index_dtype, (7, 13)) + + call = op(src0, src1, idx0, idx1) + + assert isinstance(call.type, ir.TupleType) + value_type, index_type = call.type.types + assert isinstance(value_type, ir.TileType) + assert isinstance(index_type, ir.TileType) + assert value_type.dtype == value_dtype + assert index_type.dtype == index_dtype + assert _const_values(value_type.get_effective_tile_view().valid_shape) == [8, 16] + assert _const_values(index_type.get_effective_tile_view().valid_shape) == [8, 16] + + +def test_partial_arg_result_layout_follows_dominating_source_pair(): + src0 = _tile("src0", DataType.FP32, (7, 13), ir.TileLayout.col_major) + src1 = _tile("src1", DataType.FP32, (8, 16)) + idx0 = _tile("idx0", DataType.INT32, (7, 13), ir.TileLayout.col_major) + idx1 = _tile("idx1", DataType.INT32, (8, 16)) + + result_type = tile.part_argmax(src0, src1, idx0, idx1).type + assert isinstance(result_type, ir.TupleType) + value_type, index_type = result_type.types + assert isinstance(value_type, ir.TileType) + assert isinstance(index_type, ir.TileType) + + assert value_type.get_effective_tile_view().blayout == ir.TileLayout.row_major + assert index_type.get_effective_tile_view().blayout == ir.TileLayout.row_major + + +@pytest.mark.parametrize("op", [tile.part_argmax, tile.part_argmin]) +def test_partial_arg_contract_rejects_mismatched_pairs_and_crossing_valid_shapes(op): + src0 = _tile("src0", DataType.FP32, (8, 12)) + src1 = _tile("src1", DataType.FP32, (7, 16)) + idx0 = _tile("idx0", DataType.INT32, (8, 12)) + idx1 = _tile("idx1", DataType.INT32, (7, 16)) + + with pytest.raises(ValueError, match="contain the other"): + op(src0, src1, idx0, idx1) + + full = _tile("full", DataType.FP32, (8, 16)) + bad_idx = _tile("bad_idx", DataType.INT32, (7, 16)) + with pytest.raises(ValueError, match="src0 and src0_idx"): + op(full, full, bad_idx, _tile("idx", DataType.INT32, (8, 16))) + + +@pytest.mark.parametrize("byte", [0, 1]) +def test_histogram_contract_uint16(byte): + src = _tile("src", DataType.UINT16, (7, 13)) + idx = _tile("idx", DataType.UINT8, (7, 1), ir.TileLayout.col_major, shape=(8, 1)) + + call = tile.histogram(src, idx, byte=byte) + + assert isinstance(call.type, ir.TileType) + assert call.type.dtype == DataType.UINT32 + assert _const_values(call.type.shape) == [8, 256] + assert _const_values(call.type.get_effective_tile_view().valid_shape) == [7, 256] + + +@pytest.mark.parametrize("byte,rows", [(0, 3), (1, 2), (2, 1), (3, 1)]) +def test_histogram_contract_uint32(byte, rows): + src = _tile("src", DataType.UINT32, (7, 13)) + idx = _tile("idx", DataType.UINT8, (rows, 13), shape=(rows, 16)) + + call = tile.histogram(src, idx, byte=byte) + + assert isinstance(call.type, ir.TileType) + assert call.type.dtype == DataType.UINT32 + assert _const_values(call.type.shape) == [8, 256] + + +def test_histogram_contract_rejects_invalid_byte_and_index_layout(): + src = _tile("src", DataType.UINT16, (7, 13)) + row_major_idx = _tile("idx", DataType.UINT8, (7, 1), shape=(8, 1)) + + with pytest.raises(ValueError, match=r"\[0, 3\]"): + tile.histogram(src, row_major_idx, byte=4) + with pytest.raises(ValueError, match="col_major"): + tile.histogram(src, row_major_idx, byte=1) diff --git a/tests/ut/ir/transforms/test_memory_reuse.py b/tests/ut/ir/transforms/test_memory_reuse.py index 1d56c9d1be..b010af9ccd 100644 --- a/tests/ut/ir/transforms/test_memory_reuse.py +++ b/tests/ut/ir/transforms/test_memory_reuse.py @@ -3851,6 +3851,38 @@ def main( f"row_sum output must not alias its tmp buffer, but both bind to {bases['s']}" ) + def test_tuple_outputs_do_not_alias_inputs_of_inplace_unsafe_op(self): + """Each part-arg tuple output must stay off all four input buffers.""" + + @pl.program + class Before: + @pl.function + def main( + self, + value_a: pl.Tensor[[8, 8], pl.FP32], + value_b: pl.Tensor[[8, 8], pl.FP32], + index_a: pl.Tensor[[8, 8], pl.INT32], + index_b: pl.Tensor[[8, 8], pl.INT32], + value_out: pl.Out[pl.Tensor[[8, 8], pl.FP32]], + index_out: pl.Out[pl.Tensor[[8, 8], pl.INT32]], + ) -> tuple[pl.Tensor[[8, 8], pl.FP32], pl.Tensor[[8, 8], pl.INT32]]: + va = pl.load(value_a, [0, 0], [8, 8]) + vb = pl.load(value_b, [0, 0], [8, 8]) + ia = pl.load(index_a, [0, 0], [8, 8]) + ib = pl.load(index_b, [0, 0], [8, 8]) + value, index = pl.tile.part_argmax(va, vb, ia, ib) + value_result = pl.store(value, [0, 0], value_out) + index_result = pl.store(index, [0, 0], index_out) + return value_result, index_result + + After = _run_pipeline(Before) + bases = _collect_tile_memref_bases(After) + inputs = {bases[name] for name in ("va", "vb", "ia", "ib")} + + assert bases["value"] not in inputs + assert bases["index"] not in inputs + assert bases["value"] != bases["index"] + def test_forbidden_input_reached_through_view_is_honored(self): """A not_inplace_safe op reading a VIEW of its input must still not alias it.