From dbf29d58643326b10420c2f8396abaaf1fd192a3 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 1/3] 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 0320f4148a..19b518aed2 100644 --- a/src/ascend/codegen/codegen_ascend.cc +++ b/src/ascend/codegen/codegen_ascend.cc @@ -44,6 +44,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()) { @@ -2569,6 +2582,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 4e88a735fa..365e4a2706 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 " @@ -6754,6 +7150,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; @@ -6779,14 +7183,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); @@ -7386,6 +7782,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()) @@ -7582,6 +7992,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 234c42f62e..8105eab341 100644 --- a/src/ascend/codegen/codegen_pto.h +++ b/src/ascend/codegen/codegen_pto.h @@ -189,6 +189,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 91b34e240d..dbeee6de5a 100644 --- a/tilelang/ascend/language/__init__.py +++ b/tilelang/ascend/language/__init__.py @@ -65,6 +65,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.*) # Ascend owns its debug surface: device_assert lowers through the toolkit's # assert() macro, print gates by the NPU execution model (no CUDA-style @@ -120,6 +121,7 @@ "rng_rand", "rng_rand_float", "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 ed166469e227ed8ed00e595c9d8fef5db74a4a89 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 2/3] 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 1a235d9fec..88a1b5e7b8 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 8d9d2ec06c15ece531875199e429079fde3c6949 Mon Sep 17 00:00:00 2001 From: qukelin Date: Thu, 15 Oct 2026 10:00:24 +0800 Subject: [PATCH 3/3] 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)