diff --git a/src/ascend/codegen/codegen_ascend.cc b/src/ascend/codegen/codegen_ascend.cc index e7951c7ab7..44b772d606 100644 --- a/src/ascend/codegen/codegen_ascend.cc +++ b/src/ascend/codegen/codegen_ascend.cc @@ -45,6 +45,19 @@ bool IsSimdOp(const CallNode *op, std::string *op_name = nullptr) { return false; } +bool IsVmiOp(const CallNode *op, std::string *op_name = nullptr) { + if (auto opt_call_op = op->op.as()) { + const std::string &name = opt_call_op.value()->name; + if (name.rfind("tl.vmi.", 0) == 0) { + if (op_name != nullptr) { + *op_name = name; + } + return true; + } + } + return false; +} + bool IsSimdMergingCall(const CallNode *op, std::string *intrinsic_name = nullptr) { if (!op->dtype.is_void() || op->args.empty()) { @@ -2580,6 +2593,13 @@ void CodeGenTileLangAscend::VisitExpr_(const CallNode *op, std::ostream &os) { << "Unsupported RNG distribution on Ascend: " << dist; os << "tl::philox_rand_" << dist << "(&" << ascend_rng_state_var_ << ")"; } else { + std::string op_name; + // Evaluate statements print their value through this CallNode visitor, so + // statement-form VMI calls such as vstore are rejected here as well. + if (IsVmiOp(op, &op_name)) { + LOG(FATAL) << "Ascend CCE codegen does not support " << op_name + << "; use target='pto' for T.vmi.*"; + } CodeGenC::VisitExpr_(op, os); } } diff --git a/src/ascend/codegen/codegen_pto.cc b/src/ascend/codegen/codegen_pto.cc index 69ab33eaf8..9b985d6cf1 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); @@ -411,11 +470,44 @@ bool IsSupportedSIMTLocalStorageType(DataType t) { return IsSimtLocalBufferDtype(t) || (t.is_scalar() && t == DataType::Int(32)); } +bool IsSimtLocalIntegerStorageDtype(DataType t) { + return t.is_scalar() && (t.is_int() || t.is_uint()) && + (t.bits() == 8 || t.bits() == 32) && + IsSupportedSIMTLocalStorageType(t); +} + +std::string SimtLocalIntegerStorageTypeName(DataType t) { + ICHECK(IsSimtLocalIntegerStorageDtype(t)); + return "pto.i" + std::to_string(t.bits()); +} + +// SIMT local integer storage is signless because LLVM stack accesses require +// signless integer element types. TileLang's logical signedness is restored +// at the load boundary and folded back to the backing type at the store +// boundary. These adapters are only consumed by the SIMT local branches of +// ScalarLoad/EmitScalarStore/EmitScalarizedLoad/EmitScalarizedStore. +std::string SimtLocalLoadAdapter(const std::string &value, DataType t) { + if (!IsSimtLocalIntegerStorageDtype(t)) { + return value; + } + return "pto.cast(" + value + ", " + DataTypeName(t) + ")"; +} + +std::string SimtLocalStoreAdapter(const std::string &value, DataType t) { + if (!IsSimtLocalIntegerStorageDtype(t)) { + return value; + } + return "pto.cast(" + value + ", " + SimtLocalIntegerStorageTypeName(t) + ")"; +} + std::string SIMTLocalStorageTypeName(DataType t) { ICHECK(IsSupportedSIMTLocalStorageType(t)); // LLVM stack allocations require signless integer element types. Preserve // TileLang's signedness at load/store boundaries instead. - return IsInteger32(t) ? "pto.i32" : DataTypeName(t); + if (IsSimtLocalIntegerStorageDtype(t)) { + return SimtLocalIntegerStorageTypeName(t); + } + return DataTypeName(t); } bool IsSupportedSIMTFP8ContiguousLaneCount(int lanes) { @@ -554,8 +646,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) @@ -2010,6 +2103,7 @@ void CodeGenTileLangPTO::ResetFunctionState_() { mixed_entry_snapshot_ids_.clear(); loop_condition_snapshot_ids_.clear(); simd_pair_vars_.clear(); + value_defs_.clear(); gemm_emit_contexts_.clear(); gemm_emit_context_by_call_.clear(); gemm_zero_addr_emitted_ = false; @@ -2116,8 +2210,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(); @@ -2569,6 +2668,7 @@ std::string CodeGenTileLangPTO::PrintCondition(const PrimExpr &condition) { (dtype.is_bool() || dtype.is_int() || dtype.is_uint())) << "PTO control-flow condition must be a scalar integer predicate, got " << dtype; + CheckPtoScalarOperand_(condition, "PTO control-flow condition"); std::string value = RemoveOutermostParentheses(PrintExpr_(condition)); // TIR follows C truthiness for all integer conditions. Always compare at // the use site because a bool-typed expression may be materialized in an @@ -2596,10 +2696,7 @@ std::string CodeGenTileLangPTO::ScalarLoad(const BufferNode *buffer, return "tl.fp8_byte_load(" + GetVarID(buffer->data.get()) + ", " + index_str + ", " + FP8TypeName(buffer->dtype) + ")"; } - if (IsInteger32(buffer->dtype)) { - return "pto.cast(" + value + ", " + DataTypeName(buffer->dtype) + ")"; - } - return value; + return SimtLocalLoadAdapter(value, buffer->dtype); } if (scope == "shared" || scope == "shared.dyn" || scope == "global" || @@ -2623,20 +2720,24 @@ std::string CodeGenTileLangPTO::ScalarLoad(const BufferNode *buffer, std::string CodeGenTileLangPTO::LocalVarStoreValue(const PrimExpr &value, DataType dtype) { - std::string value_expr = RemoveOutermostParentheses(PrintExpr_(value)); if (dtype.is_handle()) { - return value_expr; + return RemoveOutermostParentheses(PrintExpr_(value)); } bool is_literal = value.as() != nullptr || value.as() != nullptr; if (dtype.is_scalar() && is_literal) { - return "pto.const(" + value_expr + ", dtype=" + DataTypeName(dtype) + ")"; + return "pto.const(" + RemoveOutermostParentheses(PrintExpr_(value)) + + ", dtype=" + DataTypeName(dtype) + ")"; } if (dtype.is_bool() || dtype.bits() == 1 || dtype.is_int() || dtype.is_uint()) { - return ScalarCastExpr(value_expr, dtype, "PTO local.var store"); + // Check provenance before printing/adapting the scalar (covers bool and + // integer local.var stores; other local.var branches keep their route). + CheckPtoScalarOperand_(value, "PTO local.var integer store"); + return ScalarCastExpr(RemoveOutermostParentheses(PrintExpr_(value)), dtype, + "PTO local.var store"); } - return value_expr; + return RemoveOutermostParentheses(PrintExpr_(value)); } void CodeGenTileLangPTO::EmitScalarStore(const BufferNode *buffer, @@ -2664,9 +2765,7 @@ void CodeGenTileLangPTO::EmitScalarStore(const BufferNode *buffer, << FP8TypeName(buffer->dtype) << ")\n"; return; } - if (IsInteger32(buffer->dtype)) { - store_value = "pto.cast(" + store_value + ", pto.i32)"; - } + store_value = SimtLocalStoreAdapter(store_value, buffer->dtype); stream << "pto.store(" << store_value << ", " << GetVarID(buffer->data.get()) << ", " << index_str << ")\n"; return; @@ -2707,6 +2806,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. @@ -2771,13 +2879,25 @@ void CodeGenTileLangPTO::EmitBufferAllocation(const Buffer &buffer) { "float32, int8, uint8, int32, uint32, and " "float8_e4m3fn/float8_e5m2, got " << buffer->dtype; + // Both allocation sub-branches use the same signless integer element + // backing for the admitted SIMT local integer types: scalar-element + // buffers through SIMTLocalStorageTypeName, vector-element register + // slots through the element type (without changing admission or + // claiming a supported slot read/write path). + auto storage_element_type = [&]() -> std::string { + if (!buffer->dtype.is_vector()) { + return SIMTLocalStorageTypeName(buffer->dtype); + } + if (tl::IsAscendVectorizableFP8(buffer->dtype)) { + return DataTypeName(buffer->dtype); + } + if (IsSimtLocalIntegerStorageDtype(elem)) { + return SimtLocalIntegerStorageTypeName(elem); + } + return ScalarType(elem); + }; stream << vid << " = pto.alloc_buffer((" << opt_size.value() << ",), " - << (buffer->dtype.is_vector() - ? (tl::IsAscendVectorizableFP8(buffer->dtype) - ? DataTypeName(buffer->dtype) - : ScalarType(elem)) - : SIMTLocalStorageTypeName(buffer->dtype)) - << ")\n"; + << storage_element_type() << ")\n"; } else if (persistent) { ICHECK(IsFloat32(buffer->dtype)) << "PTO persistent local allocation currently supports float32 " @@ -2799,13 +2919,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 +2931,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; } @@ -3358,6 +3477,7 @@ std::string CodeGenTileLangPTO::GetLocalByteAddrExpr( ", dtype=pto.int64)"; } + CheckPtoScalarOperand_(elem_index, context.c_str()); std::string index_expr = RemoveOutermostParentheses(PrintExpr_(elem_index)); std::string coerced = ScalarCastExpr(index_expr, DataType::Int(64), context); if (elem_bytes == 1) { @@ -3493,6 +3613,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 +4829,776 @@ 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(); +} + +static std::string VmiDtypeAnnotation(DataType element) { + if (element.is_int()) { + return PtoSignedIntegerTypeName(element); + } + if (element.is_uint()) { + return "pto.ui" + std::to_string(element.bits()); + } + return DataTypeName(element); +} + +const char *CodeGenTileLangPTO::PtoValueKindName(PtoValueKind kind) { + switch (kind) { + case PtoValueKind::kScalar: + return "scalar"; + case PtoValueKind::kVmiData: + return "VMI data"; + case PtoValueKind::kVmiMask: + return "VMI mask"; + case PtoValueKind::kVmiPair: + return "VMI pair"; + case PtoValueKind::kVoid: + return "void"; + case PtoValueKind::kHandle: + return "handle"; + case PtoValueKind::kOtherVector: + return "other vector"; + case PtoValueKind::kUnknown: + return "unknown"; + } + return "unknown"; +} + +struct CodeGenTileLangPTO::PtoValueQueryContext { + // Query-local lexical overrides (Let bindings seen by the active query). + std::unordered_map local_defs; + // Vars on the current definition chain; re-entering one is a cycle. + std::unordered_set active; + // Query-local memo for expression/definition classification. Node identity + // is stable, but a result computed under an active Let override is only + // valid for that lexical scope, so entries carry the scope generation. + std::unordered_map> + memo; + size_t scope_generation{0}; +}; + +// Classifies each immediate semantic operand of `expr` exactly once without +// descending into operand subtrees: ExprVisitor's per-node dispatch walks one +// level, and each child visit is intercepted below and handed to the full +// (memoized) classification. This keeps ordinary scalar expressions linear +// even for deeply nested or shared DAG-shaped terms. +struct CodeGenTileLangPTO::PtoOperandScanner final : public tirx::ExprVisitor { + PtoOperandScanner(const CodeGenTileLangPTO *codegen, + PtoValueQueryContext *context) + : codegen_(codegen), context_(context) {} + + bool has_unknown_origin{false}; + + void VisitExpr(const PrimExpr &expr) final { + if (!expr.defined() || has_unknown_origin) { + return; + } + PtoValueKind kind = codegen_->ClassifyPtoValue_(expr, context_); + has_unknown_origin = + kind == PtoValueKind::kVmiData || kind == PtoValueKind::kVmiMask || + kind == PtoValueKind::kVmiPair || kind == PtoValueKind::kVoid || + kind == PtoValueKind::kUnknown; + } + + // Walks the immediate semantic operands of the root expression. The + // qualified base call bypasses this class's VisitExpr override for the root + // only; the per-node handlers reach the children through the virtual + // VisitExpr above. + void ScanImmediateOperands(const PrimExpr &expr) { + tirx::ExprVisitor::VisitExpr(expr); + } + +private: + const CodeGenTileLangPTO *codegen_; + PtoValueQueryContext *context_; +}; + +namespace { + +bool IsVmiOpName(const std::string &name) { + return name.rfind("tl.vmi.", 0) == 0; +} + +} // namespace + +CodeGenTileLangPTO::PtoValueKind +CodeGenTileLangPTO::ClassifyPtoValue_(const PrimExpr &expr) const { + PtoValueQueryContext context; + return ClassifyPtoValue_(expr, &context); +} + +CodeGenTileLangPTO::PtoValueKind +CodeGenTileLangPTO::ClassifyPtoValue_(const PrimExpr &expr, + PtoValueQueryContext *context) const { + if (!expr.defined() || context == nullptr) { + return PtoValueKind::kUnknown; + } + const Object *node = expr.get(); + auto memo_it = context->memo.find(node); + if (memo_it != context->memo.end() && + memo_it->second.first == context->scope_generation) { + return memo_it->second.second; + } + PtoValueKind kind = ClassifyPtoValueUncached_(expr, context); + context->memo[node] = std::make_pair(context->scope_generation, kind); + return kind; +} + +CodeGenTileLangPTO::PtoValueKind CodeGenTileLangPTO::ClassifyPtoValueUncached_( + const PrimExpr &expr, PtoValueQueryContext *context) const { + if (const auto *var = expr.as()) { + const PrimExpr *definition = nullptr; + PrimExpr stored; + auto local_it = context->local_defs.find(var); + if (local_it != context->local_defs.end()) { + stored = local_it->second; + definition = &stored; + } else { + auto it = value_defs_.find(var); + if (it != value_defs_.end()) { + stored = it->second; + definition = &stored; + } + } + if (definition != nullptr) { + ICHECK(context->active.insert(var).second) + << "PTO SSA definition cycle detected for variable " + << var->name_hint; + PtoValueKind kind = ClassifyPtoValue_(*definition, context); + context->active.erase(var); + return kind; + } + // Unbound Var: parameter/loop variable contract. + if (expr.dtype().is_handle()) { + return PtoValueKind::kHandle; + } + if (expr.dtype().is_scalar()) { + return PtoValueKind::kScalar; + } + return PtoValueKind::kOtherVector; + } + + if (expr.as() != nullptr || expr.as() != nullptr) { + return PtoValueKind::kScalar; + } + + if (const auto *let = expr.as()) { + // Query-local lexical scope: bind the RHS in the local override layer and + // classify the body with the binding visible. The persistent table is not + // modified. The RHS is classified lazily when the body resolves the bound + // Var: classifying it eagerly here would traverse the same RHS a second + // time under the bumped scope generation and make a linear nested-Let + // chain exponential. Unused bindings do not affect the body's result + // classification; emission still handles and validates the RHS. + auto previous = context->local_defs.find(let->var.get()); + bool had_previous = previous != context->local_defs.end(); + PrimExpr saved; + if (had_previous) { + saved = previous->second; + } + context->local_defs[let->var.get()] = let->value; + ++context->scope_generation; + PtoValueKind kind = ClassifyPtoValue_(let->body, context); + ++context->scope_generation; + if (had_previous) { + context->local_defs[let->var.get()] = saved; + } else { + context->local_defs.erase(let->var.get()); + } + return kind; + } + + if (const auto *cast = expr.as()) { + PtoValueKind source = ClassifyPtoValue_(cast->value, context); + if (cast->value.dtype() == cast->dtype) { + return source; + } + // Non-identity casts of VMI values require VMI conversion operations. + // A scalar result dtype alone does not make a VMI or unknown source + // a scalar value. + if (source == PtoValueKind::kVmiData || source == PtoValueKind::kVmiMask || + source == PtoValueKind::kVmiPair || source == PtoValueKind::kVoid || + source == PtoValueKind::kUnknown) { + return PtoValueKind::kUnknown; + } + if (expr.dtype().is_handle()) { + return PtoValueKind::kHandle; + } + if (expr.dtype().is_scalar()) { + return PtoValueKind::kScalar; + } + return PtoValueKind::kOtherVector; + } + + if (const auto *load = expr.as()) { + return ClassifyBufferLoad_(load); + } + + if (const auto *call = expr.as()) { + auto opt_call_op = call->op.as(); + if (opt_call_op.has_value()) { + const std::string name = opt_call_op.value()->name; + if (IsVmiOpName(name)) { + return ClassifyVmiCallResult_(call, context); + } + if (call->op.same_as(tirx::builtin::reinterpret())) { + ICHECK_EQ(call->args.size(), 1U) + << "tirx.reinterpret expects 1 argument (value)"; + PtoValueKind source = ClassifyPtoValue_(call->args[0], context); + if (source == PtoValueKind::kVmiData) { + DataType source_dtype = call->args[0].dtype(); + DataType target_dtype = expr.dtype(); + if (source_dtype.bits() == target_dtype.bits() && + source_dtype.lanes() == target_dtype.lanes()) { + return PtoValueKind::kVmiData; + } + return PtoValueKind::kUnknown; + } + if (source == PtoValueKind::kVmiMask || + source == PtoValueKind::kVmiPair || source == PtoValueKind::kVoid || + source == PtoValueKind::kUnknown) { + return PtoValueKind::kUnknown; + } + if (expr.dtype().is_handle()) { + return PtoValueKind::kHandle; + } + if (expr.dtype().is_scalar()) { + return PtoValueKind::kScalar; + } + return PtoValueKind::kOtherVector; + } + } + return ClassifyOrdinaryExpr_(expr, context); + } + + return ClassifyOrdinaryExpr_(expr, context); +} + +CodeGenTileLangPTO::PtoValueKind +CodeGenTileLangPTO::ClassifyOrdinaryExpr_(const PrimExpr &expr, + PtoValueQueryContext *context) const { + // Ordinary scalar expressions keep their contract unless a semantic operand + // carries a VMI kind or an unknown VMI origin. Only the immediate operands + // are classified (once each, memoized); the query never prints or emits. + PtoOperandScanner scanner(this, context); + scanner.ScanImmediateOperands(expr); + if (scanner.has_unknown_origin) { + return PtoValueKind::kUnknown; + } + if (expr.dtype().is_handle()) { + return PtoValueKind::kHandle; + } + if (expr.dtype().is_scalar()) { + return PtoValueKind::kScalar; + } + return PtoValueKind::kOtherVector; +} + +CodeGenTileLangPTO::PtoValueKind +CodeGenTileLangPTO::ClassifyBufferLoad_(const BufferLoadNode *op) const { + // Ordinary memory reads follow the storage contract: a scalar-valued + // BufferLoad is a scalar regardless of which VMI store previously wrote the + // address. VMI local register slots are the only buffer reads classified as + // VMI data; multi-lane local.var and ordinary vectors stay on their existing + // register/vector routes. + if (IsVmiLocalRegisterBuffer(op->buffer.get())) { + return PtoValueKind::kVmiData; + } + if (op->dtype.is_scalar()) { + return PtoValueKind::kScalar; + } + return PtoValueKind::kOtherVector; +} + +CodeGenTileLangPTO::PtoValueKind CodeGenTileLangPTO::ClassifyVmiCallResult_( + const CallNode *op, PtoValueQueryContext *context) const { + auto opt_call_op = op->op.as(); + ICHECK(opt_call_op.has_value()) << "tl.vmi call has a non-Op callee"; + const std::string name = opt_call_op.value()->name; + if (name == "tl.vmi.create_mask" || name == "tl.vmi.vcmp" || + name == "tl.vmi.vcmps") { + return PtoValueKind::kVmiMask; + } + if (name == "tl.vmi.vintlv" || name == "tl.vmi.vdintlv" || + name == "tl.vmi.vunzip" || name == "tl.vmi.vmull") { + return PtoValueKind::kVmiPair; + } + if (name == "tl.vmi.vload") { + return PtoValueKind::kVmiData; + } + if (name == "tl.vmi.pair_get") { + ICHECK_EQ(op->args.size(), 2U) << "tl.vmi.pair_get expects (pair, index)"; + PtoValueKind pair_kind = ClassifyPtoValue_(op->args[0], context); + ICHECK(pair_kind == PtoValueKind::kVmiPair) + << "PTO provenance tl.vmi.pair_get source must be a VMI pair; got " + << PtoValueKindName(pair_kind) << " (" << op->args[0].dtype() + << "): " << op->args[0]; + const auto *index = op->args[1].as(); + ICHECK(index != nullptr && (index->value == 0 || index->value == 1)) + << "PTO provenance tl.vmi.pair_get index must be the constant 0 or 1, " + "got " + << op->args[1]; + return PtoValueKind::kVmiData; + } + if (name == "tl.vmi.vstore" || name == "tl.vmi.vscatter") { + return PtoValueKind::kVoid; + } + static const std::unordered_set kDataOps = { + "tl.vmi.vci", "tl.vmi.vbrc", + "tl.vmi.vzip", "tl.vmi.vadd", + "tl.vmi.vsub", "tl.vmi.vmul", + "tl.vmi.vdiv", "tl.vmi.vmax", + "tl.vmi.vmin", "tl.vmi.vand", + "tl.vmi.vor", "tl.vmi.vxor", + "tl.vmi.vshl", "tl.vmi.vshr", + "tl.vmi.vabs", "tl.vmi.vneg", + "tl.vmi.vrelu", "tl.vmi.vexp", + "tl.vmi.vln", "tl.vmi.vsqrt", + "tl.vmi.vnot", "tl.vmi.vadds", + "tl.vmi.vmuls", "tl.vmi.vmaxs", + "tl.vmi.vmins", "tl.vmi.vshls", + "tl.vmi.vshrs", "tl.vmi.vsel", + "tl.vmi.vselr", "tl.vmi.vcadd", + "tl.vmi.vcmax", "tl.vmi.vcmin", + "tl.vmi.vcvt", "tl.vmi.vinterpret_cast", + "tl.vmi.vexpdif", "tl.vmi.vaxpy", + "tl.vmi.vlrelu", "tl.vmi.vprelu", + "tl.vmi.vmula", "tl.vmi.vdhist", + "tl.vmi.vchist", "tl.vmi.vgather", + "tl.vmi.vgatherb", + }; + if (kDataOps.count(name) != 0U) { + return PtoValueKind::kVmiData; + } + // Registered operations without an explicit contract (for example a future + // addition) stay unknown so consumers diagnose instead of guessing. + return PtoValueKind::kUnknown; +} + +void CodeGenTileLangPTO::CheckPtoScalarOperand_(const PrimExpr &expr, + const char *consumer) const { + PtoValueKind kind = ClassifyPtoValue_(expr); + if (kind == PtoValueKind::kScalar) { + return; + } + std::ostringstream message; + message << consumer << " expects a scalar; got " << PtoValueKindName(kind) + << " (" << expr.dtype() << ", lanes=" << expr.dtype().lanes() + << "): " << expr; + if (kind == PtoValueKind::kUnknown) { + const auto *call = expr.as(); + if (call != nullptr) { + auto opt_call_op = call->op.as(); + if (opt_call_op.has_value() && IsVmiOpName(opt_call_op.value()->name)) { + message << "; unclassified " << opt_call_op.value()->name + << " result; add an explicit VMI return contract"; + } + } + } + LOG(FATAL) << message.str(); +} + +bool CodeGenTileLangPTO::ResolvesToMutableRegisterRead( + const PrimExpr &expr, + std::unordered_map *local_defs, + std::unordered_set *active) const { + if (!expr.defined() || local_defs == nullptr || active == nullptr) { + return false; + } + if (const auto *cast = expr.as()) { + // Same-dtype casts are transparent aliases; numeric conversions are not. + return cast->value.dtype() == cast->dtype && + ResolvesToMutableRegisterRead(cast->value, local_defs, active); + } + if (const auto *let = expr.as()) { + auto previous = local_defs->find(let->var.get()); + const bool had_previous = previous != local_defs->end(); + PrimExpr saved; + if (had_previous) { + saved = previous->second; + } + (*local_defs)[let->var.get()] = let->value; + const bool resolves = + ResolvesToMutableRegisterRead(let->body, local_defs, active); + if (had_previous) { + (*local_defs)[let->var.get()] = saved; + } else { + local_defs->erase(let->var.get()); + } + return resolves; + } + if (const auto *var = expr.as()) { + PrimExpr stored; + const PrimExpr *definition = nullptr; + auto local_it = local_defs->find(var); + if (local_it != local_defs->end()) { + stored = local_it->second; + definition = &stored; + } else { + auto it = value_defs_.find(var); + if (it != value_defs_.end()) { + stored = it->second; + definition = &stored; + } + } + if (definition == nullptr || !active->insert(var).second) { + return false; + } + const bool resolves = + ResolvesToMutableRegisterRead(*definition, local_defs, active); + active->erase(var); + return resolves; + } + if (const auto *load = expr.as()) { + return ScopeOfBuffer(load->buffer.get()) == "local.var"; + } + return false; +} + +bool CodeGenTileLangPTO::IsMutableRegisterGroupSource( + const PrimExpr &source) const { + // Accept numeric local.var reads through Var/Let aliases and same-dtype + // casts. Keep their classification unchanged: other-vector for multiple + // lanes, scalar for one lane. This helper only checks whether grouped vbrc + // may consume the expression; it does not inspect mutable stores. + // is_uint() includes TIR uint1/bool. Handle reads are not numeric. + // Other expressions, including VMI mask/pair/void/unknown results, do not + // resolve to this read pattern. + const DataType dtype = source.dtype(); + if (!dtype.is_int() && !dtype.is_uint() && !dtype.is_float() && + !dtype.is_bfloat16()) { + return false; + } + std::unordered_map local_defs; + std::unordered_set active; + return ResolvesToMutableRegisterRead(source, &local_defs, &active); +} + +void CodeGenTileLangPTO::CheckPtoVmiBridgeOperand_( + const CallNode *op, const char *op_name, PtoValueKind source_kind) const { + ICHECK(!op->args.empty()) << op_name << " expects a source operand"; + const PrimExpr &source = op->args[0]; + const bool has_group = op->annotations.find("group") != op->annotations.end(); + if (std::string(op_name) == "tl.vmi.vci") { + if (source_kind == PtoValueKind::kVmiData) { + LOG(FATAL) << "tl.vmi.vci base expects a scalar; got VMI data (" + << source.dtype() << ", lanes=" << source.dtype().lanes() + << "): " << source; + } + if (source_kind != PtoValueKind::kScalar) { + LOG(FATAL) << "tl.vmi.vci base expects a scalar; got " + << PtoValueKindName(source_kind) << " (" << source.dtype() + << ", lanes=" << source.dtype().lanes() << "): " << source; + } + return; + } + // tl.vmi.vbrc + if (has_group) { + const auto *group = op->annotations.at("group").as(); + ICHECK(group != nullptr) + << "tl.vmi.vbrc group annotation must be a constant integer"; + if (source_kind != PtoValueKind::kVmiData) { + // Also accept numeric local.var reads through pure aliases without + // changing their classification. Scalar constants and parameters do not + // match this read pattern. Require the same lane count as the group. + if (IsMutableRegisterGroupSource(source)) { + ICHECK_EQ(source.dtype().lanes(), group->value) + << "tl.vmi.vbrc source lane count must equal group; got " + << source.dtype().lanes() << " lanes and group=" << group->value; + return; + } + LOG(FATAL) << "tl.vmi.vbrc with group=... requires a VMI vreg; got " + << PtoValueKindName(source_kind) << " (" << source.dtype() + << ", lanes=" << source.dtype().lanes() << "): " << source; + } + ICHECK_EQ(source.dtype().lanes(), group->value) + << "tl.vmi.vbrc source lane count must equal group; got " + << source.dtype().lanes() << " lanes and group=" << group->value; + return; + } + if (source_kind == PtoValueKind::kVmiMask || + source_kind == PtoValueKind::kVmiPair || + source_kind == PtoValueKind::kVoid || + source_kind == PtoValueKind::kHandle || + source_kind == PtoValueKind::kUnknown) { + LOG(FATAL) << "tl.vmi.vbrc cannot broadcast " + << PtoValueKindName(source_kind) << " (" << source.dtype() + << ", lanes=" << source.dtype().lanes() << "): " << source; + } +} + +void CodeGenTileLangPTO::PrintPtoVmiCall_(const CallNode *op, + std::ostream &os) { + auto opt_call_op = op->op.as(); + 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; + } + + // Classify the VMI source *before* printing any argument. A lane-1 VMI + // vreg is a TIR scalar by dtype but remains VMI data; only a proven scalar + // may enter the integer signless bridge. vci/scalar-only consumers receive + // an explicit category diagnostic instead of an invalid PTODSL scalar op. + PtoValueKind vmi_source_kind = PtoValueKind::kUnknown; + if (!op->args.empty()) { + vmi_source_kind = ClassifyPtoValue_(op->args[0]); + } + if (op_name == "tl.vmi.vbrc" || op_name == "tl.vmi.vci") { + CheckPtoVmiBridgeOperand_(op, op_name.c_str(), vmi_source_kind); + } + // Adapt proven integer scalar sources of vci/vbrc at the VMI + // value-argument boundary. The generated scalar must be *signless*: PTOAS + // materialises per-physical-vector offsets with `arith.constant`, which + // requires signless integers, so a typed (signed) constant base fails for + // 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 = + vmi_source_kind == PtoValueKind::kScalar && !op->args.empty() && + op->args[0].dtype().is_scalar() && + (op->args[0].dtype().is_int() || op->args[0].dtype().is_uint()); + const bool integer_data_result = + op->dtype.element_of().is_int() || op->dtype.element_of().is_uint(); + const bool integer_source_literal = + !op->args.empty() && op->args[0].as() != nullptr; + const bool bridge_integer_source = + (op_name == "tl.vmi.vci" || + (op_name == "tl.vmi.vbrc" && !integer_source_literal)) && + scalar_integer_input && integer_data_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 { + 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)) { @@ -4700,6 +5612,15 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, RequiresStatementControlFlow(op->args[1]) || RequiresStatementControlFlow(op->args[2]); auto branch_value = [&](const PrimExpr &expr) { + const bool scalar_adapted_branch = + expr.as() == nullptr && + expr.as() == nullptr && + (expr.dtype().is_bool() || + (op->dtype.is_scalar() && + (op->dtype.is_int() || op->dtype.is_uint()))); + if (scalar_adapted_branch) { + CheckPtoScalarOperand_(expr, "PTO if_then_else branch"); + } std::string value = PrintExpr_(expr); if (expr.dtype().is_bool()) { // PTODSL typed constants require numeric bool literals. @@ -5079,8 +6000,35 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, if (op->op.same_as(tirx::builtin::reinterpret())) { ICHECK_EQ(op->args.size(), 1U) << "tirx.reinterpret expects 1 argument (value)"; - std::string value = PrintExpr_(op->args[0]); + PtoValueKind reinterpret_source_kind = ClassifyPtoValue_(op->args[0]); DataType target_dtype = op->dtype; + if (reinterpret_source_kind == PtoValueKind::kVmiData) { + DataType source_dtype = op->args[0].dtype(); + if (source_dtype.bits() == target_dtype.bits() && + source_dtype.lanes() == target_dtype.lanes()) { + // Same element width and lane count is a register-level re-view; use + // the explicit VMI reinterpret API, never scalar/lexical bitcasts. + os << "pto.vmi.vinterpret_cast(" << PrintExpr_(op->args[0]) + << ", to_dtype=" << VmiDtypeAnnotation(target_dtype.element_of()) + << ")"; + return; + } + LOG(FATAL) << "PTO reinterpret of VMI data requires the same element " + "width and lane count; use T.vmi.vinterpret_cast with the " + "explicit packed/shape contract (source " + << source_dtype << ", target " << target_dtype + << "): " << op->args[0]; + } + if (reinterpret_source_kind == PtoValueKind::kVmiMask || + reinterpret_source_kind == PtoValueKind::kVmiPair || + reinterpret_source_kind == PtoValueKind::kVoid || + reinterpret_source_kind == PtoValueKind::kUnknown) { + LOG(FATAL) << "PTO reinterpret cannot consume " + << PtoValueKindName(reinterpret_source_kind) << " (source " + << op->args[0].dtype() << ", target " << target_dtype + << "): " << op->args[0]; + } + std::string value = PrintExpr_(op->args[0]); if (target_dtype.lanes() > 1) { // Vector reinterpret: lower to pto.vbitcast(vec, elem_type). os << "pto.vbitcast(" << value << ", " @@ -5480,6 +6428,7 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, << "tl.simd vlrelu/vprelu expects 3 arguments"; if (op->op.same_as(tl::simd_vlrelu())) { DataType vector_dtype = op->args[0].dtype(); + CheckPtoScalarOperand_(op->args[1], "tl.simd.vlrelu scalar"); std::string alpha = WrapTypedConst( RemoveOutermostParentheses(PrintExpr_(PeelScalarCasts(op->args[1]))), vector_dtype, IsImmediateScalar(op->args[1])); @@ -5529,6 +6478,7 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, : "vshrs"; std::string src = PrintExpr_(op->args[0]); DataType result_dtype = op->args[0].dtype(); + CheckPtoScalarOperand_(op->args[1], "tl.simd vector-scalar operand"); std::string scalar = WrapTypedConst( RemoveOutermostParentheses(PrintExpr_(PeelScalarCasts(op->args[1]))), result_dtype, IsImmediateScalar(op->args[1])); @@ -5603,6 +6553,7 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, if (op->op.same_as(tl::simd_vcmps())) { ICHECK_EQ(op->args.size(), 4U) << "tl.simd.vcmps expects 4 arguments (src, scalar, mask, op)"; + CheckPtoScalarOperand_(op->args[1], "tl.simd.vcmps scalar"); std::string cmp_op = Downcast(op->args[3])->value; static const std::unordered_map kCmpModeEnum = { {"eq", "EQ"}, {"ne", "NE"}, {"gt", "GT"}, @@ -5634,6 +6585,7 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, // Match the authored vector element dtype (i16 for packed-ue8m0 lane math, // i32 for f32 gather paths). Always-i32 coercion breaks i16 vci users. DataType elem = op->dtype.element_of(); + CheckPtoScalarOperand_(op->args[0], "tl.simd.vci index"); std::string index; if (IsImmediateScalar(op->args[0])) { index = @@ -5654,6 +6606,7 @@ void CodeGenTileLangPTO::VisitExpr_(const CallNode *op, << "tl.simd.vdup expects 3 arguments (src, mask, mode)"; ValidateVecMode_(op, 2); DataType elem = op->dtype.element_of(); + CheckPtoScalarOperand_(op->args[0], "tl.simd.vdup scalar"); std::string scalar; if (IsImmediateScalar(op->args[0])) { scalar = @@ -6050,6 +7003,27 @@ 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"; + PtoValueKind pair_kind = ClassifyPtoValue_(op->args[0]); + ICHECK(pair_kind == PtoValueKind::kVmiPair) + << "tl.vmi.pair_get source must be a VMI pair; got " + << PtoValueKindName(pair_kind) << " (" << op->args[0].dtype() + << "): " << op->args[0]; + const auto *pair_index = op->args[1].as(); + ICHECK(pair_index != nullptr && + (pair_index->value == 0 || pair_index->value == 1)) + << "tl.vmi.pair_get index must be the constant 0 or 1, got " + << op->args[1]; + os << "(" << PrintExpr_(op->args[0]) << ")[" << pair_index->value + << "]"; + return; + } + 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 " @@ -6089,6 +7063,21 @@ void CodeGenTileLangPTO::VisitExpr_(const CastNode *op, return; } + PtoValueKind cast_source_kind = ClassifyPtoValue_(op->value); + if (cast_source_kind == PtoValueKind::kVmiData) { + LOG(FATAL) << "PTO integer cast cannot consume VMI data; use T.vmi.vcvt " + "or T.vmi.vinterpret_cast (source " + << op->value.dtype() << ", target " << to << "): " << op->value; + } + if (cast_source_kind == PtoValueKind::kVmiMask || + cast_source_kind == PtoValueKind::kVmiPair || + cast_source_kind == PtoValueKind::kVoid || + cast_source_kind == PtoValueKind::kUnknown) { + LOG(FATAL) << "PTO cast cannot consume " + << PtoValueKindName(cast_source_kind) << " (source " + << op->value.dtype() << ", target " << to << "): " << op->value; + } + const bool from_fp8 = tl::IsAscendVectorizableFP8(from); const bool to_fp8 = tl::IsAscendVectorizableFP8(to); if (from_fp8 || to_fp8) { @@ -6225,8 +7214,14 @@ void CodeGenTileLangPTO::VisitExpr_(const LetNode *op, << "PTO Let variable is already defined: " << op->var->name_hint; std::string vid = AllocVarID(op->var.get()); + ICHECK(!value_defs_.count(op->var.get())) + << "duplicate PTO Let definition: " << op->var->name_hint; + value_defs_.emplace(op->var.get(), op->value); std::string body = PrintExpr_(op->body); + bool removed_def = value_defs_.erase(op->var.get()); + ICHECK(removed_def) << "PTO Let definition was not registered: " + << op->var->name_hint; bool removed = var_idmap_.erase(op->var.get()); ICHECK(removed) << "PTO Let variable was not registered: " << op->var->name_hint; @@ -6423,6 +7418,9 @@ void CodeGenTileLangPTO::VisitExpr_(const NotNode *op, ICHECK(op->dtype.is_scalar() && op->a.dtype().is_bool()) << "PTO logical-not currently supports scalar bool predicates only, got " << op->a.dtype(); + // A lane-1 VMI mask is a TIR scalar bool but remains a VMI value and cannot + // enter the logical scalar adapter. + CheckPtoScalarOperand_(op->a, "PTO logical-not operand"); // tl.logical_not(value) implements logical negation as value == 0. os << "tl.logical_not("; PrintExpr_(op->a, os); @@ -6441,10 +7439,14 @@ void CodeGenTileLangPTO::VisitExpr_(const SelectNode *op, return; } if (op->dtype.is_scalar() && op->dtype.is_bool()) { + // The logical-bool adapter is scalar-only; a lane-1 VMI mask (or VMI + // data used as a predicate) must be diagnosed, not normalized. + CheckPtoScalarOperand_(expr, "PTO select branch"); os << "tl.as_logical_bool(" << PrintExpr_(expr) << ")"; return; } if (coerce_integer_branches) { + CheckPtoScalarOperand_(expr, "PTO select branch"); os << ScalarCastExpr(PrintExpr_(expr), op->dtype, "PTO select branch"); return; } @@ -6570,6 +7572,13 @@ void CodeGenTileLangPTO::VisitExpr_(const BroadcastNode *op, void CodeGenTileLangPTO::VisitStmt_(const DeclBufferNode *op) { (void)op; } void CodeGenTileLangPTO::VisitStmt_(const BindNode *op) { + // Record every SSA definition before any printing branch (including the + // early-return SIMD pair/pointer paths below). Queries only consume the + // stored IR; they never depend on the printed form. + ICHECK(!value_defs_.count(op->var.get())) + << "duplicate PTO SSA definition: " << op->var->name_hint; + value_defs_.emplace(op->var.get(), op->value); + if (const auto *call = op->value.as()) { if (IsOpName(call->op, "tl.simd.alloc")) { AllocVarID(op->var.get()); @@ -6769,6 +7778,14 @@ void CodeGenTileLangPTO::VisitStmt_(const BindNode *op) { } void CodeGenTileLangPTO::VisitStmt_(const AllocBufferNode *op) { + // All local.var buffers (scalar, handle and vector) share the unified + // allocation path with block alloc_buffers: scope registration, + // duplicate-allocation dedup, name reuse and mixed-section snapshots all + // live in EmitBufferAllocation / EmitLocalVarInitialization_. + if (GetPtrStorageScope(op->buffer->data) == "local.var") { + EmitBufferAllocation(op->buffer); + return; + } if (!current_function_is_cube_ || in_mixed_vector_section_) { EmitBufferAllocation(op->buffer); return; @@ -6794,14 +7811,6 @@ void CodeGenTileLangPTO::VisitStmt_(const AllocBufferNode *op) { auto opt_size = alloc_ref.ConstantAllocationSize(); ICHECK(opt_size.has_value()) << "PTO local.fragment allocation expects a constant size"; - } else if (scope == "local.var") { - // Scalar alloc_var is used by GEMM tile mapping; vector alloc_var is a - // mutable SIMD register and is emitted as a PTODSL Vec value. - CheckLocalVarBuffer(op->buffer.get()); - PrintIndent(); - local_var_buffers_.insert(buffer_var.get()); - EmitLocalVarInitialization_(op->buffer, AllocVarID(buffer_var.get())); - return; } RegisterHandleType_(buffer_var.get(), op->buffer->dtype); @@ -7401,6 +8410,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()) @@ -7474,8 +8497,10 @@ void CodeGenTileLangPTO::EmitScalarizedLoad(const BufferLoadNode *op, index_str = RemoveOutermostParentheses(PrintExpr_(ramp->base)); } if (scope == "local.fragment" || scope == "local") { - os << "pto.load(" << GetVarID(op->buffer->data.get()) << ", " << index_str - << ", contiguous=" << lanes << ")"; + std::string value = "pto.load(" + GetVarID(op->buffer->data.get()) + + ", " + index_str + + ", contiguous=" + std::to_string(lanes) + ")"; + os << SimtLocalLoadAdapter(value, element_dtype); return; } if (scope == "shared" || scope == "shared.dyn") { @@ -7597,6 +8622,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()) @@ -7671,8 +8710,8 @@ void CodeGenTileLangPTO::EmitScalarizedStore(const BufferStoreNode *op) { const std::string contiguous_suffix = store_lanes > 1 ? ", contiguous=" + std::to_string(store_lanes) : ""; if (scope == "local.fragment" || scope == "local") { - stream << "pto.store(" << value << ", " - << GetVarID(op->buffer->data.get()) << ", " << index_str + stream << "pto.store(" << SimtLocalStoreAdapter(value, buffer_dtype) + << ", " << GetVarID(op->buffer->data.get()) << ", " << index_str << contiguous_suffix << ")\n"; return; } diff --git a/src/ascend/codegen/codegen_pto.h b/src/ascend/codegen/codegen_pto.h index eb6c857bc8..47a2dc3934 100644 --- a/src/ascend/codegen/codegen_pto.h +++ b/src/ascend/codegen/codegen_pto.h @@ -191,6 +191,49 @@ 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); + // SSA provenance for PTO values. The table maps a bound VarNode to its RHS + // IR; only Bind/Let registration writes it. Queries are pure: they never + // print, emit, allocate names or modify the persistent table. + enum class PtoValueKind { + kScalar, + kVmiData, + kVmiMask, + kVmiPair, + kVoid, + kHandle, + kOtherVector, + kUnknown, + }; + struct PtoValueQueryContext; // query-local Let overrides + active Var set + struct PtoOperandScanner; // one-level immediate-operand classification + PtoValueKind ClassifyPtoValue_(const PrimExpr &expr) const; + PtoValueKind ClassifyPtoValue_(const PrimExpr &expr, + PtoValueQueryContext *context) const; + PtoValueKind ClassifyPtoValueUncached_(const PrimExpr &expr, + PtoValueQueryContext *context) const; + PtoValueKind ClassifyVmiCallResult_(const CallNode *op, + PtoValueQueryContext *context) const; + PtoValueKind ClassifyBufferLoad_(const BufferLoadNode *op) const; + PtoValueKind ClassifyOrdinaryExpr_(const PrimExpr &expr, + PtoValueQueryContext *context) const; + static const char *PtoValueKindName(PtoValueKind kind); + void CheckPtoScalarOperand_(const PrimExpr &expr, const char *consumer) const; + void CheckPtoVmiBridgeOperand_(const CallNode *op, const char *op_name, + PtoValueKind source_kind) const; + // Grouped vbrc also accepts numeric local.var reads through pure aliases. + // The source classification is unchanged; the consumer checks lanes == group. + // This does not infer the read's origin from mutable stores. + bool IsMutableRegisterGroupSource(const PrimExpr &source) const; + bool ResolvesToMutableRegisterRead( + const PrimExpr &expr, + std::unordered_map *local_defs, + std::unordered_set *active) const; + bool IsVmiLocalRegisterBuffer(const BufferNode *buffer) const; + void CheckVmiLocalRegisterIndex(const BufferNode *buffer, + const PrimExpr &index) const; bool HasAscendGemmL1(const PrimFunc &func) const; bool HasAscendBlockscaledGemmL1(const PrimFunc &func) const; bool HasAscendMad(const PrimFunc &func) const; @@ -240,6 +283,9 @@ class CodeGenTileLangPTO final : public CodeGenTileLangPY { // VarNode to its (low, high) SSA name pair. std::unordered_map> simd_pair_vars_; + // Persistent per-function SSA definitions used by PtoValueKind queries. + // Registered at every Bind/Let entry and cleared in ResetFunctionState_. + std::unordered_map value_defs_; std::vector gemm_emit_contexts_; std::unordered_map gemm_emit_context_by_call_; diff --git a/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..a4a65dbddb --- /dev/null +++ b/src/ascend/transform/parallel_to_pto_utils.cc @@ -0,0 +1,1626 @@ +/*! + * \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 (chunk, lane) through the inverse layout to obtain the logical + // iteration coordinates followed by the replicate coordinate. With + // replicate_size == 1, the trailing replicate coordinate is zero. + // The branches below use the logical coordinates for the loop dimensionality. + Var chunk_var("pto_chunk", loop->loop_var.dtype()); + Var lane_var("pto_lane", loop->loop_var.dtype()); + Array logical = inverse->Forward({chunk_var, lane_var}); + // 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; +} + +bool IsVmiLogicalLaneCount(int64_t lanes) { + return lanes == 64 || lanes == 128 || lanes == 256; +} + +std::string CastContextName(const CastContext &context) { + if (context.value_kind == CastValueKind::kScalar) { + return "scalar"; + } + std::ostringstream oss; + oss << "VMI, lanes=" << context.logical_lanes; + return oss.str(); +} + +bool IsHalfOrBfloat16(DataType dtype) { + return (dtype.is_float() && dtype.bits() == 16) || dtype.is_bfloat16(); +} + +/*! Identity attribute policy shared by the scalar and VMI branches: no + * conversion exists, so sat=True (the default) is normalized away, while + * sat=False or a non-empty rounding mode would claim a conversion semantic + * that does not exist. An explicit empty round key means "no request". */ +bool ValidateIdentityCastAttributes(const CastAttributes &attrs, + const std::string &pair, std::string *why) { + if (attrs.has_sat && !attrs.sat) { + *why = pair + ": an identity cast does not accept sat=False"; + return false; + } + if (attrs.has_round_mode) { + *why = pair + ": an identity cast does not accept a rounding mode"; + return false; + } + return true; +} + +/*! Integer attribute policy: rounding modes are never accepted; an explicit + * saturation flag is accepted only for narrowing and only as sat=False, which + * the direct NOSAT lowering already honours. */ +bool ValidateIntegerCastAttributes(const CastAttributes &attrs, bool narrowing, + const std::string &pair, std::string *why) { + if (attrs.has_round_mode) { + *why = pair + ": integer conversions do not accept a rounding mode"; + return false; + } + if (attrs.has_sat) { + if (!narrowing) { + *why = pair + ": this integer conversion does not accept an explicit " + "saturation flag"; + return false; + } + if (attrs.sat) { + *why = pair + ": integer narrowing accepts only sat=False (NOSAT), not " + "sat=True"; + return false; + } + } + return true; +} + +Optional PlanScalarValueCast(const CastNode *cast, + const CastAttributes &attrs, + const std::string &pair, + std::string *reason) { + const DataType from = cast->value.dtype(); + const DataType to = cast->dtype; + auto reject = [&](const std::string &why) { + if (reason != nullptr) { + *reason = why; + } + return Optional(); + }; + ScalarCastPlan keep; + keep.action = ScalarCastAction::kKeepInput; + ScalarCastPlan rebuild; + rebuild.action = ScalarCastAction::kRebuildCast; + if (from == to) { + std::string why; + if (!ValidateIdentityCastAttributes(attrs, pair, &why)) { + return reject(why); + } + return CastPlan{keep}; + } + const bool from_int = from.is_int() || from.is_uint(); + const bool to_int = to.is_int() || to.is_uint(); + if (from_int && to_int) { + std::string why; + if (!ValidateIntegerCastAttributes(attrs, to.bits() < from.bits(), pair, + &why)) { + return reject(why); + } + // Scalar integer casts keep the ordinary pto.cast/trunc lowering; the + // validated sat=False marker is already honoured and is normalized away. + return CastPlan{rebuild}; + } + if (IsSupportedFloat8(from) || IsSupportedFloat8(to)) { + return reject(pair + ": non-identity FP8 conversions have no scalar " + "lowering inside SIMD_VF; use a lane-varying value " + "so the cast takes the VMI vcvt path"); + } + const bool from_float = from.is_float() || from.is_bfloat16(); + const bool to_float = to.is_float() || to.is_bfloat16(); + if (from_float && to_float) { + if (attrs.has_round_mode) { + return reject(pair + ": this conversion does not accept a rounding mode"); + } + const bool from_f32 = from.is_float() && from.bits() == 32; + const bool to_f32 = to.is_float() && to.bits() == 32; + if (!from_f32 && to_f32) { + if (attrs.has_sat) { + return reject(pair + ": this direction does not accept a saturation " + "mode"); + } + return CastPlan{rebuild}; + } + if (from_f32 && !to_f32) { + if (attrs.has_sat && attrs.sat) { + return reject(pair + ": this conversion accepts only sat=False " + "(non-saturating)"); + } + return CastPlan{rebuild}; + } + return reject(pair + ": cross-family f16 <-> bf16 narrowing has no " + "validated rounding semantics; convert through " + "float32 explicitly"); + } + return reject(pair + ": no scalar conversion exists for this dtype pair"); +} + +Optional +PlanVectorValueCast(const CastNode *cast, const CastContext &context, + const CastAttributes &attrs, const std::string &round_token, + const std::string &pair, std::string *reason) { + const DataType from = cast->value.dtype(); + const DataType to = cast->dtype; + auto reject = [&](const std::string &why) { + if (reason != nullptr) { + *reason = why; + } + return Optional(); + }; + if (from == to) { + std::string why; + if (!ValidateIdentityCastAttributes(attrs, pair, &why)) { + return reject(why); + } + return CastPlan{VectorCastPlan{}}; + } + const bool from_int = from.is_int() || from.is_uint(); + const bool to_int = to.is_int() || to.is_uint(); + if (from_int && to_int) { + std::string why; + if (!ValidateIntegerCastAttributes(attrs, to.bits() < from.bits(), pair, + &why)) { + return reject(why); + } + VectorCastPlan plan; + if (to.bits() < from.bits()) { + plan.steps.push_back( + {VectorCastOp::kIntegerTruncate, from, to, "", "NOSAT"}); + return CastPlan{plan}; + } + if (to.bits() == from.bits()) { + plan.steps.push_back( + {VectorCastOp::kIntegerReinterpret, from, to, "", ""}); + return CastPlan{plan}; + } + if (from.is_int() == to.is_int()) { + plan.steps.push_back({VectorCastOp::kIntegerExtend, from, to, "", ""}); + return CastPlan{plan}; + } + // Cross-sign widening: extend with the source signedness to the target + // width, then re-view the same-width result as the target signedness. + const DataType intermediate = + from.is_int() ? DataType::Int(to.bits()) : DataType::UInt(to.bits()); + plan.steps.push_back( + {VectorCastOp::kIntegerExtend, from, intermediate, "", ""}); + plan.steps.push_back( + {VectorCastOp::kIntegerReinterpret, intermediate, to, "", ""}); + return CastPlan{plan}; + } + const bool from_fp8 = IsSupportedFloat8(from); + const bool to_fp8 = IsSupportedFloat8(to); + if (from_fp8 || to_fp8) { + if (from_fp8 && to_fp8) { + return reject(pair + ": conversions between the two FP8 types are not " + "supported (they are not a reinterpretation)"); + } + VectorCastPlan plan; + if (from_fp8 && to.is_float() && to.bits() == 32) { + // Exact widening: a real conversion with no options. An explicit hint + // here cannot be honoured, so it is rejected rather than dropped. + if (attrs.has_round_key) { + return reject(pair + ": this direction does not accept a rounding " + "mode"); + } + if (attrs.has_sat) { + return reject(pair + ": this direction does not accept a saturation " + "mode"); + } + plan.steps.push_back( + {VectorCastOp::kNonIntegerConvert, from, to, "", ""}); + return CastPlan{plan}; + } + if (to_fp8) { + 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(); + const std::string token = attrs.has_round_mode ? round_token : "R"; + if (allowed.count(token) == 0) { + std::ostringstream oss; + oss << pair << ": rounding mode " << token; + if (attrs.has_round_mode && attrs.round != token) { + oss << " (from round=\"" << attrs.round << "\")"; + } + oss << " is not supported by the FP8 conversion policy"; + return reject(oss.str()); + } + plan.steps.push_back({VectorCastOp::kNonIntegerConvert, from, to, token, + (attrs.has_sat && !attrs.sat) ? "NOSAT" : "SAT"}); + return CastPlan{plan}; + } + return reject(pair + ": FP8 -> FP16/BF16 and FP8 <-> integer are out of " + "scope; convert through float32 explicitly"); + } + const bool from_float = from.is_float() || from.is_bfloat16(); + const bool to_float = to.is_float() || to.is_bfloat16(); + if (from_float && to_float) { + if (attrs.has_round_mode) { + return reject(pair + ": this conversion does not accept a rounding mode"); + } + const bool from_f32 = from.is_float() && from.bits() == 32; + const bool to_f32 = to.is_float() && to.bits() == 32; + VectorCastPlan plan; + if (!from_f32 && to_f32) { + if (attrs.has_sat) { + return reject(pair + ": this direction does not accept a saturation " + "mode"); + } + // A real conversion: empty options are not a no-op. + plan.steps.push_back( + {VectorCastOp::kNonIntegerConvert, from, to, "", ""}); + return CastPlan{plan}; + } + if (from_f32 && !to_f32) { + if (attrs.has_sat && attrs.sat) { + return reject(pair + ": this conversion accepts only sat=False " + "(NOSAT)"); + } + // Pin R/NOSAT so the generated vcvt never depends on PTODSL defaults. + plan.steps.push_back( + {VectorCastOp::kNonIntegerConvert, from, to, "R", "NOSAT"}); + return CastPlan{plan}; + } + return reject(pair + ": cross-family f16 <-> bf16 narrowing has no " + "validated rounding semantics; convert through " + "float32 explicitly"); + } + return reject(pair + ": float <-> integer conversions need per-combination " + "rounding/overflow validation"); +} + +} // namespace + +Optional PlanValueCast(const CastNode *cast, + const CastContext &context, + std::string *reason) { + DataType from = cast->value.dtype(); + DataType to = cast->dtype; + auto reject = [&](const std::string &why) { + if (reason != nullptr) { + *reason = why; + } + return Optional(); + }; + std::ostringstream pair_os; + pair_os << "`" << from << " -> " << to << "` (" << CastContextName(context) + << ")"; + const std::string pair = pair_os.str(); + // The explicit kind is authoritative; lanes only validate it. Inferring + // the kind from lanes would turn a lane-1 VMI value into a scalar and a + // scalar request with vector lanes into an implicit VMI fallback. + if (context.value_kind == CastValueKind::kScalar) { + if (context.logical_lanes != 1) { + return reject(pair + + ": a scalar cast context must have exactly one " + "logical lane, got " + + std::to_string(context.logical_lanes)); + } + } else { + if (!IsVmiLogicalLaneCount(context.logical_lanes)) { + return reject(pair + + ": a VMI cast context requires logical lanes " + "64/128/256, got " + + std::to_string(context.logical_lanes)); + } + } + // Dtype legality is a precondition for every action plan, + // including identity — unsupported types (int24/FP64/...) must not slip + // through the same-type fast path. The Parallel value path is + // element-wise, so both sides must be lane-1 elements. + if (from.lanes() != 1 || to.lanes() != 1) { + return reject(pair + ": value casts require lane-1 element operands"); + } + if (!IsSupportedElementDType(from) || !IsSupportedElementDType(to)) { + return reject(pair + " is outside the supported element dtype set"); + } + + // Attributes are parsed for every Cast, not just the FP8 ones: the + // expression whitelist deliberately does not look inside a cast's + // annotations (attributes are not computation values), so an annotation the + // selected conversion cannot honour has to be rejected here — for every + // direction — or it would be dropped silently. + CastAttributes attrs; + std::string why; + if (!ReadCastAttributes(cast->annotations, &attrs, &why)) { + return reject(pair + ": " + why); + } + if (attrs.has_rbits) { + return reject(pair + ": the rbits operand (stochastic rounding) is not " + "supported"); + } + std::string round_token; + if (attrs.has_round_mode && + !NormalizeCastRounding(attrs.round, &round_token, &why)) { + return reject(pair + ": " + why); + } + + if (context.value_kind == CastValueKind::kScalar) { + return PlanScalarValueCast(cast, attrs, pair, reason); + } + return PlanVectorValueCast(cast, context, attrs, round_token, pair, reason); +} + +std::optional FindUnsupportedExprNode(const PrimExpr &expr) { + 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..c410d6cccb --- /dev/null +++ b/src/ascend/transform/parallel_to_pto_utils.h @@ -0,0 +1,386 @@ +/*! + * \file parallel_to_pto_utils.h + * \brief Shared read-only analysis helpers for PTO Parallel vectorization. + * + * VerifyParallelToPTO consumes these helpers to check the input contract; + * VectorizeParallelToPTO re-runs the same analysis to build its VMI code. + * 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 +#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); + +/*! Value category of one value-Cast position. The kind is supplied by the + * caller (Verify from AnalyzeLaneUse, Vectorize from the converted input + * shape); the planner never infers it from lanes. */ +enum class CastValueKind { + kScalar, + kVectorVmi, +}; + +/*! Per-Cast planning context. `logical_lanes` is validation information + * only; `value_kind` is the dispatch key. */ +struct CastContext { + CastValueKind value_kind = CastValueKind::kScalar; + int64_t logical_lanes = 1; +}; + +/*! Scalar value-Cast plan: either keep the input (identity) or rebuild an + * ordinary scalar Cast with the validated annotations and the original span. + * No new scalar rounding/saturation mode is introduced. */ +enum class ScalarCastAction { + kKeepInput, + kRebuildCast, +}; + +struct ScalarCastPlan { + ScalarCastAction action = ScalarCastAction::kKeepInput; + /*! Validated annotations to carry onto the rebuilt Cast (the same shape as + * CastNode::annotations). Attributes that are proven to match the + * existing scalar lowering are normalized away, so this map is empty for + * the currently admitted directions. */ + ffi::Map annotations; +}; + +/*! One ordered vector conversion action. */ +enum class VectorCastOp { + /*! Same-sign widening: one logical vcvt. */ + kIntegerExtend, + /*! Same-width signedness change / source-sign re-view: one + * vinterpret_cast. */ + kIntegerReinterpret, + /*! Narrowing including cross-sign: one logical vcvt with NOSAT. */ + kIntegerTruncate, + /*! Real non-integer conversion (f16/bf16/FP8 -> f32, f32 -> f16/bf16/FP8); + * empty options still mean a conversion, never a no-op. */ + kNonIntegerConvert, +}; + +struct VectorCastStep { + VectorCastOp op = VectorCastOp::kIntegerExtend; + DataType from; // element type consumed by this step + DataType to; // element type produced by this step + /*! PTODSL rounding annotation value, e.g. "R"; empty when N/A. */ + std::string rounding; + /*! PTODSL saturate annotation value, e.g. "NOSAT"; empty when N/A. */ + std::string saturate; +}; + +struct VectorCastPlan { + /*! Ordered actions; an empty list is a proven identity (no conversion). + * A non-empty list still may be a single real conversion with empty + * options (that is not an identity). */ + std::vector steps; +}; + +using CastPlan = std::variant; + +/*! Operation-level type support for the vector arithmetic emitted + * here. VMI vmul only supports i16/i32/f16/bf16/f32 + * elements; vadd/vsub additionally allow i8; vdiv admits f16/f32 only. + * 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 plan for a value Cast. This is the single admission + * point for every Cast on the vectorize path; Verify and Vectorize both pass + * an explicit CastContext and consume the same plan, so they cannot diverge. + * + * Dispatch is first scalar vs VMI, then integer vs non-integer: + * - scalar: identity KeepInput; integer widening/narrowing and the newly + * admitted same-width signedness change rebuild an ordinary scalar Cast; + * f16/bf16 <-> f32 keep their existing scalar lowering; non-identity FP8, + * float<->int and f16<->bf16 stay rejected; + * - VMI integer: identity 0 steps; same-width cross-sign one + * vinterpret_cast; same-sign widening one vcvt; cross-sign widening + * source-sign widen then same-width vinterpret_cast; narrowing one direct + * vcvt NOSAT. All 36 supported pairs are represented; no consumer-layout + * feasibility is guessed here; + * - VMI non-integer: a real conversion action for f16/bf16/FP8 -> f32 + * (empty options are not a no-op) and for f32 -> f16/bf16/FP8 with the + * existing rounding/saturation policy. + * + * `cast` carries the backend annotations written by the frontend cast + * helpers (rounding mode, saturation flag, stochastic rbits operand). They + * are read here — never interpreted by the expression visitors — and an + * attribute the selected conversion cannot honour is a rejection, not a + * silently dropped hint. When `reason` is non-null it receives a diagnostic + * naming the category, source/target types, lanes and the offending option. */ +Optional PlanValueCast(const CastNode *cast, + const CastContext &context, + std::string *reason = nullptr); + +/*! Check whether an expression only uses nodes from the supported + * 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..3994eadfcb --- /dev/null +++ b/src/ascend/transform/vectorize_parallel_to_pto.cc @@ -0,0 +1,1057 @@ +// 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.vmi.vinterpret_cast", + "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()) { + // Convert the input exactly once, then plan from the explicit shape of + // that converted input (the original Cast carries the element types). + PrimExpr value = ConvertExpr(cast->value, pre); + pto_analysis::CastContext context; + if (value.dtype().lanes() == 1) { + context.value_kind = pto_analysis::CastValueKind::kScalar; + context.logical_lanes = 1; + } else { + ICHECK_EQ(value.dtype().lanes(), lanes_) + << "[VectorizeParallelToPTO] internal error: Cast input has " + "unexpected lanes " + << value.dtype().lanes() << " (current VL " << lanes_ << ")"; + context.value_kind = pto_analysis::CastValueKind::kVectorVmi; + context.logical_lanes = lanes_; + } + ICHECK_EQ(value.dtype().element_of(), cast->value.dtype()) + << "[VectorizeParallelToPTO] internal error: converted Cast input " + "element type " + << value.dtype().element_of() << " does not match the original " + << cast->value.dtype(); + std::string plan_why; + auto plan = pto_analysis::PlanValueCast(cast, context, &plan_why); + ICHECK(plan.has_value()) + << "[VectorizeParallelToPTO] internal error: Cast passed Verify " + "without a conversion plan: " + << plan_why; + if (const auto *scalar = + std::get_if(&plan.value())) { + ICHECK_EQ(value.dtype().lanes(), 1) + << "[VectorizeParallelToPTO] internal error: scalar Cast plan " + "for a vector input"; + if (scalar->action == pto_analysis::ScalarCastAction::kKeepInput) { + return value; + } + // Rebuild the ordinary scalar Cast with the validated annotations and + // the original span; the scalar lowering itself is unchanged. + return Cast(cast->dtype, value, scalar->annotations, cast->span); + } + const auto &vector = std::get(plan.value()); + PrimExpr current = value; + for (const auto &step : vector.steps) { + ICHECK_EQ(current.dtype().element_of(), step.from) + << "[VectorizeParallelToPTO] internal error: Cast step expects " + << step.from << ", got " << current.dtype().element_of(); + Map attrs; + attrs.Set("to_dtype", StringImm(TvmDtypeName(step.to))); + PrimExpr call; + if (step.op == pto_analysis::VectorCastOp::kIntegerReinterpret) { + call = Call(VectorDType(step.to, lanes_), VmiOp("vinterpret_cast"), + {current}, attrs); + } else { + if (!step.saturate.empty()) { + attrs.Set("saturate", StringImm(step.saturate)); + } + if (!step.rounding.empty()) { + attrs.Set("rounding", StringImm(step.rounding)); + } + call = Call(VectorDType(step.to, lanes_), VmiOp("vcvt"), {current}, + attrs); + } + current = EmitBind("cast", call, pre); + } + // Zero steps is the proven identity; a final type check keeps every + // multi-step chain honest (source convert once, dtype/VL continuous). + ICHECK_EQ(current.dtype(), VectorDType(cast->dtype, lanes_)) + << "[VectorizeParallelToPTO] internal error: Cast plan ended at " + << current.dtype() << ", expected " + << VectorDType(cast->dtype, lanes_); + return current; + } + LOG(FATAL) << "[VectorizeParallelToPTO] internal error: unsupported " + "expression passed Verify: " + << 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..f0bd64e1c8 --- /dev/null +++ b/src/ascend/transform/verify_parallel_to_pto.cc @@ -0,0 +1,1365 @@ +// 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 { + // Classify the *input of this Cast* with the existing lane-use + // analysis, then hand the explicit kind to the shared planner. + pto_analysis::LaneUseContext ctx = parent_->MakeLaneUseContext(); + std::string lane_why; + pto_analysis::LaneUse use = + pto_analysis::AnalyzeLaneUse(op->value, &ctx, &lane_why); + if (use == pto_analysis::LaneUse::kUnknown) { + std::ostringstream oss; + oss << "Cast " << op->value.dtype() << " -> " << op->dtype + << " cannot classify its input: " << lane_why; + parent_->parent_->Fail("dtype", oss.str()); + return; + } + pto_analysis::CastContext context; + if (use == pto_analysis::LaneUse::kUniform) { + context.value_kind = pto_analysis::CastValueKind::kScalar; + context.logical_lanes = 1; + } else { + context.value_kind = pto_analysis::CastValueKind::kVectorVmi; + context.logical_lanes = parent_->mapping_.lanes; + } + std::string why; + auto plan = pto_analysis::PlanValueCast(op, context, &why); + if (!plan.has_value()) { + std::ostringstream oss; + oss << "Cast " << op->value.dtype() << " -> " << op->dtype + << " has no supported conversion"; + if (!why.empty()) { + oss << ": " << why; + } + parent_->parent_->Fail("dtype", oss.str()); + } + // Nested value casts each plan independently; annotations are not + // computation values and are validated by the planner. + VisitExpr(op->value); + } + void VisitExpr_(const BufferLoadNode *) final { + // Index casts belong to address analysis; do not descend. + } + + 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/ascend/transform/verify_pto_codegen_legality.cc b/src/ascend/transform/verify_pto_codegen_legality.cc new file mode 100644 index 0000000000..18b3d1ab5c --- /dev/null +++ b/src/ascend/transform/verify_pto_codegen_legality.cc @@ -0,0 +1,1162 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file verify_pto_codegen_legality.cc + * \brief Verify PTO requirements before code generation without modifying IR. + * + * This pass consumes the final device TIR after LowerDeviceKernelLaunch and + * the final layout and UB rewrites. It currently checks destination alignment + * for default/continuous explicit-mask VMI stores of b8/b16/b32 data. + * Ordinary masked stores require provable 32-byte alignment. Stores selected + * for 1PT_B8/B16/B32 writeback require natural element alignment (1/2/4 bytes): + * the value must have 1/2/4/8 lanes and the mask must be a recognized + * full-active seed with the same lane count. Bool data is excluded from this + * check. The checker resolves byte addresses using the Analyzer facts at each + * access. Shapes, layouts, masks and addresses are never modified. Functions + * without a PTO target pass through unchanged. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arith/ir_visitor_with_analyzer.h" +#include "tirx/transform/ir_utils.h" + +#include "../../op/builtin.h" +#include "../target_utils.h" + +namespace tvm { +namespace tl { +namespace transform { + +using namespace tirx; + +namespace { + +constexpr int kOrdinaryMaskedStoreAlignment = 32; + +enum class StoreMode { + kOrdinaryMasked, + kPointStores, +}; + +enum class Outcome { + kAligned, + kMisaligned, + kUnknown, +}; + +struct ByteAddress { + enum class Kind { + kUbRoot, + kIntegerBase, + kUnknown, + }; + Kind kind = Kind::kUnknown; + const VarNode *root = nullptr; + PrimExpr base_bytes; + PrimExpr offset_bytes; + std::string error; +}; + +struct StoreMatch { + bool checked = false; + StoreMode mode = StoreMode::kOrdinaryMasked; + const char *mode_name = ""; + int required_alignment = kOrdinaryMaskedStoreAlignment; + int elem_bytes = 0; + DataType value_dtype; + // Element type of the pointer that reaches the store after descriptor, + // address_of, typed Var/make_tensor and pure alias resolution. It is the + // unit of the instruction offset; it is never defaulted to float32. + std::optional pointer_dtype; + bool has_descriptor = false; + DataType descriptor_dtype = DataType::Float(32); + PrimExpr descriptor_offset; + PrimExpr instruction_offset; +}; + +struct SeedInfo { + bool recognized = false; + bool full_active = false; + int64_t mask_lanes = 0; +}; + +struct Proof { + Outcome outcome = Outcome::kUnknown; + PrimExpr byte_offset; + PrimExpr remainder; + std::string error; +}; + +} // namespace + +/*! Checker for explicit-mask continuous VMI stores. Traversal is scoped: + * loop/branch/attribute bodies push and pop the pointer/pure-definition + * maps, and the Analyzer provides the existing range/condition facts. */ +class PTOMaskedStoreAlignmentChecker final + : public arith::IRVisitorWithAnalyzer { +public: + explicit PTOMaskedStoreAlignmentChecker(std::string function_name) + : function_name_(std::move(function_name)) {} + + using Base = arith::IRVisitorWithAnalyzer; + using Base::VisitExpr_; + using Base::VisitStmt_; + + void Run(const PrimFunc &func) { VisitStmt(func->body); } + + // UB arena roots: PTO codegen materializes the merged dyn-shared + // allocation at UB byte address 0, and MergeUBAllocations aligns the arena + // to 32 bytes. Only allocations known to feed that arena are roots; a raw + // handle is never assumed aligned. Roots are recorded in lexical order + // during the single statement walk, so expression DAGs are never re-walked + // by a whole-body pre-scan. + void VisitStmt_(const AllocBufferNode *op) final { + const Var &data = op->buffer->data; + const std::string scope = GetPtrStorageScope(data); + if (scope == "shared.dyn" || scope == "shared") { + ub_roots_.insert(data.get()); + } + StmtExprVisitor::VisitStmt_(op); + } + + void VisitStmt_(const ForNode *op) final { + constraint_scope_.WithNewScope([&]() { + analyzer_.Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); + constraint_scope_.Current().Emplace(&analyzer_, op->extent > 0); + if (op->step.defined()) { + const auto *step = op->step.as(); + if (step != nullptr && step->value > 1) { + // The base visitor binds [min, min+extent) but drops the loop + // lattice. Record the real iteration congruence so a positive + // constant step such as i = 0, 8, 16, ... can be proven. + constraint_scope_.Current().Emplace( + &analyzer_, + floormod(op->loop_var - op->min, op->step.value()) == 0); + } + } + AliasScope saved = SaveAliases(); + StmtExprVisitor::VisitStmt_(op); + RestoreAliases(std::move(saved)); + }); + } + + void VisitStmt_(const SBlockNode *op) final { + constraint_scope_.WithNewScope([&]() { + for (const auto &iter_var : op->iter_vars) { + analyzer_.Bind(iter_var->var, iter_var->dom); + } + AliasScope saved = SaveAliases(); + StmtExprVisitor::VisitStmt_(op); + RestoreAliases(std::move(saved)); + }); + } + + void VisitStmt_(const IfThenElseNode *op) final { + constraint_scope_.WithNewScope([&]() { + this->VisitExpr(op->condition); + PrimExpr real_condition = ExtractRealCondition(op->condition); + constraint_scope_.WithNewScope([&]() { + constraint_scope_.Current().Emplace(&analyzer_, real_condition); + AliasScope saved = SaveAliases(); + this->VisitStmt(op->then_case); + RestoreAliases(std::move(saved)); + }); + if (op->else_case) { + constraint_scope_.WithNewScope([&]() { + constraint_scope_.Current().Emplace( + &analyzer_, analyzer_.rewrite_simplify(Not(real_condition))); + AliasScope saved = SaveAliases(); + this->VisitStmt(op->else_case.value()); + RestoreAliases(std::move(saved)); + }); + } + }); + } + + void VisitStmt_(const AttrStmtNode *op) final { + constraint_scope_.WithNewScope([&]() { + if (op->attr_key == tirx::attr::thread_extent || + op->attr_key == s_tir::attr::virtual_thread) { + IterVar iv = Downcast(op->node); + TVM_FFI_ICHECK_NE(iv->thread_tag.length(), 0U); + analyzer_.Bind(iv->var, Range::FromMinExtent( + IntImm(op->value->dtype, 0), op->value)); + } else if (op->attr_key == tirx::attr::tilelang_assume) { + auto condition = Downcast(op->node); + constraint_scope_.Current().Emplace(&analyzer_, condition, + /*is_assume=*/true); + } + AliasScope saved = SaveAliases(); + StmtExprVisitor::VisitStmt_(op); + RestoreAliases(std::move(saved)); + }); + } + + void VisitStmt_(const BindNode *op) final { + std::unordered_set visiting{op->var.get()}; + ExprMemo memo; + ScalarFacts facts = AnalyzeScalar(op->value, &visiting, &memo); + if (facts.needs_generic_visit) { + this->VisitExpr(op->value); + } + RegisterDefinition(op->var.get(), op->value, facts); + } + + void VisitExpr_(const LetNode *op) final { + std::unordered_set visiting{op->var.get()}; + ExprMemo value_memo; + ScalarFacts value_facts = AnalyzeScalar(op->value, &visiting, &value_memo); + if (value_facts.needs_generic_visit) { + this->VisitExpr(op->value); + } + AliasScope saved = SaveAliases(); + RegisterDefinition(op->var.get(), op->value, value_facts); + std::unordered_set body_visiting{op->var.get()}; + ExprMemo body_memo; + ScalarFacts body_facts = + AnalyzeScalar(op->body, &body_visiting, &body_memo); + if (body_facts.needs_generic_visit) { + this->VisitExpr(op->body); + } + RestoreAliases(std::move(saved)); + } + + void VisitExpr_(const CallNode *op) final { + if (op->op.same_as(tl::vmi_vstore())) { + CheckStore(op); + } + Base::VisitExpr_(op); + } + +private: + struct AliasScope { + std::unordered_map pointer_defs; + std::unordered_map value_defs; + }; + + AliasScope SaveAliases() const { + return AliasScope{pointer_defs_, value_defs_}; + } + + void RestoreAliases(AliasScope saved) { + pointer_defs_ = std::move(saved.pointer_defs); + value_defs_ = std::move(saved.value_defs); + } + + struct ScalarFacts { + bool origin_safe = false; + bool needs_generic_visit = false; + bool has_shared_node = false; + }; + using ExprMemo = std::unordered_map; + using CanonMemo = std::unordered_map; + + // Bottom-up, memoized facts for one expression / right-hand side. + // + // origin_safe: the expression is a scalar (lanes == 1, non-handle) tree of + // immediates, free variables, casts and pure scalar arithmetic, and every + // referenced bind definition is origin-safe as well (a lane-1 VMI alias, a + // memory read, a Call or a cycle is unsafe). Plain SSA chains of any length + // are classified through the per-Var cache one shallow definition at a time. + // + // needs_generic_visit: the generic statement/expression walk is required + // because the expression may contain a checked store (a void Call), a Let + // definition, branch constraints, Reduce axes or an unrecognized node. A + // pure elementwise arithmetic expression cannot contain any of these, so + // skipping the walk is safe for nested checked-store discovery. + // + // has_shared_node: a non-leaf expression node is referenced more than once + // within this right-hand side (a DAG). This only selects the canonicalized + // Analyzer binding; the definition and its equality are always kept. + ScalarFacts AnalyzeScalar(const PrimExpr &value, + std::unordered_set *visiting, + ExprMemo *memo, int depth = 0) { + // Recursion/stack-safety guard only, not a complexity or acceptance + // bound: repeated work is already bounded by the per-RHS memo and the + // per-Var classification cache. An expression nested deeper than this is + // treated conservatively (no Analyzer equality and a generic visit), so a + // dependent store is diagnosed as unprovable instead of risking the + // compiler stack. + constexpr int kMaxExpressionRecursionGuard = 4096; + if (depth > kMaxExpressionRecursionGuard) { + return ScalarFacts{false, true, false}; + } + const DataType dtype = value.dtype(); + if (value.as() || value.as() || + value.as()) { + return ScalarFacts{true, false, false}; + } + if (const auto *var = value.as()) { + if (visiting->count(var) != 0) { + // Self/cyclic reference: not a lawful SSA definition. + return ScalarFacts{false, false, false}; + } + auto safe_it = var_origin_safe_.find(var); + if (safe_it != var_origin_safe_.end()) { + return ScalarFacts{safe_it->second, false, false}; + } + auto def_it = value_defs_.find(var); + if (def_it == value_defs_.end()) { + // Function/loop/block parameter without a recorded definition. + return ScalarFacts{true, false, false}; + } + visiting->insert(var); + // Analyze the definition with its own memo: sharing inside a separate + // right-hand side must not mark the current expression as shared. + ExprMemo def_memo; + ScalarFacts def_facts = + AnalyzeScalar(def_it->second, visiting, &def_memo, depth + 1); + visiting->erase(var); + var_origin_safe_[var] = def_facts.origin_safe; + return ScalarFacts{def_facts.origin_safe, false, false}; + } + auto memo_it = memo->find(value.get()); + if (memo_it != memo->end()) { + return ScalarFacts{memo_it->second.origin_safe, + memo_it->second.needs_generic_visit, true}; + } + auto child = [&](const PrimExpr &expr) { + return AnalyzeScalar(expr, visiting, memo, depth + 1); + }; + const bool scalar_root = !dtype.is_handle() && dtype.lanes() == 1; + auto binary = [&](const PrimExpr &a, const PrimExpr &b) { + ScalarFacts fa = child(a); + ScalarFacts fb = child(b); + ScalarFacts out; + out.origin_safe = scalar_root && fa.origin_safe && fb.origin_safe; + out.needs_generic_visit = + fa.needs_generic_visit || fb.needs_generic_visit; + out.has_shared_node = fa.has_shared_node || fb.has_shared_node; + return out; + }; + auto ternary = [&](const PrimExpr &a, const PrimExpr &b, + const PrimExpr &c) { + ScalarFacts fa = child(a); + ScalarFacts fb = child(b); + ScalarFacts fc = child(c); + ScalarFacts out; + out.origin_safe = + scalar_root && fa.origin_safe && fb.origin_safe && fc.origin_safe; + out.needs_generic_visit = fa.needs_generic_visit || + fb.needs_generic_visit || + fc.needs_generic_visit; + out.has_shared_node = + fa.has_shared_node || fb.has_shared_node || fc.has_shared_node; + return out; + }; + ScalarFacts facts; + if (const auto *node = value.as()) { + ScalarFacts c = child(node->value); + facts.origin_safe = scalar_root && c.origin_safe; + facts.needs_generic_visit = c.needs_generic_visit; + facts.has_shared_node = c.has_shared_node; + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = ternary(node->condition, node->true_value, node->false_value); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + facts = binary(node->a, node->b); + } else if (const auto *node = value.as()) { + ScalarFacts c = child(node->a); + facts.origin_safe = scalar_root && c.origin_safe; + facts.needs_generic_visit = c.needs_generic_visit; + facts.has_shared_node = c.has_shared_node; + } else if (const auto *node = value.as()) { + bool needs = false; + bool shared = false; + for (const auto &index : node->indices) { + ScalarFacts f = child(index); + needs = needs || f.needs_generic_visit; + shared = shared || f.has_shared_node; + } + if (node->predicate.defined()) { + ScalarFacts f = child(node->predicate.value()); + needs = needs || f.needs_generic_visit; + shared = shared || f.has_shared_node; + } + facts = ScalarFacts{false, needs, shared}; + } else if (const auto *node = value.as()) { + ScalarFacts f = binary(node->base, node->stride); + facts = ScalarFacts{false, f.needs_generic_visit, f.has_shared_node}; + } else if (const auto *node = value.as()) { + ScalarFacts f = child(node->value); + facts = ScalarFacts{false, f.needs_generic_visit, f.has_shared_node}; + } else if (const auto *node = value.as()) { + bool needs = false; + bool shared = false; + for (const auto &vector : node->vectors) { + ScalarFacts f = child(vector); + needs = needs || f.needs_generic_visit; + shared = shared || f.has_shared_node; + } + for (const auto &index : node->indices) { + ScalarFacts f = child(index); + needs = needs || f.needs_generic_visit; + shared = shared || f.has_shared_node; + } + facts = ScalarFacts{false, needs, shared}; + } else { + // Calls (including every VMI producer and vstore), Let, Reduce and + // every other opaque node keep the generic walk and are not scalar + // equalities. + facts = ScalarFacts{false, true, false}; + } + (*memo)[value.get()] = facts; + return facts; + } + + // Rebuild an origin-safe expression with each distinct non-leaf node bound + // to a fresh shallow symbol. The Analyzer's scalar binding walks its + // argument recursively, so handing it a shared DAG (e+k; e=e+e; ...) is + // exponential even though the classification above is memoized. With one + // shallow binding per distinct node the substitution graph stays linear and + // the equality -- including its modular/alignment consequences -- is fully + // preserved. + PrimExpr CanonicalizeForAnalyzer(const PrimExpr &value, CanonMemo *memo) { + if (value.as() || value.as() || + value.as() || value.as()) { + return value; + } + auto memo_it = memo->find(value.get()); + if (memo_it != memo->end()) { + return memo_it->second; + } + auto canon = [&](const PrimExpr &expr) { + return CanonicalizeForAnalyzer(expr, memo); + }; + PrimExpr rebuilt; + if (const auto *node = value.as()) { + rebuilt = + Cast(node->dtype, canon(node->value), node->annotations, node->span); + } else if (const auto *node = value.as()) { + rebuilt = Add(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Sub(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Mul(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Div(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = FloorDiv(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Mod(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = FloorMod(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Min(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Max(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Select(canon(node->condition), canon(node->true_value), + canon(node->false_value), node->span); + } else if (const auto *node = value.as()) { + rebuilt = LE(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = LT(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = GT(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = GE(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = EQ(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = NE(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = And(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Or(canon(node->a), canon(node->b), node->span); + } else if (const auto *node = value.as()) { + rebuilt = Not(canon(node->a), node->span); + } else { + // Unreachable for a classified origin-safe expression; keep the value + // rather than inventing a different semantic. + (*memo)[value.get()] = value; + return value; + } + Var fresh("__pto_align_dag_" + std::to_string(next_dag_symbol_++), + rebuilt.dtype()); + analyzer_.Bind(fresh, rebuilt); + (*memo)[value.get()] = fresh; + return fresh; + } + + void RegisterDefinition(const VarNode *var, const PrimExpr &value, + const ScalarFacts &facts) { + if (value.dtype().is_handle()) { + pointer_defs_[var] = value; + return; + } + // Scalar/index definitions and pointer aliases are recorded separately: + // value_defs_ backs the origin-safe scalar check and the mask-seed alias + // walk, while pointer_defs_ backs address resolution. Seed matching only + // accepts recognized create_mask/pset calls and compile-time active + // counts, so a memory-read definition never classifies as a seed. + value_defs_[var] = value; + const bool safe = value.dtype().is_scalar() && !value.dtype().is_handle() && + facts.origin_safe; + var_origin_safe_[var] = safe; + if (!safe) { + return; + } + if (facts.has_shared_node) { + // Preserve the equality with a compact substitution graph instead of + // dropping it (which would lose alignment proofs) or expanding the DAG. + CanonMemo canon_memo; + PrimExpr canonical = CanonicalizeForAnalyzer(value, &canon_memo); + analyzer_.Bind(ffi::GetRef(var), canonical); + } else { + analyzer_.Bind(ffi::GetRef(var), value); + } + } + + // Follow immutable SSA aliases iteratively: the length of a plain alias + // chain is not a recursion-depth budget. A visited set detects cycles. + std::optional ResolveMaskActiveCount(const PrimExpr &expr) const { + PrimExpr current = expr; + std::unordered_set seen; + while (const auto *var = current.as()) { + if (!seen.insert(var).second) { + return std::nullopt; + } + auto it = value_defs_.find(var); + if (it == value_defs_.end()) { + return std::nullopt; + } + current = it->second; + } + if (const auto *imm = current.as()) { + return imm->value; + } + return std::nullopt; + } + + PrimExpr ResolvePureAlias(const PrimExpr &expr) const { + PrimExpr current = expr; + std::unordered_set seen; + while (const auto *var = current.as()) { + if (!seen.insert(var).second) { + return current; // cycle: not a recognized seed + } + auto it = value_defs_.find(var); + if (it == value_defs_.end()) { + return current; + } + current = it->second; + } + return current; + } + + // Mask seeds recognized by the PTO backend for 1PT_B8/B16/B32 writeback: + // + // * A default/continuous vstore with 1/2/4/8 data lanes can use point + // writeback when its create_mask(active, size) seed has the same lane + // count and a compile-time active count at least as large as size. + // This selects full-active pset_b32 PAT_ALL and one point store per + // logical element. The value need not come from a reduction. + // * A pset PAT_ALL mask is also a recognized seed. TileLang does not expose + // a pset operator in its Python API; this branch handles direct TIR calls. + // * create_mask(0), partial masks and dynamic active counts do not match. + // These remain ordinary masked stores and require 32-byte alignment. + // Mismatched value/mask lanes are rejected by the backend, and arbitrary + // all-true expressions are not recognized seeds. + SeedInfo ClassifyMaskSeed(const PrimExpr &mask) const { + SeedInfo info; + const DataType mask_dtype = mask.dtype(); + if (!mask_dtype.is_bool() || mask_dtype.lanes() <= 0) { + return info; + } + info.mask_lanes = mask_dtype.lanes(); + PrimExpr resolved = ResolvePureAlias(mask); + const auto *call = resolved.as(); + if (call == nullptr) { + return info; + } + auto opt_op = call->op.as(); + if (!opt_op.has_value()) { + return info; + } + const std::string &name = opt_op.value()->name; + if (name == "tl.vmi.create_mask") { + info.recognized = true; + if (!call->args.empty()) { + std::optional active = ResolveMaskActiveCount(call->args[0]); + info.full_active = active.has_value() && *active >= info.mask_lanes; + } + return info; + } + if (name == "tl.vmi.pset" && !call->args.empty()) { + if (const auto *dist = call->args[0].as()) { + if (dist->value == "PAT_ALL") { + info.recognized = true; + info.full_active = true; + } + } + } + return info; + } + + std::optional DecodeAccessDType(const PrimExpr &annotation) const { + const auto *call = annotation.as(); + if (call == nullptr) { + return std::nullopt; + } + std::string n; + if (!call->args.empty()) { + if (const auto *name = call->args[0].as()) { + n = name->value; + } + } + if (n.empty()) { + // The final device IR carries a `tirx.type_annotation` call whose + // result dtype is the annotated element type (the call has no + // positional arguments at this stage). Use that dtype so descriptor + // and instruction offsets are scaled exactly as codegen does. + auto opt_op = call->op.as(); + if (!opt_op.has_value() || + opt_op.value()->name != "tirx.type_annotation") { + return std::nullopt; + } + DataType annotated = call->dtype; + if (!annotated.is_scalar() || annotated.is_handle()) { + return std::nullopt; + } + return annotated; + } + if (n == "float" || n == "float32") + return DataType::Float(32); + if (n == "half" || n == "float16") + return DataType::Float(16); + if (n == "bfloat16" || n == "bfloat16_t") + return DataType::BFloat(16); + if (n == "float8_e4m3" || n == "float8_e4m3fn") { + return DataType::Float8E4M3FN(); + } + if (n == "float8_e5m2") + return DataType::Float8E5M2(); + if (n == "int8") + return DataType::Int(8); + if (n == "uint8") + return DataType::UInt(8); + if (n == "int16") + return DataType::Int(16); + if (n == "uint16") + return DataType::UInt(16); + if (n == "int32") + return DataType::Int(32); + if (n == "uint32") + return DataType::UInt(32); + if (n == "int64") + return DataType::Int(64); + if (n == "uint64") + return DataType::UInt(64); + return std::nullopt; + } + + std::optional + PointerAnnotationElementType(const VarNode *var) const { + const auto *ptr = var->type_annotation.as(); + if (ptr == nullptr) { + return std::nullopt; + } + const auto *prim = ptr->element_type.as(); + if (prim == nullptr) { + return std::nullopt; + } + DataType dtype = prim->dtype; + if (dtype.is_void() || dtype.is_handle() || dtype.lanes() != 1) { + return std::nullopt; + } + return dtype; + } + + // Element type of the pointer that the store actually receives, matching + // codegen exactly (codegen_pto.cc Bind handling): + // + // * A typed Var bound to ``reinterpret(handle, int_addr)`` (T.make_tensor) + // emits ``v = pto.castptr(addr, ptr)``: this binding's own annotation + // is a real castptr and establishes the type. + // * Every other binding emits ``v = `` with no cast. A typed + // annotation on a plain alias therefore does NOT re-type the pointer; the + // defining expression's type wins. + // * ``tvm_access_ptr``/``address_of`` print through GetPointerExpr and cast + // to their own descriptor/buffer element type, overriding any outer + // annotation. + // * A pointer type re-view (``reinterpret(handle, pointer)``) prints the + // inner expression unchanged. + // + // Plain immutable SSA aliases are followed iteratively (their length is not + // a recursion budget); only pointer expression nesting consumes the + // stack-safety guard. Only the Var directly bound to the terminal + // expression can establish a type; an untyped/unknown pointer is not + // guessed as float32. + std::optional + ResolvePointerElementType(const PrimExpr &ptr, + std::unordered_set *visiting, + int nesting = 0) const { + // Stack-safety guard only; not an alias-length or acceptance limit. + constexpr int kMaxPointerExpressionNesting = 4096; + if (nesting > kMaxPointerExpressionNesting) { + return std::nullopt; + } + const VarNode *defining_var = nullptr; + PrimExpr current = ptr; + while (const auto *var = current.as()) { + if (!visiting->insert(var).second) { + return std::nullopt; // alias cycle + } + auto def_it = pointer_defs_.find(var); + if (def_it == pointer_defs_.end()) { + // Terminal Var without a binding (parameter/root): its own annotation + // is the type of the pointer it denotes. + return PointerAnnotationElementType(var); + } + defining_var = var; + current = def_it->second; + } + if (const auto *call = current.as()) { + if (call->op.same_as(builtin::tvm_access_ptr())) { + if (!call->args.empty()) { + return DecodeAccessDType(call->args[0]); + } + return std::nullopt; + } + if (call->op.same_as(builtin::address_of())) { + if (call->args.size() == 1U) { + if (const auto *load = call->args[0].as()) { + DataType dtype = load->dtype.element_of(); + if (!dtype.is_void() && !dtype.is_handle()) { + return dtype; + } + } + } + return std::nullopt; + } + if (call->op.same_as(tirx::builtin::reinterpret()) && + call->args.size() == 1U) { + if (call->args[0].dtype().is_handle()) { + // Pointer type re-view: the inner pointer's type is the emitted + // type; an outer Var annotation does not emit a cast. + return ResolvePointerElementType(call->args[0], visiting, + nesting + 1); + } + // Integer-address make_tensor binding: the castptr uses the + // annotation of the Var directly bound to this reinterpret. + if (defining_var != nullptr) { + return PointerAnnotationElementType(defining_var); + } + return std::nullopt; + } + } + return std::nullopt; + } + + StoreMatch MatchStore(const CallNode *op) const { + StoreMatch match; + // Default continuous masked form only: no mode-changing annotations. + for (const auto &[key_ref, value] : op->annotations) { + const std::string key = key_ref; + if (key == "dist_mode") { + const auto *mode = value.as(); + if (mode == nullptr || (mode->value != "continuous")) { + return match; + } + continue; + } + if (key == "group" || key == "stride" || key == "block_stride" || + key == "repeat_stride" || key == "pmode") { + return match; + } + } + // Masked ordinary form is (value, ptr, instruction_offset, mask). The + // dintlv no-mask form also has four positional arguments but carries + // dist_mode="dintlv" and was rejected above. + if (op->args.size() != 4U) { + return match; + } + const DataType value_dtype = op->args[0].dtype(); + const DataType elem = value_dtype.element_of(); + // Boolean (mask-register) values are not part of the b8/b16/b32 data + // storage contract; their backend admission stays unchanged. + if (elem.is_bool()) { + return match; + } + const int bits = elem.bits(); + if (bits != 8 && bits != 16 && bits != 32) { + return match; + } + const DataType mask_dtype = op->args[3].dtype(); + if (!mask_dtype.is_bool()) { + return match; + } + match.checked = true; + match.value_dtype = value_dtype; + match.elem_bytes = bits / 8; + // Record the immediate descriptor only for diagnostics and for the + // inherited byte offset (ResolveAddress consumes it). The instruction + // offset unit is resolved from the *resulting* pointer expression and is + // never defaulted to float32. + if (const auto *ptr_call = op->args[1].as()) { + if (ptr_call->op.same_as(builtin::tvm_access_ptr()) && + ptr_call->args.size() == 5U) { + match.has_descriptor = true; + match.descriptor_offset = ptr_call->args[2]; + if (auto dtype = DecodeAccessDType(ptr_call->args[0])) { + match.descriptor_dtype = dtype.value(); + } + } + } + match.instruction_offset = op->args[2]; + const int64_t lanes = value_dtype.lanes(); + const bool point_candidate = + lanes == 1 || lanes == 2 || lanes == 4 || lanes == 8; + const SeedInfo seed = ClassifyMaskSeed(op->args[3]); + if (point_candidate && seed.recognized && seed.full_active && + seed.mask_lanes == lanes) { + match.mode = StoreMode::kPointStores; + match.mode_name = "point"; + match.required_alignment = bits / 8; + } else { + match.mode = StoreMode::kOrdinaryMasked; + match.mode_name = "ordinary masked"; + match.required_alignment = kOrdinaryMaskedStoreAlignment; + } + return match; + } + + ByteAddress UnknownAddress(const std::string &why) const { + ByteAddress addr; + addr.kind = ByteAddress::Kind::kUnknown; + addr.error = why; + return addr; + } + + ByteAddress ResolveAddress(const PrimExpr &ptr, + std::unordered_set *visiting, + int nesting = 0) { + // Stack-safety guard only; plain alias length is walked iteratively and + // does not consume this budget. + constexpr int kMaxPointerExpressionNesting = 4096; + if (nesting > kMaxPointerExpressionNesting) { + return UnknownAddress( + "pointer expression nesting exceeds the stack-safety guard"); + } + PrimExpr current = ptr; + while (const auto *var = current.as()) { + if (!visiting->insert(var).second) { + return UnknownAddress("pointer alias cycle at `" + var->name_hint + + "`"); + } + auto def_it = pointer_defs_.find(var); + if (def_it == pointer_defs_.end()) { + break; + } + current = def_it->second; + } + if (const auto *var = current.as()) { + if (ub_roots_.count(var) != 0) { + ByteAddress addr; + addr.kind = ByteAddress::Kind::kUbRoot; + addr.root = var; + addr.base_bytes = IntImm(DataType::Int(64), 0); + addr.offset_bytes = IntImm(DataType::Int(64), 0); + return addr; + } + return UnknownAddress("pointer origin unavailable for `" + + var->name_hint + "`"); + } + if (const auto *call = current.as()) { + if (call->op.same_as(builtin::tvm_access_ptr())) { + TVM_FFI_ICHECK_EQ(call->args.size(), 5U) + << "tvm_access_ptr expects 5 arguments"; + auto dtype = DecodeAccessDType(call->args[0]); + if (!dtype.has_value()) { + return UnknownAddress("access pointer element type is unavailable"); + } + ByteAddress base = ResolveAddress(call->args[1], visiting, nesting + 1); + if (base.kind == ByteAddress::Kind::kUnknown) { + return base; + } + base.offset_bytes = analyzer_.Simplify( + base.offset_bytes + + Cast(DataType::Int(64), + call->args[2] * + IntImm(DataType::Int(64), dtype.value().bits() / 8))); + return base; + } + if (call->op.same_as(builtin::address_of())) { + if (call->args.size() != 1U) { + return UnknownAddress("address_of expects one argument"); + } + const auto *load = call->args[0].as(); + if (load == nullptr) { + return UnknownAddress("address_of expects a BufferLoad"); + } + if (load->indices.size() != 1U) { + return UnknownAddress( + "address_of with a non-flat index is unsupported"); + } + ByteAddress base = + ResolveAddress(load->buffer->data, visiting, nesting + 1); + if (base.kind == ByteAddress::Kind::kUnknown) { + return base; + } + base.offset_bytes = analyzer_.Simplify( + base.offset_bytes + + Cast(DataType::Int(64), + load->indices[0] * + IntImm(DataType::Int(64), load->dtype.bits() / 8))); + return base; + } + if (call->op.same_as(tirx::builtin::reinterpret()) && + call->args.size() == 1U && current.dtype().is_handle()) { + if (call->args[0].dtype().is_handle()) { + // Pointer type re-view: the byte address is unchanged. + return ResolveAddress(call->args[0], visiting, nesting + 1); + } + if (call->args[0].dtype().is_scalar()) { + ByteAddress addr; + addr.kind = ByteAddress::Kind::kIntegerBase; + addr.base_bytes = Cast(DataType::Int(64), call->args[0]); + addr.offset_bytes = IntImm(DataType::Int(64), 0); + return addr; + } + } + } + return UnknownAddress("unrecognized pointer expression"); + } + + Proof Prove(const ByteAddress &address, const PrimExpr &extra_bytes, + int required_alignment) { + Proof proof; + if (required_alignment <= 1) { + proof.outcome = Outcome::kAligned; + return proof; + } + if (address.kind == ByteAddress::Kind::kUnknown) { + proof.outcome = Outcome::kUnknown; + proof.error = address.error; + return proof; + } + PrimExpr base = address.base_bytes.defined() ? address.base_bytes + : IntImm(DataType::Int(64), 0); + PrimExpr total = analyzer_.Simplify(Cast(DataType::Int(64), base) + + address.offset_bytes + extra_bytes); + proof.byte_offset = total; + PrimExpr remainder = analyzer_.Simplify( + floormod(total, IntImm(DataType::Int(64), required_alignment))); + proof.remainder = remainder; + if (analyzer_.CanProve(remainder == 0)) { + proof.outcome = Outcome::kAligned; + } else if (analyzer_.CanProve(remainder != 0)) { + proof.outcome = Outcome::kMisaligned; + } else { + proof.outcome = Outcome::kUnknown; + proof.error = "cannot prove alignment for every execution"; + } + return proof; + } + + void ReportFailure(const StoreMatch &match, size_t store_index, + const ByteAddress &address, const Proof &proof) const { + std::ostringstream oss; + oss << "[VerifyPTOCodegenLegality] kernel=" << function_name_ << ", store=#" + << store_index << ": "; + if (proof.outcome == Outcome::kMisaligned) { + oss << "destination is misaligned: byte remainder=" << proof.remainder + << ", required alignment=" << match.required_alignment; + } else { + oss << "cannot prove destination alignment: " + << (proof.error.empty() ? "alignment is not provable" : proof.error) + << ", required alignment=" << match.required_alignment; + } + oss << "\n mode=" << match.mode_name + << " vstore, dtype=" << match.value_dtype + << ", element bytes=" << match.elem_bytes; + if (match.has_descriptor) { + oss << "\n descriptor offset=" << match.descriptor_offset << " (unit " + << match.descriptor_dtype << ")"; + } else { + oss << "\n descriptor offset="; + } + oss << ", instruction offset=" << match.instruction_offset << " (unit "; + if (match.pointer_dtype.has_value()) { + oss << match.pointer_dtype.value(); + } else { + oss << ""; + } + oss << ")"; + if (proof.byte_offset.defined()) { + oss << "\n normalized byte offset=" << proof.byte_offset; + } + if (address.root != nullptr) { + oss << "\n root=" << address.root->name_hint; + } else if (address.kind == ByteAddress::Kind::kIntegerBase) { + oss << "\n base byte address=" << address.base_bytes; + } else if (!address.error.empty()) { + oss << "\n address resolution: " << address.error; + } + if (match.mode == StoreMode::kPointStores) { + oss << "\n point-writeback explicit-mask vstore requires a destination " + "aligned to the stored element width (" + << match.required_alignment + << " bytes); check the starting byte address and offsets"; + } else { + oss << "\n ordinary explicit-mask vstore requires a provably 32-byte " + "aligned destination; consider a physical row stride whose byte " + "size is a multiple of 32 and check the starting " + "column/elem_offset"; + } + LOG(FATAL) << oss.str(); + } + + void CheckStore(const CallNode *op) { + StoreMatch match = MatchStore(op); + if (!match.checked) { + return; + } + const size_t store_index = next_store_index_++; + std::unordered_set visiting; + ByteAddress address = ResolveAddress(op->args[1], &visiting); + std::unordered_set dtype_visiting; + match.pointer_dtype = + ResolvePointerElementType(op->args[1], &dtype_visiting); + Proof proof; + if (match.required_alignment > 1) { + std::string dtype_error; + if (!match.pointer_dtype.has_value()) { + dtype_error = "resulting pointer element type is unavailable"; + } else { + const DataType ptr_dtype = match.pointer_dtype.value(); + const int ptr_bits = ptr_dtype.bits(); + if (ptr_bits < 8 || ptr_bits % 8 != 0) { + dtype_error = + "resulting pointer element type is not byte-addressable"; + } else { + const PrimExpr instruction_bytes = + Cast(DataType::Int(64), + op->args[2] * IntImm(DataType::Int(64), ptr_bits / 8)); + proof = Prove(address, instruction_bytes, match.required_alignment); + } + } + if (!dtype_error.empty()) { + proof.outcome = Outcome::kUnknown; + proof.error = dtype_error; + proof.byte_offset = PrimExpr(); + } + } else { + // A one-byte natural requirement holds for every byte address. + proof.outcome = Outcome::kAligned; + } + if (proof.outcome == Outcome::kAligned) { + return; + } + ReportFailure(match, store_index, address, proof); + } + + std::string function_name_; + size_t next_store_index_{0}; + std::unordered_set ub_roots_; + std::unordered_map pointer_defs_; + std::unordered_map value_defs_; + // Bind/Let definitions are immutable SSA values with unique VarNode + // identity, and every Var a definition references is already defined in an + // enclosing scope. A definition's classification therefore cannot change + // while that Var can still be referenced; once its lexical scope exits the + // VarNode cannot appear again. The cache is intentionally not part of + // AliasScope, and a Var used before its own definition is classified as a + // free parameter (non-SSA IR is not claimed as supported). + std::unordered_map var_origin_safe_; + size_t next_dag_symbol_{0}; +}; + +tvm::transform::Pass VerifyPTOCodegenLegality() { + auto pass_func = [](PrimFunc f, const IRModule &m, + const tvm::transform::PassContext &ctx) { + auto target = f->GetAttr(tvm::attr::kTarget); + if (!target.defined() || !TargetIsPTO(target.value())) { + return f; + } + std::string name = "?"; + if (auto symbol = f->GetAttr(tvm::attr::kGlobalSymbol)) { + name = symbol.value(); + } + PTOMaskedStoreAlignmentChecker checker(name); + checker.Run(f); + return f; + }; + return tvm::tirx::transform::CreatePrimFuncPass( + pass_func, 0, "tl.VerifyPTOCodegenLegality", {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tl.transform.VerifyPTOCodegenLegality", + VerifyPTOCodegenLegality); +} + +} // namespace transform +} // namespace tl +} // namespace tvm diff --git a/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/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/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/cast_plan_context_probe.cc b/testing/ascend/language/pto/cast_plan_context_probe.cc new file mode 100644 index 0000000000..d1020e3be3 --- /dev/null +++ b/testing/ascend/language/pto/cast_plan_context_probe.cc @@ -0,0 +1,148 @@ +/* Direct PlanValueCast context probe. + * + * Calls PlanValueCast directly with CastContext combinations that the + * Verify/Vectorize passes cannot produce end-to-end, and checks both the + * rejection reasons and the emitted vector/scalar plans. + * + * The two passes can only produce the context combinations that their + * classification yields (scalar lanes=1, VMI lanes=64/128/256). The planner + * guards for {kVectorVmi, 1} and {kScalar, 128} are therefore unreachable + * end-to-end; this local probe calls the planner directly so those contracts + * are tested without adding a product FFI. + * + * Build (from the repository root; the library lives in build/lib and the + * TVM DSOs must be linked explicitly): + * c++ -std=c++17 -O0 testing/ascend/language/pto/cast_plan_context_probe.cc \ + * -Isrc -I3rdparty/tvm/include -I3rdparty/tvm/src \ + * -I3rdparty/tvm/3rdparty/dlpack/include \ + * -I3rdparty/tvm/3rdparty/tvm-ffi/include \ + * -I3rdparty/tvm/3rdparty/tvm-ffi/3rdparty/dlpack/include \ + * -Lbuild/lib -ltilelang -ltvm_compiler -ltvm_ffi -ltvm_runtime \ + * -Wl,-rpath,$PWD/build/lib -o /tmp/cast_plan_probe + * LD_LIBRARY_PATH=$PWD/build/lib /tmp/cast_plan_probe + */ +#include +#include +#include + +#include "ascend/transform/parallel_to_pto_utils.h" +#include + +namespace { + +using tvm::DataType; +using tvm::tirx::Cast; +using tvm::tirx::Var; +using tvm::tl::pto::CastContext; +using tvm::tl::pto::CastValueKind; +using tvm::tl::pto::PlanValueCast; +using tvm::tl::pto::ScalarCastAction; +using tvm::tl::pto::ScalarCastPlan; +using tvm::tl::pto::VectorCastOp; +using tvm::tl::pto::VectorCastPlan; + +int failures = 0; + +void Check(bool condition, const std::string &what) { + if (!condition) { + std::cout << "FAIL " << what << "\n"; + ++failures; + } else { + std::cout << "PASS " << what << "\n"; + } +} + +const VectorCastPlan * +AsVector(const tvm::ffi::Optional &plan, + const std::string &what) { + if (!plan.has_value()) { + std::cout << "FAIL " << what << ": no plan\n"; + ++failures; + return nullptr; + } + const auto *vector = std::get_if(&plan.value()); + if (vector == nullptr) { + std::cout << "FAIL " << what << ": not a vector plan\n"; + ++failures; + } + return vector; +} + +const ScalarCastPlan * +AsScalar(const tvm::ffi::Optional &plan, + const std::string &what) { + if (!plan.has_value()) { + std::cout << "FAIL " << what << ": no plan\n"; + ++failures; + return nullptr; + } + const auto *scalar = std::get_if(&plan.value()); + if (scalar == nullptr) { + std::cout << "FAIL " << what << ": not a scalar plan\n"; + ++failures; + } + return scalar; +} + +} // namespace + +int main() { + Var x("x", DataType::Int(8)); + Cast widen(DataType::Int(16), x); + Cast cross(DataType::UInt(16), x); + Cast same_width(DataType::UInt(8), x); + + std::string why; + CastContext vector_lane1{CastValueKind::kVectorVmi, 1}; + auto plan = PlanValueCast(widen.get(), vector_lane1, &why); + Check(!plan.has_value() && + why.find("requires logical lanes 64/128/256") != std::string::npos, + "kVectorVmi lanes=1 is rejected as VMI, never scalar: " + why); + + why.clear(); + CastContext scalar_lanes{CastValueKind::kScalar, 128}; + plan = PlanValueCast(widen.get(), scalar_lanes, &why); + Check(!plan.has_value() && + why.find("must have exactly one logical lane") != std::string::npos, + "kScalar lanes=128 is rejected, never VMI: " + why); + + why.clear(); + CastContext vector64{CastValueKind::kVectorVmi, 64}; + plan = PlanValueCast(cross.get(), vector64, &why); + if (const VectorCastPlan *steps = AsVector(plan, "int8->uint16 VMI plan")) { + Check(steps->steps.size() == 2, "int8->uint16 has 2 steps"); + if (steps->steps.size() == 2) { + Check(steps->steps[0].op == VectorCastOp::kIntegerExtend && + steps->steps[0].from == DataType::Int(8) && + steps->steps[0].to == DataType::Int(16), + "step 1 extends with the source signedness to the target width"); + Check(steps->steps[1].op == VectorCastOp::kIntegerReinterpret && + steps->steps[1].from == DataType::Int(16) && + steps->steps[1].to == DataType::UInt(16), + "step 2 reinterprets at the same width"); + } + } + + why.clear(); + CastContext scalar1{CastValueKind::kScalar, 1}; + plan = PlanValueCast(cross.get(), scalar1, &why); + if (const ScalarCastPlan *scalar = + AsScalar(plan, "int8->uint16 scalar plan")) { + Check(scalar->action == ScalarCastAction::kRebuildCast, + "int8->uint16 scalar keeps a rebuilt ordinary Cast"); + Check(scalar->annotations.empty(), + "scalar annotations are validated/empty"); + } + + why.clear(); + CastContext vector128{CastValueKind::kVectorVmi, 128}; + plan = PlanValueCast(same_width.get(), vector128, &why); + if (const VectorCastPlan *steps = + AsVector(plan, "int8->uint8 VMI same-width plan")) { + Check(steps->steps.size() == 1 && + steps->steps[0].op == VectorCastOp::kIntegerReinterpret, + "same-width cross-sign is one vinterpret_cast"); + } + + return failures == 0 ? 0 : 1; +} diff --git a/testing/ascend/language/pto/test_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_feature_interaction_run.py b/testing/ascend/language/pto/test_pto_feature_interaction_run.py new file mode 100644 index 0000000000..90ed403ac5 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_feature_interaction_run.py @@ -0,0 +1,438 @@ +"""Cross-feature interaction controls: combined end-to-end device kernels. + +The dedicated per-feature suites cover each contract in isolation. This file +exercises three combinations of features in a single device run: + +* a handwritten VMI kernel loads from a nonzero UB buffer start, performs an + integer reduction, broadcasts the lane-1 result with ``vbrc(group=1)``, + stores it through an aligned masked store and copies the result back with + DMA. Automatic scheduling stays enabled (the generated MTE2->V / V->MTE3 + synchronization is asserted). Signed-negative and unsigned high-bit inputs + are compared exactly, with physical source padding and an initialized + output. + +* an auto-vectorized ``T.Parallel`` kernel reads from a nonzero source + offset, applies a cross-sign widening Cast (source-sign extend then target + reinterpret), stores through a masked aligned store and copies back with + DMA. Both a tail (E = VL-1) and a multi-chunk extent (E = 2*VL+17) are + covered for int8->uint16 and uint8->int32. The source storage is + physically padded, the destination is initialized with sentinels, and every + logical element plus the untouched tail is checked. The final alignment + checker must accept the store (compilation would fail otherwise). + +* SIMT, VMI and Parallel kernels are compiled in alternating orders in one + process. Each compilation uses fresh inputs and an independent numeric + reference, and the generated source must retain only the expected + per-kernel features, so storage or provenance state cannot leak between + compilations. Handwritten VMI is never mixed into a Parallel VF. + +Every test executes on the device and is skipped when no NPU runtime is +available. +""" + +import re +from contextlib import contextmanager + +import pytest +import torch +import torch_npu # noqa: F401,E402 (must import before tilelang/ptodsl) + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + +tirx = tvm.tirx +START = 131 +VL = 64 +REDUCE_LANES = 128 +_ELEMENT_BYTES = { + "int8": 1, + "uint8": 1, + "int16": 2, + "uint16": 2, + "int32": 4, + "uint32": 4, +} + + +def _require_npu_runtime(): + if not hasattr(torch, "npu") or not torch.npu.is_available(): + pytest.skip("torch.npu is unavailable in the current environment") + try: + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source) + except Exception as err: # noqa: BLE001 + pytest.skip(f"NPU runtime is unavailable in the current environment: {err}") + + +def _compile(func): + return tilelang.compile(func, target="pto", out_idx=-1) + + +def _run(kernel, *inputs): + device_inputs = [value.to("npu") for value in inputs] + output = kernel(*device_inputs) + torch.npu.synchronize() + return output.cpu() + + +def _physical_tail_lanes(lanes, element_bytes): + """Whole-vector load footprint: 256 bytes or one logical vector, whichever is larger.""" + return max(lanes, 256 // element_bytes) + + +def _padded_source_len(start, extent, lanes, element_bytes): + last_chunk = ((extent - 1) // lanes) * lanes + return start + last_chunk + _physical_tail_lanes(lanes, element_bytes) + + +@contextmanager +def _capture_final_module(): + captured = {} + original = ascend_transform.VerifyPTOCodegenLegality + + def factory(): + real = original() + + def run(mod): + captured["module"] = mod + return real(mod) + + return run + + ascend_transform.VerifyPTOCodegenLegality = factory + try: + yield captured + finally: + ascend_transform.VerifyPTOCodegenLegality = original + + +def _final_vload_descriptor_offsets(func): + """Lower to the pass boundary and return every VMI load descriptor offset.""" + target = normalize_pto_target("pto") + with _capture_final_module() as captured, target: + _lower(func, target=target) + module = captured["module"] + offsets = [] + + def visit(node): + if isinstance(node, tirx.Call) and hasattr(node.op, "name") and node.op.name == "tl.vmi.vload": + offsets.append(node.args[0].args[2]) + + for func in module.functions.values(): + if isinstance(func, tirx.PrimFunc): + tirx.stmt_functor.post_order_visit(func.body, visit) + assert offsets, "expected a lowered VMI load" + return offsets + + +def _constant_term(expr): + variables = [] + + def visit(node): + if isinstance(node, tirx.Var): + variables.append(node) + + tirx.stmt_functor.post_order_visit(expr, visit) + substitution = {var: tirx.IntImm(str(var.dtype), 0) for var in variables} + simplified = tvm.arith.Analyzer().simplify(tirx.stmt_functor.substitute(expr, substitution)) + assert isinstance(simplified, tirx.IntImm), simplified + return int(simplified.value) + + +# --------------------------------------------------------------------------- +# Group 1: nonzero-start VMI load, integer reduce, group-1 broadcast, +# aligned masked store, DMA; automatic scheduling enabled. +# --------------------------------------------------------------------------- +def _vmi_reduce_kernel(dtype, lanes=REDUCE_LANES, start=START): + storage = start + lanes + 64 + + # The whole padded source is filled by one aligned DMA; the VMI load then + # reads at the nonzero start, which the backend handles with + # vldas/vldus. A DMA directly into an unaligned UB region is not a legal + # MTE non-atomic access and is not what this control is about. + @T.prim_func + def main( + A: T.Tensor((storage,), dtype), + S: T.Tensor((lanes,), dtype), + C: T.Tensor((lanes,), dtype), + ): + with T.Kernel(1): + work = T.alloc_shared((storage,), dtype) + out = T.alloc_shared((lanes,), dtype) + T.copy(A, work) + T.copy(S, out) + with T.SimdVF(): + mask = T.vmi.create_mask(lanes, size=lanes) + x = T.vmi.vload(work[start], size=lanes) + reduced = T.vmi.vcmax(x, mask) + broadcast = T.vmi.vbrc(reduced, size=lanes, group=1) + T.vmi.vstore(broadcast, out[0], mask) + T.copy(out, C) + + return main + + +def _padded_reduce_input(values, torch_dtype): + padded = torch.zeros(START + REDUCE_LANES + 64, dtype=torch_dtype) + padded[START : START + REDUCE_LANES] = values + return padded + + +def _int32_reduce_input(round_index): + generator = torch.Generator().manual_seed(0xA101 + round_index) + values = torch.randint(-(2**31), 2**31 - 1, (REDUCE_LANES,), generator=generator).to(torch.int32) + values[0] = -(2**31) + values[1] = 2**31 - 1 + values[2] = -1 + values[3] = 0 + values[4] = 1 + return values + + +def _uint32_reduce_input(round_index): + generator = torch.Generator().manual_seed(0xA102 + round_index) + values = torch.randint(0, 2**31 - 1, (REDUCE_LANES,), generator=generator).to(torch.uint32) + values[0] = 2**32 - 1 + values[1] = 2**31 # INT32_MAX + 1, high bit set + values[2] = 2**31 - 1 + values[3] = 0 + return values + + +@pytest.mark.parametrize( + ("dtype", "torch_dtype", "sentinel", "make_input"), + [ + ("int32", torch.int32, -7, _int32_reduce_input), + ("uint32", torch.uint32, 7, _uint32_reduce_input), + ], + ids=["int32-signed-negative", "uint32-high-bit"], +) +def test_nonzero_start_integer_reduce_broadcast_store(dtype, torch_dtype, sentinel, make_input): + _require_npu_runtime() + kernel = _compile(_vmi_reduce_kernel(dtype)) + source = kernel.get_kernel_source() + load_line = next(line for line in source.splitlines() if "vmi.vload" in line) + assert "addptr" in load_line and str(START) in load_line, load_line + broadcast_line = next(line for line in source.splitlines() if "vmi.vbrc" in line) + assert "pto.cast(" not in broadcast_line, broadcast_line + assert "vmi.vcmax" in source, source + assert "vmi.create_mask" in source, source + assert 'set_flag("MTE2", "V"' in source, source + assert 'wait_flag("MTE2", "V"' in source, source + assert 'set_flag("V", "MTE3"' in source, source + assert 'wait_flag("V", "MTE3"' in source, source + + for round_index in range(2): + values = make_input(round_index) + sentinels = torch.full((REDUCE_LANES,), sentinel, dtype=torch_dtype) + output = _run(kernel, _padded_reduce_input(values, torch_dtype), sentinels) + expected = torch.full( + (REDUCE_LANES,), + max(int(value) for value in values.tolist()), + dtype=torch_dtype, + ) + torch.testing.assert_close(output, expected, rtol=0.0, atol=0.0) + + +# --------------------------------------------------------------------------- +# Group 2: automatic Parallel nonzero source access, cross-sign widening +# Cast (extend then reinterpret), aligned masked store, DMA; tail and +# multi-chunk extents with padded source and tail sentinels. +# --------------------------------------------------------------------------- +_CAST_CASES = [ + ("int8", "uint16", torch.int8, torch.uint16, "pto.si16", "pto.ui16", 30000), + ("uint8", "int32", torch.uint8, torch.int32, "pto.ui32", "pto.si32", -999999), +] +_EXTENTS = [VL - 1, 2 * VL + 17] + + +def _parallel_cast_kernel(src, dst, extent, lanes=VL, start=START, read_start=None): + # Continuous vloads cover whole physical vectors: the source storage + # covers 256 bytes or one logical VL, whichever is larger, after the + # last chunk base. ``read_start`` is used only by the lower-only + # descriptor control, so both kernels keep identical shapes. + source_len = _padded_source_len(start, extent, lanes, _ELEMENT_BYTES[src]) + out_len = extent + 2 * lanes + read = start if read_start is None else read_start + + @T.prim_func + def main( + A: T.Tensor((source_len,), src), + S: T.Tensor((out_len,), dst), + C: T.Tensor((out_len,), dst), + ): + with T.Kernel(1): + a = T.alloc_shared((source_len,), src) + c = T.alloc_shared((out_len,), dst) + T.copy(A, a) + T.copy(S, c) + with T.SimdVF(lanes=lanes): + for j in T.Parallel(extent): + c[j] = T.cast(a[read + j], dst) + T.copy(c, C) + + return main + + +def _source_values(src, torch_src, source_len): + if src == "int8": + values = torch.arange(source_len) * 37 % 251 - 125 + return values.to(torch_src) + values = torch.arange(source_len) * 53 % 256 + return values.to(torch_src) + + +@pytest.mark.parametrize( + ("src", "dst", "torch_src", "torch_dst", "intermediate", "target", "sentinel"), + _CAST_CASES, + ids=["int8_to_uint16", "uint8_to_int32"], +) +@pytest.mark.parametrize("extent", _EXTENTS, ids=["tail", "multi_chunk"]) +def test_parallel_nonzero_source_cross_sign_cast_store(src, dst, torch_src, torch_dst, intermediate, target, sentinel, extent): + _require_npu_runtime() + kernel = _compile(_parallel_cast_kernel(src, dst, extent)) + source = kernel.get_kernel_source() + load_line = next(line for line in source.splitlines() if "vmi.vload" in line) + assert "addptr" in load_line, load_line + # Exact final-descriptor check: lower the same-shape kernel once with the + # source read at 0 and once at START, then verify that the constant part of + # the final merged descriptor offset differs by exactly START elements. + actual_offsets = _final_vload_descriptor_offsets(_parallel_cast_kernel(src, dst, extent)) + control_offsets = _final_vload_descriptor_offsets(_parallel_cast_kernel(src, dst, extent, read_start=0)) + assert len(actual_offsets) == len(control_offsets) + for actual, control in zip(actual_offsets, control_offsets): + assert _constant_term(actual) - _constant_term(control) == START + cast_lines = [ + line.strip() + for line in source.splitlines() + if ("vmi.vcvt" in line or "vinterpret_cast" in line) and "vmi.vbrc(" not in line and "vmi.vci(" not in line + ] + assert len(cast_lines) == 2, cast_lines + assert "vmi.vcvt" in cast_lines[0], cast_lines + assert f"to_dtype={intermediate}" in cast_lines[0], cast_lines[0] + assert "vinterpret_cast" in cast_lines[1], cast_lines + assert f"to_dtype={target}" in cast_lines[1], cast_lines[1] + assert "vmi.create_mask" in source, source + store_lines = [line for line in source.splitlines() if "vmi.vstore" in line] + assert store_lines, source + for line in store_lines: + assert "mask" in line, line + + source_len = _padded_source_len(START, extent, VL, _ELEMENT_BYTES[src]) + out_len = extent + 2 * VL + values = _source_values(src, torch_src, source_len) + sentinels = torch.full((out_len,), sentinel, dtype=torch_dst) + output = _run(kernel, values, sentinels) + logical = values[START : START + extent].to(torch.int64) + if dst == "uint16": + expected = (logical & 0xFFFF).to(torch.uint16) + else: + expected = logical.to(torch.int32) + torch.testing.assert_close(output[:extent], expected, rtol=0.0, atol=0.0) + torch.testing.assert_close(output[extent:], sentinels[extent:], rtol=0.0, atol=0.0) + + +# --------------------------------------------------------------------------- +# Group 3: SIMT and VMI/Parallel compilation succession and interleaving; +# independent numeric checks for every compilation. +# --------------------------------------------------------------------------- +def _simt_copy_kernel(dtype="int32", lanes=8): + # The destination is seeded with a distinct sentinel, so a missing SIMT + # local store cannot leave the expected output in place. + @T.prim_func + def main( + A: T.Tensor((64,), dtype), + S: T.Tensor((64,), dtype), + C: T.Tensor((64,), dtype), + ): + with T.Kernel(1): + src = T.alloc_shared((64,), dtype) + dst = T.alloc_shared((64,), dtype) + T.copy(A, src) + T.copy(S, dst) + with T.SimtVF(threads=1): + local = T.alloc_local((64,), dtype) + local[T.Ramp(0, 1, lanes)] = src[T.Ramp(0, 1, lanes)] + dst[T.Ramp(0, 1, lanes)] = local[T.Ramp(0, 1, lanes)] + T.copy(dst, C) + + return main + + +def _compile_simt(round_index): + kernel = _compile(_simt_copy_kernel()) + source = kernel.get_kernel_source() + assert "vmi.vload" not in source, source + assert "vmi.vstore" not in source, source + # Same-width signless local backing plus explicit signed logical + # adaptation on the load from and store back to the shared buffer. + assert "pto.alloc_buffer((64,), pto.i32)" in source, source + assert "pto.alloc_buffer((64,), pto.si32)" not in source, source + assert re.search( + r"pto\.cast\(pto\.load\(local, 0, contiguous=8\), pto\.si32\)", + source, + ), source + assert re.search(r"pto\.cast\(pto\.load\(.*contiguous=8\), pto\.i32\)", source), source + generator = torch.Generator().manual_seed(0x51A1 + round_index) + values = torch.randint(-1000, 1000, (64,), generator=generator).to(torch.int32) + sentinels = torch.full((64,), -777777, dtype=torch.int32) + output = _run(kernel, values, sentinels) + torch.testing.assert_close(output[:8], values[:8], rtol=0.0, atol=0.0) + torch.testing.assert_close(output[8:], sentinels[8:], rtol=0.0, atol=0.0) + + +def _compile_vmi(round_index): + kernel = _compile(_vmi_reduce_kernel("int32")) + source = kernel.get_kernel_source() + assert "vmi.vload" in source and "vmi.vstore" in source, source + assert "vmi.vbrc" in source, source + broadcast_line = next(line for line in source.splitlines() if "vmi.vbrc" in line) + assert "pto.cast(" not in broadcast_line, broadcast_line + values = _int32_reduce_input(round_index) + sentinels = torch.full((REDUCE_LANES,), -7, dtype=torch.int32) + output = _run(kernel, _padded_reduce_input(values, torch.int32), sentinels) + expected = torch.full( + (REDUCE_LANES,), + max(int(value) for value in values.tolist()), + dtype=torch.int32, + ) + torch.testing.assert_close(output, expected, rtol=0.0, atol=0.0) + + +def _compile_parallel(round_index): + extent = VL - 1 + kernel = _compile(_parallel_cast_kernel("int8", "uint16", extent)) + source = kernel.get_kernel_source() + assert "vinterpret_cast" in source and "vmi.vload" in source, source + assert "vmi.vstore" in source, source + source_len = _padded_source_len(START, extent, VL, _ELEMENT_BYTES["int8"]) + out_len = extent + 2 * VL + values = _source_values("int8", torch.int8, source_len) + sentinels = torch.full((out_len,), 30000, dtype=torch.uint16) + output = _run(kernel, values, sentinels) + expected = (values[START : START + extent].to(torch.int64) & 0xFFFF).to(torch.uint16) + torch.testing.assert_close(output[:extent], expected, rtol=0.0, atol=0.0) + torch.testing.assert_close(output[extent:], sentinels[extent:], rtol=0.0, atol=0.0) + + +def test_simt_vmi_parallel_compile_sequences_do_not_leak_state(): + _require_npu_runtime() + sequences = [ + ("simt", "vmi", "parallel"), + ("parallel", "vmi", "simt"), + ("vmi", "simt", "parallel"), + ] + builders = { + "simt": _compile_simt, + "vmi": _compile_vmi, + "parallel": _compile_parallel, + } + for round_index, sequence in enumerate(sequences): + for name in sequence: + builders[name](round_index) diff --git a/testing/ascend/language/pto/test_pto_masked_store_alignment.py b/testing/ascend/language/pto/test_pto_masked_store_alignment.py new file mode 100644 index 0000000000..97fa28b24f --- /dev/null +++ b/testing/ascend/language/pto/test_pto_masked_store_alignment.py @@ -0,0 +1,883 @@ +"""Explicit-mask PTO VMI store alignment verification. + +The read-only ``VerifyPTOCodegenLegality`` PrimFuncPass runs at the end of +the PTO pipeline, after ``LowerDeviceKernelLaunch``, the buffer-version +rewrites and UB merging. It consumes the final device TIR, resolves the real +destination byte address and proves the alignment required by PTO store +operations: + +* an ordinary default/continuous explicit-mask ``tl.vmi.vstore`` requires a + provably 32-byte-aligned destination; +* the implicit per-point writeback (``1PT_B8/B16/B32``) selected when the + value has 1/2/4/8 logical lanes and the store mask is a recognized + full-active seed with the same logical lane count requires only the stored + element's natural alignment (1/2/4 bytes). + +Covered here: + +* matcher boundaries: default and continuous full-mask stores pass; stores + with no mask, a ``dintlv``/``group`` distribution annotation or a scatter + op stay untouched; the diagnostics distinguish a provably misaligned + destination from one whose alignment cannot be proven; +* the instruction offset's actual resulting-pointer element unit: inherited + offsets stay in bytes and are counted exactly once across typed alias + layers and merged UB arena descriptors; +* mask-seed precision: a recognized full-active seed with a matching lane + count lets a 1/2/4/8-lane store use point writeback with only natural + element alignment; partial, dynamic and arbitrary all-true arithmetic + seeds stay ordinary; +* scoped Analyzer facts: loop ranges and non-zero loop mins; +* the final-pipeline integration: the pass sees the final merged UB arena, + descriptor and version-slot addresses; compact ``(3,131)`` FP16/FP32 and + flat/column-offset stores fail before codegen; padded, row-0-only, tail + and dynamic-offset controls pass; the point-writeback reductions, the + buffer-version layout and a non-32B-aligned FP16/FP32 load with a + physically padded source feeding an aligned store compile and run on the + device with numeric checks; dynamic-mask and single-active stores are + rejected. + +Every accepted direct-pass case asserts the returned IR is structurally +identical to its input. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager + +import pytest +import torch +import torch_npu # noqa: F401,E402 (must import before tilelang/ptodsl) + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + +tirx = tvm.tirx +_PASS = tvm.ffi.get_global_func("tl.transform.VerifyPTOCodegenLegality")() +_PTO_TARGET = normalize_pto_target("pto") +_MARKER = "[VerifyPTOCodegenLegality]" +_INTERNAL_ERROR = tvm.error.InternalError +_OP_VSTORE = tvm.ir.Op.get("tl.vmi.vstore") +_OP_VSCATTER = tvm.ir.Op.get("tl.vmi.vscatter") +_OP_CREATE_MASK = tvm.ir.Op.get("tl.vmi.create_mask") + + +def _require_npu_runtime(): + if not hasattr(torch, "npu") or not torch.npu.is_available(): + pytest.skip("torch.npu is unavailable in the current environment") + try: + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source) + except Exception as err: # noqa: BLE001 + pytest.skip(f"NPU runtime is unavailable in the current environment: {err}") + + +# --------------------------------------------------------------------------- +# Raw TIR builders (direct read-only pass tests) +# --------------------------------------------------------------------------- +def _imm(dtype, value): + return tirx.IntImm(dtype, value) + + +def _float_value(dtype, lanes): + base = tirx.FloatImm(dtype, 1.0) + return base if lanes == 1 else tirx.Broadcast(base, lanes) + + +def _int_value(dtype, lanes): + base = tirx.IntImm(dtype, 1) + return base if lanes == 1 else tirx.Broadcast(base, lanes) + + +def _value(dtype, lanes): + if str(dtype).startswith(("float", "bfloat")): + return _float_value(dtype, lanes) + return _int_value(dtype, lanes) + + +def _seed(lanes, active=None): + active = lanes if active is None else active + return tirx.Call( + f"boolx{lanes}", + _OP_CREATE_MASK, + [active if not isinstance(active, int) else _imm("int32", active)], + annotations={"size": _imm("int32", lanes)}, + ) + + +def _store(value, ptr, offset, mask, annotations=None): + return tirx.Evaluate(tirx.Call(value.dtype, _OP_VSTORE, [value, ptr, offset, mask], annotations=annotations)) + + +def _root(name="root", dtype="uint8", size=8192, scope="shared.dyn"): + var = tirx.Var(name, tvm.ir.PointerType(tvm.ir.PrimType(dtype), scope)) + buffer = tirx.decl_buffer((size,), dtype, data=var, scope=scope) + return var, tirx.AllocBuffer(buffer) + + +def _typed_int_base(name, dtype, base, scope="shared.dyn"): + var = tirx.Var(name, tvm.ir.PointerType(tvm.ir.PrimType(dtype), scope)) + bind = tirx.Bind(var, tirx.reinterpret("handle", tirx.IntImm("int64", base))) + return var, bind + + +def _access(dtype, data, offset, extent=64): + return tirx.tvm_access_ptr( + tirx.type_annotation(dtype), + data, + offset if not isinstance(offset, int) else _imm("int32", offset), + _imm("int32", extent), + _imm("int32", 2), + ) + + +def _func(body, params=(), target=None): + return tirx.PrimFunc(params, body).with_attr("target", target if target is not None else _PTO_TARGET) + + +def _run(func): + return _PASS(tvm.IRModule({"main_kernel": func})) + + +def _accept(func): + out = _run(func) + assert tvm.ir.structural_equal(out["main_kernel"], func), "the alignment pass must return the input IR unchanged" + return out + + +def _reject(func): + with pytest.raises(_INTERNAL_ERROR) as excinfo: + _run(func) + text = str(excinfo.value) + assert _MARKER in text, text + return text + + +def _compile_ptoas_artifacts(source): + """Compile PTODSL to PTO, VPTO and LLVM artifacts.""" + import importlib.util + import re + import subprocess + import tempfile + import uuid + + decorator = re.search(r'@pto\.jit\(name="([^"]+)"', source) + name = decorator.group(1) if decorator is not None else "main" + work = tempfile.mkdtemp(prefix="masked-store-alignment-", dir=os.environ.get("TMPDIR") or None) + source_path = os.path.join(work, "kernel.ptodsl.py") + pto_path = os.path.join(work, "kernel.pto") + vpto_path = os.path.join(work, "kernel.vpto.mlir") + llvm_path = os.path.join(work, "kernel.ll") + with open(source_path, "w", encoding="utf-8") as file: + file.write(source) + spec = importlib.util.spec_from_file_location("masked_store_alignment_" + uuid.uuid4().hex, source_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + pto_text = getattr(module, name).compile().mlir_text() + with open(pto_path, "w", encoding="utf-8") as file: + file.write(pto_text) + vpto_result = subprocess.run( + [ + "ptoas", + "--pto-arch=a5", + "--pto-backend=vpto", + "--emit-vpto", + pto_path, + "-o", + vpto_path, + ], + text=True, + capture_output=True, + ) + llvm_result = subprocess.run( + [ + "ptoas", + "--pto-arch=a5", + "--pto-backend=vpto", + "--emit-vpto-llvm-ir", + pto_path, + "-o", + llvm_path, + ], + text=True, + capture_output=True, + ) + + def read(path): + if not os.path.exists(path): + return "" + with open(path, encoding="utf-8") as file: + return file.read() + + stderr = "\n".join(error for error in (vpto_result.stderr, llvm_result.stderr) if error) + return ( + vpto_result.returncode or llvm_result.returncode, + stderr, + pto_text, + read(vpto_path), + read(llvm_path), + ) + + +def _ordinary_case(dtype="float32", lanes=64, offset=0): + """Ordinary masked store through a UB root; returns the PrimFunc.""" + root, alloc = _root(dtype="uint8") + value = _value(dtype, lanes) + mask = _seed(lanes) + ptr = _access(dtype, root, offset) + return _func(tirx.SeqStmt([alloc, _store(value, ptr, 0, mask)])) + + +# --------------------------------------------------------------------------- +# Matcher boundaries +# --------------------------------------------------------------------------- +class TestMatcher: + def test_default_and_continuous_full_mask_pass(self): + _accept(_ordinary_case(offset=8)) + root, alloc = _root() + value = _value("float32", 64) + func = _func( + tirx.SeqStmt( + [ + alloc, + _store( + value, + _access("float32", root, 8), + _imm("int32", 0), + _seed(64), + annotations={"dist_mode": tirx.StringImm("continuous")}, + ), + ] + ) + ) + _accept(func) + + def test_ordinary_store_misaligned_and_unknown_reasons_differ(self): + misaligned = _reject(_ordinary_case(offset=7)) + assert "destination is misaligned" in misaligned + assert "byte remainder" in misaligned + assert "required alignment=32" in misaligned + root, alloc = _root() + unknown_offset = tirx.Var("dynamic_offset", "int32") + func = _func( + tirx.SeqStmt( + [ + alloc, + _store( + _value("float32", 64), + _access("float32", root, unknown_offset), + _imm("int32", 0), + _seed(64), + ), + ] + ), + params=[unknown_offset], + ) + unknown = _reject(func) + assert "cannot prove destination alignment" in unknown + assert "destination is misaligned" not in unknown + assert "cannot prove alignment for every execution" in unknown + + @pytest.mark.parametrize("case", ["no_mask", "dintlv", "group", "vscatter"]) + def test_unmatched_forms_are_untouched(self, case): + root, alloc = _root() + ptr = _access("float32", root, 1) + value = _value("float32", 64) + mask = _seed(64) + statements = [alloc] + if case == "no_mask": + statements.append(tirx.Evaluate(tirx.Call(value.dtype, _OP_VSTORE, [value, ptr, _imm("int32", 1)]))) + elif case == "dintlv": + statements.append( + _store( + value, + ptr, + _imm("int32", 1), + mask, + annotations={"dist_mode": tirx.StringImm("dintlv")}, + ) + ) + elif case == "group": + statements.append( + _store( + value, + ptr, + _imm("int32", 1), + mask, + annotations={"group": _imm("int32", 2)}, + ) + ) + else: # vscatter + statements.append( + tirx.Evaluate( + tirx.Call( + value.dtype, + _OP_VSCATTER, + [value, ptr, _imm("int32", 1), _imm("int32", 0), mask], + ) + ) + ) + _accept(_func(tirx.SeqStmt(statements))) + + +# --------------------------------------------------------------------------- +# Instruction-offset element unit and address resolution +# --------------------------------------------------------------------------- +class TestInstructionUnit: + def test_cross_type_inherited_offset_counted_once(self): + root, alloc = _root() + base = tirx.Var("base", "handle") + alias = tirx.Var("alias", "handle") + prefix = [alloc, tirx.Bind(base, _access("uint8", root, 32)), tirx.Bind(alias, base)] + value = _value("float32", 64) + # 32 + 6*4 + 2*4 = 64 bytes: aligned only when every layer is counted + # exactly once with its own element unit. + aligned = _func( + tirx.SeqStmt( + prefix + + [ + _store( + value, + _access("float32", alias, 6), + _imm("int32", 2), + _seed(64), + ) + ] + ) + ) + _accept(aligned) + misaligned = _func( + tirx.SeqStmt( + prefix + + [ + _store( + value, + _access("float32", alias, 6), + _imm("int32", 0), + _seed(64), + ) + ] + ) + ) + text = _reject(misaligned) + assert "byte remainder=T.int64(24)" in text + + def test_merged_arena_offset_counted_once(self): + # A descriptor offset of 400 f32 elements = 1600 bytes is 32B aligned; + # 401 elements = 1604 bytes leaves a remainder of 4. + root, alloc = _root() + _accept( + _func( + tirx.SeqStmt( + [ + alloc, + _store( + _value("float32", 64), + _access("float32", root, 400), + _imm("int32", 0), + _seed(64), + ), + ] + ) + ) + ) + text = _reject( + _func( + tirx.SeqStmt( + [ + alloc, + _store( + _value("float32", 64), + _access("float32", root, 401), + _imm("int32", 0), + _seed(64), + ), + ] + ) + ) + ) + assert "byte remainder=T.int64(4)" in text + + +# --------------------------------------------------------------------------- +# Mask seeds and 1PT classification +# --------------------------------------------------------------------------- +class TestMaskSeed: + @pytest.mark.parametrize( + ("dtype", "natural"), + [("uint8", 1), ("float16", 2), ("float32", 4)], + ids=["b8", "b16", "b32"], + ) + def test_point_writeback_needs_only_natural_alignment(self, dtype, natural): + ptr, bind = _typed_int_base("mt", dtype, natural) + _accept(_func(tirx.SeqStmt([bind, _store(_value(dtype, 1), ptr, _imm("int32", 0), _seed(1))]))) + if natural > 1: + ptr, bind = _typed_int_base("mt", dtype, natural - 1) + _reject(_func(tirx.SeqStmt([bind, _store(_value(dtype, 1), ptr, _imm("int32", 0), _seed(1))]))) + + @pytest.mark.parametrize( + "mask_factory", + [ + lambda: _seed(4, active=2), + lambda: _seed(1, active=tirx.Var("active", "int32")), + lambda: tirx.Or(_seed(1), _seed(1)), + ], + ids=["partial", "dynamic_active", "all_true_arith"], + ) + def test_non_full_seeds_stay_ordinary(self, mask_factory): + root, alloc = _root() + mask = mask_factory() + lanes = mask.dtype.lanes + value = _value("float32", lanes) + text = _reject( + _func( + tirx.SeqStmt( + [ + alloc, + _store(value, _access("float32", root, 1), _imm("int32", 0), mask), + ] + ) + ) + ) + assert "ordinary masked vstore" in text + _accept( + _func( + tirx.SeqStmt( + [ + alloc, + _store(value, _access("float32", root, 8), _imm("int32", 0), mask), + ] + ) + ) + ) + + +# --------------------------------------------------------------------------- +# Scoped analysis facts +# --------------------------------------------------------------------------- +def _loop_with_offset(offset_builder): + root, alloc = _root() + loop_var = tirx.Var("i", "int32") + body = _store( + _value("float32", 64), + _access("float32", root, offset_builder(loop_var)), + _imm("int32", 0), + _seed(64), + ) + loop = tirx.For( + loop_var, + _imm("int32", 0), + _imm("int32", 4), + tirx.ForKind.SERIAL, + body, + ) + return _func(tirx.SeqStmt([alloc, loop])) + + +class TestAnalysisScope: + def test_loop_range_facts(self): + _accept(_loop_with_offset(lambda i: i * _imm("int32", 8))) + _reject(_loop_with_offset(lambda i: i)) + + def test_nonzero_loop_min_with_step(self): + i = tirx.Var("i", "int32") + offset = (i - _imm("int32", 4)) * _imm("int32", 2) + root, alloc = _root() + body = _store( + _value("float32", 64), + _access("float32", root, offset), + _imm("int32", 0), + _seed(64), + ) + loop = tirx.For( + i, + _imm("int32", 4), + _imm("int32", 8), + tirx.ForKind.SERIAL, + body, + step=_imm("int32", 4), + ) + _accept(_func(tirx.SeqStmt([alloc, loop]))) + + +# --------------------------------------------------------------------------- +# Final pipeline integration +# --------------------------------------------------------------------------- +@contextmanager +def _pass_boundary_capture(): + captured = {} + original = ascend_transform.VerifyPTOCodegenLegality + + def factory(): + real = original() + + def run(mod): + captured["module"] = mod + return real(mod) + + return run + + ascend_transform.VerifyPTOCodegenLegality = factory + try: + yield captured + finally: + ascend_transform.VerifyPTOCodegenLegality = original + + +def _lower_capture(func): + target = normalize_pto_target("pto") + with _pass_boundary_capture() as captured, target: + _lower(func, target=target) + return captured["module"] + + +def _lower_expect_error(func): + with pytest.raises(Exception) as excinfo: # noqa: B017 - kernel errors vary + tilelang.lower(func, target="pto") + text = str(excinfo.value) + assert _MARKER in text, text + assert "ptoas" not in text.lower() + return text + + +def _parallel_kernel(dtype, m, n, lanes=256, row=None): + if row is None: + + @T.prim_func + def main(A: T.Tensor((m, n), dtype), C: T.Tensor((m, n), dtype)): + with T.Kernel(1): + a = T.alloc_shared((m, n), dtype) + c = T.alloc_shared((m, n), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(m, n): + c[i, j] = a[i, j] + T.cast(3, dtype) + T.copy(c, C) + + return main + + @T.prim_func + def main(A: T.Tensor((m, n), dtype), C: T.Tensor((m, n), dtype)): + with T.Kernel(1): + a = T.alloc_shared((m, n), dtype) + c = T.alloc_shared((m, n), dtype) + T.copy(A[row, :], a[row, :]) + with T.SimdVF(lanes=lanes): + for j in T.Parallel(n): + c[row, j] = a[row, j] + T.cast(3, dtype) + T.copy(c[row, :], C[0, :]) + + return main + + +def _tail_kernel(dtype="float32", n=131, lanes=64): + @T.prim_func + def main(A: T.Tensor((n,), dtype), C: T.Tensor((n,), dtype)): + with T.Kernel(1): + a = T.alloc_shared((n,), dtype) + c = T.alloc_shared((n,), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for j in T.Parallel(n): + c[j] = a[j] + T.cast(3, dtype) + T.copy(c, C) + + return main + + +def _flat_kernel(m=3, n=131, dtype="float32", lanes=256): + @T.prim_func + def main(A: T.Tensor((m * n,), dtype), C: T.Tensor((m * n,), dtype)): + with T.Kernel(1): + a = T.alloc_shared((m * n,), dtype) + c = T.alloc_shared((m * n,), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i, j in T.Parallel(m, n): + c[i * n + j] = a[i * n + j] + T.cast(3, dtype) + T.copy(c, C) + + return main + + +def _column_offset_kernel(n=256, dtype="float32", lanes=256): + @T.prim_func + def main(A: T.Tensor((n,), dtype), C: T.Tensor((n + 1,), dtype)): + with T.Kernel(1): + a = T.alloc_shared((n,), dtype) + c = T.alloc_shared((n + 1,), dtype) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for j in T.Parallel(n): + c[j + 1] = a[j] + T.cast(3, dtype) + T.copy(c[1 : n + 1], C[1 : n + 1]) + + return main + + +_REDUCE_OFFSET = 4100 +_REDUCE_WINDOW = (_REDUCE_OFFSET // 64) * 64 +_REDUCE_INDEX = _REDUCE_OFFSET - _REDUCE_WINDOW + + +_SENTINEL_LANES = 128 + + +def _manual_reduce_kernel(group=None, size=64, dynamic=False): + out_lanes = 1 if group is None else group + window_start = _REDUCE_WINDOW + window_end = window_start + _SENTINEL_LANES + + if dynamic: + + @T.prim_func + def main( + A: T.Tensor((size,), "float32"), + S: T.Tensor((_SENTINEL_LANES,), "float32"), + active: T.int32, + C: T.Tensor((_SENTINEL_LANES,), "float32"), + ): + with T.Kernel(1): + work = T.alloc_shared((size,), "float32") + out = T.alloc_shared((8192,), "float32") + T.copy(A, work) + T.copy(S, out[window_start:window_end]) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + x = T.vmi.vload(work[0], size=size) + r = T.vmi.vcadd(x, mask, reassoc=True) + slot = T.alloc_var("int32") + slot = active + store_mask = T.vmi.create_mask(slot, size=out_lanes) + T.vmi.vstore(r, out[_REDUCE_OFFSET], store_mask) + T.copy(out[window_start:window_end], C) + + return main + + @T.prim_func + def main( + A: T.Tensor((size,), "float32"), + S: T.Tensor((_SENTINEL_LANES,), "float32"), + C: T.Tensor((_SENTINEL_LANES,), "float32"), + ): + with T.Kernel(1): + work = T.alloc_shared((size,), "float32") + out = T.alloc_shared((8192,), "float32") + T.copy(A, work) + # Independently initialized surrounding sentinels: a swapped, + # duplicated or out-of-footprint write inside this window changes + # a sentinel slot and fails the comparison. + T.copy(S, out[window_start:window_end]) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + x = T.vmi.vload(work[0], size=size) + if group is None: + r = T.vmi.vcadd(x, mask, reassoc=True) + else: + r = T.vmi.vcadd(x, mask, group=group, reassoc=True) + store_mask = T.vmi.create_mask(out_lanes, size=out_lanes) + T.vmi.vstore(r, out[_REDUCE_OFFSET], store_mask) + T.copy(out[window_start:window_end], C) + + return main + + +def _manual_single_active_ordinary(lanes=64): + @T.prim_func + def main(A: T.Tensor((lanes,), "float32"), C: T.Tensor((lanes,), "float32")): + with T.Kernel(1): + work = T.alloc_shared((lanes,), "float32") + out = T.alloc_shared((8192,), "float32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(1, size=lanes) + x = T.vmi.vload(work[0], size=lanes) + T.vmi.vstore(x, out[_REDUCE_OFFSET], mask) + T.copy( + out[_REDUCE_WINDOW : _REDUCE_WINDOW + lanes], + C, + ) + + return main + + +def _serial_version_kernel(): + @T.prim_func + def main(A: T.Tensor((512,), "float32"), C: T.Tensor((512,), "float32")): + with T.Kernel(1): + a = T.alloc_shared((128,), "float32") + c = T.alloc_shared((128,), "float32") + T.annotate_buffer_versions({a: 2, c: 2}) + for k in T.serial(4): + T.copy(A[k * 128 : (k + 1) * 128], a) + with T.SimdVF(lanes=128): + for j in T.Parallel(128): + c[j] = a[j] + T.cast(3, "float32") + T.copy(c, C[k * 128 : (k + 1) * 128]) + + return main + + +def _non32b_load_kernel(dtype): + source_len = 131 + 256 + + @T.prim_func + def main(S: T.Tensor((source_len,), dtype), C: T.Tensor((256,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((source_len + 128,), dtype) + out = T.alloc_shared((256,), dtype) + T.copy(S, work[0:source_len]) + with T.SimdVF(lanes=256): + mask = T.vmi.create_mask(256, size=256) + x = T.vmi.vload(work[131], size=256) + T.vmi.vstore(x, out[0], mask) + T.copy(out[0:256], C) + + return main + + +def _dynamic_offset_kernel(mult): + @T.prim_func + def main( + A: T.Tensor((64,), "float32"), + k: T.int32, + C: T.Tensor((64,), "float32"), + ): + with T.Kernel(1): + work = T.alloc_shared((1024,), "float32") + T.copy(A, work[0:64]) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + x = T.vmi.vload(work[0], size=64) + T.vmi.vstore(x, work[k * mult], mask) + T.copy(work[0:64], C) + + return main + + +class TestFinalPipeline: + @pytest.mark.parametrize("dtype", ["float32", "float16"]) + def test_compact_three_by_131_fails_before_codegen(self, dtype): + text = _lower_expect_error(_parallel_kernel(dtype, 3, 131)) + assert "ordinary masked vstore" in text + assert "required alignment=32" in text + assert "root=" in text + assert "normalized byte offset=" in text + + def test_flat_same_address_fails(self): + text = _lower_expect_error(_flat_kernel()) + assert "cannot prove destination alignment" in text + + def test_nonzero_column_offset_fails(self): + text = _lower_expect_error(_column_offset_kernel()) + assert "destination is misaligned" in text + + @pytest.mark.parametrize( + ("name", "kernel"), + [ + ("padded", _parallel_kernel("float32", 3, 136)), + ("row0_only", _parallel_kernel("float32", 3, 131, row=0)), + ("single_row_tail_vl64", _tail_kernel()), + ], + ) + def test_aligned_controls_full_compile(self, name, kernel): + compiled = tilelang.compile(kernel, target="pto", out_idx=-1) + source = compiled.get_kernel_source() + assert "vmi.vstore" in source, name + + def test_dynamic_offset_controls(self): + aligned = _dynamic_offset_kernel(8) + module = _lower_capture(aligned) + assert module is not None + tilelang.compile(aligned, target="pto", out_idx=-1) + text = _lower_expect_error(_dynamic_offset_kernel(1)) + assert "cannot prove destination alignment" in text + + @pytest.mark.parametrize( + ("group", "size", "lanes"), + [(None, 64, 1), (2, 128, 2), (4, 256, 4), (8, 256, 8)], + ) + def test_point_writeback_compiles_and_runs(self, group, size, lanes): + _require_npu_runtime() + kernel = tilelang.compile( + _manual_reduce_kernel(group=group, size=size), + target="pto", + out_idx=-1, + ) + source = kernel.get_kernel_source() + assert "vmi.vstore" in source + values = torch.arange(size, dtype=torch.float32) + sentinels = torch.arange(_SENTINEL_LANES, dtype=torch.float32) * -3.0 - 500.0 + output = kernel(values.to("npu"), sentinels.to("npu")) + torch.npu.synchronize() + got = output.cpu() + # The group reduction produces one result per consecutive + # size/lanes-element chunk. Every other slot in the initialized + # 128-element window must retain its sentinel. Swapped, duplicated or + # out-of-footprint writes must fail even when the total sum matches. + chunk = size // lanes + expected_slots = [float(values[index * chunk : (index + 1) * chunk].sum()) for index in range(lanes)] + expected = sentinels.clone() + expected[_REDUCE_INDEX : _REDUCE_INDEX + lanes] = torch.tensor(expected_slots, dtype=torch.float32) + torch.testing.assert_close(got, expected, rtol=0.0, atol=1e-3) + + def test_actual_version_layout_address_is_checked_and_runs(self): + _require_npu_runtime() + module = _lower_capture(_serial_version_kernel()) + stores = [] + + def visit(node): + if isinstance(node, tirx.Call) and hasattr(node.op, "name") and node.op.name == "tl.vmi.vstore": + stores.append(node) + + tirx.stmt_functor.post_order_visit(module["main_kernel"].body, visit) + assert stores, "expected a versioned VMI store" + for store in stores: + descriptor_offset = store.args[1].args[2] + version_terms = [] + + def offset_visit(node, terms=version_terms): + if isinstance(node, tirx.FloorMod): + terms.append(node) + + tirx.stmt_functor.post_order_visit(descriptor_offset, offset_visit) + assert version_terms, f"expected the final descriptor to contain the version-slot index, got {descriptor_offset}" + kernel = tilelang.compile(_serial_version_kernel(), target="pto", out_idx=-1) + values = torch.arange(512, dtype=torch.float32) + output = kernel(values.to("npu")) + torch.npu.synchronize() + torch.testing.assert_close(output.cpu(), values + 3.0, rtol=0.0, atol=0.0) + + @pytest.mark.parametrize("dtype", ["float16", "float32"]) + def test_non32b_aligned_load_with_padded_source_runs(self, dtype): + _require_npu_runtime() + kernel = tilelang.compile(_non32b_load_kernel(dtype), target="pto", out_idx=-1) + source = kernel.get_kernel_source() + load_line = next(line for line in source.splitlines() if "vmi.vload" in line) + assert "131" in load_line and "addptr" in load_line, load_line + returncode, stderr, pto_text, vpto_text, llvm_text = _compile_ptoas_artifacts(source) + assert returncode == 0, stderr + assert "pto.vldas" in vpto_text, vpto_text + assert "pto.vldus" in vpto_text, vpto_text + assert "131" in pto_text, pto_text + assert "llvm.hivm.vldas" in llvm_text, llvm_text + assert "llvm.hivm.vldus" in llvm_text, llvm_text + values = torch.arange(387, dtype=torch.float32).to(getattr(torch, dtype)) + output = kernel(values.to("npu")) + torch.npu.synchronize() + expected = values[131:387].float().cpu() + torch.testing.assert_close(output.cpu().float(), expected, rtol=0.0, atol=0.0) + + def test_dynamic_and_single_active_stores_reject(self): + text = _lower_expect_error(_manual_reduce_kernel(dynamic=True)) + assert "ordinary masked vstore" in text + assert "required alignment=32" in text + text = _lower_expect_error(_manual_single_active_ordinary()) + assert "ordinary masked vstore" in text + assert "required alignment=32" in text diff --git a/testing/ascend/language/pto/test_pto_parallel_cast_matrix.py b/testing/ascend/language/pto/test_pto_parallel_cast_matrix.py new file mode 100644 index 0000000000..38e8d8cad6 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_parallel_cast_matrix.py @@ -0,0 +1,702 @@ +"""Parallel value-Cast planning: action structure and attribute rules. + +Hardware independent. Each case lowers the real TileLang PTO pipeline with +``tilelang.engine.lower.lower`` (VerifyParallelToPTO + VectorizeParallelToPTO ++ PTODSL codegen; no PTOAS/device) and inspects the generated source and the +module at the Verify boundary. + +Covered contract: + +* ``PlanValueCast(cast, CastContext{value_kind, logical_lanes}, reason)`` + dispatches scalar vs VMI first, then integer vs non-integer; Verify and + Vectorize classify the current Cast input independently and consistently; +* integer VMI actions: identity 0, same-width cross-sign 1 vinterpret_cast, + same-sign widening 1 vcvt, cross-sign widening source-sign extend + same-width + vinterpret_cast (2), supported narrowing 1 vcvt NOSAT; all 36 integer pairs + at VL=128 follow the table; +* the converted input is produced once and steps are emitted in order with + continuous dtype/VL and an exact final type; vinterpret_cast participates in + the local DCE; +* scalar plans keep ordinary scalar lowering (KeepInput / validated rebuild), + no new scalar rounding/saturation mode, and no parallel planner is wired + into generic codegen; +* attribute rules: accepted conversions pin the emitted rounding/saturation + mode, and unsupported combinations are rejected with a reason (unknown keys, + illegal sat policies, float <-> integer, f16 <-> bf16, out-of-scope FP8, + lane-uniform FP8); +* classification scenarios: external scalars, uniform loads, Bind aliases, 2D + kept/selected coordinates, nested casts, and scalar cast then broadcast. + +Full backend compilation and the device numerical checks live in +``test_pto_parallel_cast_matrix_run.py``. +""" + +from __future__ import annotations + +import os +import re +import subprocess +from contextlib import contextmanager + +import pytest +import torch # noqa: F401 (torch/torch_npu must import before tilelang/ptodsl) +import torch_npu # noqa: F401 + +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + +tirx = tvm.tirx +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) + +INT_TYPES = ["int8", "uint8", "int16", "uint16", "int32", "uint32"] +CROSS_SIGN_WIDEN_INTERMEDIATE = { + ("int8", "uint16"): "int16", + ("int8", "uint32"): "int32", + ("uint8", "int16"): "uint16", + ("uint8", "int32"): "uint32", + ("int16", "uint32"): "int32", + ("uint16", "int32"): "uint32", +} +PTO_INT_NAME = { + "int8": "pto.si8", + "uint8": "pto.ui8", + "int16": "pto.si16", + "uint16": "pto.ui16", + "int32": "pto.si32", + "uint32": "pto.ui32", +} + + +def _bits(dtype): + return int("".join(ch for ch in dtype if ch.isdigit())) + + +def _signed(dtype): + return dtype.startswith("int") + + +def _pair_category(src, dst): + if src == dst: + return "identity" + if _bits(src) == _bits(dst): + return "reinterpret" + if _bits(dst) > _bits(src): + return "widen" if _signed(src) == _signed(dst) else "widen_and_reinterpret" + return "narrow" + + +# --------------------------------------------------------------------------- +# Lowering / capture helpers +# --------------------------------------------------------------------------- +def _target(): + return normalize_pto_target("pto") + + +def _lower_source(func): + return _lower(func, target=_target()).kernel_source + + +@contextmanager +def _verify_capture(): + captured = {} + original = ascend_transform.VerifyParallelToPTO + + def factory(): + verify = original() + + def run(mod): + captured["module"] = mod + return verify(mod) + + return run + + ascend_transform.VerifyParallelToPTO = factory + try: + yield captured + finally: + ascend_transform.VerifyParallelToPTO = original + + +def _lower_capture(func): + error = None + with _verify_capture() as captured: + try: + _lower(func, target=_target()) + except Exception as exc: # noqa: BLE001 - expected diagnostics + error = exc + return captured.get("module"), error + + +def _expect_lower_error(func, pattern): + _, error = _lower_capture(func) + assert error is not None, f"expected a Verify diagnostic matching {pattern!r}" + assert re.search(pattern, str(error)), f"unexpected diagnostic: {error}" + + +def _cast_lines(source): + """Lines emitted by a value-Cast plan. + + Excludes the scalar-bridge `vinterpret_cast` that wraps `vbrc` (the scalar + signedness recovery), which is broadcasting, not a cast action. + """ + lines = [] + for line in source.splitlines(): + stripped = line.strip() + if "vmi.vcvt" in stripped or ("vinterpret_cast" in stripped and "vmi.vbrc(" not in stripped and "vmi.vci(" not in stripped): + lines.append(stripped) + return lines + + +def _one(source, needle): + hits = [line.strip() for line in source.splitlines() if needle in line] + assert len(hits) == 1, f"expected one line containing {needle!r}, got {hits!r}" + return hits[0] + + +def _count(source, needle): + return source.count(needle) + + +# --------------------------------------------------------------------------- +# Kernels +# --------------------------------------------------------------------------- +def _cast_kernel(src, dst, lanes, annotations=None): + n = 2 * lanes + + @T.prim_func + def main(A: T.Buffer((n,), src), C: T.Buffer((n,), dst)): + with T.Kernel(1): + a = T.alloc_shared((n,), src) + c = T.alloc_shared((n,), dst) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + if annotations is None: + c[i] = T.cast(a[i], dst) + else: + c[i] = T.Cast(dst, a[i], annotations=annotations) + T.copy(c, C) + + return main + + +def _ann_kernel(src, dst, lanes, annotations, *, uniform=False): + n = lanes + + @T.prim_func + def main(A: T.Buffer((n,), src), C: T.Buffer((n,), dst)): + with T.Kernel(1): + a = T.alloc_shared((n,), src) + c = T.alloc_shared((n,), dst) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + value = a[0] if uniform else a[i] + c[i] = T.Cast(dst, value, annotations=annotations) + T.copy(c, C) + + return main + + +def _alias_kernel(src, dst, lanes, uniform): + n = lanes + + @T.prim_func + def main(A: T.Buffer((n,), src), C: T.Buffer((n,), dst)): + with T.Kernel(1): + a = T.alloc_shared((n,), src) + c = T.alloc_shared((n,), dst) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + x = a[0] if uniform else a[i] + y = x + c[i] = T.cast(y, dst) + T.copy(c, C) + + return main + + +def _external_scalar_kernel(vl): + n = vl + + @T.prim_func + def main(S: T.int8, C: T.Buffer((n,), "int32")): + with T.Kernel(1): + c = T.alloc_shared((n,), "int32") + with T.SimdVF(lanes=vl): + for i in T.Parallel(n): + c[i] = T.cast(S, "int32") + T.copy(c, C) + + return main + + +def _binary_scalar_cast_kernel(vl): + n = vl + + @T.prim_func + def main(A: T.Buffer((n,), "int16"), B: T.Buffer((n,), "int8"), C: T.Buffer((n,), "int16")): + with T.Kernel(1): + a = T.alloc_shared((n,), "int16") + b = T.alloc_shared((n,), "int8") + c = T.alloc_shared((n,), "int16") + T.copy(A, a) + T.copy(B, b) + with T.SimdVF(lanes=vl): + for i in T.Parallel(n): + c[i] = a[i] + T.cast(b[0], "int16") + T.copy(c, C) + + return main + + +def _coordinate_kernel(vl, which): + m, n = 4, vl + out = "int16" + + @T.prim_func + def main(A: T.Buffer((m, n), "int32"), C: T.Buffer((m, n), out)): + with T.Kernel(1): + a = T.alloc_shared((m, n), "int32") + c = T.alloc_shared((m, n), out) + T.copy(A, a) + with T.SimdVF(lanes=vl): + for i, j in T.Parallel(m, n): + if which == "kept_of_load": + c[i, j] = T.cast(a[i, 0], out) + else: + c[i, j] = T.cast(a[i, j], out) + T.copy(c, C) + + return main + + +def _nested_cast_kernel(vl): + n = vl + + @T.prim_func + def main(A: T.Buffer((n,), "int16"), C: T.Buffer((n,), "uint32")): + with T.Kernel(1): + a = T.alloc_shared((n,), "int16") + c = T.alloc_shared((n,), "uint32") + T.copy(A, a) + with T.SimdVF(lanes=vl): + for i in T.Parallel(n): + c[i] = T.cast(T.cast(a[i], "uint16"), "uint32") + T.copy(c, C) + + return main + + +def _dead_chain_kernel(vl): + n = vl + + @T.prim_func + def main(A: T.Buffer((n,), "int8"), C: T.Buffer((n,), "int8")): + with T.Kernel(1): + a = T.alloc_shared((n,), "int8") + c = T.alloc_shared((n,), "int8") + T.copy(A, a) + with T.SimdVF(lanes=vl): + for i in T.Parallel(n): + # The assignment creates the Bind whose pure conversion + # chain must be eliminated by the local DCE. + unused = T.cast(T.cast(a[i], "uint16"), "uint32") # noqa: F841 + c[i] = a[i] + T.copy(c, C) + + return main + + +def _fp8_kernel(src, dst, lanes, annotations=None, **kwargs): + n = lanes + + @T.prim_func + def main(A: T.Buffer((n,), src), C: T.Buffer((n,), dst)): + with T.Kernel(1): + a = T.alloc_shared((n,), src) + c = T.alloc_shared((n,), dst) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + if annotations is not None: + c[i] = T.Cast(dst, a[i], annotations=annotations) + else: + c[i] = T.cast(a[i], dst, **kwargs) + T.copy(c, C) + + return main + + +# --------------------------------------------------------------------------- +# A. 36 integer pairs at VL=128: action structure (lower-only, fast) +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("src", INT_TYPES) +@pytest.mark.parametrize("dst", INT_TYPES) +def test_integer_pair_action_structure(src, dst): + category = _pair_category(src, dst) + source = _lower_source(_cast_kernel(src, dst, 128)) + lines = _cast_lines(source) + assert _count(source, "vmi.vload") == 1, source + if category == "identity": + assert lines == [], lines + return + if category == "reinterpret": + assert len(lines) == 1 and "vinterpret_cast" in lines[0], lines + assert f"to_dtype={PTO_INT_NAME[dst]}" in lines[0], lines[0] + return + if category == "widen": + assert len(lines) == 1 and "vmi.vcvt" in lines[0], lines + assert f"to_dtype={PTO_INT_NAME[dst]}" in lines[0], lines[0] + assert "saturate" not in lines[0], lines[0] + return + if category == "widen_and_reinterpret": + intermediate = CROSS_SIGN_WIDEN_INTERMEDIATE[(src, dst)] + assert len(lines) == 2, lines + assert "vmi.vcvt" in lines[0], lines[0] + assert f"to_dtype={PTO_INT_NAME[intermediate]}" in lines[0], lines[0] + assert "vinterpret_cast" in lines[1], lines[1] + assert f"to_dtype={PTO_INT_NAME[dst]}" in lines[1], lines[1] + first_result = lines[0].split("=")[0].strip() + assert lines[1].count(first_result) == 1, lines + return + # Narrowing: one direct logical vcvt with NOSAT. + assert len(lines) == 1 and "vmi.vcvt" in lines[0], lines + assert 'saturate="NOSAT"' in lines[0], lines[0] + assert f"to_dtype={PTO_INT_NAME[dst]}" in lines[0], lines[0] + + +def _count_cast_nodes(module): + func = next(iter(module.functions.values())) + found = [] + + def visit(node): + if isinstance(node, tirx.Cast): + found.append(node) + + tirx.stmt_functor.post_order_visit(func.body, visit) + return found + + +@pytest.mark.parametrize("dtype", INT_TYPES) +def test_identity_cast_reaches_planner_and_emits_zero_actions(dtype): + # A plain same-dtype T.cast is folded by the frontend, so the annotated + # raw Cast is used to prove the planner identity path is reachable: the + # manual sat=True marker must be validated and normalized away, and the + # emission has zero conversion actions. + func = _ann_kernel(dtype, dtype, 128, {"sat": True}) + module, error = _lower_capture(func) + assert error is None, error + casts = [node for node in _count_cast_nodes(module) if str(node.dtype) == dtype and str(node.value.dtype) == dtype] + assert casts, "the annotated same-dtype Cast must reach Verify" + source = _lower_source(func) + assert _cast_lines(source) == [] + + +# --------------------------------------------------------------------------- +# B. Producer once, ordered binds, scope and DCE +# --------------------------------------------------------------------------- +def test_cross_sign_widening_producer_and_order(): + source = _lower_source(_cast_kernel("int8", "uint32", 256)) + lines = _cast_lines(source) + assert len(lines) == 2 + vload = _one(source, "vmi.vload") + assert vload.split("=")[0].strip() in lines[0] + assert source.index(lines[0]) < source.index(lines[1]) + assert _count(source, "vmi.vload") == 1 + + +def test_conversion_binds_stay_between_load_and_store(): + source = _lower_source(_cast_kernel("int8", "uint32", 128)) + vload_index = source.index("vmi.vload") + first_cast = source.index(_cast_lines(source)[0]) + store_index = source.index("vmi.vstore") + assert vload_index < first_cast < store_index + + +def test_dead_pure_conversion_chain_is_dropped(): + source = _lower_source(_dead_chain_kernel(128)) + assert "vmi.vcvt" not in source, source + assert "vinterpret_cast" not in source, source + # The dead chain's vload is gone too; only the live a[i] load remains. + assert _count(source, "vmi.vload") == 1, source + + +def test_nested_cast_chains_inner_then_outer(): + source = _lower_source(_nested_cast_kernel(128)) + lines = _cast_lines(source) + assert len(lines) == 2, lines + assert "vinterpret_cast" in lines[0] and "to_dtype=pto.ui16" in lines[0] + assert "vmi.vcvt" in lines[1] and "to_dtype=pto.ui32" in lines[1] + assert _count(source, "vmi.vload") == 1 + + +# --------------------------------------------------------------------------- +# C. Verify/Vectorize classification consistency +# --------------------------------------------------------------------------- +def test_external_scalar_cast_stays_scalar(): + source = _lower_source(_external_scalar_kernel(128)) + assert _cast_lines(source) == [] + assert "pto.cast(" in source + assert "vbrc" in source + + +def test_uniform_load_cast_stays_scalar(): + source = _lower_source(_ann_kernel("int8", "int32", 128, {}, uniform=True)) + assert _cast_lines(source) == [] + assert "pto.cast(" in source + + +def test_bind_alias_keeps_vmi_classification(): + source = _lower_source(_alias_kernel("int8", "int32", 128, uniform=False)) + assert len(_cast_lines(source)) == 1 + assert _count(source, "vmi.vload") == 1 + + +def test_scalar_cast_then_parent_broadcast(): + source = _lower_source(_binary_scalar_cast_kernel(128)) + assert "vbrc" in source + assert "pto.cast(" in source + assert _cast_lines(source) == [] + + +@pytest.mark.parametrize( + ("which", "expect_vector"), + [ + ("kept_of_load", False), + ("selected_of_load", True), + ], +) +def test_2d_coordinate_classification(which, expect_vector): + source = _lower_source(_coordinate_kernel(128, which)) + lines = _cast_lines(source) + if expect_vector: + assert len(lines) == 1, lines + else: + assert lines == [], lines + assert "pto.cast(" in source + + +def test_lane_one_vector_context_is_not_constructible(): + # {kVectorVmi, 1} would be a VMI value with one logical lane; the current + # Parallel VL contract cannot produce it (lanes are 64/128/256), and the + # planner rejects it explicitly. The observable frontend boundary: + with pytest.raises(ValueError, match="lanes must be one of"): + T.SimdVF(lanes=1) + + +# --------------------------------------------------------------------------- +# D. Attribute table +# --------------------------------------------------------------------------- +def _lower_ann(src, dst, annotations): + return _lower_source(_ann_kernel(src, dst, 128, annotations)) + + +def test_integer_narrowing_accepts_sat_false(): + source = _lower_ann("int16", "int8", {"sat": False}) + line = _cast_lines(source)[0] + assert 'saturate="NOSAT"' in line, line + + +def test_float_narrowing_accepts_sat_false_and_nosat(): + source = _lower_ann("float32", "float16", {"sat": False}) + line = _cast_lines(source)[0] + assert 'rounding="R"' in line and 'saturate="NOSAT"' in line, line + + +def test_fp8_narrowing_accepts_round_and_sat_false(): + source = _lower_source(_fp8_kernel("float32", "float8_e5m2", 128, {"round": "Z", "sat": False})) + line = _cast_lines(source)[0] + assert 'rounding="Z"' in line and 'saturate="NOSAT"' in line, line + + +def test_integer_narrowing_rejects_sat_true(): + _expect_lower_error( + _ann_kernel("int16", "int8", 128, {"sat": True}), + "integer narrowing accepts only sat=False", + ) + + +def test_unknown_annotation_is_rejected(): + _expect_lower_error( + _ann_kernel("int8", "int32", 128, {"mystery": 1}), + "unknown cast annotation", + ) + + +def test_float_int_conversion_is_rejected(): + _expect_lower_error( + _ann_kernel("float32", "int32", 128, {}), + "float <-> integer conversions need per-combination", + ) + + +def test_f16_bf16_cross_family_is_rejected(): + _expect_lower_error( + _ann_kernel("float16", "bfloat16", 128, {}), + "cross-family f16 <-> bf16", + ) + + +def test_fp8_to_half_is_rejected(): + _expect_lower_error( + _fp8_kernel("float8_e4m3fn", "float16", 128), + "FP8 -> FP16/BF16 and FP8 <-> integer are out of scope", + ) + + +def test_lane_uniform_fp8_conversion_is_rejected(): + _expect_lower_error( + _ann_kernel("float8_e4m3fn", "float32", 128, {}, uniform=True), + "non-identity FP8 conversions have no scalar lowering", + ) + + +# --------------------------------------------------------------------------- +# E. Scalar paths and literal boundaries (existing semantics) +# --------------------------------------------------------------------------- +def test_scalar_float_widening_stays_scalar(): + source = _lower_source(_ann_kernel("float16", "float32", 128, {}, uniform=True)) + assert _cast_lines(source) == [] + assert "vbrc" in source + + +def test_scalar_float_narrowing_stays_scalar(): + source = _lower_source(_ann_kernel("float32", "float16", 128, {}, uniform=True)) + assert _cast_lines(source) == [] + assert "vbrc" in source + + +def test_scalar_integer_narrowing_keeps_plain_cast(): + source = _lower_source(_ann_kernel("int32", "int8", 128, {}, uniform=True)) + assert _cast_lines(source) == [] + assert "pto.cast(" in source + + +def test_actual_t_cast_literal_out_of_range_is_rejected_by_frontend(): + with pytest.raises(Exception, match="exceeds maximum of int8"): + T.cast(300, "int8") + + +def test_actual_t_cast_negative_to_unsigned_is_rejected_by_frontend(): + with pytest.raises(Exception, match="cannot make uint from negative value"): + T.cast(-1, "uint8") + + +# --------------------------------------------------------------------------- +# F. Direct planner contexts and scalar rebuilt span/annotations +# --------------------------------------------------------------------------- +def test_plan_value_cast_context_matrix_directly(tmp_path): + """The two passes can only produce scalar lanes=1 and VMI lanes=64/128/256, + so {kVectorVmi, 1} and {kScalar, 128} are unreachable end-to-end. A local + C++ probe calls PlanValueCast directly for both and for positive + vector/scalar plans, without adding a product FFI.""" + binary = tmp_path / "cast_plan_probe" + source = os.path.join(os.path.dirname(__file__), "cast_plan_context_probe.cc") + command = [ + "c++", + "-std=c++17", + "-O0", + source, + "-I", + f"{REPO_ROOT}/src", + "-I", + f"{REPO_ROOT}/3rdparty/tvm/include", + "-I", + f"{REPO_ROOT}/3rdparty/tvm/src", + "-I", + f"{REPO_ROOT}/3rdparty/tvm/3rdparty/dlpack/include", + "-I", + f"{REPO_ROOT}/3rdparty/tvm/3rdparty/tvm-ffi/include", + "-I", + f"{REPO_ROOT}/3rdparty/tvm/3rdparty/tvm-ffi/3rdparty/dlpack/include", + "-L", + f"{REPO_ROOT}/build/lib", + "-ltilelang", + "-ltvm_compiler", + "-ltvm_ffi", + "-ltvm_runtime", + f"-Wl,-rpath,{REPO_ROOT}/build/lib", + "-o", + str(binary), + ] + build = subprocess.run(command, cwd=REPO_ROOT, capture_output=True, text=True, timeout=600) + assert build.returncode == 0, build.stdout + build.stderr + env = dict(os.environ, LD_LIBRARY_PATH=f"{REPO_ROOT}/build/lib") + run = subprocess.run([str(binary)], env=env, capture_output=True, text=True, timeout=120) + assert run.returncode == 0, run.stdout + run.stderr + assert "FAIL" not in run.stdout, run.stdout + assert "kVectorVmi lanes=1 is rejected" in run.stdout + assert "kScalar lanes=128 is rejected" in run.stdout + + +@contextmanager +def _vectorize_capture(): + captured = {} + original = ascend_transform.VectorizeParallelToPTO + + def factory(): + vectorize = original() + + def run(mod): + out = vectorize(mod) + captured["module"] = out + return out + + return run + + ascend_transform.VectorizeParallelToPTO = factory + try: + yield captured + finally: + ascend_transform.VectorizeParallelToPTO = original + + +def _span_scalar_kernel(vl, span): + @T.prim_func + def main(S: T.Buffer((32,), "int32"), C: T.Buffer((vl,), "int8")): + with T.Kernel(1): + s = T.alloc_shared((32,), "int32") + c = T.alloc_shared((vl,), "int8") + T.copy(S, s) + with T.SimdVF(lanes=vl): + for i in T.Parallel(vl): + c[i] = T.cast(s[0], "int8", sat=False, span=span) + T.copy(c, C) + + return main + + +def test_scalar_rebuild_preserves_span_and_drop_provenance(): + # Span(source_name, line, end_line, column, end_column). + span = tvm.ir.Span(tvm.ir.SourceName("cast_span_probe"), 3, 3, 7, 11) + with _vectorize_capture() as captured: + _lower(_span_scalar_kernel(128, span), target=_target()) + module = captured["module"] + func = next(iter(module.functions.values())) + rebuilt = [] + + def visit(node): + if isinstance(node, tirx.Cast) and str(node.dtype) == "int8": + rebuilt.append(node) + + tirx.stmt_functor.post_order_visit(func.body, visit) + assert rebuilt, "the rebuilt scalar Cast is missing" + matches = [ + node + for node in rebuilt + if node.span is not None and node.span.source_name is not None and node.span.source_name.name == "cast_span_probe" + ] + assert matches, "the rebuilt scalar Cast lost the original span" + node = matches[0] + assert node.span.line == 3 and node.span.column == 7 + assert node.span.end_line == 3 and node.span.end_column == 11 + # The validated sat=False marker is proven equivalent to the ordinary + # truncating scalar lowering and is normalized away; no new attribute + # format is introduced. + assert "sat" not in node.annotations + assert "round" not in node.annotations diff --git a/testing/ascend/language/pto/test_pto_parallel_cast_matrix_run.py b/testing/ascend/language/pto/test_pto_parallel_cast_matrix_run.py new file mode 100644 index 0000000000..4ef6037726 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_parallel_cast_matrix_run.py @@ -0,0 +1,400 @@ +"""Device acceptance: parallel value-Cast matrix. + +Each case compiles through the real TileLang -> Verify/Vectorize -> PTODSL -> +PTOAS VPTO -> LLVM pipeline and executes on the NPU. References are +computed independently on the CPU with explicit low-bit selection and +target-signedness adjustment (never from the generated kernel or backend). + +Coverage: +* 22 representative integer pair cells at VL=128, direct store: all 12 + narrowing pairs, all 6 cross-signed widening pairs, 2 same-signed widening + pairs, and 2 same-width cross-sign reinterpret pairs; +* exact named boundary values through the full pipeline; +* target-signed continued consumption (same-sign widening of the converted + target) for representative 8/16-bit sources; +* one tail extent (2VL+17) with padded physical source and output sentinels; +* scalar (lane-uniform load) integer runtime casts and a scalar literal cast; +* the scalar-float expected backend rejection. +""" + +import os +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import torch_npu # noqa: E402,F401 (import before tilelang/ptodsl) +import tilelang # noqa: E402 +import tilelang.ascend.language as T # noqa: E402 + +TORCH_DT = { + "int8": torch.int8, + "uint8": torch.uint8, + "int16": torch.int16, + "uint16": torch.uint16, + "int32": torch.int32, + "uint32": torch.uint32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, +} +SENTINEL = { + "int8": -99, + "uint8": 199, + "int16": -9999, + "uint16": 55555, + "int32": -999999, + "uint32": 4000000000, +} + + +def _require_npu_runtime(): + assert hasattr(torch, "npu"), "torch.npu is unavailable in this environment" + assert torch.npu.is_available(), "torch.npu is unavailable in the current environment" + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source), "NPU host->device copy probe mismatch" + + +@pytest.fixture(scope="module") +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +def _bits(dtype): + return int("".join(ch for ch in dtype if ch.isdigit())) + + +def _signed(dtype): + return dtype.startswith("int") + + +def _wrap(value, bits, signed): + value &= (1 << bits) - 1 + if signed and value >= (1 << (bits - 1)): + value -= 1 << bits + return value + + +def _convert(value, src, dst): + """Independent reference: select the source value, truncate to dst bits, + then interpret the result with the target signedness.""" + src_value = _wrap(int(value), _bits(src), _signed(src)) + result = src_value & ((1 << _bits(dst)) - 1) + if _signed(dst) and result >= (1 << (_bits(dst) - 1)): + result -= 1 << _bits(dst) + return result + + +def _source_values(src, count): + bits = _bits(src) + signed = _signed(src) + raw = [ + -(1 << (bits - 1)) if signed else 0, + (1 << (bits - 1)) - 1 if signed else (1 << bits) - 1, + 0, + 1, + 2, + -1 if signed else (1 << (bits - 1)), + -2 if signed else (1 << (bits - 1)) + 1, + 127, + 128, + 255, + 256, + 257, + 300, + 65535, + 65536, + 4294967295, + 0x80000000, + ] + values = [_wrap(value, bits, signed) for value in raw] + return [values[index % len(values)] for index in range(count)] + + +def _tensor(values, dtype): + return torch.tensor(values, dtype=TORCH_DT[dtype]) + + +# --------------------------------------------------------------------------- +# Kernels +# --------------------------------------------------------------------------- +def _cast_kernel(src, dst, lanes, extent=None): + n = extent if extent is not None else 2 * lanes + storage = _storage_len(src, lanes, n) + + @T.prim_func + def main(A: T.Buffer((storage,), src), C: T.Buffer((storage,), dst)): + with T.Kernel(1): + a = T.alloc_shared((storage,), src) + c = T.alloc_shared((storage,), dst) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.cast(a[i], dst) + T.copy(c, C) + + return main, storage + + +def _consumer_kernel(src, mid, out, lanes): + n = 2 * lanes + storage = _storage_len(src, lanes, n) + + @T.prim_func + def main(A: T.Buffer((storage,), src), C: T.Buffer((storage,), out)): + with T.Kernel(1): + a = T.alloc_shared((storage,), src) + c = T.alloc_shared((storage,), out) + T.copy(A, a) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.cast(T.cast(a[i], mid), out) + T.copy(c, C) + + return main + + +def _scalar_cast_kernel(src, dst, lanes): + n = lanes + + @T.prim_func + def main(S: T.Buffer((32,), src), C: T.Buffer((n,), dst)): + with T.Kernel(1): + s = T.alloc_shared((32,), src) + c = T.alloc_shared((n,), dst) + T.copy(S, s) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(n): + c[i] = T.cast(s[0], dst) + T.copy(c, C) + + return main + + +def _physical_lanes(dtype): + """Physical VMI vector width in elements (256-byte vectors).""" + return 256 * 8 // _bits(dtype) + + +def _storage_len(src, lanes, extent): + """Readable source storage for the real continuous-vload footprint. + + A logical vload of ``size`` lanes is declared with the full physical + vector extent on the last chunk, so the last chunk base must be followed + by a full physical vector even when the logical lane count is smaller. + """ + padded = ((extent + lanes - 1) // lanes) * lanes + return max(padded, padded - lanes + _physical_lanes(src)) + + +def _tail_kernel(src, dst, lanes, extent): + storage = _storage_len(src, lanes, extent) + + @T.prim_func + def main( + A: T.Buffer((storage,), src), + S: T.Buffer((storage,), dst), + C: T.Buffer((storage,), dst), + ): + with T.Kernel(1): + a = T.alloc_shared((storage,), src) + c = T.alloc_shared((storage,), dst) + T.copy(A, a) + T.copy(S, c) + with T.SimdVF(lanes=lanes): + for i in T.Parallel(extent): + c[i] = T.cast(a[i], dst) + T.copy(c, C) + + return main, storage + + +def _run_pair(npu, src, dst, lanes, extent=None): + func, storage = _cast_kernel(src, dst, lanes, extent=extent) + kernel = tilelang.compile(func, target="pto", out_idx=-1) + n = extent if extent is not None else 2 * lanes + values = _source_values(src, storage) + expected = torch.tensor([_convert(value, src, dst) for value in values], dtype=TORCH_DT[dst]) + out = kernel(_tensor(values, src).to(npu)) + torch.npu.synchronize() + got = out.cpu() + assert got.dtype == expected.dtype, (src, dst, lanes, got.dtype) + assert torch.equal(got[0:n], expected[0:n]), ( + f"{src}->{dst} VL{lanes}: value mismatch\n got={got[:16].tolist()}\n exp={expected[:16].tolist()}" + ) + assert got.numel() == expected.numel() == storage + return got, expected, storage + + +# --------------------------------------------------------------------------- +# A. Representative pair cells at VL=128: full compile + device execution +# --------------------------------------------------------------------------- +DEVICE_PAIRS = [ + # Narrowing: one vcvt with NOSAT. + ("int16", "int8"), + ("int16", "uint8"), + ("uint16", "int8"), + ("uint16", "uint8"), + ("int32", "int8"), + ("int32", "uint8"), + ("int32", "int16"), + ("int32", "uint16"), + ("uint32", "int8"), + ("uint32", "uint8"), + ("uint32", "int16"), + ("uint32", "uint16"), + # Cross-signed widening: source-signed vcvt, then vinterpret_cast. + ("int8", "uint16"), + ("int8", "uint32"), + ("uint8", "int16"), + ("uint8", "int32"), + ("int16", "uint32"), + ("uint16", "int32"), + # Same-signed widening: one vcvt. + ("int8", "int16"), + ("uint16", "uint32"), + # Same-width cross-sign: one vinterpret_cast. + ("int8", "uint8"), + ("uint32", "int32"), +] + + +@pytest.mark.parametrize(("src", "dst"), DEVICE_PAIRS) +def test_integer_pair_vl_device(npu, src, dst): + got, _, storage = _run_pair(npu, src, dst, 128) + # The returned allocation is the readable-storage contract (physical + # load padding included); only the logical [0, 2*VL) region is written + # and compared. The padding beyond it is uninitialized by design. + assert got.numel() == storage + assert storage >= 256 + + +def test_reference_examples_are_exact(npu): + cases = [ + ("int8", "uint32", -1, 4294967295), + ("uint8", "int32", 255, 255), + ("int32", "int8", 300, 44), + ("int32", "uint8", -1, 255), + ("int32", "uint8", 256, 0), + ] + for src, dst, value, expected in cases: + func, storage = _cast_kernel(src, dst, 128, extent=128) + kernel = tilelang.compile(func, target="pto", out_idx=-1) + values = [value] * storage + source = torch.tensor(values, dtype=TORCH_DT[src]) + out = kernel(source.to(npu)) + torch.npu.synchronize() + got = out.cpu() + assert got.numel() == storage + # Only the logical extent (128) was written; the physical-padding + # region is not part of the value contract. + assert int(got[0]) == expected, (src, dst, value, int(got[0])) + assert bool((got[:128] == expected).all()), (src, dst, value, got.tolist()[:8]) + + +# --------------------------------------------------------------------------- +# B. Target-signed continued consumption +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("src", "mid", "out"), + [ + ("int8", "uint16", "uint32"), + ("int8", "uint8", "uint32"), + ("uint8", "int8", "int32"), + ("int16", "uint16", "uint32"), + ], +) +def test_target_signed_continued_widening(npu, src, mid, out): + """The converted target is consumed by its own same-sign widening: a + '-' source value must not sign-extend through the unsigned target.""" + func = _consumer_kernel(src, mid, out, 128) + kernel = tilelang.compile(func, target="pto", out_idx=-1) + storage = _storage_len(src, 128, 256) + values = [-1, 0, 1, 127, -128, 255, 256, 300] * 64 + values = [_wrap(value, _bits(src), _signed(src)) for value in values[:storage]] + expected = torch.tensor( + [_convert(_convert(v, src, mid), mid, out) for v in values[:256]], + dtype=TORCH_DT[out], + ) + got = kernel(torch.tensor(values, dtype=TORCH_DT[src]).to(npu)).cpu() + torch.npu.synchronize() + assert torch.equal(got[:256], expected), f"{src}->{mid}->{out}\n got={got[:16].tolist()}\n exp={expected[:16].tolist()}" + + +# --------------------------------------------------------------------------- +# C. Tail extent with padded source and output sentinels +# --------------------------------------------------------------------------- +def test_tail_extent_with_output_sentinel(npu): + src, dst, lanes, extent = "int8", "uint32", 128, 2 * 128 + 17 + func, storage = _tail_kernel(src, dst, lanes, extent) + kernel = tilelang.compile(func, target="pto", out_idx=-1) + values = _source_values(src, storage) + sentinel = SENTINEL[dst] + c_init = torch.full((storage,), sentinel, dtype=TORCH_DT[dst]) + got = kernel(_tensor(values, src).to(npu), c_init.to(npu)).cpu() + # `kernel` returns the third argument (C); `c_init` seeds it via S. + torch.npu.synchronize() + expected = torch.tensor([_convert(values[i], src, dst) for i in range(extent)], dtype=TORCH_DT[dst]) + assert torch.equal(got[0:extent], expected), f"{src}->{dst} VL{lanes} E={extent}: values" + assert bool((got[extent:storage] == sentinel).all()), f"{src}->{dst} VL{lanes} E={extent}: tail sentinel overwritten" + + +# --------------------------------------------------------------------------- +# D. Scalar (lane-uniform) integer and float controls +# --------------------------------------------------------------------------- +SCALAR_INT_CASES = [ + ("int8", "uint32"), + ("int32", "uint8"), +] + + +@pytest.mark.parametrize(("src", "dst"), SCALAR_INT_CASES) +def test_scalar_integer_runtime_cast(npu, src, dst): + func = _scalar_cast_kernel(src, dst, 128) + kernel = tilelang.compile(func, target="pto", out_idx=-1) + for value in (-1, 0, 1, 255, 256, 300): + source_values = [_wrap(value, _bits(src), _signed(src))] * 32 + source = torch.tensor(source_values, dtype=TORCH_DT[src]) + out = kernel(source.to(npu)).cpu() + torch.npu.synchronize() + expected = _convert(source_values[0], src, dst) + assert int(out[0]) == expected, (src, dst, value, int(out[0])) + assert bool((out == expected).all()), (src, dst, value) + + +def test_scalar_integer_literal_cast(npu): + @T.prim_func + def main(C: T.Buffer((128,), "int32")): + with T.Kernel(1): + c = T.alloc_shared((128,), "int32") + with T.SimdVF(lanes=128): + for i in T.Parallel(128): + c[i] = T.cast(7, "int32") + T.copy(c, C) + + kernel = tilelang.compile(main, target="pto", out_idx=-1) + out = kernel().cpu() + torch.npu.synchronize() + assert bool((out == 7).all()), out.tolist()[:8] + + +def test_scalar_float_cast_rejected_by_device_compiler(npu): + """The BiSheng device compiler rejects scalar float casts in an AIV loop. + + PTODSL, VPTO and LLVM emission support these conversions. The device + compiler reports "Unsupported scalar instruction in AIV loop" with or + without the value-Cast planner. This test checks that rejection; the + scalar planner preserves the existing cast semantics. + """ + with pytest.raises(Exception) as excinfo: + tilelang.compile(_scalar_cast_kernel("float32", "float16", 128), target="pto", out_idx=-1) + assert "Unsupported scalar instruction in AIV loop" in str(excinfo.value) 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_pto_vmi_broadcast_composition.py b/testing/ascend/language/pto/test_pto_vmi_broadcast_composition.py new file mode 100644 index 0000000000..4e45f297c9 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_vmi_broadcast_composition.py @@ -0,0 +1,51 @@ +"""Source-level coverage for explicit VMI broadcast-load compositions.""" + +import ast + +import pytest + +import tilelang.ascend.language as T +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + + +@pytest.mark.pto +@pytest.mark.parametrize("group,stride", [(None, None), (2, 1), (4, 3)]) +def test_pto_codegen_preserves_vload_vbrc_composition(group, stride): + load_size = 1 if group is None else group + + @T.prim_func + def func(A: T.Tensor((64,), "float32"), B: T.Tensor((64,), "float32")): + with T.Kernel(1): + a_ub = T.alloc_shared((64,), "float32") + b_ub = T.alloc_shared((64,), "float32") + T.copy(A, a_ub) + with T.SimdVF(): + loaded = T.vmi.vload(a_ub[4], size=load_size, group=group, stride=stride) + value = T.vmi.vbrc(loaded, size=64, group=group) + mask = T.vmi.create_mask(64, size=64) + T.vmi.vstore(value, b_ub[0], mask) + T.copy(b_ub, B) + + target = normalize_pto_target("pto") + with target: + source = _lower(func, target=target).kernel_source + compile(source, "", "exec") + load_line = next(line for line in source.splitlines() if "pto.vmi.vload(" in line) + broadcast_line = next(line for line in source.splitlines() if "pto.vmi.vbrc(" in line) + load = ast.parse(load_line.strip()).body[0] + broadcast = ast.parse(broadcast_line.strip()).body[0] + assert isinstance(load, ast.Assign) + assert isinstance(broadcast, ast.Assign) + assert broadcast.value.args[0].id == load.targets[0].id + assert ast.literal_eval(load.value.args[1]) == 0 + load_options = {kw.arg: ast.literal_eval(kw.value) for kw in load.value.keywords} + broadcast_options = {kw.arg: ast.literal_eval(kw.value) for kw in broadcast.value.keywords} + assert "dist_mode" not in source + if group is None: + assert load_options == {"size": 1} + assert broadcast_options == {"size": 64} + else: + assert load_options == {"size": group, "group": group, "stride": stride} + assert broadcast_options == {"size": 64, "group": group} + diff --git a/testing/ascend/language/pto/test_pto_vmi_memory_access.py b/testing/ascend/language/pto/test_pto_vmi_memory_access.py new file mode 100644 index 0000000000..8188ec80a8 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_vmi_memory_access.py @@ -0,0 +1,783 @@ +"""VMI memory-access descriptors and interface adaptation (structural). + +Hardware independent. Each case lowers the real TileLang PTO pipeline (no +device execution) and inspects the produced TIR/PTODSL. Cases that the +backend can express also call ``tilelang.compile``, which runs PTODSL -> PTOAS +(VPTO) -> Bisheng; that is compile evidence, not numerical evidence (device +results live in ``test_pto_vmi_memory_access_run.py``). + +Contracts verified here: + +* ``T.vmi.vload``/``T.vmi.vstore`` access descriptors carry the true start + address and a per-access-mode effective extent, with the VMI instruction + offset fixed at 0, so the scheduler sees the actual accessed range. +* Ordinary continuous loads round up to full physical vectors; ordinary + broadcast keeps the conservative extent; ``dintlv`` loads declare both + results and a two-input ``dintlv`` store declares the full 2*N span (its + PTODSL attribute is emitted as ``intlv``); group envelopes cover all + groups (slot/broadcast count one element per group, a static negative + stride relocates the lowest address, a dynamic stride falls back to the + whole tight buffer); ``block_stride`` extents are measured in 32B blocks; + genuine FP4 descriptors stay logical and are packed exactly once (odd + starts are rejected); packed (PK) store extents stay logical target + elements; an explicit ``access_ptr`` keeps the caller-declared descriptor + plus the extra offset and the UB scope check; gather/scatter preserve + non-zero bases. +* Masks never shrink the may-write extent, and a logical ``size=80`` load + declares both physical reads. +* Removed interface parameters (``repeat_stride``, vload ``unpack`` / + ``to_dtype``, vstore ``pmode``) fail in a legal scope. + +``size=80`` VMI vectors are not constructible by the PTODSL VMI API +(lane counts are 1/2/4/8/64/128/256). Their real lowering is verified with a +direct PTOAS compile of the PTO form the descriptor describes; the TileLang +diagnostics/descriptors are verified structurally. +""" + +from __future__ import annotations + +import os +import subprocess +import tempfile +from contextlib import contextmanager + +import pytest +import torch # noqa: F401 (torch/torch_npu must import before tilelang/ptodsl) +import torch_npu # noqa: F401 + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend import transform as ascend_transform +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + +FP4 = "float4_e2m1fn" +F16 = "float16" +F32 = "float32" +F8_E4M3 = "float8_e4m3fn" + + +# --------------------------------------------------------------------------- +# Pipeline/diagnostic helpers +# --------------------------------------------------------------------------- +def _lower_pto(func): + """Lower *func* for the PTO target and return the compiled artifact.""" + target = normalize_pto_target("pto") + with target: + return _lower(func, target=target) + + +def _compile_pto(func): + """Compile through the PTODSL -> PTOAS (VPTO) -> Bisheng pipeline.""" + return tilelang.compile(func, target="pto", out_idx=-1) + + +def _iter_calls(artifact, op_name): + for func in artifact.device_mod.functions.values(): + found = [] + + def visit(node, found=found): + if isinstance(node, tvm.tirx.Call) and getattr(node.op, "name", "") == op_name: + found.append(node) + + tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + for node in found: + yield func, node + + +def _resolve_binding(expr, func): + """Resolve a local handle binding (``p = access_ptr(...)``) to its value.""" + if not isinstance(expr, tvm.tirx.Var): + return expr + found = {} + + def visit(node): + if isinstance(node, (tvm.tirx.Let, tvm.tirx.Bind)) and node.var is not None and node.var.same_as(expr): + found["value"] = node.value + + tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + return found.get("value", expr) + + +def _const(expr): + from tvm import arith + + simplified = arith.Analyzer().simplify(expr) + assert isinstance(simplified, tvm.tirx.IntImm), f"expected a constant, got {simplified!r}" + return int(simplified.value) + + +def _single(artifact, op_name): + pairs = list(_iter_calls(artifact, op_name)) + assert len(pairs) == 1, f"expected exactly one {op_name}, got {len(pairs)}" + return pairs[0] + + +def _ptr_position(call, op_name): + return 1 if op_name == "tl.vmi.vstore" else 0 + + +def _descriptor(artifact, op_name): + """Return (descriptor_offset, extent, rw_mask, instruction_offset).""" + func, call = _single(artifact, op_name) + ptr_pos = _ptr_position(call, op_name) + ptr = _resolve_binding(call.args[ptr_pos], func) + assert isinstance(ptr, tvm.tirx.Call) and ptr.op.name.endswith("tvm_access_ptr"), ( + f"{op_name} address is not an access descriptor: {ptr!r}" + ) + return ( + _const(ptr.args[2]), + _const(ptr.args[3]), + _const(ptr.args[4]), + _const(call.args[ptr_pos + 1]), + ) + + +def _vload_descriptor(func): + return _descriptor(_lower_pto(func), "tl.vmi.vload") + + +def _vstore_descriptor(func): + return _descriptor(_lower_pto(func), "tl.vmi.vstore") + + +@contextmanager +def _capture_rewrite_input(): + """Capture the module handed to RewriteFp4ToFp4x2, then delegate.""" + captured = {} + original = ascend_transform.RewriteFp4ToFp4x2 + + def factory(): + rewrite = original() + + def run(mod): + captured["module"] = mod + return rewrite(mod) + + return run + + ascend_transform.RewriteFp4ToFp4x2 = factory + try: + yield captured + finally: + ascend_transform.RewriteFp4ToFp4x2 = original + + +def _descriptor_in_module(mod, op_name): + class _Artifact: + device_mod = mod + + return _descriptor(_Artifact(), op_name) + + +def _ptoas_vpto(pto_text: str) -> str: + """Compile direct PTO IR through PTOAS and return the VPTO text.""" + work = tempfile.mkdtemp(prefix="vmi-vpto-", dir=os.environ.get("TMPDIR") or None) + pto_path = os.path.join(work, "case.pto") + vpto_path = os.path.join(work, "case.vpto.mlir") + with open(pto_path, "w", encoding="utf-8") as file: + file.write(pto_text) + result = subprocess.run( + ["ptoas", "--pto-arch=a5", "--pto-backend=vpto", "--emit-vpto", pto_path, "-o", vpto_path], + text=True, + capture_output=True, + ) + assert result.returncode == 0, f"PTOAS rejected the probe:\n{result.stderr}" + with open(vpto_path, encoding="utf-8") as file: + return file.read() + + +# --------------------------------------------------------------------------- +# Ordinary non-zero BufferLoad start +# --------------------------------------------------------------------------- +def _load_store_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((320,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work[0:256]) + T.fill(work[256:320], T.float32(-5.0)) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + v = T.vmi.vload(work[64], size=64) + T.vmi.vstore(v, work[192], mask) + T.copy(work[0:320], C) + + return main + + +def test_load_descriptor_originates_at_actual_load(): + """vload(work[64], size=64) declares [64,128) and instruction offset 0.""" + artifact = _lower_pto(_load_store_kernel()) + assert _descriptor(artifact, "tl.vmi.vload") == (64, 64, 1, 0) + + source = artifact.kernel_source + vload_line = next(ln.strip() for ln in source.splitlines() if "vmi.vload" in ln) + assert "addptr" in vload_line and ", 64)" in vload_line, vload_line + assert ", 0, size=64)" in vload_line, vload_line + + +def test_store_descriptor_originates_at_actual_destination(): + """vstore(..., work[192]) declares may-write [192,256) with offset 0.""" + artifact = _lower_pto(_load_store_kernel()) + assert _descriptor(artifact, "tl.vmi.vstore") == (192, 64, 2, 0) + + source = artifact.kernel_source + vstore_line = next(ln.strip() for ln in source.splitlines() if "vmi.vstore" in ln) + assert "addptr" in vstore_line and ", 192)" in vstore_line, vstore_line + assert "), 0, mask)" in vstore_line, vstore_line + + +def test_nonzero_load_and_store_compile_through_backend(): + kernel = _compile_pto(_load_store_kernel()) + source = kernel.get_kernel_source() + assert "pto.vmi.vload(" in source and "pto.vmi.vstore(" in source, source + + +# --------------------------------------------------------------------------- +# Buffer / BufferRegion / elem_offset +# --------------------------------------------------------------------------- +def _buffer_region_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + v = T.vmi.vload(work[0:64], size=64) + T.vmi.vstore(v, work[128], mask) + T.copy(work, C) + + return main + + +def _elem_offset_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + view = tvm.tirx.decl_buffer((128,), F32, data=work.data, elem_offset=16, scope="shared.dyn") + mask = T.vmi.create_mask(64, size=64) + v = T.vmi.vload(view[32], size=64) + T.vmi.vstore(v, view[32], mask) + T.copy(work, C) + + return main + + +def test_buffer_region_source_declares_from_region_minimum(): + assert _vload_descriptor(_buffer_region_kernel()) == (0, 64, 1, 0) + + +def test_elem_offset_counted_exactly_once(): + """A view with elem_offset=16 loaded at index 32 starts at element 48.""" + artifact = _lower_pto(_elem_offset_kernel()) + assert _descriptor(artifact, "tl.vmi.vload") == (48, 64, 1, 0) + assert _descriptor(artifact, "tl.vmi.vstore") == (48, 64, 2, 0) + + +# --------------------------------------------------------------------------- +# Ordinary broadcast +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# Physical-vector rounding of ordinary continuous loads +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("dtype", "size", "expected"), + [ + ("int8", 300, 512), # b8 VL = 256 + (F16, 200, 256), # b16 VL = 128 + (F32, 80, 128), # b32 VL = 64 (structural: PTODSL cannot build 80) + ], + ids=["b8", "b16", "b32-size80"], +) +def test_continuous_load_covers_full_physical_vectors(dtype, size, expected): + @T.prim_func + def main(A: T.Tensor((1024,), dtype), C: T.Tensor((1024,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((1024,), dtype) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[64], size=size) + T.vmi.vstore(v, work[512]) + T.copy(work, C) + + assert _vload_descriptor(main) == (64, expected, 1, 0) + assert _vstore_descriptor(main) == (512, size, 2, 0) + + +# --------------------------------------------------------------------------- +# dintlv load / two-input store +# --------------------------------------------------------------------------- + + + + +# --------------------------------------------------------------------------- +# Logical size 80 / masks never shrink the footprint +# --------------------------------------------------------------------------- +def test_size80_load_declares_two_full_physical_reads(): + """b32 size 80 loads two 64-lane physical reads, so the extent is 128.""" + + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[64], size=80) + T.vmi.vstore(v, work[128]) + T.copy(work, C) + + assert _vload_descriptor(main) == (64, 128, 1, 0) + assert _vstore_descriptor(main) == (128, 80, 2, 0) + + +def test_size80_backend_lowering_covers_second_physical_chunk(): + """Direct PTOAS compile of the size-80 PTO form the descriptor describes. + + The PTODSL VMI API cannot construct an 80-lane VMI vreg + (``_check_vmi_lane_count``), so this verifies the backend lowering itself: + two full vlds at offsets 64 and 128. + """ + pto_text = """module { + func.func @group1_n80_s80() attributes {pto.kernel} { + %start = arith.constant 64 : index + %stride = arith.constant 80 : index + %base = arith.constant 0 : i64 + %dest = arith.constant 4096 : i64 + %src = pto.castptr %base : i64 -> !pto.ptr + %dst = pto.castptr %dest : i64 -> !pto.ptr + pto.vecscope { + %v = pto.vmi.vload %src[%start], %stride {group = 1 : i64} : !pto.ptr -> !pto.vmi.vreg<80xf32> + pto.vmi.vstore %v, %dst[%start], %stride {group = 1 : i64} : !pto.vmi.vreg<80xf32>, !pto.ptr + } + return + } +} +""" + vpto = _ptoas_vpto(pto_text) + assert "pto.vlds %0[%c64]" in vpto, vpto + assert "pto.vlds %0[%c128]" in vpto, vpto + + +def test_store_mask_activity_does_not_shrink_may_write(): + """A partial or all-false mask keeps the full 64-lane may-write span.""" + + def kernel(mask_active): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(mask_active, size=64) + v = T.vmi.vload(work[0], size=64) + T.vmi.vstore(v, work[128], mask) + T.copy(work, C) + + return main + + for active in (0, 10, 64): + assert _vstore_descriptor(kernel(active)) == (128, 64, 2, 0), active + + +# --------------------------------------------------------------------------- +# Group envelopes, negative and dynamic stride +# --------------------------------------------------------------------------- +def test_group_ordinary_envelope_covers_all_groups(): + """A=64, N=128, G=2, S=128 => [64,256).""" + + @T.prim_func + def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + v = T.vmi.vload(work[64], size=128, group=2, stride=128) + T.vmi.vstore(v, work[256], mask) + T.copy(work, C) + + assert _vload_descriptor(main) == (64, 192, 1, 0) + + +def test_group_slot_uses_one_element_per_group(): + @T.prim_func + def slot(A: T.Tensor((768,), F32), C: T.Tensor((768,), F32)): + with T.Kernel(1): + # G=8/S=64: source touches index 512; the destination starting at + # 256 touches 256 + 7*64 = 704. Both need legal, initialized + # storage, so allocate/initialize 768 elements. + work = T.alloc_shared((768,), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[64], size=8, group=8, stride=64) + T.vmi.vstore(v, work[256], group=8, stride=64) + T.copy(work, C) + + assert _vload_descriptor(slot) == (64, 449, 1, 0) + assert _vstore_descriptor(slot) == (256, 449, 2, 0) + + + +def test_group_static_negative_stride_relocates_lowest_address(): + """A=192, N=128, G=2, S=-128 => envelope [64,256) and instruction offset 128.""" + + @T.prim_func + def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[192], size=128, group=2, stride=-128) + T.vmi.vstore(v, work[192], group=2, stride=-128) + T.copy(work, C) + + artifact = _lower_pto(main) + assert _descriptor(artifact, "tl.vmi.vload") == (64, 192, 1, 128) + assert _descriptor(artifact, "tl.vmi.vstore") == (64, 192, 2, 128) + + +def test_group_dynamic_stride_declares_whole_tight_buffer(): + @T.prim_func + def main(A: T.Tensor((512,), F32), C: T.Tensor((512,), F32), stride: T.int32): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[192], size=128, group=2, stride=stride) + T.vmi.vstore(v, work[192], group=2, stride=stride) + T.copy(work, C) + + artifact = _lower_pto(main) + assert _descriptor(artifact, "tl.vmi.vload") == (0, 512, 1, 192) + assert _descriptor(artifact, "tl.vmi.vstore") == (0, 512, 2, 192) + + +def test_group_negative_stride_nd_uses_whole_buffer_plan(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((4, 64), F32) + T.copy(A, work) + with T.SimdVF(): + v = T.vmi.vload(work[3, 0], size=128, group=2, stride=-64) + T.vmi.vstore(v, work[3, 0], group=2, stride=-64) + T.copy(work, C) + + artifact = _lower_pto(main) + assert _descriptor(artifact, "tl.vmi.vload") == (0, 256, 1, 192) + + +# --------------------------------------------------------------------------- +# block_stride +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("dtype", "size", "block_stride", "expected_extent"), + [ + (F32, 64, 1, 64), # one full b32 vector, contiguous blocks + (F32, 64, 2, 120), # [64, 184) + (F32, 128, 2, 248), # two full vectors: [64, 312) + (F16, 128, 2, 240), # b16: 8 blocks of 16 elements + ], + ids=["b32-single-vl-s1", "b32-single-vl-s2", "b32-double-vl-s2", "b16-single-vl-s2"], +) +def test_block_stride_envelope_uses_32b_blocks(dtype, size, block_stride, expected_extent): + # The destination at 320 spans up to 320 + expected_extent = 568 (b32 + # double-VL) / 560 (b16); allocate and initialize the real footprint. + @T.prim_func + def main(A: T.Tensor((768,), dtype), C: T.Tensor((768,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((768,), dtype) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + v = T.vmi.vload(work[64], size=size, block_stride=block_stride) + T.vmi.vstore(v, work[320], mask, block_stride=block_stride) + T.copy(work, C) + + artifact = _lower_pto(main) + assert _descriptor(artifact, "tl.vmi.vload") == (64, expected_extent, 1, 0) + assert _descriptor(artifact, "tl.vmi.vstore") == (320, expected_extent, 2, 0) + + +# --------------------------------------------------------------------------- +# Genuine FP4 +# --------------------------------------------------------------------------- +def _fp4_kernel(index=64, size=128): + @T.prim_func + def main(A: T.Tensor((256,), FP4), C: T.Tensor((256,), FP4)): + with T.Kernel(1): + work = T.alloc_shared((256,), FP4) + out = T.alloc_shared((256,), FP4) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(size // 2, size=size // 2) + v = T.vmi.vload(work[index], size=size) + T.vmi.vstore(v, out[index], mask) + T.copy(out, C) + + return main + + +def test_fp4_descriptor_stays_logical_then_packed_once(): + with _capture_rewrite_input() as captured: + artifact = _lower_pto(_fp4_kernel()) + before = captured["module"] + + # Logical FP4 element units before the rewrite. + assert _descriptor_in_module(before, "tl.vmi.vload") == (64, 128, 1, 0) + assert _descriptor_in_module(before, "tl.vmi.vstore") == (64, 128, 2, 0) + + # Exactly one physical packing: the single rewrite pass halves both. + rewritten = ascend_transform.RewriteFp4ToFp4x2()(before) + assert _descriptor_in_module(rewritten, "tl.vmi.vload") == (32, 64, 1, 0) + assert _descriptor_in_module(rewritten, "tl.vmi.vstore") == (32, 64, 2, 0) + + source = artifact.kernel_source + load_line = next(ln.strip() for ln in source.splitlines() if "vmi.vload" in ln) + assert "addptr" in load_line and "), 32)" in load_line, load_line + assert load_line.rstrip().endswith("size=(128 // 2))"), load_line + + +def test_fp4_odd_start_is_still_rejected(): + with pytest.raises(Exception, match="fp4 access_ptr offset must be even"): + _lower_pto(_fp4_kernel(index=65)) + + +# --------------------------------------------------------------------------- +# Packed (PK) stores keep logical target-element extents +# --------------------------------------------------------------------------- +def _bits(dtype) -> int: + return int(tvm.DataType(dtype).bits) + + +def _merged_target_offset(src, dst, count=512) -> int: + """Offset of a target buffer merged after ``count`` source elements. + + MergeUBAllocations packs the source span first in these two-buffer + kernels; descriptors are in destination-element units. Deriving the + expectation keeps the assertion on the absolute may-write start while + making the layout assumption explicit. + """ + source_bytes = count * _bits(src) // 8 + return source_bytes // (_bits(dst) // 8) + + +@pytest.mark.parametrize( + ("src", "dst", "vcvt_kwargs"), + [ + (F32, F16, {"rounding": "R", "saturate": "SAT"}), + (F32, F8_E4M3, {"rounding": "R", "saturate": "SAT"}), + ("int16", "int8", {"saturate": "NOSAT"}), + ], + ids=["pk-b32", "pk4-b32", "pk-b16"], +) +def test_packed_store_declares_logical_target_elements(src, dst, vcvt_kwargs): + size = 128 + + @T.prim_func + def main(A: T.Tensor((512,), src), C: T.Tensor((512,), dst)): + with T.Kernel(1): + source = T.alloc_shared((512,), src) + target = T.alloc_shared((512,), dst) + T.copy(A, source) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + v = T.vmi.vload(source[64], size=size) + narrowed = T.vmi.vcvt(v, to_dtype=dst, **vcvt_kwargs) + T.vmi.vstore(narrowed, target[64], mask) + T.copy(target, C) + + artifact = _lower_pto(main) + target_offset = _merged_target_offset(src, dst) + 64 + # The may-write span stays the logical target element count (128), not a + # physical vreg capacity or a PK-derived fraction. + assert _descriptor(artifact, "tl.vmi.vstore") == (target_offset, size, 2, 0) + src_extent = 128 if src == F32 else size + assert _descriptor(artifact, "tl.vmi.vload") == (64, src_extent, 1, 0) + + kernel = _compile_pto(main) + assert "pto.vmi.vcvt(" in kernel.get_kernel_source() + + +# --------------------------------------------------------------------------- +# Explicit access_ptr + extra offset +# --------------------------------------------------------------------------- +def _explicit_pointer_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + p = T.access_ptr(work[32], "r", extent=128) + v = T.vmi.vload(p, 16, size=64) + T.vmi.vstore(v, work[128], mask) + T.copy(work, C) + + return main + + +def test_explicit_pointer_keeps_declaration_and_extra_offset(): + artifact = _lower_pto(_explicit_pointer_kernel()) + func, call = _single(artifact, "tl.vmi.vload") + ptr = _resolve_binding(call.args[0], func) + # The caller-declared descriptor is not overwritten (offset 32, extent 128). + assert _const(ptr.args[2]) == 32 + assert _const(ptr.args[3]) == 128 + # The extra offset is preserved, not zeroed or folded twice. + assert _const(call.args[1]) == 16 + + source = artifact.kernel_source + lines = [ln.strip() for ln in source.splitlines()] + definition = next(ln for ln in lines if ln.startswith("p = ")) + assert "pto.addptr(" in definition and ", 32)" in definition, definition + # The VMI instruction offset stays 16: the base is counted once. + assert "vload(p, 16, size=64)" in source, source + + +def test_explicit_pointer_scope_validation_is_preserved(): + """Recovering the element type from an immediate access_ptr does not + bypass the UB scope check for indirect-address operations.""" + + def build(): + @T.prim_func + def main(A: T.Tensor((256,), F32), I: T.Tensor((64,), "uint32"), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + index = T.alloc_shared((64,), "uint32") + T.copy(A, work) + T.copy(I, index) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + offsets = T.vmi.vload(index[0], size=64) + p = T.access_ptr(A[0], "r", extent=64) + v = T.vmi.vgather(p, offsets, mask) + T.vmi.vstore(v, work[192], mask) + T.copy(work, C) + + return main + + with pytest.raises(TypeError, match="requires a UB pointer"): + build() + + +# --------------------------------------------------------------------------- +# Gather/scatter non-zero bases (no footprint analysis) +# --------------------------------------------------------------------------- +def _gather_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), I: T.Tensor((64,), "uint32"), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + index = T.alloc_shared((64,), "uint32") + T.copy(A, work) + T.copy(I, index) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + offsets = T.vmi.vload(index[0], size=64) + v = T.vmi.vgather(work[64], offsets, mask) + T.vmi.vstore(v, work[192], mask) + T.copy(work, C) + + return main + + +def _scatter_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), I: T.Tensor((64,), "uint32"), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + index = T.alloc_shared((64,), "uint32") + T.copy(A, work) + T.copy(I, index) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + offsets = T.vmi.vload(index[0], size=64) + value = T.vmi.vload(work[128], size=64) + T.vmi.vscatter(value, work[64], offsets, mask) + T.copy(work, C) + + return main + + +def test_gather_descriptor_preserves_nonzero_base(): + artifact = _lower_pto(_gather_kernel()) + _, call = _single(artifact, "tl.vmi.vgather") + ptr = call.args[0] + assert _const(ptr.args[2]) == 64 + # Lanes extent: generic index footprints are explicitly out of scope. + assert _const(ptr.args[3]) == 64 + assert _const(call.args[1]) == 0 + line = next(ln.strip() for ln in artifact.kernel_source.splitlines() if "vmi.vgather" in ln) + assert "addptr" in line and ", 64)" in line, line + + +def test_scatter_descriptor_preserves_nonzero_base(): + artifact = _lower_pto(_scatter_kernel()) + _, call = _single(artifact, "tl.vmi.vscatter") + assert _const(call.args[1].args[2]) == 64 + assert _const(call.args[2]) == 0 + line = next(ln.strip() for ln in artifact.kernel_source.splitlines() if "vmi.vscatter" in ln) + assert "addptr" in line and ", 64)" in line, line + + +# --------------------------------------------------------------------------- +# Removed VMI interface parameters +# --------------------------------------------------------------------------- +def _vmi_body_kernel(body): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((256,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + body(work, mask) + T.copy(work, C) + + return main + + +@pytest.mark.parametrize( + ("case_id", "body", "message"), + [ + ( + "vload-repeat-stride", + lambda work, mask: T.vmi.vload(work[0], size=64, repeat_stride=0), + "unexpected keyword argument 'repeat_stride'", + ), + ( + "vstore-repeat-stride", + lambda work, mask: T.vmi.vstore(T.vmi.vload(work[0], size=64), work[0], mask, repeat_stride=0), + "unexpected keyword argument 'repeat_stride'", + ), + ( + "vload-to-dtype", + lambda work, mask: T.vmi.vload(work[0], size=64, to_dtype=F16), + "unexpected keyword argument 'to_dtype'", + ), + ( + "vload-dist-mode", + lambda work, mask: T.vmi.vload(work[0], size=64, dist_mode="unpack"), + "unexpected keyword argument 'dist_mode'", + ), + ( + "vstore-pmode", + lambda work, mask: T.vmi.vstore(T.vmi.vload(work[0], size=64), work[0], mask, pmode="zero"), + "unexpected keyword argument 'pmode'", + ), + ], +) +def test_removed_vmi_parameters_fail_in_legal_scope(case_id, body, message): + with pytest.raises(TypeError, match=message): + _vmi_body_kernel(body) diff --git a/testing/ascend/language/pto/test_pto_vmi_memory_access_run.py b/testing/ascend/language/pto/test_pto_vmi_memory_access_run.py new file mode 100644 index 0000000000..fb0cdaa12d --- /dev/null +++ b/testing/ascend/language/pto/test_pto_vmi_memory_access_run.py @@ -0,0 +1,511 @@ +"""Device acceptance for VMI memory-access descriptors. + +Every kernel runs through the real ``tilelang.compile(target="pto")`` pipeline +with the automatic scheduler enabled (the generated PTODSL must contain the +MTE sync flags; no test disables scheduling). Inputs are built on CPU with +torch, copied to the NPU, and compared against an independent CPU reference +after ``torch.npu.synchronize()``. + +Critical dependency paths execute 20 fresh trials each with independently +prepared inputs and a fresh output tensor, and also assert the +generated-source order and MTE sync pair: + +* a DMA that writes only the non-zero ``[64,128)`` region feeding + ``vload(work[64], size=64)``; +* ``vstore(..., work[64])`` feeding a DMA that reads only ``[64,128)``. + +The remaining cases are one device-numerics representative per access mode: +group dynamic stride (positive/zero/negative), block_stride, FP4 raw packed +bytes (buffer and explicit-pointer forms), a packed (PK) narrowing store, an +explicit ``access_ptr`` plus extra offset, and gather/scatter with non-zero +bases (direct and explicit-pointer forms). The device numerical reference is +exact for copy-like kernels and an independent CPU computation (torch cast +reference) for narrowing. + +Non-VMI lane counts such as ``size=80`` cannot be constructed by the PTODSL +VMI API; their descriptor/backend evidence lives in the structural suite +(``test_pto_vmi_memory_access.py``) and is not claimed as device-passed. +""" + +import os +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import torch_npu # noqa: E402,F401 (must import before tilelang/ptodsl) +import tilelang # noqa: E402 +import tilelang.ascend.language as T # noqa: E402 +from tvm import tirx # noqa: E402 + +F16 = "float16" +F32 = "float32" +FP4 = "float4_e2m1fn" + + +# --------------------------------------------------------------------------- +# Runtime helpers +# --------------------------------------------------------------------------- +def _require_npu_runtime(): + if not hasattr(torch, "npu") or not torch.npu.is_available(): + pytest.skip("torch.npu is unavailable in the current environment") + try: + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source) + except Exception as err: # noqa: BLE001 + pytest.skip(f"NPU runtime is unavailable in the current environment: {err}") + + +@pytest.fixture(scope="module") +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +def _compile(func): + kernel = tilelang.compile(func, target="pto", out_idx=-1) + source = kernel.get_kernel_source() + assert "pto.set_flag(" in source and "pto.wait_flag(" in source, ( + "automatic scheduling must stay enabled (no MTE sync in the generated source):\n" + source + ) + return kernel, source + + +def _vpto_text(kernel_source): + """Compile the generated PTODSL through PTOAS and return the VPTO text.""" + import importlib.util + import re + import subprocess + import tempfile + import uuid + + decorator = re.search(r'@pto\.jit\(name="([^"]+)"', kernel_source) + name = decorator.group(1) if decorator is not None else "main_kernel" + work = tempfile.mkdtemp(prefix="vmi-runtime-vpto-", dir=os.environ.get("TMPDIR") or None) + source_path = os.path.join(work, "kernel.ptodsl.py") + pto_path = os.path.join(work, "kernel.pto") + vpto_path = os.path.join(work, "kernel.vpto.mlir") + with open(source_path, "w", encoding="utf-8") as file: + file.write(kernel_source) + spec = importlib.util.spec_from_file_location("vmi_runtime_" + uuid.uuid4().hex, source_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + pto_text = getattr(module, name).compile().mlir_text() + with open(pto_path, "w", encoding="utf-8") as file: + file.write(pto_text) + result = subprocess.run( + ["ptoas", "--pto-arch=a5", "--pto-backend=vpto", "--emit-vpto", pto_path, "-o", vpto_path], + text=True, + capture_output=True, + ) + assert result.returncode == 0, f"PTOAS rejected the generated kernel:\n{result.stderr}" + with open(vpto_path, encoding="utf-8") as file: + return file.read() + + +def _assert_addptr_pointer_use(source, definition_fragment, call_fragment): + lines = [line.strip() for line in source.splitlines()] + definitions = [line for line in lines if "= pto.addptr(" in line and definition_fragment in line] + assert len(definitions) == 1, f"expected one pointer definition with {definition_fragment!r}: {definitions}" + name = definitions[0].split("=")[0].strip() + assert call_fragment.format(name=name) in source, source + + +def _to_device(value, device): + return value.to(device) if isinstance(value, torch.Tensor) else value + + +def _run_trials(kernel, make_case, device, trials): + """Run *trials* independent cases; return the per-trial outputs.""" + outputs = [] + for trial in range(trials): + inputs, expected = make_case(trial) + device_inputs = [_to_device(value, device) for value in inputs] + output = kernel(*device_inputs) + torch.npu.synchronize() + got = output.cpu() + assert got.dtype == expected.dtype, f"trial {trial}: {got.dtype} != {expected.dtype}" + assert torch.equal(got, expected), f"trial {trial}: numeric mismatch\n got={got.tolist()[:16]}\n exp={expected.tolist()[:16]}" + outputs.append(got) + return outputs + + +def _line_index(lines, predicate, start=0, what="line"): + for index in range(start, len(lines)): + if predicate(lines[index]): + return index + raise AssertionError(f"missing {what} from line {start}:\n" + "\n".join(lines)) + + +def _assert_mte2_to_v_then_load(lines, dma_predicate, descriptor_needle): + dma = _line_index(lines, dma_predicate, what="producer MTE2 DMA") + set_index = _line_index(lines, lambda line: 'set_flag("MTE2", "V"' in line, dma, "MTE2->V set_flag") + wait_index = _line_index(lines, lambda line: 'wait_flag("MTE2", "V"' in line, set_index, "MTE2->V wait_flag") + vload = _line_index(lines, lambda line: "vmi.vload" in line, wait_index, "vmi.vload") + assert dma < set_index < wait_index < vload + assert descriptor_needle in lines[vload], lines[vload] + + +def _assert_store_then_v_to_mte3(lines, dma_predicate, store_needle): + store = _line_index(lines, lambda line: "vmi.vstore" in line, what="vmi.vstore") + set_index = _line_index(lines, lambda line: 'set_flag("V", "MTE3"' in line, store, "V->MTE3 set_flag") + wait_index = _line_index(lines, lambda line: 'wait_flag("V", "MTE3"' in line, set_index, "V->MTE3 wait_flag") + dma = _line_index(lines, dma_predicate, wait_index, "consumer MTE3 DMA") + assert store < set_index < wait_index < dma + assert store_needle in lines[store], lines[store] + + +def _float_inputs(shape, trial, base=0.0, step=1.0): + generator = torch.Generator().manual_seed(1000 + trial) + values = torch.randn(shape, generator=generator) + return values * step + (base + trial * 37.0) + + +# --------------------------------------------------------------------------- +# Critical path 1: DMA writing only [64,128) -> vload(work[64]) +# --------------------------------------------------------------------------- +def _dma_load_kernel(): + @T.prim_func + def main(A: T.Tensor((64,), F32), S: T.Tensor((64,), F32), C: T.Tensor((64,), F32)): + with T.Kernel(1): + work = T.alloc_shared((128,), F32) + T.copy(S, work[64:128]) # sentinel first + T.copy(A, work[64:128]) # producer DMA: only [64,128) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + v = T.vmi.vload(work[64], size=64) + T.vmi.vstore(v, work[0], mask) + T.copy(work[0:64], C) + + return main + + +def test_dma_to_nonzero_load_dependency_20_trials(npu): + kernel, source = _compile(_dma_load_kernel()) + lines = source.splitlines() + _assert_mte2_to_v_then_load( + lines, + lambda line: "mte_gm_ub" in line and "addptr" in line and ", 64)" in line, + "), 0, size=64)", + ) + + def make_case(trial): + a = _float_inputs((64,), trial) + sentinel = torch.full((64,), -777.0 - trial) + return [a, sentinel], a + + _run_trials(kernel, make_case, npu, trials=20) + + +# --------------------------------------------------------------------------- +# Critical path 2: vstore(..., work[64]) -> DMA reading only [64,128) +# --------------------------------------------------------------------------- +def _store_dma_kernel(): + @T.prim_func + def main(A: T.Tensor((64,), F32), S: T.Tensor((64,), F32), C: T.Tensor((64,), F32)): + with T.Kernel(1): + work = T.alloc_shared((128,), F32) + T.copy(S, work[64:128]) # sentinel before the store + T.copy(A, work[0:64]) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + v = T.vmi.vload(work[0], size=64) + T.vmi.vstore(v, work[64], mask) + T.copy(work[64:128], C) # consumer DMA: only [64,128) + + return main + + +def test_nonzero_store_to_dma_dependency_20_trials(npu): + kernel, source = _compile(_store_dma_kernel()) + lines = source.splitlines() + _assert_store_then_v_to_mte3( + lines, + lambda line: "mte_ub_gm" in line and "addptr" in line and ", 64)" in line, + "), 0, mask)", + ) + + def make_case(trial): + a = _float_inputs((64,), trial) + sentinel = torch.full((64,), -888.0 - trial) + return [a, sentinel], a + + _run_trials(kernel, make_case, npu, trials=20) + + +# --------------------------------------------------------------------------- +# Group dynamic stride (positive / zero / negative) +# --------------------------------------------------------------------------- +def _dynamic_stride_kernel(): + @T.prim_func + def main(A: T.Tensor((512,), F32), stride: T.int32, C: T.Tensor((128,), F32)): + with T.Kernel(1): + work = T.alloc_shared((512,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + v = T.vmi.vload(work[192], size=128, group=2, stride=stride) + T.vmi.vstore(v, work[0], mask) + T.copy(work[0:128], C) + + return main + + +@pytest.mark.parametrize("stride", [128, 0, -128], ids=["positive", "zero", "negative"]) +def test_group_dynamic_stride_load(npu, stride): + kernel, _ = _compile(_dynamic_stride_kernel()) + + def make_case(trial): + a = _float_inputs((512,), trial) + second = 192 + stride + return [a, stride], torch.cat([a[192:256], a[second : second + 64]]) + + _run_trials(kernel, make_case, npu, trials=5) + + +# --------------------------------------------------------------------------- +# Block stride: stride 1/2, single and multi VL +# --------------------------------------------------------------------------- +def _block_expected(a, base, lanes, block_stride, vector=64, block=8): + chunks = (lanes + vector - 1) // vector + values = [] + for chunk in range(chunks): + for slot in range(vector // block): + start = base + chunk * (vector // block) * block_stride * block + slot * block_stride * block + values.extend(a[start : start + block].tolist()) + return torch.tensor(values[:lanes], dtype=a.dtype) + + +def _block_load_kernel(size, block_stride): + @T.prim_func + def main(A: T.Tensor((768,), F32), C: T.Tensor((size,), F32)): + with T.Kernel(1): + work = T.alloc_shared((768,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(size, size=size) + v = T.vmi.vload(work[64], size=size, block_stride=block_stride) + T.vmi.vstore(v, work[512], mask) + T.copy(work[512 : 512 + size], C) + + return main + + +@pytest.mark.parametrize( + ("size", "block_stride"), + [(64, 1), (64, 2), (128, 2)], + ids=["single-vl-s1", "single-vl-s2", "multi-vl-s2"], +) +def test_block_stride_load(npu, size, block_stride): + kernel, _ = _compile(_block_load_kernel(size, block_stride)) + + def make_case(trial): + a = _float_inputs((768,), trial) + return [a], _block_expected(a, 64, size, block_stride) + + _run_trials(kernel, make_case, npu, trials=3) + + +# --------------------------------------------------------------------------- +# FP4 raw packed bytes: non-zero logical start and explicit pointer offset +# --------------------------------------------------------------------------- +def _fp4_view_kernel(explicit_pointer): + @T.prim_func + def main(B: T.Tensor((128,), "uint8"), Z: T.Tensor((128,), "uint8"), C: T.Tensor((128,), "uint8")): + with T.Kernel(1): + raw_in = T.alloc_shared((128,), "uint8") + raw_out = T.alloc_shared((128,), "uint8") + T.copy(B, raw_in) + T.copy(Z, raw_out) + with T.SimdVF(): + f4_in = tirx.decl_buffer((256,), FP4, data=raw_in.data, scope="shared.dyn") + f4_out = tirx.decl_buffer((256,), FP4, data=raw_out.data, scope="shared.dyn") + mask = T.vmi.create_mask(64, size=64) + if explicit_pointer: + ptr = T.access_ptr(f4_in[0], "r", extent=256) + v = T.vmi.vload(ptr, 64, size=128) + else: + v = T.vmi.vload(f4_in[64], size=128) + T.vmi.vstore(v, f4_out[128], mask) + T.copy(raw_out, C) + + return main + + +@pytest.mark.parametrize("explicit_pointer", [False, True], ids=["bufferload", "ptr-offset"]) +def test_fp4_nonzero_logical_start_raw_bytes(npu, explicit_pointer): + kernel, source = _compile(_fp4_view_kernel(explicit_pointer)) + if explicit_pointer: + # Logical extra offset 64 is packed once to 32 bytes. The pointer + # variable name is not assumed: resolve the actual definition/use. + import re + + match = re.search(r"v = pto\.vmi\.vload\((\w+), 32, size=\(128 // 2\)\)", source) + assert match is not None, source + pointer_name = match.group(1) + assert re.search(rf"\b{pointer_name} = pto\.(?:addptr|castptr)\(", source), source + else: + assert "addptr" in source and "), 32)" in source, source + + def make_case(trial): + generator = torch.Generator().manual_seed(2000 + trial) + byte_values = torch.randint(0, 256, (128,), generator=generator, dtype=torch.int64).to(torch.uint8) + zeros = torch.zeros(128, dtype=torch.uint8) + expected = zeros.clone() + expected[64:128] = byte_values[32:96] + return [byte_values, zeros], expected + + _run_trials(kernel, make_case, npu, trials=3) + + +# --------------------------------------------------------------------------- +# Packed (PK) store: logical target extent and exact float16 reference +# --------------------------------------------------------------------------- +def _pk_b32_kernel(): + @T.prim_func + def main(A: T.Tensor((128,), F32), S: T.Tensor((128,), F16), C: T.Tensor((128,), F16)): + with T.Kernel(1): + src = T.alloc_shared((128,), F32) + out = T.alloc_shared((192,), F16) + T.copy(A, src) + T.copy(S, out[64:192]) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + v = T.vmi.vload(src[0], size=128) + narrowed = T.vmi.vcvt(v, to_dtype=F16, rounding="R", saturate="SAT") + T.vmi.vstore(narrowed, out[64], mask) + T.copy(out[64:192], C) + + return main + + +def test_pk_b32_f16_store(npu): + kernel, source = _compile(_pk_b32_kernel()) + vpto = _vpto_text(source) + assert vpto.count('dist = "PK_B32"') == 2, vpto + + def make_case(trial): + a = ((torch.arange(128, dtype=torch.float32) - 64.0) * 0.25) + trial + expected = a.to(torch.float16) + sentinel = torch.full((128,), -3.0, dtype=torch.float16) + return [a, sentinel], expected + + _run_trials(kernel, make_case, npu, trials=3) + + +# --------------------------------------------------------------------------- +# Explicit pointer +# --------------------------------------------------------------------------- +def _explicit_pointer_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), F32), C: T.Tensor((64,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + ptr = T.access_ptr(work[32], "r", extent=128) + v = T.vmi.vload(ptr, 16, size=64) + T.vmi.vstore(v, work[128], mask) + T.copy(work[128:192], C) + + return main + + +def test_explicit_pointer_base32_extra16_reads_48(npu): + kernel, source = _compile(_explicit_pointer_kernel()) + _assert_addptr_pointer_use(source, ", 32)", "vmi.vload({name}, 16, size=64)") + + def make_case(trial): + a = _float_inputs((256,), trial) + return [a], a[48:112] + + _run_trials(kernel, make_case, npu, trials=5) + + +# --------------------------------------------------------------------------- +# Gather / scatter with non-zero direct and explicit-pointer bases +# --------------------------------------------------------------------------- +def _gather_kernel(explicit_pointer): + @T.prim_func + def main(A: T.Tensor((256,), F32), I: T.Tensor((64,), "uint32"), C: T.Tensor((64,), F32)): + with T.Kernel(1): + work = T.alloc_shared((256,), F32) + idx = T.alloc_shared((64,), "uint32") + out = T.alloc_shared((64,), F32) + T.copy(A, work) + T.copy(I, idx) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + offsets = T.vmi.vload(idx[0], size=64) + if explicit_pointer: + ptr = T.access_ptr(work[64], "r", extent=128) + v = T.vmi.vgather(ptr, offsets, mask) + else: + v = T.vmi.vgather(work[64], offsets, mask) + T.vmi.vstore(v, out[0], mask) + T.copy(out, C) + + return main + + +@pytest.mark.parametrize("explicit_pointer", [False, True], ids=["direct", "ptr"]) +def test_gather_nonzero_base(npu, explicit_pointer): + kernel, _ = _compile(_gather_kernel(explicit_pointer)) + + def make_case(trial): + a = _float_inputs((256,), trial) + generator = torch.Generator().manual_seed(3000 + trial) + offsets = torch.randperm(128, generator=generator)[:64].to(torch.uint32) + expected = a[64 + offsets.to(torch.int64)] + return [a, offsets], expected + + _run_trials(kernel, make_case, npu, trials=3) + + +def _scatter_kernel(explicit_pointer): + @T.prim_func + def main(A: T.Tensor((64,), F32), S: T.Tensor((256,), F32), I: T.Tensor((64,), "uint32"), C: T.Tensor((256,), F32)): + with T.Kernel(1): + src = T.alloc_shared((64,), F32) + work = T.alloc_shared((256,), F32) + idx = T.alloc_shared((64,), "uint32") + T.copy(A, src) + T.copy(S, work) + T.copy(I, idx) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + offsets = T.vmi.vload(idx[0], size=64) + v = T.vmi.vload(src[0], size=64) + if explicit_pointer: + ptr = T.access_ptr(work[64], "w", extent=128) + T.vmi.vscatter(v, ptr, offsets, mask) + else: + T.vmi.vscatter(v, work[64], offsets, mask) + T.copy(work, C) + + return main + + +@pytest.mark.parametrize("explicit_pointer", [False, True], ids=["direct", "ptr"]) +def test_scatter_nonzero_base(npu, explicit_pointer): + kernel, _ = _compile(_scatter_kernel(explicit_pointer)) + + def make_case(trial): + a = _float_inputs((64,), trial) + sentinel = torch.full((256,), -123.0 - trial) + generator = torch.Generator().manual_seed(4000 + trial) + offsets = torch.randperm(128, generator=generator)[:64].to(torch.uint32) + expected = sentinel.clone() + expected[64 + offsets.to(torch.int64)] = a + return [a, sentinel, offsets], expected + + _run_trials(kernel, make_case, npu, trials=3) diff --git a/testing/ascend/language/pto/test_pto_vmi_value_provenance.py b/testing/ascend/language/pto/test_pto_vmi_value_provenance.py new file mode 100644 index 0000000000..d4512823d6 --- /dev/null +++ b/testing/ascend/language/pto/test_pto_vmi_value_provenance.py @@ -0,0 +1,854 @@ +"""PTO VMI value provenance and scalar-only consumer boundaries. + +The codegen records every ``Bind``/``Let`` definition and classifies value +origins (scalar, VMI data/mask/pair/void, handle, other vector, unknown) with a +pure query that never prints, emits, or mutates the persistent definition +table. These tests pin: + +* lane-1 VMI data stays VMI data through aliases and ``vbrc(group=1)``: a VMI + data source is passed to ``vbrc`` directly, without the signless integer + bridge (no ``pto.cast``/``vinterpret_cast`` wrapper); +* every registered VMI producer has an explicit return contract + (data/mask/pair/void, with ``pair_get`` extracting lane-1 data from a pair); +* unknown VMI origins stay unknown through aliases and are diagnosed at + scalar consumers instead of ever defaulting to scalar; +* true scalar integer bridges (literal/dynamic, width/signedness) are + preserved; +* lexical ``Let`` scoping, same-name/different-Var identity, per-function + definition-table reset, definition-cycle detection, and linear-time + classification of deep scalar expressions; +* grouped ``vbrc`` accepts numeric local.var reads without changing their + classification. + +Direct codegen (``target.build.tilelang_pto_without_compile`` + +``inspect_source``) is used where the standard lowering pipeline would inline +the alias/``Let`` structure under test. Fatal diagnostics surface as +``tvm.error.InternalError``. Backend compilation evidence is included; +device numerics live in ``test_pto_vmi_value_provenance_run.py``. + +The "unknown VMI op" contract is exercised with a uniquely named test-only op +registered through ``ir.RegisterOp``. It is never added to the product +contract table; the codegen must diagnose its result as unknown. +""" + +from __future__ import annotations + +import re + +import pytest +import torch # noqa: F401 (torch/torch_npu must import before tilelang/ptodsl) +import torch_npu # noqa: F401 + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend.target import normalize_pto_target +from tilelang.engine.lower import lower as _lower + +tirx = tvm.tirx +_Op = tvm.ir.Op.get +_BUILD = tvm.ffi.get_global_func("target.build.tilelang_pto_without_compile") +_INTERNAL_ERROR = tvm.error.InternalError + +VMI_DATA_OPS = [ + "vci", + "vbrc", + "vzip", + "vadd", + "vsub", + "vmul", + "vdiv", + "vmax", + "vmin", + "vand", + "vor", + "vxor", + "vshl", + "vshr", + "vabs", + "vneg", + "vrelu", + "vexp", + "vln", + "vsqrt", + "vnot", + "vadds", + "vmuls", + "vmaxs", + "vmins", + "vshls", + "vshrs", + "vsel", + "vselr", + "vcadd", + "vcmax", + "vcmin", + "vcvt", + "vinterpret_cast", + "vexpdif", + "vaxpy", + "vlrelu", + "vprelu", + "vmula", + "vdhist", + "vchist", + "vgather", + "vgatherb", +] +VMI_MASK_OPS = ["create_mask", "vcmp", "vcmps"] +VMI_PAIR_OPS = ["vintlv", "vdintlv", "vunzip", "vmull"] +VMI_VOID_OPS = ["vstore", "vscatter"] +UNKNOWN_VMI_OP_NAME = "tl.vmi.provenance_test_unknown_result" + + +# --------------------------------------------------------------------------- +# Lowering / direct codegen helpers +# --------------------------------------------------------------------------- +def _target(): + return normalize_pto_target("pto") + + +def _lowered(func): + target = _target() + with target: + return _lower(func, target=target) + + +def _kernel_source(func): + return _lowered(func).kernel_source + + +def _compile(func): + return tilelang.compile(func, target="pto", out_idx=-1) + + +def _source_of(module): + target = _target() + with target: + return _BUILD(module, target).inspect_source() + + +def _raw_source(body, params=()): + """Codegen a raw PrimFunc without the TileLang lowering pipeline.""" + func = tirx.PrimFunc(list(params), body).with_attr("global_symbol", "main") + return _source_of(tvm.IRModule({"main": func})) + + +def _codegen_device_fn(artifact, fn, global_symbol="main"): + gvar = tvm.ir.GlobalVar(global_symbol) + fn = fn.with_attr("global_symbol", global_symbol) + module = tvm.IRModule({gvar: fn}) + target = _target() + with target: + module = tilelang.transform.LowerIntrin()(module) + module = tilelang.transform.HoistBroadcastValues()(module) + return _source_of(module) + + +def _lines(source, needle): + return [line.strip() for line in source.splitlines() if needle in line] + + +def _one(source, needle): + hits = _lines(source, needle) + assert len(hits) == 1, f"expected one line containing {needle!r}, got {hits!r}" + return hits[0] + + +def _expr(line): + """Strip an optional ``name = `` SSA assignment prefix from a source line.""" + match = re.match(r"^[A-Za-z_]\w* = (.*)$", line) + return match.group(1) if match else line + + +def _expect_error(pattern, build): + with pytest.raises(_INTERNAL_ERROR, match=re.escape(pattern)): + build() + + +_UNKNOWN_VMI_OP = None + + +def _unknown_vmi_op(): + """Register the test-only VMI op once with no provenance contract. + + The product table is deliberately not extended; this op only proves that a + ``tl.vmi.*`` result without an explicit return contract stays unknown. The + registration entry point and attribute write fail loudly if unavailable. + """ + global _UNKNOWN_VMI_OP + if _UNKNOWN_VMI_OP is None: + register = tvm.ffi.get_global_func("ir.RegisterOp") + register(UNKNOWN_VMI_OP_NAME, "Test-only VMI op without a contract.") + op = tvm.ir.Op.get(UNKNOWN_VMI_OP_NAME) + # Pure call effect, as required by the codegen's purity check. + op.set_attr("TCallEffectKind", 0) + _UNKNOWN_VMI_OP = op + return _UNKNOWN_VMI_OP + + +# --------------------------------------------------------------------------- +# Real PTO kernels (standard lowering + codegen or full compile) +# --------------------------------------------------------------------------- +def _reduce_broadcast(dtype, reduce_name): + reduce = getattr(T.vmi, reduce_name) + + @T.prim_func + def main(A: T.Tensor((128,), dtype), C: T.Tensor((128,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((256,), dtype) + out = T.alloc_shared((128,), dtype) + T.copy(A, work[0:128]) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + r = reduce(x, mask) + y = T.vmi.vbrc(r, size=128, group=1) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def _scalar_bridge(dtype, literal): + @T.prim_func + def main(A: T.Tensor((128,), dtype), scale: T.int32, C: T.Tensor((128,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((128,), dtype) + around = T.alloc_shared((128,), dtype) + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + base = T.int32(7) if literal else scale + y = T.vmi.vbrc(base, size=128) + z = T.vmi.vci(base, size=128) + T.vmi.vstore(T.vmi.vadd(y, z, mask), around[0], mask) + T.copy(around, C) + + return main + + +def _reinterpret_kernel(): + @T.prim_func + def main(A: T.Tensor((128,), "uint32"), C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "uint32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + y = T.reinterpret("int32x128", x) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def _vci_from_vmi(): + @T.prim_func + def main(A: T.Tensor((128,), "int32"), C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "int32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + r = T.vmi.vcmax(x, mask) + y = T.vmi.vci(r, size=128) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def _scalar_load_after_vmi_store(): + @T.prim_func + def main(A: T.Tensor((128,), "int32"), C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "int32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + T.vmi.vstore(x, work[0], mask) + s = work[0] + y = T.vmi.vbrc(s, size=128) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def _local_var_vmi_store(value_is_vmi): + @T.prim_func + def main(A: T.Tensor((128,), "int32"), C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "int32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + slot = T.alloc_var("int32") + value = T.vmi.vcmax(x, mask) if value_is_vmi else work[0] + slot = value + T.vmi.vstore(T.vmi.vbrc(slot, size=128), out[0], mask) + T.copy(out, C) + + return main + + +def _mutable_register(dtype, lanes, group): + @T.prim_func + def main(A: T.Tensor((128,), dtype), C: T.Tensor((128,), dtype)): + with T.Kernel(1): + src = T.alloc_shared((128,), dtype) + dst = T.alloc_shared((128,), dtype) + T.copy(A, src) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + slot = T.vmi.alloc_var(dtype, size=lanes) + slot = T.vmi.vload(src[0], size=lanes) + value = T.vmi.vbrc(slot, size=128, group=group) + T.vmi.vstore(value, dst[0], mask) + T.copy(dst, C) + + return main + + +# --------------------------------------------------------------------------- +# Raw producer/consumer helpers (direct codegen contracts) +# --------------------------------------------------------------------------- +def _vbrc(arg, group=1): + annotations = {"size": 128} + if group is not None: + annotations["group"] = group + return tirx.call_intrin("int32x128", _Op("tl.vmi.vbrc"), arg, annotations=annotations) + + +def _vmi_data_producer(op_name): + op = _Op(f"tl.vmi.{op_name}") + if op_name in ("vci", "vbrc"): + return tirx.call_intrin("int32", op, tirx.IntImm("int32", 0)), () + if op_name in ("vgather", "vgatherb"): + ptr = tirx.Var("ptr", "handle") + offsets = tirx.Var("offsets", "int32x128") + mask = tirx.Var("mask", "boolx128") + call = tirx.call_intrin("int32", op, ptr, tirx.IntImm("int32", 0), offsets, mask) + return call, (ptr, offsets, mask) + return tirx.call_intrin("int32", op), () + + +def _device_fn(func): + artifact = _lowered(func) + assert len(artifact.device_mod.functions) == 1 + gvar = next(iter(artifact.device_mod.functions.keys())) + return artifact, gvar, artifact.device_mod.functions[gvar] + + +def _find_vbrc_call(fn): + found = [] + + def visit(node): + if isinstance(node, tirx.Call) and getattr(node.op, "name", "") == "tl.vmi.vbrc": + found.append(node) + + tirx.stmt_functor.post_order_visit(fn.body, visit) + assert len(found) == 1, f"expected one vbrc call, got {len(found)}" + return found[0] + + +def _find_binding(fn, var): + found = [] + + def visit(node): + if isinstance(node, tirx.Bind) and node.var.same_as(var): + found.append(node) + + tirx.stmt_functor.post_order_visit(fn.body, visit) + assert len(found) == 1, f"expected one binding for {var}, got {len(found)}" + return found[0] + + +def _transform(fn, preorder): + body = tirx.stmt_functor.ir_transform(fn.body, preorder, None) + return fn.with_body(body) + + +def _replace_vbrc_source(fn, make_source, drop_producer=False): + """Replace the vbrc source with ``make_source(producer_bind, source_var)``.""" + call = _find_vbrc_call(fn) + source_var = call.args[0] + producer_bind = _find_binding(fn, source_var) + + def preorder(node): + if isinstance(node, tirx.Bind): + if node.var.same_as(source_var): + if drop_producer: + return tirx.Evaluate(tirx.IntImm("int32", 0)) + return node + if ( + isinstance(node.value, tirx.Call) + and getattr(node.value.op, "name", "") == "tl.vmi.vbrc" + and node.value.args[0].same_as(source_var) + ): + value = node.value + new_call = tirx.Call( + value.dtype, + value.op, + [make_source(producer_bind, source_var)] + list(value.args[1:]), + value.annotations, + value.span, + ) + return tirx.Bind(node.var, new_call) + return None + + return _transform(fn, preorder) + + +def _alias_chain_source(fn): + call = _find_vbrc_call(fn) + source_var = call.args[0] + alias_a = tirx.Var("alias_a", source_var.dtype) + alias_b = tirx.Var("alias_b", source_var.dtype) + + def preorder(node): + if isinstance(node, tirx.Bind): + if node.var.same_as(source_var): + return tirx.SeqStmt([node, tirx.Bind(alias_a, node.var), tirx.Bind(alias_b, alias_a)]) + if ( + isinstance(node.value, tirx.Call) + and getattr(node.value.op, "name", "") == "tl.vmi.vbrc" + and node.value.args[0].same_as(source_var) + ): + value = node.value + new_call = tirx.Call( + value.dtype, + value.op, + [alias_b] + list(value.args[1:]), + value.annotations, + value.span, + ) + return tirx.Bind(node.var, new_call) + return None + + return _transform(fn, preorder) + + +def _let_source(fn, nested): + def make_source(producer_bind, source_var): + let_var = tirx.Var("let_t", source_var.dtype) + body = tirx.Let(let_var, producer_bind.value, let_var) + if nested: + inner_var = tirx.Var("let_u", source_var.dtype) + body = tirx.Let(let_var, producer_bind.value, tirx.Let(inner_var, let_var, inner_var)) + return body + + return _replace_vbrc_source(fn, make_source, drop_producer=True) + + +def _direct_call_source(fn): + return _replace_vbrc_source(fn, lambda producer_bind, source_var: producer_bind.value, drop_producer=True) + + +# --------------------------------------------------------------------------- +# A. Core lane-1 VMI data -> vbrc(group=1), scalar bridges, reinterpret +# --------------------------------------------------------------------------- +def _assert_direct_vbrc(source, producer_name): + producer_line = _one(source, f"pto.vmi.{producer_name}(") + bound = producer_line.split("=", 1)[0].strip() + vbrc_line = _one(source, "pto.vmi.vbrc(") + assert _expr(vbrc_line) == f"pto.vmi.vbrc({bound}, group=1, size=128)" + assert source.count(f"pto.vmi.{producer_name}(") == 1 + + +def test_int32_vcmax_vbrc_uses_vmi_source_directly(): + source = _kernel_source(_reduce_broadcast("int32", "vcmax")) + _assert_direct_vbrc(source, "vcmax") + + +def test_uint32_vcmin_vbrc_uses_vmi_source_directly(): + source = _kernel_source(_reduce_broadcast("uint32", "vcmin")) + _assert_direct_vbrc(source, "vcmin") + + +def test_float32_vcmax_vbrc_control_uses_vmi_source_directly(): + source = _kernel_source(_reduce_broadcast("float32", "vcmax")) + _assert_direct_vbrc(source, "vcmax") + + +@pytest.mark.parametrize( + ("dtype", "reduce_name"), + [("int32", "vcmax"), ("uint32", "vcmin"), ("float32", "vcmax")], +) +def test_reduce_broadcast_compiles_through_backend(dtype, reduce_name): + _compile(_reduce_broadcast(dtype, reduce_name)) + + +def test_scalar_literal_bridge_is_preserved(): + source = _kernel_source(_scalar_bridge("int32", literal=True)) + assert _expr(_one(source, "pto.vmi.vbrc(pto.si32(7), size=128)")) == ("pto.vmi.vbrc(pto.si32(7), size=128)") + assert _expr(_one(source, "pto.vmi.vci(")).startswith("pto.vmi.vinterpret_cast(pto.vmi.vci(pto.i32(7), size=128), to_dtype=pto.si32)") + + +def test_scalar_dynamic_bridge_is_preserved(): + source = _kernel_source(_scalar_bridge("int32", literal=False)) + # Dynamic integer scalar sources take the signless bridge and recover the + # authored signedness on the VMI result. + assert _expr(_one(source, "pto.vmi.vbrc(")) == ( + "pto.vmi.vinterpret_cast(pto.vmi.vbrc(pto.cast(base, pto.i32), size=128), to_dtype=pto.si32)" + ) + assert _expr(_one(source, "pto.vmi.vci(")) == ( + "pto.vmi.vinterpret_cast(pto.vmi.vci(pto.cast(base, pto.i32), size=128), to_dtype=pto.si32)" + ) + + +@pytest.mark.parametrize("literal", [True, False]) +def test_scalar_bridge_compiles_through_backend(literal): + _compile(_scalar_bridge("int32", literal=literal)) + + +def test_vmi_reinterpret_uses_vinterpret_cast(): + source = _kernel_source(_reinterpret_kernel()) + assert "pto.vmi.vinterpret_cast(x, to_dtype=pto.si32)" in _one(source, "pto.vmi.vstore(") + assert "vbitcast" not in source + + +# --------------------------------------------------------------------------- +# B. Grouped-vbrc mutable local.var register compatibility +# --------------------------------------------------------------------------- +def test_mutable_register_grouped_vbrc_is_preserved(): + # Grouped vbrc accepts this numeric local.var read without classifying it + # as VMI data. Its classification stays other-vector, and the source is + # emitted directly without a scalar cast. + source = _kernel_source(_mutable_register("float32", 4, 4)) + vbrc_line = _one(source, "pto.vmi.vbrc(") + assert "pto.vmi.vbrc(slot, group=4, size=128)" in vbrc_line + assert "pto.cast(slot" not in vbrc_line + assert "slot = pto.vmi.vreg(4, pto.f32)" in source + assert "slot = pto.vmi.vload(" in source + + +# --------------------------------------------------------------------------- +# C. Aliases, Let scoping, Var identity, per-function reset, cycles +# --------------------------------------------------------------------------- +def test_multi_level_alias_chain_keeps_vmi_data(): + artifact, _, fn = _device_fn(_reduce_broadcast("int32", "vcmax")) + source = _codegen_device_fn(artifact, _alias_chain_source(fn)) + assert "alias_a = r" in source or re.search(r"alias_a = \w+", source) + assert re.search(r"alias_b = alias_a", source) + assert _expr(_one(source, "pto.vmi.vbrc(")) == ("pto.vmi.vbrc(alias_b, group=1, size=128)") + assert source.count("pto.vmi.vcmax(") == 1 + + +def test_direct_nested_producer_call_keeps_vmi_data(): + artifact, _, fn = _device_fn(_reduce_broadcast("int32", "vcmax")) + source = _codegen_device_fn(artifact, _direct_call_source(fn)) + vbrc_line = _one(source, "pto.vmi.vbrc(") + assert "pto.vmi.vbrc(pto.vmi.vcmax(" in vbrc_line + assert vbrc_line.endswith(", group=1, size=128)") + assert source.count("pto.vmi.vcmax(") == 1 + + +@pytest.mark.parametrize("nested", [False, True]) +def test_let_source_is_query_local_and_producer_runs_once(nested): + artifact, _, fn = _device_fn(_reduce_broadcast("int32", "vcmax")) + source = _codegen_device_fn(artifact, _let_source(fn, nested=nested)) + vbrc_line = _one(source, "pto.vmi.vbrc(") + if nested: + assert "(lambda let_t: (lambda let_u: let_u)(let_t))(pto.vmi.vcmax(" in vbrc_line + else: + assert "(lambda let_t: let_t)(pto.vmi.vcmax(" in vbrc_line + assert vbrc_line.endswith(", group=1, size=128)") + assert "pto.cast(" not in vbrc_line + assert source.count("pto.vmi.vcmax(") == 1 + + +def test_same_named_different_vars_are_not_conflated(): + artifact, _, fn = _device_fn(_reduce_broadcast("int32", "vcmax")) + call = _find_vbrc_call(fn) + source_var = call.args[0] + same_name = tirx.Var(source_var.name, source_var.dtype) + + def preorder(node): + if isinstance(node, tirx.Bind) and isinstance(node.value, tirx.Call) and getattr(node.value.op, "name", "") == "tl.vmi.vbrc": + second = tirx.call_intrin( + "int32x128", + _Op("tl.vmi.vbrc"), + same_name, + annotations={"size": 128}, + ) + return tirx.SeqStmt( + [ + node, + tirx.Bind(same_name, tirx.IntImm("int32", 5)), + tirx.Evaluate(second), + ] + ) + return None + + source = _codegen_device_fn(artifact, _transform(fn, preorder)) + vbrc_lines = _lines(source, "pto.vmi.vbrc(") + assert len(vbrc_lines) == 2 + assert re.search(r"pto\.vmi\.vbrc\(\w+, group=1, size=128\)", vbrc_lines[0]) + assert "group" not in vbrc_lines[1] + # The same-named scalar Var is bridged independently of the VMI data Var, + # proving the table is keyed by Var identity, not name. The printer + # disambiguates the duplicate Python name (r vs r_1). + assert re.search( + rf"pto\.vmi\.vbrc\(pto\.cast\({source_var.name}_\d+, pto\.i32\), size=128\)", + vbrc_lines[1], + ), vbrc_lines[1] + + +def _two_function_reset_module(order): + """One codegen run with the same Var defined in one function and used as a + true scalar parameter in another. Function names select the generation + order (the module map is name-ordered), so both orders are covered without + relying on insertion order.""" + artifact_a, _, fn_a = _device_fn(_reduce_broadcast("int32", "vcmax")) + shared = tirx.Var("shared_reset", "int32") + + call_a = _find_vbrc_call(fn_a) + old_a = call_a.args[0] + + def preorder_a(node): + if isinstance(node, tirx.Bind) and node.var.same_as(old_a): + return tirx.Bind(shared, node.value) + if isinstance(node, tirx.Bind) and isinstance(node.value, tirx.Call): + value = node.value + if getattr(value.op, "name", "") == "tl.vmi.vbrc" and value.args[0].same_as(old_a): + return tirx.Bind( + node.var, + tirx.Call( + value.dtype, + value.op, + [shared] + list(value.args[1:]), + value.annotations, + value.span, + ), + ) + return None + + fn_a = fn_a.with_body(tirx.stmt_functor.ir_transform(fn_a.body, preorder_a, None)) + + # The SAME Var object is an ordinary scalar parameter here: no Bind, no VMI + # producer. A leaked definition table would misclassify it as VMI data. + scalar_vbrc = tirx.call_intrin("int32x128", _Op("tl.vmi.vbrc"), shared, annotations={"size": 128}) + fn_b = tirx.PrimFunc([shared], tirx.Evaluate(scalar_vbrc)) + + if order == "vmi-first": + vmi_name, scalar_name = "a_vmi_fn", "b_scalar_fn" + else: + vmi_name, scalar_name = "b_vmi_fn", "a_scalar_fn" + target = _target() + fn_a = fn_a.with_attr("global_symbol", vmi_name).with_attr("target", target) + fn_b = fn_b.with_attr("global_symbol", scalar_name).with_attr("target", target) + vmi_entry = (tvm.ir.GlobalVar(vmi_name), fn_a) + scalar_entry = (tvm.ir.GlobalVar(scalar_name), fn_b) + entries = [vmi_entry, scalar_entry] + if order == "scalar-first": + entries.reverse() + module = tvm.IRModule(dict(entries)) + with target: + module = tilelang.transform.LowerIntrin()(module) + module = tilelang.transform.HoistBroadcastValues()(module) + return _source_of(module), vmi_name, scalar_name + + +def _function_section(source, name): + start = source.index(f"def {name}(") + end = source.find("def ", start + 1) + return source[start : end if end != -1 else len(source)] + + +def test_same_var_across_two_functions_resets_definition_table(): + observed_orders = [] + for order in ("vmi-first", "scalar-first"): + source, vmi_name, scalar_name = _two_function_reset_module(order) + section_vmi = _function_section(source, vmi_name) + section_scalar = _function_section(source, scalar_name) + assert "pto.vmi.vbrc(shared_reset, group=1, size=128)" in section_vmi + assert "pto.cast(shared_reset" not in section_vmi + assert "pto.vmi.vbrc(pto.cast(shared_reset, pto.i32), size=128)" in section_scalar + assert "pto.vmi.vbrc(shared_reset, group=1" not in section_scalar + observed_orders.append("vmi-first" if source.index(f"def {vmi_name}(") < source.index(f"def {scalar_name}(") else "scalar-first") + assert sorted(observed_orders) == ["scalar-first", "vmi-first"] + + +def test_definition_cycle_is_detected(): + x = tirx.Var("cycle_x", "int32") + body = tirx.SeqStmt( + [ + tirx.Bind(x, tirx.Add(x, tirx.IntImm("int32", 1))), + tirx.Evaluate(_vbrc(x, group=None)), + ] + ) + _expect_error("PTO SSA definition cycle detected", lambda: _raw_source(body)) + + +# --------------------------------------------------------------------------- +# D. Unknown-origin propagation and diagnosis +# --------------------------------------------------------------------------- +def test_registered_unknown_vmi_op_through_alias_stays_unknown(): + unknown = tirx.Var("unknown_x", "int32") + alias = tirx.Var("unknown_alias", "int32") + body = tirx.SeqStmt( + [ + tirx.Bind(unknown, tirx.call_intrin("int32", _unknown_vmi_op())), + tirx.Bind(alias, unknown), + tirx.Evaluate(_vbrc(alias, group=None)), + ] + ) + _expect_error("cannot broadcast unknown (int32, lanes=1)", lambda: _raw_source(body)) + + +def test_unknown_vmi_origin_diagnosed_at_vci(): + x = tirx.Var("x", "int32") + alias = tirx.Var("alias", "int32") + body = tirx.SeqStmt( + [ + tirx.Bind(x, tirx.call_intrin("int32", _Op("tl.vmi.vcmax"))), + tirx.Bind(alias, x), + tirx.Evaluate( + tirx.call_intrin( + "int32x128", + _Op("tl.vmi.vci"), + tirx.Cast("float32", alias), + annotations={"size": 128}, + ) + ), + ] + ) + _expect_error("tl.vmi.vci base expects a scalar; got unknown", lambda: _raw_source(body)) + + +# --------------------------------------------------------------------------- +# E. Registered producer return-contract matrix +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("op_name", VMI_DATA_OPS) +def test_registered_vmi_data_producer_never_enters_scalar_bridge(op_name): + producer, params = _vmi_data_producer(op_name) + if op_name in ("vgather", "vgatherb"): + # The gather printer requires a real access-pointer operand and the + # function-signature printer cannot emit a synthetic handle parameter. + # The return contract is asserted through pair_get's source validation, + # which classifies the gather result before printing it (free variables + # are only rendered by the TIR printer in the diagnostic): it must be + # VMI data, not scalar/unknown. + get = tirx.call_intrin("int32", _Op("tl.vmi.pair_get"), producer, tirx.IntImm("int32", 0)) + _expect_error( + "tl.vmi.pair_get source must be a VMI pair; got VMI data", + lambda: _raw_source(tirx.Evaluate(get)), + ) + return + source = _raw_source(tirx.Evaluate(_vbrc(producer)), params) + vbrc_line = _one(source, "pto.vmi.vbrc(") + assert "pto.cast(" not in vbrc_line + assert vbrc_line.endswith(", group=1, size=128)") + + +def test_vload_lane1_is_vmi_data(): + producer = tirx.call_intrin("int32", _Op("tl.vmi.vload")) + source = _raw_source(tirx.Evaluate(_vbrc(producer))) + assert "pto.cast(" not in _one(source, "pto.vmi.vbrc(") + + +@pytest.mark.parametrize("op_name", VMI_MASK_OPS) +def test_vmi_mask_producer_is_not_broadcast_as_data(op_name): + producer = tirx.call_intrin("boolx128", _Op(f"tl.vmi.{op_name}")) + _expect_error( + "tl.vmi.vbrc with group=... requires a VMI vreg; got VMI mask", + lambda: _raw_source(tirx.Evaluate(_vbrc(producer))), + ) + + +@pytest.mark.parametrize("op_name", VMI_PAIR_OPS) +def test_vmi_pair_producer_is_not_broadcast_as_data(op_name): + producer = tirx.call_intrin("int32x128", _Op(f"tl.vmi.{op_name}")) + _expect_error( + "tl.vmi.vbrc with group=... requires a VMI vreg; got VMI pair", + lambda: _raw_source(tirx.Evaluate(_vbrc(producer))), + ) + + +@pytest.mark.parametrize("op_name", VMI_VOID_OPS) +def test_vmi_void_producer_is_not_broadcast_as_data(op_name): + producer = tirx.call_intrin("void", _Op(f"tl.vmi.{op_name}")) + _expect_error( + "tl.vmi.vbrc with group=... requires a VMI vreg; got void", + lambda: _raw_source(tirx.Evaluate(_vbrc(producer))), + ) + + +def test_pair_get_returns_data_from_pair(): + pair = tirx.call_intrin("int32x128", _Op("tl.vmi.vintlv")) + get = tirx.call_intrin("int32", _Op("tl.vmi.pair_get"), pair, tirx.IntImm("int32", 0)) + source = _raw_source(tirx.Evaluate(_vbrc(get))) + assert _expr(_one(source, "pto.vmi.vbrc(")) == "pto.vmi.vbrc((pto.vmi.vintlv())[0], group=1, size=128)" + + +# --------------------------------------------------------------------------- +# F. Scalar-only consumers +# --------------------------------------------------------------------------- +def test_vci_rejects_vmi_data_source(): + _expect_error( + "tl.vmi.vci base expects a scalar; got VMI data", + lambda: _kernel_source(_vci_from_vmi()), + ) + + +def test_vbrc_group_rejects_true_scalar_source(): + _expect_error( + "tl.vmi.vbrc with group=... requires a VMI vreg; got scalar", + lambda: _raw_source(tirx.Evaluate(_vbrc(tirx.IntImm("int32", 0), group=1))), + ) + + +def test_ordinary_scalar_load_after_vmi_store_stays_scalar(): + source = _kernel_source(_scalar_load_after_vmi_store()) + vbrc_line = _one(source, "pto.vmi.vbrc(") + assert "pto.cast(" in vbrc_line + assert "group" not in vbrc_line + + +def test_local_var_store_rejects_vmi_data(): + _expect_error( + "PTO local.var integer store expects a scalar; got VMI data", + lambda: _kernel_source(_local_var_vmi_store(value_is_vmi=True)), + ) + + +def test_local_var_store_keeps_scalar_value(): + source = _kernel_source(_local_var_vmi_store(value_is_vmi=False)) + assert "slot = pto.const(0, dtype=pto.si32)" in source + assert re.search(r"slot = pto\.cast\(\w+, pto\.si32\)", source) + assert "pto.cast(slot, pto.i32)" in _one(source, "pto.vmi.vbrc(") + + +def test_select_branch_rejects_vmi_data(): + data = tirx.Var("data", "int32") + cond = tirx.Var("cond", "bool") + body = tirx.SeqStmt( + [ + tirx.Bind(data, tirx.call_intrin("int32", _Op("tl.vmi.vcmax"))), + tirx.Evaluate(tirx.Select(cond, data, tirx.IntImm("int32", 0))), + ] + ) + _expect_error("PTO select branch expects a scalar; got VMI data", lambda: _raw_source(body, (cond,))) + + +# --------------------------------------------------------------------------- +# G. Deep scalar expressions stay bounded +# --------------------------------------------------------------------------- +def test_deep_scalar_expression_chain_is_classified_linearly(): + x = tirx.Var("x", "int32") + expr = x + for i in range(200): + expr = tirx.Add(expr, tirx.IntImm("int32", i + 1)) + source = _raw_source(tirx.Evaluate(tirx.call_intrin("int32", _Op("tl.vmi.vci"), expr)), (x,)) + # The scalar bridge proves one classification of the 200-deep chain. + assert "pto.vmi.vinterpret_cast(pto.vmi.vci(pto.cast(" in source diff --git a/testing/ascend/language/pto/test_pto_vmi_value_provenance_run.py b/testing/ascend/language/pto/test_pto_vmi_value_provenance_run.py new file mode 100644 index 0000000000..0dbc47cfed --- /dev/null +++ b/testing/ascend/language/pto/test_pto_vmi_value_provenance_run.py @@ -0,0 +1,246 @@ +"""Device acceptance: PTO VMI value provenance and scalar bridge boundaries. + +Each case compiles through the real TileLang -> PTODSL -> PTOAS VPTO -> +Bisheng pipeline and executes on the NPU. References are computed on the CPU +from Python integers/floats, independently of the generated kernel. + +Coverage: +* int32/uint32 ``vcmax``/``vcmin`` -> ``vbrc(group=1)``: signed negatives, + exact integer boundaries, unsigned high-bit values (> ``INT32_MAX``), and + both reduction directions. 20 independent trials per integer case. +* float32 control: the same reduce/broadcast route for a non-integer VMI + source. +* true scalar integer bridges: compile-time literal and runtime scalar + sources through ``vbrc`` + ``vci``. +* ordinary scalar memory read after a VMI store (storage contract, not the + last ``BufferStore`` RHS). +""" + +import os +import re +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import torch_npu # noqa: E402,F401 (import before tilelang/ptodsl) +import tilelang # noqa: E402 +import tilelang.ascend.language as T # noqa: E402 + + +def _require_npu_runtime(): + # These device tests require an available NPU runtime and a successful + # host-to-device copy probe; a missing runtime or probe mismatch fails. + assert hasattr(torch, "npu"), "torch.npu is unavailable in this environment" + assert torch.npu.is_available(), "torch.npu is unavailable in the current environment" + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source), "NPU host->device copy probe mismatch" + + +@pytest.fixture(scope="module") +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +def _to_device(value, device): + if isinstance(value, torch.Tensor): + return value.to(device) + return value + + +def _compile(func): + kernel = tilelang.compile(func, target="pto", out_idx=-1) + return kernel, kernel.get_kernel_source() + + +def _run_trials(kernel, make_case, device, trials): + outputs = [] + for trial in range(trials): + inputs, expected = make_case(trial) + device_inputs = [_to_device(value, device) for value in inputs] + output = kernel(*device_inputs) + torch.npu.synchronize() + got = output.cpu() + assert got.dtype == expected.dtype, f"trial {trial}: {got.dtype} != {expected.dtype}" + assert torch.equal(got, expected), f"trial {trial}: numeric mismatch\n got={got.tolist()[:16]}\n exp={expected.tolist()[:16]}" + outputs.append(got) + return outputs + + +# --------------------------------------------------------------------------- +# Reduce -> vbrc(group=1): lane-1 VMI data must stay VMI data +# --------------------------------------------------------------------------- +def _reduce_broadcast(dtype, reduce_name): + reduce = getattr(T.vmi, reduce_name) + + @T.prim_func + def main(A: T.Tensor((128,), dtype), C: T.Tensor((128,), dtype)): + with T.Kernel(1): + work = T.alloc_shared((256,), dtype) + out = T.alloc_shared((128,), dtype) + T.copy(A, work[0:128]) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + r = reduce(x, mask) + y = T.vmi.vbrc(r, size=128, group=1) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def _int32_input(trial): + generator = torch.Generator().manual_seed(20261017 + trial) + values = torch.randint(-(2**31), 2**31 - 1, (128,), generator=generator).to(torch.int32) + values[0] = -(2**31) + values[1] = 2**31 - 1 + values[2] = -1 + values[3] = 0 + values[4] = 1 + values[5] = -(2**31) + 1 if trial % 2 else 2**31 - 2 + return values + + +def _uint32_input(trial): + generator = torch.Generator().manual_seed(20261017 + trial) + values = torch.randint(0, 2**31 - 1, (128,), generator=generator).to(torch.uint32) + values[0] = 2**32 - 1 + values[1] = 2**31 # INT32_MAX + 1, high bit set + values[2] = 2**31 - 1 + values[3] = 0 + values[4] = 1 + values[5] = 2**32 - 2 if trial % 2 else 2**31 + 1 + return values + + +def _float32_input(trial): + generator = torch.Generator().manual_seed(20261017 + trial) + values = torch.randn(128, generator=generator) * 1000.0 + values[0] = 3.0e38 + values[1] = -3.0e38 + values[2] = 0.0 + values[3] = -0.0 + return values + + +INPUTS = {"int32": _int32_input, "uint32": _uint32_input} +REDUCE_OPS = {"vcmax": max, "vcmin": min} +INTEGER_REDUCE_CASES = [ + ("int32", "vcmax"), + ("int32", "vcmin"), + ("uint32", "vcmax"), + ("uint32", "vcmin"), +] + + +@pytest.mark.parametrize(("dtype", "reduce_name"), INTEGER_REDUCE_CASES) +def test_integer_reduce_broadcast_is_bit_exact(npu, dtype, reduce_name): + kernel, source = _compile(_reduce_broadcast(dtype, reduce_name)) + vbrc_line = next(line.strip() for line in source.splitlines() if "pto.vmi.vbrc(" in line) + assert "pto.cast(" not in vbrc_line, vbrc_line + assert "group=1" in vbrc_line, vbrc_line + reduce = REDUCE_OPS[reduce_name] + + def make_case(trial): + a = INPUTS[dtype](trial) + expected_value = reduce(int(v) for v in a.tolist()) + expected = torch.full((128,), expected_value, dtype=a.dtype) + return [a], expected + + _run_trials(kernel, make_case, npu, trials=20) + + +def test_float32_reduce_broadcast_control(npu): + kernel, source = _compile(_reduce_broadcast("float32", "vcmax")) + vbrc_line = next(line.strip() for line in source.splitlines() if "pto.vmi.vbrc(" in line) + assert "pto.cast(" not in vbrc_line, vbrc_line + + def make_case(trial): + a = _float32_input(trial) + expected = torch.full((128,), float(a.max().item()), dtype=torch.float32) + return [a], expected + + _run_trials(kernel, make_case, npu, trials=5) + + +# --------------------------------------------------------------------------- +# True scalar integer bridges (literal + runtime scalar) +# --------------------------------------------------------------------------- +def _scalar_bridge(literal): + @T.prim_func + def main(A: T.Tensor((128,), "int32"), scale: T.int32, C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "int32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + base = T.int32(7) if literal else scale + y = T.vmi.vbrc(base, size=128) + z = T.vmi.vci(base, size=128) + T.vmi.vstore(T.vmi.vadd(y, z, mask), out[0], mask) + T.copy(out, C) + + return main + + +@pytest.mark.parametrize("literal", [True, False], ids=["literal", "runtime"]) +def test_scalar_integer_bridge_device(npu, literal): + kernel, source = _compile(_scalar_bridge(literal)) + if literal: + assert "pto.vmi.vbrc(pto.si32(7), size=128)" in source, source + else: + assert "pto.cast(base, pto.i32)" in source, source + + def make_case(trial): + a = torch.arange(128, dtype=torch.int32) + trial + scale = 7 if literal else 5 + trial * 3 + expected = (2 * scale + torch.arange(128, dtype=torch.int32)).to(torch.int32) + return [a, scale], expected + + _run_trials(kernel, make_case, npu, trials=5) + + +# --------------------------------------------------------------------------- +# Ordinary scalar memory read after a VMI store +# --------------------------------------------------------------------------- +def _scalar_load_after_vmi_store(): + @T.prim_func + def main(A: T.Tensor((128,), "int32"), C: T.Tensor((128,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((128,), "int32") + out = T.alloc_shared((128,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(128, size=128) + x = T.vmi.vload(work[0], size=128) + T.vmi.vstore(x, work[0], mask) + scalar = work[0] + y = T.vmi.vbrc(scalar, size=128) + T.vmi.vstore(y, out[0], mask) + T.copy(out, C) + + return main + + +def test_scalar_buffer_load_after_vmi_store_device(npu): + kernel, source = _compile(_scalar_load_after_vmi_store()) + vbrc_line = next(line.strip() for line in source.splitlines() if "pto.vmi.vbrc(" in line) + assert re.search(r"pto\.cast\(scalar_\d+, pto\.i32\)", vbrc_line), vbrc_line + + def make_case(trial): + a = torch.arange(128, dtype=torch.int32) * (trial + 1) - 10 + expected = torch.full((128,), int(a[0]), dtype=torch.int32) + return [a], expected + + _run_trials(kernel, make_case, npu, trials=3) diff --git a/testing/ascend/language/pto/test_simtvf_local_integer_storage.py b/testing/ascend/language/pto/test_simtvf_local_integer_storage.py new file mode 100644 index 0000000000..b26534fb01 --- /dev/null +++ b/testing/ascend/language/pto/test_simtvf_local_integer_storage.py @@ -0,0 +1,233 @@ +"""SIMT local integer buffers use signless LLVM stack storage. + +This file verifies that: + +* scalar-element SIMT local buffers of the admitted integer types + int8/uint8/int32/uint32 allocate with a same-width *signless* PTO element + type (``pto.i8`` / ``pto.i32``); +* the four access sites -- scalar load, scalar store, contiguous + builtin-vector load and contiguous builtin-vector store -- adapt between + the signless backing and TileLang's logical signed/unsigned type via + ``pto.cast``; +* the adapter is scoped to SIMT local storage only: VMI, ``local.var``, + shared and global accesses keep their existing types; +* vector-element local slots use the same-width signless element backing in + the allocation branch; +* the admitted dtype set is unchanged: int16/int64 stay rejected. +""" + +import re + +import pytest +import torch # noqa: F401 (torch/torch_npu before tilelang/ptodsl) +import torch_npu # noqa: F401 + +import tilelang +import tilelang.ascend.language as T +from tilelang import tvm +from tilelang.ascend.target import normalize_pto_target + + +def _lower(func): + """Lower with an explicit PTO target context (required for fragments).""" + target = normalize_pto_target("pto") + with target: + return tilelang.lower(func, target=target) + + +# --------------------------------------------------------------------------- +# SIMT local integer kernel: scalar and contiguous access sites +# --------------------------------------------------------------------------- +def _simt_local_integer_kernel(dtype): + @T.prim_func + def main(A: T.Tensor((8,), dtype), B: T.Tensor((8,), dtype)): + with T.Kernel(1), T.SimtVF(threads=1): + local = T.alloc_local((8,), dtype) + local[0] = A[0] + B[0] = local[0] + local[T.Ramp(1, 1, 4)] = A[T.Ramp(0, 1, 4)] + B[T.Ramp(1, 1, 4)] = local[T.Ramp(1, 1, 4)] + + return main + + +@pytest.mark.parametrize( + ("dtype", "logical_type", "backing_type"), + [ + ("int8", "pto.si8", "pto.i8"), + ("uint8", "pto.ui8", "pto.i8"), + ("int32", "pto.si32", "pto.i32"), + ("uint32", "pto.ui32", "pto.i32"), + ], +) +def test_simtvf_local_integer_scalar_and_contiguous_accesses_use_signless_storage(dtype, logical_type, backing_type): + source = _lower(_simt_local_integer_kernel(dtype)).kernel_source + + # Allocation: same-width signless backing. + assert re.search(rf"pto\.alloc_buffer\(\(8,\), {re.escape(backing_type)}\)", source) + assert "pto.alloc_buffer((8,), pto.si8)" not in source + assert "pto.alloc_buffer((8,), pto.ui8)" not in source + assert "pto.alloc_buffer((8,), pto.si32)" not in source + assert "pto.alloc_buffer((8,), pto.ui32)" not in source + + # Scalar load restores the logical type; scalar store returns to backing. + assert re.search(rf"pto\.cast\(pto\.load\(local, 0\), {re.escape(logical_type)}\)", source) + assert re.search( + rf"pto\.store\(pto\.cast\(pto\.load\(A, 0\), {re.escape(backing_type)}\), local, 0\)", + source, + ) + + # Contiguous builtin-vector load restores; contiguous store returns. + assert re.search( + rf"pto\.cast\(pto\.load\(local, 1, contiguous=4\), {re.escape(logical_type)}\)", + source, + ) + assert re.search( + rf"pto\.store\(pto\.cast\(pto\.load\(A, 0, contiguous=4\), " + rf"{re.escape(backing_type)}\), local, 1, contiguous=4\)", + source, + ) + + # Live allocation survives optimization: one allocation and the local + # buffer is actually read/written by all four access sites. + assert source.count("pto.alloc_buffer(") == 1 + assert len(re.findall(r"pto\.load\(local,", source)) == 2 + assert len(re.findall(r", local,", source)) == 2 + + +@pytest.mark.parametrize("dtype", ["int8", "uint8", "int32", "uint32"]) +def test_simtvf_local_integer_kernel_compiles_through_backend(dtype): + """All four admitted dtypes compile through the full PTO backend.""" + target = normalize_pto_target("pto") + with target: + kernel = tilelang.compile(_simt_local_integer_kernel(dtype), target="pto", out_idx=-1) + source = kernel.get_kernel_source() + backing = "pto.i8" if dtype in ("int8", "uint8") else "pto.i32" + assert f"pto.alloc_buffer((8,), {backing})" in source + + +# --------------------------------------------------------------------------- +# Vector-element local slot allocation +# --------------------------------------------------------------------------- +def _simt_local_integer_vector_kernel(alloc=T.alloc_local): + @T.prim_func + def main(A: T.Tensor((4,), "int8"), B: T.Tensor((4,), "int8")): + with T.Kernel(1), T.SimtVF(threads=1): + local = alloc((1,), "int8x4") + local[0] = T.int8x4(A[0], A[1], A[2], A[3]) + B[T.Ramp(0, 1, 4)] = local[0] + + return main + + +@pytest.mark.parametrize( + ("scope_name", "alloc"), + [("local", T.alloc_local), ("local.fragment", T.alloc_fragment)], + ids=["local", "local.fragment"], +) +def test_simtvf_vector_integer_local_allocation_is_signless(scope_name, alloc): + """Vector-element local slots allocate with the same-width signless i8 backing.""" + source = _lower(_simt_local_integer_vector_kernel(alloc)).kernel_source + + assert "pto.alloc_buffer((1,), pto.i8)" in source + assert "pto.alloc_buffer((1,), pto.si8)" not in source + + +# --------------------------------------------------------------------------- +# Admission boundaries stay closed +# --------------------------------------------------------------------------- +def _unsupported_local_kernel(dtype): + @T.prim_func + def main(A: T.Tensor((4,), dtype), B: T.Tensor((4,), dtype)): + with T.Kernel(1), T.SimtVF(threads=1): + local = T.alloc_local((4,), dtype) + local[0] = A[0] + B[0] = local[0] + + return main + + +@pytest.mark.parametrize("dtype", ["int16", "int64"]) +def test_simtvf_local_integer_admission_boundary_is_unchanged(dtype): + with pytest.raises(Exception, match="PTO SIMT local allocation supports"): + _lower(_unsupported_local_kernel(dtype)) + + +# --------------------------------------------------------------------------- +# Scope controls: local.var, shared/global and VMI keep their contracts +# --------------------------------------------------------------------------- +def _local_var_kernel(dtype): + @T.prim_func + def main(A: T.Tensor((8,), dtype), B: T.Tensor((8,), dtype)): + with T.Kernel(1): + src = T.alloc_shared((8,), dtype) + dst = T.alloc_shared((8,), dtype) + T.copy(A, src) + T.copy(A, dst) + with T.SimtVF(threads=1): + addr = T.alloc_var(dtype) + addr = src[3] + dst[3] = addr + T.copy(dst, B) + + return main + + +def test_simtvf_local_var_integer_keeps_logical_type(): + source = _lower(_local_var_kernel("int32")).kernel_source + # local.var is a Python SSA value typed with the logical dtype; it is not + # a signless-backed local buffer. + assert "addr = pto.const(0, dtype=pto.si32)" in source + assert "pto.alloc_buffer((1,), pto.i32)" not in source + + +def _shared_control_kernel(dtype): + @T.prim_func + def main(A: T.Tensor((8,), dtype), B: T.Tensor((8,), dtype)): + with T.Kernel(1): + src = T.alloc_shared((8,), dtype) + dst = T.alloc_shared((8,), dtype) + T.copy(A, src) + T.copy(A, dst) + with T.SimtVF(threads=1): + dst[0] = src[0] + T.copy(dst, B) + + return main + + +@pytest.mark.parametrize("dtype", ["uint8", "int32"]) +def test_simtvf_shared_integer_accesses_do_not_use_local_adapter(dtype): + source = _lower(_shared_control_kernel(dtype)).kernel_source + line = next(line for line in source.splitlines() if "pto.store(pto.load(" in line) + assert "pto.cast(" not in line, line + + +def _vmi_kernel(): + @T.prim_func + def main(A: T.Tensor((256,), "int32"), C: T.Tensor((256,), "int32")): + with T.Kernel(1): + work = T.alloc_shared((256,), "int32") + T.copy(A, work) + with T.SimdVF(): + mask = T.vmi.create_mask(64, size=64) + value = T.vmi.vload(work[0], size=64) + T.vmi.vstore(value, work[64], mask) + T.copy(work, C) + + return main + + +def test_vmi_integer_carrier_keeps_signed_type(): + artifact = _lower(_vmi_kernel()) + found = [] + + def visit(node): + if isinstance(node, tvm.tirx.Call) and getattr(node.op, "name", "") == "tl.vmi.vload": + found.append(node) + + for func in artifact.device_mod.functions.values(): + tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + + assert found, "no VMI load in the lowered module" + assert any(str(call.dtype) == "int32x64" for call in found), [str(call.dtype) for call in found] diff --git a/testing/ascend/language/pto/test_simtvf_local_integer_storage_run.py b/testing/ascend/language/pto/test_simtvf_local_integer_storage_run.py new file mode 100644 index 0000000000..ca8d174a09 --- /dev/null +++ b/testing/ascend/language/pto/test_simtvf_local_integer_storage_run.py @@ -0,0 +1,236 @@ +"""Device acceptance: SIMT local integer signless storage and adapters. + +Each case compiles through the real PTO pipeline (PTODSL -> PTOAS VPTO -> +Bisheng) and runs on the NPU. The reference is computed independently on +the CPU with explicit target-width wrapping, so a byte-copy-only implementation +cannot pass: the kernel performs a signed/unsigned comparison and arithmetic +after reloading the value from the local buffer, plus exact boundary values. + +Vector access offsets are 32-byte aligned (b8: 32 lanes, b32: 8 lanes), the +hardware-legal SIMT vector shape. All source lanes read by the kernel are +initialized by the GM->UB DMA and the destination is fully initialized by a +second DMA before the stores. +""" + +import os +import sys + +import pytest +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +os.environ["TILELANG_DISABLE_DATA_RACE_CHECK"] = "1" + +import torch_npu # noqa: E402,F401 (import before tilelang/ptodsl) +import tilelang # noqa: E402 +import tilelang.ascend.language as T # noqa: E402 + +CASES = { + # dtype: (torch dtype, vector lanes, boundary values, signed?, high bit) + "int8": (torch.int8, 32, (127, -1, -128, 0), True, 0), + "uint8": (torch.uint8, 32, (128, 255, 0, 1), False, 128), + "int32": (torch.int32, 8, (2147483647, -1, -2147483648, 0), True, 0), + "uint32": (torch.uint32, 8, (2147483648, 4294967295, 0, 1), False, 2147483648), +} + + +def _require_npu_runtime(): + if not hasattr(torch, "npu") or not torch.npu.is_available(): + pytest.skip("torch.npu is unavailable in the current environment") + try: + source = torch.arange(8, dtype=torch.float32) + copied = source.to(torch.device("npu")) + torch.npu.synchronize() + assert torch.equal(copied.cpu(), source) + except Exception as err: # noqa: BLE001 + pytest.skip(f"NPU runtime is unavailable in the current environment: {err}") + + +@pytest.fixture(scope="module") +def npu(): + _require_npu_runtime() + return torch.device("npu") + + +def _wrap(value, bits, signed): + value &= (1 << bits) - 1 + if signed and value >= (1 << (bits - 1)): + value -= 1 << bits + return value + + +def _condition(x, dtype, signed, highbit): + """Signedness-sensitive comparison evaluated while tracing the kernel.""" + if signed: + return x < T.cast(0, dtype) + return x >= T.cast(highbit, dtype) + + +def _local_integer_kernel(dtype, lanes, bounds, signed, highbit): + b0, b1, b2, b3 = bounds + scalar_index = lanes + 4 + + @T.prim_func + def main(A: T.Tensor((64,), dtype), C: T.Tensor((64,), dtype)): + with T.Kernel(1): + src = T.alloc_shared((64,), dtype) + dst = T.alloc_shared((64,), dtype) + T.copy(A, src) + T.copy(A, dst) # every destination lane is initialized + with T.SimtVF(threads=1): + local = T.alloc_local((64,), dtype) + # Contiguous builtin-vector store/load landing. + local[T.Ramp(0, 1, lanes)] = src[T.Ramp(0, 1, lanes)] + dst[T.Ramp(0, 1, lanes)] = local[T.Ramp(0, 1, lanes)] + # Reload one vector, load a second, and compute v + w. + local[T.Ramp(lanes, 1, lanes)] = src[T.Ramp(lanes, 1, lanes)] + left = local[T.Ramp(0, 1, lanes)] + right = local[T.Ramp(lanes, 1, lanes)] + dst[T.Ramp(0, 1, lanes)] = left + right + # Scalar store/load landing. The index is outside the + # vector destination range so the vector result stays + # observable. + local[scalar_index] = src[scalar_index] + dst[scalar_index] = local[scalar_index] + # Signed/unsigned comparison and arithmetic after reload. + value = local[scalar_index] + dst[scalar_index + 1] = T.if_then_else( + _condition(value, dtype, signed, highbit), + value + T.cast(1, dtype), + value + T.cast(2, dtype), + ) + # Exact boundary values through scalar landings. + local[40] = T.cast(b0, dtype) + dst[40] = local[40] + local[41] = T.cast(b1, dtype) + dst[41] = local[41] + local[42] = T.cast(b2, dtype) + dst[42] = local[42] + local[43] = T.cast(b3, dtype) + dst[43] = local[43] + T.copy(dst, C) + + return main + + +def _expected(a, lanes, bounds, signed): + out = a.clone() + bits = a.dtype.itemsize * 8 + for lane in range(lanes): + out[lane] = _wrap(int(a[lane]) + int(a[lanes + lane]), bits, signed) + scalar_index = lanes + 4 + out[scalar_index] = a[scalar_index] + x = int(a[scalar_index]) + takes_true = x < 0 if signed else x >= (1 << (bits - 1)) + x = x + 1 if takes_true else x + 2 + out[scalar_index + 1] = _wrap(x, bits, signed) + for index, boundary in enumerate(bounds): + out[40 + index] = _wrap(boundary, bits, signed) + return out + + +def _trial_input(torch_dtype, signed, trial, lanes): + """Fresh input per trial. + + Vector lanes at [0,4) and [lanes,lanes+4) use boundary/high-bit patterns + whose wrapped sums are computed independently by ``_expected``. The scalar + comparison value alternates between the true and false arms across trials. + """ + bits = torch_dtype.itemsize * 8 + generator = torch.Generator().manual_seed(7000 + trial) + base = torch.randint(0, 50, (64,), generator=generator, dtype=torch.int64) + trial * 3 + values = [int(value) for value in base] + maxv = (1 << (bits - 1)) - 1 if signed else (1 << bits) - 1 + minv = -(1 << (bits - 1)) if signed else 0 + high = 1 << (bits - 1) + if signed: + first_pattern = [maxv, minv, -1, 0] + second_pattern = [1, -1, 1, minv] + else: + first_pattern = [maxv, 0, high, 1] + second_pattern = [1, high, 1, maxv] + variants = [ + (first_pattern, second_pattern), + (second_pattern, first_pattern), + (list(reversed(first_pattern)), list(reversed(second_pattern))), + ] + first, second = variants[trial % len(variants)] + values[0:4] = first + values[lanes : lanes + 4] = second + scalar_index = lanes + 4 + if trial % 2 == 0: + values[scalar_index] = -5 if signed else high + else: + values[scalar_index] = 0 + return torch.tensor(values, dtype=torch_dtype) + + +@pytest.mark.parametrize("dtype", list(CASES), ids=list(CASES)) +def test_simtvf_local_integer_roundtrip_arithmetic_and_boundaries(npu, dtype): + torch_dtype, lanes, bounds, signed, highbit = CASES[dtype] + kernel = tilelang.compile( + _local_integer_kernel(dtype, lanes, bounds, signed, highbit), + target="pto", + out_idx=-1, + ) + source = kernel.get_kernel_source() + + backing = "pto.i8" if dtype in ("int8", "uint8") else "pto.i32" + logical = {"int8": "pto.si8", "uint8": "pto.ui8", "int32": "pto.si32", "uint32": "pto.ui32"}[dtype] + # Live allocation and all four adapter landings are present in the final + # generated source (they survived lowering/optimization). + assert f"pto.alloc_buffer((64,), {backing})" in source, source + scalar_index = lanes + 4 + assert f"pto.cast(pto.load(local, {scalar_index}), {logical})" in source, source + assert f"pto.store(pto.cast(pto.load(local, 0, contiguous={lanes}), {logical})" in source, source + assert f", {backing}), local, 0, contiguous={lanes})" in source, source + assert f", {backing}), local, {scalar_index})" in source, source + + arms = set() + for trial in range(3): + a = _trial_input(torch_dtype, signed, trial, lanes) + x = int(a[scalar_index]) + arms.add(x < 0 if signed else x >= highbit) + out = kernel(a.to(npu)) + torch.npu.synchronize() + got = out.cpu() + expected = _expected(a, lanes, bounds, signed) + assert torch.equal(got, expected), f"{dtype} trial {trial} mismatch\n got={got.tolist()}\n exp={expected.tolist()}" + # Guarded output: every lane outside the intended writes still holds + # the independently initialized input value. + touched = set(range(lanes)) | {scalar_index, scalar_index + 1, 40, 41, 42, 43} + for index in range(64): + if index not in touched: + assert got[index].item() == a[index].item(), (dtype, trial, index) + assert arms == {True, False}, f"{dtype} comparison arms exercised: {arms}" + + +def _shared_control_kernel(dtype): + @T.prim_func + def main(A: T.Tensor((8,), dtype), C: T.Tensor((8,), dtype)): + with T.Kernel(1): + src = T.alloc_shared((8,), dtype) + dst = T.alloc_shared((8,), dtype) + T.copy(A, src) + T.copy(A, dst) + with T.SimtVF(threads=1): + dst[0] = src[0] + T.copy(dst, C) + + return main + + +@pytest.mark.parametrize("dtype,torch_dtype", [("uint8", torch.uint8), ("int32", torch.int32)]) +def test_simtvf_shared_integer_access_control(npu, dtype, torch_dtype): + kernel = tilelang.compile(_shared_control_kernel(dtype), target="pto", out_idx=-1) + source = kernel.get_kernel_source() + line = next(line for line in source.splitlines() if "pto.store(pto.load(" in line) + assert "pto.cast(" not in line, line + + a = torch.arange(-4, 4, dtype=torch.int32).to(torch_dtype) + out = kernel(a.to(npu)) + torch.npu.synchronize() + assert torch.equal(out.cpu(), a) 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/__init__.py b/tilelang/ascend/language/__init__.py index 423d829d51..1a431d576e 100644 --- a/tilelang/ascend/language/__init__.py +++ b/tilelang/ascend/language/__init__.py @@ -66,6 +66,7 @@ from .frame import Cube, CubeFrame, SimdVF, SimdVFFrame, SimtVF, SimtVFFrame, Vector, VectorFrame # noqa: F401 from . import simd as simd # noqa: F401 (exposed as T.simd.*) +from . import vmi as vmi # noqa: F401 (exposed as T.vmi.*) # Scalar reciprocal sqrt is a frontend composition of common math ops. from .math import rsqrt # noqa: F401 @@ -126,6 +127,7 @@ "rng_rand_float", "rsqrt", "simd", + "vmi", "unroll", ) diff --git a/tilelang/ascend/language/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/language/vmi.py b/tilelang/ascend/language/vmi.py new file mode 100644 index 0000000000..7be9612093 --- /dev/null +++ b/tilelang/ascend/language/vmi.py @@ -0,0 +1,1449 @@ +"""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 _access_ptr_base_load(value): + """Return the immediate BufferLoad behind a frontend ``tl.access_ptr`` call. + + This is a local type recovery for an explicit pointer expression, not + general pointer provenance analysis: only the call's own base argument + (through a direct frontend ``Bind`` alias) is inspected. + """ + if isinstance(value, tirx.Var): + from tilelang.language.frame import has_let_value, get_let_value + + if has_let_value(value): + bound = get_let_value(value) + if bound is not None and not isinstance(bound, tirx.Var): + return _access_ptr_base_load(bound) + return None + if isinstance(value, tirx.Call): + op = value.op + if getattr(op, "name", None) == "tl.access_ptr" and value.args and isinstance(value.args[0], BufferLoad): + return value.args[0] + return None + + +def _require_ub_address(value, *, context: str): + buffer = None + base_load = _access_ptr_base_load(value) + if base_load is not None: + buffer = base_load.buffer + elif isinstance(value, (BufferLoad, BufferRegion)): + buffer = value.buffer + elif isinstance(value, Buffer): + buffer = value + 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): + base_load = _access_ptr_base_load(value) + if base_load is not None: + return _dtype(base_load.buffer.dtype) + if isinstance(value, BufferLoad): + return _dtype(value.buffer.dtype) + if isinstance(value, BufferRegion): + 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_memory_modes( + context: str, + *, + group, + stride, + block_stride, +): + if group is not None: + if block_stride is not None: + raise TypeError(f"{context} does not allow block_stride together with group") + if stride is None: + raise TypeError(f"{context} with group=... requires stride") + return + + if block_stride is not None: + 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 _is_unit_stride_1d(buffer) -> bool: + if len(buffer.shape) != 1: + return False + if buffer.strides: + return len(buffer.strides) == 1 and _constant_int(buffer.strides[0]) == 1 + return True + + +def _tight_buffer_span(buffer): + """Return ``prod(shape)`` when *buffer* is tightly packed, else ``None``. + + A whole-buffer descriptor is only sound when the logical buffer covers one + contiguous physical span. Explicit strides must equal the row-major + strides; default (empty) strides are row-major by construction. + """ + ndim = len(buffer.shape) + if ndim == 0: + return None + if buffer.strides: + if len(buffer.strides) != ndim: + return None + from tvm import arith + + analyzer = arith.Analyzer() + expected = 1 + for dim in range(ndim - 1, -1, -1): + if not analyzer.can_prove_equal(buffer.strides[dim], expected): + return None + expected = expected * buffer.shape[dim] + span = 1 + for dim in buffer.shape: + span = span * dim + return span + + +def _source_relative_offset(source): + """Element offset of *source* relative to its buffer's logical origin. + + ``Buffer.offset_of`` already folds strides and ``elem_offset``; the origin + offset (which is exactly ``elem_offset``) is subtracted so the result can + be added to an ``access_ptr(buffer, ...)`` whose lowering adds + ``elem_offset`` once. + """ + if isinstance(source, BufferLoad): + buffer = source.buffer + indices = list(source.indices) + elif isinstance(source, BufferRegion): + buffer = source.buffer + indices = [region.min for region in source.region] + elif isinstance(source, Buffer): + return _zero_offset() + else: + return None + if not buffer.shape or len(indices) != len(buffer.shape): + return None + source_offsets = buffer.offset_of(indices) + origin_offsets = buffer.offset_of([0] * len(buffer.shape)) + # A single flattened index is the only form whose linear address this + # frontend contract can relate to one descriptor. Multi-address / + # axis-separator buffers are not proven linear and keep the existing + # boundary instead of being silently reduced to offsets[0]. + if len(source_offsets) != 1 or len(origin_offsets) != 1: + return None + return source_offsets[0] - origin_offsets[0] + + +def _whole_buffer_plan(source, *, fallback_extent): + """Declare the whole known tight buffer and keep the relative origin offset. + + Used when a group/block stride is dynamic (or a static negative stride has + no directly constructible low address). When the buffer is not provably + tight (view/padding/strides), this returns the caller's fallback extent + beginning at the source origin: that is only the pre-existing boundary and + is not claimed as a complete conservative footprint. + """ + buffer = source if isinstance(source, Buffer) else source.buffer + span = _tight_buffer_span(buffer) + relative = _source_relative_offset(source) + if span is None or relative is None: + return fallback_extent, None, None + return span, buffer, relative + + +def _group_access_plan(source, *, lanes, group, stride, is_load, elem_dtype): + """May-access span of an equal-distance group access (C.2/C.3.1). + + The first group starts at the original access address ``A``. For stride + ``S`` and physical group width ``W`` the contiguous envelope is + ``[A + min(0,(G-1)*S), A + max(0,(G-1)*S) + W)``. A negative static + stride on a 1-D unit-stride buffer relocates the descriptor to the lowest + address and compensates with the VMI instruction offset. + + A group load is unmasked and therefore reads whole physical vectors: when + the logical per-group width ``N/G`` is not a multiple of the physical VL, + the last chunk still reads a full VL (e.g. f32 G=1/N=80 lowers to two + full vlds covering 128 address elements). A group store keeps the logical + width because the backend's generated tail mask bounds the writes; do not + copy the load rounding onto stores. + """ + group_const = _constant_int(group) + if group_const is None or group_const <= 0: + return _whole_buffer_plan(source, fallback_extent=lanes) + # Slot (one result/input element per group) and group broadcast read or + # write exactly one address element per group. Ordinary groups cover + # ceil(N/G) elements; the ceiling keeps a non-divisible shape conservative. + if lanes <= group_const: + per_group = 1 + else: + per_group = (lanes + group_const - 1) // group_const + bits = int(getattr(elem_dtype, "bits", 0) or 0) + if is_load and bits > 0: + vector_elements = max(1, 2048 // bits) + per_group = (per_group + vector_elements - 1) // vector_elements * vector_elements + stride_const = _constant_int(stride) + if stride_const is None: + return _whole_buffer_plan(source, fallback_extent=lanes) + span = (group_const - 1) * stride_const + low = min(0, span) + high = max(0, span) + per_group + extent = high - low + if low == 0: + return extent, None, None + if isinstance(source, BufferLoad) and len(source.indices) == 1 and _is_unit_stride_1d(source.buffer): + shifted = tirx.BufferLoad(source.buffer, [source.indices[0] + low]) + return extent, shifted, -low + return _whole_buffer_plan(source, fallback_extent=lanes) + + +def _block_access_plan(source, *, lanes, elem_dtype, block_stride): + """May-access span of a block-stride access (C.3.2). + + ``block_stride`` is measured in 32-byte blocks. A full physical vector + covers eight blocks; physical chunks beyond one vector start at the next + block group, and the tail chunk of a masked store conservatively covers + its whole block group as well. + """ + bits = int(getattr(elem_dtype, "bits", 0) or 0) + if bits <= 0: + return lanes, None, None + elements_per_block = max(1, 256 // bits) + vector_elements = max(1, 2048 // bits) + chunks = (lanes + vector_elements - 1) // vector_elements + blocks = chunks * max(1, vector_elements // elements_per_block) + block_const = _constant_int(block_stride) + if block_const is None or block_const < 0: + # A dynamic control word (and the i16 control-word encoding) has no + # statically provable span; declare the whole known tight buffer. + return _whole_buffer_plan(source, fallback_extent=lanes) + extent = (blocks - 1) * block_const * elements_per_block + elements_per_block + return extent, None, None + + +def _vmi_access_plan(source, *, lanes, elem_dtype, group, stride, block_stride, is_load): + """Compute (descriptor extent, descriptor base, instruction offset). + + ``descriptor base`` is ``None`` when the original source is already the + right ``access_ptr`` base. The instruction offset is ``None`` for zero. + Only buffer sources are analysed; explicit pointers keep their caller + declared ``ptr + offset`` ABI untouched. + """ + if not isinstance(source, (BufferLoad, Buffer, BufferRegion)): + return None, None, None + if group is not None: + return _group_access_plan( + source, + lanes=lanes, + group=group, + stride=stride, + is_load=is_load, + elem_dtype=elem_dtype, + ) + if block_stride is not None: + return _block_access_plan(source, lanes=lanes, elem_dtype=elem_dtype, block_stride=block_stride) + if str(elem_dtype) == "float4_e2m1fn": + # Genuine FP4 descriptors stay in logical FP4 element units; + # RewriteFp4ToFp4x2 performs the single physical packing. + return lanes, None, None + if is_load: + bits = int(getattr(elem_dtype, "bits", 0) or 0) + if bits <= 0: + return lanes, None, None + vector_elements = max(1, 2048 // bits) + # An unmasked load reads whole physical vectors: cover the tail chunk. + return (lanes + vector_elements - 1) // vector_elements * vector_elements, None, None + return lanes, None, None + + +def _resolve_ptr_and_offset(source, offset=None, *, access_type: str, extent=None, descriptor=None, instruction_offset=None): + """Normalize VMI address syntax into PTODSL's explicit ptr+offset ABI.""" + if isinstance(source, (BufferLoad, Buffer, BufferRegion)) and offset is not None: + raise TypeError("T.vmi buffer addresses do not accept an offset; express the offset in the buffer index") + base = descriptor if descriptor is not None else source + if isinstance(base, (BufferLoad, Buffer, BufferRegion)): + ptr = access_ptr(base, access_type, extent=extent, offset=0) + return ptr, instruction_offset if instruction_offset is not None else _zero_offset() + if offset is None: + return source, _zero_offset() + return source, offset + + +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, + stride=None, + block_stride=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_memory_modes( + "T.vmi.vload(...)", + group=group, + stride=stride, + block_stride=block_stride, + ) + if size is None: + raise TypeError("T.vmi.vload(...) requires size") + source_elem = _require_address_element_dtype(source, context="T.vmi.vload(...)") + 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 group is not None or block_stride is not None: + raise TypeError("T.vmi.vload(...) only supports contiguous packed FP4 loads") + + result_dtype = _vector_result_dtype(source, size, elem_dtype=source_elem) + extent, descriptor, instruction_offset = _vmi_access_plan( + source, + lanes=size, + elem_dtype=source_elem, + group=group, + stride=stride, + block_stride=block_stride, + is_load=True, + ) + ptr, extra_offset = _resolve_ptr_and_offset( + source, + offset, + access_type="r", + extent=extent, + descriptor=descriptor, + instruction_offset=instruction_offset, + ) + if is_packed_fp4: + extra_offset = _pack_fp4_element_offset(extra_offset, context="T.vmi.vload(...)") + attrs = { + "size": size, + "stride": stride, + "block_stride": block_stride, + "group": group, + "loc": loc, + "ip": ip, + } + return _call_vmi("vload", result_dtype, ptr, extra_offset, **attrs) + + +@_scope_guarded +def vstore( + value, + destination, + offset=None, + mask=None, + *, + stride=None, + block_stride=None, + group=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_memory_modes( + "T.vmi.vstore(...)", + group=group, + stride=stride, + block_stride=block_stride, + ) + if group is not None and mask is not None: + raise TypeError("T.vmi.vstore(...) group mode does not take a mask operand") + if isinstance(value, VmiPair) or _is_sequence(value): + raise TypeError("T.vmi.vstore(...) expects a single VMI vector; store each result separately") + + store_lanes = _lanes_of(value) + value_dtype = _element_dtype_of(value) + if str(value_dtype) == "float4_e2m1fn": + if store_lanes % 2: + raise ValueError("T.vmi.vstore(...) requires an even FP4 lane count for packed FP4 storage") + if group is not None or block_stride is not None: + raise TypeError("T.vmi.vstore(...) only supports contiguous packed FP4 stores") + elif mask is not None: + mask = _require_mask(mask, context="T.vmi.vstore(...)") + + extent, descriptor, instruction_offset = _vmi_access_plan( + destination, + lanes=store_lanes, + elem_dtype=value_dtype, + group=group, + stride=stride, + block_stride=block_stride, + is_load=False, + ) + ptr, extra_offset = _resolve_ptr_and_offset( + destination, + offset, + access_type="w", + extent=extent, + descriptor=descriptor, + instruction_offset=instruction_offset, + ) + if str(value_dtype) == "float4_e2m1fn": + extra_offset = _pack_fp4_element_offset(extra_offset, context="T.vmi.vstore(...)") + if mask is not None: + mask = _require_mask(mask, context="T.vmi.vstore(...) of FP4") + if not str(_dtype_of(mask)).startswith("bool"): + raise TypeError("T.vmi.vstore(...) of FP4 expects a VMI mask operand") + physical_lanes = store_lanes // 2 + if _lanes_of(mask) != physical_lanes: + raise TypeError( + "T.vmi.vstore(...) of FP4 requires a physical FP4x2 mask with " + f"half as many lanes as its {store_lanes}-lane logical FP4 value " + f"(expected {physical_lanes}, got {_lanes_of(mask)})" + ) + return _call_vmi( + "vstore", + "void", + value, + ptr, + extra_offset, + mask, + stride=stride, + block_stride=block_stride, + group=group, + 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 not None and not isinstance(reassoc, bool): + raise TypeError(f"{context} requires reassoc to be the Python boolean True or False; received {reassoc!r}") + if reassoc is False: + raise TypeError( + f"{context} does not support reassoc=False: the PTO backend accepts only reassoc=True (integer vectors may omit reassoc)" + ) + if reassoc is None and str(elem_dtype).startswith(("float", "bfloat")): + raise TypeError(f"{context} on floating-point vectors requires an explicit reassoc=True argument; the PTO backend rejects omission") + lanes = 1 if group is None else _normalize_lanes(group, context="T.vmi.vcadd(..., group=...)") + if group is not None and _lanes_of(source) % group != 0: + raise ValueError(f"{context} requires source lanes to be divisible by group") + 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", +] diff --git a/tilelang/ascend/pipeline.py b/tilelang/ascend/pipeline.py index c5f31833bf..862237e126 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) @@ -184,6 +194,11 @@ def AscendPassPipelineBody(mod: IRModule, target: Target) -> IRModule: mod = tilelang.transform.MakePackedAPI()(mod) mod = tilelang.transform.Simplify()(mod) mod = tilelang.transform.LowerDeviceKernelLaunch()(mod) + # Verify PTO requirements on the final device TIR before code generation. + # The current check proves masked-store destination alignment after all + # layout and UB address rewrites. + if target_is_pto(target): + mod = ascend_transform.VerifyPTOCodegenLegality()(mod) return mod diff --git a/tilelang/ascend/transform/__init__.py b/tilelang/ascend/transform/__init__.py index d1fd48fa95..f1d8f5ba1d 100644 --- a/tilelang/ascend/transform/__init__.py +++ b/tilelang/ascend/transform/__init__.py @@ -134,6 +134,50 @@ 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 VerifyPTOCodegenLegality(): + """Verify PTO requirements before code generation without modifying IR. + + Consumes the final device TIR after LowerDeviceKernelLaunch and the final + layout and UB rewrites. Currently checks destination alignment for + default/continuous explicit-mask VMI stores of b8/b16/b32 data: 32 bytes + for ordinary stores, or natural element alignment for recognized 1PT + writeback. Bool data is excluded. Shapes, layouts, masks and addresses are + unchanged. Functions without a PTO target pass through unchanged. + """ + return _ffi_api.VerifyPTOCodegenLegality() # type: ignore + + +def VectorizeParallelToPTO(): + """Rewrite verified T.Parallel loops inside PTO SIMD_VF blocks into + logical VMI code. + + 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 +289,7 @@ def RewriteFlagToBuf(): "NormalizeBufferVersion", "NormalizeControlFlowForSchedule", "NormalizeConflictHints", + "LegalizeParallelToPTO", "PrepareMultiBuffer", "ResolveCore", "RewriteAscendBufferVersionLayout", @@ -253,4 +298,7 @@ def RewriteFlagToBuf(): "RewriteFp4ToFp4x2", "RewriteFlagToBuf", "UnrollLoopSkipVF", + "VectorizeParallelToPTO", + "VerifyParallelToPTO", + "VerifyPTOCodegenLegality", ] 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)