From 91a33f2bcfb36e2b778c2e0589aee30f3090122c Mon Sep 17 00:00:00 2001 From: jimmychou <47636600+jimmychou0@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:00:43 +0800 Subject: [PATCH 01/10] feat(pto): add T.vmi intrinsics and PTO VMI codegen Add the VMI (virtual machine instruction) intrinsics and PTO codegen: - tilelang/ascend/language/vmi.py: the T.vmi.* surface, exposed as T.vmi through tilelang.ascend.language - vector loads/stores, masks, arithmetic/bitwise/shift ops, vcvt (incl. packed FP4 sources), vinterpret_cast, vzip/vunzip, precision-aware vdiv, gather/scatter, histogram ops, and register allocation helpers - 54 tl.vmi.* TIR builtins registered in src/op/builtin.{h,cc} - PTO codegen: pto.vmi.* emission with typed-literal wrapping, FP4 pair expansion, eager vreg allocation for vector local.var, the signed/unsigned integer type split (pto.si*/pto.ui*), and a clear AscendC rejection diagnostic; generated kernels import tilelang.contrib.ptodsl before ptodsl so this build's native libraries load ahead of the PTODSL MLIR modules - KernelParam.storage_shape: packed-storage shapes for sub-byte dtypes Co-authored-by: KurrinQu Co-authored-by: Jay Zhuang <80731350+learning-chip@users.noreply.github.com> Co-authored-by: Zhang Zhendong Co-authored-by: cj --- src/ascend/codegen/codegen_ascend.cc | 20 + src/ascend/codegen/codegen_pto.cc | 472 +++++++++- src/ascend/codegen/codegen_pto.h | 6 + src/op/builtin.cc | 77 ++ src/op/builtin.h | 395 ++++++++ tilelang/ascend/language/__init__.py | 2 + tilelang/ascend/language/vmi.py | 1287 ++++++++++++++++++++++++++ 7 files changed, 2235 insertions(+), 24 deletions(-) create mode 100644 tilelang/ascend/language/vmi.py diff --git a/src/ascend/codegen/codegen_ascend.cc b/src/ascend/codegen/codegen_ascend.cc index e7951c7ab7..44b772d606 100644 --- a/src/ascend/codegen/codegen_ascend.cc +++ b/src/ascend/codegen/codegen_ascend.cc @@ -45,6 +45,19 @@ bool IsSimdOp(const CallNode *op, std::string *op_name = nullptr) { return false; } +bool IsVmiOp(const CallNode *op, std::string *op_name = nullptr) { + if (auto opt_call_op = op->op.as()) { + const std::string &name = opt_call_op.value()->name; + if (name.rfind("tl.vmi.", 0) == 0) { + if (op_name != nullptr) { + *op_name = name; + } + return true; + } + } + return false; +} + bool IsSimdMergingCall(const CallNode *op, std::string *intrinsic_name = nullptr) { if (!op->dtype.is_void() || op->args.empty()) { @@ -2580,6 +2593,13 @@ void CodeGenTileLangAscend::VisitExpr_(const CallNode *op, std::ostream &os) { << "Unsupported RNG distribution on Ascend: " << dist; os << "tl::philox_rand_" << dist << "(&" << ascend_rng_state_var_ << ")"; } else { + std::string op_name; + // Evaluate statements print their value through this CallNode visitor, so + // statement-form VMI calls such as vstore are rejected here as well. + if (IsVmiOp(op, &op_name)) { + LOG(FATAL) << "Ascend CCE codegen does not support " << op_name + << "; use target='pto' for T.vmi.*"; + } CodeGenC::VisitExpr_(op, os); } } diff --git a/src/ascend/codegen/codegen_pto.cc b/src/ascend/codegen/codegen_pto.cc index 69ab33eaf8..7cc457952d 100644 --- a/src/ascend/codegen/codegen_pto.cc +++ b/src/ascend/codegen/codegen_pto.cc @@ -26,6 +26,7 @@ #include #include #include +#include #include #include @@ -137,6 +138,12 @@ std::string DataTypeName(DataType t) { if (tl::IsAscendVectorizableFP8(t)) return FP8TypeName(t); + // PTO stores E2M1 as one byte containing a pair of logical FP4 values. + // TileLang expresses FP4 logically; PTO stores each pair as one f4e2m1x2 + // element. + if (t.is_float4_e2m1fn()) { + return "pto.f4e2m1x2"; + } ICHECK(t.is_scalar()) << "PTO scalar type expected, got " << t; if (t.is_bool()) { return "pto.i1"; @@ -151,16 +158,29 @@ std::string DataTypeName(DataType t) { return "pto.f4e2m1x2"; } else if (t.is_bfloat16()) { return "pto.bf16"; - } else if (t.is_int() || t.is_uint()) { - const char *prefix = t.is_uint() ? "pto.ui" : "pto.i"; + } else if (t.is_uint()) { + if (t.bits() == 64) + return "pto.ui64"; + if (t.bits() == 32) + return "pto.ui32"; + if (t.bits() == 16) + return "pto.ui16"; + if (t.bits() == 8) + return "pto.ui8"; + if (t.bits() == 1) + return "pto.i1"; + } else if (t.is_int()) { + // TileLang/TIR distinguishes signed and unsigned integer dtypes; + // preserve that distinction in the generated PTO type so VMI integer + // sources (loads, vci bridges, casts) share one element-type space. if (t.bits() == 64) - return std::string(prefix) + "64"; + return "pto.si64"; if (t.bits() == 32) - return std::string(prefix) + "32"; + return "pto.si32"; if (t.bits() == 16) - return std::string(prefix) + "16"; + return "pto.si16"; if (t.bits() == 8) - return std::string(prefix) + "8"; + return "pto.si8"; if (t.bits() == 1) return "pto.i1"; } @@ -230,6 +250,45 @@ bool IsImmediateScalar(const PrimExpr &value) { peeled.as() != nullptr; } +std::string PtoSignedIntegerTypeName(DataType t) { + ICHECK(t.is_scalar() && t.is_int()) + << "PTO signed integer scalar type expected, got " << t; + switch (t.bits()) { + case 8: + return "pto.si8"; + case 16: + return "pto.si16"; + case 32: + return "pto.si32"; + case 64: + return "pto.si64"; + default: + LOG(FATAL) << "Unsupported PTO signed integer type: " << t; + } + return ""; +} + +// PTODSL infers the result element type of vci/vbrc from the scalar +// argument. A dynamic TIR integer prints as MLIR `index` (or a signless +// constant), which disagrees with the signed/unsigned vector element type the +// surrounding VMI ops use, so the value argument is adapted to a same-width +// signless scalar and the result is reinterpreted bit-wise (no numeric +// conversion; bit width, signedness target and lanes are preserved). +std::string PtoSignlessIntegerTypeName(DataType t) { + ICHECK(t.is_scalar() && (t.is_int() || t.is_uint())) + << "PTO signless integer scalar type expected, got " << t; + return "pto.i" + std::to_string(t.bits()); +} + +std::string PtoIntegerVectorInterpretTarget(DataType element) { + if (element.is_int()) { + return PtoSignedIntegerTypeName(element); + } + ICHECK(element.is_uint()) + << "PTO integer element type expected, got " << element; + return "pto.ui" + std::to_string(element.bits()); +} + std::string StripPipePrefix(const std::string &name) { if (name.rfind("PIPE_", 0) == 0) { return name.substr(5); @@ -554,8 +613,9 @@ void CheckLocalVarBuffer(const BufferNode *buffer) { // They are initialized to None and receive a pointer on their first store. return; } - if (IsVectorLocalVarDtype(dtype) || dtype.is_handle()) { - (void)DataTypeName(dtype.element_of()); + if (dtype.lanes() > 1) { + // VMI registers are logical vectors; integer, floating-point and + // low-precision element types are all valid here. return; } ICHECK_EQ(dtype.lanes(), 1) @@ -2116,8 +2176,13 @@ void CodeGenTileLangPTO::AddFunction(const GlobalVar &gvar, std::string CodeGenTileLangPTO::Finish() { std::ostringstream code; - code << "from ptodsl import pto\n"; + // Import the tilelang compatibility layer first: it loads this build's + // native tilelang/tvm libraries before the PTODSL MLIR modules come up. + // Loading them in the opposite order trips a static-initializer crash in + // the MLIR PassBuilder when both library families are dlopened into the + // same process (observed on the Ascend validation runners). code << "import tilelang.contrib.ptodsl as tl\n"; + code << "from ptodsl import pto\n"; code << "\n"; code << decl_stream.str(); code << stream.str(); @@ -2707,6 +2772,15 @@ void CodeGenTileLangPTO::EmitLocalVarInitialization_(const Buffer &buffer, vid = name_supply_->FreshName("tl" + vid, false); var_idmap_[buffer->data.get()] = vid; } + if (dtype.lanes() > 1) { + // Mutable vector registers start as a vreg descriptor placeholder; the + // first actual store overwrites it with a real vector value. + stream << vid << " = pto.vmi.vreg(" << dtype.lanes() << ", " + << DataTypeName(dtype.element_of()) << ")\n"; + RegisterHandleType_(buffer->data.get(), dtype); + EmitMixedEntrySnapshot(buffer->data.get()); + return; + } if (current_function_has_gemm_) { // Keep authored PTO surface values (not Python ints) so AST-rewritten // `if` / br.assign can merge tile indices across branches. @@ -2799,13 +2873,11 @@ void CodeGenTileLangPTO::EmitBufferAllocation(const Buffer &buffer) { stream << vid << " = [None] * " << opt_size.value() << "\n"; } } else if (scope == "local.var") { - PrintIndent(); - // A surrounding SIMD loop may discover this local.var as loop-carried - // state before the allocation statement is emitted. LocalVarID then - // reserves its Python name in var_idmap_. Reuse that reservation here; - // AllocVarID would incorrectly diagnose it as a second TIR definition. - // local_var_buffers_ records allocations that were actually emitted, so - // a genuinely repeated allocation does not reset an accumulator. + CheckLocalVarBuffer(buffer.get()); + // Reuse the Python name if this local.var already has one registered in + // var_idmap_; AllocVarID would diagnose it as a second TIR definition. + // local_var_buffers_ separately tracks emitted initializations, so a + // repeated allocation does not reset the local.var value. if (local_var_buffers_.count(buffer->data.get())) { return; } @@ -2813,6 +2885,7 @@ void CodeGenTileLangPTO::EmitBufferAllocation(const Buffer &buffer) { std::string vid = var_idmap_.count(buffer->data.get()) ? GetVarID(buffer->data.get()) : AllocVarID(buffer->data.get()); + PrintIndent(); EmitLocalVarInitialization_(buffer, std::move(vid)); return; } @@ -3493,6 +3566,28 @@ void CodeGenTileLangPTO::RestoreMixedSectionVariables( } } +bool CodeGenTileLangPTO::IsVmiLocalRegisterBuffer( + const BufferNode *buffer) const { + return !inside_simtvf_body_ && ScopeOfBuffer(buffer) == "local" && + buffer->dtype.lanes() > 1; +} + +void CodeGenTileLangPTO::CheckVmiLocalRegisterIndex( + const BufferNode *buffer, const PrimExpr &index) const { + bool is_constant = true; + tirx::PostOrderVisit(index, [&](const ObjectRef &node) { + if (node.as()) { + is_constant = false; + } + }); + ICHECK(is_constant) + << "PTO VMI local register buffer `" << buffer->name + << "` requires a compile-time constant index. Register arrays cannot " + "be accessed inside T.unroll(..., explicit=False) or other runtime " + "loops; use T.unroll(..., explicit=True), got " + << index; +} + const CodeGenTileLangPTO::PTOGemmEmitContext & CodeGenTileLangPTO::EnsureGemmHelper(const CallNode *op) { ICHECK_EQ(op->args.size(), 12U) @@ -4687,6 +4782,296 @@ std::string CodeGenTileLangPTO::EmitAllReduceExpr_(const std::string &func_name, return os.str(); } +std::string +CodeGenTileLangPTO::PrintVmiAnnotationValue(const std::string &key, + const ObjectRef &value) { + if (key == "to_dtype") { + if (const auto *dtype_name = value.as()) { + const std::string &name = dtype_name->value; + // PTODSL int-to-int widening requires a signed/unsigned source + // (si16/ui16), not TVM's signless i16. + if (name == "si8") + return "pto.si8"; + if (name == "si16") + return "pto.si16"; + if (name == "si32") + return "pto.si32"; + if (name == "si64") + return "pto.si64"; + return ScalarType(ParseDataType(name)); + } + } + + if (const auto *expr = value.as()) { + return PrintExpr_(GetRef(expr)); + } + + std::ostringstream os; + os << value; + return os.str(); +} + +// Returns true if a tl.vmi.vcvt narrows logical BF16 values to FP4. PTO only +// exposes FP4 as f4e2m1x2, so this call packs its BF16 input in codegen. +static bool +IsVmiVcvtToFp4(const std::vector> &kwargs) { + for (const auto &[key, value] : kwargs) { + if (key == "to_dtype") { + if (const auto *str = value.as()) { + return str->value == "float4_e2m1fn"; + } + } + } + return false; +} + +static std::optional +VmiVcvtTarget(const std::vector> &kwargs) { + for (const auto &[key, value] : kwargs) { + if (key == "to_dtype") { + if (const auto *str = value.as()) { + return str->value; + } + } + } + return std::nullopt; +} + +// TileLang models packed FP4 in logical scalar lanes, while PTO VMI models +// one physical byte as f4e2m1x2. Reverse conversion therefore first produces +// bf16x2 pair carriers and then reinterprets those carriers as logical BF16 +// lanes. FP32 is the same operation followed by the ordinary BF16 widening +// conversion. PTO's unified VMI lowering owns the Packed4 P0-P3 expansion. +static bool IsVmiVcvtFromFp4(const CallNode *op) { + return !op->args.empty() && op->args[0].dtype().is_vector() && + op->args[0].dtype().element_of().is_float4_e2m1fn(); +} + +void CodeGenTileLangPTO::PrintPtoVmiCall_(const CallNode *op, + std::ostream &os) { + auto opt_call_op = op->op.as(); + ICHECK(opt_call_op.has_value()); + std::string op_name = opt_call_op.value()->name; + ICHECK(StartsWith(op_name, "tl.vmi.")) + << "Expected a tl.vmi.* call, got " << op_name; + + std::vector> kwargs; + kwargs.reserve(op->annotations.size()); + for (const auto &[key, value] : op->annotations) { + std::string key_str = key; + if (key_str == "loc" || key_str == "ip") { + continue; + } + kwargs.emplace_back(std::move(key_str), value); + } + std::sort(kwargs.begin(), kwargs.end(), [](const auto &lhs, const auto &rhs) { + return lhs.first < rhs.first; + }); + + if (op_name == "tl.vmi.vcvt" && IsVmiVcvtToFp4(kwargs)) { + ICHECK_EQ(op->args.size(), 1U) + << "Packed FP4 vcvt expects exactly one BF16 source vector"; + const DataType source_dtype = op->args[0].dtype(); + ICHECK(source_dtype.is_vector() && source_dtype.element_of().is_bfloat16()) + << "Packed FP4 vcvt only supports bfloat16 source vectors, got " + << source_dtype; + ICHECK_EQ(source_dtype.lanes() % 2, 0) + << "Packed FP4 vcvt requires an even BF16 lane count, got " + << source_dtype.lanes(); + } + + if (op_name == "tl.vmi.vcvt" && IsVmiVcvtFromFp4(op)) { + const auto target = VmiVcvtTarget(kwargs); + ICHECK(target.has_value() && + (*target == "bfloat16" || *target == "float32")) + << "Packed FP4 source vcvt only supports bfloat16 or float32, got " + << (target.has_value() ? *target : ""); + const DataType source_dtype = op->args[0].dtype(); + ICHECK_EQ(source_dtype.lanes() % 2, 0) + << "Packed FP4 source vcvt requires an even logical lane count, got " + << source_dtype.lanes(); + + ICHECK_EQ(op->args.size(), 1U) << "Packed FP4 source vcvt expects exactly " + "one packed FP4 source vector"; + const std::string source = PrintExpr_(op->args[0]); + // Pair-to-scalar BF16 re-view is a register-level vinterpret_cast. + // Do not materialize the converted pairs through a UB scratch buffer. + const std::string bf16 = "pto.vmi.vinterpret_cast(pto.vmi.vcvt(" + source + + ", to_dtype=pto.vmi.bf16x2), to_dtype=pto.bf16)"; + if (*target == "bfloat16") { + os << bf16; + } else { + os << "pto.vmi.vcvt(" << bf16 << ", to_dtype=pto.f32)"; + } + return; + } + + // Adapt integer scalar sources of vci/vbrc at the VMI + // value-argument boundary. The generated scalar must be *signless*: PTOAS + // materialises per-physical-vector offsets with `arith.constant`, which + // requires signless integers, so a typed (signed) constant base fails for + // every input wider than one physical vector. A VCI base can be folded to + // a constant whenever the loop has a single chunk, so integer VCI always + // takes this path; VBRC constants keep the verified typed-literal path. + const bool scalar_integer_input = + !op->args.empty() && op->args[0].dtype().is_scalar() && + (op->args[0].dtype().is_int() || op->args[0].dtype().is_uint()); + const bool integer_vector_result = + op->dtype.is_vector() && + (op->dtype.element_of().is_int() || op->dtype.element_of().is_uint()); + const bool integer_source_literal = + !op->args.empty() && op->args[0].as() != nullptr; + const bool bridge_integer_source = + (op_name == "tl.vmi.vci" || + (op_name == "tl.vmi.vbrc" && !integer_source_literal)) && + scalar_integer_input && integer_vector_result; + + if (bridge_integer_source) { + os << "pto.vmi.vinterpret_cast("; + } + os << "pto." << op_name.substr(3) << "("; + bool needs_comma = false; + + auto print_scalar_literal_value = [&](const PrimExpr &arg) { + if (const auto *imm = arg.as()) { + if (imm->dtype == DataType::Bool()) { + os << (imm->value ? "True" : "False"); + } else { + os << imm->value; + } + return; + } + if (const auto *imm = arg.as()) { + os << "float.fromhex('" << FlexibleHexFormat(imm->value) << "')"; + return; + } + PrintExpr_(arg, os); + }; + + // PTODSL vgather/vgatherb/vscatter take a single pointer operand. TileLang + // lowers buffer addresses to (ptr, elem_offset); fold them with addptr here. + auto print_ptr_with_offset = [&](const PrimExpr &ptr, + const PrimExpr &offset) { + if (is_zero(offset)) { + os << PrintExpr_(ptr); + return; + } + os << "pto.addptr(" << PrintExpr_(ptr) << ", " + << RemoveOutermostParentheses(PrintExpr_(offset)) << ")"; + }; + + if (op_name == "tl.vmi.vgather" || op_name == "tl.vmi.vgatherb") { + ICHECK_EQ(op->args.size(), 4U) + << op_name << " expects (ptr, offset, offsets, mask)"; + print_ptr_with_offset(op->args[0], op->args[1]); + os << ", " << PrintExpr_(op->args[2]) << ", " << PrintExpr_(op->args[3]); + needs_comma = true; + } else if (op_name == "tl.vmi.vscatter") { + ICHECK_EQ(op->args.size(), 5U) + << op_name << " expects (value, ptr, offset, offsets, mask)"; + os << PrintExpr_(op->args[0]) << ", "; + print_ptr_with_offset(op->args[1], op->args[2]); + os << ", " << PrintExpr_(op->args[3]) << ", " << PrintExpr_(op->args[4]); + needs_comma = true; + } else if (op_name == "tl.vmi.vstore") { + auto dist_mode_it = op->annotations.find("dist_mode"); + ObjectRef dist_mode = dist_mode_it != op->annotations.end() + ? (*dist_mode_it).second + : ObjectRef(); + const bool is_dintlv_store = + dist_mode.defined() && dist_mode.as() != nullptr && + Downcast(dist_mode)->value == "dintlv"; + if (is_dintlv_store) { + ICHECK_GE(op->args.size(), 4U) + << op_name + << " with dist_mode=dintlv expects (even, odd, ptr, offset[, mask])"; + os << "(" << PrintExpr_(op->args[0]) << ", " << PrintExpr_(op->args[1]) + << ")"; + for (size_t i = 2; i < op->args.size(); ++i) { + os << ", " << PrintExpr_(op->args[i]); + } + needs_comma = true; + } else { + for (size_t i = 0; i < op->args.size(); ++i) { + const PrimExpr &arg = op->args[i]; + if (needs_comma) { + os << ", "; + } + os << PrintExpr_(arg); + needs_comma = true; + } + } + } else { + for (size_t i = 0; i < op->args.size(); ++i) { + const PrimExpr &arg = op->args[i]; + if (needs_comma) { + os << ", "; + } + const bool is_scalar_literal = + arg.as() != nullptr || arg.as() != nullptr; + const bool should_wrap_typed_literal = + (op_name == "tl.vmi.vbrc" || op_name == "tl.vmi.vci") && i == 0 && + is_scalar_literal; + if (bridge_integer_source && i == 0) { + if (integer_source_literal) { + // Constant base: same-width signless scalar holding the original + // bit pattern (never a Python default type, never forced to i32). + os << PtoSignlessIntegerTypeName(arg.dtype()) << "("; + print_scalar_literal_value(arg); + os << ")"; + } else { + os << "pto.cast(" << PrintExpr_(arg) << ", " + << PtoSignlessIntegerTypeName(arg.dtype()) << ")"; + } + } else if (should_wrap_typed_literal) { + // PTODSL needs typed literal scalars for these VMI sources; preserve + // the literal dtype instead of assuming every source is f32. + os << ScalarType(arg.dtype()) << "("; + print_scalar_literal_value(arg); + os << ")"; + } else if (op_name == "tl.vmi.vcvt" && i == 0 && IsVmiVcvtToFp4(kwargs) && + arg.dtype().is_vector() && + arg.dtype().element_of().is_bfloat16()) { + // PTOAS vcvt only accepts bf16x2 -> f4x2; pair the bf16 source + // (vinterpret_cast is a physical no-op) inline. Note: pto.vmi.bf16x2 + // is the BF16x2Type dtype, distinct from the MLIR vector pto.bf16x2. + os << "pto.vmi.vinterpret_cast(" << PrintExpr_(arg) + << ", to_dtype=pto.vmi.bf16x2)"; + } else { + os << PrintExpr_(arg); + } + needs_comma = true; + } + } + + for (const auto &[key, value] : kwargs) { + if (needs_comma) { + os << ", "; + } + os << key << "="; + if (op_name == "tl.vmi.vcvt" && key == "to_dtype" && op->dtype.is_int()) { + os << PtoSignedIntegerTypeName(op->dtype.element_of()); + } else if (op_name == "tl.vmi.vload" && key == "size" && + op->dtype.element_of().is_float4_e2m1fn()) { + // TileLang VMI exposes FP4 lane counts logically. PTO represents each + // f4e2m1x2 byte as one physical VMI element, so its vload size must be + // expressed in packed-pair units. The frontend has already converted + // the corresponding address offset to the same physical unit. + os << "(" << PrintVmiAnnotationValue(key, value) << " // 2)"; + } else { + os << PrintVmiAnnotationValue(key, value); + } + needs_comma = true; + } + // Close the VMI source first, then feed the reinterpret target: the + // wrapper must be `vinterpret_cast(vci(...), to_dtype=...)`. + os << ")"; + if (bridge_integer_source) { + os << ", to_dtype=" + << PtoIntegerVectorInterpretTarget(op->dtype.element_of()) << ")"; + } +} + void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, std::ostream &os) { // NOLINT(*) if (EmitSimdMergingCall_(op, os)) { @@ -6050,6 +6435,17 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, if (op->args.size() == 1U && TryEmitUnaryMath_(op_name, op->args[0], os)) { return; } + if (StartsWith(op_name, "tl.vmi.")) { + if (op_name == "tl.vmi.pair_get") { + ICHECK_EQ(op->args.size(), 2U) + << "tl.vmi.pair_get expects exactly 2 arguments"; + os << "(" << PrintExpr_(op->args[0]) << ")[" << PrintExpr_(op->args[1]) + << "]"; + return; + } + PrintPtoVmiCall_(op, os); + return; + } if (StartsWith(op_name, "tl.")) { LOG(FATAL) << "PTO codegen does not support TileLang op `" << op_name << "`. Please add a handler in CodeGenTileLangPTO or lower " @@ -6769,6 +7165,14 @@ void CodeGenTileLangPTO::VisitStmt_(const BindNode *op) { } void CodeGenTileLangPTO::VisitStmt_(const AllocBufferNode *op) { + // All local.var buffers (scalar, handle and vector) share the unified + // allocation path with block alloc_buffers: scope registration, + // duplicate-allocation dedup, name reuse and mixed-section snapshots all + // live in EmitBufferAllocation / EmitLocalVarInitialization_. + if (GetPtrStorageScope(op->buffer->data) == "local.var") { + EmitBufferAllocation(op->buffer); + return; + } if (!current_function_is_cube_ || in_mixed_vector_section_) { EmitBufferAllocation(op->buffer); return; @@ -6794,14 +7198,6 @@ void CodeGenTileLangPTO::VisitStmt_(const AllocBufferNode *op) { auto opt_size = alloc_ref.ConstantAllocationSize(); ICHECK(opt_size.has_value()) << "PTO local.fragment allocation expects a constant size"; - } else if (scope == "local.var") { - // Scalar alloc_var is used by GEMM tile mapping; vector alloc_var is a - // mutable SIMD register and is emitted as a PTODSL Vec value. - CheckLocalVarBuffer(op->buffer.get()); - PrintIndent(); - local_var_buffers_.insert(buffer_var.get()); - EmitLocalVarInitialization_(op->buffer, AllocVarID(buffer_var.get())); - return; } RegisterHandleType_(buffer_var.get(), op->buffer->dtype); @@ -7401,6 +7797,20 @@ void CodeGenTileLangPTO::VisitExpr_(const BufferLoadNode *op, return; } + if (IsVmiLocalRegisterBuffer(op->buffer.get())) { + ICHECK_EQ(op->indices.size(), 1U) + << "PTO VMI local register buffers must be flattened before codegen"; + ICHECK(!op->predicate.defined()) + << "PTO VMI local register buffers do not support predicated loads"; + CheckVmiLocalRegisterIndex(op->buffer.get(), op->indices[0]); + // Reuse the tree's vector-local accessor so register arrays follow the + // slot-name model (Name stores stay loop-carriable for PTODSL's AST + // liveness analysis, matching scalarized local.fragment buffers). + os << GetVectorLocalRef(op->buffer->data.get(), op->indices[0], + "PTO VMI local register buffer load"); + return; + } + ICHECK_EQ(op->indices.size(), 1) << "CodeGenTileLangPTO only supports flat buffer loads"; ICHECK(!op->predicate.defined()) @@ -7597,6 +8007,20 @@ void CodeGenTileLangPTO::VisitStmt_(const BufferStoreNode *op) { return; } + if (IsVmiLocalRegisterBuffer(op->buffer.get())) { + ICHECK_EQ(op->indices.size(), 1U) + << "PTO VMI local register buffers must be flattened before codegen"; + ICHECK(!op->predicate.defined()) + << "PTO VMI local register buffers do not support predicated stores"; + CheckVmiLocalRegisterIndex(op->buffer.get(), op->indices[0]); + PrintIndent(); + stream << GetVectorLocalRef(op->buffer->data.get(), op->indices[0], + "PTO VMI local register buffer store") + << " = " << RemoveOutermostParentheses(PrintExpr_(op->value)) + << "\n"; + return; + } + ICHECK_EQ(op->indices.size(), 1) << "CodeGenTileLangPTO only supports flat buffer stores"; ICHECK(!op->predicate.defined()) diff --git a/src/ascend/codegen/codegen_pto.h b/src/ascend/codegen/codegen_pto.h index eb6c857bc8..f442c1d949 100644 --- a/src/ascend/codegen/codegen_pto.h +++ b/src/ascend/codegen/codegen_pto.h @@ -191,6 +191,12 @@ class CodeGenTileLangPTO final : public CodeGenTileLangPY { bool IsLocalVarBuffer(const VarNode *var) const; void EmitMixedEntrySnapshot(const VarNode *var); void RestoreMixedSectionVariables(const SBlockNode *section); + std::string PrintVmiAnnotationValue(const std::string &key, + const ObjectRef &value); + void PrintPtoVmiCall_(const CallNode *op, std::ostream &os); + bool IsVmiLocalRegisterBuffer(const BufferNode *buffer) const; + void CheckVmiLocalRegisterIndex(const BufferNode *buffer, + const PrimExpr &index) const; bool HasAscendGemmL1(const PrimFunc &func) const; bool HasAscendBlockscaledGemmL1(const PrimFunc &func) const; bool HasAscendMad(const PrimFunc &func) const; diff --git a/src/op/builtin.cc b/src/op/builtin.cc index 71a91d24cb..a7486cb670 100644 --- a/src/op/builtin.cc +++ b/src/op/builtin.cc @@ -49,6 +49,24 @@ TIR_DEFINE_TL_BUILTIN(tvm_ffi_call_with_result) .set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); +// VMI builtin registration macros (mirrors TIR_DEFINE_TL_BUILTIN from +// builtin_registry.h, but under the tl.vmi.* op prefix). +#define TIR_DEFINE_TL_VMI_BUILTIN(Name) \ + const Op &vmi_##Name() { \ + static const Op &op = Op::Get("tl.vmi." #Name); \ + return op; \ + } \ + TVM_REGISTER_OP("tl.vmi." #Name) \ + .set_attr("TScriptPrinterName", "tl.vmi." #Name) + +#define TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(Name) \ + TIR_DEFINE_TL_VMI_BUILTIN(Name).set_attr( \ + "TCallEffectKind", Integer(CallEffectKind::kOpaque)) + +#define TIR_DEFINE_TL_VMI_PURE_BUILTIN(Name) \ + TIR_DEFINE_TL_VMI_BUILTIN(Name).set_attr( \ + "TCallEffectKind", Integer(CallEffectKind::kPure)) + TIR_DEFINE_TL_BUILTIN(access_ptr) .set_num_inputs(3) .set_attr("TCallEffectKind", @@ -138,6 +156,65 @@ TIR_DEFINE_TL_BUILTIN(sync_grid).set_num_inputs(0).set_attr( TIR_DEFINE_TL_BUILTIN(sync_warp).set_num_inputs(-1).set_attr( "TCallEffectKind", Integer(CallEffectKind::kOpaque)); +// VMI builtins are registered separately from the Python wrappers so later +// passes/codegen can match on stable tl.vmi.* op identities. The positional +// operand ABI plus annotation lowering rules are documented next to the +// declarations in builtin.h and correspond to tilelang/ascend/language/vmi.py. +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vload).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vstore).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(create_mask).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vci).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vbrc).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vintlv).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vdintlv).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vunzip).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vzip).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_PURE_BUILTIN(pair_get).set_num_inputs(2); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vadd).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vsub).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vmul).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vdiv).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vmax).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vmin).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vand).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vor).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vxor).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vshl).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vshr).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vabs).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vneg).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vrelu).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vexp).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vln).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vsqrt).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vnot).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vadds).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vmuls).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vmaxs).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vmins).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vshls).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vshrs).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vcmp).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vcmps).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vsel).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vselr).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vcadd).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vcmax).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vcmin).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vcvt).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vinterpret_cast).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vexpdif).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vaxpy).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vlrelu).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vprelu).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vmull).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vmula).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vdhist).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vchist).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vgather).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vgatherb).set_num_inputs(-1); +TIR_DEFINE_TL_VMI_OPAQUE_BUILTIN(vscatter).set_num_inputs(-1); + TIR_DEFINE_TL_BUILTIN(any_sync).set_num_inputs(2).set_attr( "TCallEffectKind", Integer(CallEffectKind::kOpaque)); diff --git a/src/op/builtin.h b/src/op/builtin.h index 0918c4ed6b..310b5f59fe 100644 --- a/src/op/builtin.h +++ b/src/op/builtin.h @@ -324,6 +324,401 @@ TVM_DLL const Op &rng_rand_float(); // CUDA; #2855 files it under the CUDA dialect, so redeclare for asc codegen. TVM_DLL const Op &device_assert(); TVM_DLL const Op &device_assert_with_msg(); +// Ascend VMI virtual vector intrinsics (T.vmi.* API). +// TIR ops registered as tl.vmi.. These mirror PTODSL's public pto.vmi +// surface. +// +// ABI notes: +// - The TIR positional operands only carry dataflow values such as +// ptr/offset/values/mask. Python keyword-only API parameters like size, +// to_dtype, dist_mode, group, pmode, and order are lowered as call +// annotations instead of extra positional operands. +// - vload(ptr, offset) with annotations {size, to_dtype?, stride?, +// block_stride?, repeat_stride?, dist_mode?, group?} +// - vstore(values, destination_ptr, offset, mask?) with annotations +// {stride?, block_stride?, repeat_stride?, dist_mode?, group?, pmode?} +// - create_mask(active_lanes) with annotations {size, group?} +// - vci(base) with annotations {size, order?} +// - vbrc(value) with annotations {size, group?} +// - vintlv(lhs, rhs, mask) / vdintlv(lhs, rhs, mask) +// - pair_get(pair, index) is an internal pure helper for tuple-style unpacking +// of multi-result VMI calls. + +// -- Load / Store / Predicates +TVM_DLL const Op &vmi_vload(); +TVM_DLL const Op &vmi_vstore(); +TVM_DLL const Op &vmi_create_mask(); + +// -- Index / Broadcast / Rearrange +TVM_DLL const Op &vmi_vci(); +TVM_DLL const Op &vmi_vbrc(); +TVM_DLL const Op &vmi_vintlv(); +TVM_DLL const Op &vmi_vdintlv(); +TVM_DLL const Op &vmi_vunzip(); +TVM_DLL const Op &vmi_vzip(); +TVM_DLL const Op &vmi_pair_get(); + +// -- Binary arithmetic / bitwise / shifts +TVM_DLL const Op &vmi_vadd(); +TVM_DLL const Op &vmi_vsub(); +TVM_DLL const Op &vmi_vmul(); +TVM_DLL const Op &vmi_vdiv(); +TVM_DLL const Op &vmi_vmax(); +TVM_DLL const Op &vmi_vmin(); +TVM_DLL const Op &vmi_vand(); +TVM_DLL const Op &vmi_vor(); +TVM_DLL const Op &vmi_vxor(); +TVM_DLL const Op &vmi_vshl(); +TVM_DLL const Op &vmi_vshr(); + +// -- Unary / math +TVM_DLL const Op &vmi_vabs(); +TVM_DLL const Op &vmi_vneg(); +TVM_DLL const Op &vmi_vrelu(); +TVM_DLL const Op &vmi_vexp(); +TVM_DLL const Op &vmi_vln(); +TVM_DLL const Op &vmi_vsqrt(); +TVM_DLL const Op &vmi_vnot(); + +// -- Vector-scalar arithmetic / shifts +TVM_DLL const Op &vmi_vadds(); +TVM_DLL const Op &vmi_vmuls(); +TVM_DLL const Op &vmi_vmaxs(); +TVM_DLL const Op &vmi_vmins(); +TVM_DLL const Op &vmi_vshls(); +TVM_DLL const Op &vmi_vshrs(); + +// -- Compare / select / reduce / convert +TVM_DLL const Op &vmi_vcmp(); +TVM_DLL const Op &vmi_vcmps(); +TVM_DLL const Op &vmi_vsel(); +TVM_DLL const Op &vmi_vselr(); +TVM_DLL const Op &vmi_vcadd(); +TVM_DLL const Op &vmi_vcmax(); +TVM_DLL const Op &vmi_vcmin(); +TVM_DLL const Op &vmi_vcvt(); +TVM_DLL const Op &vmi_vinterpret_cast(); + +// -- SFU / irregular / histogram +TVM_DLL const Op &vmi_vexpdif(); +TVM_DLL const Op &vmi_vaxpy(); +TVM_DLL const Op &vmi_vlrelu(); +TVM_DLL const Op &vmi_vprelu(); +TVM_DLL const Op &vmi_vmull(); +TVM_DLL const Op &vmi_vmula(); +TVM_DLL const Op &vmi_vdhist(); +TVM_DLL const Op &vmi_vchist(); +TVM_DLL const Op &vmi_vgather(); +TVM_DLL const Op &vmi_vgatherb(); +TVM_DLL const Op &vmi_vscatter(); + +/*! + * \brief Ascend SetFlag intrinsic for pipeline synchronization. + * + * ascend_set_flag(hard_event_string, event_id) + * + */ +TVM_DLL const Op &ascend_set_flag(); + +/*! + * \brief Ascend WaitFlag intrinsic for pipeline synchronization. + * + * ascend_wait_flag(hard_event_string, event_id) + * + */ +TVM_DLL const Op &ascend_wait_flag(); + +/*! + * \brief Ascend inter-core synchronization arrive intrinsic. + * + * ascend_sync_inter_arrive(pipe_string, flag_id) + * + */ +TVM_DLL const Op &ascend_sync_inter_arrive(); + +/*! + * \brief Ascend inter-core synchronization wait intrinsic. + * + * ascend_sync_inter_wait(pipe_string, flag_id) + * + */ +TVM_DLL const Op &ascend_sync_inter_wait(); + +/*! + * \brief Ascend thread fence intrinsic. + * + * ascend_threadfence() + * + */ +TVM_DLL const Op &ascend_threadfence(); + +/*! + * \brief Ascend CrossCoreSetFlag intrinsic for cross-core synchronization. + * + * ascend_cross_core_set_flag(mode_id_string, pipe_string, flag_id) + * + */ +TVM_DLL const Op &ascend_cross_core_set_flag(); + +/*! + * \brief Ascend CrossCoreWaitFlag intrinsic for cross-core synchronization. + * + * ascend_cross_core_wait_flag(flag_id) + * + */ +TVM_DLL const Op &ascend_cross_core_wait_flag(); + +/*! + * \brief Ascend DMA copy from GM to UBuf. + * + * ascend_copy_gm_to_ubuf(dst, src, sid, nBurst, burstLen, + * leftPadding, rightPadding, dataSelect, l2CacheCtl, + * burstSrcStride, burstDstStride) + * + */ +TVM_DLL const Op &ascend_copy_gm_to_ubuf(); + +/*! + * \brief Ascend set padding fill value for subsequent padded MTE copies. + * + * ascend_set_copy_pad_value(value) + * + * The single argument is a typed scalar PrimExpr; its dtype selects the + * AscendC::SetPadValue template instantiation. Configures stateful hardware + * pad value consumed by a following padded GM -> UB/L1 copy. + * + */ +TVM_DLL const Op &ascend_set_copy_pad_value(); + +/*! + * \brief Ascend DMA copy from UBuf to GM. + * + * ascend_copy_ubuf_to_gm(dst, src, sid, burst_num, burst_len, l2_cache_ctl, + * burst_dst_stride, burst_src_stride) + * + */ +TVM_DLL const Op &ascend_copy_ubuf_to_gm(); + +/*! + * \brief Ascend DMA copy from GM to CBuf (L1). + * + * ascend_copy_gm_to_cbuf(dst, src, sid, loop1_src_stride, l2_cache_ctrl, + * n_value, d_value, loop4_src_stride, smallc0_en, transpose, dst_n_value, + * physical_dtype) + * + * When transpose=0 emits copy_gm_to_cbuf_multi_nd2nz (row-major → NZ). + * When transpose!=0 emits copy_gm_to_cbuf_multi_dn2nz (col-major → NZ), + * which transposes the N/D mapping — useful for NN matmul on Ascend. + */ +TVM_DLL const Op &ascend_copy_gm_to_cbuf(); + +/*! + * \brief Fill an L1 (CBuf) NZ matrix region with a scalar value. + * + * ascend_fill_l1(dst, byte_offset, raw_value, repeat_times, block_num_32b, + * dst_gap_32b, fill_word_bits) + * + * Emits create_cbuf_matrix at dst + byte_offset. The three geometry fields map + * directly to the instruction config bit field. fill_word_bits selects a raw + * uint16_t or uint32_t destination view, and raw_value carries the repeated + * element bit pattern. + */ +TVM_DLL const Op &ascend_fill_l1(); + +/*! + * \brief Ascend load from CBuf (L1) to CA (L0A). + * + * ascend_load_cbuf_to_ca(dst, src, mStartPosition, kStartPosition, + * mStep, kStep, srcStride, dstStride, transpose) + * + * Optionally 16 args when an MX scale-factor companion load is attached; the + * extra args drive a following load_cbuf_to_ca_mx: + * [9] sf_ptr, [10] sf_x_start, + * [11] sf_y_start (y is contiguous fractal direction), + * [12] sf_x_step, + * [13] sf_y_step, + * [14] sf_src_stride, [15] sf_dst_stride. + */ +TVM_DLL const Op &ascend_load_cbuf_to_ca(); + +/*! + * \brief Ascend load from CBuf (L1) to CB (L0B). + * + * ascend_load_cbuf_to_cb(dst, src, mStartPosition, kStartPosition, + * mStep, kStep, srcStride, dstStride, transpose) + * + * Optionally 16 args when an MX scale-factor companion load is attached; the + * extra args drive a following load_cbuf_to_ca_mx: + * [9] sf_ptr, [10] sf_x_start, + * [11] sf_y_start (y is contiguous fractal direction), + * [12] sf_x_step, + * [13] sf_y_step, + * [14] sf_src_stride, [15] sf_dst_stride. + */ +TVM_DLL const Op &ascend_load_cbuf_to_cb(); + +/*! + * \brief Ascend copy matrix from CC (L0C) to UBuf. + * + * ascend_copy_matrix_cc_to_ub(dst, src, sid, n_size, m_size, + * loop_dst_stride, loop_src_stride, dual_dst_ctl, sub_blockid, clip_relu_pre, + * unit_flag_ctl, quant_pre, relu_pre, split_en, NZ2ND_en, quant_post, + * relu_post, clip_relu_post, loop_enhance_en, eltwise_op, eltwise_antq_en, + * loop_enhance_merge_en, C0_pad_en, wino_post_en, broadcast_en, NZ2DN_en) + * + */ +TVM_DLL const Op &ascend_copy_matrix_cc_to_ub(); + +/*! + * \brief Ascend DMA copy from UBuf to CBuf (L1). + * + * ascend_copy_ubuf_to_cbuf(dst, src, sub_blockid, burst_num, burst_len, + * src_gap, dst_gap) + * + */ +TVM_DLL const Op &ascend_copy_ubuf_to_cbuf(); + +/*! + * \brief Ascend ND→NZ scatter (SimdVF) from UB ND tile to UB NZ tile. + * + * ascend_nd2nz_scatter(src_ub, tmp_nz_ub, rows, cols, + * dst_dtype_str, src_dtype_str) + * + */ +TVM_DLL const Op &ascend_nd2nz_scatter(); + +/*! + * \brief Ascend post-scatter UB→L1 raw DMA with NZ-fractal stride correction. + * + * ascend_nd2nz_post_copy(dst_l1, src_ub, rows, cols, full_rows, dst_dtype_str) + * + */ +TVM_DLL const Op &ascend_nd2nz_post_copy(); + +TVM_DLL const Op &ascend_copy_matrix_cc_to_gm(); + +/*! + * \brief Ascend Cube MAD (matrix multiply-add) instruction. + * + * ascend_mad(dst, src_a, src_b, M, K, N, unit_flag_ctrl, gemv_ctrl, BTbuf_ctrl, + * zero_Cmatrix) + * + */ +TVM_DLL const Op &ascend_mad(); + +/*! + * \brief Ascend Cube block-scaled MAD (MXFP8) instruction. + * + * ascend_mad_mx(dst, src_a, src_b, M, K, N, unit_flag_ctrl, gemv_ctrl, + * BTbuf_ctrl, zero_Cmatrix) + * + * Same arguments as ascend_mad. The per-block scale factors are read from the + * L0A/L0B MX scale registers that must be loaded beforehand (e.g. via + * T.copy(l1_data, l0, sf=l1_sf) which emits load_cbuf_to_ca_mx/cb_mx). + */ +TVM_DLL const Op &ascend_mad_mx(); + +/*! + * \brief Ascend GEMM with L1-scoped inputs (auto sub-K pipeline). + * + * ascend_gemm_l1(cc_ptr, cbuf_a_ptr, cbuf_b_ptr, M, K, N, tile_k_sub, + * trans_b, clear_accum, dtype_str) + * + */ +TVM_DLL const Op &ascend_gemm_l1(); + +/*! + * \brief Ascend block-scaled GEMM with L1-scoped inputs (MXFP8). + * + * ascend_blockscaled_gemm_l1(cc_ptr, a_ptr, b_ptr, sfa_ptr, sfb_ptr, + * M, K, N, tile_k_sub, trans_b, clear_accum, in_dtype_str, sf_dtype_str, + * accum_dtype_str, buf_offset, sf_k_offset, sf_nz_stride, unit_flag_ctrl) + * + * Uses blockscaled_gemm.h template with mad_mx instruction. + */ +TVM_DLL const Op &ascend_blockscaled_gemm_l1(); + +/*! + * \brief Ascend scalar GM read bypassing dcache. + * + * ascend_read_gm_bypass_dcache(address_of(BufferLoad)) + * + * Replaces a scalar BufferLoad from a global buffer that has writes + * elsewhere. The codegen emits tl::read_gm_bypass_dcache(ptr). + */ +TVM_DLL const Op &ascend_read_gm_bypass_dcache(); + +/*! + * \brief Ascend scalar GM write bypassing dcache. + * + * ascend_write_gm_bypass_dcache(address_of(BufferLoad), value) + * + * Replaces a scalar BufferStore to a global buffer that has writes + * elsewhere. The codegen emits tl::write_gm_bypass_dcache(ptr, val). + */ +TVM_DLL const Op &ascend_write_gm_bypass_dcache(); + +/*! + * \brief Ascend get_buf for pipe buffer acquisition. + * + * ascend_get_buf(pipe_string, buf_id, mode) + * + */ +TVM_DLL const Op &ascend_get_buf(); + +/*! + * \brief Ascend rls_buf for pipe buffer release. + * + * ascend_rls_buf(pipe_string, buf_id, mode) + * + */ +TVM_DLL const Op &ascend_rls_buf(); + +/*! + * \brief Ascend set HF32 mode for fp32 matmul. + * + * ascend_set_hf32_mode(mode_int) + * mode_int: 0=disable, 1=enable nearest_zero, 2=enable nearest_even + * + */ +TVM_DLL const Op &ascend_set_hf32_mode(); + +/*! + * \brief Ascend arm a hardware store-mode atomic op for GM stores. + * + * ascend_set_atomic(op_str, typed_zero) + * op_str: "add" | "max" | "min" + * typed_zero: a constant whose dtype selects the accumulate type T + * (float / half / int16 / int32 / int8 / bfloat16). + * + * Emits AscendC::SetAtomic{Add,Max,Min}(). Arm once; subsequent L0C->GM / + * UB->GM stores reduce into GM in hardware. Clear with + * ascend_set_atomic_none(). + */ +TVM_DLL const Op &ascend_set_atomic(); + +/*! + * \brief Ascend clear the store-mode atomic flag (restore plain store). + * + * ascend_set_atomic_none() + */ +TVM_DLL const Op &ascend_set_atomic_none(); + +/*! + * \brief Programmatic dependency trigger. + * + * pdl_trigger() + * + */ +TVM_DLL const Op &pdl_trigger(); + +/*! + * \brief Programmatic grid dependency synchronization. + * + * pdl_sync() + * + */ +TVM_DLL const Op &pdl_sync(); /*! * \brief Warp-vote: non-zero if ANY active lane in the mask has a non-zero diff --git a/tilelang/ascend/language/__init__.py b/tilelang/ascend/language/__init__.py index 423d829d51..1a431d576e 100644 --- a/tilelang/ascend/language/__init__.py +++ b/tilelang/ascend/language/__init__.py @@ -66,6 +66,7 @@ from .frame import Cube, CubeFrame, SimdVF, SimdVFFrame, SimtVF, SimtVFFrame, Vector, VectorFrame # noqa: F401 from . import simd as simd # noqa: F401 (exposed as T.simd.*) +from . import vmi as vmi # noqa: F401 (exposed as T.vmi.*) # Scalar reciprocal sqrt is a frontend composition of common math ops. from .math import rsqrt # noqa: F401 @@ -126,6 +127,7 @@ "rng_rand_float", "rsqrt", "simd", + "vmi", "unroll", ) diff --git a/tilelang/ascend/language/vmi.py b/tilelang/ascend/language/vmi.py new file mode 100644 index 0000000000..506482f793 --- /dev/null +++ b/tilelang/ascend/language/vmi.py @@ -0,0 +1,1287 @@ +"""TileLang VMI namespace for Ascend PTO kernels.""" + +from __future__ import annotations + +from collections.abc import Sequence +from functools import wraps + +from tvm import tirx +from tvm.tirx import Buffer, BufferLoad, BufferRegion +from tvm.tirx.script.builder.ir import bind as _bind +from tvm.script.ir_builder import IRBuilder + +from tilelang.language.builtin import access_ptr +from tilelang.language.dtypes import dtype as _dtype +from .frame import inside_simdvf as _inside_simdvf + +_Op = tirx.op.Op.get + + +def alloc_local(shape, vreg_type): + """Allocate a statically indexed array of VMI vector registers. + + ``vreg_type`` is a vector dtype produced by :func:`vreg`. The resulting + buffer may only be indexed by compile-time values, such as induction + variables from ``T.unroll(..., explicit=True)`` on the PTO backend. + """ + require_vmi_scope("T.vmi.alloc_local(...)") + dtype = _dtype(vreg_type) + if int(getattr(dtype, "lanes", 1)) <= 1: + raise TypeError(f"T.vmi.alloc_local(...) expects a VMI vector type produced by T.vmi.vreg(...), got {dtype}") + + dims = shape if _is_sequence(shape) else (shape,) + normalized_shape = [] + for dim in dims: + value = _constant_int(dim) + if value is None: + raise TypeError(f"T.vmi.alloc_local(...) requires every shape dimension to be a compile-time integer, got {dim!r}") + if value <= 0: + raise ValueError(f"T.vmi.alloc_local(...) requires positive shape dimensions, got {value}") + normalized_shape.append(value) + + if not normalized_shape: + raise ValueError("T.vmi.alloc_local(...) requires at least one shape dimension") + + from tilelang.language.allocate import alloc_local as _alloc_local + + return _alloc_local(tuple(normalized_shape), dtype, scope="local") + + +def alloc_var(dtype, *, size): + """Allocate a mutable VMI vector register variable. + + This mirrors ``T.simd.alloc_var`` but preserves the VMI lane count in the + TIR vector dtype, allowing loop-carried VMI accumulators to be rebound + without eager-builder immutable-value diagnostics. + """ + lanes = _normalize_lanes(size, context="T.vmi.alloc_var(...)") + elem = _scalar_dtype(dtype, context="T.vmi.alloc_var(...)") + from tilelang.language.allocate import alloc_var as _alloc_var + + return _alloc_var(f"{elem}x{lanes}", scope="local.var") + + +def _normalize_lanes(lanes: int, *, context: str) -> int: + if isinstance(lanes, bool) or not isinstance(lanes, int): + raise TypeError(f"{context} expects lanes to be a positive Python int, got {type(lanes)}") + if lanes <= 0: + raise ValueError(f"{context} expects lanes to be positive, got {lanes!r}") + return lanes + + +# Formal PTODSL VMI vreg/mask lane counts (see ptodsl VMI_LANE_COUNTS). +_VMI_LANE_COUNTS = (1, 2, 4, 8, 64, 128, 256) + + +def _require_vmi_lane_count(lanes: int, *, context: str) -> int: + if lanes not in _VMI_LANE_COUNTS: + raise ValueError(f"{context} requires lanes to be one of {_VMI_LANE_COUNTS}; got {lanes}") + return lanes + + +# PTODSL signed integer names. TileLang/TIR distinguishes signed and unsigned +# integer dtypes; codegen preserves that distinction in the generated PTO type. +_PTO_SIGNED_DTYPE = { + "si8": "int8", + "si16": "int16", + "si32": "int32", + "si64": "int64", +} + + +def _scalar_dtype(dtype_like, *, context: str): + if isinstance(dtype_like, str) and dtype_like in _PTO_SIGNED_DTYPE: + dtype_like = _PTO_SIGNED_DTYPE[dtype_like] + dt = _dtype(dtype_like) + if getattr(dt, "lanes", 1) != 1: + raise ValueError(f"{context} expects a scalar element dtype, got {dt}") + return dt + + +def _pto_to_dtype_annotation(dtype_like) -> str: + if isinstance(dtype_like, str) and dtype_like in _PTO_SIGNED_DTYPE: + return dtype_like + return str(_scalar_dtype(dtype_like, context="T.vmi to_dtype")) + + +def _is_sequence(value) -> bool: + return isinstance(value, Sequence) and not isinstance(value, (str, bytes)) + + +def _wrap_pair(pair): + if not IRBuilder.is_in_scope(): + raise RuntimeError("VMI pair results require an active IRBuilder") + return VmiPair(_bind(pair)) + + +def _dtype_of(value): + if isinstance(value, VmiPair): + return value.dtype + if hasattr(value, "dtype"): + return _dtype(value.dtype) + if hasattr(value, "buffer") and hasattr(value.buffer, "dtype"): + return _dtype(value.buffer.dtype) + return _dtype(tirx.convert(value).dtype) + + +def _element_dtype_of(value): + dt = _dtype_of(value) + return _dtype(str(dt).split("x", 1)[0]) + + +def _lanes_of(value) -> int: + return int(getattr(_dtype_of(value), "lanes", 1)) + + +def _require_mask(mask, *, context: str): + if mask is None: + raise TypeError(f"{context} requires a mask operand") + return mask + + +def _require_same_vreg_type(*values, context: str, allow_packed_fp4: bool = False): + dtypes = [_dtype_of(value) for value in values] + if any(dtype != dtypes[0] for dtype in dtypes[1:]): + raise TypeError(f"{context} requires identical VMI vector types") + if not allow_packed_fp4: + _reject_packed_fp4(*values, context=context) + return dtypes[0] + + +def _reject_packed_fp4(*values, context: str): + """Reject FP4 values from VMI operations without packed-FP4 semantics.""" + if any(str(_element_dtype_of(value)) == "float4_e2m1fn" for value in values): + raise TypeError(f"{context} does not support packed FP4 vectors") + + +def _require_integer_vector(value, *, context: str): + if not str(_element_dtype_of(value)).startswith(("int", "uint")): + raise TypeError(f"{context} requires an integer VMI vector") + return _dtype_of(value) + + +def _require_compatible_mask(mask, lanes: int, *, context: str): + mask = _require_mask(mask, context=context) + if not str(_dtype_of(mask)).startswith("bool"): + raise TypeError(f"{context} expects a VMI mask operand") + if _lanes_of(mask) != lanes: + raise TypeError(f"{context} requires mask and vector lane counts to match") + return mask + + +def _require_ub_address(value, *, context: str): + buffer = None + if isinstance(value, (BufferLoad, BufferRegion)): + buffer = value.buffer + elif isinstance(value, Buffer): + buffer = value + if buffer is not None and buffer.scope() not in {"shared", "shared.dyn", "ub"}: + raise TypeError(f"{context} requires a UB pointer, got buffer scope {buffer.scope()!r}") + type_annotation = getattr(value, "type_annotation", None) + storage_scope = getattr(type_annotation, "storage_scope", None) + if storage_scope and storage_scope not in {"shared", "shared.dyn", "ub"}: + raise TypeError(f"{context} requires a UB pointer, got pointer scope {storage_scope!r}") + return value + + +def _address_element_dtype(value): + if isinstance(value, BufferLoad): + return _dtype(value.buffer.dtype) + if isinstance(value, BufferRegion): + return _dtype(value.buffer.dtype) + if isinstance(value, Buffer): + return _dtype(value.dtype) + type_annotation = getattr(value, "type_annotation", None) + element_type = getattr(type_annotation, "element_type", None) + dtype = getattr(element_type, "dtype", None) + if dtype is not None and str(dtype): + return _dtype(dtype) + value_dtype = getattr(value, "dtype", None) + if value_dtype is not None and str(value_dtype) not in {"", "handle"}: + try: + return _element_dtype_of(value) + except (TypeError, ValueError): + return None + return None + + +def _require_address_element_dtype(value, *, context: str): + dtype = _address_element_dtype(value) + if dtype is None: + raise TypeError(f"{context} requires an address with a known element dtype") + return dtype + + +def _normalize_compare_predicate(cmp, *, context: str): + allowed = {"eq", "ne", "lt", "le", "gt", "ge", "oeq", "one", "olt", "ole", "ogt", "oge"} + if not isinstance(cmp, str) or cmp not in allowed: + expected = ", ".join(sorted(allowed)) + raise ValueError(f"{context} does not support comparison {cmp!r}; expected one of {expected}") + return cmp + + +def _constant_int(value): + if isinstance(value, int) and not isinstance(value, bool): + return value + if isinstance(value, tirx.IntImm): + return int(value.value) + return None + + +def _annotation_value(value): + if isinstance(value, str): + return tirx.StringImm(value) + return value + + +def _call_vmi(op_name: str, result_dtype, *args, **annotations): + attrs = {key: _annotation_value(value) for key, value in annotations.items() if value is not None and key not in {"loc", "ip"}} + args = tuple(arg for arg in args if arg is not None) + return tirx.call_intrin(str(result_dtype), _Op(f"tl.vmi.{op_name}"), *args, annotations=attrs or None) + + +def _scope_guarded(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + require_vmi_scope(f"T.vmi.{fn.__name__}") + return fn(*args, **kwargs) + + return wrapper + + +def _normalize_pto_vcvt_rounding(mode, *, context: str, allowed=None): + token = mode + if not isinstance(token, str): + token = str(token) + if "." in token: + token = token.rsplit(".", 1)[-1] + normalized = token.strip().upper() + allowed_modes = set(allowed or {"R", "A", "H", "Z"}) + if normalized not in allowed_modes: + expected = ", ".join(sorted(allowed_modes)) + raise ValueError(f"{context} does not support rounding {mode!r}; expected one of {expected}") + return normalized + + +def _validate_pto_load_modes( + context: str, + *, + dist_mode, + group, + stride, + block_stride, + repeat_stride, + allow_group_brc: bool, + allowed_dist_modes, +): + if dist_mode is not None and dist_mode not in allowed_dist_modes: + expected = ", ".join(repr(mode) for mode in sorted(allowed_dist_modes, key=str)) + raise TypeError(f"{context} does not support dist_mode={dist_mode!r}; expected one of {expected}") + + if group is not None: + if dist_mode is not None and (not allow_group_brc or dist_mode != "brc"): + raise TypeError(f"{context} does not allow dist_mode together with group") + if block_stride is not None or repeat_stride is not None: + raise TypeError(f"{context} does not allow block_stride together with group") + if stride is None: + raise TypeError(f"{context} with group=... requires stride") + return + + if block_stride is not None or repeat_stride is not None: + if dist_mode is not None: + raise TypeError(f"{context} does not allow dist_mode together with block_stride") + if block_stride is None or repeat_stride is None: + raise TypeError(f"{context} requires block_stride and repeat_stride together") + if stride is not None: + raise TypeError(f"{context} does not allow stride together with block_stride") + return + + if stride is not None: + raise TypeError(f"{context} accepts stride only when group is provided") + + +def _zero_offset(): + return tirx.IntImm("int32", 0) + + +def _looks_like_mask(value) -> bool: + if value is None: + return False + try: + return str(_dtype_of(value)).startswith("bool") + except Exception: # pylint: disable=broad-except + return False + + +def _resolve_ptr_and_offset(source, offset=None, *, access_type: str, extent=None): + """Normalize VMI address syntax into PTODSL's explicit ptr+offset ABI.""" + if isinstance(source, (BufferLoad, Buffer, BufferRegion)) and offset is not None: + raise TypeError("T.vmi buffer addresses do not accept an offset; express the offset in the buffer index") + if isinstance(source, BufferLoad): + # Keep the linear element offset outside access_ptr. In particular, + # FP4 callers convert this logical offset to f4e2m1x2 units before + # PTO emission; leaving it embedded in access_ptr would make PTO + # interpret it as an already-packed offset. + strides = source.buffer.strides + if len(strides) != len(source.indices): + strides = [] + for index in range(len(source.indices)): + stride = 1 + for extent_dim in source.buffer.shape[index + 1 :]: + stride = stride * extent_dim + strides.append(stride) + linear_offset = None + for index, stride in zip(source.indices, strides): + term = index * stride + linear_offset = term if linear_offset is None else linear_offset + term + ptr = access_ptr(source.buffer, access_type, extent=extent, offset=0) + return ptr, linear_offset if linear_offset is not None else _zero_offset() + if isinstance(source, (Buffer, BufferRegion)): + ptr = access_ptr(source, access_type, extent=extent, offset=0) + return ptr, _zero_offset() + if offset is None: + return source, _zero_offset() + return source, offset + + +def _pair_result_dtype(value): + return _dtype_of(value) + + +def _vector_result_dtype(value, size=None, *, elem_dtype=None): + elem = _dtype(elem_dtype) if elem_dtype is not None else _element_dtype_of(value) + lanes = _lanes_of(value) if size is None else _normalize_lanes(size, context="T.vmi result size") + return _dtype(f"{elem}x{lanes}") + + +def _vgather_result_element_dtype(source_elem, offsets): + """Derive the result element type for the A5 byte-widening gather path.""" + source_elem = _dtype(source_elem) + offsets_elem = _element_dtype_of(offsets) + source_name = str(source_elem) + offsets_name = str(offsets_elem) + if source_name not in {"uint8", "int8"}: + return source_elem + + # The A5 b8->b16 gather path requires unsigned 16-bit element offsets and + # widens the source to a matching 16-bit result type. + if offsets_name != "uint16": + raise TypeError("T.vmi.vgather(...) requires uint16 offsets for 8-bit integer sources") + return _dtype("uint16" if source_name == "uint8" else "int16") + + +class VmiPair: + """Lazy pair wrapper for multi-result VMI calls.""" + + def __init__(self, pair): + self._pair = pair + + @property + def dtype(self): + return self._pair.dtype + + def __iter__(self): + yield _pair_get(self, 0) + yield _pair_get(self, 1) + + def __len__(self): + return 2 + + def __getitem__(self, index): + if index not in (0, 1): + raise IndexError("VmiPair only contains two values") + return _pair_get(self, index) + + +def inside_vmi() -> bool: + """Return True when executing inside ``with T.SimdVF():``.""" + return _inside_simdvf() + + +def require_vmi_scope(context: str = "T.vmi") -> None: + """Raise if a VMI op is used outside ``with T.SimdVF():``.""" + if not inside_vmi(): + raise RuntimeError(f"{context} operations require a `T.SimdVF()` scope") + + +def vreg(lanes: int, dtype_like): + """Return a logical VMI vector dtype descriptor.""" + lanes = _normalize_lanes(lanes, context="T.vmi.vreg(...)") + elem = _scalar_dtype(dtype_like, context="T.vmi.vreg(...)") + return _dtype(f"{elem}x{lanes}") + + +def mask(lanes: int): + """Return a logical VMI predicate dtype descriptor.""" + lanes = _normalize_lanes(lanes, context="T.vmi.mask(...)") + return _dtype(f"boolx{lanes}") + + +@_scope_guarded +def _pair_get(pair, index): + """Extract one value from a VMI pair.""" + inner = pair._pair if isinstance(pair, VmiPair) else pair + return tirx.call_intrin(str(inner.dtype), _Op("tl.vmi.pair_get"), inner, index) + + +def _pack_fp4_element_offset(value, *, context: str): + """Convert a logical FP4 element offset to packed-pair units.""" + constant = _constant_int(value) + if constant is not None: + if constant % 2: + raise ValueError(f"{context} requires an even FP4 element offset, got {constant}") + return constant // 2 + # Dynamic offsets are expressed in logical FP4 element units and lowered + # to packed-pair units. + return value // 2 + + +@_scope_guarded +def vload( + source, + offset=None, + *, + size, + to_dtype=None, + stride=None, + block_stride=None, + repeat_stride=None, + dist_mode=None, + group=None, + loc=None, + ip=None, +): + """Load a VMI vector from either `a_ub[0]`-style buffer loads or ptr+offset form. + + A VMI load always reads the full ``size`` lanes starting at the load + address: the ISA intentionally keeps load instructions unmasked for + performance, so masks apply to computation and stores, never to the + load itself. Reading a full vector past a logical tail is still safe + on device — the top 8K of UB is reserved as stack space, so the extra + read-only lanes cannot fault; they carry undefined values that the + compute and store masks must keep out of valid outputs. + """ + _validate_pto_load_modes( + "T.vmi.vload(...)", + dist_mode=dist_mode, + group=group, + stride=stride, + block_stride=block_stride, + repeat_stride=repeat_stride, + allow_group_brc=True, + allowed_dist_modes={None, "continuous", "dintlv", "unpack", "brc"}, + ) + if size is None: + raise TypeError("T.vmi.vload(...) requires size") + if to_dtype is not None and dist_mode != "unpack": + raise TypeError('T.vmi.vload(...) accepts to_dtype only when dist_mode="unpack"') + source_elem = _require_address_element_dtype(source, context="T.vmi.vload(...)") + # Keep si*/ui* spelling for PTODSL annotations. Unpack vload itself is not + # legalized on the current VPTO path; this preserves a correct annotation + # if enabled later. + to_dtype_annot = None + is_packed_fp4 = str(source_elem) == "float4_e2m1fn" + if is_packed_fp4: + if size % 2: + raise ValueError("T.vmi.vload(...) requires an even FP4 lane count for packed FP4 storage") + if to_dtype is not None or dist_mode == "unpack": + raise TypeError("T.vmi.vload(...) does not unpack or convert packed FP4 source vectors") + if dist_mode == "dintlv": + raise TypeError('T.vmi.vload(...) does not support packed FP4 with dist_mode="dintlv"') + if group is not None or block_stride is not None or repeat_stride is not None or dist_mode == "brc": + raise TypeError("T.vmi.vload(...) only supports contiguous packed FP4 loads") + if dist_mode == "unpack": + if to_dtype is None: + raise TypeError('T.vmi.vload(...) requires to_dtype when dist_mode="unpack"') + to_dtype_annot = _pto_to_dtype_annotation(to_dtype) + to_dtype = _scalar_dtype(to_dtype, context="T.vmi.vload(..., to_dtype=...)") + if getattr(to_dtype, "bits", None) != 2 * getattr(source_elem, "bits", None): + raise TypeError("T.vmi.vload(...) unpack must widen by exactly one step") + + result_dtype = _vector_result_dtype( + source, + size, + elem_dtype=to_dtype if to_dtype is not None else source_elem, + ) + ptr, extra_offset = _resolve_ptr_and_offset(source, offset, access_type="r", extent=size) + if is_packed_fp4: + extra_offset = _pack_fp4_element_offset(extra_offset, context="T.vmi.vload(...)") + attrs = { + "size": size, + "to_dtype": to_dtype_annot, + "stride": stride, + "block_stride": block_stride, + "repeat_stride": repeat_stride, + "dist_mode": dist_mode, + "group": group, + "loc": loc, + "ip": ip, + } + call = _call_vmi("vload", result_dtype, ptr, extra_offset, **attrs) + if dist_mode == "dintlv": + return _wrap_pair(call) + return call + + +@_scope_guarded +def vstore( + values, + destination, + offset=None, + mask=None, + *, + stride=None, + block_stride=None, + repeat_stride=None, + dist_mode=None, + group=None, + pmode=None, + loc=None, + ip=None, +): + """Store a VMI vector to either `a_ub[0]`-style buffer loads or ptr+offset form.""" + if mask is None and _looks_like_mask(offset): + mask = offset + offset = None + _validate_pto_load_modes( + "T.vmi.vstore(...)", + dist_mode=dist_mode, + group=group, + stride=stride, + block_stride=block_stride, + repeat_stride=repeat_stride, + allow_group_brc=False, + allowed_dist_modes={None, "continuous", "dintlv"}, + ) + if group is not None and mask is not None: + raise TypeError("T.vmi.vstore(...) group mode does not take a mask operand") + value_dtype = None + if dist_mode == "dintlv": + if isinstance(values, VmiPair): + value_dtype = _element_dtype_of(values) + values = tuple(values) + elif not _is_sequence(values) or len(values) != 2: + raise TypeError('T.vmi.vstore(...) with dist_mode="dintlv" requires an (even, odd) pair') + else: + _require_same_vreg_type(values[0], values[1], context="T.vmi.vstore(...)", allow_packed_fp4=True) + elif isinstance(values, VmiPair) or _is_sequence(values): + raise TypeError('T.vmi.vstore(...) expects a single VMI vector unless dist_mode="dintlv"') + + store_extent = _lanes_of(values[0]) if _is_sequence(values) else _lanes_of(values) + ptr, extra_offset = _resolve_ptr_and_offset(destination, offset, access_type="w", extent=store_extent) + value_args = list(values) if _is_sequence(values) else [values] + if value_dtype is None: + value_dtype = _element_dtype_of(values[0]) if _is_sequence(values) else _element_dtype_of(values) + if str(value_dtype) == "float4_e2m1fn": + if store_extent % 2: + raise ValueError("T.vmi.vstore(...) requires an even FP4 lane count for packed FP4 storage") + if dist_mode == "dintlv": + raise TypeError('T.vmi.vstore(...) does not support packed FP4 with dist_mode="dintlv"') + if group is not None or block_stride is not None or repeat_stride is not None: + raise TypeError("T.vmi.vstore(...) only supports contiguous packed FP4 stores") + extra_offset = _pack_fp4_element_offset(extra_offset, context="T.vmi.vstore(...)") + if mask is not None: + mask = _require_mask(mask, context="T.vmi.vstore(...) of FP4") + if not str(_dtype_of(mask)).startswith("bool"): + raise TypeError("T.vmi.vstore(...) of FP4 expects a VMI mask operand") + physical_lanes = store_extent // 2 + if _lanes_of(mask) != physical_lanes: + raise TypeError( + "T.vmi.vstore(...) of FP4 requires a physical FP4x2 mask with " + f"half as many lanes as its {store_extent}-lane logical FP4 value " + f"(expected {physical_lanes}, got {_lanes_of(mask)})" + ) + elif mask is not None: + mask = _require_mask(mask, context="T.vmi.vstore(...)") + return _call_vmi( + "vstore", + "void", + *value_args, + ptr, + extra_offset, + mask, + stride=stride, + block_stride=block_stride, + repeat_stride=repeat_stride, + dist_mode=dist_mode, + group=group, + pmode=pmode, + loc=loc, + ip=ip, + ) + + +@_scope_guarded +def vci(base, *, size, order=None, loc=None, ip=None): + if size is None: + raise TypeError("T.vmi.vci(...) requires size") + _reject_packed_fp4(base, context="T.vmi.vci(...)") + result_dtype = vreg(size, _element_dtype_of(base)) + return _call_vmi("vci", result_dtype, base, size=size, order=order, loc=loc, ip=ip) + + +def _binary_same_dtype(name): + @_scope_guarded + def wrapper(lhs, rhs, mask=None, *, pmode=None, loc=None, ip=None): + result_dtype = _require_same_vreg_type(lhs, rhs, context=f"T.vmi.{name}(...)") + args = [lhs, rhs] + if mask is not None: + args.append(_require_compatible_mask(mask, _lanes_of(lhs), context=f"T.vmi.{name}(...)")) + return _call_vmi(name, result_dtype, *args, pmode=pmode, loc=loc, ip=ip) + + return wrapper + + +def _unary_same_dtype(name): + @_scope_guarded + def wrapper(source, mask=None, *, pmode=None, loc=None, ip=None): + _reject_packed_fp4(source, context=f"T.vmi.{name}(...)") + args = [source] + if mask is not None: + args.append(mask) + return _call_vmi(name, _dtype_of(source), *args, pmode=pmode, loc=loc, ip=ip) + + return wrapper + + +def _vec_scalar_same_dtype(name): + @_scope_guarded + def wrapper(source, scalar, mask, *, pmode=None, loc=None, ip=None): + _reject_packed_fp4(source, context=f"T.vmi.{name}(...)") + return _call_vmi( + name, + _dtype_of(source), + source, + scalar, + _require_mask(mask, context=f"T.vmi.{name}(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + return wrapper + + +vadd = _binary_same_dtype("vadd") +vsub = _binary_same_dtype("vsub") +vmul = _binary_same_dtype("vmul") +vmax = _binary_same_dtype("vmax") +vmin = _binary_same_dtype("vmin") +vand = _binary_same_dtype("vand") +vor = _binary_same_dtype("vor") +vxor = _binary_same_dtype("vxor") +vshl = _binary_same_dtype("vshl") +vshr = _binary_same_dtype("vshr") + +vabs = _unary_same_dtype("vabs") +vneg = _unary_same_dtype("vneg") +vrelu = _unary_same_dtype("vrelu") +vexp = _unary_same_dtype("vexp") +vln = _unary_same_dtype("vln") +vsqrt = _unary_same_dtype("vsqrt") +vnot = _unary_same_dtype("vnot") + +vadds = _vec_scalar_same_dtype("vadds") +vmuls = _vec_scalar_same_dtype("vmuls") +vmaxs = _vec_scalar_same_dtype("vmaxs") +vmins = _vec_scalar_same_dtype("vmins") +vshls = _vec_scalar_same_dtype("vshls") +vshrs = _vec_scalar_same_dtype("vshrs") + + +# Hardware divide is the default. precision='exact' expands the CANN +# vdiv_0ulp_ftz_true residual search. Changing the no-kwarg default would +# rewrite every T.vmi.vdiv call. +_VDIV_PRECISION_ALIASES = { + None: "hw", + "ftz_true": "hw", + "exact": "exact", + "vdiv_0ulp_ftz_true": "exact", +} + + +def _vdiv_hw(lhs, rhs, mask=None, *, pmode=None, loc=None, ip=None): + result_dtype = _require_same_vreg_type(lhs, rhs, context="T.vmi.vdiv(...)") + args = [lhs, rhs] + if mask is not None: + args.append(_require_compatible_mask(mask, _lanes_of(lhs), context="T.vmi.vdiv(...)")) + return _call_vmi("vdiv", result_dtype, *args, pmode=pmode, loc=loc, ip=ip) + + +@_scope_guarded +def vdiv(lhs, rhs, mask=None, *, precision=None, pmode=None, loc=None, ip=None): + """Vector divide. ``precision='exact'`` expands CANN ``vdiv_0ulp_ftz_true``.""" + context = "T.vmi.vdiv(...)" + key = None if precision is None else str(precision).lower() + if key not in _VDIV_PRECISION_ALIASES: + raise ValueError(f"{context} precision {precision!r} is not valid. Valid: exact, ftz_true, vdiv_0ulp_ftz_true") + if _VDIV_PRECISION_ALIASES[key] == "hw": + return _vdiv_hw(lhs, rhs, mask, pmode=pmode, loc=loc, ip=ip) + + _require_same_vreg_type(lhs, rhs, context=context) + if str(_element_dtype_of(lhs)) != "float32": + raise ValueError(f"{context} precision={precision!r} requires float32 operands (got {_element_dtype_of(lhs)})") + if pmode is not None: + raise ValueError(f"{context} precision={precision!r} does not support pmode={pmode!r}") + mask = _require_compatible_mask(mask, _lanes_of(lhs), context=context) + lanes = _lanes_of(lhs) + z = _bind(_vdiv_hw(lhs, rhs, mask)) + z_bits = _bind(vinterpret_cast(z, "uint32")) + inf_nan = _bind(vor(z_bits, vbrc(tirx.const(0x80000000, "uint32"), size=lanes), mask)) + zero_cmp = _bind(vcmps(z, tirx.const(0.0, "float32"), mask, "eq")) + inf_nan_cmp = _bind(vcmp(inf_nan, vbrc(tirx.const(0xFF800000, "uint32"), size=lanes), mask, "ge")) + rhs_bits = _bind(vinterpret_cast(rhs, "uint32")) + exp = _bind( + vand( + vshrs(rhs_bits, tirx.const(23, "uint32"), mask), + vbrc(tirx.const(0xFF, "uint32"), size=lanes), + mask, + ) + ) + one_bits = vbrc(tirx.const(0x3F800000, "uint32"), size=lanes) + scale_bits = vshls( + vsub(vbrc(tirx.const(254, "uint32"), size=lanes), exp, mask), + tirx.const(23, "uint32"), + mask, + ) + scale_bits = vsel(vcmps(exp, tirx.const(0, "uint32"), mask, "eq"), one_bits, scale_bits) + scale_bits = vsel( + vcmp(exp, vbrc(tirx.const(253, "uint32"), size=lanes), mask, "gt"), + one_bits, + scale_bits, + ) + scale = _bind(vinterpret_cast(scale_bits, "float32")) + lhs_s = _bind(vmul(lhs, scale, mask)) + rhs_s = _bind(vmul(rhs, scale, mask)) + y = _bind(vmuls(rhs_s, tirx.const(-1.0, "float32"), mask)) + r = _bind(vabs(vmula(lhs_s, z, y, mask), mask)) + z_i = _bind(vinterpret_cast(z, "int32")) + z_pre = _bind(vinterpret_cast(vadds(z_i, tirx.const(-1, "int32"), mask), "float32")) + z_next = _bind(vinterpret_cast(vadds(z_i, tirx.const(1, "int32"), mask), "float32")) + r_pre = _bind(vabs(vmula(lhs_s, z_pre, y, mask), mask)) + r_next = _bind(vabs(vmula(lhs_s, z_next, y, mask), mask)) + z_corr = _bind(vsel(vcmp(r_pre, r, mask, "lt"), z_pre, z)) + r_best = _bind(vmin(r, r_pre, mask)) + z_corr = _bind(vsel(vcmp(r_next, r_best, mask, "lt"), z_next, z_corr)) + z_corr = _bind(vsel(inf_nan_cmp, z, z_corr)) + return vsel(zero_cmp, z, z_corr) + + +@_scope_guarded +def vcmp(lhs, rhs, seed, cmp, *, pmode=None, loc=None, ip=None): + context = "T.vmi.vcmp(...)" + _require_same_vreg_type(lhs, rhs, context=context) + seed = _require_compatible_mask(seed, _lanes_of(lhs), context=context) + cmp = _normalize_compare_predicate(cmp, context=context) + return _call_vmi("vcmp", _dtype_of(seed), lhs, rhs, seed, cmp, pmode=pmode, loc=loc, ip=ip) + + +@_scope_guarded +def vcmps(source, scalar, seed, cmp, *, pmode=None, loc=None, ip=None): + context = "T.vmi.vcmps(...)" + _reject_packed_fp4(source, context=context) + seed = _require_compatible_mask(seed, _lanes_of(source), context=context) + cmp = _normalize_compare_predicate(cmp, context=context) + return _call_vmi("vcmps", _dtype_of(seed), source, scalar, seed, cmp, pmode=pmode, loc=loc, ip=ip) + + +@_scope_guarded +def vsel(mask, true_value, false_value, *, pmode=None, loc=None, ip=None): + context = "T.vmi.vsel(...)" + result_dtype = _require_same_vreg_type(true_value, false_value, context=context) + mask = _require_compatible_mask(mask, _lanes_of(true_value), context=context) + return _call_vmi("vsel", result_dtype, mask, true_value, false_value, pmode=pmode, loc=loc, ip=ip) + + +@_scope_guarded +def vselr(source, index, *, loc=None, ip=None): + _reject_packed_fp4(source, context="T.vmi.vselr(...)") + _require_integer_vector(index, context="T.vmi.vselr(...)") + if _lanes_of(source) != _lanes_of(index): + raise TypeError("T.vmi.vselr(...) requires source and index lane counts to match") + return _call_vmi("vselr", _dtype_of(source), source, index, loc=loc, ip=ip) + + +@_scope_guarded +def vbrc(value, *, size, group=None, loc=None, ip=None): + if size is None: + raise TypeError("T.vmi.vbrc(...) requires size") + _reject_packed_fp4(value, context="T.vmi.vbrc(...)") + if group is not None: + if isinstance(group, bool) or not isinstance(group, int): + raise TypeError("T.vmi.vbrc(...) requires group to be a positive Python integer") + if group <= 0: + raise ValueError(f"T.vmi.vbrc(...) requires group to be positive, got {group!r}") + if not hasattr(_dtype_of(value), "lanes"): + raise TypeError("T.vmi.vbrc(...) with group=... requires a VMI vector input") + if _lanes_of(value) != group: + raise ValueError("T.vmi.vbrc(...) with group=... requires the input lane count to match group") + return _call_vmi( + "vbrc", + vreg(size, _element_dtype_of(value)), + value, + size=size, + group=group, + loc=loc, + ip=ip, + ) + + +@_scope_guarded +def vcadd(source, mask, *, group=None, pmode=None, reassoc=None, loc=None, ip=None): + context = "T.vmi.vcadd(...)" + _reject_packed_fp4(source, context=context) + elem_dtype = _element_dtype_of(source) + if reassoc is None and str(elem_dtype).startswith(("float", "bfloat")): + raise TypeError( + f"{context} on floating-point vectors requires an explicit reassoc argument; spell out reassoc=True or reassoc=False" + ) + if reassoc is not None and not isinstance(reassoc, bool): + raise TypeError(f"{context} requires reassoc to be the Python boolean True or False; received {reassoc!r}") + lanes = 1 if group is None else _normalize_lanes(group, context="T.vmi.vcadd(..., group=...)") + if group is not None and _lanes_of(source) % group != 0: + raise ValueError(f"{context} requires source lanes to be divisible by group") + mask = _require_compatible_mask(mask, _lanes_of(source), context=context) + return _call_vmi( + "vcadd", + vreg(lanes, elem_dtype), + source, + mask, + group=group, + pmode=pmode, + reassoc=reassoc, + loc=loc, + ip=ip, + ) + + +@_scope_guarded +def vcmax(source, mask, *, group=None, pmode=None, loc=None, ip=None): + _reject_packed_fp4(source, context="T.vmi.vcmax(...)") + lanes = 1 if group is None else _normalize_lanes(group, context="T.vmi.vcmax(..., group=...)") + if group is not None and _lanes_of(source) % group != 0: + raise ValueError("T.vmi.vcmax(...) requires source lanes to be divisible by group") + mask = _require_compatible_mask(mask, _lanes_of(source), context="T.vmi.vcmax(...)") + return _call_vmi("vcmax", vreg(lanes, _element_dtype_of(source)), source, mask, group=group, pmode=pmode, loc=loc, ip=ip) + + +@_scope_guarded +def vcmin(source, mask, *, group=None, pmode=None, loc=None, ip=None): + _reject_packed_fp4(source, context="T.vmi.vcmin(...)") + lanes = 1 if group is None else _normalize_lanes(group, context="T.vmi.vcmin(..., group=...)") + if group is not None and _lanes_of(source) % group != 0: + raise ValueError("T.vmi.vcmin(...) requires source lanes to be divisible by group") + mask = _require_compatible_mask(mask, _lanes_of(source), context="T.vmi.vcmin(...)") + return _call_vmi("vcmin", vreg(lanes, _element_dtype_of(source)), source, mask, group=group, pmode=pmode, loc=loc, ip=ip) + + +@_scope_guarded +def vcvt(source, to_dtype=None, mask=None, *, rounding=None, saturate=None, pmode=None, loc=None, ip=None): + if mask is not None: + raise TypeError("T.vmi.vcvt does not support masked form") + if to_dtype is None: + raise TypeError("T.vmi.vcvt(...) requires to_dtype") + to_dtype = _scalar_dtype(to_dtype, context="T.vmi.vcvt(..., to_dtype=...)") + src_dt = _element_dtype_of(source) + if str(to_dtype) == "float4_e2m1fn": + if str(src_dt) != "bfloat16": + raise TypeError("T.vmi.vcvt(...) supports packed FP4 only for bfloat16 to float4_e2m1fn") + if _lanes_of(source) % 2: + raise ValueError("T.vmi.vcvt(...) requires an even BF16 lane count for packed FP4") + if rounding is not None: + rounding = _normalize_pto_vcvt_rounding( + rounding, + context="T.vmi.vcvt(..., rounding=...)", + allowed={"R", "A", "F", "Z", "C"}, + ) + if saturate is not None: + raise ValueError("T.vmi.vcvt(...) does not support saturate for bfloat16 to packed FP4 conversion") + elif str(src_dt) == "float4_e2m1fn": + if str(to_dtype) not in {"bfloat16", "float32"}: + raise TypeError("T.vmi.vcvt(...) supports packed FP4 sources only for float4_e2m1fn to bfloat16 or float32") + if _lanes_of(source) % 2: + raise ValueError("T.vmi.vcvt(...) requires an even packed FP4 logical lane count") + # Result is vreg(logical, to_dtype). Check that width, not the + # physical packed-pair count (logical // 2), which is not a VMI size. + _require_vmi_lane_count( + _lanes_of(source), + context="T.vmi.vcvt(...) packed FP4 source lane count", + ) + if rounding is not None: + raise ValueError("T.vmi.vcvt(...) does not support rounding for packed FP4 source conversion") + if saturate is not None: + raise ValueError("T.vmi.vcvt(...) does not support saturate for packed FP4 source conversion") + if pmode is not None: + raise ValueError("T.vmi.vcvt(...) does not support pmode for packed FP4 source conversion") + elif rounding is not None: + rounding = _normalize_pto_vcvt_rounding(rounding, context="T.vmi.vcvt(..., rounding=...)") + return _call_vmi( + "vcvt", + vreg(_lanes_of(source), to_dtype), + source, + to_dtype=str(to_dtype), + rounding=rounding, + saturate=saturate, + pmode=pmode, + loc=loc, + ip=ip, + ) + + +@_scope_guarded +def vinterpret_cast(source, to_dtype=None, *, loc=None, ip=None): + if to_dtype is None: + raise TypeError("T.vmi.vinterpret_cast(...) requires to_dtype") + source_dt = _element_dtype_of(source) + target_dt = _scalar_dtype(to_dtype, context="T.vmi.vinterpret_cast(..., to_dtype=...)") + src_bits = getattr(source_dt, "bits", None) + tgt_bits = getattr(target_dt, "bits", None) + src_lanes = _lanes_of(source) + if src_bits is None or tgt_bits is None: + raise TypeError("T.vmi.vinterpret_cast(...) requires sized element dtypes") + # Same element width keeps lane count; otherwise require bit-total match + # (ASC vintlv pack: 128xbf16 → 64xf32). + if src_bits == tgt_bits: + out_lanes = src_lanes + else: + total = src_lanes * src_bits + if total % tgt_bits != 0: + raise TypeError("T.vmi.vinterpret_cast(...) requires source/target bit totals to match") + out_lanes = total // tgt_bits + # PTOAS rejects lane counts outside the formal VMI set (e.g. int8x128→si64 + # yields int64x16, which is not legal). + _require_vmi_lane_count(out_lanes, context="T.vmi.vinterpret_cast(...)") + return _call_vmi( + "vinterpret_cast", + vreg(out_lanes, target_dt), + source, + to_dtype=_pto_to_dtype_annotation(to_dtype), + loc=loc, + ip=ip, + ) + + +@_scope_guarded +def vexpdif(x, max_value, mask, *, pmode=None, loc=None, ip=None): + _reject_packed_fp4(x, max_value, context="T.vmi.vexpdif(...)") + return _call_vmi( + "vexpdif", _dtype_of(max_value), x, max_value, _require_mask(mask, context="T.vmi.vexpdif(...)"), pmode=pmode, loc=loc, ip=ip + ) + + +@_scope_guarded +def vaxpy(x, acc, alpha, mask, *, pmode=None, loc=None, ip=None): + _reject_packed_fp4(x, acc, context="T.vmi.vaxpy(...)") + return _call_vmi("vaxpy", _dtype_of(acc), x, acc, alpha, _require_mask(mask, context="T.vmi.vaxpy(...)"), pmode=pmode, loc=loc, ip=ip) + + +@_scope_guarded +def vlrelu(x, slope, mask, *, pmode=None, loc=None, ip=None): + _reject_packed_fp4(x, context="T.vmi.vlrelu(...)") + return _call_vmi("vlrelu", _dtype_of(x), x, slope, _require_mask(mask, context="T.vmi.vlrelu(...)"), pmode=pmode, loc=loc, ip=ip) + + +@_scope_guarded +def vprelu(x, alpha, mask, *, pmode=None, loc=None, ip=None): + _reject_packed_fp4(x, context="T.vmi.vprelu(...)") + return _call_vmi("vprelu", _dtype_of(x), x, alpha, _require_mask(mask, context="T.vmi.vprelu(...)"), pmode=pmode, loc=loc, ip=ip) + + +@_scope_guarded +def vmull(a, b, mask, *, pmode=None, loc=None, ip=None): + context = "T.vmi.vmull(...)" + result_dtype = _require_same_vreg_type(a, b, context=context) + elem_dtype = _element_dtype_of(a) + if not str(elem_dtype).startswith(("int", "uint")) or elem_dtype.bits != 32: + raise TypeError(f"{context} requires identical 32-bit integer vectors") + mask = _require_compatible_mask(mask, _lanes_of(a), context=context) + pair = _call_vmi("vmull", result_dtype, a, b, mask, pmode=pmode, loc=loc, ip=ip) + return _wrap_pair(pair) + + +@_scope_guarded +def vmula(acc, lhs, rhs, mask, *, pmode=None, loc=None, ip=None): + _reject_packed_fp4(acc, lhs, rhs, context="T.vmi.vmula(...)") + return _call_vmi("vmula", _dtype_of(acc), acc, lhs, rhs, _require_mask(mask, context="T.vmi.vmula(...)"), pmode=pmode, loc=loc, ip=ip) + + +@_scope_guarded +def vdhist(acc, source, mask, *, loc=None, ip=None): + if str(_dtype_of(acc)) != "uint16x256": + raise TypeError("T.vmi.vdhist(...) requires a uint16x256 accumulator") + if str(_element_dtype_of(source)) != "uint8": + raise TypeError("T.vmi.vdhist(...) requires a uint8 source vector") + mask = _require_compatible_mask(mask, _lanes_of(source), context="T.vmi.vdhist(...)") + return _call_vmi("vdhist", _dtype_of(acc), acc, source, mask, loc=loc, ip=ip) + + +@_scope_guarded +def vchist(acc, source, mask, *, loc=None, ip=None): + if str(_dtype_of(acc)) != "uint16x256": + raise TypeError("T.vmi.vchist(...) requires a uint16x256 accumulator") + if str(_element_dtype_of(source)) != "uint8": + raise TypeError("T.vmi.vchist(...) requires a uint8 source vector") + mask = _require_compatible_mask(mask, _lanes_of(source), context="T.vmi.vchist(...)") + return _call_vmi("vchist", _dtype_of(acc), acc, source, mask, loc=loc, ip=ip) + + +@_scope_guarded +def vgather(source, offsets, mask, *, pmode=None, loc=None, ip=None): + _require_ub_address(source, context="T.vmi.vgather(...)") + source_elem = _require_address_element_dtype(source, context="T.vmi.vgather(...)") + if str(source_elem) == "float4_e2m1fn": + raise TypeError("T.vmi.vgather(...) does not support packed FP4 source buffers") + _require_integer_vector(offsets, context="T.vmi.vgather(...)") + mask = _require_compatible_mask(mask, _lanes_of(offsets), context="T.vmi.vgather(...)") + ptr, offset = _resolve_ptr_and_offset(source, access_type="r", extent=_lanes_of(offsets)) + return _call_vmi( + "vgather", + _vector_result_dtype( + offsets, + elem_dtype=_vgather_result_element_dtype(source_elem, offsets), + ), + ptr, + offset, + offsets, + _require_mask(mask, context="T.vmi.vgather(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + +@_scope_guarded +def vgatherb(source, offsets, mask, *, pmode=None, loc=None, ip=None): + _require_ub_address(source, context="T.vmi.vgatherb(...)") + source_elem = _require_address_element_dtype(source, context="T.vmi.vgatherb(...)") + if str(source_elem) == "float4_e2m1fn": + raise TypeError("T.vmi.vgatherb(...) does not support packed FP4 source buffers") + _require_integer_vector(offsets, context="T.vmi.vgatherb(...)") + mask = _require_compatible_mask(mask, _lanes_of(offsets), context="T.vmi.vgatherb(...)") + ptr, offset = _resolve_ptr_and_offset(source, access_type="r", extent=_lanes_of(mask)) + return _call_vmi( + "vgatherb", + _vector_result_dtype(mask, elem_dtype=source_elem), + ptr, + offset, + offsets, + _require_mask(mask, context="T.vmi.vgatherb(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + +@_scope_guarded +def vscatter(value, destination, offsets, mask, *, pmode=None, loc=None, ip=None): + _require_ub_address(destination, context="T.vmi.vscatter(...)") + _reject_packed_fp4(value, context="T.vmi.vscatter(...)") + destination_dtype = _address_element_dtype(destination) + if destination_dtype is not None and destination_dtype != _element_dtype_of(value): + raise TypeError("T.vmi.vscatter(...) requires value and destination element types to match") + _require_integer_vector(offsets, context="T.vmi.vscatter(...)") + if _lanes_of(value) != _lanes_of(offsets): + raise TypeError("T.vmi.vscatter(...) requires value and offset lane counts to match") + mask = _require_compatible_mask(mask, _lanes_of(value), context="T.vmi.vscatter(...)") + ptr, offset = _resolve_ptr_and_offset(destination, access_type="w", extent=_lanes_of(value)) + return _call_vmi( + "vscatter", + "void", + value, + ptr, + offset, + offsets, + _require_mask(mask, context="T.vmi.vscatter(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + +@_scope_guarded +def create_mask(active_lanes, *, size, group=None, loc=None, ip=None): + if size is None: + raise TypeError("T.vmi.create_mask(...) requires size") + if group is None: + return _call_vmi("create_mask", mask(size), active_lanes, size=size, loc=loc, ip=ip) + if isinstance(group, bool) or not isinstance(group, int): + raise TypeError("T.vmi.create_mask(...) requires group to be a positive Python integer") + if group <= 0: + raise ValueError(f"T.vmi.create_mask(...) requires group to be positive, got {group!r}") + if size % group != 0: + raise ValueError("T.vmi.create_mask(...) requires size to be divisible by group") + group_size = size // group + active_lanes_const = _constant_int(active_lanes) + if active_lanes_const is not None and active_lanes_const > group_size: + raise ValueError("T.vmi.create_mask(...) requires active_lanes to be <= the inferred group_size") + return _call_vmi( + "create_mask", + mask(size), + active_lanes, + group=group, + size=size, + loc=loc, + ip=ip, + ) + + +@_scope_guarded +def vintlv(lhs, rhs, mask, *, pmode=None, loc=None, ip=None): + context = "T.vmi.vintlv(...)" + result_dtype = _require_same_vreg_type(lhs, rhs, context=context) + mask = _require_compatible_mask(mask, _lanes_of(lhs), context=context) + pair = _call_vmi("vintlv", result_dtype, lhs, rhs, mask, pmode=pmode, loc=loc, ip=ip) + return _wrap_pair(pair) + + +@_scope_guarded +def vdintlv(lhs, rhs, mask, *, pmode=None, loc=None, ip=None): + context = "T.vmi.vdintlv(...)" + result_dtype = _require_same_vreg_type(lhs, rhs, context=context) + mask = _require_compatible_mask(mask, _lanes_of(lhs), context=context) + pair = _call_vmi("vdintlv", result_dtype, lhs, rhs, mask, pmode=pmode, loc=loc, ip=ip) + return _wrap_pair(pair) + + +@_scope_guarded +def vunzip(source, to_dtype=None, *, loc=None, ip=None): + """Unzip every element into its low and high half-width halves. + + Returns (low, high) with the source lane count preserved and the element + width halved; lane i of each result is the low/high half of source lane i. + to_dtype selects the half element type and defaults to the unsigned + half-width integer of the source element (uint16 for a 32-bit source). It is + forwarded to the emitted PTO call, where PTODSL derives the half-width result + element type from it and re-checks it against the source width, so it must + agree with the result dtype this wrapper declares. + """ + context = "T.vmi.vunzip(...)" + source_dtype = _dtype_of(source) + lanes = int(getattr(source_dtype, "lanes", 1)) + src_bits = getattr(_element_dtype_of(source), "bits", None) + if src_bits is None: + raise TypeError(context + " requires a sized source element dtype") + if src_bits not in (16, 32): + raise ValueError(context + " requires a 16- or 32-bit source element so the half is 8 or 16 bits; got " + str(src_bits)) + half_bits = int(src_bits) // 2 + if to_dtype is None: + half_elem = _scalar_dtype("uint" + str(half_bits), context=context) + else: + half_elem = _scalar_dtype(to_dtype, context=context) + if getattr(half_elem, "bits", None) != half_bits: + raise TypeError(context + " requires the half element width to be exactly half the source element width") + pair = _call_vmi( + "vunzip", + vreg(lanes, half_elem), + source, + to_dtype=_pto_to_dtype_annotation(half_elem), + loc=loc, + ip=ip, + ) + return _wrap_pair(pair) + + +@_scope_guarded +def vzip(low, high, to_dtype=None, *, loc=None, ip=None): + """Zip low/high half-width halves back (inverse of vunzip). + + to_dtype selects the wide element type and defaults to the unsigned + double-width integer of the half element (uint32 for uint16 halves). As in + vunzip it is forwarded to the emitted PTO call and must agree with the result + dtype this wrapper declares. + """ + context = "T.vmi.vzip(...)" + lanes = _lanes_of(low) + high_lanes = _lanes_of(high) + low_elem = _element_dtype_of(low) + high_elem = _element_dtype_of(high) + if high_lanes != lanes or high_elem != low_elem: + raise ValueError( + context + " requires low and high to share one lane count and " + "element type; got " + str(lanes) + "x" + str(low_elem) + " and " + str(high_lanes) + "x" + str(high_elem) + ) + half_bits = getattr(low_elem, "bits", None) + if half_bits is None: + raise TypeError(context + " requires a sized half element dtype") + if half_bits not in (8, 16): + raise ValueError(context + " requires an 8- or 16-bit half element so the wide element is 16 or 32 bits; got " + str(half_bits)) + wide_bits = int(half_bits) * 2 + if to_dtype is None: + wide_elem = _scalar_dtype("uint" + str(wide_bits), context=context) + else: + wide_elem = _scalar_dtype(to_dtype, context=context) + if getattr(wide_elem, "bits", None) != wide_bits: + raise TypeError(context + " requires the wide element width to be exactly twice the half element width") + result_dtype = vreg(lanes, wide_elem) + return _call_vmi( + "vzip", + result_dtype, + low, + high, + to_dtype=_pto_to_dtype_annotation(wide_elem), + loc=loc, + ip=ip, + ) + + +__all__ = [ + "VmiPair", + "alloc_local", + "alloc_var", + "vzip", + "vunzip", + "create_mask", + "inside_vmi", + "mask", + "require_vmi_scope", + "vabs", + "vadd", + "vadds", + "vaxpy", + "vbrc", + "vchist", + "vcadd", + "vcmax", + "vcmin", + "vcmp", + "vci", + "vcmps", + "vcvt", + "vdiv", + "vdhist", + "vdintlv", + "vexp", + "vexpdif", + "vgather", + "vgatherb", + "vintlv", + "vinterpret_cast", + "vln", + "vlrelu", + "vload", + "vmax", + "vmaxs", + "vmin", + "vmins", + "vmul", + "vmula", + "vmull", + "vmuls", + "vneg", + "vnot", + "vor", + "vprelu", + "vscatter", + "vsel", + "vselr", + "vshl", + "vshls", + "vshr", + "vshrs", + "vstore", + "vrelu", + "vreg", + "vxor", + "vand", + "vsqrt", +] From 664a56f5ce04325c20d0a9e59574a1088180d0d4 Mon Sep 17 00:00:00 2001 From: jimmychou <47636600+jimmychou0@users.noreply.github.com> Date: Wed, 14 Oct 2026 14:27:03 +0800 Subject: [PATCH 02/10] feat(pto-parallel): vectorize T.Parallel inside SimdVF to PTO VMI Add the T.Parallel -> PTO VMI vectorization pipeline: - Verify -> Legalize -> Vectorize pass pipeline ahead of AscendSimdVFLowerParallel, with pass registration, the Python interface, and routing tests (SIMT and hand-written VMI regions are bypassed inside the passes; control-flow rewriting stays out of scope) - Verify: shared region scan and inverse lane-mapping analysis, the LayoutInference-for-PTO branch, 1D/2D lane mapping, and T.SimdVF(lanes=...) frontend support, with contract tests - Vectorize: chunk loops with dynamic prefix tail masks, remaining state reset per row/column and per sequential unit, and dtype/cast codegen for element-wise add/sub/mul/div, unary abs/exp/log/sqrt, and f32 <-> f16/bf16 conversions; runtime lane-uniform float computations are rejected (they would stay scalar inside SIMD_VF) - Device tests: element-wise, tail-mask, and FP8 suites Co-authored-by: qukelin Co-authored-by: jimmychou <47636600+jimmychou0@users.noreply.github.com> --- src/ascend/ir.cc | 26 +- src/ascend/target_utils.cc | 18 + src/ascend/target_utils.h | 3 + src/ascend/transform/layout_inference.cc | 161 +- .../transform/legalize_parallel_to_pto.cc | 35 + src/ascend/transform/parallel_to_pto_utils.cc | 1483 +++++++++++++++++ src/ascend/transform/parallel_to_pto_utils.h | 334 ++++ .../transform/vectorize_parallel_to_pto.cc | 1015 +++++++++++ .../transform/verify_parallel_to_pto.cc | 1371 +++++++++++++++ src/layout/layout.h | 4 + src/op/operator.h | 5 + src/op/parallel.cc | 406 +++++ src/op/parallel.h | 5 + src/op/pto_index_analysis.cc | 204 +++ src/op/pto_index_analysis.h | 71 + .../pto/test_parallel_pipeline_routing.py | 233 +++ .../test_pto_parallel_vectorize_structure.py | 566 +++++++ .../pto/test_pto_parallel_verify_contract.py | 517 ++++++ .../language/pto/test_vectorize_fp8_run.py | 609 +++++++ .../pto/test_vectorize_parallel_pto_run.py | 934 +++++++++++ .../pto/test_vectorize_tail_mask_run.py | 808 +++++++++ tilelang/ascend/language/frame.py | 15 +- tilelang/ascend/pipeline.py | 12 +- tilelang/ascend/transform/__init__.py | 34 + 24 files changed, 8860 insertions(+), 9 deletions(-) create mode 100644 src/ascend/transform/legalize_parallel_to_pto.cc create mode 100644 src/ascend/transform/parallel_to_pto_utils.cc create mode 100644 src/ascend/transform/parallel_to_pto_utils.h create mode 100644 src/ascend/transform/vectorize_parallel_to_pto.cc create mode 100644 src/ascend/transform/verify_parallel_to_pto.cc create mode 100644 src/op/pto_index_analysis.cc create mode 100644 src/op/pto_index_analysis.h create mode 100644 testing/ascend/language/pto/test_parallel_pipeline_routing.py create mode 100644 testing/ascend/language/pto/test_pto_parallel_vectorize_structure.py create mode 100644 testing/ascend/language/pto/test_pto_parallel_verify_contract.py create mode 100644 testing/ascend/language/pto/test_vectorize_fp8_run.py create mode 100644 testing/ascend/language/pto/test_vectorize_parallel_pto_run.py create mode 100644 testing/ascend/language/pto/test_vectorize_tail_mask_run.py diff --git a/src/ascend/ir.cc b/src/ascend/ir.cc index 8523f230bb..65326ef042 100644 --- a/src/ascend/ir.cc +++ b/src/ascend/ir.cc @@ -258,6 +258,7 @@ class SimdVFFrameNode : public TIRFrameNode { size_t parent_alloc_count_{0}; int64_t vf_latency_{0}; int64_t source_index_{0}; + int64_t lanes_{128}; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; @@ -291,6 +292,10 @@ class SimdVFFrameNode : public TIRFrameNode { ffi::Map simdvf_annotations; simdvf_annotations.Set("tl.vf_source_index", IntImm(DataType::Int(64), source_index_)); + // Logical lane count for the PTO SIMD vector path (64/128/256). AscendC + // lowering ignores it and keeps its physical register width. + simdvf_annotations.Set("tl.simdvf_lanes", + IntImm(DataType::Int(64), lanes_)); if (vf_latency_ > 0) { simdvf_annotations.Set("tl.vf_latency", IntImm(DataType::Int(64), vf_latency_)); @@ -325,10 +330,22 @@ class SimdVFFrame : public TIRFrame { SimdVFFrameNode); }; -SimdVFFrame SimdVF(int64_t vf_latency, int64_t source_index) { +// Shared validation for the lane count. The Python frontend checks the same +// set before calling; this is the backstop for hand-written TIR and other +// FFI callers. Keep the value list in sync with frame.py's +// _SIMDVF_VALID_LANES and PushPtoSimdVFContext in layout_inference.cc. +bool SimdVFLanesValid(int64_t lanes) { + return lanes == 64 || lanes == 128 || lanes == 256; +} + +SimdVFFrame SimdVF(int64_t vf_latency, int64_t source_index, int64_t lanes) { + ICHECK(SimdVFLanesValid(lanes)) + << "ValueError: T.SimdVF lanes must be one of {64, 128, 256}, got " + << lanes; ObjectPtr n = tvm::ffi::make_object(); n->vf_latency_ = vf_latency; n->source_index_ = source_index; + n->lanes_ = lanes; return SimdVFFrame(n); } @@ -473,10 +490,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = reflection; refl::GlobalDef() .def("tl.SimtVF", SimtVF) - .def("tl.SimdVF", SimdVF) .def("tl.Cube", Cube) .def("tl.Vector", Vector) - .def("tl.MixedKernelLaunch", MixedKernelLaunch); + .def("tl.MixedKernelLaunch", MixedKernelLaunch) + .def("tl.SimdVF", + [](int64_t vf_latency, int64_t source_index, int64_t lanes) { + return SimdVF(vf_latency, source_index, lanes); + }); SimtVFFrameNode::RegisterReflection(); SimdVFFrameNode::RegisterReflection(); CubeFrameNode::RegisterReflection(); diff --git a/src/ascend/target_utils.cc b/src/ascend/target_utils.cc index 1121a92ae9..709f43e26e 100644 --- a/src/ascend/target_utils.cc +++ b/src/ascend/target_utils.cc @@ -10,10 +10,28 @@ namespace tvm { namespace tl { +using ffi::String; + bool TargetIsAscend(Target target) { return target.defined() && target->kind->name == "ascend"; } +bool TargetIsPTO(Target target) { + // PTO shares the "ascend" target kind with AscendC; "pto" in target keys + // selects the PTO codegen (see tilelang/ascend/target.py). Keep this the + // single C++ decision point so LayoutInference, Verify/Vectorize passes + // and any later consumer cannot drift apart on the key spelling. + if (!TargetIsAscend(target)) { + return false; + } + for (const String &key : target->keys) { + if (key == "pto") { + return true; + } + } + return false; +} + bool IsAscendVectorizableFP8(DataType dtype) { // NOTE: E8M0 is a special type of FP8 which is not handled here. // We only handle FP8 types which can be represented with diff --git a/src/ascend/target_utils.h b/src/ascend/target_utils.h index f124339a7e..c985a59360 100644 --- a/src/ascend/target_utils.h +++ b/src/ascend/target_utils.h @@ -13,6 +13,9 @@ namespace tvm { namespace tl { bool TargetIsAscend(Target target); +// Whether the target selects PTO codegen within the Ascend architecture +// ("ascend" kind + "pto" in target keys). Mirrors Python target_is_pto. +bool TargetIsPTO(Target target); bool IsAscendVectorizableFP8(DataType dtype); } // namespace tl diff --git a/src/ascend/transform/layout_inference.cc b/src/ascend/transform/layout_inference.cc index 17425a378a..b8312d30a1 100644 --- a/src/ascend/transform/layout_inference.cc +++ b/src/ascend/transform/layout_inference.cc @@ -26,6 +26,7 @@ #include "arith/ir_mutator_with_analyzer.h" #include "arith/ir_visitor_with_analyzer.h" #include "ascend/op/utils.h" +#include "ascend/target_utils.h" #include "ascend/transform/vf_regions.h" #include "backend/common/target_utils.h" #include "config.h" @@ -144,6 +145,10 @@ struct LayoutInferenceResult { Map for_map; Map predicate_map; Map padding_guard_map; + // Parallel loops that live inside a PTO SIMD_VF block. Their for_map entry + // is a lane Fragment that must be written back under the PTO-specific + // annotation key instead of the SIMT parallel_loop_layout. + std::unordered_set pto_simd_parallel_roots; }; /*! \brief Everything the inference engine knows about reducer dst-steering, @@ -273,6 +278,14 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { ICHECK(thread_bounds.defined()) << "thread_bounds_vec_[" << cur_infer_id << "] is not defined."; + // Ops collected inside a PTO SIMD_VF block infer against the [0, L) lane + // range from tl.simdvf_lanes; the enclosing SIMT thread bounds stay + // irrelevant for lane distribution. + bool in_pto_simd_vf = in_pto_simd_vf_vec_[cur_infer_id]; + if (in_pto_simd_vf) { + thread_bounds = pto_lane_bounds_[cur_infer_id]; + } + const int64_t *extent_ptr = as_const_int(thread_bounds->extent); ICHECK(extent_ptr != nullptr) << "thread_bounds->extent is not a constant integer, which is " @@ -289,7 +302,8 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { bind_var_to_expr_, false, strict_layout_map, - candidate_vector_size_limit}, + candidate_vector_size_limit, + in_pto_simd_vf}, level); } catch (const std::bad_optional_access &e) { LOG(FATAL) << "bad_optional_access while inferring layout for op " @@ -686,6 +700,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { Map for_map; Map predicate_map; Map padding_guard_map; + std::unordered_set pto_simd_parallel_roots; ICHECK(infer_list_.size() == thread_index_vec_.size()) << "infer_list_ and thread_index_vec_ size mismatch"; for (int i = 0; i < infer_list_.size(); i++) { @@ -704,6 +719,12 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { if (for_infer->LoopLayoutRequiresPaddingGuard()) { padding_guard_map.Set(for_infer->GetRoot(), Bool(true)); } + if (in_pto_simd_vf_vec_[i]) { + // PTO lane layouts carry no SIMT thread predicate; the vectorize + // pass consumes the lane fragment directly. + pto_simd_parallel_roots.insert(for_infer->GetRoot().get()); + continue; + } // thread_index should be defined if we rely on it ICHECK(thread_index.defined()) << "thread_index is not defined. Cannot retrieve predicate."; @@ -714,7 +735,8 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { } } - return {layout_map, for_map, predicate_map, padding_guard_map}; + return {layout_map, for_map, predicate_map, padding_guard_map, + pto_simd_parallel_roots}; } void Collect(const PrimFunc &f) { @@ -732,6 +754,9 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { ICHECK(target.defined()) << "Layout_Inference: Require the target attribute"; target_ = target.value(); + // Shared C++ decision point: "pto" in keys selects the PTO codegen and + // the SIMD_VF lane layout path (see tilelang/ascend/target.py). + target_is_pto_ = TargetIsPTO(target_); this->operator()(f->body); // Compute floating fragment buffers after collection ComputeFloatingFragmentBuffers(f->body); @@ -746,6 +771,62 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { Range thread_bounds; }; + // PTO SIMD_VF lane context: parallel loops inside the block distribute + // logical points over [0, L) lanes from tl.simdvf_lanes instead of SIMT + // threads. Only the lane bounds are carried: the lane placeholder Var + // comes from the loop layout's own inverse fragment, and the vectorize + // pass recovers real addresses from that same fragment. + // Returns true when the lane context was pushed. A PTO SIMD_VF block + // without any T.Parallel is a pure hand-written VMI region: it needs no + // lanes annotation and gets no lane context. + bool PushPtoSimdVFContext(const SBlockNode *op) { + bool has_parallel = false; + PostOrderVisit(op->body, [&](const ObjectRef &node) { + if (auto *for_node = node.as()) { + if (for_node->kind == ForKind::kParallel) { + has_parallel = true; + } + } + }); + if (!has_parallel) { + return false; + } + auto lanes_anno = op->annotations.Get("tl.simdvf_lanes"); + ICHECK(lanes_anno.has_value()) + << "[PTO layout inference] SIMD_VF block containing T.Parallel must " + "carry the tl.simdvf_lanes annotation (written by " + "T.SimdVF(lanes=...))"; + auto lanes_imm = lanes_anno.value().as(); + ICHECK(lanes_imm != nullptr && lanes_imm->dtype == DataType::Int(64)) + << "[PTO layout inference] tl.simdvf_lanes must be an Int64 " + "constant, got " + << lanes_anno.value(); + int64_t lanes = lanes_imm->value; + // Validate here, not just in the Python frontend: hand-written TIR, + // serialized IR or other upstream passes can also write the annotation. + // VerifyParallelToPTO re-checks this as the final input contract. + ICHECK(lanes == 64 || lanes == 128 || lanes == 256) + << "[PTO layout inference] T.SimdVF lanes must be one of {64, 128, " + "256}, got " + << lanes; + Range lane_bounds = Range::FromMinExtent(make_zero(DataType::Int(32)), + IntImm(DataType::Int(32), lanes)); + pto_simd_ctx_stack_.push_back({lane_bounds}); + return true; + } + + void PopPtoSimdVFContext() { + ICHECK(!pto_simd_ctx_stack_.empty()); + pto_simd_ctx_stack_.pop_back(); + } + + bool InsidePtoSimdVF() const { return !pto_simd_ctx_stack_.empty(); } + + const Range &CurrentPtoLaneBounds() const { + ICHECK(!pto_simd_ctx_stack_.empty()); + return pto_simd_ctx_stack_.back().lane_bounds; + } + Map GetBufferMap() const { Map buffer_map; for (const auto &[var, buffers] : buffer_data_to_buffers_) { @@ -822,6 +903,10 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { thread_index_vec_.push_back(CurrentThreadIndex()); thread_bounds_vec_.push_back(CurrentThreadBounds()); analyzer_vec_.push_back(analyzer_.Clone()); + bool in_pto_simd = InsidePtoSimdVF(); + in_pto_simd_vf_vec_.push_back(in_pto_simd); + pto_lane_bounds_.push_back(in_pto_simd ? CurrentPtoLaneBounds() + : Range()); // Add the tile operator to infer_list_ infer_list_stmt_.push_back(GetRef(op)); @@ -981,6 +1066,10 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { thread_index_vec_.push_back(CurrentThreadIndex()); thread_bounds_vec_.push_back(CurrentThreadBounds()); analyzer_vec_.push_back(analyzer_.Clone()); + bool in_pto_simd = InsidePtoSimdVF(); + in_pto_simd_vf_vec_.push_back(in_pto_simd); + pto_lane_bounds_.push_back(in_pto_simd ? CurrentPtoLaneBounds() + : Range()); } else { IRVisitorWithAnalyzer::VisitStmt(op->body); } @@ -996,7 +1085,17 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { region_thread_scope_stack_.push_back({scope->first, scope->second}); pushed_region_scope = true; } + // PTO SIMD_VF: push the lane context before visiting the body so + // nested ParallelOps are collected with [0, L) lane bounds; the region + // itself still contributes nothing to the worklist. + bool pushed_pto_ctx = false; + if (target_is_pto_ && op->name_hint == "SIMD_VF") { + pushed_pto_ctx = PushPtoSimdVFContext(op); + } IRVisitorWithAnalyzer::VisitStmt(op->body); + if (pushed_pto_ctx) { + PopPtoSimdVFContext(); + } // Post-load kLayoutMap from the region block annotation so explicitly // annotated layouts are available during InferLayout. // Must happen after visiting body so buffer_data_to_buffers_ is @@ -1375,6 +1474,20 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer { // [0, 1) — no synthetic fallback Var is ever created. IterVar thread_binding_; std::vector region_thread_scope_stack_; + // PTO SIMD_VF nesting stack. Only non-empty on the PTO target: entering a + // SIMD_VF block under a non-PTO target keeps the stack empty so the + // AscendC SIMD path is untouched. + struct PtoSimdVFContext { + Range lane_bounds; + }; + std::vector pto_simd_ctx_stack_; + // Whether the whole function targets PTO (keys contain "pto"). + bool target_is_pto_ = false; + // Per-infer-list flag: this op was collected inside a PTO SIMD_VF block. + std::vector in_pto_simd_vf_vec_; + // Lane bounds [0, L) captured when the op was collected (valid only when + // the matching in_pto_simd_vf_vec_ entry is true). + std::vector pto_lane_bounds_; std::vector thread_index_vec_; std::vector thread_bounds_vec_; std::vector> analyzer_vec_; @@ -1713,7 +1826,23 @@ class LayoutInferencer : public IRMutatorWithAnalyzer { */ Stmt VisitStmt_(const SBlockNode *op) final { if (op->name_hint == "SIMD_VF") { - return ffi::GetRef(op); + // Gate on this block, not on the function: only a SIMD_VF block that + // actually holds lane-annotated T.Parallel roots needs the descent (to + // write the lane fragment onto those loops). Every other block — a + // non-PTO target, or a hand-written VMI region without T.Parallel — + // keeps the baseline early return, so the AscendC SIMT path sees the + // same IR it saw before this feature and one region never decides for + // another. + if (!ContainsPtoSimdRoot(op->body)) { + return ffi::GetRef(op); + } + // Descend for the lane annotations; only the kLayoutMap writeback is + // skipped for the block itself. + SBlock block = Downcast(IRMutatorWithAnalyzer::VisitStmt_(op)); + if (block.same_as(GetRef(op))) { + return ffi::GetRef(op); + } + return block; } SBlock block = Downcast(IRMutatorWithAnalyzer::VisitStmt_(op)); @@ -1738,6 +1867,26 @@ class LayoutInferencer : public IRMutatorWithAnalyzer { * * @return The For statement with layout annotations attached */ + /*! Whether this SIMD_VF body holds a Parallel loop that the PTO lane + * path annotated (the roots collected from the PTO lane context). */ + bool ContainsPtoSimdRoot(const Stmt &body) const { + if (result_.pto_simd_parallel_roots.empty()) { + return false; + } + bool found = false; + tirx::PostOrderVisit(body, [&](const ObjectRef &node) { + if (found) { + return; + } + if (const auto *for_node = node.as()) { + if (result_.pto_simd_parallel_roots.count(for_node) != 0) { + found = true; + } + } + }); + return found; + } + Stmt VisitStmt_(const ForNode *op) final { if (!result_.for_map.count(GetRef(op))) { return IRMutatorWithAnalyzer::VisitStmt_(op); @@ -1750,6 +1899,12 @@ class LayoutInferencer : public IRMutatorWithAnalyzer { // Store the loop layout as an annotation on the For node (outermost) auto for_ptr = for_node.CopyOnWrite(); + if (result_.pto_simd_parallel_roots.count(op)) { + // PTO SIMD lane fragment: separate key so SIMT lowering (LowerTileOp, + // AscendSimdVFLowerParallel) never consumes it as a thread layout. + for_ptr->annotations.Set(attr::kPtoParallelLoopLayout, loop_layout); + return for_node; + } for_ptr->annotations.Set(attr::kParallelLoopLayout, loop_layout); if (result_.padding_guard_map.count(root)) { for_ptr->annotations.Set(attr::kParallelLoopRequiresPaddingGuard, diff --git a/src/ascend/transform/legalize_parallel_to_pto.cc b/src/ascend/transform/legalize_parallel_to_pto.cc new file mode 100644 index 0000000000..470c339c7d --- /dev/null +++ b/src/ascend/transform/legalize_parallel_to_pto.cc @@ -0,0 +1,35 @@ +/*! + * \file legalize_parallel_to_pto.cc + * \brief Placeholder for the PTO Parallel legalization stage. + * + * Only the registration exists for now, so that the formal PTO pipeline + * order (VerifyParallelToPTO -> LegalizeParallelToPTO -> + * VectorizeParallelToPTO) is routing-testable. The placeholder preserves + * the input IR and claims no semantic support: it neither rewrites nor + * rejects any construct. The safe scalar If/Select/Min/Max rewrites are + * left to a later change. + */ + +#include +#include + +namespace tvm { +namespace tl { + +using namespace tirx; +using namespace tirx::transform; + +tvm::transform::Pass LegalizeParallelToPTO() { + auto pass_func = [](PrimFunc func, const IRModule &mod, + const tvm::transform::PassContext &ctx) { return func; }; + return CreatePrimFuncPass(pass_func, 0, "tl.LegalizeParallelToPTO", {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tl.transform.LegalizeParallelToPTO", + LegalizeParallelToPTO); +} + +} // namespace tl +} // namespace tvm diff --git a/src/ascend/transform/parallel_to_pto_utils.cc b/src/ascend/transform/parallel_to_pto_utils.cc new file mode 100644 index 0000000000..a5fa77faea --- /dev/null +++ b/src/ascend/transform/parallel_to_pto_utils.cc @@ -0,0 +1,1483 @@ +/*! + * \file parallel_to_pto_utils.cc + * \brief Shared read-only analysis for VerifyParallelToPTO and + * VectorizeParallelToPTO. + */ + +#include "parallel_to_pto_utils.h" + +#include +#include +#include + +#include "../../layout/utils.h" + +#include +#include + +namespace tvm { +namespace tl { +namespace pto { + +using namespace tirx; +using ffi::Array; +using ffi::GetRef; +using ffi::Map; + +std::string AccessPatternToString(AccessPattern pattern) { + switch (pattern) { + case AccessPattern::kContinuous: + return "continuous"; + case AccessPattern::kLaneUniform: + return "lane-uniform"; + default: + return "unknown"; + } +} + +// SubstituteVar, ResolveIndexAliases and IsSupportedAddressExpr moved to +// src/op/pto_index_analysis.cc: the layout-planning code in +// src/op and the consumers here share that single implementation. + +PtoInverseMapping InvertPtoLaneLayout(const For &loop) { + auto anno = loop->annotations.Get(attr::kPtoParallelLoopLayout); + if (!anno.has_value()) { + throw PtoAnalysisError( + "pto_parallel_loop_layout is missing on the outermost T.Parallel"); + } + auto fragment_any = anno.value().try_cast(); + if (!fragment_any.has_value()) { + throw PtoAnalysisError( + "pto_parallel_loop_layout must be a tl.Fragment annotation"); + } + Fragment fragment = fragment_any.value(); + + // A 2D unit is the outermost For of a two-layer + // adjacent Parallel nest. The annotation — and the fragment — lives on + // the outer For and describes the full (i, j) input space; the input + // coordinate order stays (i, j) whichever dimension is vectorized. + const ForNode *inner_node = nullptr; + if (const auto *body_for = loop->body.as()) { + if (body_for->kind == ForKind::kParallel) { + inner_node = body_for; + } + } + const bool is_2d = inner_node != nullptr; + const int expected_dim = is_2d ? 2 : 1; + if (fragment->InputDim() != expected_dim) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout input_size must be " << expected_dim + << "D for a " + << (is_2d ? "two-layer nested T.Parallel" : "1D T.Parallel") + << ", got input dim " << fragment->InputDim(); + throw PtoAnalysisError(oss.str()); + } + if (fragment->GetForwardIndex().size() != 1) { + // Extra per-lane position dimensions are a legal-but-unsupported + // layout shape; reject with a layout diagnostic instead of an internal + // assertion inside Forward. + std::ostringstream oss; + oss << "pto_parallel_loop_layout must have exactly one internal " + "position dimension, got " + << fragment->GetForwardIndex().size(); + throw PtoAnalysisError(oss.str()); + } + if (!is_one(fragment->ReplicateExtent())) { + throw PtoAnalysisError( + "pto_parallel_loop_layout requires replicate_size == 1"); + } + + const int64_t *lanes_ptr = as_const_int(fragment->ThreadExtent()); + if (lanes_ptr == nullptr || *lanes_ptr <= 0) { + throw PtoAnalysisError( + "pto_parallel_loop_layout thread extent must be a positive constant"); + } + int64_t lanes = *lanes_ptr; + if (lanes != 64 && lanes != 128 && lanes != 256) { + std::ostringstream oss; + oss << "T.SimdVF lanes must be one of {64, 128, 256}, got " << lanes; + throw PtoAnalysisError(oss.str()); + } + + // Shape contract: the fragment input shape must match the loop extents + // (1D: the single extent; 2D: (outer, inner)). + const int64_t *extent_ptr = as_const_int(loop->extent); + if (extent_ptr == nullptr) { + // The helper must not assert on inputs the callers did not + // pre-validate; a symbolic extent is a layout diagnostic (same rule + // as the inner extent below). + std::ostringstream oss; + oss << "pto_parallel_loop_layout loop extent must be a compile-time " + "constant, got " + << loop->extent; + throw PtoAnalysisError(oss.str()); + } + int64_t outer_extent = *extent_ptr; + int64_t inner_extent = 1; + if (is_2d) { + const int64_t *inner_extent_ptr = as_const_int(inner_node->extent); + if (inner_extent_ptr == nullptr) { + // The helper must not assert on inputs the callers did not + // pre-validate; a symbolic inner extent is a layout diagnostic. + std::ostringstream oss; + oss << "pto_parallel_loop_layout inner loop extent must be a " + "compile-time constant, got " + << inner_node->extent; + throw PtoAnalysisError(oss.str()); + } + inner_extent = *inner_extent_ptr; + } + // Shape contract: the fragment input space is padded on the + // vectorized dimension. 1D: the single input extent must be the padded + // extent P = ceil(E / L) * L. 2D: the kept dimension keeps its original + // loop extent and the selected dimension is padded to P — but which + // dimension is selected is only known jointly with the canonical mapping + // after the inversion, so the 2D shape check runs post-inversion. Here + // the 1D case is checked directly; an unpadded [E] (E % L != 0), an + // over-padded extent, and (for divisible E) any extent != E are layout + // diagnostics. + Array input_shape = fragment->InputShape(); + if (!is_2d) { + if (input_shape.size() != 1u) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout input shape " << input_shape + << " must be 1D for a 1D T.Parallel"; + throw PtoAnalysisError(oss.str()); + } + const auto [pad_1d, padded_1d] = CheckedPtoPadding(outer_extent, lanes); + (void)pad_1d; + const int64_t *frag_extent = as_const_int(input_shape[0]); + if (frag_extent == nullptr || *frag_extent != padded_1d) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout input shape " << input_shape + << " does not match the padded extent " << padded_1d + << " (loop extent " << outer_extent << ", lanes " << lanes << "); " + << "the vectorized dimension must be padded to ceil(E/L)*L"; + throw PtoAnalysisError(oss.str()); + } + } else { + if (input_shape.size() != 2u) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout input shape " << input_shape + << " must be 2D for a two-layer nested T.Parallel"; + throw PtoAnalysisError(oss.str()); + } + // Structural presence check only; the per-dimension contract (kept = + // original, selected = padded) is a joint check with the canonical + // mapping below. + for (const auto &dim : input_shape) { + if (as_const_int(dim) == nullptr) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout input shape " << input_shape + << " must be compile-time constants, got " << input_shape; + throw PtoAnalysisError(oss.str()); + } + } + } + + // Thread-range contract: [0, L). + Range thread_range = fragment->ThreadRange(); + if (thread_range.defined()) { + const int64_t *tr_min = as_const_int(thread_range->min); + const int64_t *tr_ext = as_const_int(thread_range->extent); + if (tr_min == nullptr || *tr_min != 0 || tr_ext == nullptr || + *tr_ext != lanes) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout thread_range must be [0, " << lanes + << "), got [" << thread_range->min << ", " << thread_range->extent + << ")"; + throw PtoAnalysisError(oss.str()); + } + } + + // The fragment is built over int32 placeholder vars (the only + // dtype LayoutInference produces); a loop var of any other dtype (e.g. + // int64) crashes the tvm substitution machinery with a raw assertion + // during inversion. Reject it here with a layout diagnostic instead. + if (loop->loop_var.dtype() != DataType::Int(32)) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout requires an int32 loop variable (the " + "dtype LayoutInference produces), got " + << loop->loop_var.dtype(); + throw PtoAnalysisError(oss.str()); + } + if (is_2d && inner_node->loop_var.dtype() != DataType::Int(32)) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout requires an int32 inner loop variable " + "(the dtype LayoutInference produces), got " + << inner_node->loop_var.dtype(); + throw PtoAnalysisError(oss.str()); + } + + // Invert through the real inverse layout (single source of truth for the + // mapping). InverseWithLevel appends the lane (thread) input after the + // per-lane position input. It raises when the fragment is not a valid + // bijection — e.g. a hand-built layout whose per-thread position extent + // does not cover the input space: Layout::InverseWithLevel throws + // NormalizeIterException when the iteration map has errors, and the + // iteration-map machinery itself reports tvm errors as ffi::Error. Both are + // known input failures here, so they become layout diagnostics. Only this + // call is wrapped, so a genuine implementation error anywhere else still + // surfaces as one instead of being reported as bad input. + auto invert = [&]() { + try { + return fragment->InverseWithLevel(/*require_padding_guard=*/false); + } catch (const NormalizeIterException &err) { + throw PtoAnalysisError(std::string("pto_parallel_loop_layout cannot be " + "inverted: ") + + err.what()); + } catch (const ffi::Error &err) { + throw PtoAnalysisError(std::string("pto_parallel_loop_layout cannot be " + "inverted: ") + + err.what()); + } + }; + auto inverse_info = invert(); + if (inverse_info.second != arith::IterMapLevel::Bijective) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout is not bijectively invertible (level " + << static_cast(inverse_info.second) + << "): " << fragment->DebugOutput(); + throw PtoAnalysisError(oss.str()); + } + Layout inverse = inverse_info.first; + + // Forward the (chunk, lane) coordinates through the inverse layout to + // obtain the logical coordinate expression. The inverse layout has two + // outputs (the logical coordinate and the replicate coordinate); with + // replicate_size == 1 the second output is a constant zero and the + // design-doc convention is to take the logical part only. + Var chunk_var("pto_chunk", loop->loop_var.dtype()); + Var lane_var("pto_lane", loop->loop_var.dtype()); + Array logical = inverse->Forward({chunk_var, lane_var}); + // The output count follows the input dimensionality (1D: (i, replicate); + // 2D: (i, j, replicate)); each branch below checks its own arity. + + // Only the canonical mappings are accepted. Any other + // invertible form (e.g. a reversed lane assignment) is rejected instead + // of being silently replaced with our own formula. + // 1D: i = chunk * L + lane + // 2D (j): j = (chunk * L + lane) % N, i = (chunk * L + lane) / N + // 2D (i): i = (chunk * L + lane) % M, j = (chunk * L + lane) / M + PtoInverseMapping mapping; + mapping.lane_var = lane_var; + mapping.chunk_var = chunk_var; + mapping.lanes = lanes; + + if (!is_2d) { + ICHECK_EQ(logical.size(), 2u); + // The fragment input space is padded to P, so the chunk + // position range is [0, P/L) — the padded chunk count Q. + const int64_t padded = CheckedPtoPadding(outer_extent, lanes).second; + arith::Analyzer probe; + probe.Bind(chunk_var, + Range::FromMinExtent(make_zero(chunk_var->dtype), + IntImm(chunk_var->dtype, padded / lanes))); + probe.Bind(lane_var, Range::FromMinExtent(make_zero(lane_var->dtype), + IntImm(lane_var->dtype, lanes))); + PrimExpr canonical = chunk_var * IntImm(chunk_var->dtype, lanes) + lane_var; + if (!probe.CanProveEqual(logical[0], canonical)) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout does not use the canonical 1D " + "mapping i = chunk * L + lane (got " + << logical[0] << " from " << fragment->DebugOutput() << ")"; + throw PtoAnalysisError(oss.str()); + } + mapping.logical_extent = outer_extent; + mapping.padded_extent = padded; + mapping.chunk_count = padded / lanes; + mapping.tail_lanes = outer_extent % lanes; + mapping.index_expr = logical[0]; + return mapping; + } + + ICHECK_EQ(logical.size(), 3u); + const int64_t M = outer_extent; + const int64_t N = inner_extent; + // The fragment input space is padded on the vectorized dimension, so + // the flattened position t = kept * P_sel + sel lives in [0, K * P_sel) + // and the global chunk count is K * ceil(E / L) (each row padded + // independently — not ceil(K * E / L), which is not equivalent). This + // probe domain is for the canonical mapping check only; the + // address-proof domain that AnalyzeBufferAccess uses binds the per-row + // q range instead. + // Each interpretation gets its own probe with the exact chunk domain + // [0, kept * P_sel / L): a domain bound larger than the fragment's real + // chunk count could make an identity that holds on the real domain + // unprovable, and a smaller one could admit a mapping that only looks + // canonical on part of the space. The lane range [0, L) is shared. + const auto bind_probe = [&](arith::Analyzer *probe, int64_t chunk_bound) { + probe->Bind(chunk_var, + Range::FromMinExtent(make_zero(chunk_var->dtype), + IntImm(chunk_var->dtype, chunk_bound))); + probe->Bind(lane_var, Range::FromMinExtent(make_zero(lane_var->dtype), + IntImm(lane_var->dtype, lanes))); + }; + PrimExpr t = chunk_var * IntImm(chunk_var->dtype, lanes) + lane_var; + // The vectorized coordinate is t % P_sel, the kept one t / P_sel — both + // the modulus and the divisor use the *padded* selected extent. The + // interpretation is accepted only when it also agrees with the fragment + // input shape: the kept dimension keeps its original loop extent and + // the selected dimension is padded up to P_sel (unpadded / over-padded + // / wrongly-modified kept dims each fail here). A degenerate dimension + // (extent 1) makes both canonical forms provable; the shape agreement + // breaks the tie exactly when only one dim was padded, and the inner + // preference remains for fully divisible shapes. + const auto matches = [&](bool select_inner) -> bool { + const int64_t sel_extent = select_inner ? N : M; + const auto [pad, padded] = CheckedPtoPadding(sel_extent, lanes); + const int64_t kept_extent = select_inner ? M : N; + const int64_t *frag_kept = as_const_int(input_shape[select_inner ? 0 : 1]); + const int64_t *frag_sel = as_const_int(input_shape[select_inner ? 1 : 0]); + if (frag_kept == nullptr || frag_sel == nullptr || + *frag_kept != kept_extent || *frag_sel != padded) { + return false; + } + (void)pad; + arith::Analyzer probe; + bind_probe(&probe, kept_extent * padded / lanes); + const PrimExpr &sel = select_inner ? logical[1] : logical[0]; + const PrimExpr &kept = select_inner ? logical[0] : logical[1]; + PrimExpr sel_imm = IntImm(t.dtype(), padded); + return probe.CanProveEqual(sel, FloorMod(t, sel_imm)) && + probe.CanProveEqual(kept, FloorDiv(t, sel_imm)); + }; + const bool inner_ok = matches(/*select_inner=*/true); + const bool outer_ok = matches(/*select_inner=*/false); + if (!inner_ok && !outer_ok) { + std::ostringstream oss; + oss << "pto_parallel_loop_layout does not use a canonical 2D mapping " + "(j = t % P_N, i = t / P_N or i = t % P_M, j = t / P_M with " + "t = chunk * L + lane, P_* = padded selected extent; the kept " + "dim keeps its loop extent in the fragment shape); got (i, j) = (" + << logical[0] << ", " << logical[1] << "), shape " << input_shape + << ", loop extents (" << M << ", " << N << ") from " + << fragment->DebugOutput(); + throw PtoAnalysisError(oss.str()); + } + // Both dims matching is degenerate (e.g. M == N == 1); prefer the inner + // dimension, consistent with the layout-selection rule. + const bool select_inner = inner_ok; + const int64_t sel_extent = select_inner ? N : M; + const auto [pad_sel, P_sel] = CheckedPtoPadding(sel_extent, lanes); + (void)pad_sel; + mapping.is_2d = true; + mapping.outer_var = loop->loop_var; + mapping.inner_var = inner_node->loop_var; + mapping.outer_extent = M; + mapping.inner_extent = N; + mapping.select_inner = select_inner; + mapping.logical_extent = sel_extent; + mapping.padded_extent = P_sel; + mapping.chunk_count = P_sel / lanes; + mapping.tail_lanes = sel_extent % lanes; + mapping.index_expr = t; + return mapping; +} + +MemoryAccess AnalyzeBufferAccess(const Buffer &buffer, + const Array &indices, bool is_write, + const Var &loop_var, + const PtoInverseMapping &mapping, + arith::Analyzer *analyzer, + const Map &bind_env) { + MemoryAccess access; + access.buffer = buffer; + access.elem_dtype = buffer->dtype; + access.is_write = is_write; + + // Resolve Bind aliases in the indices first: a bound index var + // must follow its definition, not be treated as lane-independent. The + // resolver rejects chains that exceed the limit instead of silently + // stopping. + // + // Only *pure coordinate* aliases (loop vars, arithmetic over + // loop vars/consts) may be expanded. A Bind whose definition contains a + // BufferLoad holds an SSA value: re-expanding it into a memory read can + // cross a same-point store and wrongly cancel (the loaded value and the + // current buffer state differ). Addresses that depend on loaded values + // have no supported lowering and are rejected here. + Array resolved; + for (const auto &idx : indices) { + auto opt = ResolveIndexAliases(idx, bind_env); + if (!opt.has_value()) { + std::ostringstream oss; + oss << "buffer `" << buffer->name + << "`: bind alias chain in index is too deep, cyclic, or " + "depends on a loaded value (data-dependent addresses are " + "not yet supported)"; + throw PtoAnalysisError(oss.str()); + } + resolved.push_back(opt.value()); + } + // Address expressions must only contain supported + // nodes; opaque calls in addresses have no lowering rule and would be + // silently collapsed to one evaluation per chunk. + for (const auto &idx : resolved) { + if (!IsSupportedAddressExpr(idx)) { + std::ostringstream oss; + oss << "buffer `" << buffer->name + << "`: address expression contains a loaded value or an " + "unsupported call (loaded values and opaque/extern calls " + "cannot appear in addresses; they would be evaluated once per " + "chunk instead of per element)"; + throw PtoAnalysisError(oss.str()); + } + } + + // indices_after_inverse: substitute loop_var with the inverse mapping's + // own expression (mapping.index_expr — the single source of truth; do + // not rebuild chunk * L + lane here). + Array indices_after_inverse; + for (const auto &idx : resolved) { + indices_after_inverse.push_back(analyzer->Simplify( + SubstituteVar(idx, loop_var.get(), mapping.index_expr))); + } + access.indices_after_inverse = indices_after_inverse; + + // start_indices: bind lane = 0. + Array start_indices; + for (const auto &idx : indices_after_inverse) { + start_indices.push_back(analyzer->Simplify( + SubstituteVar(idx, mapping.lane_var.get(), make_zero(idx.dtype())))); + } + access.start_indices = start_indices; + + // data_offset via ElemOffset (handles strides / elem_offset / versions). + Array elem_offsets = buffer->ElemOffset(indices_after_inverse); + if (elem_offsets.size() != 1) { + std::ostringstream oss; + oss << "buffer `" << buffer->name + << "`: ElemOffset must return exactly one physical axis, got " + << elem_offsets.size(); + throw PtoAnalysisError(oss.str()); + } + access.data_offset = analyzer->Simplify(elem_offsets[0]); + + Array start_elem_offsets = buffer->ElemOffset(start_indices); + ICHECK_EQ(start_elem_offsets.size(), 1u); + access.start_offset = analyzer->Simplify(start_elem_offsets[0]); + + // Classify with a *complete* proof: continuous means + // data_offset == start_offset + lane over the whole [0, L) lane range, + // not a two-point difference. Bind the ranges and require the analyzer + // to prove the identity; anything it cannot prove is rejected. + // The lane-uniform check shares this range-bound proof + // analyzer — with an unbound analyzer a chunk-uniform address such as + // (i // L) * L cannot be proven equal to its lane=0 substitution (the + // proof needs 0 <= lane < L), so legal inputs were mis-rejected. + arith::Analyzer proof; + proof.Bind( + mapping.lane_var, + Range::FromMinExtent(make_zero(mapping.lane_var->dtype), + IntImm(mapping.lane_var->dtype, mapping.lanes))); + // The chunk proof domain is the per-row padded chunk range + // [0, Q) with Q = P / L on the vectorized dimension — for a 2D unit the + // *in-row* chunk count, never the global position K*Q (the global + // flattened position only appears in the Fragment forward/inverse + // layout checks, not in address proofs). + proof.Bind(mapping.chunk_var, + Range::FromMinExtent( + make_zero(mapping.chunk_var->dtype), + IntImm(mapping.chunk_var->dtype, mapping.chunk_count))); + // In a 2D unit the kept coordinate is bounded by its own loop + // extent — a fact the layout candidate already uses when it accepts + // the unit. Binding it here keeps the final proof domain + // identical to the candidate's, so range-dependent but legal addresses + // (e.g. a term multiplied by i // 4 with i in [0, 4)) prove continuous. + // This adds no facts beyond the loop definitions; it proves nothing + // about buffer bounds or cross-point dependencies. + if (mapping.is_2d && mapping.serial_var() != nullptr) { + const VarNode *serial = mapping.serial_var(); + const int64_t kept_extent = + mapping.select_inner ? mapping.outer_extent : mapping.inner_extent; + proof.Bind(GetRef(serial), + Range::FromMinExtent(make_zero(serial->dtype), + IntImm(serial->dtype, kept_extent))); + } + { + PrimExpr identity = + access.data_offset - (access.start_offset + mapping.lane_var); + if (proof.CanProveEqual(identity, make_zero(identity.dtype()))) { + access.pattern = AccessPattern::kContinuous; + return access; + } + } + + // Lane-uniform: the offset does not depend on lane at the current + // program point (version/outer-coordinate indices may stay). Proven + // over the full lane range with the same bound analyzer. + { + PrimExpr at_zero = + proof.Simplify(SubstituteVar(access.data_offset, mapping.lane_var.get(), + make_zero(access.data_offset.dtype()))); + if (proof.CanProveEqual(access.data_offset, at_zero)) { + access.pattern = AccessPattern::kLaneUniform; + return access; + } + } + + std::ostringstream oss; + oss << "buffer `" << buffer->name + << "`: cannot prove the address is continuous " + "(data_offset == start + lane over the full lane range) or " + "lane-uniform; only these two patterns are supported"; + throw PtoAnalysisError(oss.str()); +} + +bool IsSupportedElementDType(DataType dtype) { + // No packed sub-byte element types (FP4 etc.), no + // buffer-element dtypes that already carry vector lanes, and no FP64 + // (the current PTO codegen has no FP64 scalar-type mapping). + if (dtype.bits() < 8 || dtype.lanes() != 1) { + // FP4 (e2m1) and every other sub-byte element type land here: an FP4 + // element is 4 bits wide whether packed or not, so a separate FP4 check + // would be unreachable. + return false; + } + // The PTO codegen only maps 8/16/32-bit scalars; unusual widths (int24) and + // FP64 (no scalar-type mapping) are rejected for buffers and arithmetic + // alike by the width rule below, not just inside Cast planning. + if (!dtype.is_bool() && dtype.bits() != 8 && dtype.bits() != 16 && + dtype.bits() != 32) { + return false; + } + return dtype.is_float() || dtype.is_bfloat16() || dtype.is_int() || + dtype.is_uint() || dtype.is_bool() || IsSupportedFloat8(dtype); +} + +bool IsSupportedFloat8(DataType dtype) { + // Only these two take part in the FP8 conversion matrix. float8_e4m3 is a + // distinct TIR dtype from float8_e4m3fn and is not admitted by name + // similarity; E8M0/HiF8/FNUZ are outside the supported conversion range, + // while FP4 is a packed sub-byte type with no unified element addressing. + return dtype.is_float8_e4m3fn() || dtype.is_float8_e5m2(); +} + +namespace { + +// Whitelist check for *value* expressions. A +// BufferLoad's indices belong to address analysis, not value lowering +// the version-dim FloorMod/FloorDiv inside a load index must not +// be rejected here. The value whitelist does not descend into indices. +class WhitelistVisitor : public ExprVisitor { +public: + std::optional unsupported; + + void VisitExpr_(const BufferLoadNode *) final {} + void VisitExpr_(const VarNode *) final {} + void VisitExpr_(const IntImmNode *) final {} + void VisitExpr_(const FloatImmNode *) final {} + void VisitExpr_(const StringImmNode *) final { + // String constants are metadata markers only. The base + // ExprVisitor declares StringImm and accepts it as a leaf, so the + // VisitExprDefault_ fallback never fired and `Evaluate(StringImm)` + // slipped through Verify into Vectorize's internal error. Reject it + // explicitly here, with the same wording as the lane-use analysis. + if (!unsupported.has_value()) { + unsupported = "string constants are metadata-only; they cannot " + "appear in computation values"; + } + } + void VisitExpr_(const AddNode *op) final { ExprVisitor::VisitExpr_(op); } + void VisitExpr_(const SubNode *op) final { ExprVisitor::VisitExpr_(op); } + void VisitExpr_(const MulNode *op) final { ExprVisitor::VisitExpr_(op); } + void VisitExpr_(const DivNode *op) final { ExprVisitor::VisitExpr_(op); } + void VisitExpr_(const CastNode *op) final { + // Only the converted value is a computation value. The base visitor also + // walks PrimExpr entries of the cast annotations, which would put a + // StringImm rounding hint in front of the value whitelist and reject an + // otherwise legal cast ("string constants are metadata-only"). Attribute + // legality is decided by PlanValueCast, from the raw annotations. + VisitExpr(op->value); + } + // The base ExprVisitor overrides VisitExpr_ for every standard node and + // silently descends, so non-whitelist nodes never reach + // VisitExprDefault_. Reject them explicitly. + void MarkUnsupported(const char *kind) { + if (!unsupported.has_value()) { + unsupported = std::string("unsupported expression node `") + kind + + "` (only buffer loads, constants, " + "scalars, loop indices, Add, Sub, Mul, Div, " + "fabs/exp/log/sqrt and Cast)"; + } + } + void VisitExpr_(const ModNode *) final { MarkUnsupported("Mod"); } + void VisitExpr_(const FloorDivNode *) final { MarkUnsupported("FloorDiv"); } + void VisitExpr_(const FloorModNode *) final { MarkUnsupported("FloorMod"); } + void VisitExpr_(const MinNode *) final { MarkUnsupported("Min"); } + void VisitExpr_(const MaxNode *) final { MarkUnsupported("Max"); } + void VisitExpr_(const EQNode *) final { MarkUnsupported("EQ"); } + void VisitExpr_(const NENode *) final { MarkUnsupported("NE"); } + void VisitExpr_(const LTNode *) final { MarkUnsupported("LT"); } + void VisitExpr_(const LENode *) final { MarkUnsupported("LE"); } + void VisitExpr_(const GTNode *) final { MarkUnsupported("GT"); } + void VisitExpr_(const GENode *) final { MarkUnsupported("GE"); } + void VisitExpr_(const AndNode *) final { MarkUnsupported("And"); } + void VisitExpr_(const OrNode *) final { MarkUnsupported("Or"); } + void VisitExpr_(const NotNode *) final { MarkUnsupported("Not"); } + void VisitExpr_(const SelectNode *) final { + if (!unsupported.has_value()) { + unsupported = + "Select is not vectorizable yet (integer T.abs lowers to Select; " + "it arrives with the PR-2 Select/cmp support)"; + } + } + void VisitExpr_(const RampNode *) final { MarkUnsupported("Ramp"); } + void VisitExpr_(const BroadcastNode *) final { MarkUnsupported("Broadcast"); } + void VisitExpr_(const ShuffleNode *) final { MarkUnsupported("Shuffle"); } + void VisitExpr_(const ReduceNode *) final { MarkUnsupported("Reduce"); } + void VisitExpr_(const CallNode *op) final { + if (auto call_op = op->op.as()) { + const std::string &name = call_op.value()->name; + // Hardware-backed unary math. fabs/exp/log/sqrt map to the VMI + // vabs/vexp/vln/vsqrt instructions; the dtype matrix mirrors the + // ISA documentation (see IsSupportedVectorUnaryOp). + if (name == "tirx.fabs" || name == "tirx.exp" || name == "tirx.log" || + name == "tirx.sqrt") { + if (op->args.size() != 1U) { + if (!unsupported.has_value()) { + unsupported = name + " expects exactly one operand"; + } + return; + } + const DataType arg = op->args[0].dtype(); + if (!IsSupportedVectorUnaryOp(name, arg)) { + if (!unsupported.has_value()) { + std::ostringstream oss; + oss << arg; + unsupported = name + " vectorization does not support " + + oss.str() + + " operands (per the VMI operation matrix)"; + } + return; + } + ExprVisitor::VisitExpr_(op); + return; + } + } + { + std::string name = "call"; + if (const auto *op_node = op->op.as()) { + name = op_node->name; + } + if (!unsupported.has_value()) { + unsupported = "unsupported Call `" + name + + "` (only tirx.fabs/exp/log/sqrt are vectorizable)"; + } + } + } + + void VisitExprDefault_(const ffi::Object *op) final { + std::string key = op->GetTypeKey(); + MarkUnsupported(key.c_str()); + } +}; + +} // namespace + +// --------------------------------------------------------------------------- +// Shared SIMD_VF region scan: one implementation for both passes. +// --------------------------------------------------------------------------- +namespace { + +void CollectAllParallelDeep(const Stmt &stmt, + std::vector &out) { + if (const auto *for_node = stmt.as()) { + if (for_node->kind == ForKind::kParallel) { + out.push_back(for_node); + } + CollectAllParallelDeep(for_node->body, out); + return; + } + if (const auto *seq = stmt.as()) { + for (const auto &s : seq->seq) { + CollectAllParallelDeep(s, out); + } + return; + } + if (const auto *attr = stmt.as()) { + CollectAllParallelDeep(attr->body, out); + return; + } + if (const auto *block = stmt.as()) { + CollectAllParallelDeep(block->body, out); + return; + } + if (const auto *realize = stmt.as()) { + CollectAllParallelDeep(realize->block->body, out); + return; + } + if (const auto *ite = stmt.as()) { + CollectAllParallelDeep(ite->then_case, out); + if (ite->else_case.defined()) { + CollectAllParallelDeep(ite->else_case.value(), out); + } + return; + } + if (const auto *while_node = stmt.as()) { + CollectAllParallelDeep(while_node->body, out); + return; + } +} + +void DiscoverUnitsShared(const Stmt &stmt, PtoRegionScan *scan) { + if (const auto *for_node = stmt.as()) { + if (for_node->kind == ForKind::kParallel) { + // A two-layer adjacent Parallel nest is one + // 2D unit; the layout annotation lives on the outermost For. A + // third Parallel layer (or a Parallel deeper inside) stays + // unsupported. + const auto *adjacent_inner = for_node->body.as(); + const bool inner_is_parallel = adjacent_inner != nullptr && + adjacent_inner->kind == ForKind::kParallel; + std::vector deeper; + CollectAllParallelDeep( + inner_is_parallel ? adjacent_inner->body : for_node->body, deeper); + if (!deeper.empty()) { + scan->has_unsupported_wrapper = true; + // Two distinct illegal shapes: a genuine third parallel layer + // inside an adjacent two-layer nest, or an outer Parallel whose + // body mixes other statements with an inner Parallel (the nest is + // not adjacent and cannot form a pure 2D unit). + scan->issue = + inner_is_parallel + ? "three or more nested parallel layers are not yet " + "supported" + : "a T.Parallel whose body mixes other statements with an " + "inner T.Parallel (a non-adjacent nest) is not yet " + "supported"; + } + scan->units.push_back(GetRef(for_node)); + return; + } + if (for_node->kind == ForKind::kVectorized || + for_node->kind == ForKind::kThreadBinding) { + // Only sequential wrappers may surround a unit: a re-vectorizing loop + // contradicts the vectorization this conversion already performs, and + // a thread binding contradicts the no-thread-binding contract of the + // converting region. A region without T.Parallel under them stays a + // legitimate hand-written VMI region and is skipped entirely. + std::vector under; + CollectAllParallelDeep(stmt, under); + if (!under.empty()) { + scan->has_unsupported_wrapper = true; + if (for_node->kind == ForKind::kThreadBinding) { + std::ostringstream oss; + oss << "a thread-bound loop (ForKind::kThreadBinding"; + if (const auto *iv = for_node->thread_binding.as()) { + oss << ", thread_binding `" << iv->var->name_hint << "`"; + } + oss << ") around T.Parallel is not supported: " + "thread bindings are not allowed inside a converting " + "SIMD_VF block"; + scan->issue = oss.str(); + } else { + scan->issue = + "a vectorized loop (ForKind::kVectorized) around T.Parallel is " + "not yet supported"; + } + } + return; + } + // A serial loop (with or without an unroll hint) is a sequential + // wrapper: its iterations run in order and each iteration keeps the + // unit's semantics, so it stays accepted. + DiscoverUnitsShared(for_node->body, scan); + return; + } + if (const auto *seq = stmt.as()) { + for (const auto &s : seq->seq) { + DiscoverUnitsShared(s, scan); + } + return; + } + if (const auto *attr = stmt.as()) { + DiscoverUnitsShared(attr->body, scan); + return; + } + if (const auto *block = stmt.as()) { + DiscoverUnitsShared(block->body, scan); + return; + } + if (const auto *realize = stmt.as()) { + DiscoverUnitsShared(realize->block->body, scan); + return; + } + if (stmt.as() != nullptr || stmt.as() != nullptr) { + // Control-flow wrappers: a Parallel under them is rejected; without + // one the region may still be a legitimate hand-written VMI region. + std::vector under; + CollectAllParallelDeep(stmt, under); + if (!under.empty()) { + scan->has_unsupported_wrapper = true; + scan->issue = stmt.as() != nullptr + ? "T.Parallel under a conditional branch is not yet " + "supported" + : "T.Parallel under a while loop is not yet " + "supported"; + } + return; + } +} + +} // namespace + +PtoRegionScan ScanPtoRegion(const Stmt ®ion) { + PtoRegionScan scan; + DiscoverUnitsShared(region, &scan); + std::vector all; + CollectAllParallelDeep(region, all); + scan.has_any_parallel = !all.empty(); + if (scan.has_any_parallel && scan.units.empty()) { + // A Parallel exists but only under an unsupported wrapper. + scan.has_unsupported_wrapper = true; + if (scan.issue.empty()) { + scan.issue = "T.Parallel exists only under an unsupported wrapper"; + } + } + return scan; +} + +std::string LaneUseToString(LaneUse use) { + switch (use) { + case LaneUse::kUniform: + return "uniform"; + case LaneUse::kVarying: + return "varying"; + default: + return "unknown"; + } +} + +namespace { + +// Shared per-node classification used by both the value mode and the +// metadata mode: value mode whitelists Add/Mul/Cast; metadata +// mode additionally propagates the pure scalar operators and accepts +// StringImm. +LaneUse AnalyzeLaneUseImpl(const PrimExpr &expr, LaneUseContext *ctx, + std::string *reason, bool metadata_mode) { + if (expr.as() || expr.as()) { + return LaneUse::kUniform; + } + if (expr.as() != nullptr) { + // String constants are legal metadata (markers); they never + // appear as computation values (the value whitelist rejects them). + if (metadata_mode) { + return LaneUse::kUniform; + } + if (reason != nullptr) { + // Every Unknown carries a concrete reason. + *reason = "string constants are metadata-only; they cannot appear " + "in computation values"; + } + return LaneUse::kUnknown; + } + if (const auto *var = expr.as()) { + if (var == ctx->loop_var.get()) { + return LaneUse::kVarying; + } + if (ctx->serial_var.defined() && var == ctx->serial_var.get()) { + // The coordinate of the non-vectorized + // dimension is a uniform scalar at the program point. + return LaneUse::kUniform; + } + if (ctx->bind_env.count(GetRef(var))) { + // Use the recorded classification of the bind (SSA identity). + if (ctx->classified.count(GetRef(var))) { + Integer v = ctx->classified[GetRef(var)]; + return v->value == 0 ? LaneUse::kUniform : LaneUse::kVarying; + } + // On-demand analysis with cycle guard + budget. + if (ctx->visiting.count(var) != 0) { + if (reason != nullptr) { + *reason = "bind self/forward reference detected in lane-use " + "analysis"; + } + return LaneUse::kUnknown; + } + if (ctx->budget <= 0) { + if (reason != nullptr) { + *reason = "lane-use analysis recursion budget exhausted"; + } + return LaneUse::kUnknown; + } + ctx->visiting.insert(var); + ctx->budget -= 1; + const PrimExpr &def = ctx->bind_env[GetRef(var)]; + LaneUse use = AnalyzeLaneUseImpl(def, ctx, reason, metadata_mode); + ctx->visiting.erase(var); + if (use == LaneUse::kUnknown) { + return use; + } + // Record only when the definition is fully analyzed; a cycle hit + // above leaves it unrecorded so later references re-report. + ctx->classified.Set(GetRef(var), + Integer(use == LaneUse::kUniform ? 0 : 1)); + return use; + } + // Only *known* external scalar vars are Uniform; anything + // else is undefined -> Unknown. + if (ctx->external_defs.count(var) != 0) { + // Being in scope is not enough. An external var consumed as + // a computation value (or as numeric metadata) must be scalar; a + // vector-typed var has no scalar input contract for this pass and + // must not be waved through as Uniform. Buffer data pointers are + // handle-typed (lanes == 1) and keep flowing through the address + // and Buffer paths. + if (var->dtype.lanes() != 1) { + if (reason != nullptr) { + std::ostringstream oss; + oss << "variable `" << var->name_hint << "` has vector dtype " + << var->dtype + << "; only scalar inputs are supported as computation values"; + *reason = oss.str(); + } + return LaneUse::kUnknown; + } + return LaneUse::kUniform; + } + if (reason != nullptr) { + std::ostringstream oss; + oss << "variable `" << var->name_hint + << "` is not defined in this unit or any outer scope"; + // The diagnostic was built but never written back, leaving + // metadata errors with an empty reason. + *reason = oss.str(); + } + return LaneUse::kUnknown; + } + if (const auto *load = expr.as()) { + try { + auto access = AnalyzeBufferAccess( + load->buffer, load->indices, /*is_write=*/false, ctx->loop_var, + ctx->mapping, ctx->analyzer, ctx->bind_env); + if (access.pattern == AccessPattern::kLaneUniform) { + return LaneUse::kUniform; + } + return LaneUse::kVarying; + } catch (const PtoAnalysisError &err) { + if (reason != nullptr) { + *reason = err.what(); + } + return LaneUse::kUnknown; + } + } + // Hardware-backed unary math shares the whitelist's contract (float32 + // tirx.fabs/exp/log/sqrt): recurse into the single operand. + if (const auto *call = expr.as()) { + if (auto call_op = call->op.as()) { + const std::string &name = call_op.value()->name; + if (name == "tirx.fabs" || name == "tirx.exp" || name == "tirx.log" || + name == "tirx.sqrt") { + return AnalyzeLaneUseImpl(call->args[0], ctx, reason, metadata_mode); + } + } + } + // Binary/relational/boolean operators: both value and metadata modes + // recurse into operands; the *type* checks live in the callers. + LaneUse lhs, rhs; + const PrimExpr *a = nullptr; + const PrimExpr *b = nullptr; + bool known_op = false; + if (const auto *add = expr.as()) { + a = &add->a; + b = &add->b; + known_op = true; + } else if (const auto *sub = expr.as()) { + a = &sub->a; + b = &sub->b; + known_op = true; + } else if (const auto *mul = expr.as()) { + a = &mul->a; + b = &mul->b; + known_op = true; + } else if (const auto *div = expr.as()) { + a = &div->a; + b = &div->b; + known_op = true; + } else if (const auto *mod = expr.as()) { + a = &mod->a; + b = &mod->b; + known_op = metadata_mode; + } else if (const auto *fdiv = expr.as()) { + a = &fdiv->a; + b = &fdiv->b; + known_op = metadata_mode; + } else if (const auto *fmod = expr.as()) { + a = &fmod->a; + b = &fmod->b; + known_op = metadata_mode; + } else if (const auto *minn = expr.as()) { + a = &minn->a; + b = &minn->b; + known_op = metadata_mode; + } else if (const auto *maxx = expr.as()) { + a = &maxx->a; + b = &maxx->b; + known_op = metadata_mode; + } else if (const auto *lt = expr.as()) { + a = <->a; + b = <->b; + known_op = metadata_mode; + } else if (const auto *len = expr.as()) { + a = &len->a; + b = &len->b; + known_op = metadata_mode; + } else if (const auto *gt = expr.as()) { + a = >->a; + b = >->b; + known_op = metadata_mode; + } else if (const auto *ge = expr.as()) { + a = &ge->a; + b = &ge->b; + known_op = metadata_mode; + } else if (const auto *eq = expr.as()) { + a = &eq->a; + b = &eq->b; + known_op = metadata_mode; + } else if (const auto *ne = expr.as()) { + a = &ne->a; + b = &ne->b; + known_op = metadata_mode; + } else if (const auto *andn = expr.as()) { + a = &andn->a; + b = &andn->b; + known_op = metadata_mode; + } else if (const auto *orn = expr.as()) { + a = &orn->a; + b = &orn->b; + known_op = metadata_mode; + } + if (known_op) { + lhs = AnalyzeLaneUseImpl(*a, ctx, reason, metadata_mode); + rhs = AnalyzeLaneUseImpl(*b, ctx, reason, metadata_mode); + if (lhs == LaneUse::kUnknown || rhs == LaneUse::kUnknown) { + if (reason != nullptr && reason->empty()) { + *reason = "unknown operand in binary expression"; + } + return LaneUse::kUnknown; + } + return (lhs == LaneUse::kUniform && rhs == LaneUse::kUniform) + ? LaneUse::kUniform + : LaneUse::kVarying; + } + if (!metadata_mode) { + if (const auto *notn = expr.as()) { + // Not is a value-mode whitelist reject too; keep symmetric. + } + } else { + if (const auto *notn = expr.as()) { + return AnalyzeLaneUseImpl(notn->a, ctx, reason, metadata_mode); + } + } + if (const auto *cast = expr.as()) { + return AnalyzeLaneUseImpl(cast->value, ctx, reason, metadata_mode); + } + if (reason != nullptr) { + std::ostringstream oss; + oss << "unsupported " << (metadata_mode ? "metadata" : "value") << " node `" + << expr->GetTypeKey() << "` in lane-use analysis"; + *reason = oss.str(); + } + return LaneUse::kUnknown; +} + +} // namespace + +LaneUse AnalyzeLaneUse(const PrimExpr &expr, LaneUseContext *ctx, + std::string *reason) { + return AnalyzeLaneUseImpl(expr, ctx, reason, /*metadata_mode=*/false); +} + +LaneUse AnalyzeLaneUseMeta(const PrimExpr &expr, LaneUseContext *ctx, + std::string *reason) { + return AnalyzeLaneUseImpl(expr, ctx, reason, /*metadata_mode=*/true); +} + +bool IsSupportedVectorUnaryOp(const std::string &tirx_op, DataType dtype) { + // Per the VMI ISA documentation: vabs admits f16/bf16/f32 (the vabs + // instruction also takes si8/si16/si32, but the frontend abs() lowers + // integers to Select — PR-2 territory — and returns unsigned inputs + // unchanged, so integer dtypes never reach this gate); vexp/vln/vsqrt + // admit f16/f32 only. + const int bits = dtype.bits(); + if (tirx_op == "tirx.fabs") { + return (dtype.is_float() || dtype.is_bfloat16()) && + (bits == 16 || bits == 32); + } + if (tirx_op == "tirx.exp" || tirx_op == "tirx.log" || + tirx_op == "tirx.sqrt") { + return dtype.is_float() && (bits == 16 || bits == 32); + } + return false; +} + +bool ExprUsesVar(const PrimExpr &expr, const VarNode *var) { + bool uses = false; + PostOrderVisit(expr, [&](const ObjectRef &n) { + if (n.as() == var) { + uses = true; + } + }); + return uses; +} + +bool ContainsRuntimeValue(const PrimExpr &expr) { + // True when the expression reads any runtime value: a Var (external + // scalar, serial coordinate, scalar buffer base) or a BufferLoad. + // A pure IntImm/FloatImm tree is a compile-time constant. + bool has_runtime = false; + PostOrderVisit(expr, [&](const ObjectRef &n) { + if (n.as() != nullptr || n.as() != nullptr) { + has_runtime = true; + } + }); + return has_runtime; +} + +bool IsFloatFamily(DataType dtype) { + return dtype.is_float() || dtype.is_bfloat16(); +} + +bool IsSupportedVectorBinaryOp(const std::string &vmi_op, DataType dtype) { + // Operation-level matrix (per the VMI ISA documentation): + // vmul requires i16/i32/f16/bf16/f32; vadd/vsub additionally allow i8; + // vdiv admits f16/f32 only (no bf16, no integers). + int bits = dtype.bits(); + const bool is_int_like = dtype.is_int() || dtype.is_uint(); + const bool is_fp_like = + (dtype.is_float() || dtype.is_bfloat16()) && (bits == 16 || bits == 32); + if (vmi_op == "vmul") { + return (is_int_like && (bits == 16 || bits == 32)) || is_fp_like; + } + if (vmi_op == "vadd" || vmi_op == "vsub") { + return (is_int_like && (bits == 8 || bits == 16 || bits == 32)) || + is_fp_like; + } + if (vmi_op == "vdiv") { + return dtype.is_float() && (bits == 16 || bits == 32); + } + return false; +} + +namespace { + +/*! Cast attributes read from a Cast node's backend annotations. The + * frontend writes them in two shapes — TIR Imm objects (the `T.cast` helper + * lowers `round`/`sat` to StringImm/IntImm) and plain Python values (a raw + * `annotations={...}` map) — so both spellings are accepted here. */ +struct CastAttributes { + /*! The `round` annotation was present (even when its value is empty). */ + bool has_round_key = false; + /*! A non-empty rounding mode was requested. An empty string is the + * frontend's way of saying "use the backend default", which is not the same + * as an unusable mode. */ + bool has_round_mode = false; + std::string round; // raw spelling, e.g. "rn", "R", "A" + bool has_sat = false; + bool sat = true; + bool has_rbits = false; +}; + +/*! Read the attribute hints of a Cast. Returns false with *why set when an + * annotation uses a shape this version cannot read; an unknown key is + * itself a rejection, because silently ignoring an attribute would accept + * a conversion whose semantics were never checked. */ +bool ReadCastAttributes(const ffi::Map &annotations, + CastAttributes *attrs, std::string *why) { + for (const auto &[key_ref, value] : annotations) { + const std::string key = key_ref; + if (key == "round") { + attrs->has_round_key = true; + if (const auto *imm = value.as()) { + attrs->round = imm->value; + } else if (auto str = value.as()) { + attrs->round = std::string(*str); + } else { + *why = "annotation `round` must be a string"; + return false; + } + attrs->has_round_mode = !attrs->round.empty(); + } else if (key == "sat") { + attrs->has_sat = true; + if (const auto *imm = value.as()) { + // Only 0/1 are accepted, in the TIR Imm form as well. + if (imm->value != 0 && imm->value != 1) { + *why = "annotation `sat` must be a boolean (got " + + std::to_string(imm->value) + ")"; + return false; + } + attrs->sat = imm->value != 0; + } else if (auto flag = value.as()) { + attrs->sat = *flag; + } else if (auto flag = value.as()) { + // Only 0/1 are accepted: any other integer is a malformed flag, not a + // saturating request. + if (*flag != 0 && *flag != 1) { + *why = "annotation `sat` must be a boolean (got " + + std::to_string(*flag) + ")"; + return false; + } + attrs->sat = *flag != 0; + } else { + *why = "annotation `sat` must be a boolean"; + return false; + } + } else if (key == "rbits") { + attrs->has_rbits = true; + } else { + *why = "unknown cast annotation `" + key + "`"; + return false; + } + } + return true; +} + +/*! Normalize a rounding spelling to the canonical PTO token. The frontend + * aliases are PTX-style; the canonical tokens are the ones the FP8 + * conversion policy reasons about. Returns false with *why set for + * aliases with no PTO meaning (stochastic rounding) and for tokens no + * conversion in this version accepts. */ +bool NormalizeCastRounding(const std::string &raw, std::string *token, + std::string *why) { + if (raw == "rn") { + *token = "R"; + return true; + } + if (raw == "rz") { + *token = "Z"; + return true; + } + if (raw == "rp") { + *token = "C"; + return true; + } + if (raw == "rm") { + *token = "F"; + return true; + } + if (raw == "rs") { + *why = "stochastic rounding (round=\"rs\") is not supported"; + return false; + } + static const std::set kCanonical = {"R", "A", "H", + "Z", "C", "F"}; + if (kCanonical.count(raw) != 0) { + *token = raw; + return true; + } + *why = "rounding mode \"" + raw + "\" has no PTO vcvt meaning"; + return false; +} + +/*! Rounding tokens the FP8 narrowing policy accepts. The conversion + * instructions the FP8 target exposes are FP32 <-> FP8; the matrix is the + * intersection of the PTODSL vcvt interface with that support, so the source + * is always float32 and only the rounding modes below are expressible. */ +const std::set &Fp8AllowedRoundings() { + static const std::set kAllowed = {"R", "A", "H", "Z"}; + return kAllowed; +} + +} // namespace + +Optional PlanValueCast(const CastNode *cast, std::string *reason) { + DataType from = cast->value.dtype(); + DataType to = cast->dtype; + auto reject = [&](const std::string &why) { + if (reason != nullptr) { + *reason = why; + } + return Optional(); + }; + std::ostringstream pair_os; + pair_os << "`" << from << " -> " << to << "`"; + const std::string pair = pair_os.str(); + // Dtype legality is a precondition for every action plan, + // including identity — unsupported types (int24/FP64/...) must not slip + // through the same-type fast path. + if (!IsSupportedElementDType(from) || !IsSupportedElementDType(to)) { + return reject(pair + " is outside the supported element dtype set"); + } + + // Attributes are parsed for every Cast, not just the FP8 ones: the + // expression whitelist deliberately does not look inside a cast's + // annotations (attributes are not computation values), so an annotation the + // selected conversion cannot honour has to be rejected here — for every + // direction — or it would be dropped silently. + CastAttributes attrs; + std::string why; + if (!ReadCastAttributes(cast->annotations, &attrs, &why)) { + return reject(pair + ": " + why); + } + if (attrs.has_rbits) { + return reject(pair + ": the rbits operand (stochastic rounding) is not " + "supported"); + } + std::string round_token; + if (attrs.has_round_mode && + !NormalizeCastRounding(attrs.round, &round_token, &why)) { + return reject(pair + ": " + why); + } + + if (from == to) { + // Identity, FP8 or not: no conversion is emitted, so an explicit sat=True + // is the default and is normalized away; any other explicit attribute + // would claim a conversion semantic that does not exist here. + if (attrs.has_sat && !attrs.sat) { + return reject(pair + ": an identity cast does not accept sat=False"); + } + if (attrs.has_round_mode) { + return reject(pair + + ": an identity cast does not accept a rounding mode"); + } + CastPlan identity; + identity.identity = true; + return identity; + } + + const bool from_fp8 = IsSupportedFloat8(from); + const bool to_fp8 = IsSupportedFloat8(to); + if (from_fp8 || to_fp8) { + if (from_fp8 && to_fp8) { + return reject(pair + ": conversions between the two FP8 types are not " + "supported (they are not a reinterpretation)"); + } + if (from_fp8 && to.is_float() && to.bits() == 32) { + // Widening: exact, no rounding or saturation involved. An explicit + // hint here cannot be honoured, so it is rejected rather than dropped. + // The direction contract rejects an attribute that is *carried* here, + // even an empty one: widening has no rounding to configure. + if (attrs.has_round_key) { + return reject(pair + + ": this direction does not accept a rounding mode"); + } + if (attrs.has_sat) { + return reject(pair + + ": this direction does not accept a saturation mode"); + } + return CastPlan{}; + } + if (to_fp8) { + // The FP8 conversion instructions the target exposes are FP32 <-> FP8; + // there is no half -> FP8 conversion, so a float16/bfloat16 source is + // rejected here. This pass never adds a conversion node of its own: a + // user who wants the float32 route writes it explicitly, so the + // intermediate rounding is their expressed choice rather than a silent + // step inserted here. + if (!(from.is_float() && from.bits() == 32)) { + return reject(pair + ": the FP8 conversion instructions take a float32 " + "source; convert to float32 explicitly first"); + } + const std::set &allowed = Fp8AllowedRoundings(); + // An empty round selects the documented default (R with SAT below). + const std::string token = attrs.has_round_mode ? round_token : "R"; + if (allowed.count(token) == 0) { + std::ostringstream oss; + oss << "rounding mode " << token; + if (attrs.has_round_mode && attrs.round != token) { + oss << " (from round=\"" << attrs.round << "\")"; + } + oss << " is not supported by the FP8 conversion policy for `" << from + << " -> " << to << "`"; + return reject(oss.str()); + } + CastPlan plan; + // The PTO default is R/SAT; pin both explicitly so the generated VMI + // call never depends on PTODSL defaults. + plan.rounding = token; + plan.saturate = (attrs.has_sat && !attrs.sat) ? "NOSAT" : "SAT"; + return plan; + } + return reject(pair + + " is not a supported FP8 conversion (FP8 -> FP16/BF16 and " + "FP8 <-> integer are out of scope; convert through " + "float32 explicitly)"); + } + + // Non-FP8 conversion: the plan is fixed by the TIR semantics of the + // operation, so an explicit annotation must agree with it. `finish` is the + // one place that decides this, so a hint is accepted exactly when the plan + // already does what it asks for and rejected otherwise — never dropped. + auto finish = [&](const CastPlan &plan) -> Optional { + if (attrs.has_round_mode) { + return reject(pair + ": this conversion does not accept a rounding mode"); + } + if (attrs.has_sat && (attrs.sat || plan.saturate != "NOSAT")) { + return reject(pair + ": this conversion does not accept an explicit " + "saturation flag"); + } + return plan; + }; + + // Same-width signed/unsigned reinterprets have no validated lowering + // (VMI vcvt requires a bit-width change); reject. + if (from.bits() == to.bits() && from.code() == DataType::kInt && + to.code() == DataType::kUInt) { + return reject("same-width int -> uint is a reinterpretation, which VMI " + "vcvt does not provide"); + } + if (from.bits() == to.bits() && from.code() == DataType::kUInt && + to.code() == DataType::kInt) { + return reject("same-width uint -> int is a reinterpretation, which VMI " + "vcvt does not provide"); + } + // No width re-check here: IsSupportedElementDType at function entry already + // restricts both sides to 8/16/32-bit lane-1 elements (and excludes FP64). + CastPlan plan; + // Every integer combination shares one rule: TVM integer casts truncate on + // narrowing and extend on widening, PTODSL defaults narrowing to saturate, + // and NOSAT restores the truncating semantics. The same-width int/uint + // reinterpretations are rejected above. + if ((from.is_int() || from.is_uint()) && (to.is_int() || to.is_uint())) { + if (to.bits() < from.bits()) { + plan.saturate = "NOSAT"; + } + return finish(plan); + } + if (from.is_float() && (to.is_float() || to.is_bfloat16())) { + if (to.bits() > from.bits() && to.is_float()) { + // f16 -> f32 widening is exact. + return finish(plan); + } + // f32 -> f16/bf16 narrowing: TIR Cast rounds to nearest; pin R so the + // VMI vcvt never depends on PTODSL defaults. NOSAT keeps IEEE + // overflow-to-inf semantics instead of clamping to the largest finite. + if (from.bits() == 32 && to.bits() == 16) { + plan.rounding = "R"; + plan.saturate = "NOSAT"; + return finish(plan); + } + // The only float narrowing left here is the cross-family f16 <-> bf16 + // step; its rounding semantics are not validated. + return reject( + "cross-family f16 <-> bf16 narrowing has no validated rounding " + "semantics; convert through float32 explicitly"); + } + // Bfloat16's is_float() is false, so bf16 -> f32 fell into the + // reject branch below even though it is the same exact widening as + // f16 -> f32 (and the PTO backend supports the vcvt). Allow exactly + // this one bfloat combination; every other stays rejected as before. + if (from.is_bfloat16() && to.is_float() && to.bits() == 32) { + return finish(plan); + } + if (from.is_float() != to.is_float()) { + // float <-> int conversions need per-combination rounding/overflow + // validation against PTOAS; currently rejected. + return reject("float <-> integer conversions need per-combination " + "rounding/overflow validation"); + } + return reject("no proven VMI-equivalent conversion for this dtype pair"); +} + +std::optional FindUnsupportedExprNode(const PrimExpr &expr) { + WhitelistVisitor visitor; + visitor(expr); + return visitor.unsupported; +} + +} // namespace pto +} // namespace tl +} // namespace tvm diff --git a/src/ascend/transform/parallel_to_pto_utils.h b/src/ascend/transform/parallel_to_pto_utils.h new file mode 100644 index 0000000000..453c3df7b2 --- /dev/null +++ b/src/ascend/transform/parallel_to_pto_utils.h @@ -0,0 +1,334 @@ +/*! + * \file parallel_to_pto_utils.h + * \brief Shared read-only analysis helpers for the PTO Parallel vectorize + * path. See docs/vectorize_parallel_to_pto_design.md. + * + * VerifyParallelToPTO consumes these helpers to check the input contract; + * VectorizeParallelToPTO re-runs the same analysis to build its VMI code. + * Keeping one implementation guarantees the two passes classify the same + * access identically. + */ + +#ifndef TVM_TL_ASCEND_TRANSFORM_PARALLEL_TO_PTO_UTILS_H_ +#define TVM_TL_ASCEND_TRANSFORM_PARALLEL_TO_PTO_UTILS_H_ + +#include "support/check.h" +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "../../layout/layout.h" +#include "../../op/builtin.h" +#include "../../op/pto_index_analysis.h" +#include "../target_utils.h" + +namespace tvm { +namespace tl { + +using namespace tirx; +using ffi::Array; +using ffi::GetRef; +using ffi::Map; +using ffi::ObjectRef; +using ffi::Optional; + +namespace pto { + +/*! + * \brief Classification of one buffer access against the lane mapping. + * + * The classification describes how the + * *address* varies over lanes/chunk, not the loaded data. + */ +enum class AccessPattern { + kUnknown = 0, + /*! Address = start + lane for lane in [0, L): continuous vector. */ + kContinuous, + /*! Address independent of lane at the current program point. */ + kLaneUniform, +}; + +std::string AccessPatternToString(AccessPattern pattern); + +/*! + * \brief Result of the shared single-access analysis. + * + * Verify consumes pattern legality; Vectorize additionally consumes the + * start indices/offset to build tl.access_ptr. + */ +struct MemoryAccess { + Buffer buffer; + DataType elem_dtype; + bool is_write = false; + + /*! Full element offset relative to buffer->data, incl. elem_offset, + * version dims, outer coordinates and chunk offset (lane term included). + */ + PrimExpr data_offset; + /*! indices_after_inverse: buffer indices with the inverse-mapped logical + * coordinates (lane appears via the lane Var in one index). */ + Array indices_after_inverse; + /*! indices with lane bound to 0: the start element of this vector access. + */ + Array start_indices; + /*! data_offset with lane bound to 0. */ + PrimExpr start_offset; + AccessPattern pattern = AccessPattern::kUnknown; +}; + +// PtoAnalysisError, SubstituteVar, ResolveIndexAliases and +// IsSupportedAddressExpr live in src/op/pto_index_analysis.h (included +// above) so the layout-planning code in src/op and the consumers here +// share one implementation. + +/*! + * \brief The inverse lane mapping of a 1D PTO parallel loop. + * + * For a 1D loop of extent E with lanes L (E % L == 0): the fragment is + * forward_thread = i % L, forward_index = [i // L]. Its inverse recovers + * i = p * L + lane + * where p is the per-lane chunk position. + */ +struct PtoInverseMapping { + /*! Lane variable of the enclosing SIMD_VF (analysis placeholder). */ + Var lane_var; + /*! Chunk position variable (per-lane column index). */ + Var chunk_var; + int64_t lanes = 0; + /*! Logical extent E of the vectorized dimension (the original loop + * extent, without padding). */ + int64_t logical_extent = 0; + /*! Padded extent P = ceil(E / L) * L of the fragment input space on the + * vectorized dimension. E == P for divisible extents. */ + int64_t padded_extent = 0; + /*! Chunk count Q = P / L — the per-row chunk range of the vectorized + * dimension (never the global position K*Q). */ + int64_t chunk_count = 0; + /*! Tail lanes E % L (0 when E is divisible by L). */ + int64_t tail_lanes = 0; + /*! Inverse expression of the vectorized coordinate: + * logical index = chunk * L + lane. */ + PrimExpr index_expr; + + // ------------------------------------------------------------------ + // 2D units. A 2D unit is a two-layer adjacent + // Parallel nest; the layout annotation sits on the outermost For and + // the Fragment input coordinates stay (i, j) regardless of which + // dimension is vectorized. + // ------------------------------------------------------------------ + /*! True when the unit is a 2D nest. */ + bool is_2d = false; + /*! Original loop vars of the nest (i = outer, j = inner). */ + Var outer_var; + Var inner_var; + /*! Extents of the two iterations (i and j). */ + int64_t outer_extent = 0; + int64_t inner_extent = 0; + /*! Vectorized dimension: true selects the inner (j) dim. */ + bool select_inner = true; + + /*! The loop variable substituted into addresses (the vectorized one). */ + const VarNode *vector_var() const { + if (!is_2d) { + return nullptr; // 1D: the caller passes the single loop var + } + return select_inner ? inner_var.get() : outer_var.get(); + } + /*! The loop variable kept as a serial loop (2D only). */ + const VarNode *serial_var() const { + if (!is_2d) { + return nullptr; + } + return select_inner ? outer_var.get() : inner_var.get(); + } +}; + +/*! Invert the pto_parallel_loop_layout fragment of a Parallel unit. + * Accepts a 1D Parallel For, or the outermost For of a two-layer adjacent + * Parallel nest (2D). Throws PtoAnalysisError when the fragment is + * missing, has an unsupported shape, or cannot be inverted. */ +PtoInverseMapping InvertPtoLaneLayout(const For &loop); + +/*! Classify one buffer access (load or store indices) against the inverse + * mapping. \p indices are the buffer indices in the original loop body + * expressed over the loop variable; \p loop_var is the original parallel + * loop variable. \p bind_env maps Bind vars in scope to their defining + * expressions so index aliases are resolved before classification + * pass an empty map when no binds are in scope. + * + * The helper resolves bind aliases, substitutes + * loop_var -> chunk * L + lane_var, computes ElemOffset, and proves the + * access pattern with the analyzer over the full lane range (continuous: + * data_offset == start_offset + lane as an identity; lane-uniform: the + * offset does not depend on lane). Returns a fully populated MemoryAccess + * or throws PtoAnalysisError with a reason. */ +MemoryAccess AnalyzeBufferAccess(const Buffer &buffer, + const Array &indices, bool is_write, + const Var &loop_var, + const PtoInverseMapping &mapping, + arith::Analyzer *analyzer, + const Map &bind_env = {}); + +/*! Check whether a dtype is supported by the unified element + * addressing convention: the buffer element dtype itself must not carry + * vector lanes and must not be a packed sub-byte type (FP4 etc.). + * + * E4M3FN and E5M2 are supported element types; the other FP8 variants + * (float8_e4m3 with its own TIR dtype, E8M0, FNUZ, HiF8) are not, even + * where the PTO codegen would map them to the same target type — only the + * two dtypes listed in IsSupportedFloat8 take part in the FP8 conversion + * matrix. */ +bool IsSupportedElementDType(DataType dtype); + +/*! The two FP8 element types the conversion matrix supports: float8_e4m3fn and + * float8_e5m2. Deliberately narrower than the codegen's FP8 name mapping: + * admitting a type here means its conversion directions were validated, not + * just that a target name exists. */ +bool IsSupportedFloat8(DataType dtype); + +/*! Cast conversion descriptor. + * empty() means the combination has no proven-equivalent VMI mapping and + * must be rejected by Verify. */ +struct CastPlan { + /*! Identity conversion (from == to): keep the original value, emit no + * vcvt (VMI rejects width-preserving conversions). */ + bool identity = false; + /*! PTODSL saturate annotation value, e.g. "NOSAT"; empty when the + * conversion must not rely on PTODSL defaults. */ + std::string saturate; + /*! PTODSL rounding annotation value, e.g. "Z"; empty when N/A. */ + std::string rounding; +}; + +/*! Operation-level type support for the vector arithmetic emitted + * here. VMI vmul only supports i16/i32/f16/bf16/f32 + * elements; vadd/vsub additionally allow i8; vdiv admits f16/f32 only. + * Unsupported combinations are rejected instead of synthesizing + * widening chains. */ +bool IsSupportedVectorUnaryOp(const std::string &tirx_op, DataType dtype); +bool IsSupportedVectorBinaryOp(const std::string &vmi_op, DataType dtype); +bool ContainsRuntimeValue(const PrimExpr &expr); +bool IsFloatFamily(DataType dtype); + +/*! Read-only lane-use classification of a value expression. + * Shared by the metadata check and the operation type check so both + * consume the same source semantics. Mirrors the structural conversion + * the Vectorize pass performs (no new algebraic simplification). */ +enum class LaneUse { + kUniform, ///< stays scalar under the current conversion rules + kVarying, ///< needs a vector value (incl. the loop index as a value) + kUnknown, ///< cannot analyze; caller must reject with the reason +}; + +std::string LaneUseToString(LaneUse use); + +/*! Analysis context for AnalyzeLaneUse: the unit's loop var, inverse + * mapping and analyzer, the bind-definition environment, a cache of + * already-classified bind vars, the set of known external scalar vars + * (params / buffer data pointers / outer-scope definitions), and the + * recursion guards of the classification walk. */ +struct LaneUseContext { + Var loop_var; + /*! In a 2D unit the coordinate of the non-vectorized dimension + * is a uniform scalar at the program point. Unset for 1D. */ + Var serial_var; + PtoInverseMapping mapping; + arith::Analyzer *analyzer = nullptr; + Map bind_env; + Map classified; // bind var -> 0 uniform / 1 varying + /*! Known external scalar vars: vars outside the unit that are + * legitimately defined. A var that is neither in bind_env nor here is + * *undefined* -> Unknown, not Uniform. */ + std::set external_defs; + /*! Bind vars currently being analyzed (cycle guard). */ + std::set visiting; + /*! Recursion budget for on-demand bind analysis. */ + int budget = 256; +}; + +/*! Classify \p expr bottom-up. Constants and *known* + * external scalar vars are Uniform; the loop var as a value is Varying; + * BufferLoad goes through AnalyzeBufferAccess (uniform load -> Uniform, + * continuous -> Varying); Add/Mul/Cast propagate operands; original Bind + * vars use their recorded classification (on-demand analysis is guarded + * by the visiting set and budget; cycles and exhaustion -> Unknown). + * \p reason is filled on kUnknown. */ +LaneUse AnalyzeLaneUse(const PrimExpr &expr, LaneUseContext *ctx, + std::string *reason); + +/*! Metadata-mode lane classification: same Var/Bind/BufferLoad + * rules as AnalyzeLaneUse, but the *node table* is the metadata table — + * pure scalar Sub/Div/Mod/FloorDiv/FloorMod/Min/Max/comparison/boolean + * operators and StringImm constants propagate their operands' lane use + * instead of being rejected by the value whitelist. Unknown nodes still + * produce kUnknown. */ +LaneUse AnalyzeLaneUseMeta(const PrimExpr &expr, LaneUseContext *ctx, + std::string *reason); + +/*! Decide the conversion attributes for a value Cast. This is the single + * admission point for every Cast on the vectorize path; Verify and Vectorize + * both consume its result, so they cannot diverge. + * + * Same-width reinterpretation-free identity is impossible in + * TIR Cast, so the proven set is: + * - integer narrowing/widening with truncating semantics: TVM integer + * casts truncate on narrowing (codegen_llvm CreateIntCast), PTODSL + * needs saturate="NOSAT" for narrowing to keep that semantics; + * - float widening (f16->f32, bf16->f32): exact; + * - the FP8 matrix: E4M3FN/E5M2 <-> f32. The conversion instructions the + * target exposes are FP32 <-> FP8, so a narrowing conversion carries an + * explicit rounding and saturation attribute (the FP8 conversion policy in + * the implementation file) and a half (f16/bf16) source is rejected rather + * than routed through float32 internally. FP8 -> FP16/BF16, FP8 <-> integer, + * E4M3FN <-> E5M2 and every other FP8 variant stay out of scope. + * Everything else (float narrowing, float<->int, unusual bit widths) + * returns nullopt -> reject. + * + * `cast` carries the backend annotations written by the frontend cast + * helpers (rounding mode, saturation flag, stochastic rbits operand). They + * are read here — never interpreted by the expression visitors — and an + * attribute the selected conversion cannot honour is a rejection, not a + * silently dropped hint. When `reason` is non-null it receives a diagnostic + * naming the source type, target type and the offending option. */ +Optional PlanValueCast(const CastNode *cast, + std::string *reason = nullptr); + +/*! Check whether an expression only uses nodes from the supported + * whitelist: BufferLoad, constants, scalar Vars, loop + * indices, Add, Mul, limited Cast, Bind values. Returns nullopt when + * supported; otherwise a short description of the offending node. */ +std::optional FindUnsupportedExprNode(const PrimExpr &expr); + +/*! Result of scanning a SIMD_VF region for parallel units. Shared by + * Verify and Vectorize so both passes agree on skip/reject/convert + * decisions. */ +struct PtoRegionScan { + /*! Sequential 1D parallel units found (outermost For of each chain). */ + std::vector units; + /*! True when a Parallel exists under a wrapper we cannot currently + * convert (IfThenElse / While / nested 2D+ parallel). Verify + * rejects with \p issue; Vectorize treats it as an internal error. */ + bool has_unsupported_wrapper = false; + std::string issue; + /*! True when the region contains any T.Parallel at all (directly or + * under wrappers). A region with none is a pure hand-written VMI + * region: both passes skip it entirely. */ + bool has_any_parallel = false; +}; + +/*! Scan a SIMD_VF region body (the SBlock's body, including any scope + * AttrStmt wrapper) for parallel units. One implementation shared by + * Verify and Vectorize. */ +PtoRegionScan ScanPtoRegion(const Stmt ®ion); + +} // namespace pto +} // namespace tl +} // namespace tvm + +#endif // TVM_TL_ASCEND_TRANSFORM_PARALLEL_TO_PTO_UTILS_H_ diff --git a/src/ascend/transform/vectorize_parallel_to_pto.cc b/src/ascend/transform/vectorize_parallel_to_pto.cc new file mode 100644 index 0000000000..826010819c --- /dev/null +++ b/src/ascend/transform/vectorize_parallel_to_pto.cc @@ -0,0 +1,1015 @@ +// Copyright (c) Tile-AI Corporation. +// Licensed under the MIT License. + +/*! + * \file vectorize_parallel_to_pto.cc + * \brief VectorizeParallelToPTO: rewrite verified T.Parallel loops inside + * PTO SIMD_VF blocks into logical VMI code. + * + * Consumes the pto_parallel_loop_layout fragment written by + * LayoutInference and verified by VerifyParallelToPTO. Per parallel unit: + * 1. invert the lane fragment through the shared helper (the inverse + * layout's own Forward result is the single source of truth); + * 2. build the chunk loop for q in [0, E/L), mapping the loop variable + * to q * L + lane; + * 3. classify every access via the shared MemoryAccess analysis + * (continuous -> vload/vstore via tl.access_ptr(start, L, rw); + * lane-uniform -> scalar BufferLoad with lane bound to 0); + * 4. convert value expressions (value roles tracked by result dtype); + * every VMI with a return value is bound through EmitBind before use; + * 5. delete the consumed layout annotation, keep the SIMD_VF block + * annotations, and leave no unconverted Parallel behind. + * + * Supported shapes: 1D parallel units and two-layer adjacent Parallel nests + * (2D — the nest converts as one unit with the layout + * annotation on its outermost For); sequential units allowed; static + * divisible extents; full masks. A third parallel layer or any other + * unsupported structure is rejected by Verify. + * + * Non-divisible extents emit a tail chunk whose vload reads the full L + * lanes from its chunk base: the ISA intentionally keeps load instructions + * unmasked for performance, so the per-chunk tail mask constrains + * computation and stores only, never the load. Reading the full L lanes + * past the logical tail is still safe on device — the top 8K of UB is + * reserved as stack space, so the extra read-only lanes cannot fault; they + * carry undefined values that the compute and store masks keep out of the + * valid outputs. Compact outputs of exactly E elements remain safe because + * the masked store never writes past E. + */ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "parallel_to_pto_utils.h" + +namespace tvm { +namespace tl { + +using namespace tirx; +using namespace ffi; +namespace pto_analysis = ::tvm::tl::pto; + +namespace { + +DataType VectorDType(DataType elem, int64_t lanes) { + return elem.with_lanes(static_cast(lanes)); +} + +const Op &VmiOp(const std::string &tail) { return Op::Get("tl.vmi." + tail); } + +// This version keeps local DCE in Vectorize so +// its directly emitted TIR/PTODSL is cleaner and easier to inspect. +// Cleanup may later move to PTOAS canonicalization/DCE to reduce +// TileLang-specific code and maintenance; clean-output checks would then +// target the optimized PTOAS IR. The emitted IR must be valid and preserve +// program semantics even before DCE. +// +// Dead vector-chain elimination on one converted chunk body. The +// conversion can leave Bind statements whose Var is never referenced: +// source binds consumed only through the address alias environment +// (bind_env_) never use their converted Var, and an unused source bind +// strands its whole value chain (vci/vbrc/vadd/vload). +// +// The traversal is a scope-aware *reverse liveness* pass — a +// dead Bind must not survive just because it was wrapped in an AttrStmt +// or appears as the only statement of a body. The rules are: +// 1. A SeqStmt is processed back to front; `live` always holds the +// vars everything *after* the current statement needs. +// 2. A Bind's LHS var is a definition, never a use. When the result is +// not live and the RHS is discardable, the Bind is deleted and its +// RHS inputs do NOT become live (so the dead chain collapses in one +// pass); otherwise it is kept and its RHS inputs join `live`. +// 3. An AttrStmt body is cleaned in its own nested scope with a fresh +// liveness set; only the vars actually referenced inside propagate +// outward. Attribute node/value references are always uses, and the +// attribute structure and evaluation position are preserved. +// 4. A bare Bind follows the same rules; when everything dies the +// remains are a legal no-op (Evaluate(0)). +// Nothing is merged, reordered or otherwise rescheduled; +// referenced loads stay exactly where they were (vecadd keeps its three +// vloads), and Bind-free (frontend-inlined) bodies are untouched. +class DeadBindElim { +public: + Stmt operator()(const Stmt &stmt) { + std::set live; + Stmt out = Run(stmt, &live); + return out.defined() ? out : NoOp(); + } + +private: + static Stmt NoOp() { return Evaluate(make_zero(DataType::Int(32))); } + + static void CollectExprVars(const PrimExpr &expr, + std::set *live) { + PostOrderVisit(expr, [&](const ObjectRef &n) { + if (const auto *v = n.as()) { + live->insert(v); + } + }); + } + + static void CollectStmtVars(const Stmt &stmt, + std::set *live) { + PostOrderVisit(stmt, [&](const ObjectRef &n) { + if (const auto *v = n.as()) { + live->insert(v); + } + }); + } + + static bool DiscardableRhs(const PrimExpr &rhs) { + // Positive "allowed to discard" set: a Call may disappear + // only when it is one of the tl.vmi.* ops this pass actually emits as + // a Bind RHS and that is pure/read-only by construction — vload + // (read), vbrc, vci, vadd, vmul and vcvt (pure), enumerated from the + // ConvertUnit/VmiOp call sites — or the pure tl.access_ptr address + // descriptor nested inside a vload operand (already whitelisted as a + // pure index intrinsic by IsSupportedAddressExpr). vstore never + // appears as a Bind RHS; any other Call — call_extern, other tl.* or + // an unknown op — keeps its Bind. Plain UB BufferLoads are + // discardable: they are ordinary reads with no write effect (the + // scalar "sload" binds). That "ordinary" premise is enforced by + // VerifyParallelToPTO, not assumed here: accesses to + // tirx.volatile-annotated allocations and volatile_scope attributes + // are rejected there, and opaque calls are already excluded + // from addresses. + static const std::set kDiscardable = { + "tl.vmi.vload", "tl.vmi.vbrc", "tl.vmi.vci", "tl.vmi.vadd", + "tl.vmi.vsub", "tl.vmi.vmul", "tl.vmi.vdiv", "tl.vmi.vabs", + "tl.vmi.vexp", "tl.vmi.vln", "tl.vmi.vsqrt", "tl.vmi.vcvt", + "tl.access_ptr", + }; + bool discardable = true; + PostOrderVisit(rhs, [&](const ObjectRef &n) { + if (const auto *call = n.as()) { + const auto *op = call->op.as(); + const std::string name = op != nullptr ? op->name : ""; + if (kDiscardable.count(name) == 0) { + discardable = false; + } + } + }); + return discardable; + } + + // Returns the kept statement, or an *undefined* Stmt when the statement + // was deleted (only a Bind can be deleted). `live` is threaded from + // later statements back to earlier ones. + Stmt Run(const Stmt &stmt, std::set *live) { + if (const auto *bind = stmt.as()) { + if (live->count(bind->var.get()) == 0 && DiscardableRhs(bind->value)) { + // Dead and safe to drop: the RHS inputs stay out of `live`, so a + // chain feeding only this Bind collapses in the same pass. + return Stmt(); + } + // Kept: classic liveness — subtract the definition, add the uses. + live->erase(bind->var.get()); + CollectExprVars(bind->value, live); + return stmt; + } + if (const auto *seq = stmt.as()) { + Array kept; + for (auto it = seq->seq.rbegin(); it != seq->seq.rend(); ++it) { + Stmt sub = Run(*it, live); + if (sub.defined()) { + kept.push_back(sub); + } + } + if (kept.empty()) { + return NoOp(); + } + // std::reverse does not accept the ffi Array iterators on this + // toolchain; re-append in reverse instead. + Array forward; + for (auto it = kept.rbegin(); it != kept.rend(); ++it) { + forward.push_back(*it); + } + return SeqStmt::Flatten(forward); + } + if (const auto *attr = stmt.as()) { + // Nested scope: clean the body with a fresh liveness set, then + // propagate only the references that survive to its entry (external + // uses; vars defined and used only inside never escape because the + // converted Vars are unique objects). Attribute node/value + // references are always uses and keep the structure. + std::set body_live; + Stmt body = Run(attr->body, &body_live); + if (!body.defined()) { + body = NoOp(); + } + live->insert(body_live.begin(), body_live.end()); + if (auto ne = attr->node.try_cast()) { + if (const auto *pe = ne.value().as()) { + CollectExprVars(GetRef(pe), live); + } + } + CollectExprVars(attr->value, live); + if (body.same_as(attr->body)) { + return stmt; + } + return AttrStmt(attr->node, attr->attr_key, attr->value, body); + } + // Any other statement is always kept; all its Var occurrences are + // uses. + CollectStmtVars(stmt, live); + return stmt; + } +}; + +class VectorizeParallelToPTOImpl : public StmtExprMutator { +public: + static PrimFunc Substitute(PrimFunc f) { + VectorizeParallelToPTOImpl mutator; + PrimFuncNode *fptr = f.CopyOnWrite(); + fptr->body = mutator.VisitStmt(f->body); + return GetRef(fptr); + } + +private: + // ------------------------------------------------------------------ + // Unified unit discovery (shared rules with Verify). + // A SIMD_VF region contains zero or more *sequential* parallel units. + // Each unit is the outermost parallel For of a 1D loop or of a + // two-layer adjacent Parallel nest (2D); a third parallel layer (or a + // Parallel under an unsupported wrapper) is rejected by Verify and + // defensively rejected here too. + // ------------------------------------------------------------------ + // ------------------------------------------------------------------ + // SIMD_VF block entry + // ------------------------------------------------------------------ + Stmt VisitStmt_(const SBlockNode *op) final { + if (op->name_hint != "SIMD_VF") { + return StmtExprMutator::VisitStmt_(op); + } + // Shared region scan: the same implementation Verify uses, so + // skip/reject/convert decisions cannot drift apart. + const Stmt ®ion = op->body; + pto_analysis::PtoRegionScan scan = pto_analysis::ScanPtoRegion(region); + if (!scan.has_any_parallel) { + // Pure hand-written VMI region (possibly with control flow): skipped + // by both passes; no lanes annotation needed. + return StmtExprMutator::VisitStmt_(op); + } + ICHECK(!scan.has_unsupported_wrapper) + << "[VectorizeParallelToPTO] internal error: unsupported wrapper " + "structure passed Verify: " + << scan.issue; + ICHECK(!scan.units.empty()) + << "[VectorizeParallelToPTO] internal error: parallel units lost " + "after scan"; + + Stmt new_body; + { + // One inverse per unit: the same mapping feeds the lanes read, the + // shared-mask decision and the unit conversion below, so a region + // with n units inverts n times instead of 2n+1. The inversion is + // the expensive part (layout reshape + analyzer proofs); the + // mappings are small value structs. + std::vector mappings; + mappings.reserve(scan.units.size()); + for (const For &unit : scan.units) { + mappings.push_back(pto_analysis::InvertPtoLaneLayout(unit)); + } + // Lanes come from the first unit's layout; all units in one VF + // share the same L (Verify enforces the annotation consistency). + int64_t lanes = mappings[0].lanes; + // A full mask binding per VF, shared by the divisible units. A + // region whose units all have non-divisible extents never uses it + // (every unit builds its own per-chunk dynamic mask), so the + // binding is only emitted when at least one divisible unit needs + // it. + bool any_divisible = false; + for (const pto_analysis::PtoInverseMapping &mapping : mappings) { + if (mapping.tail_lanes == 0) { + any_divisible = true; + break; + } + } + Stmt new_region = region; + if (any_divisible) { + shared_mask_ = Var("mask", DataType::Bool(static_cast(lanes))); + PrimExpr mask_value = Call( + DataType::Bool(static_cast(lanes)), VmiOp("create_mask"), + {IntImm(DataType::Int(32), static_cast(lanes))}, + {{"size", IntImm(DataType::Int(32), static_cast(lanes))}}); + // Define the shared mask inside the simdvf scope attr when the + // region is wrapped by one: VPTO treats vecscope-region data as + // region-local and rejects pset values crossing the boundary. + const auto *scope_attr = new_region.as(); + if (scope_attr != nullptr && + scope_attr->attr_key == "tl.simdvf_scope") { + new_region = AttrStmt(scope_attr->node, scope_attr->attr_key, + scope_attr->value, + SeqStmt::Flatten(Bind(shared_mask_, mask_value), + scope_attr->body)); + } else { + new_region = + SeqStmt::Flatten(Bind(shared_mask_, mask_value), new_region); + } + } + + // Convert each sequential unit in order; unit state (bind + // maps, counters) is reset per unit; the shared full mask, when + // emitted, is visible to every unit. + for (size_t i = 0; i < scan.units.size(); ++i) { + Stmt converted = ConvertUnit(scan.units[i], mappings[i]); + new_region = ReplaceStmt(new_region, scan.units[i], converted); + } + new_body = new_region; + } + + // new_body is (mask bind + region with all units converted); the + // scope AttrStmt wrapper, if any, is inside it already. + return SBlock(op->iter_vars, op->reads, op->writes, op->name_hint, new_body, + op->init, op->alloc_buffers, op->match_buffers, + op->annotations); + } + + // Replace one statement (the unit's For) with its conversion, recursing + // through the same wrapper kinds DiscoverUnits understands. + Stmt ReplaceStmt(const Stmt &body, const For &target, + const Stmt &replacement) { + if (const auto *for_node = body.as()) { + if (for_node == target.get()) { + return replacement; + } + Stmt nb = ReplaceStmt(for_node->body, target, replacement); + if (nb.same_as(for_node->body)) { + return body; + } + For nf = GetRef(for_node); + nf.CopyOnWrite()->body = nb; + return nf; + } + if (const auto *seq = body.as()) { + Array new_seq; + bool changed = false; + for (const auto &s : seq->seq) { + Stmt ns = ReplaceStmt(s, target, replacement); + changed |= !ns.same_as(s); + new_seq.push_back(ns); + } + return changed ? SeqStmt::Flatten(new_seq) : body; + } + if (const auto *attr = body.as()) { + Stmt nb = ReplaceStmt(attr->body, target, replacement); + if (nb.same_as(attr->body)) { + return body; + } + return AttrStmt(attr->node, attr->attr_key, attr->value, nb); + } + if (const auto *block = body.as()) { + Stmt nb = ReplaceStmt(block->body, target, replacement); + if (nb.same_as(block->body)) { + return body; + } + return SBlock(block->iter_vars, block->reads, block->writes, + block->name_hint, nb, block->init, block->alloc_buffers, + block->match_buffers, block->annotations); + } + if (const auto *realize = body.as()) { + // nb is the replaced *body* of the block (a SeqStmt/For/etc.), not an + // SBlock: rebuild the block with the new body first. + SBlock block = realize->block; + Stmt nb = ReplaceStmt(block->body, target, replacement); + if (nb.same_as(block->body)) { + return body; + } + SBlock new_block = + SBlock(block->iter_vars, block->reads, block->writes, + block->name_hint, nb, block->init, block->alloc_buffers, + block->match_buffers, block->annotations); + return SBlockRealize(realize->iter_values, realize->predicate, new_block); + } + return body; + } + + // ------------------------------------------------------------------ + // Per-unit conversion + // ------------------------------------------------------------------ + Stmt ConvertUnit(const For &loop, + const pto_analysis::PtoInverseMapping &mapping) { + // The inverse mapping comes from the region entry: it is the shared + // source of truth for lanes, the mask plan and the body conversion + // (see InvertPtoLaneLayout). + // Reset per-unit conversion state. + bind_env_.clear(); + value_map_.clear(); + tmp_counter_ = 0; + // Keep the complete mapping (2D fields included) for the + // per-access analysis calls in ConvertBody. + unit_mapping_ = mapping; + + // For a 2D unit the vectorized coordinate is the selected + // dimension; the other one stays a serial loop var. The converted + // body is the innermost compute body. + const ForNode *inner = nullptr; + if (const auto *body_for = loop->body.as()) { + if (body_for->kind == ForKind::kParallel) { + inner = body_for; + } + } + const Var vector_var = + mapping.is_2d + ? (mapping.select_inner ? mapping.inner_var : mapping.outer_var) + : loop->loop_var; + loop_var_ = vector_var; + lanes_ = mapping.lanes; + // The chunk loop covers the padded chunk count Q = P / L + // (mapping.chunk_count), which includes the tail chunk of a + // non-divisible extent. + nchunks_ = mapping.chunk_count; + chunk_var_ = Var("q", vector_var.dtype()); + lane_var_ = mapping.lane_var; // analysis placeholder; never escapes + chunk_analysis_ = mapping.chunk_var; + + // The substitution loop_var -> chunk_analysis_ * L + lane_var_ drives + // all address/index rewriting (same expression the shared helper + // uses). + lane_indexed_expr_ = mapping.index_expr; + + // A divisible extent uses the caller's shared full mask (constant + // across chunks). A non-divisible extent gets a per-unit scalar + // `remaining` that starts at the logical extent E and decreases by L + // per chunk; each chunk binds a dynamic create_mask(remaining, L) + // before the body and updates remaining at the chunk end. The tail + // chunk's mask zeroes the lanes past E, which protects the store + // (loads still read the full padded range — the caller guarantees + // its readability; the store mask never widens a write past E). + // The mask var must exist before the body conversion: the + // converted vstore/vadd reference it while the body is built. + const bool has_tail = mapping.tail_lanes != 0; + if (!has_tail) { + mask_var_ = shared_mask_; + } else { + // decl_buffer shapes the Var with a local.var pointer type; the + // codegen reads the scope from there and emits the scalar + // local-variable form the PTODSL surface expects. + remaining_buffer_ = + tirx::decl_buffer({IntImm(DataType::Int(32), 1)}, DataType::Int(32), + "remaining", "local.var"); + // The per-chunk mask bind is emitted inside the chunk loop below; + // this var is what it binds and what the body references. + mask_var_ = Var("mask", DataType::Bool(static_cast(lanes_))); + } + + // Convert the unit body statements in order. For a 2D unit this is + // the innermost compute body; the outer (kept) coordinate is + // re-emitted as a serial loop around the chunk loop. + Array unit_stmts; + ConvertBody(inner != nullptr ? static_cast(inner->body) + : static_cast(loop->body), + &unit_stmts); + + // Drop binds whose value is never used (address-only aliases and + // unused source binds strand vci/vbrc/vadd/vload chains) before the + // chunk loop is built. + Stmt eliminated = DeadBindElim()(SeqStmt::Flatten(unit_stmts)); + Array chunk_seq; + if (!has_tail) { + chunk_seq.push_back(std::move(eliminated)); + } else { + // Per-chunk mask bind + body + remaining update, in that order. + // mask_var_ was created before the body conversion. + PrimExpr remaining_load = + BufferLoad(remaining_buffer_, {make_zero(DataType::Int(32))}); + PrimExpr mask_value = + Call(DataType::Bool(static_cast(lanes_)), VmiOp("create_mask"), + {remaining_load}, + {{"size", IntImm(DataType::Int(32), static_cast(lanes_))}}); + chunk_seq.push_back(Bind(mask_var_, mask_value)); + chunk_seq.push_back(std::move(eliminated)); + PrimExpr next_remaining = + remaining_load - + IntImm(DataType::Int(32), static_cast(lanes_)); + chunk_seq.push_back(BufferStore(remaining_buffer_, next_remaining, + {make_zero(DataType::Int(32))})); + } + Stmt chunk_body = SeqStmt::Flatten(chunk_seq); + Stmt chunk_loop = + For(chunk_var_, make_zero(chunk_var_->dtype), + IntImm(chunk_var_->dtype, static_cast(nchunks_)), + ForKind::kSerial, chunk_body); + if (inner == nullptr) { + if (!has_tail) { + return chunk_loop; + } + // remaining lives across the whole chunk loop: declare it, set it + // to E once before the loop. + Stmt init = + BufferStore(remaining_buffer_, + IntImm(DataType::Int(32), + static_cast(mapping.logical_extent)), + {make_zero(DataType::Int(32))}); + return SeqStmt::Flatten(AllocBuffer(remaining_buffer_), + SeqStmt::Flatten(init, chunk_loop)); + } + // 2D: keep the non-vectorized coordinate as the outer + // serial loop; the vectorized dimension becomes the chunk loop. + const Var kept_var = + mapping.select_inner ? loop->loop_var : inner->loop_var; + const int64_t kept_extent = + mapping.select_inner ? mapping.outer_extent : mapping.inner_extent; + Stmt kept_body = chunk_loop; + if (has_tail) { + // remaining resets to E inside the kept loop: every row starts a + // fresh recursion, so rows are independent and the tail chunk of + // each row masks exactly that row's tail lanes. + Stmt init = + BufferStore(remaining_buffer_, + IntImm(DataType::Int(32), + static_cast(mapping.logical_extent)), + {make_zero(DataType::Int(32))}); + kept_body = SeqStmt::Flatten(init, chunk_loop); + } + Stmt kept_loop = + For(kept_var, make_zero(kept_var->dtype), + IntImm(kept_var->dtype, static_cast(kept_extent)), + ForKind::kSerial, kept_body); + if (!has_tail) { + return kept_loop; + } + return SeqStmt::Flatten(AllocBuffer(remaining_buffer_), kept_loop); + } + + // Substitute the analysis chunk var with the real loop var q (and keep + // the lane placeholder for lane-binding steps). + PrimExpr BindChunkToQ(const PrimExpr &expr) const { + class Replacer : public ExprMutator { + public: + Replacer(const VarNode *from, const Var &to) : from_(from), to_(to) {} + PrimExpr VisitExpr_(const VarNode *op) final { + if (op == from_) { + return to_; + } + return ffi::GetRef(op); + } + + private: + const VarNode *from_; + Var to_; + }; + return Replacer(chunk_analysis_.get(), chunk_var_)(expr); + } + + // Rewrite Var references inside metadata objects through value_map_ + // Follows containers (Array/Map) structurally; PrimExprs get a + // var substitution. + ObjectRef RemapMetadataVars(const ObjectRef &obj) const { + if (!obj.defined()) { + return obj; + } + if (const auto *expr = obj.as()) { + class RemapVars : public ExprMutator { + public: + RemapVars(const Map &vmap) : vmap_(vmap) {} + PrimExpr VisitExpr_(const VarNode *op) final { + auto it = vmap_.find(GetRef(op)); + if (it != vmap_.end()) { + return (*it).second; + } + return GetRef(op); + } + + private: + const Map &vmap_; + }; + return RemapVars(value_map_)(GetRef(expr)); + } + // Containers are rejected by Verify (metadata must be a PrimExpr); + // reaching here with one is an internal error. + ICHECK(obj.as() == nullptr) + << "[VectorizeParallelToPTO] internal error: container metadata " + "passed Verify"; + return obj; + } + + PrimExpr SubstituteLaneZero(const PrimExpr &expr) const { + class ZeroLane : public ExprMutator { + public: + ZeroLane(const Var &lane) : lane_(lane) {} + PrimExpr VisitExpr_(const VarNode *op) final { + if (op == lane_.get()) { + return make_zero(lane_->dtype); + } + return ffi::GetRef(op); + } + + private: + Var lane_; + }; + return ZeroLane(lane_var_)(expr); + } + + // ------------------------------------------------------------------ + // Statement conversion (pre-binds first, then the statement) + // ------------------------------------------------------------------ + void ConvertBody(const Stmt &body, Array *out) { + if (const auto *seq = body.as()) { + for (const auto &s : seq->seq) { + ConvertBody(s, out); + } + return; + } + if (const auto *bind = body.as()) { + // Whitelist is Verify's job; defensive check here. + Array pre; + PrimExpr value = ConvertExpr(bind->value, &pre); + // Pre-binds first, then the Bind itself. + for (const auto &s : pre) { + out->push_back(s); + } + Var new_var(bind->var->name_hint, value.dtype()); + bind_env_.Set(bind->var, bind->value); // analysis environment + value_map_.Set(bind->var, new_var); // value environment + out->push_back(Bind(new_var, value)); + return; + } + if (const auto *store = body.as()) { + Array pre; + PrimExpr value = ConvertExpr(store->value, &pre); + auto access = pto_analysis::AnalyzeBufferAccess( + store->buffer, store->indices, /*is_write=*/true, loop_var_, + CurrentMapping(), &analyzer_, bind_env_); + if (access.pattern == pto_analysis::AccessPattern::kContinuous) { + if (value.dtype().lanes() == 1) { + value = EmitBind( + "vbrc_val", + Call(VectorDType(value.dtype(), lanes_), VmiOp("vbrc"), {value}, + {{"size", + IntImm(DataType::Int(32), static_cast(lanes_))}}), + &pre); + } + Array start_idx; + for (const auto &idx : access.start_indices) { + start_idx.push_back( + analyzer_.Simplify(BindChunkToQ(SubstituteLaneZero(idx)))); + } + PrimExpr ptr = + Call(DataType::Handle(), tl::access_ptr(), + {BufferLoad(store->buffer, start_idx), + IntImm(DataType::Int(32), static_cast(lanes_)), + IntImm(DataType::Int(32), 2)}); + for (const auto &s : pre) { + out->push_back(s); + } + out->push_back(Evaluate( + Call(DataType::Void(), VmiOp("vstore"), + {value, ptr, IntImm(DataType::Int(32), 0), mask_var_}))); + } else { + // Verify rejects lane-uniform stores (the current + // support boundary), so this branch is unreachable. Treat any + // survivor as an internal error, same style as the wrapper + // checks above. + ICHECK(false) + << "[VectorizeParallelToPTO] internal error: lane-uniform store `" + << store->buffer->name + << "` passed Verify; lane-uniform stores must be rejected by " + "VerifyParallelToPTO"; + } + return; + } + if (const auto *attr = body.as()) { + // Verify guarantees the metadata has no lane-dependent semantics + // (direct or via bind aliases). Rewrite its Var references through + // value_map_ so already-converted binds are not left dangling + // The body is converted underneath. + Array stmts; + // The body is a nested statement scope — keep the + // source-bind and value maps scoped with it, so a bind defined + // inside cannot remap a later sibling's reference (Verify rejects + // that input; this keeps the transformation state consistent with + // the unit visitor's scope handling). + Map saved_env = bind_env_; + Map saved_map = value_map_; + ConvertBody(attr->body, &stmts); + bind_env_ = std::move(saved_env); + value_map_ = std::move(saved_map); + Stmt inner = SeqStmt::Flatten(stmts); + // Default to the *original* node. Only PrimExpr metadata is + // remapped; a non-object Any (int/str atom metadata) must forward + // unchanged — the old code wrote back a default-constructed (empty) + // node when the try_cast failed. + ffi::Any node = attr->node; + if (auto ne = attr->node.try_cast()) { + if (const auto *pe = ne.value().as()) { + node = RemapMetadataVars(GetRef(pe)); + } + } + PrimExpr value = attr->value; + if (const auto *ve = attr->value.as()) { + value = RemapMetadataVars(GetRef(ve)).as().value(); + } + out->push_back(AttrStmt(node, attr->attr_key, value, inner)); + return; + } + if (const auto *eval = body.as()) { + // Convert the expression with the same machinery so no stale loop + // variable survives; the whitelist is Verify's job. + Array pre; + PrimExpr value = ConvertExpr(eval->value, &pre); + for (const auto &s : pre) { + out->push_back(s); + } + out->push_back(Evaluate(value)); + return; + } + LOG(FATAL) << "[VectorizeParallelToPTO] internal error: unsupported " + "statement passed Verify: " + << body->GetTypeKey(); + } + + pto_analysis::PtoInverseMapping CurrentMapping() const { + // Return the *complete* inverse mapping saved for this unit. + // Rebuilding a 1D-only mapping here dropped is_2d, the two loop vars, + // their extents and select_inner, so the shared access proof lost the + // kept coordinate's range and Vectorize rejected units Verify had + // accepted. The analysis placeholders (chunk/lane vars) are exactly + // the ones saved from InvertPtoLaneLayout; only the emitted chunk + // loop var q is separate (see BindChunkToQ). + return unit_mapping_; + } + + // ------------------------------------------------------------------ + // Expression conversion (index vars resolve through bind_env_) + // ------------------------------------------------------------------ + PrimExpr ConvertExpr(const PrimExpr &expr, Array *pre) { + if (const auto *load = expr.as()) { + auto access = pto_analysis::AnalyzeBufferAccess( + load->buffer, load->indices, /*is_write=*/false, loop_var_, + CurrentMapping(), &analyzer_, bind_env_); + if (access.pattern == pto_analysis::AccessPattern::kContinuous) { + Array start_idx; + for (const auto &idx : access.start_indices) { + start_idx.push_back( + analyzer_.Simplify(BindChunkToQ(SubstituteLaneZero(idx)))); + } + PrimExpr ptr = + Call(DataType::Handle(), tl::access_ptr(), + {BufferLoad(load->buffer, start_idx), + IntImm(DataType::Int(32), static_cast(lanes_)), + IntImm(DataType::Int(32), 1)}); + // VMI with a return value goes through EmitBind. + return EmitBind("vload", + Call(VectorDType(load->dtype, lanes_), VmiOp("vload"), + {ptr, IntImm(DataType::Int(32), 0)}, + {{"size", IntImm(DataType::Int(32), + static_cast(lanes_))}}), + pre); + } + // lane-uniform scalar load: keep the BufferLoad (lane bound to 0, + // chunk bound to q — indices resolve through the bind env + // inside the shared helper already) and bind it to a scalar Var in + // the current chunk; later uses (vbrc broadcast, scalar arithmetic, + // store values) reference that Var. Each textual occurrence binds + // separately — no cross-access load reuse. The Var dtype is the + // buffer's own scalar dtype. + Array indices; + for (const auto &idx : access.indices_after_inverse) { + indices.push_back( + analyzer_.Simplify(BindChunkToQ(SubstituteLaneZero(idx)))); + } + return EmitBind("sload", BufferLoad(load->buffer, indices), pre); + } + if (const auto *var = expr.as()) { + auto it = value_map_.find(GetRef(var)); + if (it != value_map_.end()) { + return (*it).second; + } + if (var == loop_var_.get()) { + // lane_indexed used as a value: explicit vci vector. + PrimExpr base = analyzer_.Simplify( + BindChunkToQ(SubstituteLaneZero(lane_indexed_expr_))); + return EmitBind( + "vci", + Call(VectorDType(var->dtype, lanes_), VmiOp("vci"), {base}, + {{"size", + IntImm(DataType::Int(32), static_cast(lanes_))}}), + pre); + } + return GetRef(var); + } + if (expr.as() || expr.as()) { + return expr; + } + if (const auto *add = expr.as()) { + return ConvertBinary(add->a, add->b, "vadd", pre); + } + if (const auto *sub = expr.as()) { + return ConvertBinary(sub->a, sub->b, "vsub", pre); + } + if (const auto *mul = expr.as()) { + return ConvertBinary(mul->a, mul->b, "vmul", pre); + } + if (const auto *div = expr.as()) { + return ConvertBinary(div->a, div->b, "vdiv", pre); + } + if (const auto *call = expr.as()) { + // Hardware-backed unary math (fabs/exp/log/sqrt); the dtype matrix + // mirrors the ISA documentation (see IsSupportedVectorUnaryOp). + static const std::unordered_map kUnaryMap = { + {"tirx.fabs", "vabs"}, + {"tirx.exp", "vexp"}, + {"tirx.log", "vln"}, + {"tirx.sqrt", "vsqrt"}}; + auto call_op = call->op.as(); + ICHECK(call_op.has_value()) + << "[VectorizeParallelToPTO] internal error: non-Op call passed " + "Verify: " + << expr; + auto it = kUnaryMap.find(call_op.value()->name); + ICHECK(it != kUnaryMap.end()) + << "[VectorizeParallelToPTO] internal error: call passed Verify " + "without a VMI mapping: " + << call_op.value()->name; + PrimExpr a = ConvertExpr(call->args[0], pre); + if (a.dtype().lanes() == 1) { + return Call(expr.dtype(), call->op, {a}); + } + return EmitBind(it->second, + Call(VectorDType(expr.dtype().element_of(), lanes_), + VmiOp(it->second), {a, mask_var_}), + pre); + } + if (const auto *cast = expr.as()) { + PrimExpr value = ConvertExpr(cast->value, pre); + // Emit the conversion attributes PlanValueCast proved + // equivalent; never rely on PTODSL defaults (they saturate integer + // narrowing). + std::string plan_why; + auto plan = pto_analysis::PlanValueCast(cast, &plan_why); + ICHECK(plan.has_value()) + << "[VectorizeParallelToPTO] internal error: Cast passed Verify " + "without a conversion plan: " + << plan_why; + if (plan.value().identity) { + // Same-type cast is a no-op: keep the converted value. + return value; + } + if (value.dtype().lanes() > 1) { + Map attrs; + attrs.Set("to_dtype", StringImm(TvmDtypeName(cast->dtype))); + if (!plan.value().saturate.empty()) { + attrs.Set("saturate", StringImm(plan.value().saturate)); + } + if (!plan.value().rounding.empty()) { + attrs.Set("rounding", StringImm(plan.value().rounding)); + } + return EmitBind("vcvt", + Call(VectorDType(cast->dtype, value.dtype().lanes()), + VmiOp("vcvt"), {value}, attrs), + pre); + } + // A lane-uniform FP8 conversion has no scalar PTO lowering inside a + // SIMD_VF body; Verify rejects that form, so reaching it here means the + // two passes disagree. + ICHECK(!pto_analysis::IsSupportedFloat8(cast->dtype) && + !pto_analysis::IsSupportedFloat8(cast->value.dtype())) + << "[VectorizeParallelToPTO] internal error: lane-uniform FP8 Cast " + "passed Verify: " + << cast->value.dtype() << " -> " << cast->dtype; + return Cast(cast->dtype, value); + } + LOG(FATAL) << "[VectorizeParallelToPTO] internal error: unsupported " + "expression passed Verify: " + << expr->GetTypeKey(); + } + + PrimExpr ConvertBinary(const PrimExpr &lhs, const PrimExpr &rhs, + const std::string &vmi_name, Array *pre) { + PrimExpr a = ConvertExpr(lhs, pre); + PrimExpr b = ConvertExpr(rhs, pre); + bool a_vec = a.dtype().lanes() > 1; + bool b_vec = b.dtype().lanes() > 1; + if (!a_vec && !b_vec) { + // Lane-uniform folds stay scalar with the TIR semantics of the node. + if (vmi_name == "vadd") { + return a + b; + } + if (vmi_name == "vsub") { + return a - b; + } + if (vmi_name == "vmul") { + return a * b; + } + if (vmi_name == "vdiv") { + return div(a, b); + } + ICHECK(false) << "[VectorizeParallelToPTO] internal error: unknown " + "binary op " + << vmi_name; + } + if (a_vec && !b_vec) { + b = EmitBind( + "brc", + Call(VectorDType(b.dtype(), lanes_), VmiOp("vbrc"), {b}, + {{"size", IntImm(DataType::Int(32), static_cast(lanes_))}}), + pre); + } else if (!a_vec && b_vec) { + a = EmitBind( + "brc", + Call(VectorDType(a.dtype(), lanes_), VmiOp("vbrc"), {a}, + {{"size", IntImm(DataType::Int(32), static_cast(lanes_))}}), + pre); + } + return EmitBind(vmi_name, + Call(VectorDType(a.dtype().element_of(), lanes_), + VmiOp(vmi_name), {a, b, mask_var_}), + pre); + } + + // TVM dtype name identical to Python str(dtype) (vmi.py ABI): + // float32 -> "float32", bfloat16 -> "bfloat16", int8 -> "int8", ... + static std::string TvmDtypeName(DataType dtype) { + std::ostringstream os; + if (dtype.is_bfloat16()) { + return "bfloat16"; + } + // The FP8 spellings the PTO dtype parser recognizes; both are the TIR + // dtype names, not the PTO target names. + if (dtype.is_float8_e4m3fn()) { + return "float8_e4m3fn"; + } + if (dtype.is_float8_e5m2()) { + return "float8_e5m2"; + } + if (dtype.is_float()) { + os << "float" << dtype.bits(); + } else if (dtype.is_int()) { + os << "int" << dtype.bits(); + } else if (dtype.is_uint()) { + os << "uint" << dtype.bits(); + } else if (dtype.is_bool()) { + return "bool"; + } else { + os << dtype; + } + return os.str(); + } + + Var EmitBind(const std::string &name, const PrimExpr &value, + Array *pre) { + Var result(name + "_" + std::to_string(tmp_counter_++), value.dtype()); + pre->push_back(Bind(result, value)); + return result; + } + + arith::Analyzer analyzer_; + Var shared_mask_; // one full mask per VF, shared by all units + + // Per-unit state. + Var loop_var_; + Var chunk_var_; // real emitted loop var (q) + Var chunk_analysis_; // analysis placeholder from the inverse mapping + Var lane_var_; // analysis placeholder; never escapes + PrimExpr lane_indexed_expr_; + /*! The complete inverse mapping of the unit being converted + * (2D fields included), consumed by CurrentMapping(). */ + pto_analysis::PtoInverseMapping unit_mapping_; + int64_t lanes_ = 0; + int64_t nchunks_ = 0; + Var mask_var_; + /*! Per-unit scalar buffer holding `remaining` for a non-divisible + * extent (invalid for divisible units). */ + tirx::Buffer remaining_buffer_{}; + int tmp_counter_ = 0; + Map bind_env_; // analysis env: original var -> def expr + Map value_map_; // value env: original var -> converted var +}; + +} // namespace + +namespace transform { + +tvm::transform::Pass VectorizeParallelToPTO() { + auto pass_func = [](PrimFunc f, const IRModule &m, + const tvm::transform::PassContext &ctx) { + auto target = f->GetAttr(tvm::attr::kTarget); + if (!target.defined() || !TargetIsPTO(target.value())) { + return f; + } + return VectorizeParallelToPTOImpl::Substitute(std::move(f)); + }; + return tvm::tirx::transform::CreatePrimFuncPass( + pass_func, 0, "tl.VectorizeParallelToPTO", {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tl.transform.VectorizeParallelToPTO", + VectorizeParallelToPTO); +} + +} // namespace transform + +} // namespace tl +} // namespace tvm diff --git a/src/ascend/transform/verify_parallel_to_pto.cc b/src/ascend/transform/verify_parallel_to_pto.cc new file mode 100644 index 0000000000..1e03801dae --- /dev/null +++ b/src/ascend/transform/verify_parallel_to_pto.cc @@ -0,0 +1,1371 @@ +// Copyright (c) Tile-AI Corporation. +// Licensed under the MIT License. + +/*! + * \file verify_parallel_to_pto.cc + * \brief Read-only input-contract verification for the PTO Parallel + * vectorize path. Runs on the PTO target only. + * + * The pass checks every SIMD_VF block that contains a T.Parallel unit: + * lanes validity, loop structure/ranges, lane fragment presence and + * invertibility, per-access patterns (continuous / lane-uniform), the + * value-expression whitelist, element dtype constraints and region-level + * rules (no GM access, no mixed hand-written VMI, no threadIdx binding). + * References are checked against a lexical definition environment: + * params and buffer data pointers at entry; outer loop vars, + * thread_extent bindings, flat binds and buffer allocations registered + * at their statement position and restored on scope exit, so a unit + * sees exactly what is defined where it appears. The first error + * terminates compilation with a `[VerifyParallelToPTO]` diagnostic. + * + * Pure hand-written VMI regions (no T.Parallel) and SIMT_VF blocks are + * skipped entirely; a region without T.Parallel does not need the lanes + * annotation. + */ + +#include +#include + +#include +#include +#include +#include + +#include "../../transform/common/attr.h" +#include "parallel_to_pto_utils.h" + +namespace tvm { +namespace tl { + +using namespace tirx; +using namespace ffi; +namespace pto_analysis = ::tvm::tl::pto; + +namespace { + +class VerifyParallelToPTOImpl : public StmtExprVisitor { +public: + void Verify(const PrimFunc &f) { + // The definition environment starts at function entry with + // the params and the buffer data pointers. Every other definition is + // registered at its lexical position during the traversal below and + // removed when its scope exits. + for (const auto &p : f->params) { + defs_.insert(p.get()); + } + for (const auto &kv : f->buffer_map) { + defs_.insert(kv.second->data.get()); + } + CollectVolatileAllocs(f->body); + VisitStmt(f->body); + } + +private: + // Function-level pre-scan for AllocBuffer nodes carrying the + // tirx.volatile annotation. Presence of the annotation marks the + // allocation volatile (PTO codegen semantics: + // `annotations.count(tirx::attr::kVolatile)`); a False value is NOT + // interpreted as opt-out. The buffer *data Var* identity is recorded, + // so any other Buffer view over the same data is covered as well. + // Lexical visibility of the var stays with the existing definition + // checks — this set only answers "is this data volatile". + void CollectVolatileAllocs(const Stmt &body) { + PostOrderVisit(body, [&](const ObjectRef &n) { + if (const auto *alloc = n.as()) { + if (alloc->annotations.count(tirx::attr::kVolatile)) { + volatile_allocs_.insert(alloc->buffer->data.get()); + } + } + }); + } + +private: + [[noreturn]] void Fail(const std::string &category, + const std::string &message) { + std::ostringstream oss; + oss << "[VerifyParallelToPTO] " << category << ": " << message; + // Diagnostics carry the SIMD_VF/Parallel position. The prefix format + // above is unchanged; the position is appended so existing + // category/detail matching keeps working. + if (!unit_context_.empty()) { + oss << " [at " << unit_context_ << "]"; + } + LOG(FATAL) << oss.str(); + } + + // Region-level context — failures between/around units report + // at least the region (and the lane count once it is known). + void SetRegionContext() { + std::ostringstream oss; + oss << "region #" << current_region_; + if (current_lanes_ > 0) { + oss << " (lanes " << current_lanes_ << ")"; + } + unit_context_ = oss.str(); + } + + // ------------------------------------------------------------------ + // Unified unit discovery — same rules as the vectorize pass: + // sequential units are allowed; each unit is a 1D Parallel or a + // two-layer adjacent Parallel nest (2D). A third parallel + // layer, a Parallel under IfThenElse/While, and other unsupported + // structures are rejected explicitly. + // ------------------------------------------------------------------ + int64_t ReadLanes(const SBlockNode *block) { + auto anno = block->annotations.Get("tl.simdvf_lanes"); + if (!anno.has_value()) { + Fail("lanes", "SIMD_VF block with T.Parallel must carry the " + "tl.simdvf_lanes annotation (written by " + "T.SimdVF(lanes=...))"); + } + auto imm = anno.value().try_cast(); + if (!imm.has_value() || imm.value()->dtype != DataType::Int(64)) { + Fail("lanes", "tl.simdvf_lanes must be an Int64 constant"); + } + int64_t lanes = imm.value()->value; + if (lanes != 64 && lanes != 128 && lanes != 256) { + std::ostringstream oss; + oss << "T.SimdVF lanes must be one of {64, 128, 256}, got " << lanes; + Fail("lanes", oss.str()); + } + return lanes; + } + + // ------------------------------------------------------------------ + // Region-level checks: everything inside a converting SIMD_VF region + // must obey the region rules, not just the parallel body. + // ------------------------------------------------------------------ + class RegionVisitor : public StmtExprVisitor { + public: + RegionVisitor(VerifyParallelToPTOImpl *parent) : parent_(parent) {} + + void VisitStmt_(const BufferStoreNode *op) final { + CheckScope(op->buffer); + // The base visitor skips BufferStore.predicate + // (stmt_functor.cc:85); accesses inside it must reach the same + // admission check (the field sits outside the converting Parallel + // unit, so the unit-level predicate rule cannot cover it). + if (op->predicate.defined()) { + VisitExpr(op->predicate.value()); + } + StmtExprVisitor::VisitStmt_(op); + } + void VisitExpr_(const BufferLoadNode *op) final { + CheckScope(op->buffer); + // Same gap for BufferLoad.predicate (expr_functor.cc:36); + // indices stay with the base visitor. + if (op->predicate.defined()) { + VisitExpr(op->predicate.value()); + } + StmtExprVisitor::VisitExpr_(op); + } + void VisitExpr_(const CallNode *op) final { + if (const auto *op_node = op->op.as()) { + std::string name = op_node->name; + if (name.rfind("tl.vmi.", 0) == 0) { + parent_->Fail( + "mixing", + "hand-written tl.vmi.* calls cannot be mixed with T.Parallel " + "inside a converting SIMD_VF block"); + } + if (name == "tl.tileop.reduce") { + parent_->Fail("expression", "cross-lane reduce is not supported"); + } + } + StmtExprVisitor::VisitExpr_(op); + } + void VisitStmt_(const AttrStmtNode *op) final { + if (op->attr_key == tirx::attr::thread_extent) { + // Name the binding that triggered the rejection: a region can carry + // several thread_extent scopes and the message must say which one. + // The prefix says "thread bindings" — this branch also rejects + // blockIdx bindings, not only threadIdx. + std::ostringstream oss; + oss << "thread bindings are not allowed inside a converting " + "SIMD_VF block (found `thread_extent` on binding `"; + if (const auto *iv = op->node.as()) { + oss << iv->var->name_hint; + } else { + oss << op->node; + } + oss << "`)"; + parent_->Fail("region", oss.str()); + } + // The legacy volatile_scope marker declares volatile access + // semantics that are not currently supported; reject it under the + // same rule instead of letting it pass as ordinary metadata. Scope + // note: only markers *inside a converting region* are covered here + // — a volatile_scope wrapping a whole SIMD_VF block from outside + // never reaches this visitor and stays accepted. + if (op->attr_key == tl::attr::volatile_scope) { + parent_->Fail("memory access", + "volatile_scope attribute inside a converting region " + "marks accesses volatile; volatile accesses are not " + "yet supported"); + } + // The base visitor only descends into value/body. A + // PrimExpr held by `node` (metadata) can itself contain buffer + // accesses (a bare BufferLoad or a composite Add/Cast over one, a + // view over the same volatile data Var); run the same admission + // walk over the full expression so those loads reach CheckScope. + // Non-PrimExpr markers keep the existing contract (the unit-level + // metadata check rejects containers). + if (auto ne = op->node.try_cast()) { + if (const auto *pe = ne.value().as()) { + VisitExpr(GetRef(pe)); + } + } + StmtExprVisitor::VisitStmt_(op); + } + void VisitStmt_(const ForNode *op) final { + // A For-form thread binding (the public T.thread_binding API + // constructs exactly this shape) must obey the same region-wide ban + // as the thread_extent AttrStmt form above — the shared unit scan + // only rejects wrappers *around* a Parallel, not sibling bindings. + // This visitor only runs on regions that contain a Parallel, so + // hand-written VMI regions without one are unaffected. Vectorized + // sibling loops are intentionally NOT rejected here (out of scope). + if (op->kind == ForKind::kThreadBinding) { + std::ostringstream oss; + oss << "thread bindings are not allowed inside a converting " + "SIMD_VF block (found ThreadBinding loop on `" + << op->loop_var->name_hint << "`"; + if (op->thread_binding.defined()) { + oss << " bound to " << op->thread_binding.value()->thread_tag; + } + oss << ")"; + parent_->Fail("region", oss.str()); + } + StmtExprVisitor::VisitStmt_(op); + } + + private: + void CheckScope(const Buffer &buffer) { + // Volatile accesses are outside the supported contract + // (an unused volatile read must not be DCE-deleted, and a live + // uniform volatile read must not change its read count through the + // sload+broadcast lowering). The data-Var identity is matched, so + // any Buffer view over a volatile allocation is covered. + if (parent_->volatile_allocs_.count(buffer->data.get()) != 0) { + std::ostringstream oss; + oss << "buffer `" << buffer->name + << "` is allocated with the tirx.volatile annotation; volatile " + "accesses are not yet supported"; + parent_->Fail("memory access", oss.str()); + } + const std::string &scope = buffer.scope(); + if (scope == "global") { + std::ostringstream oss; + oss << "SIMD_VF block must not access GM directly (buffer `" + << buffer->name << "` has scope global); stage through UB first"; + parent_->Fail("memory access", oss.str()); + } + // Only UB accesses are currently supported. local / + // local.fragment / shared.l1 (L1) and other scopes are rejected. + if (scope != "shared" && scope != "shared.dyn") { + std::ostringstream oss; + oss << "buffer `" << buffer->name << "` has scope `" << scope + << "`; only UB (shared/shared.dyn) accesses are supported"; + parent_->Fail("memory access", oss.str()); + } + if (!pto_analysis::IsSupportedElementDType(buffer->dtype)) { + std::ostringstream oss; + oss << "buffer `" << buffer->name << "` element dtype " << buffer->dtype + << " is not supported by the unified element " + "addressing (no sub-byte packed types, no " + "pre-vectorized elements, no FP64, 8/16/32-bit widths only)"; + parent_->Fail("dtype", oss.str()); + } + } + VerifyParallelToPTOImpl *parent_; + }; + + // ------------------------------------------------------------------ + // Unit-level checks + // ------------------------------------------------------------------ + class UnitVisitor : public StmtExprVisitor { + public: + UnitVisitor(VerifyParallelToPTOImpl *parent, const Var &loop_var, + const pto_analysis::PtoInverseMapping &mapping, + const std::set &external_defs) + : parent_(parent), loop_var_(loop_var), mapping_(mapping), + external_defs_(external_defs) { + // In a 2D unit `loop_var` is the vectorized coordinate; the + // other one is kept as a serial loop and classifies as uniform. + if (mapping.is_2d) { + serial_var_ = + mapping.select_inner ? mapping.outer_var : mapping.inner_var; + } + } + + void VisitStmt_(const AttrStmtNode *op) final { + // Analyze the *whole* metadata expression with the shared + // lane-use analysis — composite expressions (Add/Cast over loads) + // classify correctly, bind identity is preserved, and unknown forms + // produce a diagnostic instead of a crash. Containers are still + // rejected. + auto check_meta_expr = [&](const PrimExpr &meta) { + // The unified reference check runs *first*, so a + // ghost anywhere in the metadata expression (including load + // indices, through any admitted wrapper node) reports the + // definition diagnostic before address analysis can mis-route + // it. The base class descends into every child expression, so + // the earlier hand-written node table's gap (Div/Mod/Min/Max/ + // comparisons treated as leaves, loop-var references slipping + // through them) cannot recur. The metadata position additionally + // rejects direct loop-var references and vector-classified + // binds: the metadata path only remaps scalar var references. + RefCheck checker(this, /*metadata=*/true); + checker(meta); + // Metadata-mode classification (Sub/FloorMod/compare/ + // StringImm propagate instead of being value-whitelist rejects). + pto_analysis::LaneUseContext ctx = MakeLaneUseContext(); + std::string why; + pto_analysis::LaneUse use = + pto_analysis::AnalyzeLaneUseMeta(meta, &ctx, &why); + if (use == pto_analysis::LaneUse::kVarying) { + parent_->Fail("expression", + "AttrStmt metadata depends on lane-varying values; " + "metadata cannot yet be given lane semantics"); + } + if (use == pto_analysis::LaneUse::kUnknown) { + std::ostringstream oss; + oss << "AttrStmt metadata cannot be analyzed: " << why; + parent_->Fail("expression", oss.str()); + } + }; + // op->node may be an Any holding a PrimExpr; containers and other + // object kinds are rejected. + if (auto ne = op->node.try_cast()) { + if (const auto *pe = ne.value().as()) { + check_meta_expr(GetRef(pe)); + } else { + parent_->Fail("expression", + "AttrStmt metadata must be a PrimExpr (containers " + "cannot be given lane semantics)"); + } + } + // Non-object metadata (int/str atoms) carries no Var references. + check_meta_expr(op->value); + // An AttrStmt body is a nested statement scope — a Bind + // defined inside it must not leak to later siblings, while outer + // binds stay visible inside. + Map saved_env = bind_env_; + Map saved_classified = classified_; + StmtExprVisitor::VisitStmt_(op); + bind_env_ = std::move(saved_env); + classified_ = std::move(saved_classified); + } + // The StmtVisitor base class already dispatches several node + // kinds (AllocBuffer/DeclBuffer/...) so VisitStmtDefault_ alone is + // not a whitelist. Reject those kinds explicitly. + + void VisitStmt_(const ForNode *op) final { + // No loops of any kind inside T.Parallel. + parent_->Fail("loop structure", + "loops (serial/pipelined/while) inside T.Parallel are " + "not yet supported"); + } + void VisitStmt_(const IfThenElseNode *) final { + parent_->Fail("loop structure", + "conditional branches inside T.Parallel are not yet " + "supported"); + } + void VisitStmt_(const WhileNode *) final { + parent_->Fail("loop structure", + "while loops inside T.Parallel are not yet supported"); + } + void VisitStmt_(const SBlockNode *) final { + parent_->Fail("loop structure", + "nested blocks inside T.Parallel are not yet " + "supported"); + } + void VisitStmt_(const SBlockRealizeNode *) final { + parent_->Fail("loop structure", + "nested blocks inside T.Parallel are not yet " + "supported"); + } + void VisitStmt_(const AllocBufferNode *) final { + parent_->Fail("loop structure", + "buffer allocation inside T.Parallel is not yet " + "supported"); + } + void VisitStmt_(const DeclBufferNode *) final { + parent_->Fail("loop structure", + "buffer declaration inside T.Parallel is not yet " + "supported"); + } + void VisitStmtDefault_(const ffi::Object *op) final { + // Backstop for kinds the base does not dispatch. + std::ostringstream oss; + oss << "statement `" << op->GetTypeKey() + << "` inside T.Parallel is not yet supported"; + parent_->Fail("loop structure", oss.str()); + } + void VisitStmt_(const AssertStmtNode *) final { + parent_->Fail("loop structure", + "assert statements inside T.Parallel are not yet " + "supported"); + } + void VisitStmt_(const BindNode *op) final { + // Self/forward reference is rejected before registering. + bool self_ref = false; + PostOrderVisit(op->value, [&](const ObjectRef &n) { + if (auto *v = n.as()) { + if (v == op->var.get()) { + self_ref = true; + } + } + }); + if (self_ref) { + std::ostringstream oss; + oss << "bind `" << op->var->name_hint + << "` references itself in its definition"; + parent_->Fail("expression", oss.str()); + } + // Bind right-hand sides run the shared + // compute-value admission (whitelist, cast plans, op types, + // references, lane classification) — the same entry as Evaluate + // and store values. + pto_analysis::LaneUse use = CheckComputeValue( + op->value, "bind `" + op->var->name_hint + "` value"); + // Register only after the analysis succeeded. + bind_env_.Set(op->var, op->value); + classified_.Set(op->var, + Integer(use == pto_analysis::LaneUse::kUniform ? 0 : 1)); + StmtExprVisitor::VisitStmt_(op); + } + void VisitStmt_(const EvaluateNode *op) final { + // An Evaluate value is a compute value — same admission as + // Bind RHS and store values (it previously ran no lane + // classification, so a StringImm reached Vectorize's internal + // error). + CheckComputeValue(op->value, "evaluate value"); + StmtExprVisitor::VisitStmt_(op); + } + void VisitStmt_(const BufferStoreNode *op) final { + // The store value runs the shared compute-value + // admission (whitelist, cast plans, op types, references, lane + // classification). + CheckComputeValue(op->value, "store value"); + // Store indices get the same unified reference check — the + // address continuity proof can cancel a common ghost offset and + // must not double as a definition check. + CheckDefinedVars(op->indices); + VerifyAccess(op->buffer, op->indices, /*is_write=*/true, op->predicate); + StmtExprVisitor::VisitStmt_(op); + } + void VisitExpr_(const BufferLoadNode *op) final { + // Load indices are checked against the same definition + // environment as values and store indices. + CheckDefinedVars(op->indices); + VerifyAccess(op->buffer, op->indices, /*is_write=*/false, op->predicate); + StmtExprVisitor::VisitExpr_(op); + } + + private: + // One construction point for the lane-use context. The + // metadata, operation-type and bind-definition entries share the same + // definition environment, classification cache and guards. + pto_analysis::LaneUseContext MakeLaneUseContext() { + pto_analysis::LaneUseContext ctx; + ctx.loop_var = loop_var_; + ctx.serial_var = serial_var_; + ctx.mapping = mapping_; + ctx.analyzer = &analyzer_; + ctx.bind_env = bind_env_; + ctx.external_defs = external_defs_; + ctx.classified = classified_; + return ctx; + } + + // The unified reference check. One StmtExprVisitor + // covers every expression position — values, load/store indices and + // metadata node/value — relying on the base class's complete descent + // (BufferLoad indices and every wrapper node the admission tables + // accept are all visited). Every Var must be the loop var, a bind + // defined earlier in this unit, or a definition in the lexical + // snapshot; the metadata position additionally rejects direct + // loop-var references and vector-classified binds because the + // metadata path only remaps scalar var references. + class RefCheck : public StmtExprVisitor { + public: + RefCheck(UnitVisitor *parent, bool metadata) + : parent_(parent), metadata_(metadata) {} + void VisitExpr_(const VarNode *v) final { + if (parent_->serial_var_.defined() && v == parent_->serial_var_.get()) { + // The kept coordinate becomes the new outer serial loop + // variable. It is a uniform scalar and is not rewritten by the + // metadata remap, so a direct metadata reference stays legal. + return; + } + if (v == parent_->loop_var_.get()) { + if (metadata_) { + parent_->parent_->Fail( + "expression", + "AttrStmt metadata directly references the loop variable " + "(even if algebraically uniform); the metadata path cannot " + "rewrite it"); + } + return; + } + if (parent_->bind_env_.count(GetRef(v))) { + if (metadata_ && parent_->classified_.count(GetRef(v))) { + Integer cls = parent_->classified_[GetRef(v)]; + if (cls->value != 0) { + std::ostringstream oss; + oss << "AttrStmt metadata references bind `" << v->name_hint + << "` whose value converts to a vector; metadata must " + "stay scalar"; + parent_->parent_->Fail("expression", oss.str()); + } + } + return; + } + if (parent_->external_defs_.count(v) != 0) { + return; + } + std::ostringstream oss; + oss << "variable `" << v->name_hint + << "` is used but never defined in this kernel"; + parent_->parent_->Fail("expression", oss.str()); + } + + private: + UnitVisitor *parent_; + bool metadata_; + }; + + void CheckDefinedVars(const PrimExpr &expr) { + RefCheck checker(this, /*metadata=*/false); + checker(expr); + } + void CheckDefinedVars(const Array &indices) { + RefCheck checker(this, /*metadata=*/false); + for (const auto &idx : indices) { + checker(idx); + } + } + + // One admission entry for every compute-value position (Bind + // RHS, Evaluate, store value): whitelist, cast plans, operation + // types, reference check and lane classification. Before this only + // the Bind path ran the lane classification, so a value the + // whitelist lets through but that is not a legal scalar computation + // (e.g. an external vector Var) reached Vectorize through Evaluate or + // a store. `what` is the diagnostic prefix for the position. + pto_analysis::LaneUse CheckComputeValue(const PrimExpr &value, + const std::string &what) { + if (auto reason = pto_analysis::FindUnsupportedExprNode(value)) { + parent_->Fail("expression", what + ": " + reason.value()); + } + CheckValueCasts(value); + CheckBinaryOpTypes(value); + CheckDefinedVars(value); + pto_analysis::LaneUseContext ctx = MakeLaneUseContext(); + std::string why; + pto_analysis::LaneUse use = + pto_analysis::AnalyzeLaneUse(value, &ctx, &why); + if (use == pto_analysis::LaneUse::kUnknown) { + std::ostringstream oss; + oss << what << " cannot be classified: " << why; + parent_->Fail("expression", oss.str()); + } + return use; + } + + // Per-operation type rules driven by the operands' + // lane classification. Uniform+Uniform stays scalar (only the base + // dtype and PTO scalar-support are required); any Varying operand + // makes the node a vector operation and the VMI op matrix applies. + // Inner sub-expressions are classified individually, so an int8 + // scalar Mul inside a Varying Add is legal while an int8 vector Mul + // stays rejected. + void CheckBinaryOpTypes(const PrimExpr &expr) { + std::function check = [&](const PrimExpr &e) { + if (const auto *add = e.as()) { + check(add->a); + check(add->b); + CheckOneOp("vadd", add->dtype, add->a, add->b); + return; + } + if (const auto *sub = e.as()) { + check(sub->a); + check(sub->b); + CheckOneOp("vsub", sub->dtype, sub->a, sub->b); + return; + } + if (const auto *mul = e.as()) { + check(mul->a); + check(mul->b); + CheckOneOp("vmul", mul->dtype, mul->a, mul->b); + return; + } + if (const auto *div = e.as()) { + check(div->a); + check(div->b); + CheckOneOp("vdiv", div->dtype, div->a, div->b); + return; + } + if (const auto *cast = e.as()) { + check(cast->value); + return; + } + if (const auto *call = e.as()) { + if (auto call_op = call->op.as()) { + const std::string &name = call_op.value()->name; + if (name == "tirx.fabs" || name == "tirx.exp" || + name == "tirx.log" || name == "tirx.sqrt") { + check(call->args[0]); + CheckOneUnaryOp(name.c_str(), call->dtype, call->args[0]); + return; + } + } + } + if (const auto *load = e.as()) { + // Index arithmetic is address analysis, not value arithmetic. + (void)load; + return; + } + // Vars/constants carry no operation to check. + }; + check(expr); + } + + void CheckOneUnaryOp(const char *op_name, DataType dtype, + const PrimExpr &arg) { + pto_analysis::LaneUseContext ctx = MakeLaneUseContext(); + std::string why; + pto_analysis::LaneUse a = pto_analysis::AnalyzeLaneUse(arg, &ctx, &why); + if (a == pto_analysis::LaneUse::kUnknown) { + return; + } + const bool vector_op = a == pto_analysis::LaneUse::kVarying; + if (pto_analysis::IsSupportedFloat8(dtype)) { + std::ostringstream oss; + oss << op_name << " on " + << (vector_op ? "lane-varying " : "lane-uniform ") << dtype + << " is not supported (FP8 arithmetic is out of scope; convert " + "to float32 explicitly first)"; + parent_->Fail("dtype", oss.str()); + } + if (vector_op) { + if (!pto_analysis::IsSupportedVectorUnaryOp(op_name, dtype)) { + std::ostringstream oss; + oss << op_name << " on " << dtype + << " elements is not supported by the VMI operation matrix"; + parent_->Fail("dtype", oss.str()); + } + return; + } + if (pto_analysis::IsFloatFamily(dtype) && + pto_analysis::ContainsRuntimeValue(arg)) { + std::ostringstream oss; + oss << op_name << " on " << dtype + << " is a runtime lane-uniform float computation; it stays " + "scalar inside SIMD_VF (no scalar float ALU by design). " + "Compute it outside the SimdVF region, or wait for the " + "gather/scatter milestone for vector write-back"; + parent_->Fail("dtype", oss.str()); + } + } + + void CheckOneOp(const char *op_name, DataType dtype, const PrimExpr &lhs, + const PrimExpr &rhs) { + pto_analysis::LaneUseContext ctx = MakeLaneUseContext(); + std::string why; + pto_analysis::LaneUse a = pto_analysis::AnalyzeLaneUse(lhs, &ctx, &why); + pto_analysis::LaneUse b = pto_analysis::AnalyzeLaneUse(rhs, &ctx, &why); + if (a == pto_analysis::LaneUse::kUnknown || + b == pto_analysis::LaneUse::kUnknown) { + // The operand's access analysis failed (stride/pattern/etc.); the + // dedicated BufferLoad visitor reports that with the correct + // category later in this same unit walk. Do not double-report it + // as a dtype error here. + return; + } + bool vector_op = (a == pto_analysis::LaneUse::kVarying || + b == pto_analysis::LaneUse::kVarying); + // FP8 arithmetic is rejected in both forms, and deliberately before the + // generic matrix: the conversion matrix admits FP8 operands, so "not in + // the VMI op matrix" alone would be an opaque reason. No implicit + // FP8 -> FP32 promotion is inserted either. + if (pto_analysis::IsSupportedFloat8(dtype)) { + std::ostringstream oss; + oss << op_name << " on " + << (vector_op ? "lane-varying " : "lane-uniform ") << dtype + << " is not supported (FP8 arithmetic is not supported; convert " + "to float32 explicitly first)"; + parent_->Fail("dtype", oss.str()); + } + if (vector_op) { + // The node emits a VMI op: apply the operation matrix. + if (!pto_analysis::IsSupportedVectorBinaryOp(op_name, dtype)) { + std::ostringstream oss; + oss << op_name << " on " << dtype + << " elements is not supported by the VMI operation matrix"; + parent_->Fail("dtype", oss.str()); + } + return; + } + // Lane-uniform stays scalar by design, so a runtime-valued float + // computation here runs on the scalar float ALU inside SIMD_VF. + // Compile-time constants fold and integer scalars have no such + // premise, so only float-family runtime values are rejected. + if (pto_analysis::IsFloatFamily(dtype) && + (pto_analysis::ContainsRuntimeValue(lhs) || + pto_analysis::ContainsRuntimeValue(rhs))) { + std::ostringstream oss; + oss << op_name << " on " << dtype + << " is a runtime lane-uniform float computation; it stays " + "scalar inside SIMD_VF (no scalar float ALU by design). " + "Compute it outside the SimdVF region, or wait for the " + "gather/scatter milestone for vector write-back"; + parent_->Fail("dtype", oss.str()); + } + // Pure scalar op: base dtype legality only (int8 scalar mul is + // fine; int24/FP64 still rejected by the base dtype rule). + if (!pto_analysis::IsSupportedElementDType(dtype)) { + std::ostringstream oss; + oss << "scalar " << op_name << " on " << dtype + << " is not supported by the base element dtype rule"; + parent_->Fail("dtype", oss.str()); + } + } + + public: + // Unified value-Cast check: called for Bind, Evaluate + // and store values alike. A Cast is only accepted when PlanValueCast + // proves a VMI-equivalent conversion (attributes included); Cast + // nodes inside BufferLoad *indices* are index arithmetic handled by + // address analysis and are not checked here. + void CheckValueCasts(const PrimExpr &expr) { + class CastCheck : public StmtExprVisitor { + public: + CastCheck(UnitVisitor *parent) : parent_(parent) {} + void VisitExpr_(const CastNode *op) final { + std::string why; + auto plan = pto_analysis::PlanValueCast(op, &why); + if (!plan.has_value()) { + std::ostringstream oss; + oss << "Cast " << op->value.dtype() << " -> " << op->dtype + << " has no proven VMI-equivalent conversion"; + if (!why.empty()) { + oss << ": " << why; + } + parent_->parent_->Fail("dtype", oss.str()); + } + // A lane-uniform Cast is emitted as a scalar Cast, while the + // conversion plan above proves the VMI-equivalent *vector* form. + // FP8 has no scalar form at all, so a lane-uniform FP8 Cast has to + // be read per lane; the non-FP8 pairs keep whatever lowering the + // codegen gives them. + if (plan.value().identity) { + VisitExpr(op->value); // annotations are not computation values + return; + } + if (!pto_analysis::IsSupportedFloat8(op->value.dtype()) && + !pto_analysis::IsSupportedFloat8(op->dtype)) { + VisitExpr(op->value); // annotations are not computation values + return; + } + pto_analysis::LaneUseContext ctx = parent_->MakeLaneUseContext(); + std::string lane_why; + pto_analysis::LaneUse use = + pto_analysis::AnalyzeLaneUse(op->value, &ctx, &lane_why); + if (use != pto_analysis::LaneUse::kVarying) { + std::ostringstream oss; + oss << "Cast " << op->value.dtype() << " -> " << op->dtype + << " stays lane-uniform"; + if (use == pto_analysis::LaneUse::kUnknown) { + oss << " (lane use cannot be classified: " << lane_why << ")"; + } + oss << "; FP8 conversions are only supported as a vector vcvt " + "over the full lane range"; + parent_->parent_->Fail("dtype", oss.str()); + } + VisitExpr(op->value); // annotations are not computation values + } + void VisitExpr_(const BufferLoadNode *) final { + // Index casts belong to address analysis; do not descend. + } + + private: + UnitVisitor *parent_; + }; + CastCheck checker(this); + checker(expr); + } + + void VerifyAccess(const Buffer &buffer, const Array &indices, + bool is_write, const Optional &predicate) { + // Predicated accesses have no supported lowering; only + // absent or provably-true predicates are accepted. + if (predicate.defined()) { + // The predicate is an expression position too — a + // provably-true predicate (e.g. `True || ghost > 0`) could hide + // undefined vars that Vectorize then silently drops with the + // predicate. Run the same reference check as every other + // position before the true-check. + CheckDefinedVars(predicate.value()); + arith::Analyzer probe; + PrimExpr simplified = probe.Simplify(predicate.value()); + if (!is_one(simplified)) { + std::ostringstream oss; + oss << "buffer `" << buffer->name + << "` has a non-trivial access predicate; predicated " + "load/store is not yet supported"; + parent_->Fail("memory access", oss.str()); + } + } + try { + auto access = pto_analysis::AnalyzeBufferAccess( + buffer, indices, is_write, loop_var_, mapping_, &analyzer_, + bind_env_); + if (access.pattern != pto_analysis::AccessPattern::kContinuous && + access.pattern != pto_analysis::AccessPattern::kLaneUniform) { + std::ostringstream oss; + oss << "buffer `" << buffer->name << "` access pattern " + << pto_analysis::AccessPatternToString(access.pattern) + << " is not supported"; + parent_->Fail("memory access", oss.str()); + } + if (access.pattern == pto_analysis::AccessPattern::kLaneUniform && + is_write) { + // Lane-uniform is only a load pattern. + std::ostringstream oss; + oss << "buffer `" << buffer->name + << "`: lane-uniform stores are not yet supported"; + parent_->Fail("memory access", oss.str()); + } + } catch (const pto_analysis::PtoAnalysisError &err) { + parent_->Fail("memory access", err.what()); + } + } + + VerifyParallelToPTOImpl *parent_; + const Var &loop_var_; + /*! 2D only — the non-vectorized coordinate (uniform). */ + Var serial_var_; + const pto_analysis::PtoInverseMapping &mapping_; + arith::Analyzer analyzer_; + Map bind_env_; + // Classification recorded at each Bind definition + // (0 = uniform, 1 = varying) for this unit. + Map classified_; + // The lexical definition snapshot at the unit's position — + // params, buffer data pointers, and every definition (outer loop + // vars, thread_extent bindings, binds, allocations) that is in scope + // where the unit appears. + std::set external_defs_; + }; + + // ------------------------------------------------------------------ + // Scope tracking: every definition — outer serial loop vars, + // thread_extent bindings, flat binds, buffer allocations/declarations + // — enters the environment at its lexical position and leaves it when + // the enclosing scope exits. A unit is verified against a snapshot of + // this environment at the unit's position, so a sibling scope's + // definitions can no longer leak in and legitimately-defined outer + // scalars (blockIdx, outer binds) are no longer rejected. + // ------------------------------------------------------------------ + void VisitStmt_(const ForNode *op) final { + if (op->kind != ForKind::kParallel) { + std::set saved = defs_; + defs_.insert(op->loop_var.get()); + StmtExprVisitor::VisitStmt_(op); + defs_ = std::move(saved); + return; + } + // A parallel For outside a converting region is not this pass's + // concern; its iteration var has no ordered single-assignment + // definition to track here. + StmtExprVisitor::VisitStmt_(op); + } + + void VisitStmt_(const AttrStmtNode *op) final { + if (op->attr_key == tirx::attr::thread_extent) { + if (const auto *iv = op->node.as()) { + // blockIdx/cthread bindings (e.g. `with T.Kernel(2) as bx`) are + // legitimately-defined outer scalars for the attr body. The + // no-threadIdx-inside-the-converting-region rule is + // unaffected: it is enforced by the RegionVisitor inside the + // region, not here. + std::set saved = defs_; + defs_.insert(iv->var.get()); + StmtExprVisitor::VisitStmt_(op); + defs_ = std::move(saved); + return; + } + } + StmtExprVisitor::VisitStmt_(op); + } + + void VisitStmt_(const BindNode *op) final { + // Flat sequence semantics: the bound var is visible to all + // subsequent statements of the enclosing scope; the scope owners + // above restore the environment on exit. + defs_.insert(op->var.get()); + StmtExprVisitor::VisitStmt_(op); + } + + void VisitStmt_(const AllocBufferNode *op) final { + // Flat statement (no body): the allocation is visible to the rest of + // the enclosing scope. + defs_.insert(op->buffer->data.get()); + StmtExprVisitor::VisitStmt_(op); + } + + void VisitStmt_(const DeclBufferNode *op) final { + defs_.insert(op->buffer->data.get()); + StmtExprVisitor::VisitStmt_(op); + } + + void VisitStmt_(const IfThenElseNode *op) final { + // Branch bodies are independent scopes. + StmtExprVisitor::VisitExpr(op->condition); + { + std::set saved = defs_; + VisitStmt(op->then_case); + defs_ = std::move(saved); + } + if (op->else_case.defined()) { + std::set saved = defs_; + VisitStmt(op->else_case.value()); + defs_ = std::move(saved); + } + } + + void VisitStmt_(const WhileNode *op) final { + StmtExprVisitor::VisitExpr(op->condition); + std::set saved = defs_; + StmtExprVisitor::VisitStmt(op->body); + defs_ = std::move(saved); + } + + void VisitStmt_(const SBlockNode *op) final { + if (op->name_hint != "SIMD_VF") { + std::set saved = defs_; + for (const auto &iv : op->iter_vars) { + defs_.insert(iv->var.get()); + } + for (const auto &buf : op->alloc_buffers) { + defs_.insert(buf->data.get()); + } + StmtExprVisitor::VisitStmt_(op); + defs_ = std::move(saved); + return; + } + HandleSimdVF(op); + } + + // ------------------------------------------------------------------ + // SIMD_VF region entry (converting regions only) + // ------------------------------------------------------------------ + void HandleSimdVF(const SBlockNode *op) { + // Block-scoped definitions (allocations, iter vars) are in scope for + // the whole region (supersedes the earlier simd_vf_allocs_ side + // table). + std::set saved = defs_; + for (const auto &iv : op->iter_vars) { + defs_.insert(iv->var.get()); + } + for (const auto &buf : op->alloc_buffers) { + defs_.insert(buf->data.get()); + } + + // Shared region scan: same implementation as Vectorize. + const Stmt ®ion = op->body; + pto_analysis::PtoRegionScan scan = pto_analysis::ScanPtoRegion(region); + if (!scan.has_any_parallel) { + // Pure hand-written VMI region (possibly with control flow): skipped + // entirely; no lanes annotation required. + defs_ = std::move(saved); + return; + } + // Number the converting region for diagnostics. The lane count is + // reset to unknown first: a failure before this region's own lanes + // are read must not inherit the previous region's value. + current_region_ = region_seq_++; + current_lanes_ = 0; + SetRegionContext(); + if (scan.has_unsupported_wrapper) { + Fail("loop structure", scan.issue); + } + const std::vector &units = scan.units; + + // Only now is the lanes annotation required. + int64_t lanes = ReadLanes(op); + current_lanes_ = lanes; + SetRegionContext(); // refresh the region context with the lane count + + // Region-level checks across the whole converting region. + RegionVisitor region_visitor(this); + region_visitor(region); + + // Per-unit checks in order (multiple sequential units OK) against the + // lexical definition snapshot at each unit's position. + size_t next_unit = 0; + WalkRegion(region, units, lanes, &next_unit); + ICHECK(next_unit == units.size()) + << "[VerifyParallelToPTO] internal error: region walk and unit " + "scan disagree on the unit count"; + defs_ = std::move(saved); + } + + // ------------------------------------------------------------------ + // Per-unit checks (structure, ranges, layout) + // ------------------------------------------------------------------ + void VerifyUnit(const For &unit, int64_t lanes, size_t unit_index) { + // Full per-unit context for diagnostics. + { + std::ostringstream oss; + oss << "region #" << current_region_ << ", unit #" << unit_index + << " (loop var `" << unit->loop_var->name_hint << "`, extent " + << unit->extent << ", lanes " << lanes << ")"; + unit_context_ = oss.str(); + } + if (!is_zero(unit->min)) { + std::ostringstream oss; + oss << "loop min must be 0, got " << unit->min; + Fail("loop range", oss.str()); + } + // The lane layout assumes unit stride; a non-unit step + // changes the set of iteration points and has no supported + // lowering. + if (unit->step.defined() && + !is_one(analyzer_.Simplify(unit->step.value()))) { + std::ostringstream oss; + oss << "T.Parallel step must be 1, got " << unit->step; + Fail("loop range", oss.str()); + } + const int64_t *extent = as_const_int(unit->extent); + if (extent == nullptr || *extent <= 0) { + std::ostringstream oss; + oss << "loop extent must be a positive compile-time constant, got " + << unit->extent; + Fail("loop range", oss.str()); + } + if (!unit->annotations.count(attr::kPtoParallelLoopLayout)) { + Fail("layout", "pto_parallel_loop_layout is missing on the " + "outermost T.Parallel"); + } + + // Identify the full loop nest and check every member loop + // *before* the inversion helper (which relies on constant extents) so + // an illegal inner loop gets its loop-range diagnostic instead of a + // raw assertion. + const ForNode *inner_probe = nullptr; + if (const auto *body_for = unit->body.as()) { + if (body_for->kind == ForKind::kParallel) { + inner_probe = body_for; + } + } + if (inner_probe != nullptr) { + if (!is_zero(inner_probe->min)) { + std::ostringstream oss; + oss << "inner loop min must be 0, got " << inner_probe->min; + Fail("loop range", oss.str()); + } + if (inner_probe->step.defined() && + !is_one(analyzer_.Simplify(inner_probe->step.value()))) { + std::ostringstream oss; + oss << "inner T.Parallel step must be 1, got " << inner_probe->step; + Fail("loop range", oss.str()); + } + if (as_const_int(inner_probe->extent) == nullptr) { + std::ostringstream oss; + oss << "inner loop extent must be a compile-time constant, got " + << inner_probe->extent; + Fail("loop range", oss.str()); + } else if (*as_const_int(inner_probe->extent) <= 0) { + std::ostringstream oss; + oss << "inner loop extent must be a positive compile-time constant, " + "got " + << inner_probe->extent; + Fail("loop range", oss.str()); + } + if (inner_probe->loop_var.dtype() != DataType::Int(32)) { + std::ostringstream oss; + oss << "inner loop variable must be int32 (the dtype " + "LayoutInference produces), got " + << inner_probe->loop_var.dtype(); + Fail("loop range", oss.str()); + } + } + + pto_analysis::PtoInverseMapping mapping = [&] { + try { + return pto_analysis::InvertPtoLaneLayout(unit); + } catch (const pto_analysis::PtoAnalysisError &err) { + Fail("layout", err.what()); + return pto_analysis::PtoInverseMapping{}; + } + }(); + if (mapping.lanes != lanes) { + std::ostringstream oss; + oss << "layout lanes (" << mapping.lanes + << ") disagree with tl.simdvf_lanes (" << lanes << ")"; + Fail("lanes", oss.str()); + } + + // A 2D unit is the outermost For of a two-layer adjacent + // Parallel nest; the unit visitor runs on the innermost compute body + // (the layout annotation only sits on the outer For). The + // inner loop's range contract was already enforced by the + // pre-inversion block above (min/step/extent/dtype), so no per-loop + // check is repeated here. + const ForNode *inner = nullptr; + if (const auto *body_for = unit->body.as()) { + if (body_for->kind == ForKind::kParallel) { + inner = body_for; + } + } + + // The vectorized coordinate is the selected dimension of a 2D mapping + // and the single loop var in 1D (the other 2D coordinate stays a + // serial loop var and classifies as uniform). + const Var vector_var = + mapping.is_2d + ? (mapping.select_inner ? mapping.inner_var : mapping.outer_var) + : unit->loop_var; + // The definition snapshot at the unit's lexical position. + const Stmt &body_to_verify = inner != nullptr ? inner->body : unit->body; + UnitVisitor unit_visitor(this, vector_var, mapping, defs_); + unit_visitor(body_to_verify); + } + + // Region-level statements (between/around the Parallel units) + // previously only *registered* definitions without scanning their + // expressions, so an undefined var in a region-level Bind RHS (or + // AttrStmt metadata, or serial-For bounds) leaked into the output + // unchecked. Check that every Var is defined in the current lexical + // environment; call sites run this *before* registering new + // definitions so lexical order is preserved (a self reference is + // undefined too). + void CheckRegionExprDefined(const PrimExpr &expr, const char *what) { + PostOrderVisit(expr, [&](const ObjectRef &node) { + if (const auto *v = node.as()) { + if (defs_.count(v) == 0) { + std::ostringstream oss; + oss << what << " references variable `" << v->name_hint + << "` which is not defined in the lexical scope of this " + "SIMD_VF region"; + Fail("expression", oss.str()); + } + } + }); + } + + // Walk the converting region in statement order, mirroring the wrapper + // kinds DiscoverUnits descends (For/SeqStmt/AttrStmt/SBlock/ + // SBlockRealize). Definitions before a unit (flat binds, serial loop + // vars, allocations) are registered as encountered so each unit is + // verified against the environment at its position. Parallel Fors are + // the discovered units themselves; IfThenElse/While cannot contain + // units (the scan rejected that), and unit bodies belong to the + // UnitVisitor. + void WalkRegion(const Stmt &stmt, const std::vector &units, + int64_t lanes, size_t *next_unit) { + if (const auto *for_node = stmt.as()) { + if (for_node->kind == ForKind::kParallel) { + ICHECK(*next_unit < units.size() && units[*next_unit].get() == for_node) + << "[VerifyParallelToPTO] internal error: region walk and " + "unit scan disagree on unit order"; + VerifyUnit(units[*next_unit], lanes, *next_unit); + ++*next_unit; + return; + } + // The serial loop bounds are expression positions too. + // The step is one as well (the base statement visitors do + // not descend into For::step). + SetRegionContext(); + CheckRegionExprDefined(for_node->min, "serial loop min"); + CheckRegionExprDefined(for_node->extent, "serial loop extent"); + if (for_node->step.defined()) { + CheckRegionExprDefined(for_node->step.value(), "serial loop step"); + } + std::set saved = defs_; + defs_.insert(for_node->loop_var.get()); + WalkRegion(for_node->body, units, lanes, next_unit); + defs_ = std::move(saved); + return; + } + if (const auto *assertion = stmt.as()) { + // A region-level assertion is an expression position: its condition + // must reference defined values. The node has no body (only the + // condition, error kind and message parts), so nothing else is + // walked; an assertion *inside* a unit stays rejected by the unit + // statement visitor. + SetRegionContext(); + CheckRegionExprDefined(assertion->condition, + "region-level assert condition"); + return; + } + if (const auto *seq = stmt.as()) { + for (const auto &s : seq->seq) { + WalkRegion(s, units, lanes, next_unit); + } + return; + } + if (const auto *attr = stmt.as()) { + // Region-level metadata (node/value) is scanned for + // undefined references before descending. + SetRegionContext(); + if (auto ne = attr->node.try_cast()) { + if (const auto *pe = ne.value().as()) { + CheckRegionExprDefined(GetRef(pe), + "region-level AttrStmt metadata"); + } + } + CheckRegionExprDefined(attr->value, "region-level AttrStmt metadata"); + // The body is a nested statement scope. + { + std::set saved = defs_; + WalkRegion(attr->body, units, lanes, next_unit); + defs_ = std::move(saved); + } + return; + } + if (const auto *block = stmt.as()) { + std::set saved = defs_; + for (const auto &iv : block->iter_vars) { + defs_.insert(iv->var.get()); + } + for (const auto &buf : block->alloc_buffers) { + defs_.insert(buf->data.get()); + } + WalkRegion(block->body, units, lanes, next_unit); + defs_ = std::move(saved); + return; + } + if (const auto *realize = stmt.as()) { + std::set saved = defs_; + for (const auto &iv : realize->block->iter_vars) { + defs_.insert(iv->var.get()); + } + for (const auto &buf : realize->block->alloc_buffers) { + defs_.insert(buf->data.get()); + } + WalkRegion(realize->block->body, units, lanes, next_unit); + defs_ = std::move(saved); + return; + } + if (const auto *bind = stmt.as()) { + // The region-level Bind RHS is checked against the + // pre-registration environment (lexical order), then the + // definition itself is registered. + SetRegionContext(); + CheckRegionExprDefined(bind->value, "region-level bind"); + defs_.insert(bind->var.get()); + return; + } + if (const auto *alloc = stmt.as()) { + defs_.insert(alloc->buffer->data.get()); + return; + } + if (const auto *decl = stmt.as()) { + defs_.insert(decl->buffer->data.get()); + return; + } + // Statements that contribute no definitions are still + // expression positions. A store, an Evaluate, a branch or a while + // outside the Parallel units previously ended the walk silently, so + // an undefined variable in any of their expression fields passed + // Verify and stayed in the output. + if (const auto *store = stmt.as()) { + SetRegionContext(); + CheckRegionExprDefined(store->value, "region-level store value"); + for (const auto &idx : store->indices) { + CheckRegionExprDefined(idx, "region-level store index"); + } + if (store->predicate.defined()) { + CheckRegionExprDefined(store->predicate.value(), + "region-level store predicate"); + } + return; + } + if (const auto *eval = stmt.as()) { + SetRegionContext(); + CheckRegionExprDefined(eval->value, "region-level evaluate value"); + return; + } + if (const auto *ite = stmt.as()) { + // A Parallel under control flow is rejected by ScanPtoRegion, so + // the branches carry plain statements only; branch-local + // definitions must not leak to later siblings. + SetRegionContext(); + CheckRegionExprDefined(ite->condition, "region-level branch condition"); + { + std::set saved = defs_; + WalkRegion(ite->then_case, units, lanes, next_unit); + defs_ = std::move(saved); + } + if (ite->else_case.defined()) { + std::set saved = defs_; + WalkRegion(ite->else_case.value(), units, lanes, next_unit); + defs_ = std::move(saved); + } + return; + } + if (const auto *while_node = stmt.as()) { + SetRegionContext(); + CheckRegionExprDefined(while_node->condition, + "region-level while condition"); + std::set saved = defs_; + WalkRegion(while_node->body, units, lanes, next_unit); + defs_ = std::move(saved); + return; + } + } + + arith::Analyzer analyzer_; + // The lexical definition environment. Supersedes the earlier + // function-wide external-def union, which missed scoped definitions + // (blockIdx, outer/VF-internal binds) and leaked sibling-scope loop + // vars into every unit. + std::set defs_; + // Data Vars of AllocBuffers annotated tirx.volatile (collected + // function-wide by CollectVolatileAllocs; identity-matched at every + // access inside a converting region). + std::set volatile_allocs_; + // Position of the converting region/unit being checked, for + // diagnostics. + int region_seq_ = 0; + int current_region_ = -1; + std::string unit_context_; + // Lane count of the converting region under check (0 when unknown); part of + // the region/unit diagnostic context. + int64_t current_lanes_ = 0; +}; + +} // namespace + +namespace transform { + +tvm::transform::Pass VerifyParallelToPTO() { + auto pass_func = [](PrimFunc f, const IRModule &m, + const tvm::transform::PassContext &ctx) { + auto target = f->GetAttr(tvm::attr::kTarget); + if (!target.defined() || !TargetIsPTO(target.value())) { + return f; + } + VerifyParallelToPTOImpl verifier; + verifier.Verify(f); + return f; + }; + return tvm::tirx::transform::CreatePrimFuncPass(pass_func, 0, + "tl.VerifyParallelToPTO", {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tl.transform.VerifyParallelToPTO", + VerifyParallelToPTO); +} + +} // namespace transform + +} // namespace tl + +} // namespace tvm diff --git a/src/layout/layout.h b/src/layout/layout.h index dabcee8eeb..5a560c0d4a 100644 --- a/src/layout/layout.h +++ b/src/layout/layout.h @@ -465,6 +465,10 @@ namespace attr { constexpr const char *kLayoutMap = "layout_map"; // ForAttr, Containing the parallel loop layout for a parallel for loop constexpr const char *kParallelLoopLayout = "parallel_loop_layout"; +// ForAttr, PTO SIMD lane layout for a parallel for loop inside a PTO +// SIMD_VF block. Distinct from kParallelLoopLayout so SIMT lowering never +// misreads a PTO lane fragment as a thread layout. +constexpr const char *kPtoParallelLoopLayout = "pto_parallel_loop_layout"; // ForAttr, Containing the predicate for a parallel for loop constexpr const char *kParallelLoopPredicate = "parallel_loop_predicate"; // ForAttr, Marks a ragged SIMT loop layout that needs guarded inverse lowering diff --git a/src/op/operator.h b/src/op/operator.h index aed41be715..f9ed2bf4bd 100644 --- a/src/op/operator.h +++ b/src/op/operator.h @@ -176,6 +176,11 @@ struct LayoutInferArgs { // the vectorizer's width; scalar reducer-root attempts pass one. This is // a search option, not operator state or a constraint on inferred layouts. int candidate_vector_size_limit = 0; + // Whether the current TileOp sits inside a PTO SIMD_VF block whose parallel + // loops are vectorized over logical lanes instead of SIMT threads. When + // true, thread_bounds is the [0, L) lane range from tl.simdvf_lanes and + // ParallelOp must build its lane Fragment with vectorize_size=1. + bool in_pto_simd_vf = false; }; class TileOperator; diff --git a/src/op/parallel.cc b/src/op/parallel.cc index b747bc87b4..f6819197d0 100644 --- a/src/op/parallel.cc +++ b/src/op/parallel.cc @@ -10,8 +10,12 @@ #include #include +#include +#include +#include #include #include +#include #include #include "../layout/layout.h" @@ -21,6 +25,7 @@ #include "arith/int_operator.h" #include "backend/common/target_utils.h" #include "builtin.h" +#include "pto_index_analysis.h" #include "reducer.h" #include "span_utils.h" #include "utils.h" @@ -354,6 +359,15 @@ LayoutMap ParallelOpNode::InferLayout(const LayoutInferArgs &layout_args, layout_args.bind_var_to_expr); } + // PTO SIMD path: parallel loops inside a PTO SIMD_VF block distribute + // over [0, L) logical lanes. Build the lane fragment directly + // (vectorize_size = 1, no SIMT vector-width selection) and skip the + // fragment-buffer inference machinery — only shared/global buffers with + // continuous access are supported, checked later by VerifyParallelToPTO. + if (layout_args.in_pto_simd_vf) { + return InferPtoSimdLayout(layout_args); + } + if (level == InferLevel::kStrict) { LayoutMap results; // Deduce buffers that should be complicated replicated. @@ -996,6 +1010,398 @@ Fragment ParallelOpNode::ComputeLoopLayoutFromBuffer( return result; } +// Pick the vectorized dimension of a two-layer +// Parallel nest from the buffer accesses. Both dimensions are tried as +// layout candidates under the same lane/chunk proof rules the consumers +// (Verify/Vectorize) use; a candidate is valid when its extent divides +// L and every access is continuous (or, for loads, lane-uniform) under +// it. The inner dimension wins when both are valid; no valid candidate +// is a layout error naming the failing access. +namespace { + +class PtoAccessCollector : public tirx::StmtExprVisitor { +public: + struct Access { + Buffer buffer; + Array indices; + bool is_write; + /*! The definition environment at this access's program point + * (binds in scope, mirrors the consumers' lexical rules). */ + Map bind_env; + }; + std::vector accesses; + + void VisitStmt_(const BufferStoreNode *op) final { + accesses.push_back({op->buffer, op->indices, /*is_write=*/true, bind_env_}); + tirx::StmtExprVisitor::VisitStmt_(op); + } + void VisitExpr_(const BufferLoadNode *op) final { + accesses.push_back( + {op->buffer, op->indices, /*is_write=*/false, bind_env_}); + tirx::StmtExprVisitor::VisitExpr_(op); + } + void VisitStmt_(const BindNode *op) final { + // Flat sequence semantics (same as the consumers): the value is + // visited *before* the var is registered, so a self reference is not + // in scope; every bind — including load-valued ones — is registered + // so ResolveIndexAliases can apply its loaded-value identity rule + // along alias chains. + tirx::StmtExprVisitor::VisitStmt_(op); + bind_env_.Set(op->var, op->value); + } + void VisitStmt_(const AttrStmtNode *op) final { + // Nested statement scope (mirror of the consumers' rule): + // binds defined inside an AttrStmt body must not leak to later + // siblings, while outer binds stay visible inside. + Map saved = bind_env_; + tirx::StmtExprVisitor::VisitStmt_(op); + bind_env_ = std::move(saved); + } + +private: + Map bind_env_; +}; + +// 2D layout candidate. Both +// dimensions are tried with their canonical mapping — selected coordinate +// `q * L + lane`, the other coordinate left as-is — over the bound +// lane/chunk/kept ranges. Alias resolution, the loaded-value identity +// rule, the address-node whitelist and the resource limits are the shared +// read-only helpers in pto_index_analysis.h — the same implementation +// AnalyzeBufferAccess uses, so planning and consumption cannot diverge. +// An access passes when its flattened offset is continuous +// (`offset == start + lane`); a load may instead be lane-uniform +// (`offset == start`), while a store must be continuous (lane-uniform +// stores are rejected downstream). A candidate is valid once every +// access proof passes over the padded chunk domain (`q in [0, Q)`, +// `lane in [0, L)`); when both candidates are valid, the selection rules +// below (non-degenerate-dimension preference, then the inner preference) +// apply. +// Layout-inference code reports construct-time failures as +// LoopLayoutInjectiveException (the type LayoutInference catches to discard +// an attempt), while the shared padding helper speaks PtoAnalysisError; +// translate here so all three call sites share one helper. +std::pair CheckedPaddingForInference(int64_t extent, + int64_t lanes) { + try { + return pto::CheckedPtoPadding(extent, lanes); + } catch (const pto::PtoAnalysisError &err) { + throw LoopLayoutInjectiveException(err.what()); + } +} + +int PtoRequiredVectorDim(const For &root, const Var &iv, const Var &jv, + int64_t M, int64_t N, int64_t lanes, + std::string *why) { + PtoAccessCollector collector; + collector(root->body); + + auto candidate_ok = [&](bool select_inner) -> bool { + const int64_t sel_extent = select_inner ? N : M; + const Var &sel_var = select_inner ? jv : iv; + const Var &kept_var = select_inner ? iv : jv; + const int64_t kept_extent = select_inner ? M : N; + // A non-divisible extent is padded to P = ceil(E/L)*L when the + // fragment is built; the candidate proof below works over the padded + // domain (Q = P/L chunks), which includes the tail chunk's full lane + // range, and the analyzer proves continuity over that whole padded + // domain. What remains a caller guarantee is only the *readability* + // of the padded range past E on the selected axis (tail loads read + // it); the candidate neither proves nor rejects that. + // Shared checked padding (rejects extents above the int32 index space + // before the addition; the same helper the fragment construction uses). + const auto [pad, padded_extent] = + CheckedPaddingForInference(sel_extent, lanes); + (void)pad; + const int64_t chunk_count = padded_extent / lanes; + DataType dt = sel_var.dtype(); + Var lane("probe_lane", dt); + Var q("probe_q", dt); + PrimExpr sel_expr = q * IntImm(dt, lanes) + lane; + for (const auto &access : collector.accesses) { + // Resolve Bind aliases per index with the shared resolver (cycle / + // depth / loaded-value failure, 4096-node budget, error conversion) + // against this access's program-point environment, then apply the + // shared address-node whitelist — the same order AnalyzeBufferAccess + // uses. + Array resolved; + for (const auto &idx : access.indices) { + Optional opt; + try { + opt = pto::ResolveIndexAliases(idx, access.bind_env); + } catch (const pto::PtoAnalysisError &err) { + *why = err.what(); + return false; + } + if (!opt.has_value()) { + std::ostringstream oss; + oss << "buffer `" << access.buffer->name + << "`: bind alias chain in index is too deep, cyclic, or " + "depends on a loaded value (data-dependent addresses are " + "not yet supported)"; + *why = oss.str(); + return false; + } + resolved.push_back(opt.value()); + } + for (const auto &idx : resolved) { + if (!pto::IsSupportedAddressExpr(idx)) { + std::ostringstream oss; + oss << "buffer `" << access.buffer->name + << "`: address expression contains a loaded value or an " + "unsupported call (loaded values and opaque/extern calls " + "cannot appear in addresses; they would be evaluated once " + "per chunk instead of per element)"; + *why = oss.str(); + return false; + } + } + // indices under the candidate mapping, then the flattened element + // offset (same construction as the consumers). + Array indices_after; + for (const auto &idx : resolved) { + indices_after.push_back( + pto::SubstituteVar(idx, sel_var.get(), sel_expr)); + } + Array raw = access.buffer->ElemOffset(indices_after); + if (raw.size() != 1) { + *why = "multi-axis element offsets are not supported"; + return false; + } + PrimExpr off = raw[0]; + Array start_indices; + for (const auto &idx : indices_after) { + start_indices.push_back( + pto::SubstituteVar(idx, lane.get(), make_zero(dt))); + } + PrimExpr start = access.buffer->ElemOffset(start_indices)[0]; + arith::Analyzer probe; + probe.Bind(lane, Range::FromMinExtent(make_zero(dt), IntImm(dt, lanes))); + // In-row chunk range only: q is the per-row chunk index in [0, Q), + // never the global position (K*Q for a kept extent K). The kept + // coordinate is separately bounded below, so the proof domain is + // exactly (q, lane, kept) — the same domain the final consumers' + // proof uses. + probe.Bind(q, + Range::FromMinExtent(make_zero(dt), IntImm(dt, chunk_count))); + // The kept coordinate is bounded by its loop extent; the + // final consumers' proof binds the same fact. + probe.Bind(kept_var, + Range::FromMinExtent(make_zero(dt), IntImm(dt, kept_extent))); + const bool continuous = probe.CanProveEqual(off, start + lane); + const bool uniform = probe.CanProveEqual( + off, + probe.Simplify(pto::SubstituteVar(off, lane.get(), make_zero(dt)))); + if (!continuous && !(uniform && !access.is_write)) { + std::ostringstream oss; + oss << "buffer `" << access.buffer->name << "` (" + << (access.is_write ? "store" : "load") + << ") is neither continuous nor lane-uniform under the " + << (select_inner ? "inner (j)" : "outer (i)") << " mapping"; + *why = oss.str(); + return false; + } + } + return true; + }; + + std::string inner_why; + std::string outer_why; + const bool inner_ok = candidate_ok(/*select_inner=*/true); + if (!inner_ok) { + inner_why = *why; + } + const bool outer_ok = candidate_ok(/*select_inner=*/false); + if (!outer_ok) { + outer_why = *why; + } + if (inner_ok && outer_ok && N == 1 && M > 1) { + // Both dims classify and only the outer is non-degenerate (design + // 4.6.1): selecting the inner (extent-1) dim would emit M chunks with + // one valid lane each, while the outer dim covers the same logical + // points in ceil(M/L) chunks carrying M valid lanes in total. The + // candidate proofs above have already checked the actual buffer + // indices and strides, so the preference only breaks the tie between + // two already-legal candidates and never overrides an access + // constraint. The symmetric M == 1 && N > 1 falls through to the + // inner preference below. + return 0; + } + if (inner_ok) { + // Inner preference; a degenerate dimension is + // only an interpretation when its accesses classify under the padded + // domain, which candidate_ok already enforces. + return 1; + } + if (outer_ok) { + return 0; + } + std::ostringstream oss; + oss << "no vectorizable dimension: inner (j) -> " << inner_why + << "; outer (i) -> " << outer_why; + *why = oss.str(); + return -1; +} + +} // namespace + +// PTO SIMD lane layout: logical points of the parallel loop are distributed +// over [0, L) lanes. The flattened logical index is +// 1D: t = i +// 2D: t = i * N + j (j vectorized) or t = j * M + i (i vectorized) +// with lane = t % L and per-lane chunk position p = t // L; forward_index +// holds p. vectorize_size is fixed at 1 — the per-thread SIMT vector-width +// selection does not apply. +// +// Scope note: this function only *builds* the lane fragment. The full input +// contract (access patterns, dtype, expression whitelist) belongs to +// VerifyParallelToPTO; the checks below are only the minimal conditions +// needed to construct a valid layout, reported as layout-inference errors. +LayoutMap +ParallelOpNode::InferPtoSimdLayout(const LayoutInferArgs &layout_args) const { + const int64_t *lanes_ptr = as_const_int(layout_args.thread_bounds->extent); + ICHECK(lanes_ptr != nullptr && *lanes_ptr > 0) + << "[PTO layout inference] lane count must be a positive compile-time " + "constant, got " + << layout_args.thread_bounds->extent; + int64_t lanes = *lanes_ptr; + + // Loop extents must be compile-time constants. Keep the + // int64 value: narrowing here would silently truncate large static + // extents. + std::vector extents; + for (const auto &iv : loop_vars_) { + const int64_t *ext = as_const_int(iv->dom->extent); + ICHECK(ext != nullptr) + << "[PTO layout inference] requires compile-time constant loop " + "extents, got " + << iv->dom->extent << " for loop " << iv->var->name_hint; + ICHECK(is_zero(iv->dom->min)) + << "[PTO layout inference] requires loop min = 0, got " << iv->dom->min + << " for loop " << iv->var->name_hint; + extents.push_back(*ext); + } + ICHECK(loop_vars_.size() == 1u || loop_vars_.size() == 2u) + << "[PTO layout inference] supports a 1D T.Parallel or a two-layer " + "Parallel nest; got " + << loop_vars_.size() << " layers"; + + // 1D mapping: t = i, lane = i % L, + // forward_index = [i // L]. A non-divisible extent E is padded to + // P = ceil(E/L)*L in the *fragment input space only* — the IterVar extent + // below is the padded one, while buffer addresses keep using the real + // extents/strides (the tail's readable range is a caller precondition). + DataType idx_dtype = loop_vars_[0]->var.dtype(); + PrimExpr flattened; + // The forward vars handed to the Fragment constructor: the selected + // dimension's IterVar is padded up to P so the thread extent L times the + // chunk extent P/L covers the padded input space bijectively; the kept + // dimension keeps its original extent. + Array frag_vars; + if (loop_vars_.size() == 1u) { + const int64_t E = extents[0]; + const auto [pad, P] = CheckedPaddingForInference(E, lanes); + frag_vars = loop_vars_; + if (pad != 0) { + // Reuse the loop var; only the input-space extent is padded. + frag_vars.Set(0, IterVar(Range::FromMinExtent(make_zero(idx_dtype), + IntImm(idx_dtype, P)), + loop_vars_[0]->var, IterVarType::kDataPar)); + } + flattened = InputPlaceholder(0); + } else { + // 2D mapping: the vectorized dimension is + // chosen from the buffer accesses (indices + strides), then the + // logical space is flattened as t = i * P_N + j (j vectorized) or + // t = j * P_M + i (i vectorized), where P_* is the *padded* selected + // extent — the flattened global position p lives in [0, K*Q*L) over the + // padded input space. Both extents stay compile-time constants. + const int64_t M = extents[0]; + const int64_t N = extents[1]; + const Var &i_var = loop_vars_[0]->var; + const Var &j_var = loop_vars_[1]->var; + std::string candidate_why; + int required = + PtoRequiredVectorDim(root_, i_var, j_var, M, N, lanes, &candidate_why); + if (required < 0) { + // Both candidates failed (an access that is neither continuous nor a + // lane-uniform load under the padded domain). Report the concrete + // access and reason — the candidate check already applies the + // inner-dimension preference, so -1 never means "no constraint". + std::ostringstream oss; + oss << "[PTO layout inference] the 2D unit has no vectorizable " + "dimension (M=" + << M << ", N=" << N << ", L=" << lanes << "): " << candidate_why; + throw LoopLayoutInjectiveException(oss.str()); + } + const bool select_inner = required == 1; + const int64_t sel_extent = select_inner ? N : M; + const int64_t kept_extent = select_inner ? M : N; + // Shared checked padding for the selected dimension. + const auto [pad, P_sel] = CheckedPaddingForInference(sel_extent, lanes); + // The flattened product K*P_sel indexes the int32 loop-var space; both + // operands of every t = kept * P_sel + sel evaluation must stay + // representable. + if (P_sel > std::numeric_limits::max() || + kept_extent > std::numeric_limits::max() / P_sel) { + std::ostringstream oss; + oss << "[PTO layout inference] padded flattened extent " << kept_extent + << "*" << P_sel + << " overflows the int32 index " + "space for lanes " + << lanes; + throw LoopLayoutInjectiveException(oss.str()); + } + frag_vars = loop_vars_; + if (pad != 0) { + frag_vars.Set(select_inner ? 1 : 0, + IterVar(Range::FromMinExtent(make_zero(idx_dtype), + IntImm(idx_dtype, P_sel)), + select_inner ? j_var : i_var, + IterVarType::kDataPar)); + } + flattened = select_inner ? InputPlaceholder(0) * IntImm(idx_dtype, P_sel) + + InputPlaceholder(1) + : InputPlaceholder(1) * IntImm(idx_dtype, P_sel) + + InputPlaceholder(0); + } + PrimExpr lane = FloorMod(flattened, IntImm(idx_dtype, lanes)); + PrimExpr chunk = FloorDiv(flattened, IntImm(idx_dtype, lanes)); + Var rep_var; + IterVar rep_iter = + IterVar(Range::FromMinExtent(make_zero(idx_dtype), IntImm(idx_dtype, 1)), + rep_var, IterVarType::kDataPar); + auto fragment = Fragment(frag_vars, {chunk}, lane, rep_iter) + ->BindThreadRange(layout_args.thread_bounds); + + // The fragment must not only be forward-injective: Vectorize consumes it + // through Fragment::InverseWithLevel, so a layout that cannot be inverted + // with the expected level has to fail here, closest to its construction, + // not deep inside the vectorize pass. + try { + auto inverse_info = fragment->InverseWithLevel( + /*require_padding_guard=*/false); + ICHECK(inverse_info.second == arith::IterMapLevel::Bijective) + << "[PTO layout inference] PTO lane layout must invert at " + "Bijective level, got " + << static_cast(inverse_info.second) << " for " + << fragment->DebugOutput(); + } catch (const Error &err) { + std::ostringstream oss; + oss << "[PTO layout inference] lane layout cannot be inverted by " + "InverseWithLevel: " + << fragment->DebugOutput() << "\n error: " << err.what() + << "\n loop AST: " << root_; + throw LoopLayoutInjectiveException(oss.str()); + } + + loop_layout_ = fragment; + loop_layout_inferred_ = true; + // No buffer layout updates: PTO SIMD loops access + // shared/global buffers only, which have no per-lane fragment layouts. + return {}; +} + Fragment ParallelOpNode::ComputePlanCandidate(const LayoutInferArgs &layout_args) const { // Vectorize Size must be aware of the buffer_remap diff --git a/src/op/parallel.h b/src/op/parallel.h index 1b77d821d3..69af4247bc 100644 --- a/src/op/parallel.h +++ b/src/op/parallel.h @@ -221,6 +221,11 @@ class ParallelOpNode : public TileOperatorNode { Fragment ComputeLoopLayoutFromBuffer(const Buffer &buffer, const LayoutInferArgs &layout_args) const; + // PTO SIMD lane layout inference (inside a PTO SIMD_VF block): builds the + // [0, L) lane Fragment with vectorize_size = 1; no SIMT vector-width + // selection. Returns an empty layout map (no fragment + // buffers in scope). + LayoutMap InferPtoSimdLayout(const LayoutInferArgs &layout_args) const; // Compute plan-based loop layout candidate using vectorization and thread // bounds. Fragment ComputePlanCandidate(const LayoutInferArgs &layout_args) const; diff --git a/src/op/pto_index_analysis.cc b/src/op/pto_index_analysis.cc new file mode 100644 index 0000000000..1c0b55330e --- /dev/null +++ b/src/op/pto_index_analysis.cc @@ -0,0 +1,204 @@ +/*! + * \file pto_index_analysis.cc + * \brief Low-level read-only index/address analysis shared by the PTO + * Parallel layout planning and the Verify/Vectorize consumers + * (see pto_index_analysis.h). + */ + +#include "pto_index_analysis.h" + +#include +#include + +#include +#include +#include +#include + +namespace tvm { +namespace tl { +namespace pto { + +using namespace tirx; +using ffi::GetRef; +using ffi::Map; +using ffi::ObjectRef; + +namespace { + +// Substitute `from` -> `to` inside an expression tree. +class VarSubstituter : public ExprMutator { +public: + VarSubstituter(const VarNode *from, PrimExpr to) + : from_(from), to_(std::move(to)) {} + +private: + PrimExpr VisitExpr_(const VarNode *op) final { + if (op == from_) { + return to_; + } + return ffi::GetRef(op); + } + const VarNode *from_; + PrimExpr to_; +}; + +} // namespace + +PrimExpr SubstituteVar(const PrimExpr &expr, const VarNode *from, PrimExpr to) { + return VarSubstituter(from, std::move(to))(expr); +} + +std::pair CheckedPtoPadding(int64_t extent, int64_t lanes) { + // The PTO lane loop indexes with int32, so any extent above the int32 + // range can never execute; rejecting it here (before the padding + // addition) also makes an int64 overflow of extent + pad impossible: + // with extent <= INT32_MAX and lanes <= 256, extent + pad stays far + // below INT64_MAX. + constexpr int64_t kMaxIndex = std::numeric_limits::max(); + if (extent > kMaxIndex) { + std::ostringstream oss; + oss << "PTO extent " << extent << " exceeds the int32 index space (max " + << kMaxIndex << ")"; + throw PtoAnalysisError(oss.str()); + } + const int64_t pad = extent % lanes == 0 ? 0 : lanes - extent % lanes; + const int64_t padded = extent + pad; + // Classified re-check: an extent within int32 still gets a distinct + // diagnostic when the padded value itself crosses the boundary. + if (padded > kMaxIndex) { + std::ostringstream oss; + oss << "PTO padded extent " << padded + << " exceeds the int32 index space (max " << kMaxIndex << ")"; + throw PtoAnalysisError(oss.str()); + } + return {pad, padded}; +} + +Optional ResolveIndexAliases(const PrimExpr &index, + const Map &bind_env) { + if (bind_env.empty()) { + return index; + } + // A Bind definition containing a BufferLoad is an SSA + // *value*, not a coordinate alias — never expand it into a memory read. + // The check runs at *every* substitution step (not only on the original + // index) so the value-identity rule propagates through coordinate + // alias chains (alias = saved; ... A[i + alias - IDX[i]]). + auto def_is_pure_coordinate = [&](const PrimExpr &def_expr) { + bool has_load = false; + PostOrderVisit(def_expr, [&](const ObjectRef &node) { + if (node.as()) { + has_load = true; + } + }); + return !has_load; + }; + constexpr int kMaxDepth = 64; + // Expression-size budget. The depth cap alone is not enough — + // every substitution round can double the expression (x1 = i + i; + // x2 = x1 + x1; ...), so a 26-deep chain already exploded to 2^26 nodes + // and died minutes later on a raw tvm int32-literal assertion instead of + // a diagnostic. Count the nodes after each round and reject once the + // budget is exceeded. Legal chains (plain var aliases, (i + x) - x + // style cancellations) stay far below the budget and resolve quickly. + constexpr int64_t kMaxNodes = 4096; + auto count_nodes = [](const PrimExpr &expr) { + int64_t nodes = 0; + PostOrderVisit(expr, [&](const ObjectRef &) { ++nodes; }); + return nodes; + }; + bool exhausted = true; + std::function resolve = [&](PrimExpr &expr, + int depth) -> bool { + if (depth > kMaxDepth) { + exhausted = false; + return false; + } + const VarNode *found = nullptr; + PostOrderVisit(expr, [&](const ObjectRef &node) { + if (found == nullptr) { + if (auto *var_node = node.as()) { + if (bind_env.count(GetRef(var_node))) { + found = var_node; + } + } + } + }); + if (found == nullptr) { + return true; + } + // Check the definition *before* substituting; a load-valued + // definition reached through any alias depth is rejected. + const PrimExpr &def = bind_env[GetRef(found)]; + if (!def_is_pure_coordinate(def)) { + exhausted = false; + return false; + } + expr = SubstituteVar(expr, found, def); + if (count_nodes(expr) > kMaxNodes) { + std::ostringstream oss; + oss << "bind alias resolution exceeded the expression-size budget (" + << kMaxNodes + << " nodes); the alias chain expands too much " + "(e.g. doubling definitions like x = y + y)"; + throw PtoAnalysisError(oss.str()); + } + return resolve(expr, depth + 1); + }; + PrimExpr result = index; + try { + if (!resolve(result, 0) || !exhausted) { + return Optional(); + } + } catch (const ffi::Error &err) { + // A tvm-internal failure during substitution (e.g. an int32 + // literal overflowing while an expression is rebuilt) must surface as + // an analysis diagnostic, not a raw Check failed crash. + throw PtoAnalysisError( + std::string("bind alias resolution hit a tvm-internal error: ") + + err.what()); + } + return result; +} + +bool IsSupportedAddressExpr(const PrimExpr &expr) { + // Addresses may contain coordinate arithmetic, constants, casts + // and pure intrinsics — anything whose per-chunk evaluation is + // equivalent. Opaque/extern calls (call_extern is kOpaque in TVM) have + // no lowering rule; evaluating them once per chunk instead of per + // logical point changes the program, so they are rejected. A load + // inside an index makes the address data-dependent: the alias resolver + // already rejects Bind-introduced loaded indices, and a direct load + // must follow the same rule. + class AddressCheck : public ExprVisitor { + public: + bool supported = true; + void VisitExpr_(const CallNode *op) final { + if (!supported) { + return; + } + const auto *op_node = op->op.as(); + std::string name = op_node != nullptr ? op_node->name : "call"; + // Pure index intrinsics with defined semantics may appear; only + // known-pure builtins are currently allowed. + // call_extern/call_pure_extern/call_llvm_intrin etc. are rejected: + // their effect is opaque to this pass. + if (name == "tir.tvm_access_ptr" || name == "tl.access_ptr" || + name == "tir.address_of") { + ExprVisitor::VisitExpr_(op); + return; + } + supported = false; + } + void VisitExpr_(const BufferLoadNode *) final { supported = false; } + void VisitExprDefault_(const ffi::Object *) final { /* fine */ } + }; + AddressCheck checker; + checker(expr); + return checker.supported; +} + +} // namespace pto +} // namespace tl +} // namespace tvm diff --git a/src/op/pto_index_analysis.h b/src/op/pto_index_analysis.h new file mode 100644 index 0000000000..42e4e90b30 --- /dev/null +++ b/src/op/pto_index_analysis.h @@ -0,0 +1,71 @@ +/*! + * \file pto_index_analysis.h + * \brief Low-level read-only index/address analysis shared by the PTO + * Parallel layout planning (src/op/parallel.cc) and the Verify/Vectorize + * consumers (src/ascend/transform/parallel_to_pto_utils.cc). + * + * These helpers live in src/op so both layers can depend on them without + * an op -> ascend/transform reverse include. They are the single + * implementation of the bind-alias resolution rules (cycle/depth failure, + * expression-size budget, loaded-value identity) and the address-node + * whitelist; planning and consuming code must not re-derive their own + * variants. + */ + +#ifndef TVM_TL_OP_PTO_INDEX_ANALYSIS_H_ +#define TVM_TL_OP_PTO_INDEX_ANALYSIS_H_ + +#include +#include + +#include +#include + +namespace tvm { +namespace tl { +namespace pto { + +using ffi::Map; +using ffi::Optional; +using tirx::Var; +using tirx::VarNode; + +/*! Error type thrown by the shared analysis on unsupported input; both + * passes translate it into their own diagnostics. */ +class PtoAnalysisError : public std::runtime_error { +public: + explicit PtoAnalysisError(const std::string &msg) : std::runtime_error(msg) {} +}; + +/*! Substitute `from` -> `to` inside an expression tree. */ +PrimExpr SubstituteVar(const PrimExpr &expr, const VarNode *from, PrimExpr to); + +/*! Resolve Bind aliases inside an index expression recursively with a + * visited set (cycle-safe) and a depth guard that *rejects* instead of + * silently stopping. Returns nullopt when a chain exceeds the + * resolution limit; the caller must reject the access. */ +Optional ResolveIndexAliases(const PrimExpr &index, + const Map &bind_env); + +/*! Check whether an address/index expression only uses supported + * nodes: coordinate arithmetic over loop/bound vars, constants, + * casts and pure intrinsics are allowed; loads inside the index + * (data-dependent addresses) and opaque or side-effecting calls + * (call_extern etc.) are rejected because collapsing them to one + * evaluation per chunk changes the program's evaluation behavior. */ +bool IsSupportedAddressExpr(const PrimExpr &expr); + +/*! Compute the padded fragment-input extent P = ceil(E / lanes) * lanes with + * the shared overflow guard. The final execution space is int32-bounded + * (the PTO lane loop uses int32 indices), so an extent above the int32 + * range is rejected before the padding addition could overflow int64. + * Returns the pair (pad, padded); throws PtoAnalysisError with a + * classified reason on rejection. The callers translate the error into + * their own diagnostics. */ +std::pair CheckedPtoPadding(int64_t extent, int64_t lanes); + +} // namespace pto +} // namespace tl +} // namespace tvm + +#endif // TVM_TL_OP_PTO_INDEX_ANALYSIS_H_ diff --git a/testing/ascend/language/pto/test_parallel_pipeline_routing.py b/testing/ascend/language/pto/test_parallel_pipeline_routing.py new file mode 100644 index 0000000000..ec482f795f --- /dev/null +++ b/testing/ascend/language/pto/test_parallel_pipeline_routing.py @@ -0,0 +1,233 @@ +"""Routing tests for the PTO Parallel pipeline skeleton. + +Scope: this file proves routing only. + +* a PTO ``SimdVF + Parallel`` kernel is routed through + ``VerifyParallelToPTO -> LegalizeParallelToPTO -> VectorizeParallelToPTO`` + in that order, and never through ``AscendSimdVFLowerParallel``; +* the AscendC SIMD_VF route keeps ``AscendSimdVFLowerParallel`` and never sees + the three PTO Parallel passes; +* PTO ``SimtVF`` and hand-written PTO VMI modules without a ``T.Parallel`` + unit enter the same three PTO pass entry points, but are bypassed inside + those passes and accepted without conversion; +* ``PlanPtoSlotStorage`` is neither registered nor called. + +This file deliberately asserts routing only: no PTO Parallel semantics, no +lowering output, and no Fragment/slot, control-flow, or Reduce support. Use +``TL_PTO_ROUTING_TRACE_DIR=`` to dump the per-case ordered pass trace and +the lowered kernel source as review artifacts. + +The lowering runs through :func:`tilelang.engine.lower.lower`: TileLang IR +lowering only, with no PTOAS/Bisheng device compilation and no device +execution. +""" + +from __future__ import annotations + +import os + +import pytest + +import tilelang.ascend.language as T +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import ( + normalize_ascend_target, + normalize_asc_target, + normalize_pto_target, +) +from tilelang.engine.lower import lower as _lower +from tilelang.instrumentation import ( + PassEventObserver, + PassInstrumentationTool, + StackedPassInstrument, + compile_pass_instrumentation, + create_pass_instruments, +) +import tvm + +_PTO_PARALLEL_PASSES = ( + "tl.VerifyParallelToPTO", + "tl.LegalizeParallelToPTO", + "tl.VectorizeParallelToPTO", +) +_ASCENDC_SIMDVF_PASS = "tl.AscendSimdVFLowerParallel" +_SLOT_STORAGE_PASS = "PlanPtoSlotStorage" + + +class _PassRecorder(PassEventObserver): + """Records the ordered pass names observed in one lowering.""" + + def __init__(self) -> None: + self.names: list[str] = [] + + def pass_started(self, mod, event): + self.names.append(event.name) + return None + + +class _PassTraceTool(PassInstrumentationTool): + """Compile-session tool that installs the routing recorder.""" + + def __init__(self, recorder: _PassRecorder) -> None: + self._recorder = recorder + + def create_pass_instrument(self): + return StackedPassInstrument(self._recorder) + + +def _resolve_target(name: str): + if name == "pto": + return normalize_pto_target("pto") + if name == "ascend": + return normalize_ascend_target("ascend") + if name == "asc": + return normalize_asc_target("asc") + raise ValueError(f"unsupported target name in routing test: {name}") + + +def _lower_with_trace(func, target_name: str): + """Lower *func* for *target_name* while recording the executed pass names.""" + recorder = _PassRecorder() + # The session must be active before the instruments are created, so the two + # context managers are entered in this order. + with ( + compile_pass_instrumentation(name="pto-routing-test", tools=[_PassTraceTool(recorder)]), + tvm.transform.PassContext(instruments=create_pass_instruments(context="pto-routing-test")), + ): + target = _resolve_target(target_name) + with target: + artifact = _lower(func, target=target) + return artifact, recorder.names + + +def _maybe_dump_artifacts(case: str, names: list[str], artifact) -> None: + """Dump the pass trace and lowered source when the artifact dir is set.""" + out_dir = os.environ.get("TL_PTO_ROUTING_TRACE_DIR") + if not out_dir: + return + os.makedirs(out_dir, exist_ok=True) + with open(os.path.join(out_dir, f"{case}.pass-trace.txt"), "w") as handle: + handle.write("\n".join(names) + "\n") + source = getattr(artifact, "kernel_source", None) + if source: + with open(os.path.join(out_dir, f"{case}.source.txt"), "w") as handle: + handle.write(source) + + +def _assert_ordered(names: list[str], expected: tuple[str, ...]) -> None: + for pass_name in expected: + assert pass_name in names, f"{pass_name} was not executed; trace tail: {names[-12:]}" + positions = [names.index(pass_name) for pass_name in expected] + assert positions == sorted(positions), f"expected order {' -> '.join(expected)}, got positions {positions}" + + +def _simdvf_parallel_kernel(n: int = 256): + @T.prim_func + def main( + A: T.Tensor((n,), "float32"), + B: T.Tensor((n,), "float32"), + C: T.Tensor((n,), "float32"), + ): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "float32") + b_ub = T.alloc_shared((n,), "float32") + c_ub = T.alloc_shared((n,), "float32") + T.copy(A, a_ub) + T.copy(B, b_ub) + with T.SimdVF(): + for i in T.Parallel(n): + c_ub[i] = a_ub[i] + b_ub[i] + T.copy(c_ub, C) + + return main + + +def _handwritten_vmi_kernel(n: int = 64): + """SimdVF with hand-written VMI calls and no ``T.Parallel`` unit.""" + + @T.prim_func + def main( + A: T.Tensor((n,), "float32"), + B: T.Tensor((n,), "float32"), + C: T.Tensor((n,), "float32"), + ): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "float32") + b_ub = T.alloc_shared((n,), "float32") + c_ub = T.alloc_shared((n,), "float32") + T.copy(A, a_ub) + T.copy(B, b_ub) + T.copy(C, c_ub) + with T.SimdVF(): + mask = T.vmi.create_mask(n, size=n) + c = T.vmi.vload(c_ub[0], size=n) + a = T.vmi.vload(a_ub[0], size=n) + b = T.vmi.vload(b_ub[0], size=n) + out = T.vmi.vadd(T.vmi.vmul(c, a, mask), b, mask) + T.vmi.vstore(out, c_ub[0], mask) + T.copy(c_ub, C) + + return main + + +def _simt_kernel(n: int = 2): + @T.prim_func + def main( + A: T.Tensor((n,), "float32"), + B: T.Tensor((n,), "float32"), + C: T.Tensor((n,), "float32"), + ): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "float32") + b_ub = T.alloc_shared((n,), "float32") + c_ub = T.alloc_shared((n,), "float32") + T.copy(A, a_ub) + T.copy(B, b_ub) + with T.SimtVF(threads=1): + for i in T.vectorized(n): + c_ub[i] = a_ub[i] + b_ub[i] + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_pto_simdvf_parallel_routes_verify_legalize_vectorize_in_order(): + artifact, names = _lower_with_trace(_simdvf_parallel_kernel(), "pto") + _assert_ordered(names, _PTO_PARALLEL_PASSES) + assert _ASCENDC_SIMDVF_PASS not in names, "the PTO route must not run the AscendC SIMD_VF lowering" + _maybe_dump_artifacts("pto-simdvf-parallel", names, artifact) + + +def test_ascendc_simdvf_keeps_ascend_simdvf_lower_parallel(): + artifact, names = _lower_with_trace(_simdvf_parallel_kernel(), "ascend") + assert _ASCENDC_SIMDVF_PASS in names, f"the AscendC route must keep AscendSimdVFLowerParallel; trace tail: {names[-12:]}" + for pass_name in _PTO_PARALLEL_PASSES: + assert pass_name not in names, f"the AscendC route must not run {pass_name}" + _maybe_dump_artifacts("ascendc-simdvf-parallel", names, artifact) + + +@pytest.mark.pto +def test_pto_simt_is_accepted_by_the_pto_route(): + artifact, names = _lower_with_trace(_simt_kernel(), "pto") + for pass_name in _PTO_PARALLEL_PASSES: + assert pass_name in names, f"{pass_name} missing on the PTO route; tail: {names[-12:]}" + _maybe_dump_artifacts("pto-simt", names, artifact) + + +@pytest.mark.pto +def test_pto_handwritten_vmi_without_parallel_is_accepted(): + artifact, names = _lower_with_trace(_handwritten_vmi_kernel(), "pto") + for pass_name in _PTO_PARALLEL_PASSES: + assert pass_name in names, f"{pass_name} missing on the PTO route; tail: {names[-12:]}" + _maybe_dump_artifacts("pto-handwritten-vmi", names, artifact) + + +@pytest.mark.pto +def test_plan_pto_slot_storage_is_not_registered_or_called(): + assert not hasattr(ascend_transform, _SLOT_STORAGE_PASS), ( + "PlanPtoSlotStorage belongs to a later pipeline stage and must not be registered yet" + ) + artifact, names = _lower_with_trace(_simdvf_parallel_kernel(), "pto") + assert not any(_SLOT_STORAGE_PASS in name for name in names), "no PlanPtoSlotStorage step may run" + _maybe_dump_artifacts("pto-slot-storage-absence", names, artifact) diff --git a/testing/ascend/language/pto/test_pto_parallel_vectorize_structure.py b/testing/ascend/language/pto/test_pto_parallel_vectorize_structure.py new file mode 100644 index 0000000000..07547321fd --- /dev/null +++ b/testing/ascend/language/pto/test_pto_parallel_vectorize_structure.py @@ -0,0 +1,566 @@ +"""PTO Parallel vectorization structure tests. + +Hardware independent: each case lowers TileLang IR for the PTO target with +:func:`tilelang.engine.lower.lower` (TileLang IR only, no PTOAS/Bisheng device +compilation and no device execution) and inspects the module produced by +``VectorizeParallelToPTO`` through a capture wrapper. + +Covered contract: + +* a divisible unit reuses the shared full mask (``create_mask(L)``) and emits + no ``remaining`` state; +* a non-divisible unit allocates a scalar ``remaining`` (scope ``local.var``), + initializes it to the logical extent ``E`` once, binds a per-chunk + ``create_mask(remaining)`` before the body and decrements it by ``L``; +* a 2D unit keeps the non-selected coordinate as an outer serial loop and + resets ``remaining`` inside it, so rows are independent; +* the real buffer indices (including a padded row stride) survive + vectorization: addresses keep the original buffers and stride expressions; +* the region emits ``vload`` / ``vstore`` / ``create_mask`` VMI calls with the + chunk mask on the store; +* explicit padding (``T.Parallel(P)`` with an ``if`` guard) is currently + rejected: guard legalization is left to a later change. +""" + +from __future__ import annotations + +from contextlib import contextmanager + +import re + +import pytest + +import tilelang.ascend.language as T +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower +import tvm + +_REMAINING = "remaining" + + +@contextmanager +def _verify_boundary_capture(): + """Capture the module handed to VerifyParallelToPTO, then delegate. + + Self-contained copy of the contract suite's helper so this file runs + standalone (cross-test-module imports are fragile under pytest's + rootdir-dependent module resolution). + """ + captured = {} + original = ascend_transform.VerifyParallelToPTO + + def capturing_factory(): + verify_pass = original() + + def run(mod): + captured["module"] = mod + return verify_pass(mod) + + return run + + ascend_transform.VerifyParallelToPTO = capturing_factory + try: + yield captured + finally: + ascend_transform.VerifyParallelToPTO = original + + +def _ceil_to_lanes(extent: int, lanes: int) -> int: + return (extent + lanes - 1) // lanes * lanes + + +def _as_int(expr) -> int: + if isinstance(expr, int): + return expr + value = getattr(expr, "value", None) + assert value is not None, f"expected an integer constant, got {expr!r}" + return int(value) + + +@contextmanager +def _vectorize_capture(): + """Capture the module returned by VectorizeParallelToPTO, then delegate.""" + captured = {} + original = ascend_transform.VectorizeParallelToPTO + + def capturing_factory(): + vectorize_pass = original() + + def run(mod): + result = vectorize_pass(mod) + captured["module"] = result + return result + + return run + + ascend_transform.VectorizeParallelToPTO = capturing_factory + try: + yield captured + finally: + ascend_transform.VectorizeParallelToPTO = original + + +def _lower_capture(func, *, expect_error=False): + """Lower *func* for PTO and return (module_after_vectorize, error).""" + target = normalize_pto_target("pto") + error = None + with _vectorize_capture() as captured: + try: + with target: + _lower(func, target=target) + except Exception as exc: # noqa: BLE001 - surfaced to the caller + error = exc + if error is not None and not expect_error: + raise error + return captured.get("module"), error + + +def _body_of(mod): + assert len(mod.functions) == 1, f"expected a single function, got {list(mod.functions)}" + return next(iter(mod.functions.values())).body + + +def _collect(mod, predicate): + found = [] + + def visit(node): + if predicate(node): + found.append(node) + + tvm.tirx.stmt_functor.post_order_visit(_body_of(mod), visit) + return found + + +def _vmi_calls(mod): + calls = {} + + def visit(node): + if isinstance(node, tvm.tirx.Call): + op = getattr(node, "op", None) + name = getattr(op, "name", None) + if isinstance(name, str) and name.startswith("tl.vmi."): + calls.setdefault(name, []).append(node) + + tvm.tirx.stmt_functor.post_order_visit(_body_of(mod), visit) + return calls + + +def _buffers(mod): + buffers = {} + + def visit(node): + if isinstance(node, tvm.tirx.AllocBuffer): + buffers[node.buffer.name] = node.buffer + if isinstance(node, (tvm.tirx.BufferStore, tvm.tirx.BufferLoad)): + buffers[node.buffer.name] = node.buffer + + tvm.tirx.stmt_functor.post_order_visit(_body_of(mod), visit) + return buffers + + +def _remaining_buffers(mod): + return [buf for name, buf in _buffers(mod).items() if name.startswith(_REMAINING)] + + +def _remaining_stores(mod, buffer): + stores = [] + + def visit(node): + if isinstance(node, tvm.tirx.BufferStore) and node.buffer.same_as(buffer): + stores.append(node) + + tvm.tirx.stmt_functor.post_order_visit(_body_of(mod), visit) + return stores + + +def _subexpr_contains_constant(expr, value: int) -> bool: + """Return whether *expr* contains an IntImm equal to *value*.""" + found = [] + + def visit(node): + if isinstance(node, tvm.tirx.IntImm) and int(node.value) == value: + found.append(node) + + tvm.tirx.stmt_functor.post_order_visit(tvm.tirx.Evaluate(expr), visit) + return bool(found) + + +# --------------------------------------------------------------------------- +# Kernels +# --------------------------------------------------------------------------- + + +def _k1d(E, out_extent, lanes, dtype="float32"): + P = _ceil_to_lanes(E, lanes) + + @T.prim_func + def main( + A: T.Tensor((P,), dtype), + B_init: T.Tensor((out_extent,), dtype), + B: T.Tensor((out_extent,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((P,), dtype) + b = T.alloc_shared((out_extent,), dtype) + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(E): + b[i] = a[i] + a[i] + T.copy(b, B) + + return main + + +def _k2d_row(M, N, lanes, dtype="float32"): + P = _ceil_to_lanes(N, lanes) + + @T.prim_func + def main( + A: T.Tensor((M * P,), dtype), + B_init: T.Tensor((M * P,), dtype), + B: T.Tensor((M * P,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((M * P,), dtype) + b = T.alloc_shared((M * P,), dtype) + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes), T.Parallel(M, N) as (i, j): + b[i * P + j] = a[i * P + j] + a[i * P + j] + T.copy(b, B) + + return main + + +def _k_explicit_padding(E, lanes, dtype="float32"): + """Explicit ``T.Parallel(P)`` with a user ``if`` guard: at this stage the + guard is still rejected; its legalization belongs to the control-flow + stage.""" + + P = _ceil_to_lanes(E, lanes) + + @T.prim_func + def main( + A: T.Tensor((P,), dtype), + B: T.Tensor((P,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((P,), dtype) + b = T.alloc_shared((P,), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(P): + if i < E: + b[i] = a[i] + a[i] + else: + b[i] = a[i] + T.copy(b, B) + + return main + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.pto +def test_divisible_unit_uses_the_shared_full_mask_without_remaining(): + E, lanes = 256, 128 + mod, _ = _lower_capture(_k1d(E, E, lanes)) + + assert not _remaining_buffers(mod), "a divisible unit must not allocate remaining" + calls = _vmi_calls(mod) + assert "tl.vmi.create_mask" in calls, "the region must create the shared full mask" + mask_sizes = [_as_int(call.args[0]) for call in calls["tl.vmi.create_mask"]] + assert mask_sizes == [lanes], f"the shared mask must be created with L={lanes}, got {mask_sizes}" + assert "tl.vmi.vload" in calls and "tl.vmi.vstore" in calls + + fors = _collect(mod, lambda n: isinstance(n, tvm.tirx.For)) + extents = sorted(_as_int(f.extent) for f in fors) + assert extents == [2], f"only the Q=2 chunk loop may remain, got {extents}" + + +@pytest.mark.pto +def test_nondivisible_unit_builds_remaining_prefix_mask(): + E, lanes = 150, 128 + P = _ceil_to_lanes(E, lanes) + mod, _ = _lower_capture(_k1d(E, P, lanes)) + + buffers = _remaining_buffers(mod) + assert len(buffers) == 1, f"expected exactly one remaining buffer, got {buffers}" + remaining = buffers[0] + assert remaining.scope() == "local.var", f"remaining must be a scalar local, got {remaining.scope()}" + + stores = _remaining_stores(mod, remaining) + values = [_as_int(store.value) for store in stores if isinstance(store.value, tvm.tirx.IntImm)] + assert E in values, f"remaining must be initialized to the logical extent {E}, got {values}" + + calls = _vmi_calls(mod) + # The per-chunk mask size lives in the Call kwargs (`size`), which the + # Python FFI does not expose; the observable contract is the value + # argument: exactly one create_mask per unit whose source is the remaining + # scalar load. + prefix_masks = [ + call + for call in calls.get("tl.vmi.create_mask", []) + if isinstance(call.args[0], tvm.tirx.BufferLoad) and call.args[0].buffer.same_as(remaining) + ] + assert len(prefix_masks) == 1, f"expected one prefixed create_mask(remaining), got {len(prefix_masks)}" + + fors = _collect(mod, lambda n: isinstance(n, tvm.tirx.For)) + assert [_as_int(f.extent) for f in fors] == [2], f"expected the Q=2 chunk loop, got {fors}" + + +@pytest.mark.pto +def test_2d_unit_resets_remaining_inside_the_kept_loop(): + M, N, lanes = 4, 150, 128 + mod, _ = _lower_capture(_k2d_row(M, N, lanes)) + + buffers = _remaining_buffers(mod) + assert len(buffers) == 1, f"expected exactly one remaining buffer, got {buffers}" + remaining = buffers[0] + + fors = _collect(mod, lambda n: isinstance(n, tvm.tirx.For)) + kept = [f for f in fors if _as_int(f.extent) == M] + chunk = [f for f in fors if _as_int(f.extent) == 2] + assert len(kept) == 1 and len(chunk) == 1, f"expected one kept loop and one chunk loop, got {fors}" + + def _inside(outer, inner) -> bool: + seen = [] + + def visit(node): + if node.same_as(inner): + seen.append(node) + + tvm.tirx.stmt_functor.post_order_visit(outer, visit) + return bool(seen) + + assert _inside(kept[0], chunk[0]), "the chunk loop must be nested inside the kept loop" + init_stores = [ + store for store in _remaining_stores(mod, remaining) if isinstance(store.value, tvm.tirx.IntImm) and int(store.value.value) == N + ] + assert init_stores, "remaining must be initialized to the selected extent N" + assert _inside(kept[0], init_stores[0]), "the 2D reset must live inside the kept loop" + + +@pytest.mark.pto +def test_real_stride_addresses_survive_vectorization(): + M, N, lanes = 4, 150, 128 + P = _ceil_to_lanes(N, lanes) + mod, _ = _lower_capture(_k2d_row(M, N, lanes)) + + calls = _vmi_calls(mod) + assert "tl.vmi.vload" in calls and "tl.vmi.vstore" in calls + + strides = [] + for call in calls["tl.vmi.vload"] + calls["tl.vmi.vstore"]: + for arg in call.args: + if isinstance(arg, tvm.tirx.PrimExpr) and _subexpr_contains_constant(arg, P): + strides.append(call) + break + assert strides, f"the padded row stride {P} must appear in the vector addresses" + + buffers = _buffers(mod) + for name in ("a", "b"): + assert name in buffers, f"buffer {name} must survive" + shape = [_as_int(dim) for dim in buffers[name].shape] + assert shape == [M * P], f"buffer {name} shape must stay as allocated, got {shape}" + + +@pytest.mark.pto +def test_explicit_padding_is_rejected_at_this_stage(): + E, lanes = 150, 128 + _, error = _lower_capture(_k_explicit_padding(E, lanes), expect_error=True) + assert error is not None, "an if-guarded T.Parallel must still be rejected at this stage" + message = str(error) + assert "VerifyParallelToPTO" in message or "conditional" in message or "unsupported" in message, ( + f"expected the guard rejection diagnostic, got: {message}" + ) + + +# --------------------------------------------------------------------------- +# Sub/div/unary emission and lane-uniform folding +# --------------------------------------------------------------------------- + + +def _sub_kernel(n, lanes): + @T.prim_func + def main(A: T.Tensor((n,), "float32"), B: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "float32") + b_ub = T.alloc_shared((n,), "float32") + c_ub = T.alloc_shared((n,), "float32") + T.copy(A, a_ub) + T.copy(B, b_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c_ub[i] = a_ub[i] - b_ub[i] + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_vectorize_emits_vsub(): + source = _lower(_sub_kernel(256, 128), target="pto").kernel_source + assert "pto.vmi.vsub(" in source + + +def _int_uniform_fold_kernel(n, lanes): + @T.prim_func + def main(A: T.Tensor((n,), "int32"), S: T.Tensor((1,), "int32"), C: T.Tensor((n,), "int32")): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "int32") + s_ub = T.alloc_shared((1,), "int32") + c_ub = T.alloc_shared((n,), "int32") + T.copy(A, a_ub) + T.copy(S, s_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + # The scalar bind pins the uniform Sub into its own + # statement (an inline `a - (s-1)` is algebraically + # restructured into `(a+1) - s` by the analyzer). The + # bound value must stay a scalar Sub feeding the + # broadcast, not become a vsub of its own. + d: T.int32 = s_ub[0] - 1 + c_ub[i] = a_ub[i] - d + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_vectorize_int_uniform_sub_folds_scalar(): + + from tilelang.ascend import transform as atx + + # Run Vectorize directly on the captured post-Verify module and check + # the emitted TIR: the inner uniform Sub must remain a scalar SubNode + # inside the vbrc operand, not become a tl.vmi.vsub call of its own. + with _verify_boundary_capture() as capture: + _lower(_int_uniform_fold_kernel(256, 128), target="pto") + mod = capture["module"] + out = atx.VectorizeParallelToPTO()(mod) + script = out["main"].script() + assert script.count("tl.vmi.vsub") == 1, "expected exactly one vsub (the varying outer sub)\n" + script + scalar_sub = re.search(r"d: T\.int32 = .* - 1", script) + assert scalar_sub is not None, "expected the runtime-uniform sub to stay a scalar bind\n" + script + + +def _runtime_uniform_float_kernel(n, lanes): + @T.prim_func + def main(S: T.Tensor((2,), "float32"), C: T.Tensor((n,), "float32")): + with T.Kernel(1): + s_ub = T.alloc_shared((2,), "float32") + c_ub = T.alloc_shared((n,), "float32") + T.copy(S, s_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + # Both operands are runtime lane-uniform floats: this + # stays scalar and must be rejected at Verify. + c_ub[i] = s_ub[0] / s_ub[1] + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_verify_runtime_uniform_float_div_rejected(): + with pytest.raises( + Exception, + match=r"runtime lane-uniform float computation.*scalar float ALU", + ): + _lower(_runtime_uniform_float_kernel(256, 128), target="pto") + + +def _unary_kernel(n, lanes, kind): + @T.prim_func + def main(A: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "float32") + c_ub = T.alloc_shared((n,), "float32") + T.copy(A, a_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + if kind == 0: + c_ub[i] = T.abs(a_ub[i] - 0.5) + else: + c_ub[i] = T.sqrt(T.abs(a_ub[i]) + 0.25) + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_vectorize_emits_vabs(): + source = _lower(_unary_kernel(256, 128, 0), target="pto").kernel_source + assert "pto.vmi.vabs(" in source + + +@pytest.mark.pto +def test_vectorize_emits_vsqrt(): + source = _lower(_unary_kernel(256, 128, 1), target="pto").kernel_source + assert "pto.vmi.vsqrt(" in source + + +# --------------------------------------------------------------------------- +# Lane-use semantics for the uniform-float rejection +# --------------------------------------------------------------------------- + + +def _varying_bind_then_unary_kernel(n, lanes): + @T.prim_func + def main(A: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "float32") + c_ub = T.alloc_shared((n,), "float32") + T.copy(A, a_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + # x is lane-varying through the bind; the uniform-float + # rejection must NOT fire for exp(x) even though the + # expression text does not mention the loop var. + x = a_ub[i] * 2.0 + c_ub[i] = T.exp(x) + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_varying_bind_unary_accepted(): + source = _lower(_varying_bind_then_unary_kernel(256, 128), target="pto").kernel_source + assert "pto.vmi.vexp(" in source + assert "pto.vmi.vmul(" in source + + +def _chunk_uniform_float_kernel(n, lanes): + @T.prim_func + def main(A: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "float32") + c_ub = T.alloc_shared((n,), "float32") + T.copy(A, a_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + # The address is uniform within a chunk (i // 128 * 128) + # but the expression mentions the loop var. The + # lane-use analysis classifies the load as uniform, so + # the float arithmetic must be rejected even though a + # syntactic loop-var check would have skipped it. + c_ub[i] = a_ub[(i // 128) * 128] + 1.0 + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_chunk_uniform_float_rejected(): + with pytest.raises( + Exception, + match=r"runtime lane-uniform float computation.*scalar float ALU", + ): + _lower(_chunk_uniform_float_kernel(256, 128), target="pto") diff --git a/testing/ascend/language/pto/test_pto_parallel_verify_contract.py b/testing/ascend/language/pto/test_pto_parallel_verify_contract.py new file mode 100644 index 0000000000..86971756bc --- /dev/null +++ b/testing/ascend/language/pto/test_pto_parallel_verify_contract.py @@ -0,0 +1,517 @@ +"""PTO Parallel mapping-contract tests. + +Hardware independent: each case lowers TileLang IR for the PTO target with +:func:`tilelang.engine.lower.lower` (TileLang IR only, no PTOAS/Bisheng device +compilation and no device execution) and inspects the module exactly as +``VerifyParallelToPTO`` receives it: the module is captured at the Verify +boundary instead of the final lowered IR. + +Covered contract: + +* the logical ``T.Parallel`` extent ``E`` stays on the For and the accesses + keep using the original loop variable (no For/Buffer rewrite); +* the fragment executing dimension is padded to ``P = ceil(E / L) * L`` and + is written under the PTO-specific ``pto_parallel_loop_layout`` annotation + (never the SIMT ``parallel_loop_layout`` key); +* ``Q = P / L`` and ``tail_lanes = E % L`` follow from ``E`` and ``L``; +* 2D units pad only the selected dimension and keep the other extent; +* a real (padded) row stride is accepted and buffer shapes are untouched; +* lanes ``{64, 128, 256}`` are accepted and agree with the annotation; +* the rejection diagnostics fire: gather access, invalid lanes, a missing or + invalid lane annotation, and hand-written VMI mixed with ``T.Parallel``. +""" + +from __future__ import annotations + +from contextlib import contextmanager + +import pytest + +import tilelang.ascend.language as T +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower +from tilelang.layout.fragment import Fragment +import tvm + +_PTO_LAYOUT_ATTR = "pto_parallel_loop_layout" +_SIMT_LAYOUT_ATTR = "parallel_loop_layout" +_LANES_ATTR = "tl.simdvf_lanes" + + +def _ceil_to_lanes(extent: int, lanes: int) -> int: + return (extent + lanes - 1) // lanes * lanes + + +def _as_int(expr) -> int: + if isinstance(expr, int): + return expr + value = getattr(expr, "value", None) + assert value is not None, f"expected an integer constant, got {expr!r}" + return int(value) + + +@contextmanager +def _verify_boundary_capture(): + """Capture the module handed to VerifyParallelToPTO, then delegate.""" + captured = {} + original = ascend_transform.VerifyParallelToPTO + + def capturing_factory(): + verify_pass = original() + + def run(mod): + captured["module"] = mod + return verify_pass(mod) + + return run + + ascend_transform.VerifyParallelToPTO = capturing_factory + try: + yield captured + finally: + ascend_transform.VerifyParallelToPTO = original + + +def _lower_capture(func, *, expect_error=False): + """Lower *func* for PTO and return (module_at_verify, error).""" + target = normalize_pto_target("pto") + error = None + with _verify_boundary_capture() as captured: + try: + with target: + _lower(func, target=target) + except Exception as exc: # noqa: BLE001 - surfaced to the caller + error = exc + if error is not None and not expect_error: + raise error + return captured.get("module"), error + + +def _body_of(mod): + assert len(mod.functions) == 1, f"expected a single function, got {list(mod.functions)}" + return next(iter(mod.functions.values())).body + + +def _pto_parallel_loops(mod): + loops = [] + + def visit(node): + if isinstance(node, tvm.tirx.For) and _PTO_LAYOUT_ATTR in node.annotations: + loops.append(node) + + tvm.tirx.stmt_functor.post_order_visit(_body_of(mod), visit) + return loops + + +def _single_pto_parallel_loop(mod): + loops = _pto_parallel_loops(mod) + assert len(loops) == 1, f"expected exactly one pto_parallel_loop_layout unit, got {len(loops)}" + return loops[0] + + +def _fragment_of(loop) -> Fragment: + fragment = loop.annotations[_PTO_LAYOUT_ATTR] + assert isinstance(fragment, Fragment), f"annotation is not a Fragment: {type(fragment)}" + return fragment + + +def _simd_vf_block(mod): + blocks = [] + + def visit(node): + if isinstance(node, tvm.tirx.SBlock) and _LANES_ATTR in node.annotations: + blocks.append(node) + + tvm.tirx.stmt_functor.post_order_visit(_body_of(mod), visit) + assert len(blocks) == 1, f"expected exactly one lane-annotated SIMD_VF block, got {len(blocks)}" + return blocks[0] + + +def _input_shape(fragment): + return [_as_int(dim) for dim in fragment.get_input_shape()] + + +def _first_store(body): + stores = [] + + def visit(node): + if isinstance(node, tvm.tirx.BufferStore) and not stores: + stores.append(node) + + tvm.tirx.stmt_functor.post_order_visit(body, visit) + assert stores, "expected a BufferStore in the parallel unit" + return stores[0] + + +def _rewrite_pto_layout_attr(mod, mutate): + """Return *mod* with each pto_parallel_loop_layout annotation rewritten. + + ``mutate(value)`` returns the replacement value; returning ``None`` drops + the annotation. Used to drive the annotation-level rejection diagnostics. + """ + + def postorder(node): + if isinstance(node, tvm.tirx.For) and _PTO_LAYOUT_ATTR in node.annotations: + annotations = dict(node.annotations) + replacement = mutate(annotations[_PTO_LAYOUT_ATTR]) + if replacement is None: + del annotations[_PTO_LAYOUT_ATTR] + else: + annotations[_PTO_LAYOUT_ATTR] = replacement + return tvm.tirx.For( + node.loop_var, + node.min, + node.extent, + node.kind, + node.body, + thread_binding=node.thread_binding, + annotations=annotations, + ) + return None + + func = next(iter(mod.functions.values())) + new_body = tvm.tirx.stmt_functor.ir_transform(func.body, None, postorder, ["tirx.For"]) + new_mod = tvm.IRModule({func.attrs["global_symbol"]: func.with_body(new_body)}) + return new_mod + + +def _run_verify_directly(mod): + target = normalize_pto_target("pto") + with target: + return ascend_transform.VerifyParallelToPTO()(mod) + + +# --------------------------------------------------------------------------- +# Kernels +# --------------------------------------------------------------------------- + + +def _k1d(E, out_extent, lanes, dtype="float32"): + """1D scale-by-2 with a padded source and a padded-or-compact output UB.""" + + P = _ceil_to_lanes(E, lanes) + + @T.prim_func + def main( + A: T.Tensor((P,), dtype), + B_init: T.Tensor((out_extent,), dtype), + B: T.Tensor((out_extent,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((P,), dtype) + b = T.alloc_shared((out_extent,), dtype) + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(E): + b[i] = a[i] + a[i] + T.copy(b, B) + + return main + + +def _k2d_row(M, N, lanes, dtype="float32"): + """2D row scale-by-2: j (inner, extent N) is selected and non-divisible; + both UBs use the padded row stride P, a real stride != N.""" + + P = _ceil_to_lanes(N, lanes) + + @T.prim_func + def main( + A: T.Tensor((M * P,), dtype), + B_init: T.Tensor((M * P,), dtype), + B: T.Tensor((M * P,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((M * P,), dtype) + b = T.alloc_shared((M * P,), dtype) + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes), T.Parallel(M, N) as (i, j): + b[i * P + j] = a[i * P + j] + a[i * P + j] + T.copy(b, B) + + return main + + +def _k2d_col(M, N, lanes, dtype="float32"): + """2D column scale-by-2: i (extent M) is selected and non-divisible, j is + kept; row-major (N, P_M) storage accessed c[j * P + i].""" + + P = _ceil_to_lanes(M, lanes) + + @T.prim_func + def main( + A: T.Tensor((N * P,), dtype), + B_init: T.Tensor((N * P,), dtype), + B: T.Tensor((N * P,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((N * P,), dtype) + b = T.alloc_shared((N * P,), dtype) + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes), T.Parallel(M, N) as (i, j): + b[j * P + i] = a[j * P + i] + a[j * P + i] + T.copy(b, B) + + return main + + +def _k_gather(E, lanes, dtype="float32"): + """Gather access: a[2 * i] is neither continuous nor lane-uniform.""" + + P = _ceil_to_lanes(E, lanes) + + @T.prim_func + def main( + A: T.Tensor((P,), dtype), + B: T.Tensor((P,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((P,), dtype) + b = T.alloc_shared((P,), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(E // 2): + b[i] = a[i * 2] + a[i * 2] + T.copy(b, B) + + return main + + +def _k_mixed_vmi(E, lanes, dtype="float32"): + """Hand-written VMI next to a T.Parallel unit in the same SIMD_VF block.""" + + P = _ceil_to_lanes(E, lanes) + + @T.prim_func + def main( + A: T.Tensor((P,), dtype), + B: T.Tensor((P,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((P,), dtype) + b = T.alloc_shared((P,), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + mask = T.vmi.create_mask(lanes, size=lanes) + value = T.vmi.vload(a[0], size=lanes) + T.vmi.vstore(value, b[0], mask) + for i in T.Parallel(E): + b[i] = a[i] + a[i] + T.copy(b, B) + + return main + + +# --------------------------------------------------------------------------- +# Mapping contract +# --------------------------------------------------------------------------- + + +@pytest.mark.pto +def test_1d_nondivisible_keeps_logical_extent_and_pads_fragment(): + E, lanes = 150, 128 + P = _ceil_to_lanes(E, lanes) + mod, _ = _lower_capture(_k1d(E, P, lanes)) + loop = _single_pto_parallel_loop(mod) + + assert _as_int(loop.extent) == E, "the For must keep the logical extent E" + assert _SIMT_LAYOUT_ATTR not in loop.annotations, "the PTO lane layout must not use the SIMT key" + + fragment = _fragment_of(loop) + assert _input_shape(fragment) == [P], "the selected dimension must be padded to P" + assert _as_int(fragment.get_thread_size()) == lanes, "the fragment thread extent must be L" + assert P // lanes == 2, "Q = P / L" + assert E % lanes == 22, "tail_lanes = E % L" + + store = _first_store(loop.body) + assert store.indices[0].same_as(loop.loop_var), "the store must keep the original loop variable and buffer" + + +@pytest.mark.pto +def test_divisible_extent_keeps_e_and_has_no_tail(): + E, lanes = 256, 128 + P = _ceil_to_lanes(E, lanes) + mod, _ = _lower_capture(_k1d(E, P, lanes)) + loop = _single_pto_parallel_loop(mod) + + assert _as_int(loop.extent) == E + fragment = _fragment_of(loop) + assert _input_shape(fragment) == [P] + assert P == E and E % lanes == 0, "divisible extents keep P == E and tail_lanes == 0" + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [64, 128, 256]) +def test_lane_matrix_agrees_with_the_annotation(lanes): + E = 150 + P = _ceil_to_lanes(E, lanes) + mod, _ = _lower_capture(_k1d(E, P, lanes)) + loop = _single_pto_parallel_loop(mod) + + assert _as_int(loop.extent) == E + block = _simd_vf_block(mod) + assert _as_int(block.annotations[_LANES_ATTR]) == lanes, "the lanes annotation must match the request" + fragment = _fragment_of(loop) + assert _input_shape(fragment) == [P] + assert _as_int(fragment.get_thread_size()) == lanes + + +@pytest.mark.pto +def test_2d_selected_dimension_padded_and_kept_dimension_unchanged(): + M, N, lanes = 4, 150, 128 + P = _ceil_to_lanes(N, lanes) + mod, _ = _lower_capture(_k2d_row(M, N, lanes)) + loop = _single_pto_parallel_loop(mod) + + inner = loop.body + assert isinstance(inner, tvm.tirx.For), "expected the two-layer Parallel nest" + assert _as_int(loop.extent) == M, "the outer loop keeps its extent" + assert _as_int(inner.extent) == N, "the selected inner loop keeps its logical extent" + + fragment = _fragment_of(loop) + assert _input_shape(fragment) == [M, P], "only the selected dimension is padded" + assert _as_int(fragment.get_thread_size()) == lanes + + store = _first_store(inner.body) + expected_index = loop.loop_var * P + inner.loop_var + assert tvm.ir.structural_equal(store.indices[0], expected_index), "the store keeps the row-major real stride index i * P + j" + + +@pytest.mark.pto +def test_2d_column_selection_pads_the_selected_dimension_only(): + M, N, lanes = 150, 4, 128 + P = _ceil_to_lanes(M, lanes) + mod, _ = _lower_capture(_k2d_col(M, N, lanes)) + loop = _single_pto_parallel_loop(mod) + + inner = loop.body + assert _as_int(loop.extent) == M + assert _as_int(inner.extent) == N + + fragment = _fragment_of(loop) + assert _input_shape(fragment) == [P, N], "the selected i dimension is padded, j keeps its extent" + assert _as_int(fragment.get_thread_size()) == lanes + + +# --------------------------------------------------------------------------- +# Rejection diagnostics +# --------------------------------------------------------------------------- + + +@pytest.mark.pto +def test_gather_access_is_rejected(): + E, lanes = 150, 128 + _, error = _lower_capture(_k_gather(E, lanes), expect_error=True) + assert error is not None, "a gather access must be rejected" + message = str(error) + assert "continuous" in message or "pto_parallel_loop_layout" in message, f"expected a continuity/layout diagnostic, got: {message}" + + +@pytest.mark.pto +def test_invalid_lanes_is_rejected_by_the_frontend(): + with pytest.raises(ValueError, match="lanes must be one of"): + T.SimdVF(lanes=96) + + +@pytest.mark.pto +def test_missing_pto_lane_layout_is_rejected(): + E, lanes = 150, 128 + P = _ceil_to_lanes(E, lanes) + mod, _ = _lower_capture(_k1d(E, P, lanes)) + stripped_mod = _rewrite_pto_layout_attr(mod, lambda value: None) + with pytest.raises(Exception, match="pto_parallel_loop_layout is missing"): + _run_verify_directly(stripped_mod) + + +@pytest.mark.pto +def test_invalid_lane_annotation_is_rejected_by_verify(): + E, lanes = 150, 128 + P = _ceil_to_lanes(E, lanes) + mod, _ = _lower_capture(_k1d(E, P, lanes)) + assert _as_int(_simd_vf_block(mod).annotations[_LANES_ATTR]) == lanes + + func = next(iter(mod.functions.values())) + + def postorder(node): + if isinstance(node, tvm.tirx.SBlock) and _LANES_ATTR in node.annotations: + annotations = dict(node.annotations) + annotations[_LANES_ATTR] = tvm.tirx.IntImm("int64", 96) + return tvm.tirx.SBlock( + node.iter_vars, + node.reads, + node.writes, + node.name_hint, + node.body, + init=node.init, + alloc_buffers=node.alloc_buffers, + match_buffers=node.match_buffers, + annotations=annotations, + ) + return None + + body = tvm.tirx.stmt_functor.ir_transform(func.body, None, postorder, ["tirx.SBlock"]) + broken_mod = tvm.IRModule({func.attrs["global_symbol"]: func.with_body(body)}) + with pytest.raises(Exception, match="lanes must be one of"): + _run_verify_directly(broken_mod) + + +@pytest.mark.pto +def test_handwritten_vmi_mixed_with_parallel_is_rejected(): + E, lanes = 150, 128 + _, error = _lower_capture(_k_mixed_vmi(E, lanes), expect_error=True) + assert error is not None, "mixing hand-written VMI with T.Parallel must be rejected" + message = str(error) + assert "mixing" in message or "cannot be mixed" in message, f"expected the hand-written VMI mixing diagnostic, got: {message}" + + +# --------------------------------------------------------------------------- +# Op-level dtype matrix for sub/div/unary +# --------------------------------------------------------------------------- + + +def _binary_kernel(n, lanes, dtype, op): + @T.prim_func + def main(A: T.Tensor((n,), dtype), B: T.Tensor((n,), dtype), C: T.Tensor((n,), dtype)): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), dtype) + b_ub = T.alloc_shared((n,), dtype) + c_ub = T.alloc_shared((n,), dtype) + T.copy(A, a_ub) + T.copy(B, b_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + if op == "sub": + c_ub[i] = a_ub[i] - b_ub[i] + else: + c_ub[i] = a_ub[i] / b_ub[i] + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_verify_bf16_div_rejected(): + with pytest.raises(Exception, match=r"vdiv on .*bfloat16.*operation matrix"): + _lower(_binary_kernel(128, 128, "bfloat16", "div"), target="pto") + + +@pytest.mark.pto +def test_verify_uint16_div_rejected(): + with pytest.raises(Exception, match=r"vdiv on uint16.*operation matrix"): + _lower(_binary_kernel(128, 128, "uint16", "div"), target="pto") + + +@pytest.mark.pto +def test_verify_fp8_sub_rejected(): + with pytest.raises(Exception, match=r"FP8 arithmetic is not supported"): + _lower(_binary_kernel(128, 128, "float8_e4m3fn", "sub"), target="pto") + + +@pytest.mark.pto +def test_verify_fp8_div_rejected(): + with pytest.raises(Exception, match=r"FP8 arithmetic is not supported"): + _lower(_binary_kernel(128, 128, "float8_e5m2", "div"), target="pto") diff --git a/testing/ascend/language/pto/test_vectorize_fp8_run.py b/testing/ascend/language/pto/test_vectorize_fp8_run.py new file mode 100644 index 0000000000..28be2f3e5f --- /dev/null +++ b/testing/ascend/language/pto/test_vectorize_fp8_run.py @@ -0,0 +1,609 @@ +"""Device tests for FP8 casts on the Parallel->VMI path (NPU). + +Two things are checked here, and they are reported separately: the attributes +the vectorize path emits (rounded/saturation mode spelled out in the generated +PTODSL), and the numeric result on the device. + +Reference values come from a host-side E4M3FN/E5M2 code table evaluated with +the *selected* mode — round-to-nearest-even or toward-zero, with or without +saturation — not from torch's own fp8 cast: the contract under test is the PTO +conversion mode, and torch has no toward-zero mode to compare against. NaN +encoding is deliberately not asserted in this batch; the overflow cases assert +either the documented saturation result (SAT) or nothing but the attribute and +successful execution (NOSAT). + +The FP8 conversion instructions the target exposes are FP32 <-> FP8, so a +half source is rejected by Verify and the supported way to convert one is the +explicit two-step cast written in the kernel (half -> float32 -> FP8), which +carries the numeric coverage here. +""" + +import math +import os +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import tilelang # noqa: E402 +import tilelang.ascend.language as T # noqa: E402 + +F8_E4M3 = T.float8_e4m3fn +F8_E5M2 = T.float8_e5m2 +TORCH_F8 = {F8_E4M3: torch.float8_e4m3fn, F8_E5M2: torch.float8_e5m2} + + +def _require_npu_runtime(): + if not hasattr(torch, "npu"): + pytest.skip("torch.npu is unavailable in the current environment") + try: + _ = torch.randn(1, device=torch.device("npu")) + torch.npu.synchronize() + except Exception as err: # noqa: BLE001 + pytest.skip(f"NPU runtime is unavailable: {err}") + + +@pytest.fixture() +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +# --------------------------------------------------------------------------- +# Host-side E4M3FN / E5M2 code tables +# --------------------------------------------------------------------------- +def _decode_e4m3(code): + """E4M3FN: 1 sign, 4 exponent (bias 7), 3 mantissa; exp=15/m=7 is NaN.""" + sign = -1.0 if code & 0x80 else 1.0 + exp = (code >> 3) & 0xF + mant = code & 0x7 + if exp == 0: + return sign * mant * 2.0**-9 + if exp == 0xF and mant == 0x7: + return None # NaN + return sign * (1.0 + mant / 8.0) * 2.0 ** (exp - 7) + + +def _decode_e5m2(code): + """E5M2: 1 sign, 5 exponent (bias 15), 2 mantissa; exp=31 is inf/NaN.""" + sign = -1.0 if code & 0x80 else 1.0 + exp = (code >> 2) & 0x1F + mant = code & 0x3 + if exp == 0: + return sign * mant * 2.0**-16 + if exp == 0x1F: + return None # inf / NaN + return sign * (1.0 + mant / 4.0) * 2.0 ** (exp - 15) + + +def _table(decode): + """All finite codes of a format, sorted by magnitude, positive side.""" + vals = [] + for code in range(256): + value = decode(code) + if value is not None and (value > 0 or code == 0): + vals.append((code, value)) + vals.sort(key=lambda pair: pair[1]) + return vals + + +TABLES = {F8_E4M3: _table(_decode_e4m3), F8_E5M2: _table(_decode_e5m2)} + + +def _max_finite(f8): + return TABLES[f8][-1][1] + + +def _encode(value, f8, mode): + """Encode one float with the given PTO mode (R = ties-to-even, Z = toward + zero), saturating.""" + table = TABLES[f8] + limit = table[-1][1] + sign_bit = 0x80 if math.copysign(1.0, value) < 0 else 0 + value = max(-limit, min(limit, abs(value))) + mag = abs(value) + if mode == "R": + # Exact nearest with ties-to-even: the magnitudes are dyadic rationals, + # so the deltas are exact and comparisons need no tolerance. + code = table[0][0] + best_delta = abs(table[0][1] - mag) + for cand, magnitude in table[1:]: + delta = abs(magnitude - mag) + if delta < best_delta or (delta == best_delta and (cand & 1) == 0 and (code & 1) != 0): + code, best_delta = cand, delta + elif mode == "Z": + # Truncation: the largest magnitude that is <= the input, compared + # exactly (a tolerance would round nextafter(1.0, 0) up to 1.0). + code = table[0][0] + for cand, magnitude in table: + if magnitude <= mag: + code = cand + else: + break + else: + raise ValueError(mode) + return code | sign_bit + + +def _encode_mask(values, f8, mode): + return bytes(_encode(float(v), f8, mode) for v in values) + + +def _assert_codes_with_overflow(out, ref_values, f8): + """Per-lane byte check: in-range lanes keep exact RNE codes; overflow + lanes (|v| > max finite) accept the saturated or NaN encoding. Under the + global CTRL saturation mode forced by pto.init_core() the instruction's + satMode bit is overridden, so the overflow encoding is + environment-defined; no existing fp8 test in the repo covers SAT + overflow, so neither do we.""" + limit = _max_finite(f8) + got = out.cpu().view(torch.uint8).flatten().tolist() + assert len(got) == len(ref_values) + for i, v in enumerate(ref_values): + if abs(v) > limit: + # e4m3fn overflow encodings: saturated ±MAX (0x7E/0xFE) or NaN + # (0x7F/0xFF; NaN does not necessarily preserve the sign). + assert f8 == F8_E4M3, "tolerant overflow codes are e4m3fn-specific" + assert got[i] in (0x7E, 0x7F, 0xFE, 0xFF), f"lane {i}: overflow encoding {got[i]:#x}" + else: + assert got[i] == _encode(float(v), f8, "R"), f"lane {i}: {got[i]:#x} vs {float(v)}" + + +# --------------------------------------------------------------------------- +# Kernels +# --------------------------------------------------------------------------- +def _f32_to_f8(f8, n, lanes, **kwargs): + ann = kwargs.pop("annotations", None) + + @T.prim_func + def main(A: T.Buffer((n,), "float32"), C: T.Buffer((n,), f8)): + with T.Kernel(1): + a = T.alloc_shared((n,), "float32") + c = T.alloc_shared((n,), f8) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.Cast(f8, a[i], annotations=ann) if ann is not None else T.cast(a[i], f8, **kwargs) + T.copy(c, C) + + return main + + +def _f8_to_f32(f8, n, lanes): + @T.prim_func + def main(A: T.Buffer((n,), f8), C: T.Buffer((n,), "float32")): + with T.Kernel(1): + a = T.alloc_shared((n,), f8) + c = T.alloc_shared((n,), "float32") + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.cast(a[i], "float32") + T.copy(c, C) + + return main + + +def _run(func, npu, source=None): + kernel = tilelang.compile(func, target="pto", out_idx=-1) + src = kernel.get_kernel_source() + if source is not None: + source(src) + return src + + +def _bytes(tensor): + return bytes(tensor.cpu().view(torch.uint8).flatten().tolist()) + + +LANES = 64 + + +def _pad(values, fill=0.0): + """The PTO path requires the parallel extent to divide L, so the value + lists are padded to a multiple of LANES; padding encodes deterministically + (0.0 -> code 0).""" + if len(values) % LANES: + values = list(values) + [fill] * (LANES - len(values) % LANES) + return values + + +# --------------------------------------------------------------------------- +# Attributes in the generated PTODSL +# --------------------------------------------------------------------------- +@pytest.mark.pto +@pytest.mark.parametrize("f8", [F8_E4M3, F8_E5M2], ids=["e4m3fn", "e5m2"]) +def test_generated_attributes_default(npu, f8): + n, lanes = 256, 128 + src = _run(_f32_to_f8(f8, n, lanes), npu) + target = "pto.f8e4m3" if f8 == F8_E4M3 else "pto.f8e5m2" + line = [ln.strip() for ln in src.split("\n") if "vmi.vcvt" in ln and "=" in ln] + assert line, src + assert 'rounding="R"' in line[0] and 'saturate="SAT"' in line[0], line[0] + assert f"to_dtype={target}" in line[0], line[0] + assert "to_dtype=pto.f32," not in line[0], line[0] + + +@pytest.mark.pto +@pytest.mark.parametrize("mode,rounding", [("rn", "R"), ("rz", "Z")]) +def test_generated_attributes_rounding(npu, mode, rounding): + src = _run(_f32_to_f8(F8_E4M3, 256, 128, round=mode), npu) + line = [ln.strip() for ln in src.split("\n") if "vmi.vcvt" in ln and "=" in ln][0] + assert f'rounding="{rounding}"' in line, line + + +@pytest.mark.pto +def test_generated_attributes_nosat(npu): + src = _run(_f32_to_f8(F8_E4M3, 256, 128, sat=False), npu) + line = [ln.strip() for ln in src.split("\n") if "vmi.vcvt" in ln and "=" in ln][0] + assert 'saturate="NOSAT"' in line, line + + +@pytest.mark.pto +def test_generated_attributes_widening_has_none(npu): + src = _run(_f8_to_f32(F8_E4M3, 256, 128), npu) + line = [ln.strip() for ln in src.split("\n") if "vmi.vcvt" in ln and "=" in ln][0] + assert "rounding" not in line and "saturate" not in line, line + assert "to_dtype=pto.f32" in line, line + + +# --------------------------------------------------------------------------- +# Numeric results +# --------------------------------------------------------------------------- +def _midpoints(f8): + """Values exactly halfway between adjacent magnitudes — the tie cases that + separate ties-to-even from ties-away, plus ordinary values.""" + table = TABLES[f8] + values = [-1.0, 0.0, 1.0, -2.5, 3.75] + for (_, lo), (_, hi) in zip(table, table[1:]): + values.append((lo + hi) / 2.0) + values.append(-(lo + hi) / 2.0) + return values + + +@pytest.mark.pto +@pytest.mark.parametrize("f8", [F8_E4M3, F8_E5M2], ids=["e4m3fn", "e5m2"]) +def test_narrow_default_rne_sat(npu, f8): + """Default mode: ties-to-even with saturation, checked on tie cases.""" + values = _pad(_midpoints(f8)) + x = torch.tensor(values, dtype=torch.float32) + out = tilelang.compile(_f32_to_f8(f8, len(values), LANES), target="pto", out_idx=-1)(x.to(npu)) + torch.npu.synchronize() + assert _bytes(out) == _encode_mask(values, f8, "R") + + +@pytest.mark.pto +def test_narrow_toward_zero(npu): + values = _pad(_midpoints(F8_E4M3)) + x = torch.tensor(values, dtype=torch.float32) + out = tilelang.compile(_f32_to_f8(F8_E4M3, len(values), LANES, round="rz"), target="pto", out_idx=-1)(x.to(npu)) + torch.npu.synchronize() + assert _bytes(out) == _encode_mask(values, F8_E4M3, "Z") + + +def _f16_to_f8_two_step(f8, n, lanes): + """Explicit two-step half -> float32 -> FP8: both conversions are + supported directions, so this is the accepted way from a half source.""" + + @T.prim_func + def main(A: T.Buffer((n,), "float16"), C: T.Buffer((n,), f8)): + with T.Kernel(1): + a = T.alloc_shared((n,), "float16") + c = T.alloc_shared((n,), f8) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.cast(T.cast(a[i], "float32"), f8) + T.copy(c, C) + + return main + + +@pytest.mark.pto +def test_half_source_two_step_numeric(npu): + """half -> float32 -> FP8 (written explicitly) matches the ties-to-even + reference.""" + values = _pad(_midpoints(F8_E4M3)) + x = torch.tensor(values, dtype=torch.float32).half() + out = tilelang.compile(_f16_to_f8_two_step(F8_E4M3, len(values), LANES), target="pto", out_idx=-1)(x.to(npu)) + torch.npu.synchronize() + assert _bytes(out) == _encode_mask([float(v) for v in x.tolist()], F8_E4M3, "R") + + +@pytest.mark.pto +@pytest.mark.parametrize("f8", [F8_E4M3, F8_E5M2], ids=["e4m3fn", "e5m2"]) +def test_widen_is_exact(npu, f8): + """Every finite code decodes exactly: fp8 -> fp32 is a widening.""" + decode = _decode_e4m3 if f8 == F8_E4M3 else _decode_e5m2 + codes = [c for c in range(256) if decode(c) is not None] + codes = _pad(codes, fill=0) + raw = torch.tensor(codes, dtype=torch.uint8).view(TORCH_F8[f8]) + expected = torch.tensor([decode(c) for c in codes], dtype=torch.float32) + out = tilelang.compile(_f8_to_f32(f8, len(codes), LANES), target="pto", out_idx=-1)(raw.to(npu)) + torch.npu.synchronize() + assert torch.equal(out.cpu(), expected), (out.cpu()[:8], expected[:8]) + + +@pytest.mark.pto +@pytest.mark.parametrize("f8,below_code", [(F8_E4M3, 0x37), (F8_E5M2, 0x3B)], ids=["e4m3fn", "e5m2"]) +def test_toward_zero_on_fp32_predecessor(npu, f8, below_code): + """Z applied to the true float32 predecessor of 1.0. + + The input has to be built in float32: `math.nextafter` works in double and + its result rounds back to exactly 1.0 when stored as float32, which would + exercise nothing. The bit pattern is asserted so the case cannot silently + regress into a no-op, and the expectation is the single code below 1.0 + (0x37 for E4M3FN, 0x3B for E5M2) with the negative side mirrored. + """ + one = torch.tensor(1.0, dtype=torch.float32) + below = torch.nextafter(one, torch.tensor(0.0, dtype=torch.float32)) + assert below.view(torch.int32).item() == 0x3F7FFFFF, hex(below.view(torch.int32).item()) + assert below.item() < 1.0 + assert below.half().float().item() != below.item() # not on the f16 grid + + values = _pad([below.item(), -below.item()]) + x = torch.tensor(values, dtype=torch.float32) + assert x[0].view(torch.int32).item() == 0x3F7FFFFF + + kernel = tilelang.compile(_f32_to_f8(f8, len(values), LANES, round="rz"), target="pto", out_idx=-1) + line = [ln.strip() for ln in kernel.get_kernel_source().split("\n") if "vmi.vcvt" in ln and "=" in ln][0] + assert 'rounding="Z"' in line, line + + out = kernel(x.to(npu)) + torch.npu.synchronize() + got = _bytes(out) + assert got[0] == below_code, f"expected {below_code:#x}, got {got[0]:#x}" + assert got[1] == below_code | 0x80, f"expected {below_code | 0x80:#x}, got {got[1]:#x}" + + +@pytest.mark.pto +def test_empty_round_annotation_uses_default(npu): + """`{"round": ""}` is the frontend's "backend default", not an unknown + mode: the narrowing keeps R/SAT, and the values match the default run.""" + values = _pad(_midpoints(F8_E4M3)) + x = torch.tensor(values, dtype=torch.float32) + kernel = tilelang.compile( + _f32_to_f8(F8_E4M3, len(values), LANES, annotations={"round": ""}), + target="pto", + out_idx=-1, + ) + line = [ln.strip() for ln in kernel.get_kernel_source().split("\n") if "vmi.vcvt" in ln and "=" in ln][0] + assert 'rounding="R"' in line and 'saturate="SAT"' in line, line + out = kernel(x.to(npu)) + torch.npu.synchronize() + assert _bytes(out) == _encode_mask(values, F8_E4M3, "R") + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [64, 128, 256]) +def test_narrow_lane_widths_rne_sat(npu, lanes): + """The numeric contract holds at every supported lane width, including the + multi-chunk case (extent = 2 * lanes).""" + values = _midpoints(F8_E4M3) + values = values * (1 + (2 * lanes) // max(len(values), 1)) if len(values) < 2 * lanes else values + values = _pad(values[: 2 * lanes], 0.0) + x = torch.tensor(values, dtype=torch.float32) + out = tilelang.compile(_f32_to_f8(F8_E4M3, len(values), lanes), target="pto", out_idx=-1)(x.to(npu)) + torch.npu.synchronize() + assert _bytes(out) == _encode_mask(values, F8_E4M3, "R") + + +def _f32_to_fp8_2d_with_ij(f8, m, n, lanes): + """2D column-continuous element-wise unit (buffers (n, m), accessed + [j, i]) that mixes coordinate data before narrowing to FP8: + c[j, i] = fp8(a[j, i] * s[i] + t[j]). + + `s[i]` is lane-varying data along the vectorized coordinate, `t[j]` is + lane-uniform data along the kept coordinate, and both arithmetic steps plus + the narrowing are written in the source order — no intermediate dtype. + """ + + @T.prim_func + def main( + A: T.Buffer((n, m), "float32"), + S: T.Buffer((m,), "float32"), + Tt: T.Buffer((n,), "float32"), + C: T.Buffer((n, m), f8), + ): + with T.Kernel(1): + a = T.alloc_shared((n, m), "float32") + s = T.alloc_shared((m,), "float32") + t = T.alloc_shared((n,), "float32") + c = T.alloc_shared((n, m), f8) + T.copy(A, a) + T.copy(S, s) + T.copy(Tt, t) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(m, n): + c[j, i] = T.cast(a[j, i] * s[i] + t[j], f8) + T.copy(c, C) + + return main + + +def _grid_values(count, step, low): + """Values whose products and sums stay exactly representable in both f16 + and f32, so the narrowing result does not depend on an internal precision + hop (the Z-mode finding stays out of this case).""" + return [low + step * k for k in range(count)] + + +@pytest.mark.pto +@pytest.mark.parametrize("f8", [F8_E4M3, F8_E5M2], ids=["e4m3fn", "e5m2"]) +def test_2d_elementwise_with_coordinate_data(npu, f8): + """_column_kernel shape, element-wise mul-add with coordinate data, then a + narrowing cast to FP8; in-range lanes are compared against the strict + RNE+SAT reference, overflow lanes accept the environment-defined + saturated/NaN encoding (see _assert_codes_with_overflow).""" + m, n, lanes = 256, 4, 128 + a_flat = _grid_values(n * m, 0.25, -6.0) # (n, m), step 1/4 + s_vals = _grid_values(m, 0.25, 0.5) # (m,), step 1/4 + t_vals = _grid_values(n, 0.5, -2.0) # (n,), step 1/2 + + a = torch.tensor(a_flat, dtype=torch.float32).reshape(n, m) + s_t = torch.tensor(s_vals, dtype=torch.float32) + t_t = torch.tensor(t_vals, dtype=torch.float32) + + kernel = tilelang.compile(_f32_to_fp8_2d_with_ij(f8, m, n, lanes), target="pto", out_idx=-1) + src = kernel.get_kernel_source() + target = "pto.f8e4m3" if f8 == F8_E4M3 else "pto.f8e5m2" + vcvt = [ln.strip() for ln in src.split("\n") if "vmi.vcvt" in ln and "=" in ln] + assert len(vcvt) == 1, f"expected exactly one conversion:\n{src}" + assert f"to_dtype={target}" in vcvt[0], vcvt[0] + assert 'rounding="R"' in vcvt[0] and 'saturate="SAT"' in vcvt[0], vcvt[0] + assert "vmi.vmul" in src and "vmi.vadd" in src and "vmi.vbrc" in src, src + + out = kernel(a.to(npu), s_t.to(npu), t_t.to(npu)) + torch.npu.synchronize() + + # Reference in float32, in the source's operation order, then RNE+SAT + # (overflow lanes tolerant, see _assert_codes_with_overflow). + ref = (a * s_t.reshape(1, m) + t_t.reshape(n, 1)).reshape(-1).tolist() + _assert_codes_with_overflow(out, ref, f8) + + +@pytest.mark.pto +def test_offset_counts_once(npu): + """A constant source offset (the chunk start) must be counted exactly + once: c[i] = f8(a[i + 8]) over a longer source buffer.""" + n, lanes, shift = 2 * 64, 64, 8 + values = [float(i) * 0.5 - 16.0 for i in range(n + shift)] + x = torch.tensor(values, dtype=torch.float32) + src = _offset_kernel(F8_E4M3, n, lanes, shift) + kernel = tilelang.compile(src, target="pto", out_idx=-1) + line = [ln.strip() for ln in kernel.get_kernel_source().split("\n") if "vmi.vload" in ln and "=" in ln][0] + assert line.count("+ 8") == 1, line + out = kernel(x.to(npu)) + torch.npu.synchronize() + expected = [values[i + shift] for i in range(n)] + assert _bytes(out) == _encode_mask(expected, F8_E4M3, "R") + + +def _offset_kernel(f8, n, lanes, shift): + """1D narrowing reading the source at a constant offset.""" + + @T.prim_func + def main(A: T.Buffer((n + shift,), "float32"), C: T.Buffer((n,), f8)): + with T.Kernel(1): + a = T.alloc_shared((n + shift,), "float32") + c = T.alloc_shared((n,), f8) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.cast(a[i + shift], f8) + T.copy(c, C) + + return main + + +def _f32_to_f8_2d(f8, m, n, lanes, column): + """2D narrowing, j-continuous (buffer (m, n)) or i-continuous (buffer + (n, m), accessed [j, i]).""" + + if column: + + @T.prim_func + def main(A: T.Buffer((n, m), "float32"), C: T.Buffer((n, m), f8)): + with T.Kernel(1): + a = T.alloc_shared((n, m), "float32") + c = T.alloc_shared((n, m), f8) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(m, n): + c[j, i] = T.cast(a[j, i], f8) + T.copy(c, C) + + else: + + @T.prim_func + def main(A: T.Buffer((m, n), "float32"), C: T.Buffer((m, n), f8)): + with T.Kernel(1): + a = T.alloc_shared((m, n), "float32") + c = T.alloc_shared((m, n), f8) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(m, n): + c[i, j] = T.cast(a[i, j], f8) + T.copy(c, C) + + return main + + +@pytest.mark.pto +@pytest.mark.parametrize("column", [False, True], ids=["row", "col"]) +def test_2d_numeric(npu, column): + """Both continuous directions narrow correctly with fp8 (L=128).""" + m, n, lanes = (256, 4, 128) if column else (4, 256, 128) + shape = (n, m) if column else (m, n) + values = [((i % 512) - 256) / 16.0 for i in range(shape[0] * shape[1])] + x = torch.tensor(values, dtype=torch.float32).reshape(shape) + out = tilelang.compile(_f32_to_f8_2d(F8_E4M3, m, n, lanes, column), target="pto", out_idx=-1)(x.to(npu)) + torch.npu.synchronize() + assert _bytes(out) == _encode_mask(values, F8_E4M3, "R") + + +@pytest.mark.pto +@pytest.mark.parametrize("token", ["A", "H"]) +def test_rounding_tokens_compile_and_run(npu, token): + """A and H are carried through to the VMI call; this build compiles and + executes them. Their numeric contract belongs to PTOAS, so only the + attribute and successful execution are asserted.""" + values = _pad(_midpoints(F8_E4M3)) + x = torch.tensor(values, dtype=torch.float32) + kernel = tilelang.compile( + _f32_to_f8(F8_E4M3, len(values), LANES, annotations={"round": token}), + target="pto", + out_idx=-1, + ) + line = [ln.strip() for ln in kernel.get_kernel_source().split("\n") if "vmi.vcvt" in ln and "=" in ln][0] + assert f'rounding="{token}"' in line, line + out = kernel(x.to(npu)) + torch.npu.synchronize() + assert out.cpu().view(torch.uint8).numel() == len(values) + + +@pytest.mark.pto +def test_sat_clamps_finite_overflow(npu): + """Overflow lanes: encoding is environment-defined, in-range lanes exact. + + Saturation on overflow depends on the inherited CTRL[60]/CTRL[48] mode: + pto.init_core() forces CTRL[60]=1 (global mode), under which the + instruction's satMode bit is overridden by the global flag, so an + overflowing lane may encode as saturated ±MAX (0x7E/0xFE) or NaN + (0x7F/0xFF) depending on the environment. No existing fp8 test in the + repo covers SAT overflow, so this case asserts only that overflow lanes + produce one of the two legal encodings; padding lanes stay exact. + """ + limit = _max_finite(F8_E4M3) + values = _pad([limit * 2.0, limit * 100.0, -limit * 2.0, -limit * 100.0]) + x = torch.tensor(values, dtype=torch.float32) + out = tilelang.compile(_f32_to_f8(F8_E4M3, len(values), LANES), target="pto", out_idx=-1)(x.to(npu)) + torch.npu.synchronize() + _assert_codes_with_overflow(out, values, F8_E4M3) + + +@pytest.mark.pto +def test_nosat_overflow_attributes_only(npu): + """With sat=False the overflow encoding is backend-defined in this batch: + assert the emitted attribute and that the kernel executes, not the bytes.""" + limit = _max_finite(F8_E4M3) + values = _pad([limit * 2.0, -limit * 2.0]) + x = torch.tensor(values, dtype=torch.float32) + kernel = tilelang.compile(_f32_to_f8(F8_E4M3, len(values), LANES, sat=False), target="pto", out_idx=-1) + line = [ln.strip() for ln in kernel.get_kernel_source().split("\n") if "vmi.vcvt" in ln and "=" in ln][0] + assert 'saturate="NOSAT"' in line, line + out = kernel(x.to(npu)) + torch.npu.synchronize() + assert out.cpu().view(torch.uint8).numel() == len(values) + + +if __name__ == "__main__": + import tilelang.testing + + tilelang.testing.main() diff --git a/testing/ascend/language/pto/test_vectorize_parallel_pto_run.py b/testing/ascend/language/pto/test_vectorize_parallel_pto_run.py new file mode 100644 index 0000000000..7e612dd030 --- /dev/null +++ b/testing/ascend/language/pto/test_vectorize_parallel_pto_run.py @@ -0,0 +1,934 @@ +"""End-to-end NPU tests for the Parallel->VMI path. + +These kernels contain ``T.Parallel`` inside ``T.SimdVF(lanes=...)`` and run +through the real ``tilelang.compile(target="pto")`` pipeline, which routes +them to ``VerifyParallelToPTO`` + ``VectorizeParallelToPTO`` instead of +``AscendSimdVFLowerParallel``. Auto-schedule stays enabled: the non-autosched +pipeline has no InsertSync (MTE sync), so numeric results would be garbage. +""" + +import os +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import tilelang +import tilelang.ascend.language as T + +from tilelang import tvm +from tilelang.engine.lower import lower + + +def _require_npu_runtime(): + if not hasattr(torch, "npu"): + pytest.skip("torch.npu is unavailable in the current environment") + + try: + _ = torch.randn(1, device=torch.device("npu")) + torch.npu.synchronize() + except Exception as err: + pytest.skip(f"NPU runtime is unavailable in the current environment: {err}") + + +def _vl_padded(n, lanes): + """Extent padded to whole VL chunks. Tail vloads read the full VL from + their chunk base (loads are unmasked by ISA design; the reserved top 8K + of UB keeps the extra read-only lanes fault-free), and the padded input + makes this test's readable storage explicit. The Parallel extent and + the compact output stay at ``n``.""" + return (n + lanes - 1) // lanes * lanes + + +def _vecadd(n, lanes): + """C = A * (A + B) over a 1D Parallel, multi-chunk when n > lanes.""" + + @T.prim_func + def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32"), C: T.Buffer((n,), "float32")): + with T.Kernel(1): + t1 = T.alloc_shared((n,), "float32") + t2 = T.alloc_shared((n,), "float32") + t3 = T.alloc_shared((n,), "float32") + T.copy(A, t1) + T.copy(B, t2) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + t3[i] = t1[i] * (t1[i] + t2[i]) + T.copy(t3, C) + + return main + + +def _i32_index(n, lanes): + """C[i] = A[i] + i: int32 loop index feeds the value stream, exercising the + integer VCI signless bridge (vinterpret_cast) on the new path.""" + + @T.prim_func + def main(A: T.Buffer((n,), "int32"), C: T.Buffer((n,), "int32")): + with T.Kernel(1): + a = T.alloc_shared((n,), "int32") + c = T.alloc_shared((n,), "int32") + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = a[i] + i + T.copy(c, C) + + return main + + +def _row_kernel(m, n, lanes): + """2D j-continuous: c[i, j] = a[i, j] + i + j over T.Parallel(m, n).""" + + @T.prim_func + def main(A: T.Buffer((m, n), "int32"), C: T.Buffer((m, n), "int32")): + with T.Kernel(1): + a = T.alloc_shared((m, n), "int32") + c = T.alloc_shared((m, n), "int32") + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(m, n): + c[i, j] = a[i, j] + i + j + T.copy(c, C) + + return main + + +def _column_kernel(m, n, lanes): + """2D i-continuous: buffers are (n, m) and accessed [j, i].""" + + @T.prim_func + def main(A: T.Buffer((n, m), "int32"), C: T.Buffer((n, m), "int32")): + with T.Kernel(1): + a = T.alloc_shared((n, m), "int32") + c = T.alloc_shared((n, m), "int32") + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(m, n): + c[j, i] = a[j, i] + i + j + T.copy(c, C) + + return main + + +def _cast_f16_muladd_1d(n, lanes): + """1D widening cast: c[i] = f32(a[i]) * (f32(a[i]) + f32(b[i])).""" + + @T.prim_func + def main(A: T.Buffer((n,), "float16"), B: T.Buffer((n,), "float16"), C: T.Buffer((n,), "float32")): + with T.Kernel(1): + a = T.alloc_shared((n,), "float16") + b = T.alloc_shared((n,), "float16") + c = T.alloc_shared((n,), "float32") + T.copy(A, a) + T.copy(B, b) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.Cast("float32", a[i]) * (T.Cast("float32", a[i]) + T.Cast("float32", b[i])) + T.copy(c, C) + + return main + + +def _cast_bf16_add_1d(n, lanes): + """1D bf16 -> f32: c[i] = f32(a[i]) + f32(a[i]).""" + + @T.prim_func + def main(A: T.Buffer((n,), "bfloat16"), C: T.Buffer((n,), "float32")): + with T.Kernel(1): + a = T.alloc_shared((n,), "bfloat16") + c = T.alloc_shared((n,), "float32") + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.Cast("float32", a[i]) + T.Cast("float32", a[i]) + T.copy(c, C) + + return main + + +def _cast_i16_square_1d(n, lanes): + """1D int16 -> int32 widening: c[i] = i32(a[i]) * i32(a[i]).""" + + @T.prim_func + def main(A: T.Buffer((n,), "int16"), C: T.Buffer((n,), "int32")): + with T.Kernel(1): + a = T.alloc_shared((n,), "int16") + c = T.alloc_shared((n,), "int32") + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.Cast("int32", a[i]) * T.Cast("int32", a[i]) + T.copy(c, C) + + return main + + +def _cast_f16_add_2d(m, n, lanes): + """2D j-continuous widening cast: c[i, j] = f32(a[i, j]) + f32(b[i, j]).""" + + @T.prim_func + def main(A: T.Buffer((m, n), "float16"), B: T.Buffer((m, n), "float16"), C: T.Buffer((m, n), "float32")): + with T.Kernel(1): + a = T.alloc_shared((m, n), "float16") + b = T.alloc_shared((m, n), "float16") + c = T.alloc_shared((m, n), "float32") + T.copy(A, a) + T.copy(B, b) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(m, n): + c[i, j] = T.Cast("float32", a[i, j]) + T.Cast("float32", b[i, j]) + T.copy(c, C) + + return main + + +def _fp8_cast_1d(n, lanes): + """Cast to fp8: c[i] = f8(a[i]). Outside the supported dtype set.""" + + @T.prim_func + def main(A: T.Buffer((n,), "float32"), C: T.Buffer((n,), "float8_e4m3")): + with T.Kernel(1): + a = T.alloc_shared((n,), "float32") + c = T.alloc_shared((n,), "float8_e4m3") + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.Cast("float8_e4m3", a[i]) + T.copy(c, C) + + return main + + +def _i8_add_1d(n, lanes): + """8-bit control: c[i] = a[i] + b[i] over int8, a whitelisted category.""" + + @T.prim_func + def main(A: T.Buffer((n,), "int8"), B: T.Buffer((n,), "int8"), C: T.Buffer((n,), "int8")): + with T.Kernel(1): + a = T.alloc_shared((n,), "int8") + b = T.alloc_shared((n,), "int8") + c = T.alloc_shared((n,), "int8") + T.copy(A, a) + T.copy(B, b) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = a[i] + b[i] + T.copy(c, C) + + return main + + +def _new_path_signatures(source): + """Signatures proving the kernel took the new Parallel->VMI path.""" + return { + "vload": source.count("pto.vmi.vload("), + "vstore": source.count("pto.vmi.vstore("), + "create_mask": source.count("pto.vmi.create_mask("), + "vlds": source.count(".vlds("), + "set_flag": source.count("set_flag"), + "wait_flag": source.count("wait_flag"), + } + + +def _assert_new_path(source, *, expect_vinterpret_cast=False, expect_scalar_cast=False, expect_vcvt_to=None): + """The compiled PTODSL must come from VectorizeParallelToPTO: VMI ops with + full mask, MTE sync from auto-schedule, and no legacy simd vlds.""" + sig = _new_path_signatures(source) + assert sig["vload"] > 0 and sig["vstore"] > 0, f"no VMI memory ops in source:\n{source}" + assert sig["create_mask"] > 0, f"no create_mask in source:\n{source}" + assert sig["vlds"] == 0, f"legacy simd vlds leaked into source:\n{source}" + assert sig["set_flag"] > 0 and sig["wait_flag"] > 0, ( + f"no MTE sync (set_flag/wait_flag) in source; auto-schedule must stay enabled for numeric correctness:\n{source}" + ) + if expect_vinterpret_cast: + assert "pto.vmi.vinterpret_cast(" in source, f"integer VCI bridge (vinterpret_cast) missing:\n{source}" + if expect_scalar_cast: + # PR269 emitted `scalar.cast(...)` here; the current PTOAS PTODSL surface + # exposes the same cast as `pto.cast(...)`, so the bridge marker is the + # cast nested directly in the reinterpreted VMI source. + bridged = "vinterpret_cast(pto.vmi.vci(pto.cast(" in source or "vinterpret_cast(pto.vmi.vbrc(pto.cast(" in source + assert bridged, f"dynamic integer bridge (pto.cast inside vinterpret_cast) missing:\n{source}" + if expect_vcvt_to is not None: + assert "pto.vmi.vcvt(" in source, f"element-wise conversion (vcvt) missing:\n{source}" + assert f"to_dtype=pto.{expect_vcvt_to}" in source, f"vcvt target {expect_vcvt_to} missing:\n{source}" + + +@pytest.fixture() +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +def _run_and_check(func, inputs, expected, source_checks=None, atol=None): + kernel = tilelang.compile(func, target="pto", out_idx=-1) + source = kernel.get_kernel_source() + if source_checks is not None: + source_checks(source) + out = kernel(*inputs) + torch.npu.synchronize() + got = out.cpu() + if atol is None: + assert torch.equal(got, expected), f"numeric mismatch: max_diff={torch.max(torch.abs((got.float() - expected.float()))).item()}" + else: + max_diff = torch.max(torch.abs(got.float() - expected.float())).item() + assert torch.allclose(got.float(), expected.float(), rtol=0, atol=atol), f"numeric mismatch: max_diff={max_diff}" + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [64, 128, 256]) +def test_vecadd_1d_multichunk(npu, lanes): + """1D float32, n=512: multi-chunk (q loop) across all three lane widths.""" + n = 512 + a = torch.randn(n) + b = torch.randn(n) + _run_and_check( + _vecadd(n, lanes), + (a.to(npu), b.to(npu)), + a * (a + b), + source_checks=lambda s: _assert_new_path(s), + ) + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [128, 256]) +def test_vecadd_1d_singlechunk(npu, lanes): + """1D float32, n=lanes: single chunk, VCI base folds to a constant.""" + n = lanes + a = torch.randn(n) + b = torch.randn(n) + _run_and_check( + _vecadd(n, lanes), + (a.to(npu), b.to(npu)), + a * (a + b), + source_checks=lambda s: _assert_new_path(s), + ) + + +@pytest.mark.pto +def test_i32_index_multichunk(npu): + """1D int32 index in the value stream, multi-chunk: dynamic base bridge.""" + n = 512 + a = (torch.randn(n) * 1000).to(torch.int32) + idx = torch.arange(n, dtype=torch.int32) + _run_and_check( + _i32_index(n, 128), + (a.to(npu),), + a + idx, + source_checks=lambda s: _assert_new_path(s, expect_vinterpret_cast=True, expect_scalar_cast=True), + ) + + +@pytest.mark.pto +def test_i32_index_singlechunk(npu): + """1D int32 index, single chunk: constant base goes through the bridge too.""" + n = 128 + a = (torch.randn(n) * 1000).to(torch.int32) + idx = torch.arange(n, dtype=torch.int32) + _run_and_check( + _i32_index(n, 128), + (a.to(npu),), + a + idx, + source_checks=lambda s: _assert_new_path(s, expect_vinterpret_cast=True), + ) + + +@pytest.mark.pto +def test_2d_row_continuous(npu): + """2D j-continuous unit over T.Parallel(4, 256).""" + m, n = 4, 256 + a = (torch.randn(m, n) * 100).to(torch.int32) + i = torch.arange(m, dtype=torch.int32).unsqueeze(1) + j = torch.arange(n, dtype=torch.int32).unsqueeze(0) + _run_and_check( + _row_kernel(m, n, 128), + (a.to(npu),), + a + i + j, + source_checks=lambda s: _assert_new_path(s, expect_vinterpret_cast=True, expect_scalar_cast=True), + ) + + +@pytest.mark.pto +def test_2d_column_continuous(npu): + """2D i-continuous unit over T.Parallel(256, 4).""" + m, n = 256, 4 + a = (torch.randn(n, m) * 100).to(torch.int32) + i = torch.arange(m, dtype=torch.int32).unsqueeze(0) + j = torch.arange(n, dtype=torch.int32).unsqueeze(1) + _run_and_check( + _column_kernel(m, n, 128), + (a.to(npu),), + a + j + i, + source_checks=lambda s: _assert_new_path(s, expect_vinterpret_cast=True, expect_scalar_cast=True), + ) + + +@pytest.mark.pto +def test_cast_f16_to_f32_1d(npu): + """1D f16 inputs widened with vcvt before the element-wise mul-add.""" + n = 256 + a = torch.randn(n).half() + b = torch.randn(n).half() + _run_and_check( + _cast_f16_muladd_1d(n, 128), + (a.to(npu), b.to(npu)), + a.float() * (a.float() + b.float()), + source_checks=lambda s: _assert_new_path(s, expect_vcvt_to="f32"), + atol=1e-3, + ) + + +@pytest.mark.pto +def test_cast_bf16_to_f32_1d(npu): + """1D bf16 widened to f32; bf16 does not satisfy ``is_float()`` here, so the + widening rule must whitelist it explicitly.""" + n = 256 + a = torch.randn(n).bfloat16() + _run_and_check( + _cast_bf16_add_1d(n, 128), + (a.to(npu),), + a.float() + a.float(), + source_checks=lambda s: _assert_new_path(s, expect_vcvt_to="f32"), + atol=1e-2, + ) + + +@pytest.mark.pto +def test_cast_i16_to_i32_1d(npu): + """1D int16 -> int32 widening: exact integer arithmetic after vcvt.""" + n = 256 + a = torch.randint(-30000, 30000, (n,), dtype=torch.int16) + _run_and_check( + _cast_i16_square_1d(n, 128), + (a.to(npu),), + a.int() * a.int(), + source_checks=lambda s: _assert_new_path(s, expect_vcvt_to="si32"), + ) + + +@pytest.mark.pto +def test_cast_f16_to_f32_2d(npu): + """2D j-continuous f16 -> f32 widening over T.Parallel(4, 256).""" + m, n = 4, 256 + a = torch.randn(m, n).half() + b = torch.randn(m, n).half() + _run_and_check( + _cast_f16_add_2d(m, n, 128), + (a.to(npu), b.to(npu)), + a.float() + b.float(), + source_checks=lambda s: _assert_new_path(s, expect_vcvt_to="f32"), + atol=1e-3, + ) + + +@pytest.mark.pto +def test_fp8_element_type_rejected(): + """fp8 is not supported, and the reason is the element + *category*, not the bit width: fp8 carries its own dtype code, so + ``is_float()`` is false for it and ``IsSupportedElementDType``'s category + whitelist (float / bfloat16 / int / uint / bool) does not admit it. Verify + rejects the buffer at the dtype check, before any Cast lowering is planned; + the same holds when an fp8 buffer is a load source rather than a cast + target. Supporting fp8 needs its rounding and saturation semantics plus + backend support validated first — a whitelist entry alone is not enough. + See ``test_int8_control_compiles`` for the same bit width inside the + whitelisted categories.""" + with pytest.raises(Exception) as exc: + tilelang.compile(_fp8_cast_1d(128, 128), target="pto") + message = str(exc.value) + assert "[VerifyParallelToPTO] dtype:" in message, message + assert "float8_e4m3" in message, message + + +@pytest.mark.pto +def test_int8_control_compiles(): + """8-bit control: int8 sits in a whitelisted category, so the same width + lowers through the new path — which is what makes the fp8 rejection above a + category question rather than a sub-byte/width one.""" + tilelang.compile(_i8_add_1d(128, 128), target="pto") + + +@pytest.mark.pto +def test_lane_varying_f16_widening_numeric(npu): + """Reading the half value per lane makes it a vector conversion, and the + widening is exact.""" + n = 256 + + @T.prim_func + def main(A: T.Buffer((n,), "float16"), C: T.Buffer((n,), "float32")): + with T.Kernel(1): + a = T.alloc_shared((n,), "float16") + c = T.alloc_shared((n,), "float32") + T.copy(A, a) + with T.SimdVF(lanes=128): + for i in T.Parallel(n): + c[i] = T.cast(a[i], "float32") + T.copy(c, C) + + x = (torch.arange(n, dtype=torch.float32) - 96.0) / 8.0 + kernel = tilelang.compile(main, target="pto", out_idx=-1) + assert "pto.vmi.vcvt(" in kernel.get_kernel_source() + out = kernel(x.to(torch.float16).to(npu)) + torch.npu.synchronize() + assert torch.equal(out.cpu(), x.to(torch.float16).float()) + + +@pytest.mark.pto +def test_lane_uniform_int_cast_numeric(npu): + """Integer scalar cast: int32 -> int16 with a lane-uniform source puts + int16(a[0]) in every lane.""" + n = 256 + + @T.prim_func + def main(A: T.Buffer((n,), "int32"), C: T.Buffer((n,), "int16")): + with T.Kernel(1): + a = T.alloc_shared((n,), "int32") + c = T.alloc_shared((n,), "int16") + T.copy(A, a) + with T.SimdVF(lanes=128): + for i in T.Parallel(n): + c[i] = T.cast(a[0], "int16") + T.copy(c, C) + + x = torch.arange(n, dtype=torch.int32) - 96 + out = tilelang.compile(main, target="pto", out_idx=-1)(x.to(npu)) + torch.npu.synchronize() + expected = torch.full((n,), int(x[0].item()), dtype=torch.int16) + assert torch.equal(out.cpu(), expected), (out.cpu()[:4], expected[:4]) + + +@pytest.mark.pto +def _sub_1d(n, lanes): + p = _vl_padded(n, lanes) + + @T.prim_func + def main(A: T.Tensor((p,), "float32"), B: T.Tensor((p,), "float32"), C: T.Tensor((n,), "float32")): + with T.Kernel(1): + a_ub = T.alloc_shared((p,), "float32") + b_ub = T.alloc_shared((p,), "float32") + c_ub = T.alloc_shared((n,), "float32") + T.copy(A, a_ub) + T.copy(B, b_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c_ub[i] = a_ub[i] - b_ub[i] - 1.0 + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [64, 128]) +def test_sub_1d_multichunk(npu, lanes): + _require_npu_runtime() + n = 300 + p = _vl_padded(n, lanes) + a = torch.randn(p, device=npu) + b = torch.randn(p, device=npu) + kernel = tilelang.compile(_sub_1d(n, lanes), target="pto", out_idx=-1) + out = kernel(a, b) + torch.npu.synchronize() + torch.testing.assert_close(out, a[:n] - b[:n] - 1.0, rtol=1e-5, atol=1e-5) + + +def _div_1d(n, lanes, dtype): + p = _vl_padded(n, lanes) + + @T.prim_func + def main(A: T.Tensor((p,), dtype), B: T.Tensor((p,), dtype), C: T.Tensor((n,), dtype)): + with T.Kernel(1): + a_ub = T.alloc_shared((p,), dtype) + b_ub = T.alloc_shared((p,), dtype) + c_ub = T.alloc_shared((n,), dtype) + T.copy(A, a_ub) + T.copy(B, b_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c_ub[i] = a_ub[i] / (b_ub[i] + T.cast(2, dtype)) + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_div_f32_1d(npu): + _require_npu_runtime() + n = 256 + a = torch.randn(n, device=npu) + b = torch.randn(n, device=npu) + kernel = tilelang.compile(_div_1d(n, 128, "float32"), target="pto", out_idx=-1) + out = kernel(a, b) + torch.npu.synchronize() + torch.testing.assert_close(out, a / (b + 2.0), rtol=1e-4, atol=1e-4) + + +def _unary_1d(n, lanes, kind): + p = _vl_padded(n, lanes) + + @T.prim_func + def main(A: T.Tensor((p,), "float32"), C: T.Tensor((n,), "float32")): + with T.Kernel(1): + a_ub = T.alloc_shared((p,), "float32") + c_ub = T.alloc_shared((n,), "float32") + T.copy(A, a_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + if kind == 0: + c_ub[i] = T.abs(a_ub[i] - 0.5) + elif kind == 1: + c_ub[i] = T.exp(a_ub[i]) + elif kind == 2: + c_ub[i] = T.log(T.abs(a_ub[i]) + 1.0) + else: + c_ub[i] = T.sqrt(T.abs(a_ub[i]) + 0.25) + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +@pytest.mark.parametrize("kind", [0, 1, 2, 3]) +def test_unary_1d(npu, kind): + _require_npu_runtime() + n = 256 + a = torch.randn(n, device=npu) + kernel = tilelang.compile(_unary_1d(n, 128, kind), target="pto", out_idx=-1) + out = kernel(a) + torch.npu.synchronize() + if kind == 0: + ref = (a - 0.5).abs() + elif kind == 1: + ref = a.exp() + elif kind == 2: + ref = (a.abs() + 1.0).log() + else: + ref = (a.abs() + 0.25).sqrt() + torch.testing.assert_close(out, ref, rtol=1e-4, atol=1e-4) + + +def _f32_narrow(n, lanes, to_dtype): + @T.prim_func + def main(A: T.Tensor((n,), "float32"), C: T.Tensor((n,), to_dtype)): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "float32") + c_ub = T.alloc_shared((n,), to_dtype) + T.copy(A, a_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c_ub[i] = T.cast(a_ub[i] * 0.5 - 0.25, to_dtype) + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +@pytest.mark.parametrize("to_dtype", ["float16", "bfloat16"]) +def test_f32_to_fp16_bf16_narrow(npu, to_dtype): + _require_npu_runtime() + n = 256 + a = torch.randn(n, device=npu) * 4.0 + kernel = tilelang.compile(_f32_narrow(n, 128, to_dtype), target="pto", out_idx=-1) + out = kernel(a) + torch.npu.synchronize() + ref = (a * 0.5 - 0.25).to(getattr(torch, to_dtype)) + torch.testing.assert_close(out.float(), ref.float(), rtol=1e-3, atol=1e-3) + + +def test_pipeline_routes_to_new_passes(npu): + """Pass-instrument check: the pto pipeline must run Verify/Vectorize and + must NOT run AscendSimdVFLowerParallel; auto-schedule must inject sync.""" + + names = [] + + @tvm.instrument.pass_instrument + class Capture: + def run_after_pass(self, mod, info): + names.append(info.name.split(".")[-1]) + + with tvm.transform.PassContext(opt_level=3, instruments=[Capture()]): + lower(_vecadd(512, 128), target="pto") + + assert "VerifyParallelToPTO" in names, f"VerifyParallelToPTO not in pipeline: {names}" + assert "VectorizeParallelToPTO" in names, f"VectorizeParallelToPTO not in pipeline: {names}" + assert "AscendSimdVFLowerParallel" not in names, f"legacy AscendC pass ran on target=pto: {names}" + assert "InsertSync" in names, f"InsertSync not in pipeline (auto-schedule off?): {names}" + + +if __name__ == "__main__": + import tilelang.testing + + tilelang.testing.main() + + +# --------------------------------------------------------------------------- +# Dtype-matrix positives and tail coverage +# --------------------------------------------------------------------------- + + +@pytest.mark.pto +def test_div_f16_1d(npu): + _require_npu_runtime() + n = 256 + a = torch.randn(n, device=npu).half() + b = (torch.randn(n, device=npu) + 2.5).half() + kernel = tilelang.compile(_div_1d(n, 128, "float16"), target="pto", out_idx=-1) + out = kernel(a, b) + torch.npu.synchronize() + # The device adds the constant in f16: the denominator rounds to f16 + # *before* the division. Mirror that in the reference. + denominator = (b.float() + 2.0).half() + ref = (a.float() / denominator.float()).half() + torch.testing.assert_close(out.float(), ref.float(), rtol=1e-3, atol=1e-3) + + +@pytest.mark.pto +def test_sub_bf16_1d(npu): + _require_npu_runtime() + n = 256 + a = torch.randn(n, device=npu).bfloat16() + b = torch.randn(n, device=npu).bfloat16() + kernel = tilelang.compile(_sub_1d_dtype(n, 128, "bfloat16"), target="pto", out_idx=-1) + out = kernel(a, b) + torch.npu.synchronize() + torch.testing.assert_close(out.float(), (a.float() - b.float()), rtol=1e-2, atol=1e-2) + + +@pytest.mark.pto +def test_sub_i32_1d(npu): + _require_npu_runtime() + n = 256 + a = torch.randint(-1000, 1000, (n,), device=npu, dtype=torch.int32) + b = torch.randint(-1000, 1000, (n,), device=npu, dtype=torch.int32) + kernel = tilelang.compile(_sub_1d_dtype(n, 128, "int32"), target="pto", out_idx=-1) + out = kernel(a, b) + torch.npu.synchronize() + torch.testing.assert_close(out, a - b) + + +@pytest.mark.pto +def test_div_f32_tail(npu): + _require_npu_runtime() + n = 300 + p = _vl_padded(n, 128) + a = torch.randn(p, device=npu) + b = torch.randn(p, device=npu) + kernel = tilelang.compile(_div_1d(n, 128, "float32"), target="pto", out_idx=-1) + out = kernel(a, b) + torch.npu.synchronize() + torch.testing.assert_close(out, a[:n] / (b[:n] + 2.0), rtol=1e-4, atol=1e-4) + + +@pytest.mark.pto +def test_unary_f32_tail(npu): + _require_npu_runtime() + n = 300 + p = _vl_padded(n, 128) + a = torch.randn(p, device=npu) + kernel = tilelang.compile(_unary_1d(n, 128, 1), target="pto", out_idx=-1) + out = kernel(a) + torch.npu.synchronize() + torch.testing.assert_close(out, a[:n].exp(), rtol=1e-4, atol=1e-4) + + +def _unary_f16_kernel(n, lanes): + @T.prim_func + def main(A: T.Tensor((n,), "float16"), C: T.Tensor((n,), "float16")): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "float16") + c_ub = T.alloc_shared((n,), "float16") + T.copy(A, a_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c_ub[i] = T.abs(a_ub[i] - T.cast(0.5, "float16")) + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_unary_f16_abs(npu): + _require_npu_runtime() + n = 256 + a = torch.randn(n, device=npu).half() + kernel = tilelang.compile(_unary_f16_kernel(n, 128), target="pto", out_idx=-1) + out = kernel(a) + torch.npu.synchronize() + ref = (a.float() - 0.5).abs().half() + torch.testing.assert_close(out.float(), ref.float(), rtol=1e-3, atol=1e-3) + + +def _sub_1d_dtype(n, lanes, dtype): + @T.prim_func + def main(A: T.Tensor((n,), dtype), B: T.Tensor((n,), dtype), C: T.Tensor((n,), dtype)): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), dtype) + b_ub = T.alloc_shared((n,), dtype) + c_ub = T.alloc_shared((n,), dtype) + T.copy(A, a_ub) + T.copy(B, b_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c_ub[i] = a_ub[i] - b_ub[i] + T.copy(c_ub, C) + + return main + + +# --------------------------------------------------------------------------- +# Narrowing-cast attribute/boundary coverage, f16 unary +# --------------------------------------------------------------------------- + + +def _narrow_kernel(n, lanes, to_dtype): + p = _vl_padded(n, lanes) + + @T.prim_func + def main(A: T.Tensor((p,), "float32"), C: T.Tensor((n,), to_dtype)): + with T.Kernel(1): + a_ub = T.alloc_shared((p,), "float32") + c_ub = T.alloc_shared((n,), to_dtype) + T.copy(A, a_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c_ub[i] = T.cast(a_ub[i], to_dtype) + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +def test_narrow_cast_attributes(npu): + kernel = tilelang.compile(_narrow_kernel(256, 128, "float16"), target="pto", out_idx=-1) + line = [ln.strip() for ln in kernel.get_kernel_source().split("\n") if "vmi.vcvt(" in ln and "=" in ln][0] + assert "to_dtype=pto.f16" in line, line + assert 'rounding="R"' in line, line + assert 'saturate="NOSAT"' in line, line + + +@pytest.mark.pto +def test_narrow_cast_f16_boundary_values(npu): + _require_npu_runtime() + # f16: 2^-14 is the smallest normal; 2049/2051 are ties-to-even + # midpoints; 65519.99 rounds to 65504 (largest finite) under R; 65520 + # overflows to +inf under NOSAT. + values = [ + 0.0, + -0.0, + 5.960464477539063e-08, # smallest f16 subnormal (2^-24) + 6.097555160522461e-05, # mid-subnormal, nonzero + 6.103515625e-05, # 2^-14, smallest f16 normal + 2049.0, # tie between 2048 and 2052 -> even (2048) + 2051.0, # tie between 2048 and 2052 -> even (2052) + 65504.0, # largest finite + 65519.99, # rounds to 65504 + 65520.0, # overflow -> +inf (NOSAT) + -65520.0, # overflow -> -inf (NOSAT) + ] + n = 256 + x = torch.zeros(n, dtype=torch.float32) + x[: len(values)] = torch.tensor(values, dtype=torch.float32) + kernel = tilelang.compile(_narrow_kernel(n, 128, "float16"), target="pto", out_idx=-1) + out = kernel(x.to(npu)) + torch.npu.synchronize() + got = out.float()[: len(values)].cpu() + + ref = x[: len(values)].half().float() + torch.testing.assert_close(got, ref, rtol=0, atol=0, equal_nan=True) + assert got[2] != 0.0 and got[3] != 0.0, "f16 subnormals must not flush to zero" + assert got[1] == 0.0 and torch.signbit(got[1]), "negative zero must keep the sign bit" + assert torch.isinf(got[9]) and got[9] > 0, "NOSAT f16 overflow must be +inf" + assert torch.isinf(got[10]) and got[10] < 0, "NOSAT f16 negative overflow must be -inf" + + +@pytest.mark.pto +def test_narrow_cast_bf16_boundary_values(npu): + _require_npu_runtime() + # bf16: 2^-126 is the smallest normal; 7-bit mantissa means ties at + # half-ULP of the 128-spaced binade. f32 sources cannot exceed bf16 + # max finite (~3.39e38 vs f32 max ~3.40e38), so no bf16 overflow case. + values = [ + 0.0, + -0.0, + 9.183549615799121e-41, # smallest bf16 subnormal (2^-133) + 1.1754943508222875e-38, # 2^-126, smallest bf16 normal + 1.9921875, # tie -> rounds to 2.0 (even) + 2.0078125, # tie -> rounds to 2.0 (even) + 3.3895313892515355e38, # near bf16 max finite + ] + n = 256 + x = torch.zeros(n, dtype=torch.float32) + x[: len(values)] = torch.tensor(values, dtype=torch.float32) + kernel = tilelang.compile(_narrow_kernel(n, 128, "bfloat16"), target="pto", out_idx=-1) + out = kernel(x.to(npu)) + torch.npu.synchronize() + got = out.float()[: len(values)].cpu() + + ref = x[: len(values)].bfloat16().float() + torch.testing.assert_close(got, ref, rtol=0, atol=0, equal_nan=True) + assert got[2] != 0.0, "bf16 subnormal must not flush to zero" + assert got[1] == 0.0 and torch.signbit(got[1]), "negative zero must keep the sign bit" + + +@pytest.mark.pto +def test_narrow_cast_bf16_attributes(npu): + kernel = tilelang.compile(_narrow_kernel(256, 128, "bfloat16"), target="pto", out_idx=-1) + line = [ln.strip() for ln in kernel.get_kernel_source().split("\n") if "vmi.vcvt(" in ln and "=" in ln][0] + assert "to_dtype=pto.bf16" in line, line + assert 'rounding="R"' in line, line + assert 'saturate="NOSAT"' in line, line + + +@pytest.mark.pto +def test_narrow_cast_bf16_tail(npu): + _require_npu_runtime() + n = 300 + p = _vl_padded(n, 128) + a = torch.randn(p, device=npu) * 1000 + kernel = tilelang.compile(_narrow_kernel(n, 128, "bfloat16"), target="pto", out_idx=-1) + out = kernel(a) + torch.npu.synchronize() + torch.testing.assert_close(out.float(), a[:n].bfloat16().float(), rtol=1e-3, atol=1e-3) + + +def _unary_f16_math_kernel(n, lanes, kind): + @T.prim_func + def main(A: T.Tensor((n,), "float16"), C: T.Tensor((n,), "float16")): + with T.Kernel(1): + a_ub = T.alloc_shared((n,), "float16") + c_ub = T.alloc_shared((n,), "float16") + T.copy(A, a_ub) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + if kind == 0: + c_ub[i] = T.exp(a_ub[i]) + elif kind == 1: + c_ub[i] = T.log(T.abs(a_ub[i]) + T.cast(1.0, "float16")) + else: + c_ub[i] = T.sqrt(T.abs(a_ub[i]) + T.cast(0.25, "float16")) + T.copy(c_ub, C) + + return main + + +@pytest.mark.pto +@pytest.mark.parametrize("kind", [0, 1, 2]) +def test_unary_f16_math(npu, kind): + _require_npu_runtime() + n = 256 + a = (torch.randn(n, device=npu) * 2).half() + kernel = tilelang.compile(_unary_f16_math_kernel(n, 128, kind), target="pto", out_idx=-1) + out = kernel(a) + torch.npu.synchronize() + if kind == 0: + ref = a.exp() + elif kind == 1: + ref = (a.abs() + 1.0).log() + else: + ref = (a.abs() + 0.25).sqrt() + torch.testing.assert_close(out.float(), ref.float(), rtol=1e-2, atol=1e-2) diff --git a/testing/ascend/language/pto/test_vectorize_tail_mask_run.py b/testing/ascend/language/pto/test_vectorize_tail_mask_run.py new file mode 100644 index 0000000000..b729a6912b --- /dev/null +++ b/testing/ascend/language/pto/test_vectorize_tail_mask_run.py @@ -0,0 +1,808 @@ +"""End-to-end NPU numeric acceptance for non-divisible Parallel extents. + +The kernels use ``T.Parallel`` extents that are not multiples of the +SIMD_VF lanes, so VectorizeParallelToPTO emits the per-unit ``remaining`` +scalar, the per-chunk dynamic create_mask, and the masked tail store. +Checks cover: + +- padded output UB + sentinel preservation past E (lane discipline of + the predicated store); +- compact output UB of exactly E elements (no out-of-bounds write); +- 2D per-row reset and per-row sentinel (row independence); +- mixed divisible/non-divisible units in one VF; +- the three lane widths and the dtypes of the design acceptance matrix. +""" + +import os +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import tilelang +import tilelang.ascend.language as T + + +def _require_npu_runtime(): + if not hasattr(torch, "npu"): + pytest.skip("torch.npu is unavailable in the current environment") + + try: + _ = torch.randn(1, device=torch.device("npu")) + torch.npu.synchronize() + except Exception as err: + pytest.skip(f"NPU runtime is unavailable in the current environment: {err}") + + +@pytest.fixture(scope="module") +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +SENTINEL = -12345.0 + +TORCH_DT = { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, +} + + +def _ceil_to_lanes(e, lanes): + return (e + lanes - 1) // lanes * lanes + + +# --------------------------------------------------------------------------- +# Kernels +# --------------------------------------------------------------------------- + + +def _k1d(E, out_extent, lanes, dtype): + """1D scale-by-2 with padded or compact output UB (out_extent = P or + E); B_init pre-fills the output so the sentinel check observes which + lanes the masked store actually wrote.""" + + P = _ceil_to_lanes(E, lanes) + + @T.prim_func + def main( + A: T.Tensor((P,), dtype), + B_init: T.Tensor((out_extent,), dtype), + B: T.Tensor((out_extent,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((P,), dtype) + b = T.alloc_shared((out_extent,), dtype) + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(E): + b[i] = a[i] + a[i] + T.copy(b, B) + + return main + + +def _k2d(M, N, lanes, dtype): + """2D row scale-by-2: j is vectorized and non-divisible. Both the + source and the output UB use the padded row stride P: the tail + vload reads [row*N, row*N + L) which crosses into the next row's + head, so the buffer must be readable through the padded range AND + the load base (row stride) must stay vector-aligned — a logical + stride of N elements would make later rows' vload bases unaligned + (the design keeps alignment a caller responsibility, same as tail + readability).""" + + P = _ceil_to_lanes(N, lanes) + + @T.prim_func + def main( + A: T.Tensor((M * P,), dtype), + B_init: T.Tensor((M * P,), dtype), + B: T.Tensor((M * P,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((M * P,), dtype) + b = T.alloc_shared((M * P,), dtype) + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes), T.Parallel(M, N) as (i, j): + b[i * P + j] = a[i * P + j] + a[i * P + j] + T.copy(b, B) + + return main + + +def _k_mixed(E_div, E_non, lanes, dtype): + """Divisible unit then non-divisible unit in one VF.""" + + Pd = _ceil_to_lanes(E_div, lanes) + Pn = _ceil_to_lanes(E_non, lanes) + + @T.prim_func + def main( + A1: T.Tensor((Pd,), dtype), + B1_init: T.Tensor((Pd,), dtype), + B1: T.Tensor((Pd,), dtype), + A2: T.Tensor((Pn,), dtype), + B2_init: T.Tensor((Pn,), dtype), + B2: T.Tensor((Pn,), dtype), + ): + with T.Kernel(1): + a1 = T.alloc_shared((Pd,), dtype) + b1 = T.alloc_shared((Pd,), dtype) + a2 = T.alloc_shared((Pn,), dtype) + b2 = T.alloc_shared((Pn,), dtype) + T.copy(A1, a1) + T.copy(B1_init, b1) + T.copy(A2, a2) + T.copy(B2_init, b2) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(E_div): + b1[i] = a1[i] + a1[i] + for i in T.Parallel(E_non): + b2[i] = a2[i] + a2[i] + T.copy(b1, B1) + T.copy(b2, B2) + + return main + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [64, 128, 256]) +@pytest.mark.parametrize("dtype", ["float32", "float16", "bfloat16"]) +def test_1d_padded_sentinel(npu, lanes, dtype): + """Padded output UB: [0,E) holds 2*A, [E,P) keeps the sentinel.""" + E = 150 if lanes <= 128 else 150 # tail shape independent of lanes + P = _ceil_to_lanes(E, lanes) + a = (torch.randn(P) * 2).to(TORCH_DT[dtype]) + b_init = torch.full((P,), SENTINEL, dtype=TORCH_DT[dtype]) + kernel = tilelang.compile(_k1d(E, P, lanes, dtype), target="pto", out_idx=-1) + out = kernel(a.to(npu), b_init.to(npu)) + torch.npu.synchronize() + got = out.cpu() + assert torch.equal(got[0:E], (a[0:E] + a[0:E])), f"{dtype} L={lanes}: values" + assert bool((got[E:P] == SENTINEL).all()), f"{dtype} L={lanes}: sentinel" + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [64, 128, 256]) +def test_1d_compact_output(npu, lanes): + """Compact output UB of exactly E elements: no write past E.""" + E = 100 if lanes == 64 else 150 + P = _ceil_to_lanes(E, lanes) + a = (torch.randn(P) * 2).to(torch.float32) + b_init = torch.full((E,), SENTINEL, dtype=torch.float32) + kernel = tilelang.compile(_k1d(E, E, lanes, "float32"), target="pto", out_idx=-1) + out = kernel(a.to(npu), b_init.to(npu)) + torch.npu.synchronize() + got = out.cpu() + assert torch.equal(got, a[0:E] + a[0:E]) + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [64, 128, 256]) +def test_2d_row_reset_sentinel(npu, lanes): + """(4, 100) j-vectorized: every row's [0,N) is correct and each row's + [N,P) tail keeps its own sentinel — rows are independent (per-row + remaining reset) and no row writes past its own E lanes.""" + M, N = 4, 100 + P = _ceil_to_lanes(N, lanes) + a_rows = torch.randn(M, N) * 2 + a = torch.zeros(M, P) + a[:, 0:N] = a_rows # padding lanes read as zeros; never stored (mask) + b_init = torch.full((M * P,), SENTINEL, dtype=torch.float32) + kernel = tilelang.compile(_k2d(M, N, lanes, "float32"), target="pto", out_idx=-1) + out = kernel(a.to(npu).flatten(), b_init.to(npu).flatten()) + torch.npu.synchronize() + got = out.cpu().reshape(M, P) + expect_row = a_rows + a_rows + assert torch.equal(got[:, 0:N], expect_row), "row values" + assert bool((got[:, N:P] == SENTINEL).all()), "row tail sentinel" + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [64, 128, 256]) +def test_mixed_units_both_masks(npu, lanes): + """Divisible unit (full mask) and non-divisible unit (dynamic mask) + in one VF: both produce correct values, and the non-divisible tail + keeps the sentinel.""" + E_div = lanes * 2 + E_non = 150 + Pn = _ceil_to_lanes(E_non, lanes) + a1 = (torch.randn(E_div) * 2).to(torch.float32) + b1_init = torch.full((E_div,), SENTINEL, dtype=torch.float32) + a2 = (torch.randn(Pn) * 2).to(torch.float32) + b2_init = torch.full((Pn,), SENTINEL, dtype=torch.float32) + kernel = tilelang.compile(_k_mixed(E_div, E_non, lanes, "float32"), target="pto", out_idx=[2, 5]) + out1, out2 = kernel(a1.to(npu), b1_init.to(npu), a2.to(npu), b2_init.to(npu)) + torch.npu.synchronize() + g1, g2 = out1.cpu(), out2.cpu() + assert torch.equal(g1, a1 + a1), "divisible unit values" + assert torch.equal(g2[0:E_non], a2[0:E_non] + a2[0:E_non]), "nondiv unit values" + assert bool((g2[E_non:Pn] == SENTINEL).all()), "nondiv unit sentinel" + + +@pytest.mark.pto +@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) +def test_2d_half_precision(npu, dtype): + """2D row reset under half precision.""" + M, N, lanes = 4, 100, 128 + P = _ceil_to_lanes(N, lanes) + a_rows = (torch.randn(M, N) * 2).to(TORCH_DT[dtype]) + a = torch.zeros(M, P, dtype=TORCH_DT[dtype]) + a[:, 0:N] = a_rows + b_init = torch.full((M * P,), SENTINEL, dtype=TORCH_DT[dtype]) + kernel = tilelang.compile(_k2d(M, N, lanes, dtype), target="pto", out_idx=-1) + out = kernel(a.to(npu).flatten(), b_init.to(npu).flatten()) + torch.npu.synchronize() + got = out.cpu().reshape(M, P) + expect_row = a_rows + a_rows + assert torch.equal(got[:, 0:N], expect_row), f"{dtype} row values" + assert bool((got[:, N:P] == SENTINEL).all()), f"{dtype} row tail sentinel" + + +# --------------------------------------------------------------------------- +# Design 9.2.4 length matrix +# --------------------------------------------------------------------------- + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [64, 128, 256]) +@pytest.mark.parametrize( + "e_kind", + ["lt_l_minus_1", "l_plus_1", "two_l_minus_1", "two_l_plus_17", "three_l"], + ids=["L-1", "L+1", "2L-1", "2L+17", "3L"], +) +def test_1d_length_matrix(npu, lanes, e_kind): + """Design 9.2.4 matrix: 1L non-divisible (L+1, 2L-1, 2L+17) and E>L divisible (3L); padded + output UB with sentinel past E. E=L and 2L are covered by + test_vecadd_1d_singlechunk and the mixed units test.""" + E = { + "lt_l_minus_1": lanes - 1, + "l_plus_1": lanes + 1, + "two_l_minus_1": 2 * lanes - 1, + "two_l_plus_17": 2 * lanes + 17, + "three_l": 3 * lanes, + }[e_kind] + P = _ceil_to_lanes(E, lanes) + a = (torch.randn(P) * 2).to(torch.float32) + b_init = torch.full((P,), SENTINEL, dtype=torch.float32) + kernel = tilelang.compile(_k1d(E, P, lanes, "float32"), target="pto", out_idx=-1) + out = kernel(a.to(npu), b_init.to(npu)) + torch.npu.synchronize() + got = out.cpu() + assert torch.equal(got[0:E], a[0:E] + a[0:E]), f"E={E} L={lanes}: values" + assert bool((got[E:P] == SENTINEL).all()), f"E={E} L={lanes}: sentinel" + + +def _k2d_col(M, N, lanes, dtype): + """i-continuous 2D: row-major (N, P_M) storage accessed c[j * P_M + i] + — i (extent M) is the selected non-divisible dim, j (extent N) the + kept coordinate; each row resets remaining and keeps its own tail + sentinel.""" + + P = _ceil_to_lanes(M, lanes) + + @T.prim_func + def main( + A: T.Tensor((N * P,), dtype), + B_init: T.Tensor((N * P,), dtype), + B: T.Tensor((N * P,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((N * P,), dtype) + b = T.alloc_shared((N * P,), dtype) + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes), T.Parallel(M, N) as (i, j): + b[j * P + i] = a[j * P + i] + a[j * P + i] + T.copy(b, B) + + return main + + +@pytest.mark.pto +@pytest.mark.parametrize("lanes", [64, 128, 256]) +def test_2d_i_continuous_nondivisible(npu, lanes): + """Design 9.2.4: the i-continuous direction with a non-divisible M, + four kept coordinates with distinct row data — per-row remaining + reset and per-row sentinel past M.""" + M, N = 100, 4 + P = _ceil_to_lanes(M, lanes) + a_rows = torch.randn(N, M) * 2 + a = torch.zeros(N, P) + a[:, 0:M] = a_rows # padding lanes read as zeros; never stored (mask) + b_init = torch.full((N * P,), SENTINEL, dtype=torch.float32) + kernel = tilelang.compile(_k2d_col(M, N, lanes, "float32"), target="pto", out_idx=-1) + out = kernel(a.to(npu).flatten(), b_init.to(npu)) + torch.npu.synchronize() + got = out.cpu().reshape(N, P) + expect_row = a_rows + a_rows + assert torch.equal(got[:, 0:M], expect_row), "column values" + assert bool((got[:, M:P] == SENTINEL).all()), "column tail sentinel" + + +# --------------------------------------------------------------------------- +# Non-divisible mixed-width / index-as-value acceptance +# --------------------------------------------------------------------------- + + +def _k1d_cast_widen(E, P, lanes, src_dtype): + """c[i] = f32(a[i]) + f32(a[i]) — widening with a non-divisible E.""" + + @T.prim_func + def main( + A: T.Tensor((P,), src_dtype), + B_init: T.Tensor((P,), "float32"), + B: T.Tensor((P,), "float32"), + ): + with T.Kernel(1): + a = T.alloc_shared((P,), src_dtype) + b = T.alloc_shared((P,), "float32") + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(E): + b[i] = T.Cast("float32", a[i]) + T.Cast("float32", a[i]) + T.copy(b, B) + + return main + + +def _k1d_fp8_narrow(E, P, lanes): + """c[i] = fp8(a[i]) — FP32 -> E4M3FN narrowing with a non-divisible E.""" + + @T.prim_func + def main( + A: T.Tensor((P,), "float32"), + C_init: T.Tensor((P,), T.float8_e4m3fn), + C: T.Tensor((P,), T.float8_e4m3fn), + ): + with T.Kernel(1): + a = T.alloc_shared((P,), "float32") + c = T.alloc_shared((P,), T.float8_e4m3fn) + T.copy(A, a) + T.copy(C_init, c) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(E): + c[i] = T.Cast(T.float8_e4m3fn, a[i]) + T.copy(c, C) + + return main + + +def _k1d_fp8_widen(E, P, lanes): + """c[i] = f32(a[i]) — E4M3FN -> FP32 widening with a non-divisible E.""" + + @T.prim_func + def main( + A: T.Tensor((P,), T.float8_e4m3fn), + B_init: T.Tensor((P,), "float32"), + B: T.Tensor((P,), "float32"), + ): + with T.Kernel(1): + a = T.alloc_shared((P,), T.float8_e4m3fn) + b = T.alloc_shared((P,), "float32") + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(E): + b[i] = T.Cast("float32", a[i]) + T.copy(b, B) + + return main + + +def _k1d_index_value(E, P, lanes): + """c[i] = a[i] + i — the int32 loop index feeds the value stream + (int->float casts are outside the supported conversion set).""" + + @T.prim_func + def main( + A: T.Tensor((P,), "int32"), + B_init: T.Tensor((P,), "int32"), + B: T.Tensor((P,), "int32"), + ): + with T.Kernel(1): + a = T.alloc_shared((P,), "int32") + b = T.alloc_shared((P,), "int32") + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(E): + b[i] = a[i] + i + T.copy(b, B) + + return main + + +def _decode_e4m3(code): + """E4M3FN: 1 sign, 4 exponent (bias 7), 3 mantissa; exp=15/m=7 is NaN.""" + sign = -1.0 if code & 0x80 else 1.0 + exp = (code >> 3) & 0xF + mant = code & 0x7 + if exp == 0: + return sign * mant * 2.0**-9 + if exp == 0xF and mant == 0x7: + return None # NaN + return sign * (1.0 + mant / 8.0) * 2.0 ** (exp - 7) + + +def _e4m3_table(): + vals = [] + for code in range(256): + value = _decode_e4m3(code) + if value is not None and (value > 0 or code == 0): + vals.append((code, value)) + vals.sort(key=lambda pair: pair[1]) + return vals + + +_E4M3_TABLE = _e4m3_table() + + +def _encode_rne(value): + """Nearest E4M3FN code with ties-to-even, saturating (host reference).""" + limit = _E4M3_TABLE[-1][1] + sign_bit = 0x80 if value < 0 else 0x00 + mag = min(abs(value), limit) + code = _E4M3_TABLE[0][0] + best = abs(_E4M3_TABLE[0][1] - mag) + for cand, magnitude in _E4M3_TABLE[1:]: + delta = abs(magnitude - mag) + if delta < best or (delta == best and (cand & 1) == 0 and (code & 1) != 0): + code, best = cand, delta + return code | sign_bit + + +@pytest.mark.pto +@pytest.mark.parametrize("src_dtype", ["float16", "bfloat16"]) +def test_1d_nondiv_cast_widen(npu, src_dtype): + """E=150, L=128 (Q=2, tail 22): FP16/BF16 -> FP32 widening under the + per-chunk dynamic mask — valid lanes exact, tail sentinel kept.""" + E, lanes = 150, 128 + P = _ceil_to_lanes(E, lanes) + a = (torch.randn(P) * 2).to(TORCH_DT[src_dtype]) + b_init = torch.full((P,), SENTINEL, dtype=torch.float32) + kernel = tilelang.compile(_k1d_cast_widen(E, P, lanes, src_dtype), target="pto", out_idx=-1) + out = kernel(a.to(npu), b_init.to(npu)) + torch.npu.synchronize() + got = out.cpu() + expect = a[0:E].float() + a[0:E].float() + assert torch.equal(got[0:E], expect), f"{src_dtype}: values" + assert bool((got[E:P] == SENTINEL).all()), f"{src_dtype}: tail sentinel" + + +@pytest.mark.pto +def test_1d_nondiv_fp8_narrow(npu): + """E=150, L=128: FP32 -> E4M3FN narrowing under the dynamic mask — + exact RNE+SAT codes on the valid lanes, tail sentinel byte kept.""" + E, lanes = 150, 128 + P = _ceil_to_lanes(E, lanes) + values = [-16.0 + i * 0.125 for i in range(P)] + a = torch.tensor(values, dtype=torch.float32) + c_init = torch.full((P,), 0x7F, dtype=torch.uint8).view(torch.float8_e4m3fn) + kernel = tilelang.compile(_k1d_fp8_narrow(E, P, lanes), target="pto", out_idx=-1) + out = kernel(a.to(npu), c_init.to(npu)) + torch.npu.synchronize() + raw = out.cpu().view(torch.uint8) + expect = torch.tensor([_encode_rne(v) for v in values[:E]], dtype=torch.uint8) + assert torch.equal(raw[0:E], expect), "narrowing codes" + assert bool((raw[E:P] == 0x7F).all()), "narrowing tail sentinel" + + +@pytest.mark.pto +def test_1d_nondiv_fp8_widen(npu): + """E=150, L=128: E4M3FN -> FP32 widening under the dynamic mask — + every finite code decodes exactly, tail sentinel kept.""" + E, lanes = 150, 128 + P = _ceil_to_lanes(E, lanes) + all_codes = [c for c in range(256) if c not in (0x7F, 0xFF)] + codes = all_codes[:75] + all_codes[-75:] + [0] * (P - E) + a = torch.tensor(codes, dtype=torch.uint8).view(torch.float8_e4m3fn) + b_init = torch.full((P,), SENTINEL, dtype=torch.float32) + kernel = tilelang.compile(_k1d_fp8_widen(E, P, lanes), target="pto", out_idx=-1) + out = kernel(a.to(npu), b_init.to(npu)) + torch.npu.synchronize() + got = out.cpu() + expect = torch.tensor([_decode_e4m3(c) for c in codes[:E]], dtype=torch.float32) + assert torch.equal(got[0:E], expect), "widening values" + assert bool((got[E:P] == SENTINEL).all()), "widening tail sentinel" + + +@pytest.mark.pto +def test_1d_nondiv_index_value(npu): + """E=150, L=128: the loop index feeds the value stream (a[i] + i) + under the dynamic mask — per-lane index values exact, tail sentinel + kept.""" + E, lanes = 150, 128 + P = _ceil_to_lanes(E, lanes) + a = (torch.randn(P) * 1000).to(torch.int32) + b_init = torch.full((P,), int(SENTINEL), dtype=torch.int32) + kernel = tilelang.compile(_k1d_index_value(E, P, lanes), target="pto", out_idx=-1) + out = kernel(a.to(npu), b_init.to(npu)) + torch.npu.synchronize() + got = out.cpu() + expect = a[0:E] + torch.arange(E, dtype=torch.int32) + assert torch.equal(got[0:E], expect), "index values" + assert bool((got[E:P] == int(SENTINEL)).all()), "index tail sentinel" + + +# --------------------------------------------------------------------------- +# Real Buffer stride S = P + L (distinct from the padded row stride P) +# --------------------------------------------------------------------------- + + +def _k2d_stride_sp(M, N, lanes, S, dtype): + """b[i * S + j] with the real row stride S > P: distinct row data and a + fully initialized store catch an implementation that rewrote the + address stride to P (reads/writes would land on wrong elements).""" + + @T.prim_func + def main( + A: T.Tensor((M * S,), dtype), + B_init: T.Tensor((M * S,), dtype), + B: T.Tensor((M * S,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((M * S,), dtype) + b = T.alloc_shared((M * S,), dtype) + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes), T.Parallel(M, N) as (i, j): + b[i * S + j] = a[i * S + j] + a[i * S + j] + T.copy(b, B) + + return main + + +def _k2d_col_stride_sp(M, N, lanes, S, dtype): + """b[j * S + i] with the real row stride S > P_M (i-continuous).""" + + @T.prim_func + def main( + A: T.Tensor((N * S,), dtype), + B_init: T.Tensor((N * S,), dtype), + B: T.Tensor((N * S,), dtype), + ): + with T.Kernel(1): + a = T.alloc_shared((N * S,), dtype) + b = T.alloc_shared((N * S,), dtype) + T.copy(A, a) + T.copy(B_init, b) + with T.SimdVF(lanes=lanes), T.Parallel(M, N) as (i, j): + b[j * S + i] = a[j * S + i] + a[j * S + i] + T.copy(b, B) + + return main + + +@pytest.mark.pto +def test_2d_j_continuous_stride_sp(npu): + """(4,100) j-continuous with real row stride S = P + L = 256: reads and + writes must use S (not the fragment's P), rows carry distinct data and + each row's [N, S) keeps its sentinel.""" + M, N, lanes = 4, 100, 128 + P = _ceil_to_lanes(N, lanes) + S = P + lanes + a_rows = torch.randn(M, N) * 2 + a = torch.zeros(M, S) + a[:, 0:N] = a_rows # padding readable as zeros + b_init = torch.full((M * S,), SENTINEL, dtype=torch.float32) + kernel = tilelang.compile(_k2d_stride_sp(M, N, lanes, S, "float32"), target="pto", out_idx=-1) + out = kernel(a.to(npu).flatten(), b_init.to(npu)) + torch.npu.synchronize() + got = out.cpu().reshape(M, S) + expect_row = a_rows + a_rows + assert torch.equal(got[:, 0:N], expect_row), "row values" + assert bool((got[:, N:S] == SENTINEL).all()), "row sentinel through S" + + +@pytest.mark.pto +def test_2d_i_continuous_stride_sp(npu): + """(100,4) i-continuous with real row stride S = P_M + L = 256: reads + and writes must use S, rows carry distinct data and each row's (M, S) + tail keeps its sentinel.""" + M, N, lanes = 100, 4, 128 + P = _ceil_to_lanes(M, lanes) + S = P + lanes + a_rows = torch.randn(N, M) * 2 + a = torch.zeros(N, S) + a[:, 0:M] = a_rows + b_init = torch.full((N * S,), SENTINEL, dtype=torch.float32) + kernel = tilelang.compile(_k2d_col_stride_sp(M, N, lanes, S, "float32"), target="pto", out_idx=-1) + out = kernel(a.to(npu).flatten(), b_init.to(npu)) + torch.npu.synchronize() + got = out.cpu().reshape(N, S) + expect_row = a_rows + a_rows + assert torch.equal(got[:, 0:M], expect_row), "column values" + assert bool((got[:, M:S] == SENTINEL).all()), "column sentinel through S" + + +# --------------------------------------------------------------------------- +# 2D non-divisible mixed-width with index-as-value +# --------------------------------------------------------------------------- + + +def _k2d_nondiv_mixed_j(M, N, lanes): + """c[i, j] = i32(a[i, j]) + i + j — j-vectorized, int16->int32 + widening, both loop indices as values on a non-divisible j.""" + + P = _ceil_to_lanes(N, lanes) + + @T.prim_func + def main( + A: T.Tensor((M * P,), "int16"), + C_init: T.Tensor((M * P,), "int32"), + C: T.Tensor((M * P,), "int32"), + ): + with T.Kernel(1): + a = T.alloc_shared((M * P,), "int16") + c = T.alloc_shared((M * P,), "int32") + T.copy(A, a) + T.copy(C_init, c) + with T.SimdVF(lanes=lanes), T.Parallel(M, N) as (i, j): + c[i * P + j] = T.Cast("int32", a[i * P + j]) + i + j + T.copy(c, C) + + return main + + +def _k2d_nondiv_mixed_i(M, N, lanes): + """c[j, i] = i32(a[j, i]) + i + j — i-continuous, non-divisible M, + same widening + index-value combination.""" + + P = _ceil_to_lanes(M, lanes) + + @T.prim_func + def main( + A: T.Tensor((N * P,), "int16"), + C_init: T.Tensor((N * P,), "int32"), + C: T.Tensor((N * P,), "int32"), + ): + with T.Kernel(1): + a = T.alloc_shared((N * P,), "int16") + c = T.alloc_shared((N * P,), "int32") + T.copy(A, a) + T.copy(C_init, c) + with T.SimdVF(lanes=lanes), T.Parallel(M, N) as (i, j): + c[j * P + i] = T.Cast("int32", a[j * P + i]) + i + j + T.copy(c, C) + + return main + + +def _k2d_nondiv_fp8_narrow_ij(f8, M, N, lanes): + """c[j, i] = fp8(a[j, i] * s[i] + t[j]) — FP32 -> FP8 narrowing on a + non-divisible vectorized (i) extent, with lane-varying s[i] and the + lane-uniform t[j] broadcast.""" + + P = _ceil_to_lanes(M, lanes) + + @T.prim_func + def main( + A: T.Tensor((N, P), "float32"), + S: T.Tensor((P,), "float32"), + Tt: T.Tensor((N,), "float32"), + C_init: T.Tensor((N, P), f8), + C: T.Tensor((N, P), f8), + ): + with T.Kernel(1): + a = T.alloc_shared((N, P), "float32") + s = T.alloc_shared((P,), "float32") + t = T.alloc_shared((N,), "float32") + c = T.alloc_shared((N, P), f8) + T.copy(A, a) + T.copy(S, s) + T.copy(Tt, t) + T.copy(C_init, c) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(M, N): + c[j, i] = T.Cast(f8, a[j, i] * s[i] + t[j]) + T.copy(c, C) + + return main + + +@pytest.mark.pto +def test_2d_nondiv_mixed_width_index_j(npu): + """(4,150) j-vectorized, L=128 (Q=2, tail 22): int16->int32 widening + with both indices as values under the per-chunk dynamic mask — valid + lanes exact, per-row tail sentinel kept.""" + M, N, lanes = 4, 150, 128 + P = _ceil_to_lanes(N, lanes) + a_rows = torch.randint(-1000, 1000, (M, N), dtype=torch.int16) + a = torch.zeros(M, P, dtype=torch.int16) + a[:, 0:N] = a_rows # padding lanes read as zeros; never stored (mask) + c_init = torch.full((M * P,), int(SENTINEL), dtype=torch.int32) + kernel = tilelang.compile(_k2d_nondiv_mixed_j(M, N, lanes), target="pto", out_idx=-1) + out = kernel(a.to(npu).flatten(), c_init.to(npu)) + torch.npu.synchronize() + got = out.cpu().reshape(M, P) + expect = a_rows.to(torch.int32) + torch.arange(M, dtype=torch.int32)[:, None] + torch.arange(N, dtype=torch.int32)[None, :] + assert torch.equal(got[:, 0:N], expect), "row values" + assert bool((got[:, N:P] == int(SENTINEL)).all()), "row tail sentinel" + + +@pytest.mark.pto +def test_2d_nondiv_mixed_width_index_i(npu): + """(150,4) i-continuous, L=128 (Q=2, tail 22): the symmetric direction + — valid lanes exact, per-row tail sentinel kept.""" + M, N, lanes = 150, 4, 128 + P = _ceil_to_lanes(M, lanes) + a_rows = torch.randint(-1000, 1000, (N, M), dtype=torch.int16) + a = torch.zeros(N, P, dtype=torch.int16) + a[:, 0:M] = a_rows + c_init = torch.full((N * P,), int(SENTINEL), dtype=torch.int32) + kernel = tilelang.compile(_k2d_nondiv_mixed_i(M, N, lanes), target="pto", out_idx=-1) + out = kernel(a.to(npu).flatten(), c_init.to(npu)) + torch.npu.synchronize() + got = out.cpu().reshape(N, P) + expect = a_rows.to(torch.int32) + torch.arange(M, dtype=torch.int32)[None, :] + torch.arange(N, dtype=torch.int32)[:, None] + assert torch.equal(got[:, 0:M], expect), "column values" + assert bool((got[:, M:P] == int(SENTINEL)).all()), "column tail sentinel" + + +@pytest.mark.pto +def test_2d_nondiv_fp8_narrow(npu): + """M=150, L=128 (Q=2, tail 22) on the vectorized i dim, N=4 rows: + FP32 -> E4M3FN narrowing with lane-varying s[i] and a lane-uniform t[j] + broadcast — exact RNE codes on in-range valid lanes, overflow lanes + accept the saturated or NaN encoding (environment-defined under the + global CTRL saturation mode), per-row tail sentinel byte kept.""" + M, N, lanes = 150, 4, 128 + P = _ceil_to_lanes(M, lanes) + + a_vals = [-6.0 + 0.25 * k for k in range(N * M)] + s_vals = [0.5 + 0.25 * k for k in range(M)] + t_vals = [-2.0 + 0.5 * k for k in range(N)] + + a = torch.zeros(N, P, dtype=torch.float32) + a[:, :M] = torch.tensor(a_vals, dtype=torch.float32).reshape(N, M) + s_t = torch.zeros(P, dtype=torch.float32) + s_t[:M] = torch.tensor(s_vals, dtype=torch.float32) + t_t = torch.tensor(t_vals, dtype=torch.float32) + c_init = torch.full((N, P), 0x7F, dtype=torch.uint8).view(torch.float8_e4m3fn) + + kernel = tilelang.compile(_k2d_nondiv_fp8_narrow_ij(T.float8_e4m3fn, M, N, lanes), target="pto", out_idx=-1) + src = kernel.get_kernel_source() + vcvt = [ln.strip() for ln in src.split("\n") if "vmi.vcvt" in ln and "=" in ln] + assert len(vcvt) == 1, f"expected exactly one conversion:\n{src}" + assert "to_dtype=pto.f8e4m3" in vcvt[0], vcvt[0] + assert "vmi.vbrc" in src, src + assert "create_mask" in src and "remaining" in src, src + + out = kernel(a.to(npu), s_t.to(npu), t_t.to(npu), c_init.to(npu)) + torch.npu.synchronize() + raw = out.cpu().view(torch.uint8) + + ref_t = a[:, :M] * s_t[:M].reshape(1, M) + t_t.reshape(N, 1) + expect = torch.tensor([_encode_rne(v) for v in ref_t.reshape(-1).tolist()], dtype=torch.uint8).reshape(N, M) + got = raw[:, :M] + # Overflow lanes (|v| > 448, the e4m3fn max finite) are environment- + # defined under the global CTRL saturation mode forced by + # pto.init_core(): accept the saturated (0x7E/0xFE) or NaN (0x7F/0xFF) + # encoding (NaN does not necessarily preserve the sign); in-range lanes + # keep exact RNE+SAT codes. No existing fp8 test in the repo covers SAT + # overflow, so neither does this one. + overflow = ref_t.abs() > 448.0 + assert torch.equal(got[~overflow], expect[~overflow]), "in-range narrowing codes" + if bool(overflow.any()): + ov = got[overflow] + ok = (ov == 0x7E) | (ov == 0x7F) | (ov == 0xFE) | (ov == 0xFF) + assert bool(ok.all()), "overflow lanes must be saturated or NaN encodings" + assert bool((raw[:, M:P] == 0x7F).all()), "per-row tail sentinel" diff --git a/tilelang/ascend/language/frame.py b/tilelang/ascend/language/frame.py index 29f1189674..49a9da6a99 100644 --- a/tilelang/ascend/language/frame.py +++ b/tilelang/ascend/language/frame.py @@ -176,22 +176,33 @@ def inside_simdvf() -> bool: return _inside_simdvf[0] > 0 -def SimdVF(latency: int = 0) -> SimdVFFrame: +_SIMDVF_VALID_LANES = (64, 128, 256) + + +def SimdVF(latency: int = 0, lanes: int = 128) -> SimdVFFrame: """Construct a SimdVF region for Ascend NPU MicroAPI operations. Parameters ---------- latency : int Measured latency in cycles. Zero uses the scheduler's latency model. + lanes : int + Logical lane count for the PTO SIMD vector path. Must be one of + 64, 128, or 256; the default is 128. AscendC lowering ignores this + parameter and keeps its physical register width. Returns ------- res : SimdVFFrame The SimdVFFrame used to denote a register-level SIMD region. """ + if isinstance(lanes, bool) or not isinstance(lanes, int): + raise TypeError(f"T.SimdVF requires lanes to be an int, got {type(lanes)}") + if lanes not in _SIMDVF_VALID_LANES: + raise ValueError(f"T.SimdVF lanes must be one of {_SIMDVF_VALID_LANES}, got {lanes}") source_index = _next_vf_source_index("simdvf") return _ffi_api.SimdVF( # type: ignore[attr-defined] # pylint: disable=no-member - int(latency), source_index + int(latency), source_index, int(lanes) ) diff --git a/tilelang/ascend/pipeline.py b/tilelang/ascend/pipeline.py index c5f31833bf..7426db1248 100644 --- a/tilelang/ascend/pipeline.py +++ b/tilelang/ascend/pipeline.py @@ -15,6 +15,7 @@ ) from . import transform as ascend_transform +from .target import target_is_pto _MERGE_UB_ALIGNMENT = 32 @@ -118,7 +119,16 @@ def AscendPassPipelineBody(mod: IRModule, target: Target) -> IRModule: mod = tilelang.transform.Simplify()(mod) mod = ascend_transform.NormalizeBufferVersion()(mod) - mod = ascend_transform.AscendSimdVFLowerParallel()(mod) + # The PTO target splits here: the formal PTO Parallel pipeline runs + # Verify -> Legalize -> Vectorize (Legalize is currently a pass-through + # placeholder). The AscendC SIMD_VF lowering only + # handles the AscendC target and must not see PTO regions. + if target_is_pto(target): + mod = ascend_transform.VerifyParallelToPTO()(mod) + mod = ascend_transform.LegalizeParallelToPTO()(mod) + mod = ascend_transform.VectorizeParallelToPTO()(mod) + else: + mod = ascend_transform.AscendSimdVFLowerParallel()(mod) mod = ascend_transform.AscendLowerTileOp()(mod) mod = tilelang.transform.VerifyReducerConsumed()(mod) diff --git a/tilelang/ascend/transform/__init__.py b/tilelang/ascend/transform/__init__.py index d1fd48fa95..56282ac40b 100644 --- a/tilelang/ascend/transform/__init__.py +++ b/tilelang/ascend/transform/__init__.py @@ -134,6 +134,37 @@ def AscendSimdVFLowerParallel(): return _ffi_api.AscendSimdVFLowerParallel() # type: ignore +def VerifyParallelToPTO(): + """Read-only verification of the PTO Parallel input contract. + + Runs on the PTO target only: SIMD_VF blocks whose T.Parallel units fall + outside the supported mapping contract are rejected, while functions + without a PTO target pass through unchanged. + """ + return _ffi_api.VerifyParallelToPTO() # type: ignore + + +def LegalizeParallelToPTO(): + """Placeholder for the PTO Parallel legalization stage. + + Only the registration exists for now, so that the formal PTO pipeline + order is fixed; the safe scalar If/Select/Min/Max rewrites are left to + a later change. The placeholder preserves the input IR and rejects + nothing. + """ + return _ffi_api.LegalizeParallelToPTO() # type: ignore + + +def VectorizeParallelToPTO(): + """Rewrite verified T.Parallel loops inside PTO SIMD_VF blocks into + logical VMI code. + + Runs on the PTO target only; functions without a PTO target pass + through unchanged. + """ + return _ffi_api.VectorizeParallelToPTO() # type: ignore + + def AscendThreadSync(storage_scope: str): """Insert thread-storage synchronization independently within each SIMT_VF.""" return _ffi_api.AscendThreadSync(storage_scope) @@ -245,6 +276,7 @@ def RewriteFlagToBuf(): "NormalizeBufferVersion", "NormalizeControlFlowForSchedule", "NormalizeConflictHints", + "LegalizeParallelToPTO", "PrepareMultiBuffer", "ResolveCore", "RewriteAscendBufferVersionLayout", @@ -253,4 +285,6 @@ def RewriteFlagToBuf(): "RewriteFp4ToFp4x2", "RewriteFlagToBuf", "UnrollLoopSkipVF", + "VectorizeParallelToPTO", + "VerifyParallelToPTO", ] From 5be19dd9a75879fbee1184154327e4b80e80a62d Mon Sep 17 00:00:00 2001 From: qukelin Date: Thu, 15 Oct 2026 10:00:24 +0800 Subject: [PATCH 03/10] fix(pto): type allreduce initial-value constants by carrier signedness The si32 reducer strategy serves both signed si32 and signless i32 carriers, but created its initial-value constant as signless pto.i32. With TIR signedness now carried into PTODSL, scalar.select mixing an si32 carrier with an i32 constant is rejected by current PTOAS (pto.select integer arms require matching types). Both creation points (_emit_warp_hw_reduce and _emit_cross_warp_reduce) now share _identity_const: the strategy key still selects the constant's numeric value, while its integer type follows the carrier value's own signedness (signed carriers get pto.si32, signless keep pto.i32). --- tilelang/contrib/ptodsl/allreduce.py | 35 ++++++++++++++++++++-------- 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/tilelang/contrib/ptodsl/allreduce.py b/tilelang/contrib/ptodsl/allreduce.py index e21ca2f1ac..5246224894 100644 --- a/tilelang/contrib/ptodsl/allreduce.py +++ b/tilelang/contrib/ptodsl/allreduce.py @@ -145,8 +145,9 @@ def shuffle_bfly(value, offset): "min": _redux_min, } -# Run the si32 strategy on the signless carrier with identity constants typed -# to match, instead of reinterpreting to si32 and back. +# Default constant dtypes for the reduce initial values. The si32 strategy +# serves both signed and signless carriers; _identity_const retypes the +# constant to the carrier's own signedness at creation time. _REDUCER_IDENTITY_DTYPE = { "f32": _f32_dtype, "f16": _f16_dtype, @@ -155,6 +156,26 @@ def shuffle_bfly(value, offset): } +def _identity_const(x, *, reducer: str, dtype: str): + """Reduce initial-value constant whose integer type follows the carrier. + + TIR integer types are signed, so the PTO codegen emits them as pto.siN + values and the "si32" strategy normally serves an si32 carrier (a + signless i32 carrier remains accepted for hand-written PTODSL). The + ``dtype`` strategy key selects the constant's numeric value, while its + integer type follows the carrier value ``x``: scalar.select/combine + must not mix signed with signless integers (current PTOAS rejects it). + Signed carriers get pto.si32; signless carriers keep pto.i32. + """ + const_dtype = _REDUCER_IDENTITY_DTYPE[dtype] + if dtype == "si32" and _integer_signedness(unwrap_surface_value(x).type) == "signed": + const_dtype = pto.si32 + return _const( + _REDUCER_IDENTITY[reducer][dtype], + dtype=_resolve(const_dtype), + ) + + # ── butterfly ────────────────────────────────────────────────────────────────── @@ -180,10 +201,7 @@ def _emit_warp_hw_reduce(x, *, threads: int, lane_in_warp, dtype: str, reducer: if groups == 1: return redux_fn(x) - c_identity = _const( - _REDUCER_IDENTITY[reducer][dtype], - dtype=_resolve(_REDUCER_IDENTITY_DTYPE[dtype]), - ) + c_identity = _identity_const(x, reducer=reducer, dtype=dtype) my_group = lane_in_warp // threads for g in range(groups): @@ -335,10 +353,7 @@ def _emit_cross_warp_reduce(x, scratch, *, dtype, threads, scale, thread_offset, One call can cover several consecutive groups of ``threads`` threads. """ num_warps = threads // 32 - c_identity = _const( - _REDUCER_IDENTITY[reducer][dtype], - dtype=_resolve(_REDUCER_IDENTITY_DTYPE[dtype]), - ) + c_identity = _identity_const(x, reducer=reducer, dtype=dtype) tx, wid, lid = _simt_warp_indexing(thread_offset=thread_offset) warp_val = _warp_partial_reduce(x, scale=scale, reducer=reducer) _write_warp_partials(warp_val, scratch, wid=wid, lid=lid, scale=scale) From adb5b339c01fa9079fc48623dc59ae239ba8c25c Mon Sep 17 00:00:00 2001 From: qukelin Date: Sat, 17 Oct 2026 11:01:17 +0800 Subject: [PATCH 04/10] fix(pto): describe actual VMI memory access footprints --- src/ascend/codegen/codegen_pto.cc | 10 +- .../pto/test_pto_vmi_memory_access.py | 867 ++++++++++++++++++ .../pto/test_pto_vmi_memory_access_run.py | 511 +++++++++++ tilelang/ascend/language/vmi.py | 359 ++++++-- 4 files changed, 1673 insertions(+), 74 deletions(-) create mode 100644 testing/ascend/language/pto/test_pto_vmi_memory_access.py create mode 100644 testing/ascend/language/pto/test_pto_vmi_memory_access_run.py diff --git a/src/ascend/codegen/codegen_pto.cc b/src/ascend/codegen/codegen_pto.cc index 7cc457952d..2d3e832e84 100644 --- a/src/ascend/codegen/codegen_pto.cc +++ b/src/ascend/codegen/codegen_pto.cc @@ -5049,7 +5049,15 @@ void CodeGenTileLangPTO::PrintPtoVmiCall_(const CallNode *op, os << ", "; } os << key << "="; - if (op_name == "tl.vmi.vcvt" && key == "to_dtype" && op->dtype.is_int()) { + if (op_name == "tl.vmi.vstore" && key == "dist_mode") { + // TileLang uses "dintlv" for two-input stores; PTODSL uses "intlv". + // Load operations keep "dintlv". Map the store mode only at emission; + // tuple recognition above uses the original annotation. + const auto *mode = value.as(); + ICHECK(mode != nullptr) << "tl.vmi.vstore dist_mode must be a string"; + os << "\"" << (mode->value == "dintlv" ? "intlv" : mode->value) << "\""; + } else if (op_name == "tl.vmi.vcvt" && key == "to_dtype" && + op->dtype.is_int()) { os << PtoSignedIntegerTypeName(op->dtype.element_of()); } else if (op_name == "tl.vmi.vload" && key == "size" && op->dtype.element_of().is_float4_e2m1fn()) { diff --git a/testing/ascend/language/pto/test_pto_vmi_memory_access.py b/testing/ascend/language/pto/test_pto_vmi_memory_access.py new file mode 100644 index 0000000000..9897236894 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_vmi_memory_access.py @@ -0,0 +1,867 @@ +"""VMI memory-access descriptors and interface adaptation (structural). + +Hardware independent. Each case lowers the real TileLang PTO pipeline (no +device execution) and inspects the produced TIR/PTODSL. Cases that the +backend can express also call ``tilelang.compile``, which runs PTODSL -> PTOAS +(VPTO) -> Bisheng; that is compile evidence, not numerical evidence (device +results live in ``test_pto_vmi_memory_access_run.py``). + +Contracts verified here: + +* ``T.vmi.vload``/``T.vmi.vstore`` access descriptors carry the true start + address and a per-access-mode effective extent, with the VMI instruction + offset fixed at 0, so the scheduler sees the actual accessed range. +* Ordinary continuous loads round up to full physical vectors; ordinary + broadcast keeps the conservative extent; ``dintlv`` loads declare both + results and a two-input ``dintlv`` store declares the full 2*N span (its + PTODSL attribute is emitted as ``intlv``); group envelopes cover all + groups (slot/broadcast count one element per group, a static negative + stride relocates the lowest address, a dynamic stride falls back to the + whole tight buffer); ``block_stride`` extents are measured in 32B blocks; + genuine FP4 descriptors stay logical and are packed exactly once (odd + starts are rejected); packed (PK) store extents stay logical target + elements; an explicit ``access_ptr`` keeps the caller-declared descriptor + plus the extra offset and the UB scope check; gather/scatter preserve + non-zero bases. +* Masks never shrink the may-write extent, and a logical ``size=80`` load + declares both physical reads. +* Removed interface parameters (``repeat_stride``, vload ``unpack`` / + ``to_dtype``, vstore ``pmode``) fail in a legal scope. + +``size=80`` VMI vectors are not constructible by the PTODSL VMI API +(lane counts are 1/2/4/8/64/128/256). Their real lowering is verified with a +direct PTOAS compile of the PTO form the descriptor describes; the TileLang +diagnostics/descriptors are verified structurally. +""" + +from __future__ import annotations + +import os +import subprocess +import tempfile +from contextlib import contextmanager + +import pytest +import torch # noqa: F401 (torch/torch_npu must import before tilelang/ptodsl) +import torch_npu # noqa: F401 + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + +FP4 = "float4_e2m1fn" +F16 = "float16" +F32 = "float32" +F8_E4M3 = "float8_e4m3fn" + + +# --------------------------------------------------------------------------- +# Pipeline/diagnostic helpers +# --------------------------------------------------------------------------- +def _lower_pto(func): + """Lower *func* for the PTO target and return the compiled artifact.""" + target = normalize_pto_target("pto") + with target: + return _lower(func, target=target) + + +def _compile_pto(func): + """Compile through the PTODSL -> PTOAS (VPTO) -> Bisheng pipeline.""" + return tilelang.compile(func, target="pto", out_idx=-1) + + +def _iter_calls(artifact, op_name): + for func in artifact.device_mod.functions.values(): + found = [] + + def visit(node, found=found): + if isinstance(node, tvm.tirx.Call) and getattr(node.op, "name", "") == op_name: + found.append(node) + + tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + for node in found: + yield func, node + + +def _resolve_binding(expr, func): + """Resolve a local handle binding (``p = access_ptr(...)``) to its value.""" + if not isinstance(expr, tvm.tirx.Var): + return expr + found = {} + + def visit(node): + if isinstance(node, (tvm.tirx.Let, tvm.tirx.Bind)) and node.var is not None and node.var.same_as(expr): + found["value"] = node.value + + tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + return found.get("value", expr) + + +def _const(expr): + from tvm import arith + + simplified = arith.Analyzer().simplify(expr) + assert isinstance(simplified, tvm.tirx.IntImm), f"expected a constant, got {simplified!r}" + return int(simplified.value) + + +def _single(artifact, op_name): + pairs = list(_iter_calls(artifact, op_name)) + assert len(pairs) == 1, f"expected exactly one {op_name}, got {len(pairs)}" + return pairs[0] + + +def _ptr_position(call, op_name): + if op_name == "tl.vmi.vstore" and "dist_mode" in call.annotations: + return 2 + return 1 if op_name == "tl.vmi.vstore" else 0 + + +def _descriptor(artifact, op_name): + """Return (descriptor_offset, extent, rw_mask, instruction_offset).""" + func, call = _single(artifact, op_name) + ptr_pos = _ptr_position(call, op_name) + ptr = _resolve_binding(call.args[ptr_pos], func) + assert isinstance(ptr, tvm.tirx.Call) and ptr.op.name.endswith("tvm_access_ptr"), ( + f"{op_name} address is not an access descriptor: {ptr!r}" + ) + return ( + _const(ptr.args[2]), + _const(ptr.args[3]), + _const(ptr.args[4]), + _const(call.args[ptr_pos + 1]), + ) + + +def _vload_descriptor(func): + return _descriptor(_lower_pto(func), "tl.vmi.vload") + + +def _vstore_descriptor(func): + return _descriptor(_lower_pto(func), "tl.vmi.vstore") + + +@contextmanager +def _capture_rewrite_input(): + """Capture the module handed to RewriteFp4ToFp4x2, then delegate.""" + captured = {} + original = ascend_transform.RewriteFp4ToFp4x2 + + def factory(): + rewrite = original() + + def run(mod): + captured["module"] = mod + return rewrite(mod) + + return run + + ascend_transform.RewriteFp4ToFp4x2 = factory + try: + yield captured + finally: + ascend_transform.RewriteFp4ToFp4x2 = original + + +def _descriptor_in_module(mod, op_name): + class _Artifact: + device_mod = mod + + return _descriptor(_Artifact(), op_name) + + +def _ptoas_vpto(pto_text: str) -> str: + """Compile direct PTO IR through PTOAS and return the VPTO text.""" + work = tempfile.mkdtemp(prefix="vmi-vpto-", dir=os.environ.get("TMPDIR") or None) + pto_path = os.path.join(work, "case.pto") + vpto_path = os.path.join(work, "case.vpto.mlir") + with open(pto_path, "w", encoding="utf-8") as file: + file.write(pto_text) + result = subprocess.run( + ["ptoas", "--pto-arch=a5", "--pto-backend=vpto", "--emit-vpto", pto_path, "-o", vpto_path], + text=True, + capture_output=True, + ) + assert result.returncode == 0, f"PTOAS rejected the probe:\n{result.stderr}" + with open(vpto_path, encoding="utf-8") as file: + return file.read() + + +# --------------------------------------------------------------------------- +# Ordinary non-zero BufferLoad start +# --------------------------------------------------------------------------- +def _load_store_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((320,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work[0:256]) + T.fill(work[256:320], T.float32(-5.0)) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + v = T.vmi.vload(work[64], size=64) + T.vmi.vstore(v, work[192], mask) + T.copy(work[0:320], C) + + return main + + +def test_load_descriptor_originates_at_actual_load(): + """vload(work[64], size=64) declares [64,128) and instruction offset 0.""" + artifact = _lower_pto(_load_store_kernel()) + assert _descriptor(artifact, "tl.vmi.vload") == (64, 64, 1, 0) + + source = artifact.kernel_source + vload_line = next(ln.strip() for ln in source.splitlines() if "vmi.vload" in ln) + assert "addptr" in vload_line and ", 64)" in vload_line, vload_line + assert ", 0, size=64)" in vload_line, vload_line + + +def test_store_descriptor_originates_at_actual_destination(): + """vstore(..., work[192]) declares may-write [192,256) with offset 0.""" + artifact = _lower_pto(_load_store_kernel()) + assert _descriptor(artifact, "tl.vmi.vstore") == (192, 64, 2, 0) + + source = artifact.kernel_source + vstore_line = next(ln.strip() for ln in source.splitlines() if "vmi.vstore" in ln) + assert "addptr" in vstore_line and ", 192)" in vstore_line, vstore_line + assert "), 0, mask)" in vstore_line, vstore_line + + +def test_nonzero_load_and_store_compile_through_backend(): + kernel = _compile_pto(_load_store_kernel()) + source = kernel.get_kernel_source() + assert "pto.vmi.vload(" in source and "pto.vmi.vstore(" in source, source + + +# --------------------------------------------------------------------------- +# Buffer / BufferRegion / elem_offset +# --------------------------------------------------------------------------- +def _buffer_region_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + v = T.vmi.vload(work[0:64], size=64) + T.vmi.vstore(v, work[128], mask) + T.copy(work, C) + + return main + + +def _elem_offset_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + view = tvm.tirx.decl_buffer((128,), F32, data=work.data, elem_offset=16, scope="shared.dyn") + mask = T.vmi.create_mask(64, size=64) + v = T.vmi.vload(view[32], size=64) + T.vmi.vstore(v, view[32], mask) + T.copy(work, C) + + return main + + +def test_buffer_region_source_declares_from_region_minimum(): + assert _vload_descriptor(_buffer_region_kernel()) == (0, 64, 1, 0) + + +def test_elem_offset_counted_exactly_once(): + """A view with elem_offset=16 loaded at index 32 starts at element 48.""" + artifact = _lower_pto(_elem_offset_kernel()) + assert _descriptor(artifact, "tl.vmi.vload") == (48, 64, 1, 0) + assert _descriptor(artifact, "tl.vmi.vstore") == (48, 64, 2, 0) + + +# --------------------------------------------------------------------------- +# Ordinary broadcast +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("dtype", "size"), + [("int8", 256), (F16, 128), (F32, 128)], + ids=["b8", "b16", "b32"], +) +def test_ordinary_broadcast_keeps_conservative_extent(dtype, size): + """brc reads one source element; the conservative ``size`` extent is kept.""" + + # b8 full-VL broadcast at 64 reads [64,320) and stores [128,384); size + # the storage (and its initialization) for the real footprint. + @T.prim_func + def main(A: T.Tensor((512,), dtype), C: T.Tensor((512,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((512,), dtype) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + v = T.vmi.vload(work[64], size=size, dist_mode="brc") + T.vmi.vstore(v, work[128], mask) + T.copy(work, C) + + artifact = _lower_pto(main) + # The declaration is not re-expressed as a broadcast output width: it stays + # `size` address elements (a conservative over-report). + assert _descriptor(artifact, "tl.vmi.vload") == (64, size, 1, 0) + assert 'dist_mode="brc"' in artifact.kernel_source + + kernel = _compile_pto(main) + assert "pto.vmi.vload(" in kernel.get_kernel_source() + + +# --------------------------------------------------------------------------- +# Physical-vector rounding of ordinary continuous loads +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("dtype", "size", "expected"), + [ + ("int8", 300, 512), # b8 VL = 256 + (F16, 200, 256), # b16 VL = 128 + (F32, 80, 128), # b32 VL = 64 (structural: PTODSL cannot build 80) + ], + ids=["b8", "b16", "b32-size80"], +) +def test_continuous_load_covers_full_physical_vectors(dtype, size, expected): + @T.prim_func + def main(A: T.Tensor((1024,), dtype), C: T.Tensor((1024,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((1024,), dtype) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[64], size=size) + T.vmi.vstore(v, work[512]) + T.copy(work, C) + + assert _vload_descriptor(main) == (64, expected, 1, 0) + assert _vstore_descriptor(main) == (512, size, 2, 0) + + +# --------------------------------------------------------------------------- +# dintlv load / two-input store +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("size", [64, 128], ids=["single-vl", "multi-vl"]) +def test_dintlv_load_declares_both_results(size): + @T.prim_func + def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + # Each result has `size` lanes: the store mask must match. + mask = T.vmi.create_mask(size, size=size) + even, odd = T.vmi.vload(work[64], size=size, dist_mode="dintlv") + T.vmi.vstore(even, work[320], mask) + T.vmi.vstore(odd, work[384], mask) + T.copy(work, C) + + artifact = _lower_pto(main) + assert _descriptor(artifact, "tl.vmi.vload") == (64, 2 * size, 1, 0) + + +@pytest.mark.parametrize("mask_lanes", [0, 10, 64], ids=["all-false", "partial", "all-true"]) +def test_dintlv_store_declares_full_two_input_span(mask_lanes): + """The store may-write span is 2*N regardless of the ineffective mask.""" + + @T.prim_func + def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(mask_lanes, size=64) + even, odd = T.vmi.vload(work[64], size=64, dist_mode="dintlv") + T.vmi.vstore((even, odd), work[192], mask, dist_mode="dintlv") + T.copy(work, C) + + artifact = _lower_pto(main) + assert _descriptor(artifact, "tl.vmi.vstore") == (192, 128, 2, 0) + + source = artifact.kernel_source + vstore_line = next(ln.strip() for ln in source.splitlines() if "vmi.vstore" in ln) + # IR/API keep "dintlv"; only the PTODSL attribute is mapped to "intlv". + assert 'dist_mode="intlv"' in vstore_line, vstore_line + assert 'dist_mode="dintlv"' not in vstore_line, vstore_line + _, store_call = _single(artifact, "tl.vmi.vstore") + assert store_call.annotations["dist_mode"].value == "dintlv" + + +# --------------------------------------------------------------------------- +# Logical size 80 / masks never shrink the footprint +# --------------------------------------------------------------------------- +def test_size80_load_declares_two_full_physical_reads(): + """b32 size 80 loads two 64-lane physical reads, so the extent is 128.""" + + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[64], size=80) + T.vmi.vstore(v, work[128]) + T.copy(work, C) + + assert _vload_descriptor(main) == (64, 128, 1, 0) + assert _vstore_descriptor(main) == (128, 80, 2, 0) + + +def test_size80_backend_lowering_covers_second_physical_chunk(): + """Direct PTOAS compile of the size-80 PTO form the descriptor describes. + + The PTODSL VMI API cannot construct an 80-lane VMI vreg + (``_check_vmi_lane_count``), so this verifies the backend lowering itself: + two full vlds at offsets 64 and 128. + """ + pto_text = """module { + func.func @group1_n80_s80() attributes {pto.kernel} { + %start = arith.constant 64 : index + %stride = arith.constant 80 : index + %base = arith.constant 0 : i64 + %dest = arith.constant 4096 : i64 + %src = pto.castptr %base : i64 -> !pto.ptr + %dst = pto.castptr %dest : i64 -> !pto.ptr + pto.vecscope { + %v = pto.vmi.vload %src[%start], %stride {group = 1 : i64} : !pto.ptr -> !pto.vmi.vreg<80xf32> + pto.vmi.vstore %v, %dst[%start], %stride {group = 1 : i64} : !pto.vmi.vreg<80xf32>, !pto.ptr + } + return + } +} +""" + vpto = _ptoas_vpto(pto_text) + assert "pto.vlds %0[%c64]" in vpto, vpto + assert "pto.vlds %0[%c128]" in vpto, vpto + + +def test_store_mask_activity_does_not_shrink_may_write(): + """A partial or all-false mask keeps the full 64-lane may-write span.""" + + def kernel(mask_active): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(mask_active, size=64) + v = T.vmi.vload(work[0], size=64) + T.vmi.vstore(v, work[128], mask) + T.copy(work, C) + + return main + + for active in (0, 10, 64): + assert _vstore_descriptor(kernel(active)) == (128, 64, 2, 0), active + + +# --------------------------------------------------------------------------- +# Group envelopes, negative and dynamic stride +# --------------------------------------------------------------------------- +def test_group_ordinary_envelope_covers_all_groups(): + """A=64, N=128, G=2, S=128 => [64,256).""" + + @T.prim_func + def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + v = T.vmi.vload(work[64], size=128, group=2, stride=128) + T.vmi.vstore(v, work[256], mask) + T.copy(work, C) + + assert _vload_descriptor(main) == (64, 192, 1, 0) + + +def test_group_slot_and_broadcast_use_one_element_per_group(): + @T.prim_func + def slot(A: T.Tensor((768,), F32), C: T.Tensor((768,), F32)): + with T.Kernel(1): + # G=8/S=64: source touches index 512; the destination starting at + # 256 touches 256 + 7*64 = 704. Both need legal, initialized + # storage, so allocate/initialize 768 elements. + work = T.alloc_shared((768,), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[64], size=8, group=8, stride=64) + T.vmi.vstore(v, work[256], group=8, stride=64) + T.copy(work, C) + + assert _vload_descriptor(slot) == (64, 449, 1, 0) + assert _vstore_descriptor(slot) == (256, 449, 2, 0) + + @T.prim_func + def group_brc(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[64], size=128, group=2, stride=128, dist_mode="brc") + T.vmi.vstore(v, work[256], group=2, stride=128) + T.copy(work, C) + + assert _vload_descriptor(group_brc) == (64, 129, 1, 0) + + +def test_group_static_negative_stride_relocates_lowest_address(): + """A=192, N=128, G=2, S=-128 => envelope [64,256) and instruction offset 128.""" + + @T.prim_func + def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[192], size=128, group=2, stride=-128) + T.vmi.vstore(v, work[192], group=2, stride=-128) + T.copy(work, C) + + artifact = _lower_pto(main) + assert _descriptor(artifact, "tl.vmi.vload") == (64, 192, 1, 128) + assert _descriptor(artifact, "tl.vmi.vstore") == (64, 192, 2, 128) + + +def test_group_dynamic_stride_declares_whole_tight_buffer(): + @T.prim_func + def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32), stride: T.int32): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[192], size=128, group=2, stride=stride) + T.vmi.vstore(v, work[192], group=2, stride=stride) + T.copy(work, C) + + artifact = _lower_pto(main) + assert _descriptor(artifact, "tl.vmi.vload") == (0, 512, 1, 192) + assert _descriptor(artifact, "tl.vmi.vstore") == (0, 512, 2, 192) + + +def test_group_negative_stride_nd_uses_whole_buffer_plan(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((4, 64), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[3, 0], size=128, group=2, stride=-64) + T.vmi.vstore(v, work[3, 0], group=2, stride=-64) + T.copy(work, C) + + artifact = _lower_pto(main) + assert _descriptor(artifact, "tl.vmi.vload") == (0, 256, 1, 192) + + +# --------------------------------------------------------------------------- +# block_stride +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("dtype", "size", "block_stride", "expected_extent"), + [ + (F32, 64, 1, 64), # one full b32 vector, contiguous blocks + (F32, 64, 2, 120), # [64, 184) + (F32, 128, 2, 248), # two full vectors: [64, 312) + (F16, 128, 2, 240), # b16: 8 blocks of 16 elements + ], + ids=["b32-single-vl-s1", "b32-single-vl-s2", "b32-double-vl-s2", "b16-single-vl-s2"], +) +def test_block_stride_envelope_uses_32b_blocks(dtype, size, block_stride, expected_extent): + # The destination at 320 spans up to 320 + expected_extent = 568 (b32 + # double-VL) / 560 (b16); allocate and initialize the real footprint. + @T.prim_func + def main(A: T.Tensor((768,), dtype), C: T.Tensor((768,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((768,), dtype) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + v = T.vmi.vload(work[64], size=size, block_stride=block_stride) + T.vmi.vstore(v, work[320], mask, block_stride=block_stride) + T.copy(work, C) + + artifact = _lower_pto(main) + assert _descriptor(artifact, "tl.vmi.vload") == (64, expected_extent, 1, 0) + assert _descriptor(artifact, "tl.vmi.vstore") == (320, expected_extent, 2, 0) + + +# --------------------------------------------------------------------------- +# Genuine FP4 +# --------------------------------------------------------------------------- +def _fp4_kernel(index=64, size=128): + @T.prim_func + def main(A: T.Tensor((256,), FP4), C: T.Tensor((256,), FP4)): + with T.Kernel(1): + work = T.alloc_shared((256,), FP4) + out = T.alloc_shared((256,), FP4) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(size // 2, size=size // 2) + v = T.vmi.vload(work[index], size=size) + T.vmi.vstore(v, out[index], mask) + T.copy(out, C) + + return main + + +def test_fp4_descriptor_stays_logical_then_packed_once(): + with _capture_rewrite_input() as captured: + artifact = _lower_pto(_fp4_kernel()) + before = captured["module"] + + # Logical FP4 element units before the rewrite. + assert _descriptor_in_module(before, "tl.vmi.vload") == (64, 128, 1, 0) + assert _descriptor_in_module(before, "tl.vmi.vstore") == (64, 128, 2, 0) + + # Exactly one physical packing: the single rewrite pass halves both. + rewritten = ascend_transform.RewriteFp4ToFp4x2()(before) + assert _descriptor_in_module(rewritten, "tl.vmi.vload") == (32, 64, 1, 0) + assert _descriptor_in_module(rewritten, "tl.vmi.vstore") == (32, 64, 2, 0) + + source = artifact.kernel_source + load_line = next(ln.strip() for ln in source.splitlines() if "vmi.vload" in ln) + assert "addptr" in load_line and "), 32)" in load_line, load_line + assert load_line.rstrip().endswith("size=(128 // 2))"), load_line + + +def test_fp4_odd_start_is_still_rejected(): + with pytest.raises(Exception, match="fp4 access_ptr offset must be even"): + _lower_pto(_fp4_kernel(index=65)) + + +# --------------------------------------------------------------------------- +# Packed (PK) stores keep logical target-element extents +# --------------------------------------------------------------------------- +def _bits(dtype) -> int: + return int(tvm.DataType(dtype).bits) + + +def _merged_target_offset(src, dst, count=512) -> int: + """Offset of a target buffer merged after ``count`` source elements. + + MergeUBAllocations packs the source span first in these two-buffer + kernels; descriptors are in destination-element units. Deriving the + expectation keeps the assertion on the absolute may-write start while + making the layout assumption explicit. + """ + source_bytes = count * _bits(src) // 8 + return source_bytes // (_bits(dst) // 8) + + +@pytest.mark.parametrize( + ("src", "dst", "vcvt_kwargs"), + [ + (F32, F16, {"rounding": "R", "saturate": "SAT"}), + (F32, F8_E4M3, {"rounding": "R", "saturate": "SAT"}), + ("int16", "int8", {"saturate": "NOSAT"}), + ], + ids=["pk-b32", "pk4-b32", "pk-b16"], +) +def test_packed_store_declares_logical_target_elements(src, dst, vcvt_kwargs): + size = 128 + + @T.prim_func + def main(A: T.Tensor((512,), src), C: T.Tensor((512,), dst)): + with T.Kernel(1): + source = T.alloc_shared((512,), src) + target = T.alloc_shared((512,), dst) + T.copy(A, source) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + v = T.vmi.vload(source[64], size=size) + narrowed = T.vmi.vcvt(v, to_dtype=dst, **vcvt_kwargs) + T.vmi.vstore(narrowed, target[64], mask) + T.copy(target, C) + + artifact = _lower_pto(main) + target_offset = _merged_target_offset(src, dst) + 64 + # The may-write span stays the logical target element count (128), not a + # physical vreg capacity or a PK-derived fraction. + assert _descriptor(artifact, "tl.vmi.vstore") == (target_offset, size, 2, 0) + src_extent = 128 if src == F32 else size + assert _descriptor(artifact, "tl.vmi.vload") == (64, src_extent, 1, 0) + + kernel = _compile_pto(main) + assert "pto.vmi.vcvt(" in kernel.get_kernel_source() + + +# --------------------------------------------------------------------------- +# Explicit access_ptr + extra offset +# --------------------------------------------------------------------------- +def _explicit_pointer_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + p = T.access_ptr(work[32], "r", extent=128) + v = T.vmi.vload(p, 16, size=64) + T.vmi.vstore(v, work[128], mask) + T.copy(work, C) + + return main + + +def test_explicit_pointer_keeps_declaration_and_extra_offset(): + artifact = _lower_pto(_explicit_pointer_kernel()) + func, call = _single(artifact, "tl.vmi.vload") + ptr = _resolve_binding(call.args[0], func) + # The caller-declared descriptor is not overwritten (offset 32, extent 128). + assert _const(ptr.args[2]) == 32 + assert _const(ptr.args[3]) == 128 + # The extra offset is preserved, not zeroed or folded twice. + assert _const(call.args[1]) == 16 + + source = artifact.kernel_source + lines = [ln.strip() for ln in source.splitlines()] + definition = next(ln for ln in lines if ln.startswith("p = ")) + assert "pto.addptr(" in definition and ", 32)" in definition, definition + # The VMI instruction offset stays 16: the base is counted once. + assert "vload(p, 16, size=64)" in source, source + + +def test_explicit_pointer_scope_validation_is_preserved(): + """Recovering the element type from an immediate access_ptr does not + bypass the UB scope check for indirect-address operations.""" + + def build(): + @T.prim_func + def main(A: T.Tensor((256,), F32), I: T.Tensor((64,), "uint32"), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + index = T.alloc_shared((64,), "uint32") + T.copy(A, work) + T.copy(I, index) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + offsets = T.vmi.vload(index[0], size=64) + p = T.access_ptr(A[0], "r", extent=64) + v = T.vmi.vgather(p, offsets, mask) + T.vmi.vstore(v, work[192], mask) + T.copy(work, C) + + return main + + with pytest.raises(TypeError, match="requires a UB pointer"): + build() + + +# --------------------------------------------------------------------------- +# Gather/scatter non-zero bases (no footprint analysis) +# --------------------------------------------------------------------------- +def _gather_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), I: T.Tensor((64,), "uint32"), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + index = T.alloc_shared((64,), "uint32") + T.copy(A, work) + T.copy(I, index) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + offsets = T.vmi.vload(index[0], size=64) + v = T.vmi.vgather(work[64], offsets, mask) + T.vmi.vstore(v, work[192], mask) + T.copy(work, C) + + return main + + +def _scatter_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), I: T.Tensor((64,), "uint32"), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + index = T.alloc_shared((64,), "uint32") + T.copy(A, work) + T.copy(I, index) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + offsets = T.vmi.vload(index[0], size=64) + value = T.vmi.vload(work[128], size=64) + T.vmi.vscatter(value, work[64], offsets, mask) + T.copy(work, C) + + return main + + +def test_gather_descriptor_preserves_nonzero_base(): + artifact = _lower_pto(_gather_kernel()) + _, call = _single(artifact, "tl.vmi.vgather") + ptr = call.args[0] + assert _const(ptr.args[2]) == 64 + # Lanes extent: generic index footprints are explicitly out of scope. + assert _const(ptr.args[3]) == 64 + assert _const(call.args[1]) == 0 + line = next(ln.strip() for ln in artifact.kernel_source.splitlines() if "vmi.vgather" in ln) + assert "addptr" in line and ", 64)" in line, line + + +def test_scatter_descriptor_preserves_nonzero_base(): + artifact = _lower_pto(_scatter_kernel()) + _, call = _single(artifact, "tl.vmi.vscatter") + assert _const(call.args[1].args[2]) == 64 + assert _const(call.args[2]) == 0 + line = next(ln.strip() for ln in artifact.kernel_source.splitlines() if "vmi.vscatter" in ln) + assert "addptr" in line and ", 64)" in line, line + + +# --------------------------------------------------------------------------- +# Removed VMI interface parameters +# --------------------------------------------------------------------------- +def _vmi_body_kernel(body): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + body(work, mask) + T.copy(work, C) + + return main + + +@pytest.mark.parametrize( + ("case_id", "body", "message"), + [ + ( + "vload-repeat-stride", + lambda work, mask: T.vmi.vload(work[0], size=64, repeat_stride=0), + "unexpected keyword argument 'repeat_stride'", + ), + ( + "vstore-repeat-stride", + lambda work, mask: T.vmi.vstore(T.vmi.vload(work[0], size=64), work[0], mask, repeat_stride=0), + "unexpected keyword argument 'repeat_stride'", + ), + ( + "vload-to-dtype", + lambda work, mask: T.vmi.vload(work[0], size=64, to_dtype=F16), + "unexpected keyword argument 'to_dtype'", + ), + ( + "vload-unpack", + lambda work, mask: T.vmi.vload(work[0], size=64, dist_mode="unpack"), + "does not support dist_mode='unpack'", + ), + ( + "vstore-pmode", + lambda work, mask: T.vmi.vstore(T.vmi.vload(work[0], size=64), work[0], mask, pmode="zero"), + "unexpected keyword argument 'pmode'", + ), + ], +) +def test_removed_vmi_parameters_fail_in_legal_scope(case_id, body, message): + with pytest.raises(TypeError, match=message): + _vmi_body_kernel(body) diff --git a/testing/ascend/language/pto/test_pto_vmi_memory_access_run.py b/testing/ascend/language/pto/test_pto_vmi_memory_access_run.py new file mode 100644 index 0000000000..fb0cdaa12d --- /dev/null +++ b/testing/ascend/language/pto/test_pto_vmi_memory_access_run.py @@ -0,0 +1,511 @@ +"""Device acceptance for VMI memory-access descriptors. + +Every kernel runs through the real ``tilelang.compile(target="pto")`` pipeline +with the automatic scheduler enabled (the generated PTODSL must contain the +MTE sync flags; no test disables scheduling). Inputs are built on CPU with +torch, copied to the NPU, and compared against an independent CPU reference +after ``torch.npu.synchronize()``. + +Critical dependency paths execute 20 fresh trials each with independently +prepared inputs and a fresh output tensor, and also assert the +generated-source order and MTE sync pair: + +* a DMA that writes only the non-zero ``[64,128)`` region feeding + ``vload(work[64], size=64)``; +* ``vstore(..., work[64])`` feeding a DMA that reads only ``[64,128)``. + +The remaining cases are one device-numerics representative per access mode: +group dynamic stride (positive/zero/negative), block_stride, FP4 raw packed +bytes (buffer and explicit-pointer forms), a packed (PK) narrowing store, an +explicit ``access_ptr`` plus extra offset, and gather/scatter with non-zero +bases (direct and explicit-pointer forms). The device numerical reference is +exact for copy-like kernels and an independent CPU computation (torch cast +reference) for narrowing. + +Non-VMI lane counts such as ``size=80`` cannot be constructed by the PTODSL +VMI API; their descriptor/backend evidence lives in the structural suite +(``test_pto_vmi_memory_access.py``) and is not claimed as device-passed. +""" + +import os +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import torch_npu # noqa: E402,F401 (must import before tilelang/ptodsl) +import tilelang # noqa: E402 +import tilelang.ascend.language as T # noqa: E402 +from tvm import tirx # noqa: E402 + +F16 = "float16" +F32 = "float32" +FP4 = "float4_e2m1fn" + + +# --------------------------------------------------------------------------- +# Runtime helpers +# --------------------------------------------------------------------------- +def _require_npu_runtime(): + if not hasattr(torch, "npu") or not torch.npu.is_available(): + pytest.skip("torch.npu is unavailable in the current environment") + try: + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source) + except Exception as err: # noqa: BLE001 + pytest.skip(f"NPU runtime is unavailable in the current environment: {err}") + + +@pytest.fixture(scope="module") +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +def _compile(func): + kernel = tilelang.compile(func, target="pto", out_idx=-1) + source = kernel.get_kernel_source() + assert "pto.set_flag(" in source and "pto.wait_flag(" in source, ( + "automatic scheduling must stay enabled (no MTE sync in the generated source):\n" + source + ) + return kernel, source + + +def _vpto_text(kernel_source): + """Compile the generated PTODSL through PTOAS and return the VPTO text.""" + import importlib.util + import re + import subprocess + import tempfile + import uuid + + decorator = re.search(r'@pto\.jit\(name="([^"]+)"', kernel_source) + name = decorator.group(1) if decorator is not None else "main_kernel" + work = tempfile.mkdtemp(prefix="vmi-runtime-vpto-", dir=os.environ.get("TMPDIR") or None) + source_path = os.path.join(work, "kernel.ptodsl.py") + pto_path = os.path.join(work, "kernel.pto") + vpto_path = os.path.join(work, "kernel.vpto.mlir") + with open(source_path, "w", encoding="utf-8") as file: + file.write(kernel_source) + spec = importlib.util.spec_from_file_location("vmi_runtime_" + uuid.uuid4().hex, source_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + pto_text = getattr(module, name).compile().mlir_text() + with open(pto_path, "w", encoding="utf-8") as file: + file.write(pto_text) + result = subprocess.run( + ["ptoas", "--pto-arch=a5", "--pto-backend=vpto", "--emit-vpto", pto_path, "-o", vpto_path], + text=True, + capture_output=True, + ) + assert result.returncode == 0, f"PTOAS rejected the generated kernel:\n{result.stderr}" + with open(vpto_path, encoding="utf-8") as file: + return file.read() + + +def _assert_addptr_pointer_use(source, definition_fragment, call_fragment): + lines = [line.strip() for line in source.splitlines()] + definitions = [line for line in lines if "= pto.addptr(" in line and definition_fragment in line] + assert len(definitions) == 1, f"expected one pointer definition with {definition_fragment!r}: {definitions}" + name = definitions[0].split("=")[0].strip() + assert call_fragment.format(name=name) in source, source + + +def _to_device(value, device): + return value.to(device) if isinstance(value, torch.Tensor) else value + + +def _run_trials(kernel, make_case, device, trials): + """Run *trials* independent cases; return the per-trial outputs.""" + outputs = [] + for trial in range(trials): + inputs, expected = make_case(trial) + device_inputs = [_to_device(value, device) for value in inputs] + output = kernel(*device_inputs) + torch.npu.synchronize() + got = output.cpu() + assert got.dtype == expected.dtype, f"trial {trial}: {got.dtype} != {expected.dtype}" + assert torch.equal(got, expected), f"trial {trial}: numeric mismatch\n got={got.tolist()[:16]}\n exp={expected.tolist()[:16]}" + outputs.append(got) + return outputs + + +def _line_index(lines, predicate, start=0, what="line"): + for index in range(start, len(lines)): + if predicate(lines[index]): + return index + raise AssertionError(f"missing {what} from line {start}:\n" + "\n".join(lines)) + + +def _assert_mte2_to_v_then_load(lines, dma_predicate, descriptor_needle): + dma = _line_index(lines, dma_predicate, what="producer MTE2 DMA") + set_index = _line_index(lines, lambda line: 'set_flag("MTE2", "V"' in line, dma, "MTE2->V set_flag") + wait_index = _line_index(lines, lambda line: 'wait_flag("MTE2", "V"' in line, set_index, "MTE2->V wait_flag") + vload = _line_index(lines, lambda line: "vmi.vload" in line, wait_index, "vmi.vload") + assert dma < set_index < wait_index < vload + assert descriptor_needle in lines[vload], lines[vload] + + +def _assert_store_then_v_to_mte3(lines, dma_predicate, store_needle): + store = _line_index(lines, lambda line: "vmi.vstore" in line, what="vmi.vstore") + set_index = _line_index(lines, lambda line: 'set_flag("V", "MTE3"' in line, store, "V->MTE3 set_flag") + wait_index = _line_index(lines, lambda line: 'wait_flag("V", "MTE3"' in line, set_index, "V->MTE3 wait_flag") + dma = _line_index(lines, dma_predicate, wait_index, "consumer MTE3 DMA") + assert store < set_index < wait_index < dma + assert store_needle in lines[store], lines[store] + + +def _float_inputs(shape, trial, base=0.0, step=1.0): + generator = torch.Generator().manual_seed(1000 + trial) + values = torch.randn(shape, generator=generator) + return values * step + (base + trial * 37.0) + + +# --------------------------------------------------------------------------- +# Critical path 1: DMA writing only [64,128) -> vload(work[64]) +# --------------------------------------------------------------------------- +def _dma_load_kernel(): + @T.prim_func + def main(A: T.Tensor((64,), F32), S: T.Tensor((64,), F32), C: T.Tensor((64,), F32)): + with T.Kernel(1): + work = T.alloc_shared((128,), F32) + T.copy(S, work[64:128]) # sentinel first + T.copy(A, work[64:128]) # producer DMA: only [64,128) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + v = T.vmi.vload(work[64], size=64) + T.vmi.vstore(v, work[0], mask) + T.copy(work[0:64], C) + + return main + + +def test_dma_to_nonzero_load_dependency_20_trials(npu): + kernel, source = _compile(_dma_load_kernel()) + lines = source.splitlines() + _assert_mte2_to_v_then_load( + lines, + lambda line: "mte_gm_ub" in line and "addptr" in line and ", 64)" in line, + "), 0, size=64)", + ) + + def make_case(trial): + a = _float_inputs((64,), trial) + sentinel = torch.full((64,), -777.0 - trial) + return [a, sentinel], a + + _run_trials(kernel, make_case, npu, trials=20) + + +# --------------------------------------------------------------------------- +# Critical path 2: vstore(..., work[64]) -> DMA reading only [64,128) +# --------------------------------------------------------------------------- +def _store_dma_kernel(): + @T.prim_func + def main(A: T.Tensor((64,), F32), S: T.Tensor((64,), F32), C: T.Tensor((64,), F32)): + with T.Kernel(1): + work = T.alloc_shared((128,), F32) + T.copy(S, work[64:128]) # sentinel before the store + T.copy(A, work[0:64]) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + v = T.vmi.vload(work[0], size=64) + T.vmi.vstore(v, work[64], mask) + T.copy(work[64:128], C) # consumer DMA: only [64,128) + + return main + + +def test_nonzero_store_to_dma_dependency_20_trials(npu): + kernel, source = _compile(_store_dma_kernel()) + lines = source.splitlines() + _assert_store_then_v_to_mte3( + lines, + lambda line: "mte_ub_gm" in line and "addptr" in line and ", 64)" in line, + "), 0, mask)", + ) + + def make_case(trial): + a = _float_inputs((64,), trial) + sentinel = torch.full((64,), -888.0 - trial) + return [a, sentinel], a + + _run_trials(kernel, make_case, npu, trials=20) + + +# --------------------------------------------------------------------------- +# Group dynamic stride (positive / zero / negative) +# --------------------------------------------------------------------------- +def _dynamic_stride_kernel(): + @T.prim_func + def main(A: T.Tensor((512,), F32), stride: T.int32, C: T.Tensor((128,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + v = T.vmi.vload(work[192], size=128, group=2, stride=stride) + T.vmi.vstore(v, work[0], mask) + T.copy(work[0:128], C) + + return main + + +@pytest.mark.parametrize("stride", [128, 0, -128], ids=["positive", "zero", "negative"]) +def test_group_dynamic_stride_load(npu, stride): + kernel, _ = _compile(_dynamic_stride_kernel()) + + def make_case(trial): + a = _float_inputs((512,), trial) + second = 192 + stride + return [a, stride], torch.cat([a[192:256], a[second : second + 64]]) + + _run_trials(kernel, make_case, npu, trials=5) + + +# --------------------------------------------------------------------------- +# Block stride: stride 1/2, single and multi VL +# --------------------------------------------------------------------------- +def _block_expected(a, base, lanes, block_stride, vector=64, block=8): + chunks = (lanes + vector - 1) // vector + values = [] + for chunk in range(chunks): + for slot in range(vector // block): + start = base + chunk * (vector // block) * block_stride * block + slot * block_stride * block + values.extend(a[start : start + block].tolist()) + return torch.tensor(values[:lanes], dtype=a.dtype) + + +def _block_load_kernel(size, block_stride): + @T.prim_func + def main(A: T.Tensor((768,), F32), C: T.Tensor((size,), F32)): + with T.Kernel(1): + work = T.alloc_shared((768,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + v = T.vmi.vload(work[64], size=size, block_stride=block_stride) + T.vmi.vstore(v, work[512], mask) + T.copy(work[512 : 512 + size], C) + + return main + + +@pytest.mark.parametrize( + ("size", "block_stride"), + [(64, 1), (64, 2), (128, 2)], + ids=["single-vl-s1", "single-vl-s2", "multi-vl-s2"], +) +def test_block_stride_load(npu, size, block_stride): + kernel, _ = _compile(_block_load_kernel(size, block_stride)) + + def make_case(trial): + a = _float_inputs((768,), trial) + return [a], _block_expected(a, 64, size, block_stride) + + _run_trials(kernel, make_case, npu, trials=3) + + +# --------------------------------------------------------------------------- +# FP4 raw packed bytes: non-zero logical start and explicit pointer offset +# --------------------------------------------------------------------------- +def _fp4_view_kernel(explicit_pointer): + @T.prim_func + def main(B: T.Tensor((128,), "uint8"), Z: T.Tensor((128,), "uint8"), C: T.Tensor((128,), "uint8")): + with T.Kernel(1): + raw_in = T.alloc_shared((128,), "uint8") + raw_out = T.alloc_shared((128,), "uint8") + T.copy(B, raw_in) + T.copy(Z, raw_out) + with T.SimdVF(): + f4_in = tirx.decl_buffer((256,), FP4, data=raw_in.data, scope="shared.dyn") + f4_out = tirx.decl_buffer((256,), FP4, data=raw_out.data, scope="shared.dyn") + mask = T.vmi.create_mask(64, size=64) + if explicit_pointer: + ptr = T.access_ptr(f4_in[0], "r", extent=256) + v = T.vmi.vload(ptr, 64, size=128) + else: + v = T.vmi.vload(f4_in[64], size=128) + T.vmi.vstore(v, f4_out[128], mask) + T.copy(raw_out, C) + + return main + + +@pytest.mark.parametrize("explicit_pointer", [False, True], ids=["bufferload", "ptr-offset"]) +def test_fp4_nonzero_logical_start_raw_bytes(npu, explicit_pointer): + kernel, source = _compile(_fp4_view_kernel(explicit_pointer)) + if explicit_pointer: + # Logical extra offset 64 is packed once to 32 bytes. The pointer + # variable name is not assumed: resolve the actual definition/use. + import re + + match = re.search(r"v = pto\.vmi\.vload\((\w+), 32, size=\(128 // 2\)\)", source) + assert match is not None, source + pointer_name = match.group(1) + assert re.search(rf"\b{pointer_name} = pto\.(?:addptr|castptr)\(", source), source + else: + assert "addptr" in source and "), 32)" in source, source + + def make_case(trial): + generator = torch.Generator().manual_seed(2000 + trial) + byte_values = torch.randint(0, 256, (128,), generator=generator, dtype=torch.int64).to(torch.uint8) + zeros = torch.zeros(128, dtype=torch.uint8) + expected = zeros.clone() + expected[64:128] = byte_values[32:96] + return [byte_values, zeros], expected + + _run_trials(kernel, make_case, npu, trials=3) + + +# --------------------------------------------------------------------------- +# Packed (PK) store: logical target extent and exact float16 reference +# --------------------------------------------------------------------------- +def _pk_b32_kernel(): + @T.prim_func + def main(A: T.Tensor((128,), F32), S: T.Tensor((128,), F16), C: T.Tensor((128,), F16)): + with T.Kernel(1): + src = T.alloc_shared((128,), F32) + out = T.alloc_shared((192,), F16) + T.copy(A, src) + T.copy(S, out[64:192]) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + v = T.vmi.vload(src[0], size=128) + narrowed = T.vmi.vcvt(v, to_dtype=F16, rounding="R", saturate="SAT") + T.vmi.vstore(narrowed, out[64], mask) + T.copy(out[64:192], C) + + return main + + +def test_pk_b32_f16_store(npu): + kernel, source = _compile(_pk_b32_kernel()) + vpto = _vpto_text(source) + assert vpto.count('dist = "PK_B32"') == 2, vpto + + def make_case(trial): + a = ((torch.arange(128, dtype=torch.float32) - 64.0) * 0.25) + trial + expected = a.to(torch.float16) + sentinel = torch.full((128,), -3.0, dtype=torch.float16) + return [a, sentinel], expected + + _run_trials(kernel, make_case, npu, trials=3) + + +# --------------------------------------------------------------------------- +# Explicit pointer +# --------------------------------------------------------------------------- +def _explicit_pointer_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((64,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + ptr = T.access_ptr(work[32], "r", extent=128) + v = T.vmi.vload(ptr, 16, size=64) + T.vmi.vstore(v, work[128], mask) + T.copy(work[128:192], C) + + return main + + +def test_explicit_pointer_base32_extra16_reads_48(npu): + kernel, source = _compile(_explicit_pointer_kernel()) + _assert_addptr_pointer_use(source, ", 32)", "vmi.vload({name}, 16, size=64)") + + def make_case(trial): + a = _float_inputs((256,), trial) + return [a], a[48:112] + + _run_trials(kernel, make_case, npu, trials=5) + + +# --------------------------------------------------------------------------- +# Gather / scatter with non-zero direct and explicit-pointer bases +# --------------------------------------------------------------------------- +def _gather_kernel(explicit_pointer): + @T.prim_func + def main(A: T.Tensor((256,), F32), I: T.Tensor((64,), "uint32"), C: T.Tensor((64,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + idx = T.alloc_shared((64,), "uint32") + out = T.alloc_shared((64,), F32) + T.copy(A, work) + T.copy(I, idx) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + offsets = T.vmi.vload(idx[0], size=64) + if explicit_pointer: + ptr = T.access_ptr(work[64], "r", extent=128) + v = T.vmi.vgather(ptr, offsets, mask) + else: + v = T.vmi.vgather(work[64], offsets, mask) + T.vmi.vstore(v, out[0], mask) + T.copy(out, C) + + return main + + +@pytest.mark.parametrize("explicit_pointer", [False, True], ids=["direct", "ptr"]) +def test_gather_nonzero_base(npu, explicit_pointer): + kernel, _ = _compile(_gather_kernel(explicit_pointer)) + + def make_case(trial): + a = _float_inputs((256,), trial) + generator = torch.Generator().manual_seed(3000 + trial) + offsets = torch.randperm(128, generator=generator)[:64].to(torch.uint32) + expected = a[64 + offsets.to(torch.int64)] + return [a, offsets], expected + + _run_trials(kernel, make_case, npu, trials=3) + + +def _scatter_kernel(explicit_pointer): + @T.prim_func + def main(A: T.Tensor((64,), F32), S: T.Tensor((256,), F32), I: T.Tensor((64,), "uint32"), C: T.Tensor((256,), F32)): + with T.Kernel(1): + src = T.alloc_shared((64,), F32) + work = T.alloc_shared((256,), F32) + idx = T.alloc_shared((64,), "uint32") + T.copy(A, src) + T.copy(S, work) + T.copy(I, idx) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + offsets = T.vmi.vload(idx[0], size=64) + v = T.vmi.vload(src[0], size=64) + if explicit_pointer: + ptr = T.access_ptr(work[64], "w", extent=128) + T.vmi.vscatter(v, ptr, offsets, mask) + else: + T.vmi.vscatter(v, work[64], offsets, mask) + T.copy(work, C) + + return main + + +@pytest.mark.parametrize("explicit_pointer", [False, True], ids=["direct", "ptr"]) +def test_scatter_nonzero_base(npu, explicit_pointer): + kernel, _ = _compile(_scatter_kernel(explicit_pointer)) + + def make_case(trial): + a = _float_inputs((64,), trial) + sentinel = torch.full((256,), -123.0 - trial) + generator = torch.Generator().manual_seed(4000 + trial) + offsets = torch.randperm(128, generator=generator)[:64].to(torch.uint32) + expected = sentinel.clone() + expected[64 + offsets.to(torch.int64)] = a + return [a, sentinel, offsets], expected + + _run_trials(kernel, make_case, npu, trials=3) diff --git a/tilelang/ascend/language/vmi.py b/tilelang/ascend/language/vmi.py index 506482f793..09337cb72d 100644 --- a/tilelang/ascend/language/vmi.py +++ b/tilelang/ascend/language/vmi.py @@ -169,9 +169,34 @@ def _require_compatible_mask(mask, lanes: int, *, context: str): return mask +def _access_ptr_base_load(value): + """Return the immediate BufferLoad behind a frontend ``tl.access_ptr`` call. + + This is a local type recovery for an explicit pointer expression, not + general pointer provenance analysis: only the call's own base argument + (through a direct frontend ``Bind`` alias) is inspected. + """ + if isinstance(value, tirx.Var): + from tilelang.language.frame import has_let_value, get_let_value + + if has_let_value(value): + bound = get_let_value(value) + if bound is not None and not isinstance(bound, tirx.Var): + return _access_ptr_base_load(bound) + return None + if isinstance(value, tirx.Call): + op = value.op + if getattr(op, "name", None) == "tl.access_ptr" and value.args and isinstance(value.args[0], BufferLoad): + return value.args[0] + return None + + def _require_ub_address(value, *, context: str): buffer = None - if isinstance(value, (BufferLoad, BufferRegion)): + base_load = _access_ptr_base_load(value) + if base_load is not None: + buffer = base_load.buffer + elif isinstance(value, (BufferLoad, BufferRegion)): buffer = value.buffer elif isinstance(value, Buffer): buffer = value @@ -185,6 +210,9 @@ def _require_ub_address(value, *, context: str): def _address_element_dtype(value): + base_load = _access_ptr_base_load(value) + if base_load is not None: + return _dtype(base_load.buffer.dtype) if isinstance(value, BufferLoad): return _dtype(value.buffer.dtype) if isinstance(value, BufferRegion): @@ -270,7 +298,6 @@ def _validate_pto_load_modes( group, stride, block_stride, - repeat_stride, allow_group_brc: bool, allowed_dist_modes, ): @@ -281,17 +308,15 @@ def _validate_pto_load_modes( if group is not None: if dist_mode is not None and (not allow_group_brc or dist_mode != "brc"): raise TypeError(f"{context} does not allow dist_mode together with group") - if block_stride is not None or repeat_stride is not None: + if block_stride is not None: raise TypeError(f"{context} does not allow block_stride together with group") if stride is None: raise TypeError(f"{context} with group=... requires stride") return - if block_stride is not None or repeat_stride is not None: + if block_stride is not None: if dist_mode is not None: raise TypeError(f"{context} does not allow dist_mode together with block_stride") - if block_stride is None or repeat_stride is None: - raise TypeError(f"{context} requires block_stride and repeat_stride together") if stride is not None: raise TypeError(f"{context} does not allow stride together with block_stride") return @@ -313,32 +338,211 @@ def _looks_like_mask(value) -> bool: return False -def _resolve_ptr_and_offset(source, offset=None, *, access_type: str, extent=None): +def _is_unit_stride_1d(buffer) -> bool: + if len(buffer.shape) != 1: + return False + if buffer.strides: + return len(buffer.strides) == 1 and _constant_int(buffer.strides[0]) == 1 + return True + + +def _tight_buffer_span(buffer): + """Return ``prod(shape)`` when *buffer* is tightly packed, else ``None``. + + A whole-buffer descriptor is only sound when the logical buffer covers one + contiguous physical span. Explicit strides must equal the row-major + strides; default (empty) strides are row-major by construction. + """ + ndim = len(buffer.shape) + if ndim == 0: + return None + if buffer.strides: + if len(buffer.strides) != ndim: + return None + from tvm import arith + + analyzer = arith.Analyzer() + expected = 1 + for dim in range(ndim - 1, -1, -1): + if not analyzer.can_prove_equal(buffer.strides[dim], expected): + return None + expected = expected * buffer.shape[dim] + span = 1 + for dim in buffer.shape: + span = span * dim + return span + + +def _source_relative_offset(source): + """Element offset of *source* relative to its buffer's logical origin. + + ``Buffer.offset_of`` already folds strides and ``elem_offset``; the origin + offset (which is exactly ``elem_offset``) is subtracted so the result can + be added to an ``access_ptr(buffer, ...)`` whose lowering adds + ``elem_offset`` once. + """ + if isinstance(source, BufferLoad): + buffer = source.buffer + indices = list(source.indices) + elif isinstance(source, BufferRegion): + buffer = source.buffer + indices = [region.min for region in source.region] + elif isinstance(source, Buffer): + return _zero_offset() + else: + return None + if not buffer.shape or len(indices) != len(buffer.shape): + return None + source_offsets = buffer.offset_of(indices) + origin_offsets = buffer.offset_of([0] * len(buffer.shape)) + # A single flattened index is the only form whose linear address this + # frontend contract can relate to one descriptor. Multi-address / + # axis-separator buffers are not proven linear and keep the existing + # boundary instead of being silently reduced to offsets[0]. + if len(source_offsets) != 1 or len(origin_offsets) != 1: + return None + return source_offsets[0] - origin_offsets[0] + + +def _whole_buffer_plan(source, *, fallback_extent): + """Declare the whole known tight buffer and keep the relative origin offset. + + Used when a group/block stride is dynamic (or a static negative stride has + no directly constructible low address). When the buffer is not provably + tight (view/padding/strides), this returns the caller's fallback extent + beginning at the source origin: that is only the pre-existing boundary and + is not claimed as a complete conservative footprint. + """ + buffer = source if isinstance(source, Buffer) else source.buffer + span = _tight_buffer_span(buffer) + relative = _source_relative_offset(source) + if span is None or relative is None: + return fallback_extent, None, None + return span, buffer, relative + + +def _group_access_plan(source, *, lanes, group, stride, dist_mode, is_load, elem_dtype): + """May-access span of an equal-distance group access (C.2/C.3.1). + + The first group starts at the original access address ``A``. For stride + ``S`` and physical group width ``W`` the contiguous envelope is + ``[A + min(0,(G-1)*S), A + max(0,(G-1)*S) + W)``. A negative static + stride on a 1-D unit-stride buffer relocates the descriptor to the lowest + address and compensates with the VMI instruction offset. + + A group load is unmasked and therefore reads whole physical vectors: when + the logical per-group width ``N/G`` is not a multiple of the physical VL, + the last chunk still reads a full VL (e.g. f32 G=1/N=80 lowers to two + full vlds covering 128 address elements). A group store keeps the logical + width because the backend's generated tail mask bounds the writes; do not + copy the load rounding onto stores. + """ + group_const = _constant_int(group) + if group_const is None or group_const <= 0: + return _whole_buffer_plan(source, fallback_extent=lanes) + # Slot (one result/input element per group) and group broadcast read or + # write exactly one address element per group. Ordinary groups cover + # ceil(N/G) elements; the ceiling keeps a non-divisible shape conservative. + if dist_mode == "brc" or lanes <= group_const: + per_group = 1 + else: + per_group = (lanes + group_const - 1) // group_const + bits = int(getattr(elem_dtype, "bits", 0) or 0) + if is_load and bits > 0: + vector_elements = max(1, 2048 // bits) + per_group = (per_group + vector_elements - 1) // vector_elements * vector_elements + stride_const = _constant_int(stride) + if stride_const is None: + return _whole_buffer_plan(source, fallback_extent=lanes) + span = (group_const - 1) * stride_const + low = min(0, span) + high = max(0, span) + per_group + extent = high - low + if low == 0: + return extent, None, None + if isinstance(source, BufferLoad) and len(source.indices) == 1 and _is_unit_stride_1d(source.buffer): + shifted = tirx.BufferLoad(source.buffer, [source.indices[0] + low]) + return extent, shifted, -low + return _whole_buffer_plan(source, fallback_extent=lanes) + + +def _block_access_plan(source, *, lanes, elem_dtype, block_stride): + """May-access span of a block-stride access (C.3.2). + + ``block_stride`` is measured in 32-byte blocks. A full physical vector + covers eight blocks; physical chunks beyond one vector start at the next + block group, and the tail chunk of a masked store conservatively covers + its whole block group as well. + """ + bits = int(getattr(elem_dtype, "bits", 0) or 0) + if bits <= 0: + return lanes, None, None + elements_per_block = max(1, 256 // bits) + vector_elements = max(1, 2048 // bits) + chunks = (lanes + vector_elements - 1) // vector_elements + blocks = chunks * max(1, vector_elements // elements_per_block) + block_const = _constant_int(block_stride) + if block_const is None or block_const < 0: + # A dynamic control word (and the i16 control-word encoding) has no + # statically provable span; declare the whole known tight buffer. + return _whole_buffer_plan(source, fallback_extent=lanes) + extent = (blocks - 1) * block_const * elements_per_block + elements_per_block + return extent, None, None + + +def _vmi_access_plan(source, *, lanes, elem_dtype, dist_mode, group, stride, block_stride, is_load): + """Compute (descriptor extent, descriptor base, instruction offset). + + ``descriptor base`` is ``None`` when the original source is already the + right ``access_ptr`` base. The instruction offset is ``None`` for zero. + Only buffer sources are analysed; explicit pointers keep their caller + declared ``ptr + offset`` ABI untouched. + """ + if not isinstance(source, (BufferLoad, Buffer, BufferRegion)): + return None, None, None + if dist_mode == "dintlv": + # Two interleaved inputs/outputs: the complete 2*N span, regardless of + # the (ineffective) store mask. + return 2 * lanes, None, None + if group is not None: + return _group_access_plan( + source, + lanes=lanes, + group=group, + stride=stride, + dist_mode=dist_mode, + is_load=is_load, + elem_dtype=elem_dtype, + ) + if block_stride is not None: + return _block_access_plan(source, lanes=lanes, elem_dtype=elem_dtype, block_stride=block_stride) + if str(elem_dtype) == "float4_e2m1fn": + # Genuine FP4 descriptors stay in logical FP4 element units; + # RewriteFp4ToFp4x2 performs the single physical packing. + return lanes, None, None + if dist_mode == "brc": + # An ordinary broadcast reads one source element; keeping the old + # conservative `size` declaration is allowed and does not change the + # instruction semantics. + return lanes, None, None + if is_load: + bits = int(getattr(elem_dtype, "bits", 0) or 0) + if bits <= 0: + return lanes, None, None + vector_elements = max(1, 2048 // bits) + # An unmasked load reads whole physical vectors: cover the tail chunk. + return (lanes + vector_elements - 1) // vector_elements * vector_elements, None, None + return lanes, None, None + + +def _resolve_ptr_and_offset(source, offset=None, *, access_type: str, extent=None, descriptor=None, instruction_offset=None): """Normalize VMI address syntax into PTODSL's explicit ptr+offset ABI.""" if isinstance(source, (BufferLoad, Buffer, BufferRegion)) and offset is not None: raise TypeError("T.vmi buffer addresses do not accept an offset; express the offset in the buffer index") - if isinstance(source, BufferLoad): - # Keep the linear element offset outside access_ptr. In particular, - # FP4 callers convert this logical offset to f4e2m1x2 units before - # PTO emission; leaving it embedded in access_ptr would make PTO - # interpret it as an already-packed offset. - strides = source.buffer.strides - if len(strides) != len(source.indices): - strides = [] - for index in range(len(source.indices)): - stride = 1 - for extent_dim in source.buffer.shape[index + 1 :]: - stride = stride * extent_dim - strides.append(stride) - linear_offset = None - for index, stride in zip(source.indices, strides): - term = index * stride - linear_offset = term if linear_offset is None else linear_offset + term - ptr = access_ptr(source.buffer, access_type, extent=extent, offset=0) - return ptr, linear_offset if linear_offset is not None else _zero_offset() - if isinstance(source, (Buffer, BufferRegion)): - ptr = access_ptr(source, access_type, extent=extent, offset=0) - return ptr, _zero_offset() + base = descriptor if descriptor is not None else source + if isinstance(base, (BufferLoad, Buffer, BufferRegion)): + ptr = access_ptr(base, access_type, extent=extent, offset=0) + return ptr, instruction_offset if instruction_offset is not None else _zero_offset() if offset is None: return source, _zero_offset() return source, offset @@ -442,10 +646,8 @@ def vload( offset=None, *, size, - to_dtype=None, stride=None, block_stride=None, - repeat_stride=None, dist_mode=None, group=None, loc=None, @@ -467,51 +669,46 @@ def vload( group=group, stride=stride, block_stride=block_stride, - repeat_stride=repeat_stride, allow_group_brc=True, - allowed_dist_modes={None, "continuous", "dintlv", "unpack", "brc"}, + allowed_dist_modes={None, "continuous", "dintlv", "brc"}, ) if size is None: raise TypeError("T.vmi.vload(...) requires size") - if to_dtype is not None and dist_mode != "unpack": - raise TypeError('T.vmi.vload(...) accepts to_dtype only when dist_mode="unpack"') source_elem = _require_address_element_dtype(source, context="T.vmi.vload(...)") - # Keep si*/ui* spelling for PTODSL annotations. Unpack vload itself is not - # legalized on the current VPTO path; this preserves a correct annotation - # if enabled later. - to_dtype_annot = None is_packed_fp4 = str(source_elem) == "float4_e2m1fn" if is_packed_fp4: if size % 2: raise ValueError("T.vmi.vload(...) requires an even FP4 lane count for packed FP4 storage") - if to_dtype is not None or dist_mode == "unpack": - raise TypeError("T.vmi.vload(...) does not unpack or convert packed FP4 source vectors") if dist_mode == "dintlv": raise TypeError('T.vmi.vload(...) does not support packed FP4 with dist_mode="dintlv"') - if group is not None or block_stride is not None or repeat_stride is not None or dist_mode == "brc": + if group is not None or block_stride is not None or dist_mode == "brc": raise TypeError("T.vmi.vload(...) only supports contiguous packed FP4 loads") - if dist_mode == "unpack": - if to_dtype is None: - raise TypeError('T.vmi.vload(...) requires to_dtype when dist_mode="unpack"') - to_dtype_annot = _pto_to_dtype_annotation(to_dtype) - to_dtype = _scalar_dtype(to_dtype, context="T.vmi.vload(..., to_dtype=...)") - if getattr(to_dtype, "bits", None) != 2 * getattr(source_elem, "bits", None): - raise TypeError("T.vmi.vload(...) unpack must widen by exactly one step") - - result_dtype = _vector_result_dtype( + + result_dtype = _vector_result_dtype(source, size, elem_dtype=source_elem) + extent, descriptor, instruction_offset = _vmi_access_plan( + source, + lanes=size, + elem_dtype=source_elem, + dist_mode=dist_mode, + group=group, + stride=stride, + block_stride=block_stride, + is_load=True, + ) + ptr, extra_offset = _resolve_ptr_and_offset( source, - size, - elem_dtype=to_dtype if to_dtype is not None else source_elem, + offset, + access_type="r", + extent=extent, + descriptor=descriptor, + instruction_offset=instruction_offset, ) - ptr, extra_offset = _resolve_ptr_and_offset(source, offset, access_type="r", extent=size) if is_packed_fp4: extra_offset = _pack_fp4_element_offset(extra_offset, context="T.vmi.vload(...)") attrs = { "size": size, - "to_dtype": to_dtype_annot, "stride": stride, "block_stride": block_stride, - "repeat_stride": repeat_stride, "dist_mode": dist_mode, "group": group, "loc": loc, @@ -532,10 +729,8 @@ def vstore( *, stride=None, block_stride=None, - repeat_stride=None, dist_mode=None, group=None, - pmode=None, loc=None, ip=None, ): @@ -549,7 +744,6 @@ def vstore( group=group, stride=stride, block_stride=block_stride, - repeat_stride=repeat_stride, allow_group_brc=False, allowed_dist_modes={None, "continuous", "dintlv"}, ) @@ -567,32 +761,51 @@ def vstore( elif isinstance(values, VmiPair) or _is_sequence(values): raise TypeError('T.vmi.vstore(...) expects a single VMI vector unless dist_mode="dintlv"') - store_extent = _lanes_of(values[0]) if _is_sequence(values) else _lanes_of(values) - ptr, extra_offset = _resolve_ptr_and_offset(destination, offset, access_type="w", extent=store_extent) + store_lanes = _lanes_of(values[0]) if _is_sequence(values) else _lanes_of(values) value_args = list(values) if _is_sequence(values) else [values] if value_dtype is None: value_dtype = _element_dtype_of(values[0]) if _is_sequence(values) else _element_dtype_of(values) if str(value_dtype) == "float4_e2m1fn": - if store_extent % 2: + if store_lanes % 2: raise ValueError("T.vmi.vstore(...) requires an even FP4 lane count for packed FP4 storage") if dist_mode == "dintlv": raise TypeError('T.vmi.vstore(...) does not support packed FP4 with dist_mode="dintlv"') - if group is not None or block_stride is not None or repeat_stride is not None: + if group is not None or block_stride is not None: raise TypeError("T.vmi.vstore(...) only supports contiguous packed FP4 stores") + elif mask is not None: + mask = _require_mask(mask, context="T.vmi.vstore(...)") + + extent, descriptor, instruction_offset = _vmi_access_plan( + destination, + lanes=store_lanes, + elem_dtype=value_dtype, + dist_mode=dist_mode, + group=group, + stride=stride, + block_stride=block_stride, + is_load=False, + ) + ptr, extra_offset = _resolve_ptr_and_offset( + destination, + offset, + access_type="w", + extent=extent, + descriptor=descriptor, + instruction_offset=instruction_offset, + ) + if str(value_dtype) == "float4_e2m1fn": extra_offset = _pack_fp4_element_offset(extra_offset, context="T.vmi.vstore(...)") if mask is not None: mask = _require_mask(mask, context="T.vmi.vstore(...) of FP4") if not str(_dtype_of(mask)).startswith("bool"): raise TypeError("T.vmi.vstore(...) of FP4 expects a VMI mask operand") - physical_lanes = store_extent // 2 + physical_lanes = store_lanes // 2 if _lanes_of(mask) != physical_lanes: raise TypeError( "T.vmi.vstore(...) of FP4 requires a physical FP4x2 mask with " - f"half as many lanes as its {store_extent}-lane logical FP4 value " + f"half as many lanes as its {store_lanes}-lane logical FP4 value " f"(expected {physical_lanes}, got {_lanes_of(mask)})" ) - elif mask is not None: - mask = _require_mask(mask, context="T.vmi.vstore(...)") return _call_vmi( "vstore", "void", @@ -602,10 +815,8 @@ def vstore( mask, stride=stride, block_stride=block_stride, - repeat_stride=repeat_stride, dist_mode=dist_mode, group=group, - pmode=pmode, loc=loc, ip=ip, ) @@ -832,12 +1043,14 @@ def vcadd(source, mask, *, group=None, pmode=None, reassoc=None, loc=None, ip=No context = "T.vmi.vcadd(...)" _reject_packed_fp4(source, context=context) elem_dtype = _element_dtype_of(source) - if reassoc is None and str(elem_dtype).startswith(("float", "bfloat")): - raise TypeError( - f"{context} on floating-point vectors requires an explicit reassoc argument; spell out reassoc=True or reassoc=False" - ) if reassoc is not None and not isinstance(reassoc, bool): raise TypeError(f"{context} requires reassoc to be the Python boolean True or False; received {reassoc!r}") + if reassoc is False: + raise TypeError( + f"{context} does not support reassoc=False: the PTO backend accepts only reassoc=True (integer vectors may omit reassoc)" + ) + if reassoc is None and str(elem_dtype).startswith(("float", "bfloat")): + raise TypeError(f"{context} on floating-point vectors requires an explicit reassoc=True argument; the PTO backend rejects omission") lanes = 1 if group is None else _normalize_lanes(group, context="T.vmi.vcadd(..., group=...)") if group is not None and _lanes_of(source) % group != 0: raise ValueError(f"{context} requires source lanes to be divisible by group") From b966529ff06cf3d4ddd053ec92970e3355e7fe5a Mon Sep 17 00:00:00 2001 From: qukelin Date: Sat, 17 Oct 2026 11:29:08 +0800 Subject: [PATCH 05/10] fix(pto): use signless backing for SIMT local integers --- src/ascend/codegen/codegen_pto.cc | 78 ++++-- .../pto/test_simtvf_local_integer_storage.py | 233 +++++++++++++++++ .../test_simtvf_local_integer_storage_run.py | 236 ++++++++++++++++++ 3 files changed, 529 insertions(+), 18 deletions(-) create mode 100644 testing/ascend/language/pto/test_simtvf_local_integer_storage.py create mode 100644 testing/ascend/language/pto/test_simtvf_local_integer_storage_run.py diff --git a/src/ascend/codegen/codegen_pto.cc b/src/ascend/codegen/codegen_pto.cc index 2d3e832e84..a447913b5d 100644 --- a/src/ascend/codegen/codegen_pto.cc +++ b/src/ascend/codegen/codegen_pto.cc @@ -470,11 +470,44 @@ bool IsSupportedSIMTLocalStorageType(DataType t) { return IsSimtLocalBufferDtype(t) || (t.is_scalar() && t == DataType::Int(32)); } +bool IsSimtLocalIntegerStorageDtype(DataType t) { + return t.is_scalar() && (t.is_int() || t.is_uint()) && + (t.bits() == 8 || t.bits() == 32) && + IsSupportedSIMTLocalStorageType(t); +} + +std::string SimtLocalIntegerStorageTypeName(DataType t) { + ICHECK(IsSimtLocalIntegerStorageDtype(t)); + return "pto.i" + std::to_string(t.bits()); +} + +// SIMT local integer storage is signless because LLVM stack accesses require +// signless integer element types. TileLang's logical signedness is restored +// at the load boundary and folded back to the backing type at the store +// boundary. These adapters are only consumed by the SIMT local branches of +// ScalarLoad/EmitScalarStore/EmitScalarizedLoad/EmitScalarizedStore. +std::string SimtLocalLoadAdapter(const std::string &value, DataType t) { + if (!IsSimtLocalIntegerStorageDtype(t)) { + return value; + } + return "pto.cast(" + value + ", " + DataTypeName(t) + ")"; +} + +std::string SimtLocalStoreAdapter(const std::string &value, DataType t) { + if (!IsSimtLocalIntegerStorageDtype(t)) { + return value; + } + return "pto.cast(" + value + ", " + SimtLocalIntegerStorageTypeName(t) + ")"; +} + std::string SIMTLocalStorageTypeName(DataType t) { ICHECK(IsSupportedSIMTLocalStorageType(t)); // LLVM stack allocations require signless integer element types. Preserve // TileLang's signedness at load/store boundaries instead. - return IsInteger32(t) ? "pto.i32" : DataTypeName(t); + if (IsSimtLocalIntegerStorageDtype(t)) { + return SimtLocalIntegerStorageTypeName(t); + } + return DataTypeName(t); } bool IsSupportedSIMTFP8ContiguousLaneCount(int lanes) { @@ -2661,10 +2694,7 @@ std::string CodeGenTileLangPTO::ScalarLoad(const BufferNode *buffer, return "tl.fp8_byte_load(" + GetVarID(buffer->data.get()) + ", " + index_str + ", " + FP8TypeName(buffer->dtype) + ")"; } - if (IsInteger32(buffer->dtype)) { - return "pto.cast(" + value + ", " + DataTypeName(buffer->dtype) + ")"; - } - return value; + return SimtLocalLoadAdapter(value, buffer->dtype); } if (scope == "shared" || scope == "shared.dyn" || scope == "global" || @@ -2729,9 +2759,7 @@ void CodeGenTileLangPTO::EmitScalarStore(const BufferNode *buffer, << FP8TypeName(buffer->dtype) << ")\n"; return; } - if (IsInteger32(buffer->dtype)) { - store_value = "pto.cast(" + store_value + ", pto.i32)"; - } + store_value = SimtLocalStoreAdapter(store_value, buffer->dtype); stream << "pto.store(" << store_value << ", " << GetVarID(buffer->data.get()) << ", " << index_str << ")\n"; return; @@ -2845,13 +2873,25 @@ void CodeGenTileLangPTO::EmitBufferAllocation(const Buffer &buffer) { "float32, int8, uint8, int32, uint32, and " "float8_e4m3fn/float8_e5m2, got " << buffer->dtype; + // Both allocation sub-branches use the same signless integer element + // backing for the admitted SIMT local integer types: scalar-element + // buffers through SIMTLocalStorageTypeName, vector-element register + // slots through the element type (without changing admission or + // claiming a supported slot read/write path). + auto storage_element_type = [&]() -> std::string { + if (!buffer->dtype.is_vector()) { + return SIMTLocalStorageTypeName(buffer->dtype); + } + if (tl::IsAscendVectorizableFP8(buffer->dtype)) { + return DataTypeName(buffer->dtype); + } + if (IsSimtLocalIntegerStorageDtype(elem)) { + return SimtLocalIntegerStorageTypeName(elem); + } + return ScalarType(elem); + }; stream << vid << " = pto.alloc_buffer((" << opt_size.value() << ",), " - << (buffer->dtype.is_vector() - ? (tl::IsAscendVectorizableFP8(buffer->dtype) - ? DataTypeName(buffer->dtype) - : ScalarType(elem)) - : SIMTLocalStorageTypeName(buffer->dtype)) - << ")\n"; + << storage_element_type() << ")\n"; } else if (persistent) { ICHECK(IsFloat32(buffer->dtype)) << "PTO persistent local allocation currently supports float32 " @@ -7892,8 +7932,10 @@ void CodeGenTileLangPTO::EmitScalarizedLoad(const BufferLoadNode *op, index_str = RemoveOutermostParentheses(PrintExpr_(ramp->base)); } if (scope == "local.fragment" || scope == "local") { - os << "pto.load(" << GetVarID(op->buffer->data.get()) << ", " << index_str - << ", contiguous=" << lanes << ")"; + std::string value = "pto.load(" + GetVarID(op->buffer->data.get()) + + ", " + index_str + + ", contiguous=" + std::to_string(lanes) + ")"; + os << SimtLocalLoadAdapter(value, element_dtype); return; } if (scope == "shared" || scope == "shared.dyn") { @@ -8103,8 +8145,8 @@ void CodeGenTileLangPTO::EmitScalarizedStore(const BufferStoreNode *op) { const std::string contiguous_suffix = store_lanes > 1 ? ", contiguous=" + std::to_string(store_lanes) : ""; if (scope == "local.fragment" || scope == "local") { - stream << "pto.store(" << value << ", " - << GetVarID(op->buffer->data.get()) << ", " << index_str + stream << "pto.store(" << SimtLocalStoreAdapter(value, buffer_dtype) + << ", " << GetVarID(op->buffer->data.get()) << ", " << index_str << contiguous_suffix << ")\n"; return; } diff --git a/testing/ascend/language/pto/test_simtvf_local_integer_storage.py b/testing/ascend/language/pto/test_simtvf_local_integer_storage.py new file mode 100644 index 0000000000..b26534fb01 --- /dev/null +++ b/testing/ascend/language/pto/test_simtvf_local_integer_storage.py @@ -0,0 +1,233 @@ +"""SIMT local integer buffers use signless LLVM stack storage. + +This file verifies that: + +* scalar-element SIMT local buffers of the admitted integer types + int8/uint8/int32/uint32 allocate with a same-width *signless* PTO element + type (``pto.i8`` / ``pto.i32``); +* the four access sites -- scalar load, scalar store, contiguous + builtin-vector load and contiguous builtin-vector store -- adapt between + the signless backing and TileLang's logical signed/unsigned type via + ``pto.cast``; +* the adapter is scoped to SIMT local storage only: VMI, ``local.var``, + shared and global accesses keep their existing types; +* vector-element local slots use the same-width signless element backing in + the allocation branch; +* the admitted dtype set is unchanged: int16/int64 stay rejected. +""" + +import re + +import pytest +import torch # noqa: F401 (torch/torch_npu before tilelang/ptodsl) +import torch_npu # noqa: F401 + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend.target import normalize_pto_target + + +def _lower(func): + """Lower with an explicit PTO target context (required for fragments).""" + target = normalize_pto_target("pto") + with target: + return tilelang.lower(func, target=target) + + +# --------------------------------------------------------------------------- +# SIMT local integer kernel: scalar and contiguous access sites +# --------------------------------------------------------------------------- +def _simt_local_integer_kernel(dtype): + @T.prim_func + def main(A: T.Tensor((8,), dtype), B: T.Tensor((8,), dtype)): + with T.Kernel(1), T.SimtVF(threads=1): + local = T.alloc_local((8,), dtype) + local[0] = A[0] + B[0] = local[0] + local[T.Ramp(1, 1, 4)] = A[T.Ramp(0, 1, 4)] + B[T.Ramp(1, 1, 4)] = local[T.Ramp(1, 1, 4)] + + return main + + +@pytest.mark.parametrize( + ("dtype", "logical_type", "backing_type"), + [ + ("int8", "pto.si8", "pto.i8"), + ("uint8", "pto.ui8", "pto.i8"), + ("int32", "pto.si32", "pto.i32"), + ("uint32", "pto.ui32", "pto.i32"), + ], +) +def test_simtvf_local_integer_scalar_and_contiguous_accesses_use_signless_storage(dtype, logical_type, backing_type): + source = _lower(_simt_local_integer_kernel(dtype)).kernel_source + + # Allocation: same-width signless backing. + assert re.search(rf"pto\.alloc_buffer\(\(8,\), {re.escape(backing_type)}\)", source) + assert "pto.alloc_buffer((8,), pto.si8)" not in source + assert "pto.alloc_buffer((8,), pto.ui8)" not in source + assert "pto.alloc_buffer((8,), pto.si32)" not in source + assert "pto.alloc_buffer((8,), pto.ui32)" not in source + + # Scalar load restores the logical type; scalar store returns to backing. + assert re.search(rf"pto\.cast\(pto\.load\(local, 0\), {re.escape(logical_type)}\)", source) + assert re.search( + rf"pto\.store\(pto\.cast\(pto\.load\(A, 0\), {re.escape(backing_type)}\), local, 0\)", + source, + ) + + # Contiguous builtin-vector load restores; contiguous store returns. + assert re.search( + rf"pto\.cast\(pto\.load\(local, 1, contiguous=4\), {re.escape(logical_type)}\)", + source, + ) + assert re.search( + rf"pto\.store\(pto\.cast\(pto\.load\(A, 0, contiguous=4\), " + rf"{re.escape(backing_type)}\), local, 1, contiguous=4\)", + source, + ) + + # Live allocation survives optimization: one allocation and the local + # buffer is actually read/written by all four access sites. + assert source.count("pto.alloc_buffer(") == 1 + assert len(re.findall(r"pto\.load\(local,", source)) == 2 + assert len(re.findall(r", local,", source)) == 2 + + +@pytest.mark.parametrize("dtype", ["int8", "uint8", "int32", "uint32"]) +def test_simtvf_local_integer_kernel_compiles_through_backend(dtype): + """All four admitted dtypes compile through the full PTO backend.""" + target = normalize_pto_target("pto") + with target: + kernel = tilelang.compile(_simt_local_integer_kernel(dtype), target="pto", out_idx=-1) + source = kernel.get_kernel_source() + backing = "pto.i8" if dtype in ("int8", "uint8") else "pto.i32" + assert f"pto.alloc_buffer((8,), {backing})" in source + + +# --------------------------------------------------------------------------- +# Vector-element local slot allocation +# --------------------------------------------------------------------------- +def _simt_local_integer_vector_kernel(alloc=T.alloc_local): + @T.prim_func + def main(A: T.Tensor((4,), "int8"), B: T.Tensor((4,), "int8")): + with T.Kernel(1), T.SimtVF(threads=1): + local = alloc((1,), "int8x4") + local[0] = T.int8x4(A[0], A[1], A[2], A[3]) + B[T.Ramp(0, 1, 4)] = local[0] + + return main + + +@pytest.mark.parametrize( + ("scope_name", "alloc"), + [("local", T.alloc_local), ("local.fragment", T.alloc_fragment)], + ids=["local", "local.fragment"], +) +def test_simtvf_vector_integer_local_allocation_is_signless(scope_name, alloc): + """Vector-element local slots allocate with the same-width signless i8 backing.""" + source = _lower(_simt_local_integer_vector_kernel(alloc)).kernel_source + + assert "pto.alloc_buffer((1,), pto.i8)" in source + assert "pto.alloc_buffer((1,), pto.si8)" not in source + + +# --------------------------------------------------------------------------- +# Admission boundaries stay closed +# --------------------------------------------------------------------------- +def _unsupported_local_kernel(dtype): + @T.prim_func + def main(A: T.Tensor((4,), dtype), B: T.Tensor((4,), dtype)): + with T.Kernel(1), T.SimtVF(threads=1): + local = T.alloc_local((4,), dtype) + local[0] = A[0] + B[0] = local[0] + + return main + + +@pytest.mark.parametrize("dtype", ["int16", "int64"]) +def test_simtvf_local_integer_admission_boundary_is_unchanged(dtype): + with pytest.raises(Exception, match="PTO SIMT local allocation supports"): + _lower(_unsupported_local_kernel(dtype)) + + +# --------------------------------------------------------------------------- +# Scope controls: local.var, shared/global and VMI keep their contracts +# --------------------------------------------------------------------------- +def _local_var_kernel(dtype): + @T.prim_func + def main(A: T.Tensor((8,), dtype), B: T.Tensor((8,), dtype)): + with T.Kernel(1): + src = T.alloc_shared((8,), dtype) + dst = T.alloc_shared((8,), dtype) + T.copy(A, src) + T.copy(A, dst) + with T.SimtVF(threads=1): + addr = T.alloc_var(dtype) + addr = src[3] + dst[3] = addr + T.copy(dst, B) + + return main + + +def test_simtvf_local_var_integer_keeps_logical_type(): + source = _lower(_local_var_kernel("int32")).kernel_source + # local.var is a Python SSA value typed with the logical dtype; it is not + # a signless-backed local buffer. + assert "addr = pto.const(0, dtype=pto.si32)" in source + assert "pto.alloc_buffer((1,), pto.i32)" not in source + + +def _shared_control_kernel(dtype): + @T.prim_func + def main(A: T.Tensor((8,), dtype), B: T.Tensor((8,), dtype)): + with T.Kernel(1): + src = T.alloc_shared((8,), dtype) + dst = T.alloc_shared((8,), dtype) + T.copy(A, src) + T.copy(A, dst) + with T.SimtVF(threads=1): + dst[0] = src[0] + T.copy(dst, B) + + return main + + +@pytest.mark.parametrize("dtype", ["uint8", "int32"]) +def test_simtvf_shared_integer_accesses_do_not_use_local_adapter(dtype): + source = _lower(_shared_control_kernel(dtype)).kernel_source + line = next(line for line in source.splitlines() if "pto.store(pto.load(" in line) + assert "pto.cast(" not in line, line + + +def _vmi_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), "int32"), C: T.Tensor((256,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((256,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + value = T.vmi.vload(work[0], size=64) + T.vmi.vstore(value, work[64], mask) + T.copy(work, C) + + return main + + +def test_vmi_integer_carrier_keeps_signed_type(): + artifact = _lower(_vmi_kernel()) + found = [] + + def visit(node): + if isinstance(node, tvm.tirx.Call) and getattr(node.op, "name", "") == "tl.vmi.vload": + found.append(node) + + for func in artifact.device_mod.functions.values(): + tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + + assert found, "no VMI load in the lowered module" + assert any(str(call.dtype) == "int32x64" for call in found), [str(call.dtype) for call in found] diff --git a/testing/ascend/language/pto/test_simtvf_local_integer_storage_run.py b/testing/ascend/language/pto/test_simtvf_local_integer_storage_run.py new file mode 100644 index 0000000000..ca8d174a09 --- /dev/null +++ b/testing/ascend/language/pto/test_simtvf_local_integer_storage_run.py @@ -0,0 +1,236 @@ +"""Device acceptance: SIMT local integer signless storage and adapters. + +Each case compiles through the real PTO pipeline (PTODSL -> PTOAS VPTO -> +Bisheng) and runs on the NPU. The reference is computed independently on +the CPU with explicit target-width wrapping, so a byte-copy-only implementation +cannot pass: the kernel performs a signed/unsigned comparison and arithmetic +after reloading the value from the local buffer, plus exact boundary values. + +Vector access offsets are 32-byte aligned (b8: 32 lanes, b32: 8 lanes), the +hardware-legal SIMT vector shape. All source lanes read by the kernel are +initialized by the GM->UB DMA and the destination is fully initialized by a +second DMA before the stores. +""" + +import os +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import torch_npu # noqa: E402,F401 (import before tilelang/ptodsl) +import tilelang # noqa: E402 +import tilelang.ascend.language as T # noqa: E402 + +CASES = { + # dtype: (torch dtype, vector lanes, boundary values, signed?, high bit) + "int8": (torch.int8, 32, (127, -1, -128, 0), True, 0), + "uint8": (torch.uint8, 32, (128, 255, 0, 1), False, 128), + "int32": (torch.int32, 8, (2147483647, -1, -2147483648, 0), True, 0), + "uint32": (torch.uint32, 8, (2147483648, 4294967295, 0, 1), False, 2147483648), +} + + +def _require_npu_runtime(): + if not hasattr(torch, "npu") or not torch.npu.is_available(): + pytest.skip("torch.npu is unavailable in the current environment") + try: + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source) + except Exception as err: # noqa: BLE001 + pytest.skip(f"NPU runtime is unavailable in the current environment: {err}") + + +@pytest.fixture(scope="module") +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +def _wrap(value, bits, signed): + value &= (1 << bits) - 1 + if signed and value >= (1 << (bits - 1)): + value -= 1 << bits + return value + + +def _condition(x, dtype, signed, highbit): + """Signedness-sensitive comparison evaluated while tracing the kernel.""" + if signed: + return x < T.cast(0, dtype) + return x >= T.cast(highbit, dtype) + + +def _local_integer_kernel(dtype, lanes, bounds, signed, highbit): + b0, b1, b2, b3 = bounds + scalar_index = lanes + 4 + + @T.prim_func + def main(A: T.Tensor((64,), dtype), C: T.Tensor((64,), dtype)): + with T.Kernel(1): + src = T.alloc_shared((64,), dtype) + dst = T.alloc_shared((64,), dtype) + T.copy(A, src) + T.copy(A, dst) # every destination lane is initialized + with T.SimtVF(threads=1): + local = T.alloc_local((64,), dtype) + # Contiguous builtin-vector store/load landing. + local[T.Ramp(0, 1, lanes)] = src[T.Ramp(0, 1, lanes)] + dst[T.Ramp(0, 1, lanes)] = local[T.Ramp(0, 1, lanes)] + # Reload one vector, load a second, and compute v + w. + local[T.Ramp(lanes, 1, lanes)] = src[T.Ramp(lanes, 1, lanes)] + left = local[T.Ramp(0, 1, lanes)] + right = local[T.Ramp(lanes, 1, lanes)] + dst[T.Ramp(0, 1, lanes)] = left + right + # Scalar store/load landing. The index is outside the + # vector destination range so the vector result stays + # observable. + local[scalar_index] = src[scalar_index] + dst[scalar_index] = local[scalar_index] + # Signed/unsigned comparison and arithmetic after reload. + value = local[scalar_index] + dst[scalar_index + 1] = T.if_then_else( + _condition(value, dtype, signed, highbit), + value + T.cast(1, dtype), + value + T.cast(2, dtype), + ) + # Exact boundary values through scalar landings. + local[40] = T.cast(b0, dtype) + dst[40] = local[40] + local[41] = T.cast(b1, dtype) + dst[41] = local[41] + local[42] = T.cast(b2, dtype) + dst[42] = local[42] + local[43] = T.cast(b3, dtype) + dst[43] = local[43] + T.copy(dst, C) + + return main + + +def _expected(a, lanes, bounds, signed): + out = a.clone() + bits = a.dtype.itemsize * 8 + for lane in range(lanes): + out[lane] = _wrap(int(a[lane]) + int(a[lanes + lane]), bits, signed) + scalar_index = lanes + 4 + out[scalar_index] = a[scalar_index] + x = int(a[scalar_index]) + takes_true = x < 0 if signed else x >= (1 << (bits - 1)) + x = x + 1 if takes_true else x + 2 + out[scalar_index + 1] = _wrap(x, bits, signed) + for index, boundary in enumerate(bounds): + out[40 + index] = _wrap(boundary, bits, signed) + return out + + +def _trial_input(torch_dtype, signed, trial, lanes): + """Fresh input per trial. + + Vector lanes at [0,4) and [lanes,lanes+4) use boundary/high-bit patterns + whose wrapped sums are computed independently by ``_expected``. The scalar + comparison value alternates between the true and false arms across trials. + """ + bits = torch_dtype.itemsize * 8 + generator = torch.Generator().manual_seed(7000 + trial) + base = torch.randint(0, 50, (64,), generator=generator, dtype=torch.int64) + trial * 3 + values = [int(value) for value in base] + maxv = (1 << (bits - 1)) - 1 if signed else (1 << bits) - 1 + minv = -(1 << (bits - 1)) if signed else 0 + high = 1 << (bits - 1) + if signed: + first_pattern = [maxv, minv, -1, 0] + second_pattern = [1, -1, 1, minv] + else: + first_pattern = [maxv, 0, high, 1] + second_pattern = [1, high, 1, maxv] + variants = [ + (first_pattern, second_pattern), + (second_pattern, first_pattern), + (list(reversed(first_pattern)), list(reversed(second_pattern))), + ] + first, second = variants[trial % len(variants)] + values[0:4] = first + values[lanes : lanes + 4] = second + scalar_index = lanes + 4 + if trial % 2 == 0: + values[scalar_index] = -5 if signed else high + else: + values[scalar_index] = 0 + return torch.tensor(values, dtype=torch_dtype) + + +@pytest.mark.parametrize("dtype", list(CASES), ids=list(CASES)) +def test_simtvf_local_integer_roundtrip_arithmetic_and_boundaries(npu, dtype): + torch_dtype, lanes, bounds, signed, highbit = CASES[dtype] + kernel = tilelang.compile( + _local_integer_kernel(dtype, lanes, bounds, signed, highbit), + target="pto", + out_idx=-1, + ) + source = kernel.get_kernel_source() + + backing = "pto.i8" if dtype in ("int8", "uint8") else "pto.i32" + logical = {"int8": "pto.si8", "uint8": "pto.ui8", "int32": "pto.si32", "uint32": "pto.ui32"}[dtype] + # Live allocation and all four adapter landings are present in the final + # generated source (they survived lowering/optimization). + assert f"pto.alloc_buffer((64,), {backing})" in source, source + scalar_index = lanes + 4 + assert f"pto.cast(pto.load(local, {scalar_index}), {logical})" in source, source + assert f"pto.store(pto.cast(pto.load(local, 0, contiguous={lanes}), {logical})" in source, source + assert f", {backing}), local, 0, contiguous={lanes})" in source, source + assert f", {backing}), local, {scalar_index})" in source, source + + arms = set() + for trial in range(3): + a = _trial_input(torch_dtype, signed, trial, lanes) + x = int(a[scalar_index]) + arms.add(x < 0 if signed else x >= highbit) + out = kernel(a.to(npu)) + torch.npu.synchronize() + got = out.cpu() + expected = _expected(a, lanes, bounds, signed) + assert torch.equal(got, expected), f"{dtype} trial {trial} mismatch\n got={got.tolist()}\n exp={expected.tolist()}" + # Guarded output: every lane outside the intended writes still holds + # the independently initialized input value. + touched = set(range(lanes)) | {scalar_index, scalar_index + 1, 40, 41, 42, 43} + for index in range(64): + if index not in touched: + assert got[index].item() == a[index].item(), (dtype, trial, index) + assert arms == {True, False}, f"{dtype} comparison arms exercised: {arms}" + + +def _shared_control_kernel(dtype): + @T.prim_func + def main(A: T.Tensor((8,), dtype), C: T.Tensor((8,), dtype)): + with T.Kernel(1): + src = T.alloc_shared((8,), dtype) + dst = T.alloc_shared((8,), dtype) + T.copy(A, src) + T.copy(A, dst) + with T.SimtVF(threads=1): + dst[0] = src[0] + T.copy(dst, C) + + return main + + +@pytest.mark.parametrize("dtype,torch_dtype", [("uint8", torch.uint8), ("int32", torch.int32)]) +def test_simtvf_shared_integer_access_control(npu, dtype, torch_dtype): + kernel = tilelang.compile(_shared_control_kernel(dtype), target="pto", out_idx=-1) + source = kernel.get_kernel_source() + line = next(line for line in source.splitlines() if "pto.store(pto.load(" in line) + assert "pto.cast(" not in line, line + + a = torch.arange(-4, 4, dtype=torch.int32).to(torch_dtype) + out = kernel(a.to(npu)) + torch.npu.synchronize() + assert torch.equal(out.cpu(), a) From ca9986edc71782e334a909df84ad7cfdf78f545b Mon Sep 17 00:00:00 2001 From: qukelin Date: Sat, 17 Oct 2026 12:46:37 +0800 Subject: [PATCH 06/10] fix(pto): track VMI value provenance across scalar bridges --- src/ascend/codegen/codegen_pto.cc | 637 ++++++++++++- src/ascend/codegen/codegen_pto.h | 40 + .../pto/test_pto_vmi_value_provenance.py | 854 ++++++++++++++++++ .../pto/test_pto_vmi_value_provenance_run.py | 246 +++++ 4 files changed, 1764 insertions(+), 13 deletions(-) create mode 100644 testing/ascend/language/pto/test_pto_vmi_value_provenance.py create mode 100644 testing/ascend/language/pto/test_pto_vmi_value_provenance_run.py diff --git a/src/ascend/codegen/codegen_pto.cc b/src/ascend/codegen/codegen_pto.cc index a447913b5d..5a570ddbf2 100644 --- a/src/ascend/codegen/codegen_pto.cc +++ b/src/ascend/codegen/codegen_pto.cc @@ -2103,6 +2103,7 @@ void CodeGenTileLangPTO::ResetFunctionState_() { mixed_entry_snapshot_ids_.clear(); loop_condition_snapshot_ids_.clear(); simd_pair_vars_.clear(); + value_defs_.clear(); gemm_emit_contexts_.clear(); gemm_emit_context_by_call_.clear(); gemm_zero_addr_emitted_ = false; @@ -2667,6 +2668,7 @@ std::string CodeGenTileLangPTO::PrintCondition(const PrimExpr &condition) { (dtype.is_bool() || dtype.is_int() || dtype.is_uint())) << "PTO control-flow condition must be a scalar integer predicate, got " << dtype; + CheckPtoScalarOperand_(condition, "PTO control-flow condition"); std::string value = RemoveOutermostParentheses(PrintExpr_(condition)); // TIR follows C truthiness for all integer conditions. Always compare at // the use site because a bool-typed expression may be materialized in an @@ -2718,20 +2720,24 @@ std::string CodeGenTileLangPTO::ScalarLoad(const BufferNode *buffer, std::string CodeGenTileLangPTO::LocalVarStoreValue(const PrimExpr &value, DataType dtype) { - std::string value_expr = RemoveOutermostParentheses(PrintExpr_(value)); if (dtype.is_handle()) { - return value_expr; + return RemoveOutermostParentheses(PrintExpr_(value)); } bool is_literal = value.as() != nullptr || value.as() != nullptr; if (dtype.is_scalar() && is_literal) { - return "pto.const(" + value_expr + ", dtype=" + DataTypeName(dtype) + ")"; + return "pto.const(" + RemoveOutermostParentheses(PrintExpr_(value)) + + ", dtype=" + DataTypeName(dtype) + ")"; } if (dtype.is_bool() || dtype.bits() == 1 || dtype.is_int() || dtype.is_uint()) { - return ScalarCastExpr(value_expr, dtype, "PTO local.var store"); + // Check provenance before printing/adapting the scalar (covers bool and + // integer local.var stores; other local.var branches keep their route). + CheckPtoScalarOperand_(value, "PTO local.var integer store"); + return ScalarCastExpr(RemoveOutermostParentheses(PrintExpr_(value)), dtype, + "PTO local.var store"); } - return value_expr; + return RemoveOutermostParentheses(PrintExpr_(value)); } void CodeGenTileLangPTO::EmitScalarStore(const BufferNode *buffer, @@ -3471,6 +3477,7 @@ std::string CodeGenTileLangPTO::GetLocalByteAddrExpr( ", dtype=pto.int64)"; } + CheckPtoScalarOperand_(elem_index, context.c_str()); std::string index_expr = RemoveOutermostParentheses(PrintExpr_(elem_index)); std::string coerced = ScalarCastExpr(index_expr, DataType::Int(64), context); if (elem_bytes == 1) { @@ -4887,6 +4894,513 @@ static bool IsVmiVcvtFromFp4(const CallNode *op) { op->args[0].dtype().element_of().is_float4_e2m1fn(); } +static std::string VmiDtypeAnnotation(DataType element) { + if (element.is_int()) { + return PtoSignedIntegerTypeName(element); + } + if (element.is_uint()) { + return "pto.ui" + std::to_string(element.bits()); + } + return DataTypeName(element); +} + +const char *CodeGenTileLangPTO::PtoValueKindName(PtoValueKind kind) { + switch (kind) { + case PtoValueKind::kScalar: + return "scalar"; + case PtoValueKind::kVmiData: + return "VMI data"; + case PtoValueKind::kVmiMask: + return "VMI mask"; + case PtoValueKind::kVmiPair: + return "VMI pair"; + case PtoValueKind::kVoid: + return "void"; + case PtoValueKind::kHandle: + return "handle"; + case PtoValueKind::kOtherVector: + return "other vector"; + case PtoValueKind::kUnknown: + return "unknown"; + } + return "unknown"; +} + +struct CodeGenTileLangPTO::PtoValueQueryContext { + // Query-local lexical overrides (Let bindings seen by the active query). + std::unordered_map local_defs; + // Vars on the current definition chain; re-entering one is a cycle. + std::unordered_set active; + // Query-local memo for expression/definition classification. Node identity + // is stable, but a result computed under an active Let override is only + // valid for that lexical scope, so entries carry the scope generation. + std::unordered_map> + memo; + size_t scope_generation{0}; +}; + +// Classifies each immediate semantic operand of `expr` exactly once without +// descending into operand subtrees: ExprVisitor's per-node dispatch walks one +// level, and each child visit is intercepted below and handed to the full +// (memoized) classification. This keeps ordinary scalar expressions linear +// even for deeply nested or shared DAG-shaped terms. +struct CodeGenTileLangPTO::PtoOperandScanner final : public tirx::ExprVisitor { + PtoOperandScanner(const CodeGenTileLangPTO *codegen, + PtoValueQueryContext *context) + : codegen_(codegen), context_(context) {} + + bool has_unknown_origin{false}; + + void VisitExpr(const PrimExpr &expr) final { + if (!expr.defined() || has_unknown_origin) { + return; + } + PtoValueKind kind = codegen_->ClassifyPtoValue_(expr, context_); + has_unknown_origin = + kind == PtoValueKind::kVmiData || kind == PtoValueKind::kVmiMask || + kind == PtoValueKind::kVmiPair || kind == PtoValueKind::kVoid || + kind == PtoValueKind::kUnknown; + } + + // Walks the immediate semantic operands of the root expression. The + // qualified base call bypasses this class's VisitExpr override for the root + // only; the per-node handlers reach the children through the virtual + // VisitExpr above. + void ScanImmediateOperands(const PrimExpr &expr) { + tirx::ExprVisitor::VisitExpr(expr); + } + +private: + const CodeGenTileLangPTO *codegen_; + PtoValueQueryContext *context_; +}; + +namespace { + +bool IsVmiOpName(const std::string &name) { + return name.rfind("tl.vmi.", 0) == 0; +} + +} // namespace + +CodeGenTileLangPTO::PtoValueKind +CodeGenTileLangPTO::ClassifyPtoValue_(const PrimExpr &expr) const { + PtoValueQueryContext context; + return ClassifyPtoValue_(expr, &context); +} + +CodeGenTileLangPTO::PtoValueKind +CodeGenTileLangPTO::ClassifyPtoValue_(const PrimExpr &expr, + PtoValueQueryContext *context) const { + if (!expr.defined() || context == nullptr) { + return PtoValueKind::kUnknown; + } + const Object *node = expr.get(); + auto memo_it = context->memo.find(node); + if (memo_it != context->memo.end() && + memo_it->second.first == context->scope_generation) { + return memo_it->second.second; + } + PtoValueKind kind = ClassifyPtoValueUncached_(expr, context); + context->memo[node] = std::make_pair(context->scope_generation, kind); + return kind; +} + +CodeGenTileLangPTO::PtoValueKind CodeGenTileLangPTO::ClassifyPtoValueUncached_( + const PrimExpr &expr, PtoValueQueryContext *context) const { + if (const auto *var = expr.as()) { + const PrimExpr *definition = nullptr; + PrimExpr stored; + auto local_it = context->local_defs.find(var); + if (local_it != context->local_defs.end()) { + stored = local_it->second; + definition = &stored; + } else { + auto it = value_defs_.find(var); + if (it != value_defs_.end()) { + stored = it->second; + definition = &stored; + } + } + if (definition != nullptr) { + ICHECK(context->active.insert(var).second) + << "PTO SSA definition cycle detected for variable " + << var->name_hint; + PtoValueKind kind = ClassifyPtoValue_(*definition, context); + context->active.erase(var); + return kind; + } + // Unbound Var: parameter/loop variable contract. + if (expr.dtype().is_handle()) { + return PtoValueKind::kHandle; + } + if (expr.dtype().is_scalar()) { + return PtoValueKind::kScalar; + } + return PtoValueKind::kOtherVector; + } + + if (expr.as() != nullptr || expr.as() != nullptr) { + return PtoValueKind::kScalar; + } + + if (const auto *let = expr.as()) { + // Query-local lexical scope: bind the RHS in the local override layer and + // classify the body with the binding visible. The persistent table is not + // modified. The RHS is classified lazily when the body resolves the bound + // Var: classifying it eagerly here would traverse the same RHS a second + // time under the bumped scope generation and make a linear nested-Let + // chain exponential. Unused bindings do not affect the body's result + // classification; emission still handles and validates the RHS. + auto previous = context->local_defs.find(let->var.get()); + bool had_previous = previous != context->local_defs.end(); + PrimExpr saved; + if (had_previous) { + saved = previous->second; + } + context->local_defs[let->var.get()] = let->value; + ++context->scope_generation; + PtoValueKind kind = ClassifyPtoValue_(let->body, context); + ++context->scope_generation; + if (had_previous) { + context->local_defs[let->var.get()] = saved; + } else { + context->local_defs.erase(let->var.get()); + } + return kind; + } + + if (const auto *cast = expr.as()) { + PtoValueKind source = ClassifyPtoValue_(cast->value, context); + if (cast->value.dtype() == cast->dtype) { + return source; + } + // Non-identity casts of VMI values require VMI conversion operations. + // A scalar result dtype alone does not make a VMI or unknown source + // a scalar value. + if (source == PtoValueKind::kVmiData || source == PtoValueKind::kVmiMask || + source == PtoValueKind::kVmiPair || source == PtoValueKind::kVoid || + source == PtoValueKind::kUnknown) { + return PtoValueKind::kUnknown; + } + if (expr.dtype().is_handle()) { + return PtoValueKind::kHandle; + } + if (expr.dtype().is_scalar()) { + return PtoValueKind::kScalar; + } + return PtoValueKind::kOtherVector; + } + + if (const auto *load = expr.as()) { + return ClassifyBufferLoad_(load); + } + + if (const auto *call = expr.as()) { + auto opt_call_op = call->op.as(); + if (opt_call_op.has_value()) { + const std::string name = opt_call_op.value()->name; + if (IsVmiOpName(name)) { + return ClassifyVmiCallResult_(call, context); + } + if (call->op.same_as(tirx::builtin::reinterpret())) { + ICHECK_EQ(call->args.size(), 1U) + << "tirx.reinterpret expects 1 argument (value)"; + PtoValueKind source = ClassifyPtoValue_(call->args[0], context); + if (source == PtoValueKind::kVmiData) { + DataType source_dtype = call->args[0].dtype(); + DataType target_dtype = expr.dtype(); + if (source_dtype.bits() == target_dtype.bits() && + source_dtype.lanes() == target_dtype.lanes()) { + return PtoValueKind::kVmiData; + } + return PtoValueKind::kUnknown; + } + if (source == PtoValueKind::kVmiMask || + source == PtoValueKind::kVmiPair || source == PtoValueKind::kVoid || + source == PtoValueKind::kUnknown) { + return PtoValueKind::kUnknown; + } + if (expr.dtype().is_handle()) { + return PtoValueKind::kHandle; + } + if (expr.dtype().is_scalar()) { + return PtoValueKind::kScalar; + } + return PtoValueKind::kOtherVector; + } + } + return ClassifyOrdinaryExpr_(expr, context); + } + + return ClassifyOrdinaryExpr_(expr, context); +} + +CodeGenTileLangPTO::PtoValueKind +CodeGenTileLangPTO::ClassifyOrdinaryExpr_(const PrimExpr &expr, + PtoValueQueryContext *context) const { + // Ordinary scalar expressions keep their contract unless a semantic operand + // carries a VMI kind or an unknown VMI origin. Only the immediate operands + // are classified (once each, memoized); the query never prints or emits. + PtoOperandScanner scanner(this, context); + scanner.ScanImmediateOperands(expr); + if (scanner.has_unknown_origin) { + return PtoValueKind::kUnknown; + } + if (expr.dtype().is_handle()) { + return PtoValueKind::kHandle; + } + if (expr.dtype().is_scalar()) { + return PtoValueKind::kScalar; + } + return PtoValueKind::kOtherVector; +} + +CodeGenTileLangPTO::PtoValueKind +CodeGenTileLangPTO::ClassifyBufferLoad_(const BufferLoadNode *op) const { + // Ordinary memory reads follow the storage contract: a scalar-valued + // BufferLoad is a scalar regardless of which VMI store previously wrote the + // address. VMI local register slots are the only buffer reads classified as + // VMI data; multi-lane local.var and ordinary vectors stay on their existing + // register/vector routes. + if (IsVmiLocalRegisterBuffer(op->buffer.get())) { + return PtoValueKind::kVmiData; + } + if (op->dtype.is_scalar()) { + return PtoValueKind::kScalar; + } + return PtoValueKind::kOtherVector; +} + +CodeGenTileLangPTO::PtoValueKind CodeGenTileLangPTO::ClassifyVmiCallResult_( + const CallNode *op, PtoValueQueryContext *context) const { + auto opt_call_op = op->op.as(); + ICHECK(opt_call_op.has_value()) << "tl.vmi call has a non-Op callee"; + const std::string name = opt_call_op.value()->name; + if (name == "tl.vmi.create_mask" || name == "tl.vmi.vcmp" || + name == "tl.vmi.vcmps") { + return PtoValueKind::kVmiMask; + } + if (name == "tl.vmi.vintlv" || name == "tl.vmi.vdintlv" || + name == "tl.vmi.vunzip" || name == "tl.vmi.vmull") { + return PtoValueKind::kVmiPair; + } + if (name == "tl.vmi.vload") { + // vload is data unless the explicit dist_mode=dintlv annotation selects + // the pair-returning form. + auto it = op->annotations.find("dist_mode"); + if (it != op->annotations.end()) { + const auto *mode = (*it).second.as(); + ICHECK(mode != nullptr) + << "tl.vmi.vload dist_mode annotation must be a string"; + if (mode->value == "dintlv") { + return PtoValueKind::kVmiPair; + } + } + return PtoValueKind::kVmiData; + } + if (name == "tl.vmi.pair_get") { + ICHECK_EQ(op->args.size(), 2U) << "tl.vmi.pair_get expects (pair, index)"; + PtoValueKind pair_kind = ClassifyPtoValue_(op->args[0], context); + ICHECK(pair_kind == PtoValueKind::kVmiPair) + << "PTO provenance tl.vmi.pair_get source must be a VMI pair; got " + << PtoValueKindName(pair_kind) << " (" << op->args[0].dtype() + << "): " << op->args[0]; + const auto *index = op->args[1].as(); + ICHECK(index != nullptr && (index->value == 0 || index->value == 1)) + << "PTO provenance tl.vmi.pair_get index must be the constant 0 or 1, " + "got " + << op->args[1]; + return PtoValueKind::kVmiData; + } + if (name == "tl.vmi.vstore" || name == "tl.vmi.vscatter") { + return PtoValueKind::kVoid; + } + static const std::unordered_set kDataOps = { + "tl.vmi.vci", "tl.vmi.vbrc", + "tl.vmi.vzip", "tl.vmi.vadd", + "tl.vmi.vsub", "tl.vmi.vmul", + "tl.vmi.vdiv", "tl.vmi.vmax", + "tl.vmi.vmin", "tl.vmi.vand", + "tl.vmi.vor", "tl.vmi.vxor", + "tl.vmi.vshl", "tl.vmi.vshr", + "tl.vmi.vabs", "tl.vmi.vneg", + "tl.vmi.vrelu", "tl.vmi.vexp", + "tl.vmi.vln", "tl.vmi.vsqrt", + "tl.vmi.vnot", "tl.vmi.vadds", + "tl.vmi.vmuls", "tl.vmi.vmaxs", + "tl.vmi.vmins", "tl.vmi.vshls", + "tl.vmi.vshrs", "tl.vmi.vsel", + "tl.vmi.vselr", "tl.vmi.vcadd", + "tl.vmi.vcmax", "tl.vmi.vcmin", + "tl.vmi.vcvt", "tl.vmi.vinterpret_cast", + "tl.vmi.vexpdif", "tl.vmi.vaxpy", + "tl.vmi.vlrelu", "tl.vmi.vprelu", + "tl.vmi.vmula", "tl.vmi.vdhist", + "tl.vmi.vchist", "tl.vmi.vgather", + "tl.vmi.vgatherb", + }; + if (kDataOps.count(name) != 0U) { + return PtoValueKind::kVmiData; + } + // Registered operations without an explicit contract (for example a future + // addition) stay unknown so consumers diagnose instead of guessing. + return PtoValueKind::kUnknown; +} + +void CodeGenTileLangPTO::CheckPtoScalarOperand_(const PrimExpr &expr, + const char *consumer) const { + PtoValueKind kind = ClassifyPtoValue_(expr); + if (kind == PtoValueKind::kScalar) { + return; + } + std::ostringstream message; + message << consumer << " expects a scalar; got " << PtoValueKindName(kind) + << " (" << expr.dtype() << ", lanes=" << expr.dtype().lanes() + << "): " << expr; + if (kind == PtoValueKind::kUnknown) { + const auto *call = expr.as(); + if (call != nullptr) { + auto opt_call_op = call->op.as(); + if (opt_call_op.has_value() && IsVmiOpName(opt_call_op.value()->name)) { + message << "; unclassified " << opt_call_op.value()->name + << " result; add an explicit VMI return contract"; + } + } + } + LOG(FATAL) << message.str(); +} + +bool CodeGenTileLangPTO::ResolvesToMutableRegisterRead( + const PrimExpr &expr, + std::unordered_map *local_defs, + std::unordered_set *active) const { + if (!expr.defined() || local_defs == nullptr || active == nullptr) { + return false; + } + if (const auto *cast = expr.as()) { + // Same-dtype casts are transparent aliases; numeric conversions are not. + return cast->value.dtype() == cast->dtype && + ResolvesToMutableRegisterRead(cast->value, local_defs, active); + } + if (const auto *let = expr.as()) { + auto previous = local_defs->find(let->var.get()); + const bool had_previous = previous != local_defs->end(); + PrimExpr saved; + if (had_previous) { + saved = previous->second; + } + (*local_defs)[let->var.get()] = let->value; + const bool resolves = + ResolvesToMutableRegisterRead(let->body, local_defs, active); + if (had_previous) { + (*local_defs)[let->var.get()] = saved; + } else { + local_defs->erase(let->var.get()); + } + return resolves; + } + if (const auto *var = expr.as()) { + PrimExpr stored; + const PrimExpr *definition = nullptr; + auto local_it = local_defs->find(var); + if (local_it != local_defs->end()) { + stored = local_it->second; + definition = &stored; + } else { + auto it = value_defs_.find(var); + if (it != value_defs_.end()) { + stored = it->second; + definition = &stored; + } + } + if (definition == nullptr || !active->insert(var).second) { + return false; + } + const bool resolves = + ResolvesToMutableRegisterRead(*definition, local_defs, active); + active->erase(var); + return resolves; + } + if (const auto *load = expr.as()) { + return ScopeOfBuffer(load->buffer.get()) == "local.var"; + } + return false; +} + +bool CodeGenTileLangPTO::IsMutableRegisterGroupSource( + const PrimExpr &source) const { + // Accept numeric local.var reads through Var/Let aliases and same-dtype + // casts. Keep their classification unchanged: other-vector for multiple + // lanes, scalar for one lane. This helper only checks whether grouped vbrc + // may consume the expression; it does not inspect mutable stores. + // is_uint() includes TIR uint1/bool. Handle reads are not numeric. + // Other expressions, including VMI mask/pair/void/unknown results, do not + // resolve to this read pattern. + const DataType dtype = source.dtype(); + if (!dtype.is_int() && !dtype.is_uint() && !dtype.is_float() && + !dtype.is_bfloat16()) { + return false; + } + std::unordered_map local_defs; + std::unordered_set active; + return ResolvesToMutableRegisterRead(source, &local_defs, &active); +} + +void CodeGenTileLangPTO::CheckPtoVmiBridgeOperand_( + const CallNode *op, const char *op_name, PtoValueKind source_kind) const { + ICHECK(!op->args.empty()) << op_name << " expects a source operand"; + const PrimExpr &source = op->args[0]; + const bool has_group = op->annotations.find("group") != op->annotations.end(); + if (std::string(op_name) == "tl.vmi.vci") { + if (source_kind == PtoValueKind::kVmiData) { + LOG(FATAL) << "tl.vmi.vci base expects a scalar; got VMI data (" + << source.dtype() << ", lanes=" << source.dtype().lanes() + << "): " << source; + } + if (source_kind != PtoValueKind::kScalar) { + LOG(FATAL) << "tl.vmi.vci base expects a scalar; got " + << PtoValueKindName(source_kind) << " (" << source.dtype() + << ", lanes=" << source.dtype().lanes() << "): " << source; + } + return; + } + // tl.vmi.vbrc + if (has_group) { + const auto *group = op->annotations.at("group").as(); + ICHECK(group != nullptr) + << "tl.vmi.vbrc group annotation must be a constant integer"; + if (source_kind != PtoValueKind::kVmiData) { + // Also accept numeric local.var reads through pure aliases without + // changing their classification. Scalar constants and parameters do not + // match this read pattern. Require the same lane count as the group. + if (IsMutableRegisterGroupSource(source)) { + ICHECK_EQ(source.dtype().lanes(), group->value) + << "tl.vmi.vbrc source lane count must equal group; got " + << source.dtype().lanes() << " lanes and group=" << group->value; + return; + } + LOG(FATAL) << "tl.vmi.vbrc with group=... requires a VMI vreg; got " + << PtoValueKindName(source_kind) << " (" << source.dtype() + << ", lanes=" << source.dtype().lanes() << "): " << source; + } + ICHECK_EQ(source.dtype().lanes(), group->value) + << "tl.vmi.vbrc source lane count must equal group; got " + << source.dtype().lanes() << " lanes and group=" << group->value; + return; + } + if (source_kind == PtoValueKind::kVmiMask || + source_kind == PtoValueKind::kVmiPair || + source_kind == PtoValueKind::kVoid || + source_kind == PtoValueKind::kHandle || + source_kind == PtoValueKind::kUnknown) { + LOG(FATAL) << "tl.vmi.vbrc cannot broadcast " + << PtoValueKindName(source_kind) << " (" << source.dtype() + << ", lanes=" << source.dtype().lanes() << "): " << source; + } +} + void CodeGenTileLangPTO::PrintPtoVmiCall_(const CallNode *op, std::ostream &os) { auto opt_call_op = op->op.as(); @@ -4946,7 +5460,18 @@ void CodeGenTileLangPTO::PrintPtoVmiCall_(const CallNode *op, return; } - // Adapt integer scalar sources of vci/vbrc at the VMI + // Classify the VMI source *before* printing any argument. A lane-1 VMI + // vreg is a TIR scalar by dtype but remains VMI data; only a proven scalar + // may enter the integer signless bridge. vci/scalar-only consumers receive + // an explicit category diagnostic instead of an invalid PTODSL scalar op. + PtoValueKind vmi_source_kind = PtoValueKind::kUnknown; + if (!op->args.empty()) { + vmi_source_kind = ClassifyPtoValue_(op->args[0]); + } + if (op_name == "tl.vmi.vbrc" || op_name == "tl.vmi.vci") { + CheckPtoVmiBridgeOperand_(op, op_name.c_str(), vmi_source_kind); + } + // Adapt proven integer scalar sources of vci/vbrc at the VMI // value-argument boundary. The generated scalar must be *signless*: PTOAS // materialises per-physical-vector offsets with `arith.constant`, which // requires signless integers, so a typed (signed) constant base fails for @@ -4954,17 +5479,17 @@ void CodeGenTileLangPTO::PrintPtoVmiCall_(const CallNode *op, // a constant whenever the loop has a single chunk, so integer VCI always // takes this path; VBRC constants keep the verified typed-literal path. const bool scalar_integer_input = - !op->args.empty() && op->args[0].dtype().is_scalar() && + vmi_source_kind == PtoValueKind::kScalar && !op->args.empty() && + op->args[0].dtype().is_scalar() && (op->args[0].dtype().is_int() || op->args[0].dtype().is_uint()); - const bool integer_vector_result = - op->dtype.is_vector() && - (op->dtype.element_of().is_int() || op->dtype.element_of().is_uint()); + const bool integer_data_result = + op->dtype.element_of().is_int() || op->dtype.element_of().is_uint(); const bool integer_source_literal = !op->args.empty() && op->args[0].as() != nullptr; const bool bridge_integer_source = (op_name == "tl.vmi.vci" || (op_name == "tl.vmi.vbrc" && !integer_source_literal)) && - scalar_integer_input && integer_vector_result; + scalar_integer_input && integer_data_result; if (bridge_integer_source) { os << "pto.vmi.vinterpret_cast("; @@ -5133,6 +5658,15 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, RequiresStatementControlFlow(op->args[1]) || RequiresStatementControlFlow(op->args[2]); auto branch_value = [&](const PrimExpr &expr) { + const bool scalar_adapted_branch = + expr.as() == nullptr && + expr.as() == nullptr && + (expr.dtype().is_bool() || + (op->dtype.is_scalar() && + (op->dtype.is_int() || op->dtype.is_uint()))); + if (scalar_adapted_branch) { + CheckPtoScalarOperand_(expr, "PTO if_then_else branch"); + } std::string value = PrintExpr_(expr); if (expr.dtype().is_bool()) { // PTODSL typed constants require numeric bool literals. @@ -5512,8 +6046,35 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, if (op->op.same_as(tirx::builtin::reinterpret())) { ICHECK_EQ(op->args.size(), 1U) << "tirx.reinterpret expects 1 argument (value)"; - std::string value = PrintExpr_(op->args[0]); + PtoValueKind reinterpret_source_kind = ClassifyPtoValue_(op->args[0]); DataType target_dtype = op->dtype; + if (reinterpret_source_kind == PtoValueKind::kVmiData) { + DataType source_dtype = op->args[0].dtype(); + if (source_dtype.bits() == target_dtype.bits() && + source_dtype.lanes() == target_dtype.lanes()) { + // Same element width and lane count is a register-level re-view; use + // the explicit VMI reinterpret API, never scalar/lexical bitcasts. + os << "pto.vmi.vinterpret_cast(" << PrintExpr_(op->args[0]) + << ", to_dtype=" << VmiDtypeAnnotation(target_dtype.element_of()) + << ")"; + return; + } + LOG(FATAL) << "PTO reinterpret of VMI data requires the same element " + "width and lane count; use T.vmi.vinterpret_cast with the " + "explicit packed/shape contract (source " + << source_dtype << ", target " << target_dtype + << "): " << op->args[0]; + } + if (reinterpret_source_kind == PtoValueKind::kVmiMask || + reinterpret_source_kind == PtoValueKind::kVmiPair || + reinterpret_source_kind == PtoValueKind::kVoid || + reinterpret_source_kind == PtoValueKind::kUnknown) { + LOG(FATAL) << "PTO reinterpret cannot consume " + << PtoValueKindName(reinterpret_source_kind) << " (source " + << op->args[0].dtype() << ", target " << target_dtype + << "): " << op->args[0]; + } + std::string value = PrintExpr_(op->args[0]); if (target_dtype.lanes() > 1) { // Vector reinterpret: lower to pto.vbitcast(vec, elem_type). os << "pto.vbitcast(" << value << ", " @@ -5913,6 +6474,7 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, << "tl.simd vlrelu/vprelu expects 3 arguments"; if (op->op.same_as(tl::simd_vlrelu())) { DataType vector_dtype = op->args[0].dtype(); + CheckPtoScalarOperand_(op->args[1], "tl.simd.vlrelu scalar"); std::string alpha = WrapTypedConst( RemoveOutermostParentheses(PrintExpr_(PeelScalarCasts(op->args[1]))), vector_dtype, IsImmediateScalar(op->args[1])); @@ -5962,6 +6524,7 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, : "vshrs"; std::string src = PrintExpr_(op->args[0]); DataType result_dtype = op->args[0].dtype(); + CheckPtoScalarOperand_(op->args[1], "tl.simd vector-scalar operand"); std::string scalar = WrapTypedConst( RemoveOutermostParentheses(PrintExpr_(PeelScalarCasts(op->args[1]))), result_dtype, IsImmediateScalar(op->args[1])); @@ -6036,6 +6599,7 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, if (op->op.same_as(tl::simd_vcmps())) { ICHECK_EQ(op->args.size(), 4U) << "tl.simd.vcmps expects 4 arguments (src, scalar, mask, op)"; + CheckPtoScalarOperand_(op->args[1], "tl.simd.vcmps scalar"); std::string cmp_op = Downcast(op->args[3])->value; static const std::unordered_map kCmpModeEnum = { {"eq", "EQ"}, {"ne", "NE"}, {"gt", "GT"}, @@ -6067,6 +6631,7 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, // Match the authored vector element dtype (i16 for packed-ue8m0 lane math, // i32 for f32 gather paths). Always-i32 coercion breaks i16 vci users. DataType elem = op->dtype.element_of(); + CheckPtoScalarOperand_(op->args[0], "tl.simd.vci index"); std::string index; if (IsImmediateScalar(op->args[0])) { index = @@ -6087,6 +6652,7 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, << "tl.simd.vdup expects 3 arguments (src, mask, mode)"; ValidateVecMode_(op, 2); DataType elem = op->dtype.element_of(); + CheckPtoScalarOperand_(op->args[0], "tl.simd.vdup scalar"); std::string scalar; if (IsImmediateScalar(op->args[0])) { scalar = @@ -6487,7 +7053,17 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, if (op_name == "tl.vmi.pair_get") { ICHECK_EQ(op->args.size(), 2U) << "tl.vmi.pair_get expects exactly 2 arguments"; - os << "(" << PrintExpr_(op->args[0]) << ")[" << PrintExpr_(op->args[1]) + PtoValueKind pair_kind = ClassifyPtoValue_(op->args[0]); + ICHECK(pair_kind == PtoValueKind::kVmiPair) + << "tl.vmi.pair_get source must be a VMI pair; got " + << PtoValueKindName(pair_kind) << " (" << op->args[0].dtype() + << "): " << op->args[0]; + const auto *pair_index = op->args[1].as(); + ICHECK(pair_index != nullptr && + (pair_index->value == 0 || pair_index->value == 1)) + << "tl.vmi.pair_get index must be the constant 0 or 1, got " + << op->args[1]; + os << "(" << PrintExpr_(op->args[0]) << ")[" << pair_index->value << "]"; return; } @@ -6533,6 +7109,21 @@ void CodeGenTileLangPTO::VisitExpr_(const CastNode *op, return; } + PtoValueKind cast_source_kind = ClassifyPtoValue_(op->value); + if (cast_source_kind == PtoValueKind::kVmiData) { + LOG(FATAL) << "PTO integer cast cannot consume VMI data; use T.vmi.vcvt " + "or T.vmi.vinterpret_cast (source " + << op->value.dtype() << ", target " << to << "): " << op->value; + } + if (cast_source_kind == PtoValueKind::kVmiMask || + cast_source_kind == PtoValueKind::kVmiPair || + cast_source_kind == PtoValueKind::kVoid || + cast_source_kind == PtoValueKind::kUnknown) { + LOG(FATAL) << "PTO cast cannot consume " + << PtoValueKindName(cast_source_kind) << " (source " + << op->value.dtype() << ", target " << to << "): " << op->value; + } + const bool from_fp8 = tl::IsAscendVectorizableFP8(from); const bool to_fp8 = tl::IsAscendVectorizableFP8(to); if (from_fp8 || to_fp8) { @@ -6669,8 +7260,14 @@ void CodeGenTileLangPTO::VisitExpr_(const LetNode *op, << "PTO Let variable is already defined: " << op->var->name_hint; std::string vid = AllocVarID(op->var.get()); + ICHECK(!value_defs_.count(op->var.get())) + << "duplicate PTO Let definition: " << op->var->name_hint; + value_defs_.emplace(op->var.get(), op->value); std::string body = PrintExpr_(op->body); + bool removed_def = value_defs_.erase(op->var.get()); + ICHECK(removed_def) << "PTO Let definition was not registered: " + << op->var->name_hint; bool removed = var_idmap_.erase(op->var.get()); ICHECK(removed) << "PTO Let variable was not registered: " << op->var->name_hint; @@ -6867,6 +7464,9 @@ void CodeGenTileLangPTO::VisitExpr_(const NotNode *op, ICHECK(op->dtype.is_scalar() && op->a.dtype().is_bool()) << "PTO logical-not currently supports scalar bool predicates only, got " << op->a.dtype(); + // A lane-1 VMI mask is a TIR scalar bool but remains a VMI value and cannot + // enter the logical scalar adapter. + CheckPtoScalarOperand_(op->a, "PTO logical-not operand"); // tl.logical_not(value) implements logical negation as value == 0. os << "tl.logical_not("; PrintExpr_(op->a, os); @@ -6885,10 +7485,14 @@ void CodeGenTileLangPTO::VisitExpr_(const SelectNode *op, return; } if (op->dtype.is_scalar() && op->dtype.is_bool()) { + // The logical-bool adapter is scalar-only; a lane-1 VMI mask (or VMI + // data used as a predicate) must be diagnosed, not normalized. + CheckPtoScalarOperand_(expr, "PTO select branch"); os << "tl.as_logical_bool(" << PrintExpr_(expr) << ")"; return; } if (coerce_integer_branches) { + CheckPtoScalarOperand_(expr, "PTO select branch"); os << ScalarCastExpr(PrintExpr_(expr), op->dtype, "PTO select branch"); return; } @@ -7014,6 +7618,13 @@ void CodeGenTileLangPTO::VisitExpr_(const BroadcastNode *op, void CodeGenTileLangPTO::VisitStmt_(const DeclBufferNode *op) { (void)op; } void CodeGenTileLangPTO::VisitStmt_(const BindNode *op) { + // Record every SSA definition before any printing branch (including the + // early-return SIMD pair/pointer paths below). Queries only consume the + // stored IR; they never depend on the printed form. + ICHECK(!value_defs_.count(op->var.get())) + << "duplicate PTO SSA definition: " << op->var->name_hint; + value_defs_.emplace(op->var.get(), op->value); + if (const auto *call = op->value.as()) { if (IsOpName(call->op, "tl.simd.alloc")) { AllocVarID(op->var.get()); diff --git a/src/ascend/codegen/codegen_pto.h b/src/ascend/codegen/codegen_pto.h index f442c1d949..47a2dc3934 100644 --- a/src/ascend/codegen/codegen_pto.h +++ b/src/ascend/codegen/codegen_pto.h @@ -194,6 +194,43 @@ class CodeGenTileLangPTO final : public CodeGenTileLangPY { std::string PrintVmiAnnotationValue(const std::string &key, const ObjectRef &value); void PrintPtoVmiCall_(const CallNode *op, std::ostream &os); + // SSA provenance for PTO values. The table maps a bound VarNode to its RHS + // IR; only Bind/Let registration writes it. Queries are pure: they never + // print, emit, allocate names or modify the persistent table. + enum class PtoValueKind { + kScalar, + kVmiData, + kVmiMask, + kVmiPair, + kVoid, + kHandle, + kOtherVector, + kUnknown, + }; + struct PtoValueQueryContext; // query-local Let overrides + active Var set + struct PtoOperandScanner; // one-level immediate-operand classification + PtoValueKind ClassifyPtoValue_(const PrimExpr &expr) const; + PtoValueKind ClassifyPtoValue_(const PrimExpr &expr, + PtoValueQueryContext *context) const; + PtoValueKind ClassifyPtoValueUncached_(const PrimExpr &expr, + PtoValueQueryContext *context) const; + PtoValueKind ClassifyVmiCallResult_(const CallNode *op, + PtoValueQueryContext *context) const; + PtoValueKind ClassifyBufferLoad_(const BufferLoadNode *op) const; + PtoValueKind ClassifyOrdinaryExpr_(const PrimExpr &expr, + PtoValueQueryContext *context) const; + static const char *PtoValueKindName(PtoValueKind kind); + void CheckPtoScalarOperand_(const PrimExpr &expr, const char *consumer) const; + void CheckPtoVmiBridgeOperand_(const CallNode *op, const char *op_name, + PtoValueKind source_kind) const; + // Grouped vbrc also accepts numeric local.var reads through pure aliases. + // The source classification is unchanged; the consumer checks lanes == group. + // This does not infer the read's origin from mutable stores. + bool IsMutableRegisterGroupSource(const PrimExpr &source) const; + bool ResolvesToMutableRegisterRead( + const PrimExpr &expr, + std::unordered_map *local_defs, + std::unordered_set *active) const; bool IsVmiLocalRegisterBuffer(const BufferNode *buffer) const; void CheckVmiLocalRegisterIndex(const BufferNode *buffer, const PrimExpr &index) const; @@ -246,6 +283,9 @@ class CodeGenTileLangPTO final : public CodeGenTileLangPY { // VarNode to its (low, high) SSA name pair. std::unordered_map> simd_pair_vars_; + // Persistent per-function SSA definitions used by PtoValueKind queries. + // Registered at every Bind/Let entry and cleared in ResetFunctionState_. + std::unordered_map value_defs_; std::vector gemm_emit_contexts_; std::unordered_map gemm_emit_context_by_call_; diff --git a/testing/ascend/language/pto/test_pto_vmi_value_provenance.py b/testing/ascend/language/pto/test_pto_vmi_value_provenance.py new file mode 100644 index 0000000000..d4512823d6 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_vmi_value_provenance.py @@ -0,0 +1,854 @@ +"""PTO VMI value provenance and scalar-only consumer boundaries. + +The codegen records every ``Bind``/``Let`` definition and classifies value +origins (scalar, VMI data/mask/pair/void, handle, other vector, unknown) with a +pure query that never prints, emits, or mutates the persistent definition +table. These tests pin: + +* lane-1 VMI data stays VMI data through aliases and ``vbrc(group=1)``: a VMI + data source is passed to ``vbrc`` directly, without the signless integer + bridge (no ``pto.cast``/``vinterpret_cast`` wrapper); +* every registered VMI producer has an explicit return contract + (data/mask/pair/void, with ``pair_get`` extracting lane-1 data from a pair); +* unknown VMI origins stay unknown through aliases and are diagnosed at + scalar consumers instead of ever defaulting to scalar; +* true scalar integer bridges (literal/dynamic, width/signedness) are + preserved; +* lexical ``Let`` scoping, same-name/different-Var identity, per-function + definition-table reset, definition-cycle detection, and linear-time + classification of deep scalar expressions; +* grouped ``vbrc`` accepts numeric local.var reads without changing their + classification. + +Direct codegen (``target.build.tilelang_pto_without_compile`` + +``inspect_source``) is used where the standard lowering pipeline would inline +the alias/``Let`` structure under test. Fatal diagnostics surface as +``tvm.error.InternalError``. Backend compilation evidence is included; +device numerics live in ``test_pto_vmi_value_provenance_run.py``. + +The "unknown VMI op" contract is exercised with a uniquely named test-only op +registered through ``ir.RegisterOp``. It is never added to the product +contract table; the codegen must diagnose its result as unknown. +""" + +from __future__ import annotations + +import re + +import pytest +import torch # noqa: F401 (torch/torch_npu must import before tilelang/ptodsl) +import torch_npu # noqa: F401 + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + +tirx = tvm.tirx +_Op = tvm.ir.Op.get +_BUILD = tvm.ffi.get_global_func("target.build.tilelang_pto_without_compile") +_INTERNAL_ERROR = tvm.error.InternalError + +VMI_DATA_OPS = [ + "vci", + "vbrc", + "vzip", + "vadd", + "vsub", + "vmul", + "vdiv", + "vmax", + "vmin", + "vand", + "vor", + "vxor", + "vshl", + "vshr", + "vabs", + "vneg", + "vrelu", + "vexp", + "vln", + "vsqrt", + "vnot", + "vadds", + "vmuls", + "vmaxs", + "vmins", + "vshls", + "vshrs", + "vsel", + "vselr", + "vcadd", + "vcmax", + "vcmin", + "vcvt", + "vinterpret_cast", + "vexpdif", + "vaxpy", + "vlrelu", + "vprelu", + "vmula", + "vdhist", + "vchist", + "vgather", + "vgatherb", +] +VMI_MASK_OPS = ["create_mask", "vcmp", "vcmps"] +VMI_PAIR_OPS = ["vintlv", "vdintlv", "vunzip", "vmull"] +VMI_VOID_OPS = ["vstore", "vscatter"] +UNKNOWN_VMI_OP_NAME = "tl.vmi.provenance_test_unknown_result" + + +# --------------------------------------------------------------------------- +# Lowering / direct codegen helpers +# --------------------------------------------------------------------------- +def _target(): + return normalize_pto_target("pto") + + +def _lowered(func): + target = _target() + with target: + return _lower(func, target=target) + + +def _kernel_source(func): + return _lowered(func).kernel_source + + +def _compile(func): + return tilelang.compile(func, target="pto", out_idx=-1) + + +def _source_of(module): + target = _target() + with target: + return _BUILD(module, target).inspect_source() + + +def _raw_source(body, params=()): + """Codegen a raw PrimFunc without the TileLang lowering pipeline.""" + func = tirx.PrimFunc(list(params), body).with_attr("global_symbol", "main") + return _source_of(tvm.IRModule({"main": func})) + + +def _codegen_device_fn(artifact, fn, global_symbol="main"): + gvar = tvm.ir.GlobalVar(global_symbol) + fn = fn.with_attr("global_symbol", global_symbol) + module = tvm.IRModule({gvar: fn}) + target = _target() + with target: + module = tilelang.transform.LowerIntrin()(module) + module = tilelang.transform.HoistBroadcastValues()(module) + return _source_of(module) + + +def _lines(source, needle): + return [line.strip() for line in source.splitlines() if needle in line] + + +def _one(source, needle): + hits = _lines(source, needle) + assert len(hits) == 1, f"expected one line containing {needle!r}, got {hits!r}" + return hits[0] + + +def _expr(line): + """Strip an optional ``name = `` SSA assignment prefix from a source line.""" + match = re.match(r"^[A-Za-z_]\w* = (.*)$", line) + return match.group(1) if match else line + + +def _expect_error(pattern, build): + with pytest.raises(_INTERNAL_ERROR, match=re.escape(pattern)): + build() + + +_UNKNOWN_VMI_OP = None + + +def _unknown_vmi_op(): + """Register the test-only VMI op once with no provenance contract. + + The product table is deliberately not extended; this op only proves that a + ``tl.vmi.*`` result without an explicit return contract stays unknown. The + registration entry point and attribute write fail loudly if unavailable. + """ + global _UNKNOWN_VMI_OP + if _UNKNOWN_VMI_OP is None: + register = tvm.ffi.get_global_func("ir.RegisterOp") + register(UNKNOWN_VMI_OP_NAME, "Test-only VMI op without a contract.") + op = tvm.ir.Op.get(UNKNOWN_VMI_OP_NAME) + # Pure call effect, as required by the codegen's purity check. + op.set_attr("TCallEffectKind", 0) + _UNKNOWN_VMI_OP = op + return _UNKNOWN_VMI_OP + + +# --------------------------------------------------------------------------- +# Real PTO kernels (standard lowering + codegen or full compile) +# --------------------------------------------------------------------------- +def _reduce_broadcast(dtype, reduce_name): + reduce = getattr(T.vmi, reduce_name) + + @T.prim_func + def main(A: T.Tensor((128,), dtype), C: T.Tensor((128,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((256,), dtype) + out = T.alloc_shared((128,), dtype) + T.copy(A, work[0:128]) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + r = reduce(x, mask) + y = T.vmi.vbrc(r, size=128, group=1) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def _scalar_bridge(dtype, literal): + @T.prim_func + def main(A: T.Tensor((128,), dtype), scale: T.int32, C: T.Tensor((128,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((128,), dtype) + around = T.alloc_shared((128,), dtype) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + base = T.int32(7) if literal else scale + y = T.vmi.vbrc(base, size=128) + z = T.vmi.vci(base, size=128) + T.vmi.vstore(T.vmi.vadd(y, z, mask), around[0], mask) + T.copy(around, C) + + return main + + +def _reinterpret_kernel(): + @T.prim_func + def main(A: T.Tensor((128,), "uint32"), C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "uint32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + y = T.reinterpret("int32x128", x) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def _vci_from_vmi(): + @T.prim_func + def main(A: T.Tensor((128,), "int32"), C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "int32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + r = T.vmi.vcmax(x, mask) + y = T.vmi.vci(r, size=128) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def _scalar_load_after_vmi_store(): + @T.prim_func + def main(A: T.Tensor((128,), "int32"), C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "int32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + T.vmi.vstore(x, work[0], mask) + s = work[0] + y = T.vmi.vbrc(s, size=128) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def _local_var_vmi_store(value_is_vmi): + @T.prim_func + def main(A: T.Tensor((128,), "int32"), C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "int32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + slot = T.alloc_var("int32") + value = T.vmi.vcmax(x, mask) if value_is_vmi else work[0] + slot = value + T.vmi.vstore(T.vmi.vbrc(slot, size=128), out[0], mask) + T.copy(out, C) + + return main + + +def _mutable_register(dtype, lanes, group): + @T.prim_func + def main(A: T.Tensor((128,), dtype), C: T.Tensor((128,), dtype)): + with T.Kernel(1): + src = T.alloc_shared((128,), dtype) + dst = T.alloc_shared((128,), dtype) + T.copy(A, src) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + slot = T.vmi.alloc_var(dtype, size=lanes) + slot = T.vmi.vload(src[0], size=lanes) + value = T.vmi.vbrc(slot, size=128, group=group) + T.vmi.vstore(value, dst[0], mask) + T.copy(dst, C) + + return main + + +# --------------------------------------------------------------------------- +# Raw producer/consumer helpers (direct codegen contracts) +# --------------------------------------------------------------------------- +def _vbrc(arg, group=1): + annotations = {"size": 128} + if group is not None: + annotations["group"] = group + return tirx.call_intrin("int32x128", _Op("tl.vmi.vbrc"), arg, annotations=annotations) + + +def _vmi_data_producer(op_name): + op = _Op(f"tl.vmi.{op_name}") + if op_name in ("vci", "vbrc"): + return tirx.call_intrin("int32", op, tirx.IntImm("int32", 0)), () + if op_name in ("vgather", "vgatherb"): + ptr = tirx.Var("ptr", "handle") + offsets = tirx.Var("offsets", "int32x128") + mask = tirx.Var("mask", "boolx128") + call = tirx.call_intrin("int32", op, ptr, tirx.IntImm("int32", 0), offsets, mask) + return call, (ptr, offsets, mask) + return tirx.call_intrin("int32", op), () + + +def _device_fn(func): + artifact = _lowered(func) + assert len(artifact.device_mod.functions) == 1 + gvar = next(iter(artifact.device_mod.functions.keys())) + return artifact, gvar, artifact.device_mod.functions[gvar] + + +def _find_vbrc_call(fn): + found = [] + + def visit(node): + if isinstance(node, tirx.Call) and getattr(node.op, "name", "") == "tl.vmi.vbrc": + found.append(node) + + tirx.stmt_functor.post_order_visit(fn.body, visit) + assert len(found) == 1, f"expected one vbrc call, got {len(found)}" + return found[0] + + +def _find_binding(fn, var): + found = [] + + def visit(node): + if isinstance(node, tirx.Bind) and node.var.same_as(var): + found.append(node) + + tirx.stmt_functor.post_order_visit(fn.body, visit) + assert len(found) == 1, f"expected one binding for {var}, got {len(found)}" + return found[0] + + +def _transform(fn, preorder): + body = tirx.stmt_functor.ir_transform(fn.body, preorder, None) + return fn.with_body(body) + + +def _replace_vbrc_source(fn, make_source, drop_producer=False): + """Replace the vbrc source with ``make_source(producer_bind, source_var)``.""" + call = _find_vbrc_call(fn) + source_var = call.args[0] + producer_bind = _find_binding(fn, source_var) + + def preorder(node): + if isinstance(node, tirx.Bind): + if node.var.same_as(source_var): + if drop_producer: + return tirx.Evaluate(tirx.IntImm("int32", 0)) + return node + if ( + isinstance(node.value, tirx.Call) + and getattr(node.value.op, "name", "") == "tl.vmi.vbrc" + and node.value.args[0].same_as(source_var) + ): + value = node.value + new_call = tirx.Call( + value.dtype, + value.op, + [make_source(producer_bind, source_var)] + list(value.args[1:]), + value.annotations, + value.span, + ) + return tirx.Bind(node.var, new_call) + return None + + return _transform(fn, preorder) + + +def _alias_chain_source(fn): + call = _find_vbrc_call(fn) + source_var = call.args[0] + alias_a = tirx.Var("alias_a", source_var.dtype) + alias_b = tirx.Var("alias_b", source_var.dtype) + + def preorder(node): + if isinstance(node, tirx.Bind): + if node.var.same_as(source_var): + return tirx.SeqStmt([node, tirx.Bind(alias_a, node.var), tirx.Bind(alias_b, alias_a)]) + if ( + isinstance(node.value, tirx.Call) + and getattr(node.value.op, "name", "") == "tl.vmi.vbrc" + and node.value.args[0].same_as(source_var) + ): + value = node.value + new_call = tirx.Call( + value.dtype, + value.op, + [alias_b] + list(value.args[1:]), + value.annotations, + value.span, + ) + return tirx.Bind(node.var, new_call) + return None + + return _transform(fn, preorder) + + +def _let_source(fn, nested): + def make_source(producer_bind, source_var): + let_var = tirx.Var("let_t", source_var.dtype) + body = tirx.Let(let_var, producer_bind.value, let_var) + if nested: + inner_var = tirx.Var("let_u", source_var.dtype) + body = tirx.Let(let_var, producer_bind.value, tirx.Let(inner_var, let_var, inner_var)) + return body + + return _replace_vbrc_source(fn, make_source, drop_producer=True) + + +def _direct_call_source(fn): + return _replace_vbrc_source(fn, lambda producer_bind, source_var: producer_bind.value, drop_producer=True) + + +# --------------------------------------------------------------------------- +# A. Core lane-1 VMI data -> vbrc(group=1), scalar bridges, reinterpret +# --------------------------------------------------------------------------- +def _assert_direct_vbrc(source, producer_name): + producer_line = _one(source, f"pto.vmi.{producer_name}(") + bound = producer_line.split("=", 1)[0].strip() + vbrc_line = _one(source, "pto.vmi.vbrc(") + assert _expr(vbrc_line) == f"pto.vmi.vbrc({bound}, group=1, size=128)" + assert source.count(f"pto.vmi.{producer_name}(") == 1 + + +def test_int32_vcmax_vbrc_uses_vmi_source_directly(): + source = _kernel_source(_reduce_broadcast("int32", "vcmax")) + _assert_direct_vbrc(source, "vcmax") + + +def test_uint32_vcmin_vbrc_uses_vmi_source_directly(): + source = _kernel_source(_reduce_broadcast("uint32", "vcmin")) + _assert_direct_vbrc(source, "vcmin") + + +def test_float32_vcmax_vbrc_control_uses_vmi_source_directly(): + source = _kernel_source(_reduce_broadcast("float32", "vcmax")) + _assert_direct_vbrc(source, "vcmax") + + +@pytest.mark.parametrize( + ("dtype", "reduce_name"), + [("int32", "vcmax"), ("uint32", "vcmin"), ("float32", "vcmax")], +) +def test_reduce_broadcast_compiles_through_backend(dtype, reduce_name): + _compile(_reduce_broadcast(dtype, reduce_name)) + + +def test_scalar_literal_bridge_is_preserved(): + source = _kernel_source(_scalar_bridge("int32", literal=True)) + assert _expr(_one(source, "pto.vmi.vbrc(pto.si32(7), size=128)")) == ("pto.vmi.vbrc(pto.si32(7), size=128)") + assert _expr(_one(source, "pto.vmi.vci(")).startswith("pto.vmi.vinterpret_cast(pto.vmi.vci(pto.i32(7), size=128), to_dtype=pto.si32)") + + +def test_scalar_dynamic_bridge_is_preserved(): + source = _kernel_source(_scalar_bridge("int32", literal=False)) + # Dynamic integer scalar sources take the signless bridge and recover the + # authored signedness on the VMI result. + assert _expr(_one(source, "pto.vmi.vbrc(")) == ( + "pto.vmi.vinterpret_cast(pto.vmi.vbrc(pto.cast(base, pto.i32), size=128), to_dtype=pto.si32)" + ) + assert _expr(_one(source, "pto.vmi.vci(")) == ( + "pto.vmi.vinterpret_cast(pto.vmi.vci(pto.cast(base, pto.i32), size=128), to_dtype=pto.si32)" + ) + + +@pytest.mark.parametrize("literal", [True, False]) +def test_scalar_bridge_compiles_through_backend(literal): + _compile(_scalar_bridge("int32", literal=literal)) + + +def test_vmi_reinterpret_uses_vinterpret_cast(): + source = _kernel_source(_reinterpret_kernel()) + assert "pto.vmi.vinterpret_cast(x, to_dtype=pto.si32)" in _one(source, "pto.vmi.vstore(") + assert "vbitcast" not in source + + +# --------------------------------------------------------------------------- +# B. Grouped-vbrc mutable local.var register compatibility +# --------------------------------------------------------------------------- +def test_mutable_register_grouped_vbrc_is_preserved(): + # Grouped vbrc accepts this numeric local.var read without classifying it + # as VMI data. Its classification stays other-vector, and the source is + # emitted directly without a scalar cast. + source = _kernel_source(_mutable_register("float32", 4, 4)) + vbrc_line = _one(source, "pto.vmi.vbrc(") + assert "pto.vmi.vbrc(slot, group=4, size=128)" in vbrc_line + assert "pto.cast(slot" not in vbrc_line + assert "slot = pto.vmi.vreg(4, pto.f32)" in source + assert "slot = pto.vmi.vload(" in source + + +# --------------------------------------------------------------------------- +# C. Aliases, Let scoping, Var identity, per-function reset, cycles +# --------------------------------------------------------------------------- +def test_multi_level_alias_chain_keeps_vmi_data(): + artifact, _, fn = _device_fn(_reduce_broadcast("int32", "vcmax")) + source = _codegen_device_fn(artifact, _alias_chain_source(fn)) + assert "alias_a = r" in source or re.search(r"alias_a = \w+", source) + assert re.search(r"alias_b = alias_a", source) + assert _expr(_one(source, "pto.vmi.vbrc(")) == ("pto.vmi.vbrc(alias_b, group=1, size=128)") + assert source.count("pto.vmi.vcmax(") == 1 + + +def test_direct_nested_producer_call_keeps_vmi_data(): + artifact, _, fn = _device_fn(_reduce_broadcast("int32", "vcmax")) + source = _codegen_device_fn(artifact, _direct_call_source(fn)) + vbrc_line = _one(source, "pto.vmi.vbrc(") + assert "pto.vmi.vbrc(pto.vmi.vcmax(" in vbrc_line + assert vbrc_line.endswith(", group=1, size=128)") + assert source.count("pto.vmi.vcmax(") == 1 + + +@pytest.mark.parametrize("nested", [False, True]) +def test_let_source_is_query_local_and_producer_runs_once(nested): + artifact, _, fn = _device_fn(_reduce_broadcast("int32", "vcmax")) + source = _codegen_device_fn(artifact, _let_source(fn, nested=nested)) + vbrc_line = _one(source, "pto.vmi.vbrc(") + if nested: + assert "(lambda let_t: (lambda let_u: let_u)(let_t))(pto.vmi.vcmax(" in vbrc_line + else: + assert "(lambda let_t: let_t)(pto.vmi.vcmax(" in vbrc_line + assert vbrc_line.endswith(", group=1, size=128)") + assert "pto.cast(" not in vbrc_line + assert source.count("pto.vmi.vcmax(") == 1 + + +def test_same_named_different_vars_are_not_conflated(): + artifact, _, fn = _device_fn(_reduce_broadcast("int32", "vcmax")) + call = _find_vbrc_call(fn) + source_var = call.args[0] + same_name = tirx.Var(source_var.name, source_var.dtype) + + def preorder(node): + if isinstance(node, tirx.Bind) and isinstance(node.value, tirx.Call) and getattr(node.value.op, "name", "") == "tl.vmi.vbrc": + second = tirx.call_intrin( + "int32x128", + _Op("tl.vmi.vbrc"), + same_name, + annotations={"size": 128}, + ) + return tirx.SeqStmt( + [ + node, + tirx.Bind(same_name, tirx.IntImm("int32", 5)), + tirx.Evaluate(second), + ] + ) + return None + + source = _codegen_device_fn(artifact, _transform(fn, preorder)) + vbrc_lines = _lines(source, "pto.vmi.vbrc(") + assert len(vbrc_lines) == 2 + assert re.search(r"pto\.vmi\.vbrc\(\w+, group=1, size=128\)", vbrc_lines[0]) + assert "group" not in vbrc_lines[1] + # The same-named scalar Var is bridged independently of the VMI data Var, + # proving the table is keyed by Var identity, not name. The printer + # disambiguates the duplicate Python name (r vs r_1). + assert re.search( + rf"pto\.vmi\.vbrc\(pto\.cast\({source_var.name}_\d+, pto\.i32\), size=128\)", + vbrc_lines[1], + ), vbrc_lines[1] + + +def _two_function_reset_module(order): + """One codegen run with the same Var defined in one function and used as a + true scalar parameter in another. Function names select the generation + order (the module map is name-ordered), so both orders are covered without + relying on insertion order.""" + artifact_a, _, fn_a = _device_fn(_reduce_broadcast("int32", "vcmax")) + shared = tirx.Var("shared_reset", "int32") + + call_a = _find_vbrc_call(fn_a) + old_a = call_a.args[0] + + def preorder_a(node): + if isinstance(node, tirx.Bind) and node.var.same_as(old_a): + return tirx.Bind(shared, node.value) + if isinstance(node, tirx.Bind) and isinstance(node.value, tirx.Call): + value = node.value + if getattr(value.op, "name", "") == "tl.vmi.vbrc" and value.args[0].same_as(old_a): + return tirx.Bind( + node.var, + tirx.Call( + value.dtype, + value.op, + [shared] + list(value.args[1:]), + value.annotations, + value.span, + ), + ) + return None + + fn_a = fn_a.with_body(tirx.stmt_functor.ir_transform(fn_a.body, preorder_a, None)) + + # The SAME Var object is an ordinary scalar parameter here: no Bind, no VMI + # producer. A leaked definition table would misclassify it as VMI data. + scalar_vbrc = tirx.call_intrin("int32x128", _Op("tl.vmi.vbrc"), shared, annotations={"size": 128}) + fn_b = tirx.PrimFunc([shared], tirx.Evaluate(scalar_vbrc)) + + if order == "vmi-first": + vmi_name, scalar_name = "a_vmi_fn", "b_scalar_fn" + else: + vmi_name, scalar_name = "b_vmi_fn", "a_scalar_fn" + target = _target() + fn_a = fn_a.with_attr("global_symbol", vmi_name).with_attr("target", target) + fn_b = fn_b.with_attr("global_symbol", scalar_name).with_attr("target", target) + vmi_entry = (tvm.ir.GlobalVar(vmi_name), fn_a) + scalar_entry = (tvm.ir.GlobalVar(scalar_name), fn_b) + entries = [vmi_entry, scalar_entry] + if order == "scalar-first": + entries.reverse() + module = tvm.IRModule(dict(entries)) + with target: + module = tilelang.transform.LowerIntrin()(module) + module = tilelang.transform.HoistBroadcastValues()(module) + return _source_of(module), vmi_name, scalar_name + + +def _function_section(source, name): + start = source.index(f"def {name}(") + end = source.find("def ", start + 1) + return source[start : end if end != -1 else len(source)] + + +def test_same_var_across_two_functions_resets_definition_table(): + observed_orders = [] + for order in ("vmi-first", "scalar-first"): + source, vmi_name, scalar_name = _two_function_reset_module(order) + section_vmi = _function_section(source, vmi_name) + section_scalar = _function_section(source, scalar_name) + assert "pto.vmi.vbrc(shared_reset, group=1, size=128)" in section_vmi + assert "pto.cast(shared_reset" not in section_vmi + assert "pto.vmi.vbrc(pto.cast(shared_reset, pto.i32), size=128)" in section_scalar + assert "pto.vmi.vbrc(shared_reset, group=1" not in section_scalar + observed_orders.append("vmi-first" if source.index(f"def {vmi_name}(") < source.index(f"def {scalar_name}(") else "scalar-first") + assert sorted(observed_orders) == ["scalar-first", "vmi-first"] + + +def test_definition_cycle_is_detected(): + x = tirx.Var("cycle_x", "int32") + body = tirx.SeqStmt( + [ + tirx.Bind(x, tirx.Add(x, tirx.IntImm("int32", 1))), + tirx.Evaluate(_vbrc(x, group=None)), + ] + ) + _expect_error("PTO SSA definition cycle detected", lambda: _raw_source(body)) + + +# --------------------------------------------------------------------------- +# D. Unknown-origin propagation and diagnosis +# --------------------------------------------------------------------------- +def test_registered_unknown_vmi_op_through_alias_stays_unknown(): + unknown = tirx.Var("unknown_x", "int32") + alias = tirx.Var("unknown_alias", "int32") + body = tirx.SeqStmt( + [ + tirx.Bind(unknown, tirx.call_intrin("int32", _unknown_vmi_op())), + tirx.Bind(alias, unknown), + tirx.Evaluate(_vbrc(alias, group=None)), + ] + ) + _expect_error("cannot broadcast unknown (int32, lanes=1)", lambda: _raw_source(body)) + + +def test_unknown_vmi_origin_diagnosed_at_vci(): + x = tirx.Var("x", "int32") + alias = tirx.Var("alias", "int32") + body = tirx.SeqStmt( + [ + tirx.Bind(x, tirx.call_intrin("int32", _Op("tl.vmi.vcmax"))), + tirx.Bind(alias, x), + tirx.Evaluate( + tirx.call_intrin( + "int32x128", + _Op("tl.vmi.vci"), + tirx.Cast("float32", alias), + annotations={"size": 128}, + ) + ), + ] + ) + _expect_error("tl.vmi.vci base expects a scalar; got unknown", lambda: _raw_source(body)) + + +# --------------------------------------------------------------------------- +# E. Registered producer return-contract matrix +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("op_name", VMI_DATA_OPS) +def test_registered_vmi_data_producer_never_enters_scalar_bridge(op_name): + producer, params = _vmi_data_producer(op_name) + if op_name in ("vgather", "vgatherb"): + # The gather printer requires a real access-pointer operand and the + # function-signature printer cannot emit a synthetic handle parameter. + # The return contract is asserted through pair_get's source validation, + # which classifies the gather result before printing it (free variables + # are only rendered by the TIR printer in the diagnostic): it must be + # VMI data, not scalar/unknown. + get = tirx.call_intrin("int32", _Op("tl.vmi.pair_get"), producer, tirx.IntImm("int32", 0)) + _expect_error( + "tl.vmi.pair_get source must be a VMI pair; got VMI data", + lambda: _raw_source(tirx.Evaluate(get)), + ) + return + source = _raw_source(tirx.Evaluate(_vbrc(producer)), params) + vbrc_line = _one(source, "pto.vmi.vbrc(") + assert "pto.cast(" not in vbrc_line + assert vbrc_line.endswith(", group=1, size=128)") + + +def test_vload_lane1_is_vmi_data(): + producer = tirx.call_intrin("int32", _Op("tl.vmi.vload")) + source = _raw_source(tirx.Evaluate(_vbrc(producer))) + assert "pto.cast(" not in _one(source, "pto.vmi.vbrc(") + + +@pytest.mark.parametrize("op_name", VMI_MASK_OPS) +def test_vmi_mask_producer_is_not_broadcast_as_data(op_name): + producer = tirx.call_intrin("boolx128", _Op(f"tl.vmi.{op_name}")) + _expect_error( + "tl.vmi.vbrc with group=... requires a VMI vreg; got VMI mask", + lambda: _raw_source(tirx.Evaluate(_vbrc(producer))), + ) + + +@pytest.mark.parametrize("op_name", VMI_PAIR_OPS) +def test_vmi_pair_producer_is_not_broadcast_as_data(op_name): + producer = tirx.call_intrin("int32x128", _Op(f"tl.vmi.{op_name}")) + _expect_error( + "tl.vmi.vbrc with group=... requires a VMI vreg; got VMI pair", + lambda: _raw_source(tirx.Evaluate(_vbrc(producer))), + ) + + +@pytest.mark.parametrize("op_name", VMI_VOID_OPS) +def test_vmi_void_producer_is_not_broadcast_as_data(op_name): + producer = tirx.call_intrin("void", _Op(f"tl.vmi.{op_name}")) + _expect_error( + "tl.vmi.vbrc with group=... requires a VMI vreg; got void", + lambda: _raw_source(tirx.Evaluate(_vbrc(producer))), + ) + + +def test_pair_get_returns_data_from_pair(): + pair = tirx.call_intrin("int32x128", _Op("tl.vmi.vintlv")) + get = tirx.call_intrin("int32", _Op("tl.vmi.pair_get"), pair, tirx.IntImm("int32", 0)) + source = _raw_source(tirx.Evaluate(_vbrc(get))) + assert _expr(_one(source, "pto.vmi.vbrc(")) == "pto.vmi.vbrc((pto.vmi.vintlv())[0], group=1, size=128)" + + +# --------------------------------------------------------------------------- +# F. Scalar-only consumers +# --------------------------------------------------------------------------- +def test_vci_rejects_vmi_data_source(): + _expect_error( + "tl.vmi.vci base expects a scalar; got VMI data", + lambda: _kernel_source(_vci_from_vmi()), + ) + + +def test_vbrc_group_rejects_true_scalar_source(): + _expect_error( + "tl.vmi.vbrc with group=... requires a VMI vreg; got scalar", + lambda: _raw_source(tirx.Evaluate(_vbrc(tirx.IntImm("int32", 0), group=1))), + ) + + +def test_ordinary_scalar_load_after_vmi_store_stays_scalar(): + source = _kernel_source(_scalar_load_after_vmi_store()) + vbrc_line = _one(source, "pto.vmi.vbrc(") + assert "pto.cast(" in vbrc_line + assert "group" not in vbrc_line + + +def test_local_var_store_rejects_vmi_data(): + _expect_error( + "PTO local.var integer store expects a scalar; got VMI data", + lambda: _kernel_source(_local_var_vmi_store(value_is_vmi=True)), + ) + + +def test_local_var_store_keeps_scalar_value(): + source = _kernel_source(_local_var_vmi_store(value_is_vmi=False)) + assert "slot = pto.const(0, dtype=pto.si32)" in source + assert re.search(r"slot = pto\.cast\(\w+, pto\.si32\)", source) + assert "pto.cast(slot, pto.i32)" in _one(source, "pto.vmi.vbrc(") + + +def test_select_branch_rejects_vmi_data(): + data = tirx.Var("data", "int32") + cond = tirx.Var("cond", "bool") + body = tirx.SeqStmt( + [ + tirx.Bind(data, tirx.call_intrin("int32", _Op("tl.vmi.vcmax"))), + tirx.Evaluate(tirx.Select(cond, data, tirx.IntImm("int32", 0))), + ] + ) + _expect_error("PTO select branch expects a scalar; got VMI data", lambda: _raw_source(body, (cond,))) + + +# --------------------------------------------------------------------------- +# G. Deep scalar expressions stay bounded +# --------------------------------------------------------------------------- +def test_deep_scalar_expression_chain_is_classified_linearly(): + x = tirx.Var("x", "int32") + expr = x + for i in range(200): + expr = tirx.Add(expr, tirx.IntImm("int32", i + 1)) + source = _raw_source(tirx.Evaluate(tirx.call_intrin("int32", _Op("tl.vmi.vci"), expr)), (x,)) + # The scalar bridge proves one classification of the 200-deep chain. + assert "pto.vmi.vinterpret_cast(pto.vmi.vci(pto.cast(" in source diff --git a/testing/ascend/language/pto/test_pto_vmi_value_provenance_run.py b/testing/ascend/language/pto/test_pto_vmi_value_provenance_run.py new file mode 100644 index 0000000000..0dbc47cfed --- /dev/null +++ b/testing/ascend/language/pto/test_pto_vmi_value_provenance_run.py @@ -0,0 +1,246 @@ +"""Device acceptance: PTO VMI value provenance and scalar bridge boundaries. + +Each case compiles through the real TileLang -> PTODSL -> PTOAS VPTO -> +Bisheng pipeline and executes on the NPU. References are computed on the CPU +from Python integers/floats, independently of the generated kernel. + +Coverage: +* int32/uint32 ``vcmax``/``vcmin`` -> ``vbrc(group=1)``: signed negatives, + exact integer boundaries, unsigned high-bit values (> ``INT32_MAX``), and + both reduction directions. 20 independent trials per integer case. +* float32 control: the same reduce/broadcast route for a non-integer VMI + source. +* true scalar integer bridges: compile-time literal and runtime scalar + sources through ``vbrc`` + ``vci``. +* ordinary scalar memory read after a VMI store (storage contract, not the + last ``BufferStore`` RHS). +""" + +import os +import re +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import torch_npu # noqa: E402,F401 (import before tilelang/ptodsl) +import tilelang # noqa: E402 +import tilelang.ascend.language as T # noqa: E402 + + +def _require_npu_runtime(): + # These device tests require an available NPU runtime and a successful + # host-to-device copy probe; a missing runtime or probe mismatch fails. + assert hasattr(torch, "npu"), "torch.npu is unavailable in this environment" + assert torch.npu.is_available(), "torch.npu is unavailable in the current environment" + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source), "NPU host->device copy probe mismatch" + + +@pytest.fixture(scope="module") +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +def _to_device(value, device): + if isinstance(value, torch.Tensor): + return value.to(device) + return value + + +def _compile(func): + kernel = tilelang.compile(func, target="pto", out_idx=-1) + return kernel, kernel.get_kernel_source() + + +def _run_trials(kernel, make_case, device, trials): + outputs = [] + for trial in range(trials): + inputs, expected = make_case(trial) + device_inputs = [_to_device(value, device) for value in inputs] + output = kernel(*device_inputs) + torch.npu.synchronize() + got = output.cpu() + assert got.dtype == expected.dtype, f"trial {trial}: {got.dtype} != {expected.dtype}" + assert torch.equal(got, expected), f"trial {trial}: numeric mismatch\n got={got.tolist()[:16]}\n exp={expected.tolist()[:16]}" + outputs.append(got) + return outputs + + +# --------------------------------------------------------------------------- +# Reduce -> vbrc(group=1): lane-1 VMI data must stay VMI data +# --------------------------------------------------------------------------- +def _reduce_broadcast(dtype, reduce_name): + reduce = getattr(T.vmi, reduce_name) + + @T.prim_func + def main(A: T.Tensor((128,), dtype), C: T.Tensor((128,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((256,), dtype) + out = T.alloc_shared((128,), dtype) + T.copy(A, work[0:128]) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + r = reduce(x, mask) + y = T.vmi.vbrc(r, size=128, group=1) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def _int32_input(trial): + generator = torch.Generator().manual_seed(20261017 + trial) + values = torch.randint(-(2**31), 2**31 - 1, (128,), generator=generator).to(torch.int32) + values[0] = -(2**31) + values[1] = 2**31 - 1 + values[2] = -1 + values[3] = 0 + values[4] = 1 + values[5] = -(2**31) + 1 if trial % 2 else 2**31 - 2 + return values + + +def _uint32_input(trial): + generator = torch.Generator().manual_seed(20261017 + trial) + values = torch.randint(0, 2**31 - 1, (128,), generator=generator).to(torch.uint32) + values[0] = 2**32 - 1 + values[1] = 2**31 # INT32_MAX + 1, high bit set + values[2] = 2**31 - 1 + values[3] = 0 + values[4] = 1 + values[5] = 2**32 - 2 if trial % 2 else 2**31 + 1 + return values + + +def _float32_input(trial): + generator = torch.Generator().manual_seed(20261017 + trial) + values = torch.randn(128, generator=generator) * 1000.0 + values[0] = 3.0e38 + values[1] = -3.0e38 + values[2] = 0.0 + values[3] = -0.0 + return values + + +INPUTS = {"int32": _int32_input, "uint32": _uint32_input} +REDUCE_OPS = {"vcmax": max, "vcmin": min} +INTEGER_REDUCE_CASES = [ + ("int32", "vcmax"), + ("int32", "vcmin"), + ("uint32", "vcmax"), + ("uint32", "vcmin"), +] + + +@pytest.mark.parametrize(("dtype", "reduce_name"), INTEGER_REDUCE_CASES) +def test_integer_reduce_broadcast_is_bit_exact(npu, dtype, reduce_name): + kernel, source = _compile(_reduce_broadcast(dtype, reduce_name)) + vbrc_line = next(line.strip() for line in source.splitlines() if "pto.vmi.vbrc(" in line) + assert "pto.cast(" not in vbrc_line, vbrc_line + assert "group=1" in vbrc_line, vbrc_line + reduce = REDUCE_OPS[reduce_name] + + def make_case(trial): + a = INPUTS[dtype](trial) + expected_value = reduce(int(v) for v in a.tolist()) + expected = torch.full((128,), expected_value, dtype=a.dtype) + return [a], expected + + _run_trials(kernel, make_case, npu, trials=20) + + +def test_float32_reduce_broadcast_control(npu): + kernel, source = _compile(_reduce_broadcast("float32", "vcmax")) + vbrc_line = next(line.strip() for line in source.splitlines() if "pto.vmi.vbrc(" in line) + assert "pto.cast(" not in vbrc_line, vbrc_line + + def make_case(trial): + a = _float32_input(trial) + expected = torch.full((128,), float(a.max().item()), dtype=torch.float32) + return [a], expected + + _run_trials(kernel, make_case, npu, trials=5) + + +# --------------------------------------------------------------------------- +# True scalar integer bridges (literal + runtime scalar) +# --------------------------------------------------------------------------- +def _scalar_bridge(literal): + @T.prim_func + def main(A: T.Tensor((128,), "int32"), scale: T.int32, C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "int32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + base = T.int32(7) if literal else scale + y = T.vmi.vbrc(base, size=128) + z = T.vmi.vci(base, size=128) + T.vmi.vstore(T.vmi.vadd(y, z, mask), out[0], mask) + T.copy(out, C) + + return main + + +@pytest.mark.parametrize("literal", [True, False], ids=["literal", "runtime"]) +def test_scalar_integer_bridge_device(npu, literal): + kernel, source = _compile(_scalar_bridge(literal)) + if literal: + assert "pto.vmi.vbrc(pto.si32(7), size=128)" in source, source + else: + assert "pto.cast(base, pto.i32)" in source, source + + def make_case(trial): + a = torch.arange(128, dtype=torch.int32) + trial + scale = 7 if literal else 5 + trial * 3 + expected = (2 * scale + torch.arange(128, dtype=torch.int32)).to(torch.int32) + return [a, scale], expected + + _run_trials(kernel, make_case, npu, trials=5) + + +# --------------------------------------------------------------------------- +# Ordinary scalar memory read after a VMI store +# --------------------------------------------------------------------------- +def _scalar_load_after_vmi_store(): + @T.prim_func + def main(A: T.Tensor((128,), "int32"), C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "int32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + T.vmi.vstore(x, work[0], mask) + scalar = work[0] + y = T.vmi.vbrc(scalar, size=128) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def test_scalar_buffer_load_after_vmi_store_device(npu): + kernel, source = _compile(_scalar_load_after_vmi_store()) + vbrc_line = next(line.strip() for line in source.splitlines() if "pto.vmi.vbrc(" in line) + assert re.search(r"pto\.cast\(scalar_\d+, pto\.i32\)", vbrc_line), vbrc_line + + def make_case(trial): + a = torch.arange(128, dtype=torch.int32) * (trial + 1) - 10 + expected = torch.full((128,), int(a[0]), dtype=torch.int32) + return [a], expected + + _run_trials(kernel, make_case, npu, trials=3) From 59500cc62e10f932433407bfa37bcfc832527628 Mon Sep 17 00:00:00 2001 From: qukelin Date: Sat, 17 Oct 2026 13:55:32 +0800 Subject: [PATCH 07/10] fix(pto): plan scalar and VMI value casts explicitly --- src/ascend/transform/parallel_to_pto_utils.cc | 443 +++++++---- src/ascend/transform/parallel_to_pto_utils.h | 114 ++- .../transform/vectorize_parallel_to_pto.cc | 102 ++- .../transform/verify_parallel_to_pto.cc | 54 +- .../language/pto/cast_plan_context_probe.cc | 148 ++++ .../pto/test_pto_parallel_cast_matrix.py | 702 ++++++++++++++++++ .../pto/test_pto_parallel_cast_matrix_run.py | 400 ++++++++++ 7 files changed, 1722 insertions(+), 241 deletions(-) create mode 100644 testing/ascend/language/pto/cast_plan_context_probe.cc create mode 100644 testing/ascend/language/pto/test_pto_parallel_cast_matrix.py create mode 100644 testing/ascend/language/pto/test_pto_parallel_cast_matrix_run.py diff --git a/src/ascend/transform/parallel_to_pto_utils.cc b/src/ascend/transform/parallel_to_pto_utils.cc index a5fa77faea..a4a65dbddb 100644 --- a/src/ascend/transform/parallel_to_pto_utils.cc +++ b/src/ascend/transform/parallel_to_pto_utils.cc @@ -242,11 +242,10 @@ PtoInverseMapping InvertPtoLaneLayout(const For &loop) { } Layout inverse = inverse_info.first; - // Forward the (chunk, lane) coordinates through the inverse layout to - // obtain the logical coordinate expression. The inverse layout has two - // outputs (the logical coordinate and the replicate coordinate); with - // replicate_size == 1 the second output is a constant zero and the - // design-doc convention is to take the logical part only. + // Forward (chunk, lane) through the inverse layout to obtain the logical + // iteration coordinates followed by the replicate coordinate. With + // replicate_size == 1, the trailing replicate coordinate is zero. + // The branches below use the logical coordinates for the loop dimensionality. Var chunk_var("pto_chunk", loop->loop_var.dtype()); Var lane_var("pto_lane", loop->loop_var.dtype()); Array logical = inverse->Forward({chunk_var, lane_var}); @@ -1281,63 +1280,185 @@ const std::set &Fp8AllowedRoundings() { return kAllowed; } -} // namespace +bool IsVmiLogicalLaneCount(int64_t lanes) { + return lanes == 64 || lanes == 128 || lanes == 256; +} -Optional PlanValueCast(const CastNode *cast, std::string *reason) { - DataType from = cast->value.dtype(); - DataType to = cast->dtype; - auto reject = [&](const std::string &why) { - if (reason != nullptr) { - *reason = why; - } - return Optional(); - }; - std::ostringstream pair_os; - pair_os << "`" << from << " -> " << to << "`"; - const std::string pair = pair_os.str(); - // Dtype legality is a precondition for every action plan, - // including identity — unsupported types (int24/FP64/...) must not slip - // through the same-type fast path. - if (!IsSupportedElementDType(from) || !IsSupportedElementDType(to)) { - return reject(pair + " is outside the supported element dtype set"); +std::string CastContextName(const CastContext &context) { + if (context.value_kind == CastValueKind::kScalar) { + return "scalar"; } + std::ostringstream oss; + oss << "VMI, lanes=" << context.logical_lanes; + return oss.str(); +} - // Attributes are parsed for every Cast, not just the FP8 ones: the - // expression whitelist deliberately does not look inside a cast's - // annotations (attributes are not computation values), so an annotation the - // selected conversion cannot honour has to be rejected here — for every - // direction — or it would be dropped silently. - CastAttributes attrs; - std::string why; - if (!ReadCastAttributes(cast->annotations, &attrs, &why)) { - return reject(pair + ": " + why); +bool IsHalfOrBfloat16(DataType dtype) { + return (dtype.is_float() && dtype.bits() == 16) || dtype.is_bfloat16(); +} + +/*! Identity attribute policy shared by the scalar and VMI branches: no + * conversion exists, so sat=True (the default) is normalized away, while + * sat=False or a non-empty rounding mode would claim a conversion semantic + * that does not exist. An explicit empty round key means "no request". */ +bool ValidateIdentityCastAttributes(const CastAttributes &attrs, + const std::string &pair, std::string *why) { + if (attrs.has_sat && !attrs.sat) { + *why = pair + ": an identity cast does not accept sat=False"; + return false; } - if (attrs.has_rbits) { - return reject(pair + ": the rbits operand (stochastic rounding) is not " - "supported"); + if (attrs.has_round_mode) { + *why = pair + ": an identity cast does not accept a rounding mode"; + return false; } - std::string round_token; - if (attrs.has_round_mode && - !NormalizeCastRounding(attrs.round, &round_token, &why)) { - return reject(pair + ": " + why); + return true; +} + +/*! Integer attribute policy: rounding modes are never accepted; an explicit + * saturation flag is accepted only for narrowing and only as sat=False, which + * the direct NOSAT lowering already honours. */ +bool ValidateIntegerCastAttributes(const CastAttributes &attrs, bool narrowing, + const std::string &pair, std::string *why) { + if (attrs.has_round_mode) { + *why = pair + ": integer conversions do not accept a rounding mode"; + return false; } + if (attrs.has_sat) { + if (!narrowing) { + *why = pair + ": this integer conversion does not accept an explicit " + "saturation flag"; + return false; + } + if (attrs.sat) { + *why = pair + ": integer narrowing accepts only sat=False (NOSAT), not " + "sat=True"; + return false; + } + } + return true; +} +Optional PlanScalarValueCast(const CastNode *cast, + const CastAttributes &attrs, + const std::string &pair, + std::string *reason) { + const DataType from = cast->value.dtype(); + const DataType to = cast->dtype; + auto reject = [&](const std::string &why) { + if (reason != nullptr) { + *reason = why; + } + return Optional(); + }; + ScalarCastPlan keep; + keep.action = ScalarCastAction::kKeepInput; + ScalarCastPlan rebuild; + rebuild.action = ScalarCastAction::kRebuildCast; if (from == to) { - // Identity, FP8 or not: no conversion is emitted, so an explicit sat=True - // is the default and is normalized away; any other explicit attribute - // would claim a conversion semantic that does not exist here. - if (attrs.has_sat && !attrs.sat) { - return reject(pair + ": an identity cast does not accept sat=False"); + std::string why; + if (!ValidateIdentityCastAttributes(attrs, pair, &why)) { + return reject(why); } + return CastPlan{keep}; + } + const bool from_int = from.is_int() || from.is_uint(); + const bool to_int = to.is_int() || to.is_uint(); + if (from_int && to_int) { + std::string why; + if (!ValidateIntegerCastAttributes(attrs, to.bits() < from.bits(), pair, + &why)) { + return reject(why); + } + // Scalar integer casts keep the ordinary pto.cast/trunc lowering; the + // validated sat=False marker is already honoured and is normalized away. + return CastPlan{rebuild}; + } + if (IsSupportedFloat8(from) || IsSupportedFloat8(to)) { + return reject(pair + ": non-identity FP8 conversions have no scalar " + "lowering inside SIMD_VF; use a lane-varying value " + "so the cast takes the VMI vcvt path"); + } + const bool from_float = from.is_float() || from.is_bfloat16(); + const bool to_float = to.is_float() || to.is_bfloat16(); + if (from_float && to_float) { if (attrs.has_round_mode) { - return reject(pair + - ": an identity cast does not accept a rounding mode"); + return reject(pair + ": this conversion does not accept a rounding mode"); + } + const bool from_f32 = from.is_float() && from.bits() == 32; + const bool to_f32 = to.is_float() && to.bits() == 32; + if (!from_f32 && to_f32) { + if (attrs.has_sat) { + return reject(pair + ": this direction does not accept a saturation " + "mode"); + } + return CastPlan{rebuild}; + } + if (from_f32 && !to_f32) { + if (attrs.has_sat && attrs.sat) { + return reject(pair + ": this conversion accepts only sat=False " + "(non-saturating)"); + } + return CastPlan{rebuild}; } - CastPlan identity; - identity.identity = true; - return identity; + return reject(pair + ": cross-family f16 <-> bf16 narrowing has no " + "validated rounding semantics; convert through " + "float32 explicitly"); } + return reject(pair + ": no scalar conversion exists for this dtype pair"); +} +Optional +PlanVectorValueCast(const CastNode *cast, const CastContext &context, + const CastAttributes &attrs, const std::string &round_token, + const std::string &pair, std::string *reason) { + const DataType from = cast->value.dtype(); + const DataType to = cast->dtype; + auto reject = [&](const std::string &why) { + if (reason != nullptr) { + *reason = why; + } + return Optional(); + }; + if (from == to) { + std::string why; + if (!ValidateIdentityCastAttributes(attrs, pair, &why)) { + return reject(why); + } + return CastPlan{VectorCastPlan{}}; + } + const bool from_int = from.is_int() || from.is_uint(); + const bool to_int = to.is_int() || to.is_uint(); + if (from_int && to_int) { + std::string why; + if (!ValidateIntegerCastAttributes(attrs, to.bits() < from.bits(), pair, + &why)) { + return reject(why); + } + VectorCastPlan plan; + if (to.bits() < from.bits()) { + plan.steps.push_back( + {VectorCastOp::kIntegerTruncate, from, to, "", "NOSAT"}); + return CastPlan{plan}; + } + if (to.bits() == from.bits()) { + plan.steps.push_back( + {VectorCastOp::kIntegerReinterpret, from, to, "", ""}); + return CastPlan{plan}; + } + if (from.is_int() == to.is_int()) { + plan.steps.push_back({VectorCastOp::kIntegerExtend, from, to, "", ""}); + return CastPlan{plan}; + } + // Cross-sign widening: extend with the source signedness to the target + // width, then re-view the same-width result as the target signedness. + const DataType intermediate = + from.is_int() ? DataType::Int(to.bits()) : DataType::UInt(to.bits()); + plan.steps.push_back( + {VectorCastOp::kIntegerExtend, from, intermediate, "", ""}); + plan.steps.push_back( + {VectorCastOp::kIntegerReinterpret, intermediate, to, "", ""}); + return CastPlan{plan}; + } const bool from_fp8 = IsSupportedFloat8(from); const bool to_fp8 = IsSupportedFloat8(to); if (from_fp8 || to_fp8) { @@ -1345,131 +1466,153 @@ Optional PlanValueCast(const CastNode *cast, std::string *reason) { return reject(pair + ": conversions between the two FP8 types are not " "supported (they are not a reinterpretation)"); } + VectorCastPlan plan; if (from_fp8 && to.is_float() && to.bits() == 32) { - // Widening: exact, no rounding or saturation involved. An explicit - // hint here cannot be honoured, so it is rejected rather than dropped. - // The direction contract rejects an attribute that is *carried* here, - // even an empty one: widening has no rounding to configure. + // Exact widening: a real conversion with no options. An explicit hint + // here cannot be honoured, so it is rejected rather than dropped. if (attrs.has_round_key) { - return reject(pair + - ": this direction does not accept a rounding mode"); + return reject(pair + ": this direction does not accept a rounding " + "mode"); } if (attrs.has_sat) { - return reject(pair + - ": this direction does not accept a saturation mode"); + return reject(pair + ": this direction does not accept a saturation " + "mode"); } - return CastPlan{}; + plan.steps.push_back( + {VectorCastOp::kNonIntegerConvert, from, to, "", ""}); + return CastPlan{plan}; } if (to_fp8) { - // The FP8 conversion instructions the target exposes are FP32 <-> FP8; - // there is no half -> FP8 conversion, so a float16/bfloat16 source is - // rejected here. This pass never adds a conversion node of its own: a - // user who wants the float32 route writes it explicitly, so the - // intermediate rounding is their expressed choice rather than a silent - // step inserted here. if (!(from.is_float() && from.bits() == 32)) { - return reject(pair + ": the FP8 conversion instructions take a float32 " - "source; convert to float32 explicitly first"); + return reject(pair + ": the FP8 conversion instructions take a " + "float32 source; convert to float32 explicitly " + "first"); } const std::set &allowed = Fp8AllowedRoundings(); - // An empty round selects the documented default (R with SAT below). const std::string token = attrs.has_round_mode ? round_token : "R"; if (allowed.count(token) == 0) { std::ostringstream oss; - oss << "rounding mode " << token; + oss << pair << ": rounding mode " << token; if (attrs.has_round_mode && attrs.round != token) { oss << " (from round=\"" << attrs.round << "\")"; } - oss << " is not supported by the FP8 conversion policy for `" << from - << " -> " << to << "`"; + oss << " is not supported by the FP8 conversion policy"; return reject(oss.str()); } - CastPlan plan; - // The PTO default is R/SAT; pin both explicitly so the generated VMI - // call never depends on PTODSL defaults. - plan.rounding = token; - plan.saturate = (attrs.has_sat && !attrs.sat) ? "NOSAT" : "SAT"; - return plan; - } - return reject(pair + - " is not a supported FP8 conversion (FP8 -> FP16/BF16 and " - "FP8 <-> integer are out of scope; convert through " - "float32 explicitly)"); - } - - // Non-FP8 conversion: the plan is fixed by the TIR semantics of the - // operation, so an explicit annotation must agree with it. `finish` is the - // one place that decides this, so a hint is accepted exactly when the plan - // already does what it asks for and rejected otherwise — never dropped. - auto finish = [&](const CastPlan &plan) -> Optional { + plan.steps.push_back({VectorCastOp::kNonIntegerConvert, from, to, token, + (attrs.has_sat && !attrs.sat) ? "NOSAT" : "SAT"}); + return CastPlan{plan}; + } + return reject(pair + ": FP8 -> FP16/BF16 and FP8 <-> integer are out of " + "scope; convert through float32 explicitly"); + } + const bool from_float = from.is_float() || from.is_bfloat16(); + const bool to_float = to.is_float() || to.is_bfloat16(); + if (from_float && to_float) { if (attrs.has_round_mode) { return reject(pair + ": this conversion does not accept a rounding mode"); } - if (attrs.has_sat && (attrs.sat || plan.saturate != "NOSAT")) { - return reject(pair + ": this conversion does not accept an explicit " - "saturation flag"); + const bool from_f32 = from.is_float() && from.bits() == 32; + const bool to_f32 = to.is_float() && to.bits() == 32; + VectorCastPlan plan; + if (!from_f32 && to_f32) { + if (attrs.has_sat) { + return reject(pair + ": this direction does not accept a saturation " + "mode"); + } + // A real conversion: empty options are not a no-op. + plan.steps.push_back( + {VectorCastOp::kNonIntegerConvert, from, to, "", ""}); + return CastPlan{plan}; + } + if (from_f32 && !to_f32) { + if (attrs.has_sat && attrs.sat) { + return reject(pair + ": this conversion accepts only sat=False " + "(NOSAT)"); + } + // Pin R/NOSAT so the generated vcvt never depends on PTODSL defaults. + plan.steps.push_back( + {VectorCastOp::kNonIntegerConvert, from, to, "R", "NOSAT"}); + return CastPlan{plan}; } - return plan; + return reject(pair + ": cross-family f16 <-> bf16 narrowing has no " + "validated rounding semantics; convert through " + "float32 explicitly"); + } + return reject(pair + ": float <-> integer conversions need per-combination " + "rounding/overflow validation"); +} + +} // namespace + +Optional PlanValueCast(const CastNode *cast, + const CastContext &context, + std::string *reason) { + DataType from = cast->value.dtype(); + DataType to = cast->dtype; + auto reject = [&](const std::string &why) { + if (reason != nullptr) { + *reason = why; + } + return Optional(); }; + std::ostringstream pair_os; + pair_os << "`" << from << " -> " << to << "` (" << CastContextName(context) + << ")"; + const std::string pair = pair_os.str(); + // The explicit kind is authoritative; lanes only validate it. Inferring + // the kind from lanes would turn a lane-1 VMI value into a scalar and a + // scalar request with vector lanes into an implicit VMI fallback. + if (context.value_kind == CastValueKind::kScalar) { + if (context.logical_lanes != 1) { + return reject(pair + + ": a scalar cast context must have exactly one " + "logical lane, got " + + std::to_string(context.logical_lanes)); + } + } else { + if (!IsVmiLogicalLaneCount(context.logical_lanes)) { + return reject(pair + + ": a VMI cast context requires logical lanes " + "64/128/256, got " + + std::to_string(context.logical_lanes)); + } + } + // Dtype legality is a precondition for every action plan, + // including identity — unsupported types (int24/FP64/...) must not slip + // through the same-type fast path. The Parallel value path is + // element-wise, so both sides must be lane-1 elements. + if (from.lanes() != 1 || to.lanes() != 1) { + return reject(pair + ": value casts require lane-1 element operands"); + } + if (!IsSupportedElementDType(from) || !IsSupportedElementDType(to)) { + return reject(pair + " is outside the supported element dtype set"); + } - // Same-width signed/unsigned reinterprets have no validated lowering - // (VMI vcvt requires a bit-width change); reject. - if (from.bits() == to.bits() && from.code() == DataType::kInt && - to.code() == DataType::kUInt) { - return reject("same-width int -> uint is a reinterpretation, which VMI " - "vcvt does not provide"); - } - if (from.bits() == to.bits() && from.code() == DataType::kUInt && - to.code() == DataType::kInt) { - return reject("same-width uint -> int is a reinterpretation, which VMI " - "vcvt does not provide"); - } - // No width re-check here: IsSupportedElementDType at function entry already - // restricts both sides to 8/16/32-bit lane-1 elements (and excludes FP64). - CastPlan plan; - // Every integer combination shares one rule: TVM integer casts truncate on - // narrowing and extend on widening, PTODSL defaults narrowing to saturate, - // and NOSAT restores the truncating semantics. The same-width int/uint - // reinterpretations are rejected above. - if ((from.is_int() || from.is_uint()) && (to.is_int() || to.is_uint())) { - if (to.bits() < from.bits()) { - plan.saturate = "NOSAT"; - } - return finish(plan); - } - if (from.is_float() && (to.is_float() || to.is_bfloat16())) { - if (to.bits() > from.bits() && to.is_float()) { - // f16 -> f32 widening is exact. - return finish(plan); - } - // f32 -> f16/bf16 narrowing: TIR Cast rounds to nearest; pin R so the - // VMI vcvt never depends on PTODSL defaults. NOSAT keeps IEEE - // overflow-to-inf semantics instead of clamping to the largest finite. - if (from.bits() == 32 && to.bits() == 16) { - plan.rounding = "R"; - plan.saturate = "NOSAT"; - return finish(plan); - } - // The only float narrowing left here is the cross-family f16 <-> bf16 - // step; its rounding semantics are not validated. - return reject( - "cross-family f16 <-> bf16 narrowing has no validated rounding " - "semantics; convert through float32 explicitly"); - } - // Bfloat16's is_float() is false, so bf16 -> f32 fell into the - // reject branch below even though it is the same exact widening as - // f16 -> f32 (and the PTO backend supports the vcvt). Allow exactly - // this one bfloat combination; every other stays rejected as before. - if (from.is_bfloat16() && to.is_float() && to.bits() == 32) { - return finish(plan); - } - if (from.is_float() != to.is_float()) { - // float <-> int conversions need per-combination rounding/overflow - // validation against PTOAS; currently rejected. - return reject("float <-> integer conversions need per-combination " - "rounding/overflow validation"); - } - return reject("no proven VMI-equivalent conversion for this dtype pair"); + // Attributes are parsed for every Cast, not just the FP8 ones: the + // expression whitelist deliberately does not look inside a cast's + // annotations (attributes are not computation values), so an annotation the + // selected conversion cannot honour has to be rejected here — for every + // direction — or it would be dropped silently. + CastAttributes attrs; + std::string why; + if (!ReadCastAttributes(cast->annotations, &attrs, &why)) { + return reject(pair + ": " + why); + } + if (attrs.has_rbits) { + return reject(pair + ": the rbits operand (stochastic rounding) is not " + "supported"); + } + std::string round_token; + if (attrs.has_round_mode && + !NormalizeCastRounding(attrs.round, &round_token, &why)) { + return reject(pair + ": " + why); + } + + if (context.value_kind == CastValueKind::kScalar) { + return PlanScalarValueCast(cast, attrs, pair, reason); + } + return PlanVectorValueCast(cast, context, attrs, round_token, pair, reason); } std::optional FindUnsupportedExprNode(const PrimExpr &expr) { diff --git a/src/ascend/transform/parallel_to_pto_utils.h b/src/ascend/transform/parallel_to_pto_utils.h index 453c3df7b2..c410d6cccb 100644 --- a/src/ascend/transform/parallel_to_pto_utils.h +++ b/src/ascend/transform/parallel_to_pto_utils.h @@ -1,7 +1,6 @@ /*! * \file parallel_to_pto_utils.h - * \brief Shared read-only analysis helpers for the PTO Parallel vectorize - * path. See docs/vectorize_parallel_to_pto_design.md. + * \brief Shared read-only analysis helpers for PTO Parallel vectorization. * * VerifyParallelToPTO consumes these helpers to check the input contract; * VectorizeParallelToPTO re-runs the same analysis to build its VMI code. @@ -22,6 +21,8 @@ #include #include #include +#include +#include #include "../../layout/layout.h" #include "../../op/builtin.h" @@ -192,20 +193,71 @@ bool IsSupportedElementDType(DataType dtype); * just that a target name exists. */ bool IsSupportedFloat8(DataType dtype); -/*! Cast conversion descriptor. - * empty() means the combination has no proven-equivalent VMI mapping and - * must be rejected by Verify. */ -struct CastPlan { - /*! Identity conversion (from == to): keep the original value, emit no - * vcvt (VMI rejects width-preserving conversions). */ - bool identity = false; - /*! PTODSL saturate annotation value, e.g. "NOSAT"; empty when the - * conversion must not rely on PTODSL defaults. */ - std::string saturate; - /*! PTODSL rounding annotation value, e.g. "Z"; empty when N/A. */ +/*! Value category of one value-Cast position. The kind is supplied by the + * caller (Verify from AnalyzeLaneUse, Vectorize from the converted input + * shape); the planner never infers it from lanes. */ +enum class CastValueKind { + kScalar, + kVectorVmi, +}; + +/*! Per-Cast planning context. `logical_lanes` is validation information + * only; `value_kind` is the dispatch key. */ +struct CastContext { + CastValueKind value_kind = CastValueKind::kScalar; + int64_t logical_lanes = 1; +}; + +/*! Scalar value-Cast plan: either keep the input (identity) or rebuild an + * ordinary scalar Cast with the validated annotations and the original span. + * No new scalar rounding/saturation mode is introduced. */ +enum class ScalarCastAction { + kKeepInput, + kRebuildCast, +}; + +struct ScalarCastPlan { + ScalarCastAction action = ScalarCastAction::kKeepInput; + /*! Validated annotations to carry onto the rebuilt Cast (the same shape as + * CastNode::annotations). Attributes that are proven to match the + * existing scalar lowering are normalized away, so this map is empty for + * the currently admitted directions. */ + ffi::Map annotations; +}; + +/*! One ordered vector conversion action. */ +enum class VectorCastOp { + /*! Same-sign widening: one logical vcvt. */ + kIntegerExtend, + /*! Same-width signedness change / source-sign re-view: one + * vinterpret_cast. */ + kIntegerReinterpret, + /*! Narrowing including cross-sign: one logical vcvt with NOSAT. */ + kIntegerTruncate, + /*! Real non-integer conversion (f16/bf16/FP8 -> f32, f32 -> f16/bf16/FP8); + * empty options still mean a conversion, never a no-op. */ + kNonIntegerConvert, +}; + +struct VectorCastStep { + VectorCastOp op = VectorCastOp::kIntegerExtend; + DataType from; // element type consumed by this step + DataType to; // element type produced by this step + /*! PTODSL rounding annotation value, e.g. "R"; empty when N/A. */ std::string rounding; + /*! PTODSL saturate annotation value, e.g. "NOSAT"; empty when N/A. */ + std::string saturate; }; +struct VectorCastPlan { + /*! Ordered actions; an empty list is a proven identity (no conversion). + * A non-empty list still may be a single real conversion with empty + * options (that is not an identity). */ + std::vector steps; +}; + +using CastPlan = std::variant; + /*! Operation-level type support for the vector arithmetic emitted * here. VMI vmul only supports i16/i32/f16/bf16/f32 * elements; vadd/vsub additionally allow i8; vdiv admits f16/f32 only. @@ -271,32 +323,32 @@ LaneUse AnalyzeLaneUse(const PrimExpr &expr, LaneUseContext *ctx, LaneUse AnalyzeLaneUseMeta(const PrimExpr &expr, LaneUseContext *ctx, std::string *reason); -/*! Decide the conversion attributes for a value Cast. This is the single - * admission point for every Cast on the vectorize path; Verify and Vectorize - * both consume its result, so they cannot diverge. +/*! Decide the conversion plan for a value Cast. This is the single admission + * point for every Cast on the vectorize path; Verify and Vectorize both pass + * an explicit CastContext and consume the same plan, so they cannot diverge. * - * Same-width reinterpretation-free identity is impossible in - * TIR Cast, so the proven set is: - * - integer narrowing/widening with truncating semantics: TVM integer - * casts truncate on narrowing (codegen_llvm CreateIntCast), PTODSL - * needs saturate="NOSAT" for narrowing to keep that semantics; - * - float widening (f16->f32, bf16->f32): exact; - * - the FP8 matrix: E4M3FN/E5M2 <-> f32. The conversion instructions the - * target exposes are FP32 <-> FP8, so a narrowing conversion carries an - * explicit rounding and saturation attribute (the FP8 conversion policy in - * the implementation file) and a half (f16/bf16) source is rejected rather - * than routed through float32 internally. FP8 -> FP16/BF16, FP8 <-> integer, - * E4M3FN <-> E5M2 and every other FP8 variant stay out of scope. - * Everything else (float narrowing, float<->int, unusual bit widths) - * returns nullopt -> reject. + * Dispatch is first scalar vs VMI, then integer vs non-integer: + * - scalar: identity KeepInput; integer widening/narrowing and the newly + * admitted same-width signedness change rebuild an ordinary scalar Cast; + * f16/bf16 <-> f32 keep their existing scalar lowering; non-identity FP8, + * float<->int and f16<->bf16 stay rejected; + * - VMI integer: identity 0 steps; same-width cross-sign one + * vinterpret_cast; same-sign widening one vcvt; cross-sign widening + * source-sign widen then same-width vinterpret_cast; narrowing one direct + * vcvt NOSAT. All 36 supported pairs are represented; no consumer-layout + * feasibility is guessed here; + * - VMI non-integer: a real conversion action for f16/bf16/FP8 -> f32 + * (empty options are not a no-op) and for f32 -> f16/bf16/FP8 with the + * existing rounding/saturation policy. * * `cast` carries the backend annotations written by the frontend cast * helpers (rounding mode, saturation flag, stochastic rbits operand). They * are read here — never interpreted by the expression visitors — and an * attribute the selected conversion cannot honour is a rejection, not a * silently dropped hint. When `reason` is non-null it receives a diagnostic - * naming the source type, target type and the offending option. */ + * naming the category, source/target types, lanes and the offending option. */ Optional PlanValueCast(const CastNode *cast, + const CastContext &context, std::string *reason = nullptr); /*! Check whether an expression only uses nodes from the supported diff --git a/src/ascend/transform/vectorize_parallel_to_pto.cc b/src/ascend/transform/vectorize_parallel_to_pto.cc index 826010819c..3994eadfcb 100644 --- a/src/ascend/transform/vectorize_parallel_to_pto.cc +++ b/src/ascend/transform/vectorize_parallel_to_pto.cc @@ -143,9 +143,19 @@ class DeadBindElim { // are rejected there, and opaque calls are already excluded // from addresses. static const std::set kDiscardable = { - "tl.vmi.vload", "tl.vmi.vbrc", "tl.vmi.vci", "tl.vmi.vadd", - "tl.vmi.vsub", "tl.vmi.vmul", "tl.vmi.vdiv", "tl.vmi.vabs", - "tl.vmi.vexp", "tl.vmi.vln", "tl.vmi.vsqrt", "tl.vmi.vcvt", + "tl.vmi.vload", + "tl.vmi.vbrc", + "tl.vmi.vci", + "tl.vmi.vadd", + "tl.vmi.vsub", + "tl.vmi.vmul", + "tl.vmi.vdiv", + "tl.vmi.vabs", + "tl.vmi.vexp", + "tl.vmi.vln", + "tl.vmi.vsqrt", + "tl.vmi.vcvt", + "tl.vmi.vinterpret_cast", "tl.access_ptr", }; bool discardable = true; @@ -841,43 +851,75 @@ class VectorizeParallelToPTOImpl : public StmtExprMutator { pre); } if (const auto *cast = expr.as()) { + // Convert the input exactly once, then plan from the explicit shape of + // that converted input (the original Cast carries the element types). PrimExpr value = ConvertExpr(cast->value, pre); - // Emit the conversion attributes PlanValueCast proved - // equivalent; never rely on PTODSL defaults (they saturate integer - // narrowing). + pto_analysis::CastContext context; + if (value.dtype().lanes() == 1) { + context.value_kind = pto_analysis::CastValueKind::kScalar; + context.logical_lanes = 1; + } else { + ICHECK_EQ(value.dtype().lanes(), lanes_) + << "[VectorizeParallelToPTO] internal error: Cast input has " + "unexpected lanes " + << value.dtype().lanes() << " (current VL " << lanes_ << ")"; + context.value_kind = pto_analysis::CastValueKind::kVectorVmi; + context.logical_lanes = lanes_; + } + ICHECK_EQ(value.dtype().element_of(), cast->value.dtype()) + << "[VectorizeParallelToPTO] internal error: converted Cast input " + "element type " + << value.dtype().element_of() << " does not match the original " + << cast->value.dtype(); std::string plan_why; - auto plan = pto_analysis::PlanValueCast(cast, &plan_why); + auto plan = pto_analysis::PlanValueCast(cast, context, &plan_why); ICHECK(plan.has_value()) << "[VectorizeParallelToPTO] internal error: Cast passed Verify " "without a conversion plan: " << plan_why; - if (plan.value().identity) { - // Same-type cast is a no-op: keep the converted value. - return value; + if (const auto *scalar = + std::get_if(&plan.value())) { + ICHECK_EQ(value.dtype().lanes(), 1) + << "[VectorizeParallelToPTO] internal error: scalar Cast plan " + "for a vector input"; + if (scalar->action == pto_analysis::ScalarCastAction::kKeepInput) { + return value; + } + // Rebuild the ordinary scalar Cast with the validated annotations and + // the original span; the scalar lowering itself is unchanged. + return Cast(cast->dtype, value, scalar->annotations, cast->span); } - if (value.dtype().lanes() > 1) { + const auto &vector = std::get(plan.value()); + PrimExpr current = value; + for (const auto &step : vector.steps) { + ICHECK_EQ(current.dtype().element_of(), step.from) + << "[VectorizeParallelToPTO] internal error: Cast step expects " + << step.from << ", got " << current.dtype().element_of(); Map attrs; - attrs.Set("to_dtype", StringImm(TvmDtypeName(cast->dtype))); - if (!plan.value().saturate.empty()) { - attrs.Set("saturate", StringImm(plan.value().saturate)); - } - if (!plan.value().rounding.empty()) { - attrs.Set("rounding", StringImm(plan.value().rounding)); + attrs.Set("to_dtype", StringImm(TvmDtypeName(step.to))); + PrimExpr call; + if (step.op == pto_analysis::VectorCastOp::kIntegerReinterpret) { + call = Call(VectorDType(step.to, lanes_), VmiOp("vinterpret_cast"), + {current}, attrs); + } else { + if (!step.saturate.empty()) { + attrs.Set("saturate", StringImm(step.saturate)); + } + if (!step.rounding.empty()) { + attrs.Set("rounding", StringImm(step.rounding)); + } + call = Call(VectorDType(step.to, lanes_), VmiOp("vcvt"), {current}, + attrs); } - return EmitBind("vcvt", - Call(VectorDType(cast->dtype, value.dtype().lanes()), - VmiOp("vcvt"), {value}, attrs), - pre); + current = EmitBind("cast", call, pre); } - // A lane-uniform FP8 conversion has no scalar PTO lowering inside a - // SIMD_VF body; Verify rejects that form, so reaching it here means the - // two passes disagree. - ICHECK(!pto_analysis::IsSupportedFloat8(cast->dtype) && - !pto_analysis::IsSupportedFloat8(cast->value.dtype())) - << "[VectorizeParallelToPTO] internal error: lane-uniform FP8 Cast " - "passed Verify: " - << cast->value.dtype() << " -> " << cast->dtype; - return Cast(cast->dtype, value); + // Zero steps is the proven identity; a final type check keeps every + // multi-step chain honest (source convert once, dtype/VL continuous). + ICHECK_EQ(current.dtype(), VectorDType(cast->dtype, lanes_)) + << "[VectorizeParallelToPTO] internal error: Cast plan ended at " + << current.dtype() << ", expected " + << VectorDType(cast->dtype, lanes_); + return current; } LOG(FATAL) << "[VectorizeParallelToPTO] internal error: unsupported " "expression passed Verify: " diff --git a/src/ascend/transform/verify_parallel_to_pto.cc b/src/ascend/transform/verify_parallel_to_pto.cc index 1e03801dae..f0bd64e1c8 100644 --- a/src/ascend/transform/verify_parallel_to_pto.cc +++ b/src/ascend/transform/verify_parallel_to_pto.cc @@ -744,47 +744,41 @@ class VerifyParallelToPTOImpl : public StmtExprVisitor { public: CastCheck(UnitVisitor *parent) : parent_(parent) {} void VisitExpr_(const CastNode *op) final { - std::string why; - auto plan = pto_analysis::PlanValueCast(op, &why); - if (!plan.has_value()) { + // Classify the *input of this Cast* with the existing lane-use + // analysis, then hand the explicit kind to the shared planner. + pto_analysis::LaneUseContext ctx = parent_->MakeLaneUseContext(); + std::string lane_why; + pto_analysis::LaneUse use = + pto_analysis::AnalyzeLaneUse(op->value, &ctx, &lane_why); + if (use == pto_analysis::LaneUse::kUnknown) { std::ostringstream oss; oss << "Cast " << op->value.dtype() << " -> " << op->dtype - << " has no proven VMI-equivalent conversion"; - if (!why.empty()) { - oss << ": " << why; - } + << " cannot classify its input: " << lane_why; parent_->parent_->Fail("dtype", oss.str()); - } - // A lane-uniform Cast is emitted as a scalar Cast, while the - // conversion plan above proves the VMI-equivalent *vector* form. - // FP8 has no scalar form at all, so a lane-uniform FP8 Cast has to - // be read per lane; the non-FP8 pairs keep whatever lowering the - // codegen gives them. - if (plan.value().identity) { - VisitExpr(op->value); // annotations are not computation values return; } - if (!pto_analysis::IsSupportedFloat8(op->value.dtype()) && - !pto_analysis::IsSupportedFloat8(op->dtype)) { - VisitExpr(op->value); // annotations are not computation values - return; + pto_analysis::CastContext context; + if (use == pto_analysis::LaneUse::kUniform) { + context.value_kind = pto_analysis::CastValueKind::kScalar; + context.logical_lanes = 1; + } else { + context.value_kind = pto_analysis::CastValueKind::kVectorVmi; + context.logical_lanes = parent_->mapping_.lanes; } - pto_analysis::LaneUseContext ctx = parent_->MakeLaneUseContext(); - std::string lane_why; - pto_analysis::LaneUse use = - pto_analysis::AnalyzeLaneUse(op->value, &ctx, &lane_why); - if (use != pto_analysis::LaneUse::kVarying) { + std::string why; + auto plan = pto_analysis::PlanValueCast(op, context, &why); + if (!plan.has_value()) { std::ostringstream oss; oss << "Cast " << op->value.dtype() << " -> " << op->dtype - << " stays lane-uniform"; - if (use == pto_analysis::LaneUse::kUnknown) { - oss << " (lane use cannot be classified: " << lane_why << ")"; + << " has no supported conversion"; + if (!why.empty()) { + oss << ": " << why; } - oss << "; FP8 conversions are only supported as a vector vcvt " - "over the full lane range"; parent_->parent_->Fail("dtype", oss.str()); } - VisitExpr(op->value); // annotations are not computation values + // Nested value casts each plan independently; annotations are not + // computation values and are validated by the planner. + VisitExpr(op->value); } void VisitExpr_(const BufferLoadNode *) final { // Index casts belong to address analysis; do not descend. diff --git a/testing/ascend/language/pto/cast_plan_context_probe.cc b/testing/ascend/language/pto/cast_plan_context_probe.cc new file mode 100644 index 0000000000..d1020e3be3 --- /dev/null +++ b/testing/ascend/language/pto/cast_plan_context_probe.cc @@ -0,0 +1,148 @@ +/* Direct PlanValueCast context probe. + * + * Calls PlanValueCast directly with CastContext combinations that the + * Verify/Vectorize passes cannot produce end-to-end, and checks both the + * rejection reasons and the emitted vector/scalar plans. + * + * The two passes can only produce the context combinations that their + * classification yields (scalar lanes=1, VMI lanes=64/128/256). The planner + * guards for {kVectorVmi, 1} and {kScalar, 128} are therefore unreachable + * end-to-end; this local probe calls the planner directly so those contracts + * are tested without adding a product FFI. + * + * Build (from the repository root; the library lives in build/lib and the + * TVM DSOs must be linked explicitly): + * c++ -std=c++17 -O0 testing/ascend/language/pto/cast_plan_context_probe.cc \ + * -Isrc -I3rdparty/tvm/include -I3rdparty/tvm/src \ + * -I3rdparty/tvm/3rdparty/dlpack/include \ + * -I3rdparty/tvm/3rdparty/tvm-ffi/include \ + * -I3rdparty/tvm/3rdparty/tvm-ffi/3rdparty/dlpack/include \ + * -Lbuild/lib -ltilelang -ltvm_compiler -ltvm_ffi -ltvm_runtime \ + * -Wl,-rpath,$PWD/build/lib -o /tmp/cast_plan_probe + * LD_LIBRARY_PATH=$PWD/build/lib /tmp/cast_plan_probe + */ +#include +#include +#include + +#include "ascend/transform/parallel_to_pto_utils.h" +#include + +namespace { + +using tvm::DataType; +using tvm::tirx::Cast; +using tvm::tirx::Var; +using tvm::tl::pto::CastContext; +using tvm::tl::pto::CastValueKind; +using tvm::tl::pto::PlanValueCast; +using tvm::tl::pto::ScalarCastAction; +using tvm::tl::pto::ScalarCastPlan; +using tvm::tl::pto::VectorCastOp; +using tvm::tl::pto::VectorCastPlan; + +int failures = 0; + +void Check(bool condition, const std::string &what) { + if (!condition) { + std::cout << "FAIL " << what << "\n"; + ++failures; + } else { + std::cout << "PASS " << what << "\n"; + } +} + +const VectorCastPlan * +AsVector(const tvm::ffi::Optional &plan, + const std::string &what) { + if (!plan.has_value()) { + std::cout << "FAIL " << what << ": no plan\n"; + ++failures; + return nullptr; + } + const auto *vector = std::get_if(&plan.value()); + if (vector == nullptr) { + std::cout << "FAIL " << what << ": not a vector plan\n"; + ++failures; + } + return vector; +} + +const ScalarCastPlan * +AsScalar(const tvm::ffi::Optional &plan, + const std::string &what) { + if (!plan.has_value()) { + std::cout << "FAIL " << what << ": no plan\n"; + ++failures; + return nullptr; + } + const auto *scalar = std::get_if(&plan.value()); + if (scalar == nullptr) { + std::cout << "FAIL " << what << ": not a scalar plan\n"; + ++failures; + } + return scalar; +} + +} // namespace + +int main() { + Var x("x", DataType::Int(8)); + Cast widen(DataType::Int(16), x); + Cast cross(DataType::UInt(16), x); + Cast same_width(DataType::UInt(8), x); + + std::string why; + CastContext vector_lane1{CastValueKind::kVectorVmi, 1}; + auto plan = PlanValueCast(widen.get(), vector_lane1, &why); + Check(!plan.has_value() && + why.find("requires logical lanes 64/128/256") != std::string::npos, + "kVectorVmi lanes=1 is rejected as VMI, never scalar: " + why); + + why.clear(); + CastContext scalar_lanes{CastValueKind::kScalar, 128}; + plan = PlanValueCast(widen.get(), scalar_lanes, &why); + Check(!plan.has_value() && + why.find("must have exactly one logical lane") != std::string::npos, + "kScalar lanes=128 is rejected, never VMI: " + why); + + why.clear(); + CastContext vector64{CastValueKind::kVectorVmi, 64}; + plan = PlanValueCast(cross.get(), vector64, &why); + if (const VectorCastPlan *steps = AsVector(plan, "int8->uint16 VMI plan")) { + Check(steps->steps.size() == 2, "int8->uint16 has 2 steps"); + if (steps->steps.size() == 2) { + Check(steps->steps[0].op == VectorCastOp::kIntegerExtend && + steps->steps[0].from == DataType::Int(8) && + steps->steps[0].to == DataType::Int(16), + "step 1 extends with the source signedness to the target width"); + Check(steps->steps[1].op == VectorCastOp::kIntegerReinterpret && + steps->steps[1].from == DataType::Int(16) && + steps->steps[1].to == DataType::UInt(16), + "step 2 reinterprets at the same width"); + } + } + + why.clear(); + CastContext scalar1{CastValueKind::kScalar, 1}; + plan = PlanValueCast(cross.get(), scalar1, &why); + if (const ScalarCastPlan *scalar = + AsScalar(plan, "int8->uint16 scalar plan")) { + Check(scalar->action == ScalarCastAction::kRebuildCast, + "int8->uint16 scalar keeps a rebuilt ordinary Cast"); + Check(scalar->annotations.empty(), + "scalar annotations are validated/empty"); + } + + why.clear(); + CastContext vector128{CastValueKind::kVectorVmi, 128}; + plan = PlanValueCast(same_width.get(), vector128, &why); + if (const VectorCastPlan *steps = + AsVector(plan, "int8->uint8 VMI same-width plan")) { + Check(steps->steps.size() == 1 && + steps->steps[0].op == VectorCastOp::kIntegerReinterpret, + "same-width cross-sign is one vinterpret_cast"); + } + + return failures == 0 ? 0 : 1; +} diff --git a/testing/ascend/language/pto/test_pto_parallel_cast_matrix.py b/testing/ascend/language/pto/test_pto_parallel_cast_matrix.py new file mode 100644 index 0000000000..38e8d8cad6 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_parallel_cast_matrix.py @@ -0,0 +1,702 @@ +"""Parallel value-Cast planning: action structure and attribute rules. + +Hardware independent. Each case lowers the real TileLang PTO pipeline with +``tilelang.engine.lower.lower`` (VerifyParallelToPTO + VectorizeParallelToPTO ++ PTODSL codegen; no PTOAS/device) and inspects the generated source and the +module at the Verify boundary. + +Covered contract: + +* ``PlanValueCast(cast, CastContext{value_kind, logical_lanes}, reason)`` + dispatches scalar vs VMI first, then integer vs non-integer; Verify and + Vectorize classify the current Cast input independently and consistently; +* integer VMI actions: identity 0, same-width cross-sign 1 vinterpret_cast, + same-sign widening 1 vcvt, cross-sign widening source-sign extend + same-width + vinterpret_cast (2), supported narrowing 1 vcvt NOSAT; all 36 integer pairs + at VL=128 follow the table; +* the converted input is produced once and steps are emitted in order with + continuous dtype/VL and an exact final type; vinterpret_cast participates in + the local DCE; +* scalar plans keep ordinary scalar lowering (KeepInput / validated rebuild), + no new scalar rounding/saturation mode, and no parallel planner is wired + into generic codegen; +* attribute rules: accepted conversions pin the emitted rounding/saturation + mode, and unsupported combinations are rejected with a reason (unknown keys, + illegal sat policies, float <-> integer, f16 <-> bf16, out-of-scope FP8, + lane-uniform FP8); +* classification scenarios: external scalars, uniform loads, Bind aliases, 2D + kept/selected coordinates, nested casts, and scalar cast then broadcast. + +Full backend compilation and the device numerical checks live in +``test_pto_parallel_cast_matrix_run.py``. +""" + +from __future__ import annotations + +import os +import re +import subprocess +from contextlib import contextmanager + +import pytest +import torch # noqa: F401 (torch/torch_npu must import before tilelang/ptodsl) +import torch_npu # noqa: F401 + +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + +tirx = tvm.tirx +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) + +INT_TYPES = ["int8", "uint8", "int16", "uint16", "int32", "uint32"] +CROSS_SIGN_WIDEN_INTERMEDIATE = { + ("int8", "uint16"): "int16", + ("int8", "uint32"): "int32", + ("uint8", "int16"): "uint16", + ("uint8", "int32"): "uint32", + ("int16", "uint32"): "int32", + ("uint16", "int32"): "uint32", +} +PTO_INT_NAME = { + "int8": "pto.si8", + "uint8": "pto.ui8", + "int16": "pto.si16", + "uint16": "pto.ui16", + "int32": "pto.si32", + "uint32": "pto.ui32", +} + + +def _bits(dtype): + return int("".join(ch for ch in dtype if ch.isdigit())) + + +def _signed(dtype): + return dtype.startswith("int") + + +def _pair_category(src, dst): + if src == dst: + return "identity" + if _bits(src) == _bits(dst): + return "reinterpret" + if _bits(dst) > _bits(src): + return "widen" if _signed(src) == _signed(dst) else "widen_and_reinterpret" + return "narrow" + + +# --------------------------------------------------------------------------- +# Lowering / capture helpers +# --------------------------------------------------------------------------- +def _target(): + return normalize_pto_target("pto") + + +def _lower_source(func): + return _lower(func, target=_target()).kernel_source + + +@contextmanager +def _verify_capture(): + captured = {} + original = ascend_transform.VerifyParallelToPTO + + def factory(): + verify = original() + + def run(mod): + captured["module"] = mod + return verify(mod) + + return run + + ascend_transform.VerifyParallelToPTO = factory + try: + yield captured + finally: + ascend_transform.VerifyParallelToPTO = original + + +def _lower_capture(func): + error = None + with _verify_capture() as captured: + try: + _lower(func, target=_target()) + except Exception as exc: # noqa: BLE001 - expected diagnostics + error = exc + return captured.get("module"), error + + +def _expect_lower_error(func, pattern): + _, error = _lower_capture(func) + assert error is not None, f"expected a Verify diagnostic matching {pattern!r}" + assert re.search(pattern, str(error)), f"unexpected diagnostic: {error}" + + +def _cast_lines(source): + """Lines emitted by a value-Cast plan. + + Excludes the scalar-bridge `vinterpret_cast` that wraps `vbrc` (the scalar + signedness recovery), which is broadcasting, not a cast action. + """ + lines = [] + for line in source.splitlines(): + stripped = line.strip() + if "vmi.vcvt" in stripped or ("vinterpret_cast" in stripped and "vmi.vbrc(" not in stripped and "vmi.vci(" not in stripped): + lines.append(stripped) + return lines + + +def _one(source, needle): + hits = [line.strip() for line in source.splitlines() if needle in line] + assert len(hits) == 1, f"expected one line containing {needle!r}, got {hits!r}" + return hits[0] + + +def _count(source, needle): + return source.count(needle) + + +# --------------------------------------------------------------------------- +# Kernels +# --------------------------------------------------------------------------- +def _cast_kernel(src, dst, lanes, annotations=None): + n = 2 * lanes + + @T.prim_func + def main(A: T.Buffer((n,), src), C: T.Buffer((n,), dst)): + with T.Kernel(1): + a = T.alloc_shared((n,), src) + c = T.alloc_shared((n,), dst) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + if annotations is None: + c[i] = T.cast(a[i], dst) + else: + c[i] = T.Cast(dst, a[i], annotations=annotations) + T.copy(c, C) + + return main + + +def _ann_kernel(src, dst, lanes, annotations, *, uniform=False): + n = lanes + + @T.prim_func + def main(A: T.Buffer((n,), src), C: T.Buffer((n,), dst)): + with T.Kernel(1): + a = T.alloc_shared((n,), src) + c = T.alloc_shared((n,), dst) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + value = a[0] if uniform else a[i] + c[i] = T.Cast(dst, value, annotations=annotations) + T.copy(c, C) + + return main + + +def _alias_kernel(src, dst, lanes, uniform): + n = lanes + + @T.prim_func + def main(A: T.Buffer((n,), src), C: T.Buffer((n,), dst)): + with T.Kernel(1): + a = T.alloc_shared((n,), src) + c = T.alloc_shared((n,), dst) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + x = a[0] if uniform else a[i] + y = x + c[i] = T.cast(y, dst) + T.copy(c, C) + + return main + + +def _external_scalar_kernel(vl): + n = vl + + @T.prim_func + def main(S: T.int8, C: T.Buffer((n,), "int32")): + with T.Kernel(1): + c = T.alloc_shared((n,), "int32") + with T.SimdVF(lanes=vl): + for i in T.Parallel(n): + c[i] = T.cast(S, "int32") + T.copy(c, C) + + return main + + +def _binary_scalar_cast_kernel(vl): + n = vl + + @T.prim_func + def main(A: T.Buffer((n,), "int16"), B: T.Buffer((n,), "int8"), C: T.Buffer((n,), "int16")): + with T.Kernel(1): + a = T.alloc_shared((n,), "int16") + b = T.alloc_shared((n,), "int8") + c = T.alloc_shared((n,), "int16") + T.copy(A, a) + T.copy(B, b) + with T.SimdVF(lanes=vl): + for i in T.Parallel(n): + c[i] = a[i] + T.cast(b[0], "int16") + T.copy(c, C) + + return main + + +def _coordinate_kernel(vl, which): + m, n = 4, vl + out = "int16" + + @T.prim_func + def main(A: T.Buffer((m, n), "int32"), C: T.Buffer((m, n), out)): + with T.Kernel(1): + a = T.alloc_shared((m, n), "int32") + c = T.alloc_shared((m, n), out) + T.copy(A, a) + with T.SimdVF(lanes=vl): + for i, j in T.Parallel(m, n): + if which == "kept_of_load": + c[i, j] = T.cast(a[i, 0], out) + else: + c[i, j] = T.cast(a[i, j], out) + T.copy(c, C) + + return main + + +def _nested_cast_kernel(vl): + n = vl + + @T.prim_func + def main(A: T.Buffer((n,), "int16"), C: T.Buffer((n,), "uint32")): + with T.Kernel(1): + a = T.alloc_shared((n,), "int16") + c = T.alloc_shared((n,), "uint32") + T.copy(A, a) + with T.SimdVF(lanes=vl): + for i in T.Parallel(n): + c[i] = T.cast(T.cast(a[i], "uint16"), "uint32") + T.copy(c, C) + + return main + + +def _dead_chain_kernel(vl): + n = vl + + @T.prim_func + def main(A: T.Buffer((n,), "int8"), C: T.Buffer((n,), "int8")): + with T.Kernel(1): + a = T.alloc_shared((n,), "int8") + c = T.alloc_shared((n,), "int8") + T.copy(A, a) + with T.SimdVF(lanes=vl): + for i in T.Parallel(n): + # The assignment creates the Bind whose pure conversion + # chain must be eliminated by the local DCE. + unused = T.cast(T.cast(a[i], "uint16"), "uint32") # noqa: F841 + c[i] = a[i] + T.copy(c, C) + + return main + + +def _fp8_kernel(src, dst, lanes, annotations=None, **kwargs): + n = lanes + + @T.prim_func + def main(A: T.Buffer((n,), src), C: T.Buffer((n,), dst)): + with T.Kernel(1): + a = T.alloc_shared((n,), src) + c = T.alloc_shared((n,), dst) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + if annotations is not None: + c[i] = T.Cast(dst, a[i], annotations=annotations) + else: + c[i] = T.cast(a[i], dst, **kwargs) + T.copy(c, C) + + return main + + +# --------------------------------------------------------------------------- +# A. 36 integer pairs at VL=128: action structure (lower-only, fast) +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("src", INT_TYPES) +@pytest.mark.parametrize("dst", INT_TYPES) +def test_integer_pair_action_structure(src, dst): + category = _pair_category(src, dst) + source = _lower_source(_cast_kernel(src, dst, 128)) + lines = _cast_lines(source) + assert _count(source, "vmi.vload") == 1, source + if category == "identity": + assert lines == [], lines + return + if category == "reinterpret": + assert len(lines) == 1 and "vinterpret_cast" in lines[0], lines + assert f"to_dtype={PTO_INT_NAME[dst]}" in lines[0], lines[0] + return + if category == "widen": + assert len(lines) == 1 and "vmi.vcvt" in lines[0], lines + assert f"to_dtype={PTO_INT_NAME[dst]}" in lines[0], lines[0] + assert "saturate" not in lines[0], lines[0] + return + if category == "widen_and_reinterpret": + intermediate = CROSS_SIGN_WIDEN_INTERMEDIATE[(src, dst)] + assert len(lines) == 2, lines + assert "vmi.vcvt" in lines[0], lines[0] + assert f"to_dtype={PTO_INT_NAME[intermediate]}" in lines[0], lines[0] + assert "vinterpret_cast" in lines[1], lines[1] + assert f"to_dtype={PTO_INT_NAME[dst]}" in lines[1], lines[1] + first_result = lines[0].split("=")[0].strip() + assert lines[1].count(first_result) == 1, lines + return + # Narrowing: one direct logical vcvt with NOSAT. + assert len(lines) == 1 and "vmi.vcvt" in lines[0], lines + assert 'saturate="NOSAT"' in lines[0], lines[0] + assert f"to_dtype={PTO_INT_NAME[dst]}" in lines[0], lines[0] + + +def _count_cast_nodes(module): + func = next(iter(module.functions.values())) + found = [] + + def visit(node): + if isinstance(node, tirx.Cast): + found.append(node) + + tirx.stmt_functor.post_order_visit(func.body, visit) + return found + + +@pytest.mark.parametrize("dtype", INT_TYPES) +def test_identity_cast_reaches_planner_and_emits_zero_actions(dtype): + # A plain same-dtype T.cast is folded by the frontend, so the annotated + # raw Cast is used to prove the planner identity path is reachable: the + # manual sat=True marker must be validated and normalized away, and the + # emission has zero conversion actions. + func = _ann_kernel(dtype, dtype, 128, {"sat": True}) + module, error = _lower_capture(func) + assert error is None, error + casts = [node for node in _count_cast_nodes(module) if str(node.dtype) == dtype and str(node.value.dtype) == dtype] + assert casts, "the annotated same-dtype Cast must reach Verify" + source = _lower_source(func) + assert _cast_lines(source) == [] + + +# --------------------------------------------------------------------------- +# B. Producer once, ordered binds, scope and DCE +# --------------------------------------------------------------------------- +def test_cross_sign_widening_producer_and_order(): + source = _lower_source(_cast_kernel("int8", "uint32", 256)) + lines = _cast_lines(source) + assert len(lines) == 2 + vload = _one(source, "vmi.vload") + assert vload.split("=")[0].strip() in lines[0] + assert source.index(lines[0]) < source.index(lines[1]) + assert _count(source, "vmi.vload") == 1 + + +def test_conversion_binds_stay_between_load_and_store(): + source = _lower_source(_cast_kernel("int8", "uint32", 128)) + vload_index = source.index("vmi.vload") + first_cast = source.index(_cast_lines(source)[0]) + store_index = source.index("vmi.vstore") + assert vload_index < first_cast < store_index + + +def test_dead_pure_conversion_chain_is_dropped(): + source = _lower_source(_dead_chain_kernel(128)) + assert "vmi.vcvt" not in source, source + assert "vinterpret_cast" not in source, source + # The dead chain's vload is gone too; only the live a[i] load remains. + assert _count(source, "vmi.vload") == 1, source + + +def test_nested_cast_chains_inner_then_outer(): + source = _lower_source(_nested_cast_kernel(128)) + lines = _cast_lines(source) + assert len(lines) == 2, lines + assert "vinterpret_cast" in lines[0] and "to_dtype=pto.ui16" in lines[0] + assert "vmi.vcvt" in lines[1] and "to_dtype=pto.ui32" in lines[1] + assert _count(source, "vmi.vload") == 1 + + +# --------------------------------------------------------------------------- +# C. Verify/Vectorize classification consistency +# --------------------------------------------------------------------------- +def test_external_scalar_cast_stays_scalar(): + source = _lower_source(_external_scalar_kernel(128)) + assert _cast_lines(source) == [] + assert "pto.cast(" in source + assert "vbrc" in source + + +def test_uniform_load_cast_stays_scalar(): + source = _lower_source(_ann_kernel("int8", "int32", 128, {}, uniform=True)) + assert _cast_lines(source) == [] + assert "pto.cast(" in source + + +def test_bind_alias_keeps_vmi_classification(): + source = _lower_source(_alias_kernel("int8", "int32", 128, uniform=False)) + assert len(_cast_lines(source)) == 1 + assert _count(source, "vmi.vload") == 1 + + +def test_scalar_cast_then_parent_broadcast(): + source = _lower_source(_binary_scalar_cast_kernel(128)) + assert "vbrc" in source + assert "pto.cast(" in source + assert _cast_lines(source) == [] + + +@pytest.mark.parametrize( + ("which", "expect_vector"), + [ + ("kept_of_load", False), + ("selected_of_load", True), + ], +) +def test_2d_coordinate_classification(which, expect_vector): + source = _lower_source(_coordinate_kernel(128, which)) + lines = _cast_lines(source) + if expect_vector: + assert len(lines) == 1, lines + else: + assert lines == [], lines + assert "pto.cast(" in source + + +def test_lane_one_vector_context_is_not_constructible(): + # {kVectorVmi, 1} would be a VMI value with one logical lane; the current + # Parallel VL contract cannot produce it (lanes are 64/128/256), and the + # planner rejects it explicitly. The observable frontend boundary: + with pytest.raises(ValueError, match="lanes must be one of"): + T.SimdVF(lanes=1) + + +# --------------------------------------------------------------------------- +# D. Attribute table +# --------------------------------------------------------------------------- +def _lower_ann(src, dst, annotations): + return _lower_source(_ann_kernel(src, dst, 128, annotations)) + + +def test_integer_narrowing_accepts_sat_false(): + source = _lower_ann("int16", "int8", {"sat": False}) + line = _cast_lines(source)[0] + assert 'saturate="NOSAT"' in line, line + + +def test_float_narrowing_accepts_sat_false_and_nosat(): + source = _lower_ann("float32", "float16", {"sat": False}) + line = _cast_lines(source)[0] + assert 'rounding="R"' in line and 'saturate="NOSAT"' in line, line + + +def test_fp8_narrowing_accepts_round_and_sat_false(): + source = _lower_source(_fp8_kernel("float32", "float8_e5m2", 128, {"round": "Z", "sat": False})) + line = _cast_lines(source)[0] + assert 'rounding="Z"' in line and 'saturate="NOSAT"' in line, line + + +def test_integer_narrowing_rejects_sat_true(): + _expect_lower_error( + _ann_kernel("int16", "int8", 128, {"sat": True}), + "integer narrowing accepts only sat=False", + ) + + +def test_unknown_annotation_is_rejected(): + _expect_lower_error( + _ann_kernel("int8", "int32", 128, {"mystery": 1}), + "unknown cast annotation", + ) + + +def test_float_int_conversion_is_rejected(): + _expect_lower_error( + _ann_kernel("float32", "int32", 128, {}), + "float <-> integer conversions need per-combination", + ) + + +def test_f16_bf16_cross_family_is_rejected(): + _expect_lower_error( + _ann_kernel("float16", "bfloat16", 128, {}), + "cross-family f16 <-> bf16", + ) + + +def test_fp8_to_half_is_rejected(): + _expect_lower_error( + _fp8_kernel("float8_e4m3fn", "float16", 128), + "FP8 -> FP16/BF16 and FP8 <-> integer are out of scope", + ) + + +def test_lane_uniform_fp8_conversion_is_rejected(): + _expect_lower_error( + _ann_kernel("float8_e4m3fn", "float32", 128, {}, uniform=True), + "non-identity FP8 conversions have no scalar lowering", + ) + + +# --------------------------------------------------------------------------- +# E. Scalar paths and literal boundaries (existing semantics) +# --------------------------------------------------------------------------- +def test_scalar_float_widening_stays_scalar(): + source = _lower_source(_ann_kernel("float16", "float32", 128, {}, uniform=True)) + assert _cast_lines(source) == [] + assert "vbrc" in source + + +def test_scalar_float_narrowing_stays_scalar(): + source = _lower_source(_ann_kernel("float32", "float16", 128, {}, uniform=True)) + assert _cast_lines(source) == [] + assert "vbrc" in source + + +def test_scalar_integer_narrowing_keeps_plain_cast(): + source = _lower_source(_ann_kernel("int32", "int8", 128, {}, uniform=True)) + assert _cast_lines(source) == [] + assert "pto.cast(" in source + + +def test_actual_t_cast_literal_out_of_range_is_rejected_by_frontend(): + with pytest.raises(Exception, match="exceeds maximum of int8"): + T.cast(300, "int8") + + +def test_actual_t_cast_negative_to_unsigned_is_rejected_by_frontend(): + with pytest.raises(Exception, match="cannot make uint from negative value"): + T.cast(-1, "uint8") + + +# --------------------------------------------------------------------------- +# F. Direct planner contexts and scalar rebuilt span/annotations +# --------------------------------------------------------------------------- +def test_plan_value_cast_context_matrix_directly(tmp_path): + """The two passes can only produce scalar lanes=1 and VMI lanes=64/128/256, + so {kVectorVmi, 1} and {kScalar, 128} are unreachable end-to-end. A local + C++ probe calls PlanValueCast directly for both and for positive + vector/scalar plans, without adding a product FFI.""" + binary = tmp_path / "cast_plan_probe" + source = os.path.join(os.path.dirname(__file__), "cast_plan_context_probe.cc") + command = [ + "c++", + "-std=c++17", + "-O0", + source, + "-I", + f"{REPO_ROOT}/src", + "-I", + f"{REPO_ROOT}/3rdparty/tvm/include", + "-I", + f"{REPO_ROOT}/3rdparty/tvm/src", + "-I", + f"{REPO_ROOT}/3rdparty/tvm/3rdparty/dlpack/include", + "-I", + f"{REPO_ROOT}/3rdparty/tvm/3rdparty/tvm-ffi/include", + "-I", + f"{REPO_ROOT}/3rdparty/tvm/3rdparty/tvm-ffi/3rdparty/dlpack/include", + "-L", + f"{REPO_ROOT}/build/lib", + "-ltilelang", + "-ltvm_compiler", + "-ltvm_ffi", + "-ltvm_runtime", + f"-Wl,-rpath,{REPO_ROOT}/build/lib", + "-o", + str(binary), + ] + build = subprocess.run(command, cwd=REPO_ROOT, capture_output=True, text=True, timeout=600) + assert build.returncode == 0, build.stdout + build.stderr + env = dict(os.environ, LD_LIBRARY_PATH=f"{REPO_ROOT}/build/lib") + run = subprocess.run([str(binary)], env=env, capture_output=True, text=True, timeout=120) + assert run.returncode == 0, run.stdout + run.stderr + assert "FAIL" not in run.stdout, run.stdout + assert "kVectorVmi lanes=1 is rejected" in run.stdout + assert "kScalar lanes=128 is rejected" in run.stdout + + +@contextmanager +def _vectorize_capture(): + captured = {} + original = ascend_transform.VectorizeParallelToPTO + + def factory(): + vectorize = original() + + def run(mod): + out = vectorize(mod) + captured["module"] = out + return out + + return run + + ascend_transform.VectorizeParallelToPTO = factory + try: + yield captured + finally: + ascend_transform.VectorizeParallelToPTO = original + + +def _span_scalar_kernel(vl, span): + @T.prim_func + def main(S: T.Buffer((32,), "int32"), C: T.Buffer((vl,), "int8")): + with T.Kernel(1): + s = T.alloc_shared((32,), "int32") + c = T.alloc_shared((vl,), "int8") + T.copy(S, s) + with T.SimdVF(lanes=vl): + for i in T.Parallel(vl): + c[i] = T.cast(s[0], "int8", sat=False, span=span) + T.copy(c, C) + + return main + + +def test_scalar_rebuild_preserves_span_and_drop_provenance(): + # Span(source_name, line, end_line, column, end_column). + span = tvm.ir.Span(tvm.ir.SourceName("cast_span_probe"), 3, 3, 7, 11) + with _vectorize_capture() as captured: + _lower(_span_scalar_kernel(128, span), target=_target()) + module = captured["module"] + func = next(iter(module.functions.values())) + rebuilt = [] + + def visit(node): + if isinstance(node, tirx.Cast) and str(node.dtype) == "int8": + rebuilt.append(node) + + tirx.stmt_functor.post_order_visit(func.body, visit) + assert rebuilt, "the rebuilt scalar Cast is missing" + matches = [ + node + for node in rebuilt + if node.span is not None and node.span.source_name is not None and node.span.source_name.name == "cast_span_probe" + ] + assert matches, "the rebuilt scalar Cast lost the original span" + node = matches[0] + assert node.span.line == 3 and node.span.column == 7 + assert node.span.end_line == 3 and node.span.end_column == 11 + # The validated sat=False marker is proven equivalent to the ordinary + # truncating scalar lowering and is normalized away; no new attribute + # format is introduced. + assert "sat" not in node.annotations + assert "round" not in node.annotations diff --git a/testing/ascend/language/pto/test_pto_parallel_cast_matrix_run.py b/testing/ascend/language/pto/test_pto_parallel_cast_matrix_run.py new file mode 100644 index 0000000000..4ef6037726 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_parallel_cast_matrix_run.py @@ -0,0 +1,400 @@ +"""Device acceptance: parallel value-Cast matrix. + +Each case compiles through the real TileLang -> Verify/Vectorize -> PTODSL -> +PTOAS VPTO -> LLVM pipeline and executes on the NPU. References are +computed independently on the CPU with explicit low-bit selection and +target-signedness adjustment (never from the generated kernel or backend). + +Coverage: +* 22 representative integer pair cells at VL=128, direct store: all 12 + narrowing pairs, all 6 cross-signed widening pairs, 2 same-signed widening + pairs, and 2 same-width cross-sign reinterpret pairs; +* exact named boundary values through the full pipeline; +* target-signed continued consumption (same-sign widening of the converted + target) for representative 8/16-bit sources; +* one tail extent (2VL+17) with padded physical source and output sentinels; +* scalar (lane-uniform load) integer runtime casts and a scalar literal cast; +* the scalar-float expected backend rejection. +""" + +import os +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import torch_npu # noqa: E402,F401 (import before tilelang/ptodsl) +import tilelang # noqa: E402 +import tilelang.ascend.language as T # noqa: E402 + +TORCH_DT = { + "int8": torch.int8, + "uint8": torch.uint8, + "int16": torch.int16, + "uint16": torch.uint16, + "int32": torch.int32, + "uint32": torch.uint32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, +} +SENTINEL = { + "int8": -99, + "uint8": 199, + "int16": -9999, + "uint16": 55555, + "int32": -999999, + "uint32": 4000000000, +} + + +def _require_npu_runtime(): + assert hasattr(torch, "npu"), "torch.npu is unavailable in this environment" + assert torch.npu.is_available(), "torch.npu is unavailable in the current environment" + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source), "NPU host->device copy probe mismatch" + + +@pytest.fixture(scope="module") +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +def _bits(dtype): + return int("".join(ch for ch in dtype if ch.isdigit())) + + +def _signed(dtype): + return dtype.startswith("int") + + +def _wrap(value, bits, signed): + value &= (1 << bits) - 1 + if signed and value >= (1 << (bits - 1)): + value -= 1 << bits + return value + + +def _convert(value, src, dst): + """Independent reference: select the source value, truncate to dst bits, + then interpret the result with the target signedness.""" + src_value = _wrap(int(value), _bits(src), _signed(src)) + result = src_value & ((1 << _bits(dst)) - 1) + if _signed(dst) and result >= (1 << (_bits(dst) - 1)): + result -= 1 << _bits(dst) + return result + + +def _source_values(src, count): + bits = _bits(src) + signed = _signed(src) + raw = [ + -(1 << (bits - 1)) if signed else 0, + (1 << (bits - 1)) - 1 if signed else (1 << bits) - 1, + 0, + 1, + 2, + -1 if signed else (1 << (bits - 1)), + -2 if signed else (1 << (bits - 1)) + 1, + 127, + 128, + 255, + 256, + 257, + 300, + 65535, + 65536, + 4294967295, + 0x80000000, + ] + values = [_wrap(value, bits, signed) for value in raw] + return [values[index % len(values)] for index in range(count)] + + +def _tensor(values, dtype): + return torch.tensor(values, dtype=TORCH_DT[dtype]) + + +# --------------------------------------------------------------------------- +# Kernels +# --------------------------------------------------------------------------- +def _cast_kernel(src, dst, lanes, extent=None): + n = extent if extent is not None else 2 * lanes + storage = _storage_len(src, lanes, n) + + @T.prim_func + def main(A: T.Buffer((storage,), src), C: T.Buffer((storage,), dst)): + with T.Kernel(1): + a = T.alloc_shared((storage,), src) + c = T.alloc_shared((storage,), dst) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.cast(a[i], dst) + T.copy(c, C) + + return main, storage + + +def _consumer_kernel(src, mid, out, lanes): + n = 2 * lanes + storage = _storage_len(src, lanes, n) + + @T.prim_func + def main(A: T.Buffer((storage,), src), C: T.Buffer((storage,), out)): + with T.Kernel(1): + a = T.alloc_shared((storage,), src) + c = T.alloc_shared((storage,), out) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.cast(T.cast(a[i], mid), out) + T.copy(c, C) + + return main + + +def _scalar_cast_kernel(src, dst, lanes): + n = lanes + + @T.prim_func + def main(S: T.Buffer((32,), src), C: T.Buffer((n,), dst)): + with T.Kernel(1): + s = T.alloc_shared((32,), src) + c = T.alloc_shared((n,), dst) + T.copy(S, s) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.cast(s[0], dst) + T.copy(c, C) + + return main + + +def _physical_lanes(dtype): + """Physical VMI vector width in elements (256-byte vectors).""" + return 256 * 8 // _bits(dtype) + + +def _storage_len(src, lanes, extent): + """Readable source storage for the real continuous-vload footprint. + + A logical vload of ``size`` lanes is declared with the full physical + vector extent on the last chunk, so the last chunk base must be followed + by a full physical vector even when the logical lane count is smaller. + """ + padded = ((extent + lanes - 1) // lanes) * lanes + return max(padded, padded - lanes + _physical_lanes(src)) + + +def _tail_kernel(src, dst, lanes, extent): + storage = _storage_len(src, lanes, extent) + + @T.prim_func + def main( + A: T.Buffer((storage,), src), + S: T.Buffer((storage,), dst), + C: T.Buffer((storage,), dst), + ): + with T.Kernel(1): + a = T.alloc_shared((storage,), src) + c = T.alloc_shared((storage,), dst) + T.copy(A, a) + T.copy(S, c) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(extent): + c[i] = T.cast(a[i], dst) + T.copy(c, C) + + return main, storage + + +def _run_pair(npu, src, dst, lanes, extent=None): + func, storage = _cast_kernel(src, dst, lanes, extent=extent) + kernel = tilelang.compile(func, target="pto", out_idx=-1) + n = extent if extent is not None else 2 * lanes + values = _source_values(src, storage) + expected = torch.tensor([_convert(value, src, dst) for value in values], dtype=TORCH_DT[dst]) + out = kernel(_tensor(values, src).to(npu)) + torch.npu.synchronize() + got = out.cpu() + assert got.dtype == expected.dtype, (src, dst, lanes, got.dtype) + assert torch.equal(got[0:n], expected[0:n]), ( + f"{src}->{dst} VL{lanes}: value mismatch\n got={got[:16].tolist()}\n exp={expected[:16].tolist()}" + ) + assert got.numel() == expected.numel() == storage + return got, expected, storage + + +# --------------------------------------------------------------------------- +# A. Representative pair cells at VL=128: full compile + device execution +# --------------------------------------------------------------------------- +DEVICE_PAIRS = [ + # Narrowing: one vcvt with NOSAT. + ("int16", "int8"), + ("int16", "uint8"), + ("uint16", "int8"), + ("uint16", "uint8"), + ("int32", "int8"), + ("int32", "uint8"), + ("int32", "int16"), + ("int32", "uint16"), + ("uint32", "int8"), + ("uint32", "uint8"), + ("uint32", "int16"), + ("uint32", "uint16"), + # Cross-signed widening: source-signed vcvt, then vinterpret_cast. + ("int8", "uint16"), + ("int8", "uint32"), + ("uint8", "int16"), + ("uint8", "int32"), + ("int16", "uint32"), + ("uint16", "int32"), + # Same-signed widening: one vcvt. + ("int8", "int16"), + ("uint16", "uint32"), + # Same-width cross-sign: one vinterpret_cast. + ("int8", "uint8"), + ("uint32", "int32"), +] + + +@pytest.mark.parametrize(("src", "dst"), DEVICE_PAIRS) +def test_integer_pair_vl_device(npu, src, dst): + got, _, storage = _run_pair(npu, src, dst, 128) + # The returned allocation is the readable-storage contract (physical + # load padding included); only the logical [0, 2*VL) region is written + # and compared. The padding beyond it is uninitialized by design. + assert got.numel() == storage + assert storage >= 256 + + +def test_reference_examples_are_exact(npu): + cases = [ + ("int8", "uint32", -1, 4294967295), + ("uint8", "int32", 255, 255), + ("int32", "int8", 300, 44), + ("int32", "uint8", -1, 255), + ("int32", "uint8", 256, 0), + ] + for src, dst, value, expected in cases: + func, storage = _cast_kernel(src, dst, 128, extent=128) + kernel = tilelang.compile(func, target="pto", out_idx=-1) + values = [value] * storage + source = torch.tensor(values, dtype=TORCH_DT[src]) + out = kernel(source.to(npu)) + torch.npu.synchronize() + got = out.cpu() + assert got.numel() == storage + # Only the logical extent (128) was written; the physical-padding + # region is not part of the value contract. + assert int(got[0]) == expected, (src, dst, value, int(got[0])) + assert bool((got[:128] == expected).all()), (src, dst, value, got.tolist()[:8]) + + +# --------------------------------------------------------------------------- +# B. Target-signed continued consumption +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("src", "mid", "out"), + [ + ("int8", "uint16", "uint32"), + ("int8", "uint8", "uint32"), + ("uint8", "int8", "int32"), + ("int16", "uint16", "uint32"), + ], +) +def test_target_signed_continued_widening(npu, src, mid, out): + """The converted target is consumed by its own same-sign widening: a + '-' source value must not sign-extend through the unsigned target.""" + func = _consumer_kernel(src, mid, out, 128) + kernel = tilelang.compile(func, target="pto", out_idx=-1) + storage = _storage_len(src, 128, 256) + values = [-1, 0, 1, 127, -128, 255, 256, 300] * 64 + values = [_wrap(value, _bits(src), _signed(src)) for value in values[:storage]] + expected = torch.tensor( + [_convert(_convert(v, src, mid), mid, out) for v in values[:256]], + dtype=TORCH_DT[out], + ) + got = kernel(torch.tensor(values, dtype=TORCH_DT[src]).to(npu)).cpu() + torch.npu.synchronize() + assert torch.equal(got[:256], expected), f"{src}->{mid}->{out}\n got={got[:16].tolist()}\n exp={expected[:16].tolist()}" + + +# --------------------------------------------------------------------------- +# C. Tail extent with padded source and output sentinels +# --------------------------------------------------------------------------- +def test_tail_extent_with_output_sentinel(npu): + src, dst, lanes, extent = "int8", "uint32", 128, 2 * 128 + 17 + func, storage = _tail_kernel(src, dst, lanes, extent) + kernel = tilelang.compile(func, target="pto", out_idx=-1) + values = _source_values(src, storage) + sentinel = SENTINEL[dst] + c_init = torch.full((storage,), sentinel, dtype=TORCH_DT[dst]) + got = kernel(_tensor(values, src).to(npu), c_init.to(npu)).cpu() + # `kernel` returns the third argument (C); `c_init` seeds it via S. + torch.npu.synchronize() + expected = torch.tensor([_convert(values[i], src, dst) for i in range(extent)], dtype=TORCH_DT[dst]) + assert torch.equal(got[0:extent], expected), f"{src}->{dst} VL{lanes} E={extent}: values" + assert bool((got[extent:storage] == sentinel).all()), f"{src}->{dst} VL{lanes} E={extent}: tail sentinel overwritten" + + +# --------------------------------------------------------------------------- +# D. Scalar (lane-uniform) integer and float controls +# --------------------------------------------------------------------------- +SCALAR_INT_CASES = [ + ("int8", "uint32"), + ("int32", "uint8"), +] + + +@pytest.mark.parametrize(("src", "dst"), SCALAR_INT_CASES) +def test_scalar_integer_runtime_cast(npu, src, dst): + func = _scalar_cast_kernel(src, dst, 128) + kernel = tilelang.compile(func, target="pto", out_idx=-1) + for value in (-1, 0, 1, 255, 256, 300): + source_values = [_wrap(value, _bits(src), _signed(src))] * 32 + source = torch.tensor(source_values, dtype=TORCH_DT[src]) + out = kernel(source.to(npu)).cpu() + torch.npu.synchronize() + expected = _convert(source_values[0], src, dst) + assert int(out[0]) == expected, (src, dst, value, int(out[0])) + assert bool((out == expected).all()), (src, dst, value) + + +def test_scalar_integer_literal_cast(npu): + @T.prim_func + def main(C: T.Buffer((128,), "int32")): + with T.Kernel(1): + c = T.alloc_shared((128,), "int32") + with T.SimdVF(lanes=128): + for i in T.Parallel(128): + c[i] = T.cast(7, "int32") + T.copy(c, C) + + kernel = tilelang.compile(main, target="pto", out_idx=-1) + out = kernel().cpu() + torch.npu.synchronize() + assert bool((out == 7).all()), out.tolist()[:8] + + +def test_scalar_float_cast_rejected_by_device_compiler(npu): + """The BiSheng device compiler rejects scalar float casts in an AIV loop. + + PTODSL, VPTO and LLVM emission support these conversions. The device + compiler reports "Unsupported scalar instruction in AIV loop" with or + without the value-Cast planner. This test checks that rejection; the + scalar planner preserves the existing cast semantics. + """ + with pytest.raises(Exception) as excinfo: + tilelang.compile(_scalar_cast_kernel("float32", "float16", 128), target="pto", out_idx=-1) + assert "Unsupported scalar instruction in AIV loop" in str(excinfo.value) From 2800fbdc2a223c9945cb3ddf2d9e53d196f53811 Mon Sep 17 00:00:00 2001 From: qukelin Date: Sat, 17 Oct 2026 15:52:08 +0800 Subject: [PATCH 08/10] fix(pto): verify masked-store destination alignment --- .../transform/verify_pto_codegen_legality.cc | 1162 +++++++++++++++++ .../pto/test_pto_feature_interaction_run.py | 438 +++++++ .../pto/test_pto_masked_store_alignment.py | 883 +++++++++++++ tilelang/ascend/pipeline.py | 5 + tilelang/ascend/transform/__init__.py | 14 + 5 files changed, 2502 insertions(+) create mode 100644 src/ascend/transform/verify_pto_codegen_legality.cc create mode 100644 testing/ascend/language/pto/test_pto_feature_interaction_run.py create mode 100644 testing/ascend/language/pto/test_pto_masked_store_alignment.py diff --git a/src/ascend/transform/verify_pto_codegen_legality.cc b/src/ascend/transform/verify_pto_codegen_legality.cc new file mode 100644 index 0000000000..18b3d1ab5c --- /dev/null +++ b/src/ascend/transform/verify_pto_codegen_legality.cc @@ -0,0 +1,1162 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file verify_pto_codegen_legality.cc + * \brief Verify PTO requirements before code generation without modifying IR. + * + * This pass consumes the final device TIR after LowerDeviceKernelLaunch and + * the final layout and UB rewrites. It currently checks destination alignment + * for default/continuous explicit-mask VMI stores of b8/b16/b32 data. + * Ordinary masked stores require provable 32-byte alignment. Stores selected + * for 1PT_B8/B16/B32 writeback require natural element alignment (1/2/4 bytes): + * the value must have 1/2/4/8 lanes and the mask must be a recognized + * full-active seed with the same lane count. Bool data is excluded from this + * check. The checker resolves byte addresses using the Analyzer facts at each + * access. Shapes, layouts, masks and addresses are never modified. Functions + * without a PTO target pass through unchanged. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arith/ir_visitor_with_analyzer.h" +#include "tirx/transform/ir_utils.h" + +#include "../../op/builtin.h" +#include "../target_utils.h" + +namespace tvm { +namespace tl { +namespace transform { + +using namespace tirx; + +namespace { + +constexpr int kOrdinaryMaskedStoreAlignment = 32; + +enum class StoreMode { + kOrdinaryMasked, + kPointStores, +}; + +enum class Outcome { + kAligned, + kMisaligned, + kUnknown, +}; + +struct ByteAddress { + enum class Kind { + kUbRoot, + kIntegerBase, + kUnknown, + }; + Kind kind = Kind::kUnknown; + const VarNode *root = nullptr; + PrimExpr base_bytes; + PrimExpr offset_bytes; + std::string error; +}; + +struct StoreMatch { + bool checked = false; + StoreMode mode = StoreMode::kOrdinaryMasked; + const char *mode_name = ""; + int required_alignment = kOrdinaryMaskedStoreAlignment; + int elem_bytes = 0; + DataType value_dtype; + // Element type of the pointer that reaches the store after descriptor, + // address_of, typed Var/make_tensor and pure alias resolution. It is the + // unit of the instruction offset; it is never defaulted to float32. + std::optional pointer_dtype; + bool has_descriptor = false; + DataType descriptor_dtype = DataType::Float(32); + PrimExpr descriptor_offset; + PrimExpr instruction_offset; +}; + +struct SeedInfo { + bool recognized = false; + bool full_active = false; + int64_t mask_lanes = 0; +}; + +struct Proof { + Outcome outcome = Outcome::kUnknown; + PrimExpr byte_offset; + PrimExpr remainder; + std::string error; +}; + +} // namespace + +/*! Checker for explicit-mask continuous VMI stores. Traversal is scoped: + * loop/branch/attribute bodies push and pop the pointer/pure-definition + * maps, and the Analyzer provides the existing range/condition facts. */ +class PTOMaskedStoreAlignmentChecker final + : public arith::IRVisitorWithAnalyzer { +public: + explicit PTOMaskedStoreAlignmentChecker(std::string function_name) + : function_name_(std::move(function_name)) {} + + using Base = arith::IRVisitorWithAnalyzer; + using Base::VisitExpr_; + using Base::VisitStmt_; + + void Run(const PrimFunc &func) { VisitStmt(func->body); } + + // UB arena roots: PTO codegen materializes the merged dyn-shared + // allocation at UB byte address 0, and MergeUBAllocations aligns the arena + // to 32 bytes. Only allocations known to feed that arena are roots; a raw + // handle is never assumed aligned. Roots are recorded in lexical order + // during the single statement walk, so expression DAGs are never re-walked + // by a whole-body pre-scan. + void VisitStmt_(const AllocBufferNode *op) final { + const Var &data = op->buffer->data; + const std::string scope = GetPtrStorageScope(data); + if (scope == "shared.dyn" || scope == "shared") { + ub_roots_.insert(data.get()); + } + StmtExprVisitor::VisitStmt_(op); + } + + void VisitStmt_(const ForNode *op) final { + constraint_scope_.WithNewScope([&]() { + analyzer_.Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); + constraint_scope_.Current().Emplace(&analyzer_, op->extent > 0); + if (op->step.defined()) { + const auto *step = op->step.as(); + if (step != nullptr && step->value > 1) { + // The base visitor binds [min, min+extent) but drops the loop + // lattice. Record the real iteration congruence so a positive + // constant step such as i = 0, 8, 16, ... can be proven. + constraint_scope_.Current().Emplace( + &analyzer_, + floormod(op->loop_var - op->min, op->step.value()) == 0); + } + } + AliasScope saved = SaveAliases(); + StmtExprVisitor::VisitStmt_(op); + RestoreAliases(std::move(saved)); + }); + } + + void VisitStmt_(const SBlockNode *op) final { + constraint_scope_.WithNewScope([&]() { + for (const auto &iter_var : op->iter_vars) { + analyzer_.Bind(iter_var->var, iter_var->dom); + } + AliasScope saved = SaveAliases(); + StmtExprVisitor::VisitStmt_(op); + RestoreAliases(std::move(saved)); + }); + } + + void VisitStmt_(const IfThenElseNode *op) final { + constraint_scope_.WithNewScope([&]() { + this->VisitExpr(op->condition); + PrimExpr real_condition = ExtractRealCondition(op->condition); + constraint_scope_.WithNewScope([&]() { + constraint_scope_.Current().Emplace(&analyzer_, real_condition); + AliasScope saved = SaveAliases(); + this->VisitStmt(op->then_case); + RestoreAliases(std::move(saved)); + }); + if (op->else_case) { + constraint_scope_.WithNewScope([&]() { + constraint_scope_.Current().Emplace( + &analyzer_, analyzer_.rewrite_simplify(Not(real_condition))); + AliasScope saved = SaveAliases(); + this->VisitStmt(op->else_case.value()); + RestoreAliases(std::move(saved)); + }); + } + }); + } + + void VisitStmt_(const AttrStmtNode *op) final { + constraint_scope_.WithNewScope([&]() { + if (op->attr_key == tirx::attr::thread_extent || + op->attr_key == s_tir::attr::virtual_thread) { + IterVar iv = Downcast(op->node); + TVM_FFI_ICHECK_NE(iv->thread_tag.length(), 0U); + analyzer_.Bind(iv->var, Range::FromMinExtent( + IntImm(op->value->dtype, 0), op->value)); + } else if (op->attr_key == tirx::attr::tilelang_assume) { + auto condition = Downcast(op->node); + constraint_scope_.Current().Emplace(&analyzer_, condition, + /*is_assume=*/true); + } + AliasScope saved = SaveAliases(); + StmtExprVisitor::VisitStmt_(op); + RestoreAliases(std::move(saved)); + }); + } + + void VisitStmt_(const BindNode *op) final { + std::unordered_set visiting{op->var.get()}; + ExprMemo memo; + ScalarFacts facts = AnalyzeScalar(op->value, &visiting, &memo); + if (facts.needs_generic_visit) { + this->VisitExpr(op->value); + } + RegisterDefinition(op->var.get(), op->value, facts); + } + + void VisitExpr_(const LetNode *op) final { + std::unordered_set visiting{op->var.get()}; + ExprMemo value_memo; + ScalarFacts value_facts = AnalyzeScalar(op->value, &visiting, &value_memo); + if (value_facts.needs_generic_visit) { + this->VisitExpr(op->value); + } + AliasScope saved = SaveAliases(); + RegisterDefinition(op->var.get(), op->value, value_facts); + std::unordered_set body_visiting{op->var.get()}; + ExprMemo body_memo; + ScalarFacts body_facts = + AnalyzeScalar(op->body, &body_visiting, &body_memo); + if (body_facts.needs_generic_visit) { + this->VisitExpr(op->body); + } + RestoreAliases(std::move(saved)); + } + + void VisitExpr_(const CallNode *op) final { + if (op->op.same_as(tl::vmi_vstore())) { + CheckStore(op); + } + Base::VisitExpr_(op); + } + +private: + struct AliasScope { + std::unordered_map pointer_defs; + std::unordered_map value_defs; + }; + + AliasScope SaveAliases() const { + return AliasScope{pointer_defs_, value_defs_}; + } + + void RestoreAliases(AliasScope saved) { + pointer_defs_ = std::move(saved.pointer_defs); + value_defs_ = std::move(saved.value_defs); + } + + struct ScalarFacts { + bool origin_safe = false; + bool needs_generic_visit = false; + bool has_shared_node = false; + }; + using ExprMemo = std::unordered_map; + using CanonMemo = std::unordered_map; + + // Bottom-up, memoized facts for one expression / right-hand side. + // + // origin_safe: the expression is a scalar (lanes == 1, non-handle) tree of + // immediates, free variables, casts and pure scalar arithmetic, and every + // referenced bind definition is origin-safe as well (a lane-1 VMI alias, a + // memory read, a Call or a cycle is unsafe). Plain SSA chains of any length + // are classified through the per-Var cache one shallow definition at a time. + // + // needs_generic_visit: the generic statement/expression walk is required + // because the expression may contain a checked store (a void Call), a Let + // definition, branch constraints, Reduce axes or an unrecognized node. A + // pure elementwise arithmetic expression cannot contain any of these, so + // skipping the walk is safe for nested checked-store discovery. + // + // has_shared_node: a non-leaf expression node is referenced more than once + // within this right-hand side (a DAG). This only selects the canonicalized + // Analyzer binding; the definition and its equality are always kept. + ScalarFacts AnalyzeScalar(const PrimExpr &value, + std::unordered_set *visiting, + ExprMemo *memo, int depth = 0) { + // Recursion/stack-safety guard only, not a complexity or acceptance + // bound: repeated work is already bounded by the per-RHS memo and the + // per-Var classification cache. An expression nested deeper than this is + // treated conservatively (no Analyzer equality and a generic visit), so a + // dependent store is diagnosed as unprovable instead of risking the + // compiler stack. + constexpr int kMaxExpressionRecursionGuard = 4096; + if (depth > kMaxExpressionRecursionGuard) { + return ScalarFacts{false, true, false}; + } + const DataType dtype = value.dtype(); + if (value.as() || value.as() || + value.as()) { + return ScalarFacts{true, false, false}; + } + if (const auto *var = value.as()) { + if (visiting->count(var) != 0) { + // Self/cyclic reference: not a lawful SSA definition. + return ScalarFacts{false, false, false}; + } + auto safe_it = var_origin_safe_.find(var); + if (safe_it != var_origin_safe_.end()) { + return ScalarFacts{safe_it->second, false, false}; + } + auto def_it = value_defs_.find(var); + if (def_it == value_defs_.end()) { + // Function/loop/block parameter without a recorded definition. + return ScalarFacts{true, false, false}; + } + visiting->insert(var); + // Analyze the definition with its own memo: sharing inside a separate + // right-hand side must not mark the current expression as shared. + ExprMemo def_memo; + ScalarFacts def_facts = + AnalyzeScalar(def_it->second, visiting, &def_memo, depth + 1); + visiting->erase(var); + var_origin_safe_[var] = def_facts.origin_safe; + return ScalarFacts{def_facts.origin_safe, false, false}; + } + auto memo_it = memo->find(value.get()); + if (memo_it != memo->end()) { + return ScalarFacts{memo_it->second.origin_safe, + memo_it->second.needs_generic_visit, true}; + } + auto child = [&](const PrimExpr &expr) { + return AnalyzeScalar(expr, visiting, memo, depth + 1); + }; + const bool scalar_root = !dtype.is_handle() && dtype.lanes() == 1; + auto binary = [&](const PrimExpr &a, const PrimExpr &b) { + ScalarFacts fa = child(a); + ScalarFacts fb = child(b); + ScalarFacts out; + out.origin_safe = scalar_root && fa.origin_safe && fb.origin_safe; + out.needs_generic_visit = + fa.needs_generic_visit || fb.needs_generic_visit; + out.has_shared_node = fa.has_shared_node || fb.has_shared_node; + return out; + }; + auto ternary = [&](const PrimExpr &a, const PrimExpr &b, + const PrimExpr &c) { + ScalarFacts fa = child(a); + ScalarFacts fb = child(b); + ScalarFacts fc = child(c); + ScalarFacts out; + out.origin_safe = + scalar_root && fa.origin_safe && fb.origin_safe && fc.origin_safe; + out.needs_generic_visit = fa.needs_generic_visit || + fb.needs_generic_visit || + fc.needs_generic_visit; + out.has_shared_node = + fa.has_shared_node || fb.has_shared_node || fc.has_shared_node; + return out; + }; + ScalarFacts facts; + if (const auto *node = value.as()) { + ScalarFacts c = child(node->value); + facts.origin_safe = scalar_root && c.origin_safe; + facts.needs_generic_visit = c.needs_generic_visit; + facts.has_shared_node = c.has_shared_node; + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = ternary(node->condition, node->true_value, node->false_value); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + ScalarFacts c = child(node->a); + facts.origin_safe = scalar_root && c.origin_safe; + facts.needs_generic_visit = c.needs_generic_visit; + facts.has_shared_node = c.has_shared_node; + } else if (const auto *node = value.as()) { + bool needs = false; + bool shared = false; + for (const auto &index : node->indices) { + ScalarFacts f = child(index); + needs = needs || f.needs_generic_visit; + shared = shared || f.has_shared_node; + } + if (node->predicate.defined()) { + ScalarFacts f = child(node->predicate.value()); + needs = needs || f.needs_generic_visit; + shared = shared || f.has_shared_node; + } + facts = ScalarFacts{false, needs, shared}; + } else if (const auto *node = value.as()) { + ScalarFacts f = binary(node->base, node->stride); + facts = ScalarFacts{false, f.needs_generic_visit, f.has_shared_node}; + } else if (const auto *node = value.as()) { + ScalarFacts f = child(node->value); + facts = ScalarFacts{false, f.needs_generic_visit, f.has_shared_node}; + } else if (const auto *node = value.as()) { + bool needs = false; + bool shared = false; + for (const auto &vector : node->vectors) { + ScalarFacts f = child(vector); + needs = needs || f.needs_generic_visit; + shared = shared || f.has_shared_node; + } + for (const auto &index : node->indices) { + ScalarFacts f = child(index); + needs = needs || f.needs_generic_visit; + shared = shared || f.has_shared_node; + } + facts = ScalarFacts{false, needs, shared}; + } else { + // Calls (including every VMI producer and vstore), Let, Reduce and + // every other opaque node keep the generic walk and are not scalar + // equalities. + facts = ScalarFacts{false, true, false}; + } + (*memo)[value.get()] = facts; + return facts; + } + + // Rebuild an origin-safe expression with each distinct non-leaf node bound + // to a fresh shallow symbol. The Analyzer's scalar binding walks its + // argument recursively, so handing it a shared DAG (e+k; e=e+e; ...) is + // exponential even though the classification above is memoized. With one + // shallow binding per distinct node the substitution graph stays linear and + // the equality -- including its modular/alignment consequences -- is fully + // preserved. + PrimExpr CanonicalizeForAnalyzer(const PrimExpr &value, CanonMemo *memo) { + if (value.as() || value.as() || + value.as() || value.as()) { + return value; + } + auto memo_it = memo->find(value.get()); + if (memo_it != memo->end()) { + return memo_it->second; + } + auto canon = [&](const PrimExpr &expr) { + return CanonicalizeForAnalyzer(expr, memo); + }; + PrimExpr rebuilt; + if (const auto *node = value.as()) { + rebuilt = + Cast(node->dtype, canon(node->value), node->annotations, node->span); + } else if (const auto *node = value.as()) { + rebuilt = Add(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Sub(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Mul(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Div(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = FloorDiv(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Mod(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = FloorMod(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Min(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Max(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Select(canon(node->condition), canon(node->true_value), + canon(node->false_value), node->span); + } else if (const auto *node = value.as()) { + rebuilt = LE(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = LT(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = GT(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = GE(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = EQ(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = NE(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = And(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Or(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Not(canon(node->a), node->span); + } else { + // Unreachable for a classified origin-safe expression; keep the value + // rather than inventing a different semantic. + (*memo)[value.get()] = value; + return value; + } + Var fresh("__pto_align_dag_" + std::to_string(next_dag_symbol_++), + rebuilt.dtype()); + analyzer_.Bind(fresh, rebuilt); + (*memo)[value.get()] = fresh; + return fresh; + } + + void RegisterDefinition(const VarNode *var, const PrimExpr &value, + const ScalarFacts &facts) { + if (value.dtype().is_handle()) { + pointer_defs_[var] = value; + return; + } + // Scalar/index definitions and pointer aliases are recorded separately: + // value_defs_ backs the origin-safe scalar check and the mask-seed alias + // walk, while pointer_defs_ backs address resolution. Seed matching only + // accepts recognized create_mask/pset calls and compile-time active + // counts, so a memory-read definition never classifies as a seed. + value_defs_[var] = value; + const bool safe = value.dtype().is_scalar() && !value.dtype().is_handle() && + facts.origin_safe; + var_origin_safe_[var] = safe; + if (!safe) { + return; + } + if (facts.has_shared_node) { + // Preserve the equality with a compact substitution graph instead of + // dropping it (which would lose alignment proofs) or expanding the DAG. + CanonMemo canon_memo; + PrimExpr canonical = CanonicalizeForAnalyzer(value, &canon_memo); + analyzer_.Bind(ffi::GetRef(var), canonical); + } else { + analyzer_.Bind(ffi::GetRef(var), value); + } + } + + // Follow immutable SSA aliases iteratively: the length of a plain alias + // chain is not a recursion-depth budget. A visited set detects cycles. + std::optional ResolveMaskActiveCount(const PrimExpr &expr) const { + PrimExpr current = expr; + std::unordered_set seen; + while (const auto *var = current.as()) { + if (!seen.insert(var).second) { + return std::nullopt; + } + auto it = value_defs_.find(var); + if (it == value_defs_.end()) { + return std::nullopt; + } + current = it->second; + } + if (const auto *imm = current.as()) { + return imm->value; + } + return std::nullopt; + } + + PrimExpr ResolvePureAlias(const PrimExpr &expr) const { + PrimExpr current = expr; + std::unordered_set seen; + while (const auto *var = current.as()) { + if (!seen.insert(var).second) { + return current; // cycle: not a recognized seed + } + auto it = value_defs_.find(var); + if (it == value_defs_.end()) { + return current; + } + current = it->second; + } + return current; + } + + // Mask seeds recognized by the PTO backend for 1PT_B8/B16/B32 writeback: + // + // * A default/continuous vstore with 1/2/4/8 data lanes can use point + // writeback when its create_mask(active, size) seed has the same lane + // count and a compile-time active count at least as large as size. + // This selects full-active pset_b32 PAT_ALL and one point store per + // logical element. The value need not come from a reduction. + // * A pset PAT_ALL mask is also a recognized seed. TileLang does not expose + // a pset operator in its Python API; this branch handles direct TIR calls. + // * create_mask(0), partial masks and dynamic active counts do not match. + // These remain ordinary masked stores and require 32-byte alignment. + // Mismatched value/mask lanes are rejected by the backend, and arbitrary + // all-true expressions are not recognized seeds. + SeedInfo ClassifyMaskSeed(const PrimExpr &mask) const { + SeedInfo info; + const DataType mask_dtype = mask.dtype(); + if (!mask_dtype.is_bool() || mask_dtype.lanes() <= 0) { + return info; + } + info.mask_lanes = mask_dtype.lanes(); + PrimExpr resolved = ResolvePureAlias(mask); + const auto *call = resolved.as(); + if (call == nullptr) { + return info; + } + auto opt_op = call->op.as(); + if (!opt_op.has_value()) { + return info; + } + const std::string &name = opt_op.value()->name; + if (name == "tl.vmi.create_mask") { + info.recognized = true; + if (!call->args.empty()) { + std::optional active = ResolveMaskActiveCount(call->args[0]); + info.full_active = active.has_value() && *active >= info.mask_lanes; + } + return info; + } + if (name == "tl.vmi.pset" && !call->args.empty()) { + if (const auto *dist = call->args[0].as()) { + if (dist->value == "PAT_ALL") { + info.recognized = true; + info.full_active = true; + } + } + } + return info; + } + + std::optional DecodeAccessDType(const PrimExpr &annotation) const { + const auto *call = annotation.as(); + if (call == nullptr) { + return std::nullopt; + } + std::string n; + if (!call->args.empty()) { + if (const auto *name = call->args[0].as()) { + n = name->value; + } + } + if (n.empty()) { + // The final device IR carries a `tirx.type_annotation` call whose + // result dtype is the annotated element type (the call has no + // positional arguments at this stage). Use that dtype so descriptor + // and instruction offsets are scaled exactly as codegen does. + auto opt_op = call->op.as(); + if (!opt_op.has_value() || + opt_op.value()->name != "tirx.type_annotation") { + return std::nullopt; + } + DataType annotated = call->dtype; + if (!annotated.is_scalar() || annotated.is_handle()) { + return std::nullopt; + } + return annotated; + } + if (n == "float" || n == "float32") + return DataType::Float(32); + if (n == "half" || n == "float16") + return DataType::Float(16); + if (n == "bfloat16" || n == "bfloat16_t") + return DataType::BFloat(16); + if (n == "float8_e4m3" || n == "float8_e4m3fn") { + return DataType::Float8E4M3FN(); + } + if (n == "float8_e5m2") + return DataType::Float8E5M2(); + if (n == "int8") + return DataType::Int(8); + if (n == "uint8") + return DataType::UInt(8); + if (n == "int16") + return DataType::Int(16); + if (n == "uint16") + return DataType::UInt(16); + if (n == "int32") + return DataType::Int(32); + if (n == "uint32") + return DataType::UInt(32); + if (n == "int64") + return DataType::Int(64); + if (n == "uint64") + return DataType::UInt(64); + return std::nullopt; + } + + std::optional + PointerAnnotationElementType(const VarNode *var) const { + const auto *ptr = var->type_annotation.as(); + if (ptr == nullptr) { + return std::nullopt; + } + const auto *prim = ptr->element_type.as(); + if (prim == nullptr) { + return std::nullopt; + } + DataType dtype = prim->dtype; + if (dtype.is_void() || dtype.is_handle() || dtype.lanes() != 1) { + return std::nullopt; + } + return dtype; + } + + // Element type of the pointer that the store actually receives, matching + // codegen exactly (codegen_pto.cc Bind handling): + // + // * A typed Var bound to ``reinterpret(handle, int_addr)`` (T.make_tensor) + // emits ``v = pto.castptr(addr, ptr)``: this binding's own annotation + // is a real castptr and establishes the type. + // * Every other binding emits ``v = `` with no cast. A typed + // annotation on a plain alias therefore does NOT re-type the pointer; the + // defining expression's type wins. + // * ``tvm_access_ptr``/``address_of`` print through GetPointerExpr and cast + // to their own descriptor/buffer element type, overriding any outer + // annotation. + // * A pointer type re-view (``reinterpret(handle, pointer)``) prints the + // inner expression unchanged. + // + // Plain immutable SSA aliases are followed iteratively (their length is not + // a recursion budget); only pointer expression nesting consumes the + // stack-safety guard. Only the Var directly bound to the terminal + // expression can establish a type; an untyped/unknown pointer is not + // guessed as float32. + std::optional + ResolvePointerElementType(const PrimExpr &ptr, + std::unordered_set *visiting, + int nesting = 0) const { + // Stack-safety guard only; not an alias-length or acceptance limit. + constexpr int kMaxPointerExpressionNesting = 4096; + if (nesting > kMaxPointerExpressionNesting) { + return std::nullopt; + } + const VarNode *defining_var = nullptr; + PrimExpr current = ptr; + while (const auto *var = current.as()) { + if (!visiting->insert(var).second) { + return std::nullopt; // alias cycle + } + auto def_it = pointer_defs_.find(var); + if (def_it == pointer_defs_.end()) { + // Terminal Var without a binding (parameter/root): its own annotation + // is the type of the pointer it denotes. + return PointerAnnotationElementType(var); + } + defining_var = var; + current = def_it->second; + } + if (const auto *call = current.as()) { + if (call->op.same_as(builtin::tvm_access_ptr())) { + if (!call->args.empty()) { + return DecodeAccessDType(call->args[0]); + } + return std::nullopt; + } + if (call->op.same_as(builtin::address_of())) { + if (call->args.size() == 1U) { + if (const auto *load = call->args[0].as()) { + DataType dtype = load->dtype.element_of(); + if (!dtype.is_void() && !dtype.is_handle()) { + return dtype; + } + } + } + return std::nullopt; + } + if (call->op.same_as(tirx::builtin::reinterpret()) && + call->args.size() == 1U) { + if (call->args[0].dtype().is_handle()) { + // Pointer type re-view: the inner pointer's type is the emitted + // type; an outer Var annotation does not emit a cast. + return ResolvePointerElementType(call->args[0], visiting, + nesting + 1); + } + // Integer-address make_tensor binding: the castptr uses the + // annotation of the Var directly bound to this reinterpret. + if (defining_var != nullptr) { + return PointerAnnotationElementType(defining_var); + } + return std::nullopt; + } + } + return std::nullopt; + } + + StoreMatch MatchStore(const CallNode *op) const { + StoreMatch match; + // Default continuous masked form only: no mode-changing annotations. + for (const auto &[key_ref, value] : op->annotations) { + const std::string key = key_ref; + if (key == "dist_mode") { + const auto *mode = value.as(); + if (mode == nullptr || (mode->value != "continuous")) { + return match; + } + continue; + } + if (key == "group" || key == "stride" || key == "block_stride" || + key == "repeat_stride" || key == "pmode") { + return match; + } + } + // Masked ordinary form is (value, ptr, instruction_offset, mask). The + // dintlv no-mask form also has four positional arguments but carries + // dist_mode="dintlv" and was rejected above. + if (op->args.size() != 4U) { + return match; + } + const DataType value_dtype = op->args[0].dtype(); + const DataType elem = value_dtype.element_of(); + // Boolean (mask-register) values are not part of the b8/b16/b32 data + // storage contract; their backend admission stays unchanged. + if (elem.is_bool()) { + return match; + } + const int bits = elem.bits(); + if (bits != 8 && bits != 16 && bits != 32) { + return match; + } + const DataType mask_dtype = op->args[3].dtype(); + if (!mask_dtype.is_bool()) { + return match; + } + match.checked = true; + match.value_dtype = value_dtype; + match.elem_bytes = bits / 8; + // Record the immediate descriptor only for diagnostics and for the + // inherited byte offset (ResolveAddress consumes it). The instruction + // offset unit is resolved from the *resulting* pointer expression and is + // never defaulted to float32. + if (const auto *ptr_call = op->args[1].as()) { + if (ptr_call->op.same_as(builtin::tvm_access_ptr()) && + ptr_call->args.size() == 5U) { + match.has_descriptor = true; + match.descriptor_offset = ptr_call->args[2]; + if (auto dtype = DecodeAccessDType(ptr_call->args[0])) { + match.descriptor_dtype = dtype.value(); + } + } + } + match.instruction_offset = op->args[2]; + const int64_t lanes = value_dtype.lanes(); + const bool point_candidate = + lanes == 1 || lanes == 2 || lanes == 4 || lanes == 8; + const SeedInfo seed = ClassifyMaskSeed(op->args[3]); + if (point_candidate && seed.recognized && seed.full_active && + seed.mask_lanes == lanes) { + match.mode = StoreMode::kPointStores; + match.mode_name = "point"; + match.required_alignment = bits / 8; + } else { + match.mode = StoreMode::kOrdinaryMasked; + match.mode_name = "ordinary masked"; + match.required_alignment = kOrdinaryMaskedStoreAlignment; + } + return match; + } + + ByteAddress UnknownAddress(const std::string &why) const { + ByteAddress addr; + addr.kind = ByteAddress::Kind::kUnknown; + addr.error = why; + return addr; + } + + ByteAddress ResolveAddress(const PrimExpr &ptr, + std::unordered_set *visiting, + int nesting = 0) { + // Stack-safety guard only; plain alias length is walked iteratively and + // does not consume this budget. + constexpr int kMaxPointerExpressionNesting = 4096; + if (nesting > kMaxPointerExpressionNesting) { + return UnknownAddress( + "pointer expression nesting exceeds the stack-safety guard"); + } + PrimExpr current = ptr; + while (const auto *var = current.as()) { + if (!visiting->insert(var).second) { + return UnknownAddress("pointer alias cycle at `" + var->name_hint + + "`"); + } + auto def_it = pointer_defs_.find(var); + if (def_it == pointer_defs_.end()) { + break; + } + current = def_it->second; + } + if (const auto *var = current.as()) { + if (ub_roots_.count(var) != 0) { + ByteAddress addr; + addr.kind = ByteAddress::Kind::kUbRoot; + addr.root = var; + addr.base_bytes = IntImm(DataType::Int(64), 0); + addr.offset_bytes = IntImm(DataType::Int(64), 0); + return addr; + } + return UnknownAddress("pointer origin unavailable for `" + + var->name_hint + "`"); + } + if (const auto *call = current.as()) { + if (call->op.same_as(builtin::tvm_access_ptr())) { + TVM_FFI_ICHECK_EQ(call->args.size(), 5U) + << "tvm_access_ptr expects 5 arguments"; + auto dtype = DecodeAccessDType(call->args[0]); + if (!dtype.has_value()) { + return UnknownAddress("access pointer element type is unavailable"); + } + ByteAddress base = ResolveAddress(call->args[1], visiting, nesting + 1); + if (base.kind == ByteAddress::Kind::kUnknown) { + return base; + } + base.offset_bytes = analyzer_.Simplify( + base.offset_bytes + + Cast(DataType::Int(64), + call->args[2] * + IntImm(DataType::Int(64), dtype.value().bits() / 8))); + return base; + } + if (call->op.same_as(builtin::address_of())) { + if (call->args.size() != 1U) { + return UnknownAddress("address_of expects one argument"); + } + const auto *load = call->args[0].as(); + if (load == nullptr) { + return UnknownAddress("address_of expects a BufferLoad"); + } + if (load->indices.size() != 1U) { + return UnknownAddress( + "address_of with a non-flat index is unsupported"); + } + ByteAddress base = + ResolveAddress(load->buffer->data, visiting, nesting + 1); + if (base.kind == ByteAddress::Kind::kUnknown) { + return base; + } + base.offset_bytes = analyzer_.Simplify( + base.offset_bytes + + Cast(DataType::Int(64), + load->indices[0] * + IntImm(DataType::Int(64), load->dtype.bits() / 8))); + return base; + } + if (call->op.same_as(tirx::builtin::reinterpret()) && + call->args.size() == 1U && current.dtype().is_handle()) { + if (call->args[0].dtype().is_handle()) { + // Pointer type re-view: the byte address is unchanged. + return ResolveAddress(call->args[0], visiting, nesting + 1); + } + if (call->args[0].dtype().is_scalar()) { + ByteAddress addr; + addr.kind = ByteAddress::Kind::kIntegerBase; + addr.base_bytes = Cast(DataType::Int(64), call->args[0]); + addr.offset_bytes = IntImm(DataType::Int(64), 0); + return addr; + } + } + } + return UnknownAddress("unrecognized pointer expression"); + } + + Proof Prove(const ByteAddress &address, const PrimExpr &extra_bytes, + int required_alignment) { + Proof proof; + if (required_alignment <= 1) { + proof.outcome = Outcome::kAligned; + return proof; + } + if (address.kind == ByteAddress::Kind::kUnknown) { + proof.outcome = Outcome::kUnknown; + proof.error = address.error; + return proof; + } + PrimExpr base = address.base_bytes.defined() ? address.base_bytes + : IntImm(DataType::Int(64), 0); + PrimExpr total = analyzer_.Simplify(Cast(DataType::Int(64), base) + + address.offset_bytes + extra_bytes); + proof.byte_offset = total; + PrimExpr remainder = analyzer_.Simplify( + floormod(total, IntImm(DataType::Int(64), required_alignment))); + proof.remainder = remainder; + if (analyzer_.CanProve(remainder == 0)) { + proof.outcome = Outcome::kAligned; + } else if (analyzer_.CanProve(remainder != 0)) { + proof.outcome = Outcome::kMisaligned; + } else { + proof.outcome = Outcome::kUnknown; + proof.error = "cannot prove alignment for every execution"; + } + return proof; + } + + void ReportFailure(const StoreMatch &match, size_t store_index, + const ByteAddress &address, const Proof &proof) const { + std::ostringstream oss; + oss << "[VerifyPTOCodegenLegality] kernel=" << function_name_ << ", store=#" + << store_index << ": "; + if (proof.outcome == Outcome::kMisaligned) { + oss << "destination is misaligned: byte remainder=" << proof.remainder + << ", required alignment=" << match.required_alignment; + } else { + oss << "cannot prove destination alignment: " + << (proof.error.empty() ? "alignment is not provable" : proof.error) + << ", required alignment=" << match.required_alignment; + } + oss << "\n mode=" << match.mode_name + << " vstore, dtype=" << match.value_dtype + << ", element bytes=" << match.elem_bytes; + if (match.has_descriptor) { + oss << "\n descriptor offset=" << match.descriptor_offset << " (unit " + << match.descriptor_dtype << ")"; + } else { + oss << "\n descriptor offset="; + } + oss << ", instruction offset=" << match.instruction_offset << " (unit "; + if (match.pointer_dtype.has_value()) { + oss << match.pointer_dtype.value(); + } else { + oss << ""; + } + oss << ")"; + if (proof.byte_offset.defined()) { + oss << "\n normalized byte offset=" << proof.byte_offset; + } + if (address.root != nullptr) { + oss << "\n root=" << address.root->name_hint; + } else if (address.kind == ByteAddress::Kind::kIntegerBase) { + oss << "\n base byte address=" << address.base_bytes; + } else if (!address.error.empty()) { + oss << "\n address resolution: " << address.error; + } + if (match.mode == StoreMode::kPointStores) { + oss << "\n point-writeback explicit-mask vstore requires a destination " + "aligned to the stored element width (" + << match.required_alignment + << " bytes); check the starting byte address and offsets"; + } else { + oss << "\n ordinary explicit-mask vstore requires a provably 32-byte " + "aligned destination; consider a physical row stride whose byte " + "size is a multiple of 32 and check the starting " + "column/elem_offset"; + } + LOG(FATAL) << oss.str(); + } + + void CheckStore(const CallNode *op) { + StoreMatch match = MatchStore(op); + if (!match.checked) { + return; + } + const size_t store_index = next_store_index_++; + std::unordered_set visiting; + ByteAddress address = ResolveAddress(op->args[1], &visiting); + std::unordered_set dtype_visiting; + match.pointer_dtype = + ResolvePointerElementType(op->args[1], &dtype_visiting); + Proof proof; + if (match.required_alignment > 1) { + std::string dtype_error; + if (!match.pointer_dtype.has_value()) { + dtype_error = "resulting pointer element type is unavailable"; + } else { + const DataType ptr_dtype = match.pointer_dtype.value(); + const int ptr_bits = ptr_dtype.bits(); + if (ptr_bits < 8 || ptr_bits % 8 != 0) { + dtype_error = + "resulting pointer element type is not byte-addressable"; + } else { + const PrimExpr instruction_bytes = + Cast(DataType::Int(64), + op->args[2] * IntImm(DataType::Int(64), ptr_bits / 8)); + proof = Prove(address, instruction_bytes, match.required_alignment); + } + } + if (!dtype_error.empty()) { + proof.outcome = Outcome::kUnknown; + proof.error = dtype_error; + proof.byte_offset = PrimExpr(); + } + } else { + // A one-byte natural requirement holds for every byte address. + proof.outcome = Outcome::kAligned; + } + if (proof.outcome == Outcome::kAligned) { + return; + } + ReportFailure(match, store_index, address, proof); + } + + std::string function_name_; + size_t next_store_index_{0}; + std::unordered_set ub_roots_; + std::unordered_map pointer_defs_; + std::unordered_map value_defs_; + // Bind/Let definitions are immutable SSA values with unique VarNode + // identity, and every Var a definition references is already defined in an + // enclosing scope. A definition's classification therefore cannot change + // while that Var can still be referenced; once its lexical scope exits the + // VarNode cannot appear again. The cache is intentionally not part of + // AliasScope, and a Var used before its own definition is classified as a + // free parameter (non-SSA IR is not claimed as supported). + std::unordered_map var_origin_safe_; + size_t next_dag_symbol_{0}; +}; + +tvm::transform::Pass VerifyPTOCodegenLegality() { + auto pass_func = [](PrimFunc f, const IRModule &m, + const tvm::transform::PassContext &ctx) { + auto target = f->GetAttr(tvm::attr::kTarget); + if (!target.defined() || !TargetIsPTO(target.value())) { + return f; + } + std::string name = "?"; + if (auto symbol = f->GetAttr(tvm::attr::kGlobalSymbol)) { + name = symbol.value(); + } + PTOMaskedStoreAlignmentChecker checker(name); + checker.Run(f); + return f; + }; + return tvm::tirx::transform::CreatePrimFuncPass( + pass_func, 0, "tl.VerifyPTOCodegenLegality", {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tl.transform.VerifyPTOCodegenLegality", + VerifyPTOCodegenLegality); +} + +} // namespace transform +} // namespace tl +} // namespace tvm diff --git a/testing/ascend/language/pto/test_pto_feature_interaction_run.py b/testing/ascend/language/pto/test_pto_feature_interaction_run.py new file mode 100644 index 0000000000..90ed403ac5 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_feature_interaction_run.py @@ -0,0 +1,438 @@ +"""Cross-feature interaction controls: combined end-to-end device kernels. + +The dedicated per-feature suites cover each contract in isolation. This file +exercises three combinations of features in a single device run: + +* a handwritten VMI kernel loads from a nonzero UB buffer start, performs an + integer reduction, broadcasts the lane-1 result with ``vbrc(group=1)``, + stores it through an aligned masked store and copies the result back with + DMA. Automatic scheduling stays enabled (the generated MTE2->V / V->MTE3 + synchronization is asserted). Signed-negative and unsigned high-bit inputs + are compared exactly, with physical source padding and an initialized + output. + +* an auto-vectorized ``T.Parallel`` kernel reads from a nonzero source + offset, applies a cross-sign widening Cast (source-sign extend then target + reinterpret), stores through a masked aligned store and copies back with + DMA. Both a tail (E = VL-1) and a multi-chunk extent (E = 2*VL+17) are + covered for int8->uint16 and uint8->int32. The source storage is + physically padded, the destination is initialized with sentinels, and every + logical element plus the untouched tail is checked. The final alignment + checker must accept the store (compilation would fail otherwise). + +* SIMT, VMI and Parallel kernels are compiled in alternating orders in one + process. Each compilation uses fresh inputs and an independent numeric + reference, and the generated source must retain only the expected + per-kernel features, so storage or provenance state cannot leak between + compilations. Handwritten VMI is never mixed into a Parallel VF. + +Every test executes on the device and is skipped when no NPU runtime is +available. +""" + +import re +from contextlib import contextmanager + +import pytest +import torch +import torch_npu # noqa: F401,E402 (must import before tilelang/ptodsl) + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + +tirx = tvm.tirx +START = 131 +VL = 64 +REDUCE_LANES = 128 +_ELEMENT_BYTES = { + "int8": 1, + "uint8": 1, + "int16": 2, + "uint16": 2, + "int32": 4, + "uint32": 4, +} + + +def _require_npu_runtime(): + if not hasattr(torch, "npu") or not torch.npu.is_available(): + pytest.skip("torch.npu is unavailable in the current environment") + try: + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source) + except Exception as err: # noqa: BLE001 + pytest.skip(f"NPU runtime is unavailable in the current environment: {err}") + + +def _compile(func): + return tilelang.compile(func, target="pto", out_idx=-1) + + +def _run(kernel, *inputs): + device_inputs = [value.to("npu") for value in inputs] + output = kernel(*device_inputs) + torch.npu.synchronize() + return output.cpu() + + +def _physical_tail_lanes(lanes, element_bytes): + """Whole-vector load footprint: 256 bytes or one logical vector, whichever is larger.""" + return max(lanes, 256 // element_bytes) + + +def _padded_source_len(start, extent, lanes, element_bytes): + last_chunk = ((extent - 1) // lanes) * lanes + return start + last_chunk + _physical_tail_lanes(lanes, element_bytes) + + +@contextmanager +def _capture_final_module(): + captured = {} + original = ascend_transform.VerifyPTOCodegenLegality + + def factory(): + real = original() + + def run(mod): + captured["module"] = mod + return real(mod) + + return run + + ascend_transform.VerifyPTOCodegenLegality = factory + try: + yield captured + finally: + ascend_transform.VerifyPTOCodegenLegality = original + + +def _final_vload_descriptor_offsets(func): + """Lower to the pass boundary and return every VMI load descriptor offset.""" + target = normalize_pto_target("pto") + with _capture_final_module() as captured, target: + _lower(func, target=target) + module = captured["module"] + offsets = [] + + def visit(node): + if isinstance(node, tirx.Call) and hasattr(node.op, "name") and node.op.name == "tl.vmi.vload": + offsets.append(node.args[0].args[2]) + + for func in module.functions.values(): + if isinstance(func, tirx.PrimFunc): + tirx.stmt_functor.post_order_visit(func.body, visit) + assert offsets, "expected a lowered VMI load" + return offsets + + +def _constant_term(expr): + variables = [] + + def visit(node): + if isinstance(node, tirx.Var): + variables.append(node) + + tirx.stmt_functor.post_order_visit(expr, visit) + substitution = {var: tirx.IntImm(str(var.dtype), 0) for var in variables} + simplified = tvm.arith.Analyzer().simplify(tirx.stmt_functor.substitute(expr, substitution)) + assert isinstance(simplified, tirx.IntImm), simplified + return int(simplified.value) + + +# --------------------------------------------------------------------------- +# Group 1: nonzero-start VMI load, integer reduce, group-1 broadcast, +# aligned masked store, DMA; automatic scheduling enabled. +# --------------------------------------------------------------------------- +def _vmi_reduce_kernel(dtype, lanes=REDUCE_LANES, start=START): + storage = start + lanes + 64 + + # The whole padded source is filled by one aligned DMA; the VMI load then + # reads at the nonzero start, which the backend handles with + # vldas/vldus. A DMA directly into an unaligned UB region is not a legal + # MTE non-atomic access and is not what this control is about. + @T.prim_func + def main( + A: T.Tensor((storage,), dtype), + S: T.Tensor((lanes,), dtype), + C: T.Tensor((lanes,), dtype), + ): + with T.Kernel(1): + work = T.alloc_shared((storage,), dtype) + out = T.alloc_shared((lanes,), dtype) + T.copy(A, work) + T.copy(S, out) + with T.SimdVF(): + mask = T.vmi.create_mask(lanes, size=lanes) + x = T.vmi.vload(work[start], size=lanes) + reduced = T.vmi.vcmax(x, mask) + broadcast = T.vmi.vbrc(reduced, size=lanes, group=1) + T.vmi.vstore(broadcast, out[0], mask) + T.copy(out, C) + + return main + + +def _padded_reduce_input(values, torch_dtype): + padded = torch.zeros(START + REDUCE_LANES + 64, dtype=torch_dtype) + padded[START : START + REDUCE_LANES] = values + return padded + + +def _int32_reduce_input(round_index): + generator = torch.Generator().manual_seed(0xA101 + round_index) + values = torch.randint(-(2**31), 2**31 - 1, (REDUCE_LANES,), generator=generator).to(torch.int32) + values[0] = -(2**31) + values[1] = 2**31 - 1 + values[2] = -1 + values[3] = 0 + values[4] = 1 + return values + + +def _uint32_reduce_input(round_index): + generator = torch.Generator().manual_seed(0xA102 + round_index) + values = torch.randint(0, 2**31 - 1, (REDUCE_LANES,), generator=generator).to(torch.uint32) + values[0] = 2**32 - 1 + values[1] = 2**31 # INT32_MAX + 1, high bit set + values[2] = 2**31 - 1 + values[3] = 0 + return values + + +@pytest.mark.parametrize( + ("dtype", "torch_dtype", "sentinel", "make_input"), + [ + ("int32", torch.int32, -7, _int32_reduce_input), + ("uint32", torch.uint32, 7, _uint32_reduce_input), + ], + ids=["int32-signed-negative", "uint32-high-bit"], +) +def test_nonzero_start_integer_reduce_broadcast_store(dtype, torch_dtype, sentinel, make_input): + _require_npu_runtime() + kernel = _compile(_vmi_reduce_kernel(dtype)) + source = kernel.get_kernel_source() + load_line = next(line for line in source.splitlines() if "vmi.vload" in line) + assert "addptr" in load_line and str(START) in load_line, load_line + broadcast_line = next(line for line in source.splitlines() if "vmi.vbrc" in line) + assert "pto.cast(" not in broadcast_line, broadcast_line + assert "vmi.vcmax" in source, source + assert "vmi.create_mask" in source, source + assert 'set_flag("MTE2", "V"' in source, source + assert 'wait_flag("MTE2", "V"' in source, source + assert 'set_flag("V", "MTE3"' in source, source + assert 'wait_flag("V", "MTE3"' in source, source + + for round_index in range(2): + values = make_input(round_index) + sentinels = torch.full((REDUCE_LANES,), sentinel, dtype=torch_dtype) + output = _run(kernel, _padded_reduce_input(values, torch_dtype), sentinels) + expected = torch.full( + (REDUCE_LANES,), + max(int(value) for value in values.tolist()), + dtype=torch_dtype, + ) + torch.testing.assert_close(output, expected, rtol=0.0, atol=0.0) + + +# --------------------------------------------------------------------------- +# Group 2: automatic Parallel nonzero source access, cross-sign widening +# Cast (extend then reinterpret), aligned masked store, DMA; tail and +# multi-chunk extents with padded source and tail sentinels. +# --------------------------------------------------------------------------- +_CAST_CASES = [ + ("int8", "uint16", torch.int8, torch.uint16, "pto.si16", "pto.ui16", 30000), + ("uint8", "int32", torch.uint8, torch.int32, "pto.ui32", "pto.si32", -999999), +] +_EXTENTS = [VL - 1, 2 * VL + 17] + + +def _parallel_cast_kernel(src, dst, extent, lanes=VL, start=START, read_start=None): + # Continuous vloads cover whole physical vectors: the source storage + # covers 256 bytes or one logical VL, whichever is larger, after the + # last chunk base. ``read_start`` is used only by the lower-only + # descriptor control, so both kernels keep identical shapes. + source_len = _padded_source_len(start, extent, lanes, _ELEMENT_BYTES[src]) + out_len = extent + 2 * lanes + read = start if read_start is None else read_start + + @T.prim_func + def main( + A: T.Tensor((source_len,), src), + S: T.Tensor((out_len,), dst), + C: T.Tensor((out_len,), dst), + ): + with T.Kernel(1): + a = T.alloc_shared((source_len,), src) + c = T.alloc_shared((out_len,), dst) + T.copy(A, a) + T.copy(S, c) + with T.SimdVF(lanes=lanes): + for j in T.Parallel(extent): + c[j] = T.cast(a[read + j], dst) + T.copy(c, C) + + return main + + +def _source_values(src, torch_src, source_len): + if src == "int8": + values = torch.arange(source_len) * 37 % 251 - 125 + return values.to(torch_src) + values = torch.arange(source_len) * 53 % 256 + return values.to(torch_src) + + +@pytest.mark.parametrize( + ("src", "dst", "torch_src", "torch_dst", "intermediate", "target", "sentinel"), + _CAST_CASES, + ids=["int8_to_uint16", "uint8_to_int32"], +) +@pytest.mark.parametrize("extent", _EXTENTS, ids=["tail", "multi_chunk"]) +def test_parallel_nonzero_source_cross_sign_cast_store(src, dst, torch_src, torch_dst, intermediate, target, sentinel, extent): + _require_npu_runtime() + kernel = _compile(_parallel_cast_kernel(src, dst, extent)) + source = kernel.get_kernel_source() + load_line = next(line for line in source.splitlines() if "vmi.vload" in line) + assert "addptr" in load_line, load_line + # Exact final-descriptor check: lower the same-shape kernel once with the + # source read at 0 and once at START, then verify that the constant part of + # the final merged descriptor offset differs by exactly START elements. + actual_offsets = _final_vload_descriptor_offsets(_parallel_cast_kernel(src, dst, extent)) + control_offsets = _final_vload_descriptor_offsets(_parallel_cast_kernel(src, dst, extent, read_start=0)) + assert len(actual_offsets) == len(control_offsets) + for actual, control in zip(actual_offsets, control_offsets): + assert _constant_term(actual) - _constant_term(control) == START + cast_lines = [ + line.strip() + for line in source.splitlines() + if ("vmi.vcvt" in line or "vinterpret_cast" in line) and "vmi.vbrc(" not in line and "vmi.vci(" not in line + ] + assert len(cast_lines) == 2, cast_lines + assert "vmi.vcvt" in cast_lines[0], cast_lines + assert f"to_dtype={intermediate}" in cast_lines[0], cast_lines[0] + assert "vinterpret_cast" in cast_lines[1], cast_lines + assert f"to_dtype={target}" in cast_lines[1], cast_lines[1] + assert "vmi.create_mask" in source, source + store_lines = [line for line in source.splitlines() if "vmi.vstore" in line] + assert store_lines, source + for line in store_lines: + assert "mask" in line, line + + source_len = _padded_source_len(START, extent, VL, _ELEMENT_BYTES[src]) + out_len = extent + 2 * VL + values = _source_values(src, torch_src, source_len) + sentinels = torch.full((out_len,), sentinel, dtype=torch_dst) + output = _run(kernel, values, sentinels) + logical = values[START : START + extent].to(torch.int64) + if dst == "uint16": + expected = (logical & 0xFFFF).to(torch.uint16) + else: + expected = logical.to(torch.int32) + torch.testing.assert_close(output[:extent], expected, rtol=0.0, atol=0.0) + torch.testing.assert_close(output[extent:], sentinels[extent:], rtol=0.0, atol=0.0) + + +# --------------------------------------------------------------------------- +# Group 3: SIMT and VMI/Parallel compilation succession and interleaving; +# independent numeric checks for every compilation. +# --------------------------------------------------------------------------- +def _simt_copy_kernel(dtype="int32", lanes=8): + # The destination is seeded with a distinct sentinel, so a missing SIMT + # local store cannot leave the expected output in place. + @T.prim_func + def main( + A: T.Tensor((64,), dtype), + S: T.Tensor((64,), dtype), + C: T.Tensor((64,), dtype), + ): + with T.Kernel(1): + src = T.alloc_shared((64,), dtype) + dst = T.alloc_shared((64,), dtype) + T.copy(A, src) + T.copy(S, dst) + with T.SimtVF(threads=1): + local = T.alloc_local((64,), dtype) + local[T.Ramp(0, 1, lanes)] = src[T.Ramp(0, 1, lanes)] + dst[T.Ramp(0, 1, lanes)] = local[T.Ramp(0, 1, lanes)] + T.copy(dst, C) + + return main + + +def _compile_simt(round_index): + kernel = _compile(_simt_copy_kernel()) + source = kernel.get_kernel_source() + assert "vmi.vload" not in source, source + assert "vmi.vstore" not in source, source + # Same-width signless local backing plus explicit signed logical + # adaptation on the load from and store back to the shared buffer. + assert "pto.alloc_buffer((64,), pto.i32)" in source, source + assert "pto.alloc_buffer((64,), pto.si32)" not in source, source + assert re.search( + r"pto\.cast\(pto\.load\(local, 0, contiguous=8\), pto\.si32\)", + source, + ), source + assert re.search(r"pto\.cast\(pto\.load\(.*contiguous=8\), pto\.i32\)", source), source + generator = torch.Generator().manual_seed(0x51A1 + round_index) + values = torch.randint(-1000, 1000, (64,), generator=generator).to(torch.int32) + sentinels = torch.full((64,), -777777, dtype=torch.int32) + output = _run(kernel, values, sentinels) + torch.testing.assert_close(output[:8], values[:8], rtol=0.0, atol=0.0) + torch.testing.assert_close(output[8:], sentinels[8:], rtol=0.0, atol=0.0) + + +def _compile_vmi(round_index): + kernel = _compile(_vmi_reduce_kernel("int32")) + source = kernel.get_kernel_source() + assert "vmi.vload" in source and "vmi.vstore" in source, source + assert "vmi.vbrc" in source, source + broadcast_line = next(line for line in source.splitlines() if "vmi.vbrc" in line) + assert "pto.cast(" not in broadcast_line, broadcast_line + values = _int32_reduce_input(round_index) + sentinels = torch.full((REDUCE_LANES,), -7, dtype=torch.int32) + output = _run(kernel, _padded_reduce_input(values, torch.int32), sentinels) + expected = torch.full( + (REDUCE_LANES,), + max(int(value) for value in values.tolist()), + dtype=torch.int32, + ) + torch.testing.assert_close(output, expected, rtol=0.0, atol=0.0) + + +def _compile_parallel(round_index): + extent = VL - 1 + kernel = _compile(_parallel_cast_kernel("int8", "uint16", extent)) + source = kernel.get_kernel_source() + assert "vinterpret_cast" in source and "vmi.vload" in source, source + assert "vmi.vstore" in source, source + source_len = _padded_source_len(START, extent, VL, _ELEMENT_BYTES["int8"]) + out_len = extent + 2 * VL + values = _source_values("int8", torch.int8, source_len) + sentinels = torch.full((out_len,), 30000, dtype=torch.uint16) + output = _run(kernel, values, sentinels) + expected = (values[START : START + extent].to(torch.int64) & 0xFFFF).to(torch.uint16) + torch.testing.assert_close(output[:extent], expected, rtol=0.0, atol=0.0) + torch.testing.assert_close(output[extent:], sentinels[extent:], rtol=0.0, atol=0.0) + + +def test_simt_vmi_parallel_compile_sequences_do_not_leak_state(): + _require_npu_runtime() + sequences = [ + ("simt", "vmi", "parallel"), + ("parallel", "vmi", "simt"), + ("vmi", "simt", "parallel"), + ] + builders = { + "simt": _compile_simt, + "vmi": _compile_vmi, + "parallel": _compile_parallel, + } + for round_index, sequence in enumerate(sequences): + for name in sequence: + builders[name](round_index) diff --git a/testing/ascend/language/pto/test_pto_masked_store_alignment.py b/testing/ascend/language/pto/test_pto_masked_store_alignment.py new file mode 100644 index 0000000000..97fa28b24f --- /dev/null +++ b/testing/ascend/language/pto/test_pto_masked_store_alignment.py @@ -0,0 +1,883 @@ +"""Explicit-mask PTO VMI store alignment verification. + +The read-only ``VerifyPTOCodegenLegality`` PrimFuncPass runs at the end of +the PTO pipeline, after ``LowerDeviceKernelLaunch``, the buffer-version +rewrites and UB merging. It consumes the final device TIR, resolves the real +destination byte address and proves the alignment required by PTO store +operations: + +* an ordinary default/continuous explicit-mask ``tl.vmi.vstore`` requires a + provably 32-byte-aligned destination; +* the implicit per-point writeback (``1PT_B8/B16/B32``) selected when the + value has 1/2/4/8 logical lanes and the store mask is a recognized + full-active seed with the same logical lane count requires only the stored + element's natural alignment (1/2/4 bytes). + +Covered here: + +* matcher boundaries: default and continuous full-mask stores pass; stores + with no mask, a ``dintlv``/``group`` distribution annotation or a scatter + op stay untouched; the diagnostics distinguish a provably misaligned + destination from one whose alignment cannot be proven; +* the instruction offset's actual resulting-pointer element unit: inherited + offsets stay in bytes and are counted exactly once across typed alias + layers and merged UB arena descriptors; +* mask-seed precision: a recognized full-active seed with a matching lane + count lets a 1/2/4/8-lane store use point writeback with only natural + element alignment; partial, dynamic and arbitrary all-true arithmetic + seeds stay ordinary; +* scoped Analyzer facts: loop ranges and non-zero loop mins; +* the final-pipeline integration: the pass sees the final merged UB arena, + descriptor and version-slot addresses; compact ``(3,131)`` FP16/FP32 and + flat/column-offset stores fail before codegen; padded, row-0-only, tail + and dynamic-offset controls pass; the point-writeback reductions, the + buffer-version layout and a non-32B-aligned FP16/FP32 load with a + physically padded source feeding an aligned store compile and run on the + device with numeric checks; dynamic-mask and single-active stores are + rejected. + +Every accepted direct-pass case asserts the returned IR is structurally +identical to its input. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager + +import pytest +import torch +import torch_npu # noqa: F401,E402 (must import before tilelang/ptodsl) + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + +tirx = tvm.tirx +_PASS = tvm.ffi.get_global_func("tl.transform.VerifyPTOCodegenLegality")() +_PTO_TARGET = normalize_pto_target("pto") +_MARKER = "[VerifyPTOCodegenLegality]" +_INTERNAL_ERROR = tvm.error.InternalError +_OP_VSTORE = tvm.ir.Op.get("tl.vmi.vstore") +_OP_VSCATTER = tvm.ir.Op.get("tl.vmi.vscatter") +_OP_CREATE_MASK = tvm.ir.Op.get("tl.vmi.create_mask") + + +def _require_npu_runtime(): + if not hasattr(torch, "npu") or not torch.npu.is_available(): + pytest.skip("torch.npu is unavailable in the current environment") + try: + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source) + except Exception as err: # noqa: BLE001 + pytest.skip(f"NPU runtime is unavailable in the current environment: {err}") + + +# --------------------------------------------------------------------------- +# Raw TIR builders (direct read-only pass tests) +# --------------------------------------------------------------------------- +def _imm(dtype, value): + return tirx.IntImm(dtype, value) + + +def _float_value(dtype, lanes): + base = tirx.FloatImm(dtype, 1.0) + return base if lanes == 1 else tirx.Broadcast(base, lanes) + + +def _int_value(dtype, lanes): + base = tirx.IntImm(dtype, 1) + return base if lanes == 1 else tirx.Broadcast(base, lanes) + + +def _value(dtype, lanes): + if str(dtype).startswith(("float", "bfloat")): + return _float_value(dtype, lanes) + return _int_value(dtype, lanes) + + +def _seed(lanes, active=None): + active = lanes if active is None else active + return tirx.Call( + f"boolx{lanes}", + _OP_CREATE_MASK, + [active if not isinstance(active, int) else _imm("int32", active)], + annotations={"size": _imm("int32", lanes)}, + ) + + +def _store(value, ptr, offset, mask, annotations=None): + return tirx.Evaluate(tirx.Call(value.dtype, _OP_VSTORE, [value, ptr, offset, mask], annotations=annotations)) + + +def _root(name="root", dtype="uint8", size=8192, scope="shared.dyn"): + var = tirx.Var(name, tvm.ir.PointerType(tvm.ir.PrimType(dtype), scope)) + buffer = tirx.decl_buffer((size,), dtype, data=var, scope=scope) + return var, tirx.AllocBuffer(buffer) + + +def _typed_int_base(name, dtype, base, scope="shared.dyn"): + var = tirx.Var(name, tvm.ir.PointerType(tvm.ir.PrimType(dtype), scope)) + bind = tirx.Bind(var, tirx.reinterpret("handle", tirx.IntImm("int64", base))) + return var, bind + + +def _access(dtype, data, offset, extent=64): + return tirx.tvm_access_ptr( + tirx.type_annotation(dtype), + data, + offset if not isinstance(offset, int) else _imm("int32", offset), + _imm("int32", extent), + _imm("int32", 2), + ) + + +def _func(body, params=(), target=None): + return tirx.PrimFunc(params, body).with_attr("target", target if target is not None else _PTO_TARGET) + + +def _run(func): + return _PASS(tvm.IRModule({"main_kernel": func})) + + +def _accept(func): + out = _run(func) + assert tvm.ir.structural_equal(out["main_kernel"], func), "the alignment pass must return the input IR unchanged" + return out + + +def _reject(func): + with pytest.raises(_INTERNAL_ERROR) as excinfo: + _run(func) + text = str(excinfo.value) + assert _MARKER in text, text + return text + + +def _compile_ptoas_artifacts(source): + """Compile PTODSL to PTO, VPTO and LLVM artifacts.""" + import importlib.util + import re + import subprocess + import tempfile + import uuid + + decorator = re.search(r'@pto\.jit\(name="([^"]+)"', source) + name = decorator.group(1) if decorator is not None else "main" + work = tempfile.mkdtemp(prefix="masked-store-alignment-", dir=os.environ.get("TMPDIR") or None) + source_path = os.path.join(work, "kernel.ptodsl.py") + pto_path = os.path.join(work, "kernel.pto") + vpto_path = os.path.join(work, "kernel.vpto.mlir") + llvm_path = os.path.join(work, "kernel.ll") + with open(source_path, "w", encoding="utf-8") as file: + file.write(source) + spec = importlib.util.spec_from_file_location("masked_store_alignment_" + uuid.uuid4().hex, source_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + pto_text = getattr(module, name).compile().mlir_text() + with open(pto_path, "w", encoding="utf-8") as file: + file.write(pto_text) + vpto_result = subprocess.run( + [ + "ptoas", + "--pto-arch=a5", + "--pto-backend=vpto", + "--emit-vpto", + pto_path, + "-o", + vpto_path, + ], + text=True, + capture_output=True, + ) + llvm_result = subprocess.run( + [ + "ptoas", + "--pto-arch=a5", + "--pto-backend=vpto", + "--emit-vpto-llvm-ir", + pto_path, + "-o", + llvm_path, + ], + text=True, + capture_output=True, + ) + + def read(path): + if not os.path.exists(path): + return "" + with open(path, encoding="utf-8") as file: + return file.read() + + stderr = "\n".join(error for error in (vpto_result.stderr, llvm_result.stderr) if error) + return ( + vpto_result.returncode or llvm_result.returncode, + stderr, + pto_text, + read(vpto_path), + read(llvm_path), + ) + + +def _ordinary_case(dtype="float32", lanes=64, offset=0): + """Ordinary masked store through a UB root; returns the PrimFunc.""" + root, alloc = _root(dtype="uint8") + value = _value(dtype, lanes) + mask = _seed(lanes) + ptr = _access(dtype, root, offset) + return _func(tirx.SeqStmt([alloc, _store(value, ptr, 0, mask)])) + + +# --------------------------------------------------------------------------- +# Matcher boundaries +# --------------------------------------------------------------------------- +class TestMatcher: + def test_default_and_continuous_full_mask_pass(self): + _accept(_ordinary_case(offset=8)) + root, alloc = _root() + value = _value("float32", 64) + func = _func( + tirx.SeqStmt( + [ + alloc, + _store( + value, + _access("float32", root, 8), + _imm("int32", 0), + _seed(64), + annotations={"dist_mode": tirx.StringImm("continuous")}, + ), + ] + ) + ) + _accept(func) + + def test_ordinary_store_misaligned_and_unknown_reasons_differ(self): + misaligned = _reject(_ordinary_case(offset=7)) + assert "destination is misaligned" in misaligned + assert "byte remainder" in misaligned + assert "required alignment=32" in misaligned + root, alloc = _root() + unknown_offset = tirx.Var("dynamic_offset", "int32") + func = _func( + tirx.SeqStmt( + [ + alloc, + _store( + _value("float32", 64), + _access("float32", root, unknown_offset), + _imm("int32", 0), + _seed(64), + ), + ] + ), + params=[unknown_offset], + ) + unknown = _reject(func) + assert "cannot prove destination alignment" in unknown + assert "destination is misaligned" not in unknown + assert "cannot prove alignment for every execution" in unknown + + @pytest.mark.parametrize("case", ["no_mask", "dintlv", "group", "vscatter"]) + def test_unmatched_forms_are_untouched(self, case): + root, alloc = _root() + ptr = _access("float32", root, 1) + value = _value("float32", 64) + mask = _seed(64) + statements = [alloc] + if case == "no_mask": + statements.append(tirx.Evaluate(tirx.Call(value.dtype, _OP_VSTORE, [value, ptr, _imm("int32", 1)]))) + elif case == "dintlv": + statements.append( + _store( + value, + ptr, + _imm("int32", 1), + mask, + annotations={"dist_mode": tirx.StringImm("dintlv")}, + ) + ) + elif case == "group": + statements.append( + _store( + value, + ptr, + _imm("int32", 1), + mask, + annotations={"group": _imm("int32", 2)}, + ) + ) + else: # vscatter + statements.append( + tirx.Evaluate( + tirx.Call( + value.dtype, + _OP_VSCATTER, + [value, ptr, _imm("int32", 1), _imm("int32", 0), mask], + ) + ) + ) + _accept(_func(tirx.SeqStmt(statements))) + + +# --------------------------------------------------------------------------- +# Instruction-offset element unit and address resolution +# --------------------------------------------------------------------------- +class TestInstructionUnit: + def test_cross_type_inherited_offset_counted_once(self): + root, alloc = _root() + base = tirx.Var("base", "handle") + alias = tirx.Var("alias", "handle") + prefix = [alloc, tirx.Bind(base, _access("uint8", root, 32)), tirx.Bind(alias, base)] + value = _value("float32", 64) + # 32 + 6*4 + 2*4 = 64 bytes: aligned only when every layer is counted + # exactly once with its own element unit. + aligned = _func( + tirx.SeqStmt( + prefix + + [ + _store( + value, + _access("float32", alias, 6), + _imm("int32", 2), + _seed(64), + ) + ] + ) + ) + _accept(aligned) + misaligned = _func( + tirx.SeqStmt( + prefix + + [ + _store( + value, + _access("float32", alias, 6), + _imm("int32", 0), + _seed(64), + ) + ] + ) + ) + text = _reject(misaligned) + assert "byte remainder=T.int64(24)" in text + + def test_merged_arena_offset_counted_once(self): + # A descriptor offset of 400 f32 elements = 1600 bytes is 32B aligned; + # 401 elements = 1604 bytes leaves a remainder of 4. + root, alloc = _root() + _accept( + _func( + tirx.SeqStmt( + [ + alloc, + _store( + _value("float32", 64), + _access("float32", root, 400), + _imm("int32", 0), + _seed(64), + ), + ] + ) + ) + ) + text = _reject( + _func( + tirx.SeqStmt( + [ + alloc, + _store( + _value("float32", 64), + _access("float32", root, 401), + _imm("int32", 0), + _seed(64), + ), + ] + ) + ) + ) + assert "byte remainder=T.int64(4)" in text + + +# --------------------------------------------------------------------------- +# Mask seeds and 1PT classification +# --------------------------------------------------------------------------- +class TestMaskSeed: + @pytest.mark.parametrize( + ("dtype", "natural"), + [("uint8", 1), ("float16", 2), ("float32", 4)], + ids=["b8", "b16", "b32"], + ) + def test_point_writeback_needs_only_natural_alignment(self, dtype, natural): + ptr, bind = _typed_int_base("mt", dtype, natural) + _accept(_func(tirx.SeqStmt([bind, _store(_value(dtype, 1), ptr, _imm("int32", 0), _seed(1))]))) + if natural > 1: + ptr, bind = _typed_int_base("mt", dtype, natural - 1) + _reject(_func(tirx.SeqStmt([bind, _store(_value(dtype, 1), ptr, _imm("int32", 0), _seed(1))]))) + + @pytest.mark.parametrize( + "mask_factory", + [ + lambda: _seed(4, active=2), + lambda: _seed(1, active=tirx.Var("active", "int32")), + lambda: tirx.Or(_seed(1), _seed(1)), + ], + ids=["partial", "dynamic_active", "all_true_arith"], + ) + def test_non_full_seeds_stay_ordinary(self, mask_factory): + root, alloc = _root() + mask = mask_factory() + lanes = mask.dtype.lanes + value = _value("float32", lanes) + text = _reject( + _func( + tirx.SeqStmt( + [ + alloc, + _store(value, _access("float32", root, 1), _imm("int32", 0), mask), + ] + ) + ) + ) + assert "ordinary masked vstore" in text + _accept( + _func( + tirx.SeqStmt( + [ + alloc, + _store(value, _access("float32", root, 8), _imm("int32", 0), mask), + ] + ) + ) + ) + + +# --------------------------------------------------------------------------- +# Scoped analysis facts +# --------------------------------------------------------------------------- +def _loop_with_offset(offset_builder): + root, alloc = _root() + loop_var = tirx.Var("i", "int32") + body = _store( + _value("float32", 64), + _access("float32", root, offset_builder(loop_var)), + _imm("int32", 0), + _seed(64), + ) + loop = tirx.For( + loop_var, + _imm("int32", 0), + _imm("int32", 4), + tirx.ForKind.SERIAL, + body, + ) + return _func(tirx.SeqStmt([alloc, loop])) + + +class TestAnalysisScope: + def test_loop_range_facts(self): + _accept(_loop_with_offset(lambda i: i * _imm("int32", 8))) + _reject(_loop_with_offset(lambda i: i)) + + def test_nonzero_loop_min_with_step(self): + i = tirx.Var("i", "int32") + offset = (i - _imm("int32", 4)) * _imm("int32", 2) + root, alloc = _root() + body = _store( + _value("float32", 64), + _access("float32", root, offset), + _imm("int32", 0), + _seed(64), + ) + loop = tirx.For( + i, + _imm("int32", 4), + _imm("int32", 8), + tirx.ForKind.SERIAL, + body, + step=_imm("int32", 4), + ) + _accept(_func(tirx.SeqStmt([alloc, loop]))) + + +# --------------------------------------------------------------------------- +# Final pipeline integration +# --------------------------------------------------------------------------- +@contextmanager +def _pass_boundary_capture(): + captured = {} + original = ascend_transform.VerifyPTOCodegenLegality + + def factory(): + real = original() + + def run(mod): + captured["module"] = mod + return real(mod) + + return run + + ascend_transform.VerifyPTOCodegenLegality = factory + try: + yield captured + finally: + ascend_transform.VerifyPTOCodegenLegality = original + + +def _lower_capture(func): + target = normalize_pto_target("pto") + with _pass_boundary_capture() as captured, target: + _lower(func, target=target) + return captured["module"] + + +def _lower_expect_error(func): + with pytest.raises(Exception) as excinfo: # noqa: B017 - kernel errors vary + tilelang.lower(func, target="pto") + text = str(excinfo.value) + assert _MARKER in text, text + assert "ptoas" not in text.lower() + return text + + +def _parallel_kernel(dtype, m, n, lanes=256, row=None): + if row is None: + + @T.prim_func + def main(A: T.Tensor((m, n), dtype), C: T.Tensor((m, n), dtype)): + with T.Kernel(1): + a = T.alloc_shared((m, n), dtype) + c = T.alloc_shared((m, n), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(m, n): + c[i, j] = a[i, j] + T.cast(3, dtype) + T.copy(c, C) + + return main + + @T.prim_func + def main(A: T.Tensor((m, n), dtype), C: T.Tensor((m, n), dtype)): + with T.Kernel(1): + a = T.alloc_shared((m, n), dtype) + c = T.alloc_shared((m, n), dtype) + T.copy(A[row, :], a[row, :]) + with T.SimdVF(lanes=lanes): + for j in T.Parallel(n): + c[row, j] = a[row, j] + T.cast(3, dtype) + T.copy(c[row, :], C[0, :]) + + return main + + +def _tail_kernel(dtype="float32", n=131, lanes=64): + @T.prim_func + def main(A: T.Tensor((n,), dtype), C: T.Tensor((n,), dtype)): + with T.Kernel(1): + a = T.alloc_shared((n,), dtype) + c = T.alloc_shared((n,), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for j in T.Parallel(n): + c[j] = a[j] + T.cast(3, dtype) + T.copy(c, C) + + return main + + +def _flat_kernel(m=3, n=131, dtype="float32", lanes=256): + @T.prim_func + def main(A: T.Tensor((m * n,), dtype), C: T.Tensor((m * n,), dtype)): + with T.Kernel(1): + a = T.alloc_shared((m * n,), dtype) + c = T.alloc_shared((m * n,), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(m, n): + c[i * n + j] = a[i * n + j] + T.cast(3, dtype) + T.copy(c, C) + + return main + + +def _column_offset_kernel(n=256, dtype="float32", lanes=256): + @T.prim_func + def main(A: T.Tensor((n,), dtype), C: T.Tensor((n + 1,), dtype)): + with T.Kernel(1): + a = T.alloc_shared((n,), dtype) + c = T.alloc_shared((n + 1,), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for j in T.Parallel(n): + c[j + 1] = a[j] + T.cast(3, dtype) + T.copy(c[1 : n + 1], C[1 : n + 1]) + + return main + + +_REDUCE_OFFSET = 4100 +_REDUCE_WINDOW = (_REDUCE_OFFSET // 64) * 64 +_REDUCE_INDEX = _REDUCE_OFFSET - _REDUCE_WINDOW + + +_SENTINEL_LANES = 128 + + +def _manual_reduce_kernel(group=None, size=64, dynamic=False): + out_lanes = 1 if group is None else group + window_start = _REDUCE_WINDOW + window_end = window_start + _SENTINEL_LANES + + if dynamic: + + @T.prim_func + def main( + A: T.Tensor((size,), "float32"), + S: T.Tensor((_SENTINEL_LANES,), "float32"), + active: T.int32, + C: T.Tensor((_SENTINEL_LANES,), "float32"), + ): + with T.Kernel(1): + work = T.alloc_shared((size,), "float32") + out = T.alloc_shared((8192,), "float32") + T.copy(A, work) + T.copy(S, out[window_start:window_end]) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + x = T.vmi.vload(work[0], size=size) + r = T.vmi.vcadd(x, mask, reassoc=True) + slot = T.alloc_var("int32") + slot = active + store_mask = T.vmi.create_mask(slot, size=out_lanes) + T.vmi.vstore(r, out[_REDUCE_OFFSET], store_mask) + T.copy(out[window_start:window_end], C) + + return main + + @T.prim_func + def main( + A: T.Tensor((size,), "float32"), + S: T.Tensor((_SENTINEL_LANES,), "float32"), + C: T.Tensor((_SENTINEL_LANES,), "float32"), + ): + with T.Kernel(1): + work = T.alloc_shared((size,), "float32") + out = T.alloc_shared((8192,), "float32") + T.copy(A, work) + # Independently initialized surrounding sentinels: a swapped, + # duplicated or out-of-footprint write inside this window changes + # a sentinel slot and fails the comparison. + T.copy(S, out[window_start:window_end]) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + x = T.vmi.vload(work[0], size=size) + if group is None: + r = T.vmi.vcadd(x, mask, reassoc=True) + else: + r = T.vmi.vcadd(x, mask, group=group, reassoc=True) + store_mask = T.vmi.create_mask(out_lanes, size=out_lanes) + T.vmi.vstore(r, out[_REDUCE_OFFSET], store_mask) + T.copy(out[window_start:window_end], C) + + return main + + +def _manual_single_active_ordinary(lanes=64): + @T.prim_func + def main(A: T.Tensor((lanes,), "float32"), C: T.Tensor((lanes,), "float32")): + with T.Kernel(1): + work = T.alloc_shared((lanes,), "float32") + out = T.alloc_shared((8192,), "float32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(1, size=lanes) + x = T.vmi.vload(work[0], size=lanes) + T.vmi.vstore(x, out[_REDUCE_OFFSET], mask) + T.copy( + out[_REDUCE_WINDOW : _REDUCE_WINDOW + lanes], + C, + ) + + return main + + +def _serial_version_kernel(): + @T.prim_func + def main(A: T.Tensor((512,), "float32"), C: T.Tensor((512,), "float32")): + with T.Kernel(1): + a = T.alloc_shared((128,), "float32") + c = T.alloc_shared((128,), "float32") + T.annotate_buffer_versions({a: 2, c: 2}) + for k in T.serial(4): + T.copy(A[k * 128 : (k + 1) * 128], a) + with T.SimdVF(lanes=128): + for j in T.Parallel(128): + c[j] = a[j] + T.cast(3, "float32") + T.copy(c, C[k * 128 : (k + 1) * 128]) + + return main + + +def _non32b_load_kernel(dtype): + source_len = 131 + 256 + + @T.prim_func + def main(S: T.Tensor((source_len,), dtype), C: T.Tensor((256,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((source_len + 128,), dtype) + out = T.alloc_shared((256,), dtype) + T.copy(S, work[0:source_len]) + with T.SimdVF(lanes=256): + mask = T.vmi.create_mask(256, size=256) + x = T.vmi.vload(work[131], size=256) + T.vmi.vstore(x, out[0], mask) + T.copy(out[0:256], C) + + return main + + +def _dynamic_offset_kernel(mult): + @T.prim_func + def main( + A: T.Tensor((64,), "float32"), + k: T.int32, + C: T.Tensor((64,), "float32"), + ): + with T.Kernel(1): + work = T.alloc_shared((1024,), "float32") + T.copy(A, work[0:64]) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + x = T.vmi.vload(work[0], size=64) + T.vmi.vstore(x, work[k * mult], mask) + T.copy(work[0:64], C) + + return main + + +class TestFinalPipeline: + @pytest.mark.parametrize("dtype", ["float32", "float16"]) + def test_compact_three_by_131_fails_before_codegen(self, dtype): + text = _lower_expect_error(_parallel_kernel(dtype, 3, 131)) + assert "ordinary masked vstore" in text + assert "required alignment=32" in text + assert "root=" in text + assert "normalized byte offset=" in text + + def test_flat_same_address_fails(self): + text = _lower_expect_error(_flat_kernel()) + assert "cannot prove destination alignment" in text + + def test_nonzero_column_offset_fails(self): + text = _lower_expect_error(_column_offset_kernel()) + assert "destination is misaligned" in text + + @pytest.mark.parametrize( + ("name", "kernel"), + [ + ("padded", _parallel_kernel("float32", 3, 136)), + ("row0_only", _parallel_kernel("float32", 3, 131, row=0)), + ("single_row_tail_vl64", _tail_kernel()), + ], + ) + def test_aligned_controls_full_compile(self, name, kernel): + compiled = tilelang.compile(kernel, target="pto", out_idx=-1) + source = compiled.get_kernel_source() + assert "vmi.vstore" in source, name + + def test_dynamic_offset_controls(self): + aligned = _dynamic_offset_kernel(8) + module = _lower_capture(aligned) + assert module is not None + tilelang.compile(aligned, target="pto", out_idx=-1) + text = _lower_expect_error(_dynamic_offset_kernel(1)) + assert "cannot prove destination alignment" in text + + @pytest.mark.parametrize( + ("group", "size", "lanes"), + [(None, 64, 1), (2, 128, 2), (4, 256, 4), (8, 256, 8)], + ) + def test_point_writeback_compiles_and_runs(self, group, size, lanes): + _require_npu_runtime() + kernel = tilelang.compile( + _manual_reduce_kernel(group=group, size=size), + target="pto", + out_idx=-1, + ) + source = kernel.get_kernel_source() + assert "vmi.vstore" in source + values = torch.arange(size, dtype=torch.float32) + sentinels = torch.arange(_SENTINEL_LANES, dtype=torch.float32) * -3.0 - 500.0 + output = kernel(values.to("npu"), sentinels.to("npu")) + torch.npu.synchronize() + got = output.cpu() + # The group reduction produces one result per consecutive + # size/lanes-element chunk. Every other slot in the initialized + # 128-element window must retain its sentinel. Swapped, duplicated or + # out-of-footprint writes must fail even when the total sum matches. + chunk = size // lanes + expected_slots = [float(values[index * chunk : (index + 1) * chunk].sum()) for index in range(lanes)] + expected = sentinels.clone() + expected[_REDUCE_INDEX : _REDUCE_INDEX + lanes] = torch.tensor(expected_slots, dtype=torch.float32) + torch.testing.assert_close(got, expected, rtol=0.0, atol=1e-3) + + def test_actual_version_layout_address_is_checked_and_runs(self): + _require_npu_runtime() + module = _lower_capture(_serial_version_kernel()) + stores = [] + + def visit(node): + if isinstance(node, tirx.Call) and hasattr(node.op, "name") and node.op.name == "tl.vmi.vstore": + stores.append(node) + + tirx.stmt_functor.post_order_visit(module["main_kernel"].body, visit) + assert stores, "expected a versioned VMI store" + for store in stores: + descriptor_offset = store.args[1].args[2] + version_terms = [] + + def offset_visit(node, terms=version_terms): + if isinstance(node, tirx.FloorMod): + terms.append(node) + + tirx.stmt_functor.post_order_visit(descriptor_offset, offset_visit) + assert version_terms, f"expected the final descriptor to contain the version-slot index, got {descriptor_offset}" + kernel = tilelang.compile(_serial_version_kernel(), target="pto", out_idx=-1) + values = torch.arange(512, dtype=torch.float32) + output = kernel(values.to("npu")) + torch.npu.synchronize() + torch.testing.assert_close(output.cpu(), values + 3.0, rtol=0.0, atol=0.0) + + @pytest.mark.parametrize("dtype", ["float16", "float32"]) + def test_non32b_aligned_load_with_padded_source_runs(self, dtype): + _require_npu_runtime() + kernel = tilelang.compile(_non32b_load_kernel(dtype), target="pto", out_idx=-1) + source = kernel.get_kernel_source() + load_line = next(line for line in source.splitlines() if "vmi.vload" in line) + assert "131" in load_line and "addptr" in load_line, load_line + returncode, stderr, pto_text, vpto_text, llvm_text = _compile_ptoas_artifacts(source) + assert returncode == 0, stderr + assert "pto.vldas" in vpto_text, vpto_text + assert "pto.vldus" in vpto_text, vpto_text + assert "131" in pto_text, pto_text + assert "llvm.hivm.vldas" in llvm_text, llvm_text + assert "llvm.hivm.vldus" in llvm_text, llvm_text + values = torch.arange(387, dtype=torch.float32).to(getattr(torch, dtype)) + output = kernel(values.to("npu")) + torch.npu.synchronize() + expected = values[131:387].float().cpu() + torch.testing.assert_close(output.cpu().float(), expected, rtol=0.0, atol=0.0) + + def test_dynamic_and_single_active_stores_reject(self): + text = _lower_expect_error(_manual_reduce_kernel(dynamic=True)) + assert "ordinary masked vstore" in text + assert "required alignment=32" in text + text = _lower_expect_error(_manual_single_active_ordinary()) + assert "ordinary masked vstore" in text + assert "required alignment=32" in text diff --git a/tilelang/ascend/pipeline.py b/tilelang/ascend/pipeline.py index 7426db1248..862237e126 100644 --- a/tilelang/ascend/pipeline.py +++ b/tilelang/ascend/pipeline.py @@ -194,6 +194,11 @@ def AscendPassPipelineBody(mod: IRModule, target: Target) -> IRModule: mod = tilelang.transform.MakePackedAPI()(mod) mod = tilelang.transform.Simplify()(mod) mod = tilelang.transform.LowerDeviceKernelLaunch()(mod) + # Verify PTO requirements on the final device TIR before code generation. + # The current check proves masked-store destination alignment after all + # layout and UB address rewrites. + if target_is_pto(target): + mod = ascend_transform.VerifyPTOCodegenLegality()(mod) return mod diff --git a/tilelang/ascend/transform/__init__.py b/tilelang/ascend/transform/__init__.py index 56282ac40b..f1d8f5ba1d 100644 --- a/tilelang/ascend/transform/__init__.py +++ b/tilelang/ascend/transform/__init__.py @@ -155,6 +155,19 @@ def LegalizeParallelToPTO(): return _ffi_api.LegalizeParallelToPTO() # type: ignore +def VerifyPTOCodegenLegality(): + """Verify PTO requirements before code generation without modifying IR. + + Consumes the final device TIR after LowerDeviceKernelLaunch and the final + layout and UB rewrites. Currently checks destination alignment for + default/continuous explicit-mask VMI stores of b8/b16/b32 data: 32 bytes + for ordinary stores, or natural element alignment for recognized 1PT + writeback. Bool data is excluded. Shapes, layouts, masks and addresses are + unchanged. Functions without a PTO target pass through unchanged. + """ + return _ffi_api.VerifyPTOCodegenLegality() # type: ignore + + def VectorizeParallelToPTO(): """Rewrite verified T.Parallel loops inside PTO SIMD_VF blocks into logical VMI code. @@ -287,4 +300,5 @@ def RewriteFlagToBuf(): "UnrollLoopSkipVF", "VectorizeParallelToPTO", "VerifyParallelToPTO", + "VerifyPTOCodegenLegality", ] From eac39889226b0b403bdc112a7217f85a5127f62c Mon Sep 17 00:00:00 2001 From: Zhang Zhendong Date: Wed, 14 Oct 2026 01:52:42 +0000 Subject: [PATCH 09/10] Adapt TileLang VMI APIs to PTOAS compositions --- src/ascend/codegen/codegen_pto.cc | 48 +--------- .../pto/test_pto_vmi_memory_access.py | 90 +------------------ tilelang/ascend/language/vmi.py | 81 ++++------------- 3 files changed, 19 insertions(+), 200 deletions(-) diff --git a/src/ascend/codegen/codegen_pto.cc b/src/ascend/codegen/codegen_pto.cc index 5a570ddbf2..9b985d6cf1 100644 --- a/src/ascend/codegen/codegen_pto.cc +++ b/src/ascend/codegen/codegen_pto.cc @@ -5187,17 +5187,6 @@ CodeGenTileLangPTO::PtoValueKind CodeGenTileLangPTO::ClassifyVmiCallResult_( return PtoValueKind::kVmiPair; } if (name == "tl.vmi.vload") { - // vload is data unless the explicit dist_mode=dintlv annotation selects - // the pair-returning form. - auto it = op->annotations.find("dist_mode"); - if (it != op->annotations.end()) { - const auto *mode = (*it).second.as(); - ICHECK(mode != nullptr) - << "tl.vmi.vload dist_mode annotation must be a string"; - if (mode->value == "dintlv") { - return PtoValueKind::kVmiPair; - } - } return PtoValueKind::kVmiData; } if (name == "tl.vmi.pair_get") { @@ -5538,34 +5527,6 @@ void CodeGenTileLangPTO::PrintPtoVmiCall_(const CallNode *op, print_ptr_with_offset(op->args[1], op->args[2]); os << ", " << PrintExpr_(op->args[3]) << ", " << PrintExpr_(op->args[4]); needs_comma = true; - } else if (op_name == "tl.vmi.vstore") { - auto dist_mode_it = op->annotations.find("dist_mode"); - ObjectRef dist_mode = dist_mode_it != op->annotations.end() - ? (*dist_mode_it).second - : ObjectRef(); - const bool is_dintlv_store = - dist_mode.defined() && dist_mode.as() != nullptr && - Downcast(dist_mode)->value == "dintlv"; - if (is_dintlv_store) { - ICHECK_GE(op->args.size(), 4U) - << op_name - << " with dist_mode=dintlv expects (even, odd, ptr, offset[, mask])"; - os << "(" << PrintExpr_(op->args[0]) << ", " << PrintExpr_(op->args[1]) - << ")"; - for (size_t i = 2; i < op->args.size(); ++i) { - os << ", " << PrintExpr_(op->args[i]); - } - needs_comma = true; - } else { - for (size_t i = 0; i < op->args.size(); ++i) { - const PrimExpr &arg = op->args[i]; - if (needs_comma) { - os << ", "; - } - os << PrintExpr_(arg); - needs_comma = true; - } - } } else { for (size_t i = 0; i < op->args.size(); ++i) { const PrimExpr &arg = op->args[i]; @@ -5614,14 +5575,7 @@ void CodeGenTileLangPTO::PrintPtoVmiCall_(const CallNode *op, os << ", "; } os << key << "="; - if (op_name == "tl.vmi.vstore" && key == "dist_mode") { - // TileLang uses "dintlv" for two-input stores; PTODSL uses "intlv". - // Load operations keep "dintlv". Map the store mode only at emission; - // tuple recognition above uses the original annotation. - const auto *mode = value.as(); - ICHECK(mode != nullptr) << "tl.vmi.vstore dist_mode must be a string"; - os << "\"" << (mode->value == "dintlv" ? "intlv" : mode->value) << "\""; - } else if (op_name == "tl.vmi.vcvt" && key == "to_dtype" && + if (op_name == "tl.vmi.vcvt" && key == "to_dtype" && op->dtype.is_int()) { os << PtoSignedIntegerTypeName(op->dtype.element_of()); } else if (op_name == "tl.vmi.vload" && key == "size" && diff --git a/testing/ascend/language/pto/test_pto_vmi_memory_access.py b/testing/ascend/language/pto/test_pto_vmi_memory_access.py index 9897236894..8188ec80a8 100644 --- a/testing/ascend/language/pto/test_pto_vmi_memory_access.py +++ b/testing/ascend/language/pto/test_pto_vmi_memory_access.py @@ -115,8 +115,6 @@ def _single(artifact, op_name): def _ptr_position(call, op_name): - if op_name == "tl.vmi.vstore" and "dist_mode" in call.annotations: - return 2 return 1 if op_name == "tl.vmi.vstore" else 0 @@ -285,35 +283,6 @@ def test_elem_offset_counted_exactly_once(): # --------------------------------------------------------------------------- # Ordinary broadcast # --------------------------------------------------------------------------- -@pytest.mark.parametrize( - ("dtype", "size"), - [("int8", 256), (F16, 128), (F32, 128)], - ids=["b8", "b16", "b32"], -) -def test_ordinary_broadcast_keeps_conservative_extent(dtype, size): - """brc reads one source element; the conservative ``size`` extent is kept.""" - - # b8 full-VL broadcast at 64 reads [64,320) and stores [128,384); size - # the storage (and its initialization) for the real footprint. - @T.prim_func - def main(A: T.Tensor((512,), dtype), C: T.Tensor((512,), dtype)): - with T.Kernel(1): - work = T.alloc_shared((512,), dtype) - T.copy(A, work) - with T.SimdVF(): - mask = T.vmi.create_mask(size, size=size) - v = T.vmi.vload(work[64], size=size, dist_mode="brc") - T.vmi.vstore(v, work[128], mask) - T.copy(work, C) - - artifact = _lower_pto(main) - # The declaration is not re-expressed as a broadcast output width: it stays - # `size` address elements (a conservative over-report). - assert _descriptor(artifact, "tl.vmi.vload") == (64, size, 1, 0) - assert 'dist_mode="brc"' in artifact.kernel_source - - kernel = _compile_pto(main) - assert "pto.vmi.vload(" in kernel.get_kernel_source() # --------------------------------------------------------------------------- @@ -346,50 +315,8 @@ def main(A: T.Tensor((1024,), dtype), C: T.Tensor((1024,), dtype)): # --------------------------------------------------------------------------- # dintlv load / two-input store # --------------------------------------------------------------------------- -@pytest.mark.parametrize("size", [64, 128], ids=["single-vl", "multi-vl"]) -def test_dintlv_load_declares_both_results(size): - @T.prim_func - def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): - with T.Kernel(1): - work = T.alloc_shared((512,), F32) - T.copy(A, work) - with T.SimdVF(): - # Each result has `size` lanes: the store mask must match. - mask = T.vmi.create_mask(size, size=size) - even, odd = T.vmi.vload(work[64], size=size, dist_mode="dintlv") - T.vmi.vstore(even, work[320], mask) - T.vmi.vstore(odd, work[384], mask) - T.copy(work, C) - artifact = _lower_pto(main) - assert _descriptor(artifact, "tl.vmi.vload") == (64, 2 * size, 1, 0) - - -@pytest.mark.parametrize("mask_lanes", [0, 10, 64], ids=["all-false", "partial", "all-true"]) -def test_dintlv_store_declares_full_two_input_span(mask_lanes): - """The store may-write span is 2*N regardless of the ineffective mask.""" - - @T.prim_func - def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): - with T.Kernel(1): - work = T.alloc_shared((512,), F32) - T.copy(A, work) - with T.SimdVF(): - mask = T.vmi.create_mask(mask_lanes, size=64) - even, odd = T.vmi.vload(work[64], size=64, dist_mode="dintlv") - T.vmi.vstore((even, odd), work[192], mask, dist_mode="dintlv") - T.copy(work, C) - artifact = _lower_pto(main) - assert _descriptor(artifact, "tl.vmi.vstore") == (192, 128, 2, 0) - - source = artifact.kernel_source - vstore_line = next(ln.strip() for ln in source.splitlines() if "vmi.vstore" in ln) - # IR/API keep "dintlv"; only the PTODSL attribute is mapped to "intlv". - assert 'dist_mode="intlv"' in vstore_line, vstore_line - assert 'dist_mode="dintlv"' not in vstore_line, vstore_line - _, store_call = _single(artifact, "tl.vmi.vstore") - assert store_call.annotations["dist_mode"].value == "dintlv" # --------------------------------------------------------------------------- @@ -481,7 +408,7 @@ def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): assert _vload_descriptor(main) == (64, 192, 1, 0) -def test_group_slot_and_broadcast_use_one_element_per_group(): +def test_group_slot_uses_one_element_per_group(): @T.prim_func def slot(A: T.Tensor((768,), F32), C: T.Tensor((768,), F32)): with T.Kernel(1): @@ -498,17 +425,6 @@ def slot(A: T.Tensor((768,), F32), C: T.Tensor((768,), F32)): assert _vload_descriptor(slot) == (64, 449, 1, 0) assert _vstore_descriptor(slot) == (256, 449, 2, 0) - @T.prim_func - def group_brc(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): - with T.Kernel(1): - work = T.alloc_shared((512,), F32) - T.copy(A, work) - with T.SimdVF(): - v = T.vmi.vload(work[64], size=128, group=2, stride=128, dist_mode="brc") - T.vmi.vstore(v, work[256], group=2, stride=128) - T.copy(work, C) - - assert _vload_descriptor(group_brc) == (64, 129, 1, 0) def test_group_static_negative_stride_relocates_lowest_address(): @@ -851,9 +767,9 @@ def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): "unexpected keyword argument 'to_dtype'", ), ( - "vload-unpack", + "vload-dist-mode", lambda work, mask: T.vmi.vload(work[0], size=64, dist_mode="unpack"), - "does not support dist_mode='unpack'", + "unexpected keyword argument 'dist_mode'", ), ( "vstore-pmode", diff --git a/tilelang/ascend/language/vmi.py b/tilelang/ascend/language/vmi.py index 09337cb72d..7be9612093 100644 --- a/tilelang/ascend/language/vmi.py +++ b/tilelang/ascend/language/vmi.py @@ -291,23 +291,14 @@ def _normalize_pto_vcvt_rounding(mode, *, context: str, allowed=None): return normalized -def _validate_pto_load_modes( +def _validate_pto_memory_modes( context: str, *, - dist_mode, group, stride, block_stride, - allow_group_brc: bool, - allowed_dist_modes, ): - if dist_mode is not None and dist_mode not in allowed_dist_modes: - expected = ", ".join(repr(mode) for mode in sorted(allowed_dist_modes, key=str)) - raise TypeError(f"{context} does not support dist_mode={dist_mode!r}; expected one of {expected}") - if group is not None: - if dist_mode is not None and (not allow_group_brc or dist_mode != "brc"): - raise TypeError(f"{context} does not allow dist_mode together with group") if block_stride is not None: raise TypeError(f"{context} does not allow block_stride together with group") if stride is None: @@ -315,8 +306,6 @@ def _validate_pto_load_modes( return if block_stride is not None: - if dist_mode is not None: - raise TypeError(f"{context} does not allow dist_mode together with block_stride") if stride is not None: raise TypeError(f"{context} does not allow stride together with block_stride") return @@ -421,7 +410,7 @@ def _whole_buffer_plan(source, *, fallback_extent): return span, buffer, relative -def _group_access_plan(source, *, lanes, group, stride, dist_mode, is_load, elem_dtype): +def _group_access_plan(source, *, lanes, group, stride, is_load, elem_dtype): """May-access span of an equal-distance group access (C.2/C.3.1). The first group starts at the original access address ``A``. For stride @@ -443,7 +432,7 @@ def _group_access_plan(source, *, lanes, group, stride, dist_mode, is_load, elem # Slot (one result/input element per group) and group broadcast read or # write exactly one address element per group. Ordinary groups cover # ceil(N/G) elements; the ceiling keeps a non-divisible shape conservative. - if dist_mode == "brc" or lanes <= group_const: + if lanes <= group_const: per_group = 1 else: per_group = (lanes + group_const - 1) // group_const @@ -490,7 +479,7 @@ def _block_access_plan(source, *, lanes, elem_dtype, block_stride): return extent, None, None -def _vmi_access_plan(source, *, lanes, elem_dtype, dist_mode, group, stride, block_stride, is_load): +def _vmi_access_plan(source, *, lanes, elem_dtype, group, stride, block_stride, is_load): """Compute (descriptor extent, descriptor base, instruction offset). ``descriptor base`` is ``None`` when the original source is already the @@ -500,17 +489,12 @@ def _vmi_access_plan(source, *, lanes, elem_dtype, dist_mode, group, stride, blo """ if not isinstance(source, (BufferLoad, Buffer, BufferRegion)): return None, None, None - if dist_mode == "dintlv": - # Two interleaved inputs/outputs: the complete 2*N span, regardless of - # the (ineffective) store mask. - return 2 * lanes, None, None if group is not None: return _group_access_plan( source, lanes=lanes, group=group, stride=stride, - dist_mode=dist_mode, is_load=is_load, elem_dtype=elem_dtype, ) @@ -520,11 +504,6 @@ def _vmi_access_plan(source, *, lanes, elem_dtype, dist_mode, group, stride, blo # Genuine FP4 descriptors stay in logical FP4 element units; # RewriteFp4ToFp4x2 performs the single physical packing. return lanes, None, None - if dist_mode == "brc": - # An ordinary broadcast reads one source element; keeping the old - # conservative `size` declaration is allowed and does not change the - # instruction semantics. - return lanes, None, None if is_load: bits = int(getattr(elem_dtype, "bits", 0) or 0) if bits <= 0: @@ -648,7 +627,6 @@ def vload( size, stride=None, block_stride=None, - dist_mode=None, group=None, loc=None, ip=None, @@ -663,14 +641,11 @@ def vload( read-only lanes cannot fault; they carry undefined values that the compute and store masks must keep out of valid outputs. """ - _validate_pto_load_modes( + _validate_pto_memory_modes( "T.vmi.vload(...)", - dist_mode=dist_mode, group=group, stride=stride, block_stride=block_stride, - allow_group_brc=True, - allowed_dist_modes={None, "continuous", "dintlv", "brc"}, ) if size is None: raise TypeError("T.vmi.vload(...) requires size") @@ -679,9 +654,7 @@ def vload( if is_packed_fp4: if size % 2: raise ValueError("T.vmi.vload(...) requires an even FP4 lane count for packed FP4 storage") - if dist_mode == "dintlv": - raise TypeError('T.vmi.vload(...) does not support packed FP4 with dist_mode="dintlv"') - if group is not None or block_stride is not None or dist_mode == "brc": + if group is not None or block_stride is not None: raise TypeError("T.vmi.vload(...) only supports contiguous packed FP4 loads") result_dtype = _vector_result_dtype(source, size, elem_dtype=source_elem) @@ -689,7 +662,6 @@ def vload( source, lanes=size, elem_dtype=source_elem, - dist_mode=dist_mode, group=group, stride=stride, block_stride=block_stride, @@ -709,27 +681,22 @@ def vload( "size": size, "stride": stride, "block_stride": block_stride, - "dist_mode": dist_mode, "group": group, "loc": loc, "ip": ip, } - call = _call_vmi("vload", result_dtype, ptr, extra_offset, **attrs) - if dist_mode == "dintlv": - return _wrap_pair(call) - return call + return _call_vmi("vload", result_dtype, ptr, extra_offset, **attrs) @_scope_guarded def vstore( - values, + value, destination, offset=None, mask=None, *, stride=None, block_stride=None, - dist_mode=None, group=None, loc=None, ip=None, @@ -738,38 +705,22 @@ def vstore( if mask is None and _looks_like_mask(offset): mask = offset offset = None - _validate_pto_load_modes( + _validate_pto_memory_modes( "T.vmi.vstore(...)", - dist_mode=dist_mode, group=group, stride=stride, block_stride=block_stride, - allow_group_brc=False, - allowed_dist_modes={None, "continuous", "dintlv"}, ) if group is not None and mask is not None: raise TypeError("T.vmi.vstore(...) group mode does not take a mask operand") - value_dtype = None - if dist_mode == "dintlv": - if isinstance(values, VmiPair): - value_dtype = _element_dtype_of(values) - values = tuple(values) - elif not _is_sequence(values) or len(values) != 2: - raise TypeError('T.vmi.vstore(...) with dist_mode="dintlv" requires an (even, odd) pair') - else: - _require_same_vreg_type(values[0], values[1], context="T.vmi.vstore(...)", allow_packed_fp4=True) - elif isinstance(values, VmiPair) or _is_sequence(values): - raise TypeError('T.vmi.vstore(...) expects a single VMI vector unless dist_mode="dintlv"') - - store_lanes = _lanes_of(values[0]) if _is_sequence(values) else _lanes_of(values) - value_args = list(values) if _is_sequence(values) else [values] - if value_dtype is None: - value_dtype = _element_dtype_of(values[0]) if _is_sequence(values) else _element_dtype_of(values) + if isinstance(value, VmiPair) or _is_sequence(value): + raise TypeError("T.vmi.vstore(...) expects a single VMI vector; store each result separately") + + store_lanes = _lanes_of(value) + value_dtype = _element_dtype_of(value) if str(value_dtype) == "float4_e2m1fn": if store_lanes % 2: raise ValueError("T.vmi.vstore(...) requires an even FP4 lane count for packed FP4 storage") - if dist_mode == "dintlv": - raise TypeError('T.vmi.vstore(...) does not support packed FP4 with dist_mode="dintlv"') if group is not None or block_stride is not None: raise TypeError("T.vmi.vstore(...) only supports contiguous packed FP4 stores") elif mask is not None: @@ -779,7 +730,6 @@ def vstore( destination, lanes=store_lanes, elem_dtype=value_dtype, - dist_mode=dist_mode, group=group, stride=stride, block_stride=block_stride, @@ -809,13 +759,12 @@ def vstore( return _call_vmi( "vstore", "void", - *value_args, + value, ptr, extra_offset, mask, stride=stride, block_stride=block_stride, - dist_mode=dist_mode, group=group, loc=loc, ip=ip, From 9b28c67de738f96d23da55e4dd39ed7cbee225c9 Mon Sep 17 00:00:00 2001 From: mouliangyu Date: Thu, 22 Oct 2026 02:52:51 +0000 Subject: [PATCH 10/10] test(pto): cover explicit VMI broadcast-load composition --- .../pto/test_pto_vmi_broadcast_composition.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 testing/ascend/language/pto/test_pto_vmi_broadcast_composition.py diff --git a/testing/ascend/language/pto/test_pto_vmi_broadcast_composition.py b/testing/ascend/language/pto/test_pto_vmi_broadcast_composition.py new file mode 100644 index 0000000000..4e45f297c9 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_vmi_broadcast_composition.py @@ -0,0 +1,51 @@ +"""Source-level coverage for explicit VMI broadcast-load compositions.""" + +import ast + +import pytest + +import tilelang.ascend.language as T +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + + +@pytest.mark.pto +@pytest.mark.parametrize("group,stride", [(None, None), (2, 1), (4, 3)]) +def test_pto_codegen_preserves_vload_vbrc_composition(group, stride): + load_size = 1 if group is None else group + + @T.prim_func + def func(A: T.Tensor((64,), "float32"), B: T.Tensor((64,), "float32")): + with T.Kernel(1): + a_ub = T.alloc_shared((64,), "float32") + b_ub = T.alloc_shared((64,), "float32") + T.copy(A, a_ub) + with T.SimdVF(): + loaded = T.vmi.vload(a_ub[4], size=load_size, group=group, stride=stride) + value = T.vmi.vbrc(loaded, size=64, group=group) + mask = T.vmi.create_mask(64, size=64) + T.vmi.vstore(value, b_ub[0], mask) + T.copy(b_ub, B) + + target = normalize_pto_target("pto") + with target: + source = _lower(func, target=target).kernel_source + compile(source, "", "exec") + load_line = next(line for line in source.splitlines() if "pto.vmi.vload(" in line) + broadcast_line = next(line for line in source.splitlines() if "pto.vmi.vbrc(" in line) + load = ast.parse(load_line.strip()).body[0] + broadcast = ast.parse(broadcast_line.strip()).body[0] + assert isinstance(load, ast.Assign) + assert isinstance(broadcast, ast.Assign) + assert broadcast.value.args[0].id == load.targets[0].id + assert ast.literal_eval(load.value.args[1]) == 0 + load_options = {kw.arg: ast.literal_eval(kw.value) for kw in load.value.keywords} + broadcast_options = {kw.arg: ast.literal_eval(kw.value) for kw in broadcast.value.keywords} + assert "dist_mode" not in source + if group is None: + assert load_options == {"size": 1} + assert broadcast_options == {"size": 64} + else: + assert load_options == {"size": group, "group": group, "stride": stride} + assert broadcast_options == {"size": 64, "group": group} +