From 8ffdd4942483d4d7a7e8aeef36012aa8acff68fb Mon Sep 17 00:00:00 2001 From: tqchen Date: Mon, 21 Sep 2026 21:34:59 +0000 Subject: [PATCH 1/5] [REFACTOR][TIR] Represent For thread binding as a semantic annotation --- include/tvm/tirx/stmt.h | 37 ++- python/tvm/s_tir/dlight/gpu/fallback.py | 2 +- python/tvm/s_tir/schedule/schedule.py | 6 + python/tvm/tirx/stmt.py | 19 +- .../transform/split_call_tir_by_pattern.cc | 2 +- .../sblock_buffer_access_lca_detector.cc | 4 +- .../feature_extractor/per_store_feature.cc | 22 +- .../rewrite_parallel_vectorize_unroll.cc | 2 +- src/s_tir/schedule/analysis/analysis.cc | 4 +- src/s_tir/schedule/ir_comparator.cc | 21 +- src/s_tir/schedule/primitive/annotate.cc | 8 + .../schedule/primitive/blockize_tensorize.cc | 2 +- .../schedule/primitive/compute_inline.cc | 2 +- .../schedule/primitive/decompose_padding.cc | 2 +- src/s_tir/schedule/primitive/for_kind.cc | 35 ++- .../schedule/primitive/loop_transformation.cc | 8 +- src/s_tir/schedule/primitive/reduction.cc | 4 +- src/s_tir/schedule/utils.h | 4 +- src/s_tir/transform/compact_buffer_region.cc | 7 +- src/s_tir/transform/default_gpu_schedule.cc | 2 +- .../transform/inject_software_pipeline.cc | 2 + src/s_tir/transform/lift_thread_binding.cc | 24 +- src/s_tir/transform/loop_partition.cc | 2 +- .../transform/lower_cross_thread_reduction.cc | 30 +-- src/s_tir/transform/lower_opaque_block.cc | 9 +- src/s_tir/transform/memhammer_coalesce.cc | 9 +- .../transform/memhammer_intermediate_stage.cc | 8 +- .../transform/memhammer_lower_auto_copy.cc | 11 +- .../transform/profile_instrumentation.cc | 8 +- src/s_tir/transform/unify_thread_binding.cc | 9 +- src/target/llvm/codegen_cpu.cc | 2 +- src/tirx/ir/data_type_rewriter.cc | 6 +- src/tirx/ir/stmt.cc | 61 +++-- src/tirx/script/builder/ir.cc | 2 +- src/tirx/script/printer/for_loop.cc | 25 +- src/tirx/transform/bind_target.cc | 4 +- src/tirx/transform/lower_tirx_opaque.cc | 9 +- src/tirx/transform/lower_tvm_builtin.cc | 2 +- src/tirx/transform/make_packed_api.cc | 4 +- src/tirx/transform/storage_rewrite.cc | 4 +- src/tirx/transform/unroll_loop.cc | 5 + tests/python/relax/test_pipeline.py | 2 +- .../test_tir_schedule_binding_annotation.py | 119 +++++++++ .../tirx-base/test_tir_for_thread_binding.py | 229 ++++++++++++++++++ .../test_tvmscript_ir_builder_tir.py | 2 +- .../tvmscript/test_tvmscript_roundtrip.py | 4 +- 46 files changed, 603 insertions(+), 182 deletions(-) create mode 100644 tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py create mode 100644 tests/python/tirx-base/test_tir_for_thread_binding.py diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h index c612f25b0b43..c311bf7e1d1b 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -559,7 +559,7 @@ class IfThenElse : public Stmt { enum class ForKind : int { /*! \brief default semantics -- serial execution. */ kSerial = 0, - /*! \brief Parallel execution on CPU. */ + /*! \brief Parallel execution, optionally bound to an execution thread. */ kParallel = 1, /*! * \brief Vector SIMD loop. @@ -567,14 +567,7 @@ enum class ForKind : int { */ kVectorized = 2, /*! \brief The loop body must be unrolled. */ - kUnrolled = 3, - /*! - * \brief The loop variable is bound to a thread in - * an environment. In the final stage of lowering, - * the loop is simply removed and the loop variable is - * mapped to the corresponding context thread. - */ - kThreadBinding = 4 + kUnrolled = 3 }; /*! @@ -599,18 +592,12 @@ class ForNode : public StmtNode { ForKind kind; /*! \brief The body of the for loop. */ Stmt body; - /*! - * \brief Only valid when kind == ForKind::kThreadBinding - * The context thread that this loop variable bounds to. - */ - ffi::Optional thread_binding; /*! * \brief Additional annotations about the loop. * - * These annotations can be used as auxiliary hint - * to future transformations. An annotation should - * not change the control flow semantics of the loop - * and can be ignored in most passes. + * Most annotations are auxiliary transformation hints. The reserved + * thread_binding annotation is semantic: its IterVar binds a parallel loop + * to an execution thread and must be preserved until binding is lowered. */ ffi::Map annotations; /*! @@ -626,7 +613,6 @@ class ForNode : public StmtNode { .def_ro("extent", &ForNode::extent) .def_ro("kind", &ForNode::kind) .def_ro("body", &ForNode::body) - .def_ro("thread_binding", &ForNode::thread_binding) .def_ro("annotations", &ForNode::annotations) .def_ro("step", &ForNode::step); } @@ -634,6 +620,15 @@ class ForNode : public StmtNode { /*! \brief Check it is a loop without nontrivial loop step. */ bool HasTrivialStep() const; + /*! \brief Get the validated semantic thread binding, if present. */ + TVM_DLL ffi::Optional GetThreadBinding() const; + /*! \brief Set or remove the semantic binding on a parallel loop. */ + TVM_DLL void SetThreadBinding(ffi::Optional binding); + /*! \brief Whether this parallel loop is bound to an execution thread. */ + bool IsThreadBinding() const { return GetThreadBinding().has_value(); } + /*! \brief Whether this is an ordinary, unbound CPU parallel loop. */ + bool IsParallel() const { return kind == ForKind::kParallel && !IsThreadBinding(); } + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.For", ForNode, StmtNode); }; @@ -799,6 +794,8 @@ class ScopeIdDefStmt : public Stmt { /*! \brief namespace of possible attributes in AttrStmt.attr_key */ namespace attr { +/*! \brief Semantic IterVar binding of a parallel For loop to an execution thread. */ +constexpr const char* thread_binding = "thread_binding"; /*! * \brief Mark the scope as when computation start to happen. * This can hint some code generator to create a new function for compute. @@ -886,8 +883,6 @@ inline const char* ForKind2String(ForKind t) { return "vectorized"; case ForKind::kUnrolled: return "unroll"; - case ForKind::kThreadBinding: - return "thread_binding"; } TVM_FFI_THROW(InternalError) << "Unknown ForKind" << t; TVM_FFI_UNREACHABLE(); diff --git a/python/tvm/s_tir/dlight/gpu/fallback.py b/python/tvm/s_tir/dlight/gpu/fallback.py index e64f8544ecf1..a194b07eb8f0 100644 --- a/python/tvm/s_tir/dlight/gpu/fallback.py +++ b/python/tvm/s_tir/dlight/gpu/fallback.py @@ -42,7 +42,7 @@ def visit_attr(node: tirx.AttrStmt): def visit_for(node: tirx.For): nonlocal found - if node.kind == tirx.ForKind.THREAD_BINDING: + if node.is_thread_binding(): found = True tvm_ffi.structural_walk( diff --git a/python/tvm/s_tir/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py index 42ac24df3dce..f1c15fe5d090 100644 --- a/python/tvm/s_tir/schedule/schedule.py +++ b/python/tvm/s_tir/schedule/schedule.py @@ -3128,6 +3128,9 @@ def annotate( ) -> None: """Annotate a block/loop with a key value pair + The semantic loop annotation ``thread_binding`` must be set with + :meth:`bind`, which checks scheduling legality. + Parameters ---------- block_or_loop: SBlockRV | LoopRV @@ -3184,6 +3187,9 @@ def after_annotate(a: T.handle, b: T.handle) -> None: def unannotate(self, block_or_loop: SBlockRV | LoopRV, ann_key: str) -> None: """Unannotate a block/loop's annotation with key ann_key + Remove a loop's semantic ``thread_binding`` with :meth:`parallel`, + which checks whether ordinary parallel execution is legal. + Parameters ---------- block_or_loop: SBlockRV | LoopRV diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 70ac359a3e04..3502c03e744a 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -171,7 +171,6 @@ class ForKind(IntEnum): PARALLEL = 1 VECTORIZED = 2 UNROLLED = 3 - THREAD_BINDING = 4 # pylint: disable=invalid-name @tvm_ffi.register_object("tirx.For") @@ -197,14 +196,14 @@ class For(Stmt): thread_binding: Optional[tirx.IterVar] The thread this loop binds to. Only valid - if kind is ThreadBinding + if kind is PARALLEL. Stored in the semantic thread_binding annotation. step : Expr The loop step. Default to none which represent one. annotations: Optional[Mapping[str, Object]] - Additional annotation hints. + Additional annotations, including the semantic thread_binding IterVar. span : Optional[Span] The location of the stmt in the source code. @@ -215,7 +214,6 @@ class For(Stmt): extent: Expr kind: ForKind body: Stmt - thread_binding: IterVar | None annotations: Mapping[str, Object] step: Expr | None span: Span | None @@ -246,6 +244,19 @@ def __init__( span, ) + @property + def thread_binding(self) -> IterVar | None: + """The semantic thread binding of this parallel loop, if present.""" + return self.annotations.get("thread_binding") + + def is_thread_binding(self) -> bool: + """Whether this loop is bound to an execution thread.""" + return self.thread_binding is not None + + def is_parallel(self) -> bool: + """Whether this is an ordinary unbound CPU parallel loop.""" + return self.kind == ForKind.PARALLEL and not self.is_thread_binding() + @tvm_ffi.register_object("tirx.While") class While(Stmt): diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index 53a0812001aa..cf4548ab58e5 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -266,7 +266,7 @@ class ForMatcher : public TensorizeComparator { if (!DefEqual(op->loop_var, rhs->loop_var)) return false; // Only handle the case where the loop start from 0 if (!is_zero(op->min) || !is_zero(rhs->min)) return false; - if (op->thread_binding.has_value() || rhs->thread_binding.has_value()) return false; + if (op->GetThreadBinding().has_value() || rhs->GetThreadBinding().has_value()) return false; if (op->kind != ForKind::kSerial || op->kind != rhs->kind) return false; if (!op->annotations.empty() || !rhs->annotations.empty()) return false; // Match the extents of loops diff --git a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc index 0d73fb07655e..1581b8e22b49 100644 --- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc +++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc @@ -99,9 +99,9 @@ class LCADetector : public s_tir::StmtExprVisitor { const ScopeInfo* parent_scope = ancestor_scopes_.back(); auto* current_scope = arena_.make(parent_scope, op, n); - if (op->thread_binding.has_value()) { + if (op->GetThreadBinding().has_value()) { const runtime::ThreadScope& scope = - runtime::ThreadScope::Create(op->thread_binding.value()->thread_tag); + runtime::ThreadScope::Create(op->GetThreadBinding().value()->thread_tag); if (scope.rank == 0) { blockidx_scopes_.push_back(current_scope); } diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc index 939a136eecb9..f6fa67f4d163 100644 --- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc +++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc @@ -342,16 +342,16 @@ struct LoopNest { int64_t prod = 1; // The product of the extents of all the loops ForVec loops; // All the loops IntVec auto_unroll; // The loops with auto unroll pragma - ForVec parallel; // The loops whose ForKind are kParallel + ForVec parallel; // The CPU parallel loops ForVec vectorize; // The loops whose ForKind are kVectorized ForVec unroll; // The loops whose ForKind are kUnrolled - ForVec blockIdx_x; // The loops whose ForKind are kThreadBinding to blockIdx.x - ForVec blockIdx_y; // The loops whose ForKind are kThreadBinding to blockIdx.y - ForVec blockIdx_z; // The loops whose ForKind are kThreadBinding to blockIdx.z - ForVec threadIdx_x; // The loops whose ForKind are kThreadBinding to threadIdx.x - ForVec threadIdx_y; // The loops whose ForKind are kThreadBinding to threadIdx.y - ForVec threadIdx_z; // The loops whose ForKind are kThreadBinding to threadIdx.z - ForVec vthread; // The loops whose ForKind are kThreadBinding to vthread.* + ForVec blockIdx_x; // The loops bound to blockIdx.x + ForVec blockIdx_y; // The loops bound to blockIdx.y + ForVec blockIdx_z; // The loops bound to blockIdx.z + ForVec threadIdx_x; // The loops bound to threadIdx.x + ForVec threadIdx_y; // The loops bound to threadIdx.y + ForVec threadIdx_z; // The loops bound to threadIdx.z + ForVec vthread; // The loops bound to vthread.* /*! * \brief Push a new loop into the loop nest @@ -370,14 +370,14 @@ struct LoopNest { this->auto_unroll.push_back(*auto_unroll_attr); } ForVec* ref_loops = nullptr; - if (loop->kind == ForKind::kParallel) { + if (loop->IsParallel()) { ref_loops = ∥ } else if (loop->kind == ForKind::kVectorized) { ref_loops = &vectorize; } else if (loop->kind == ForKind::kUnrolled) { ref_loops = &unroll; - } else if (loop->kind == ForKind::kThreadBinding) { - std::string thread_tag = loop->thread_binding.value()->thread_tag; + } else if (loop->IsThreadBinding()) { + std::string thread_tag = loop->GetThreadBinding().value()->thread_tag; if (thread_tag == "blockIdx.x") { ref_loops = &blockIdx_x; } else if (thread_tag == "blockIdx.y") { diff --git a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc index fbfe65128c99..55014ecf9de7 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc @@ -33,7 +33,7 @@ using namespace tvm::tirx; * \return Whether the loop has any annotation */ inline bool HasAnnOrBinding(const ForNode* loop) { - return loop->kind == ForKind::kThreadBinding || !loop->annotations.empty(); + return loop->IsThreadBinding() || !loop->annotations.empty(); } /*! \brief The visitor for extracting the stride of a var in a PrimExpr. */ diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index 8c9194958b72..82f21a8c98ef 100644 --- a/src/s_tir/schedule/analysis/analysis.cc +++ b/src/s_tir/schedule/analysis/analysis.cc @@ -703,8 +703,8 @@ ffi::Map LoopDomainOfSRefTreePath(const StmtSRef& low_inclusive, if (extra_relax_scope.rank != runtime::StorageRank::kGlobal) { for (; p; p = p->parent) { if (const ForNode* loop = p->StmtAs()) { - if (loop->kind == ForKind::kThreadBinding) { - const ffi::String& thread_tag = loop->thread_binding.value()->thread_tag; + if (loop->IsThreadBinding()) { + const ffi::String& thread_tag = loop->GetThreadBinding().value()->thread_tag; if (CanRelaxStorageUnderThread(extra_relax_scope, runtime::ThreadScope::Create(thread_tag))) { result.Set(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent)); diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc index fa0f4c074e34..3263c4add139 100644 --- a/src/s_tir/schedule/ir_comparator.cc +++ b/src/s_tir/schedule/ir_comparator.cc @@ -200,20 +200,16 @@ bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) { } return false; } - if (op->thread_binding.has_value() != rhs->thread_binding.has_value()) { + if (op->GetThreadBinding().has_value() != rhs->GetThreadBinding().has_value()) { if (assert_mode_) { std::ostringstream os; - os << "ForNode thread_bindings do not match: op->thread_binding.has_value()=" - << op->thread_binding.has_value() - << " vs rhs->thread_binding.has_value()=" << rhs->thread_binding.has_value(); + os << "ForNode thread_bindings do not match: op->GetThreadBinding().has_value()=" + << op->GetThreadBinding().has_value() + << " vs rhs->GetThreadBinding().has_value()=" << rhs->GetThreadBinding().has_value(); EmitError(os.str()); } return false; } - if (op->thread_binding.has_value() && - !Dispatch(op->thread_binding.value(), rhs->thread_binding.value())) { - return false; - } if (op->kind != rhs->kind) { if (assert_mode_) { std::ostringstream os; @@ -446,6 +442,15 @@ bool TensorizeComparator::CompareAnnotation(const std::pair(); + IterVar rhs_iter = rhs.second.as_or_throw(); + return CompareIterVar(lhs_iter, rhs_iter) && lhs_iter->thread_tag == rhs_iter->thread_tag && + lhs_iter->dom.defined() == rhs_iter->dom.defined() && + (!lhs_iter->dom.defined() || CompareRange(lhs_iter->dom, rhs_iter->dom)); + } // handle expr values if (lhs.second.as() && rhs.second.as()) { return Dispatch(lhs.second.as_or_throw(), rhs.second.as_or_throw()); diff --git a/src/s_tir/schedule/primitive/annotate.cc b/src/s_tir/schedule/primitive/annotate.cc index 9fd1bc275246..29efcd0c6bf1 100644 --- a/src/s_tir/schedule/primitive/annotate.cc +++ b/src/s_tir/schedule/primitive/annotate.cc @@ -27,6 +27,8 @@ using namespace tvm::tirx; void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_key, const Any& ann_val) { + TVM_FFI_CHECK(!sref->StmtAs() || ann_key != tirx::attr::thread_binding, ValueError) + << "thread_binding is a semantic annotation; use Schedule.bind to set it"; // Extract annotation const ffi::Map* annotations = nullptr; if (const auto* loop = sref->StmtAs()) { @@ -47,6 +49,8 @@ void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_k if (const auto* loop = sref->StmtAs()) { ffi::ObjectPtr n = ffi::make_object(*loop); n->annotations = std::move(new_ann); + // Validate semantic loop annotations before installing the replacement. + n->GetThreadBinding(); self->Replace(sref, For(n), {}); } else if (const auto* block = sref->StmtAs()) { ffi::ObjectPtr n = ffi::make_object(*block); @@ -60,6 +64,8 @@ void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_k } void Unannotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_key) { + TVM_FFI_CHECK(!sref->StmtAs() || ann_key != tirx::attr::thread_binding, ValueError) + << "thread_binding is a semantic annotation; use Schedule.parallel to remove it"; // Extract annotation const ffi::Map* annotations = nullptr; if (const auto* loop = sref->StmtAs()) { @@ -78,6 +84,8 @@ void Unannotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann if (const auto* loop = sref->StmtAs()) { ffi::ObjectPtr n = ffi::make_object(*loop); n->annotations = std::move(new_ann); + // Validate semantic loop annotations before installing the replacement. + n->GetThreadBinding(); self->Replace(sref, For(n), {}); } else if (const auto* block = sref->StmtAs()) { ffi::ObjectPtr n = ffi::make_object(*block); diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index 6b815827878c..561fca780ec5 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -759,7 +759,7 @@ class BlockizeRewriter : public StmtExprMutator { UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { if (loop == lca_->stmt) { return For(loop->loop_var, loop->min, loop->extent, loop->kind, RewriteSeq(loop->body), - loop->thread_binding, loop->annotations, loop->step, loop->span); + loop->GetThreadBinding(), loop->annotations, loop->step, loop->span); } return StmtExprMutator::Mutate_(loop, inplace_mode); } diff --git a/src/s_tir/schedule/primitive/compute_inline.cc b/src/s_tir/schedule/primitive/compute_inline.cc index b93e29f5b812..c8e900a917d9 100644 --- a/src/s_tir/schedule/primitive/compute_inline.cc +++ b/src/s_tir/schedule/primitive/compute_inline.cc @@ -1796,7 +1796,7 @@ class SingleBlockFusionReplacer : public StmtExprMutator { } return For(loop->loop_var, loop->min, loop->extent, loop->kind, mutated_body, - loop->thread_binding, loop->annotations); + loop->GetThreadBinding(), loop->annotations); } UnchangedOr Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final { diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc b/src/s_tir/schedule/primitive/decompose_padding.cc index 0ad27fc00f73..4d4d6efc70bc 100644 --- a/src/s_tir/schedule/primitive/decompose_padding.cc +++ b/src/s_tir/schedule/primitive/decompose_padding.cc @@ -374,7 +374,7 @@ static std::pair CreateInBoundBlock(const SBlockRealizeNode PrimExpr min = it == new_loop_ranges.end() ? loop->min : (*it).second->min; PrimExpr extent = it == new_loop_ranges.end() ? loop->extent : (*it).second->extent; nest_stmt_root = For(loop->loop_var, min, extent, loop->kind, nest_stmt_root, - loop->thread_binding, loop->annotations, loop->step, loop->span); + loop->GetThreadBinding(), loop->annotations, loop->step, loop->span); if (loop.same_as(highest_pos_inclusive)) { break; } diff --git a/src/s_tir/schedule/primitive/for_kind.cc b/src/s_tir/schedule/primitive/for_kind.cc index 558aff410d43..cd1346b80b8b 100644 --- a/src/s_tir/schedule/primitive/for_kind.cc +++ b/src/s_tir/schedule/primitive/for_kind.cc @@ -28,11 +28,11 @@ using namespace tvm::tirx; class WrongBlockIterTypeError : public ScheduleErrorContextObj { public: - explicit WrongBlockIterTypeError(IRModule mod, ForKind for_kind, Var loop_var, SBlock block) + explicit WrongBlockIterTypeError(IRModule mod, ForKind for_kind, bool is_thread_binding, + Var loop_var, SBlock block) : mod_(std::move(mod)), loop_var_(std::move(loop_var)), block_(std::move(block)) { - op_str_ = for_kind == ForKind::kParallel - ? "parallel" - : (for_kind == ForKind::kVectorized ? "vectorize" : "bind"); + op_str_ = + is_thread_binding ? "bind" : (for_kind == ForKind::kParallel ? "parallel" : "vectorize"); } ffi::String FastErrorString() const final { std::ostringstream os; @@ -73,7 +73,7 @@ class WrongBlockIterTypeError : public ScheduleErrorContextObj { * - the block iter is a reduction block iter, and the input `thread_tag` starts with "threadIdx" * in case of cross-thread reduction. * \param self The schedule state - * \param for_kind The desired ForKind (only `kParallel`, `kVectorized` and `kThreadBinding` are + * \param for_kind The desired ForKind (only `kParallel` and `kVectorized` are * allowed) * \param loop_var The loop variable of the loop to be checked * \param block_realize The block-realize of the block to be checked @@ -112,7 +112,8 @@ void CheckLoopParallelizableInBlock(const ScheduleState& self, ForKind for_kind, IterVarType iter_type = iter_var->iter_type; if (!(iter_type == kDataPar || (iter_type == kCommReduce && thread_scope.rank == 1 && thread_scope.dim_index != -1))) { - throw MakeScheduleError(self->mod, for_kind, loop_var, block); + throw MakeScheduleError(self->mod, for_kind, thread_scope.rank != -1, + loop_var, block); } } } @@ -122,7 +123,7 @@ void CheckLoopParallelizableInBlock(const ScheduleState& self, ForKind for_kind, * parallelized/vectorized/bound with regard to the block * \param self The schedule state * \param loop The loop to be parallelized/vectorized/bound - * \param for_kind The desired ForKind (only `kParallel`, `kVectorized` and `kThreadBinding` are + * \param for_kind The desired ForKind (only `kParallel` and `kVectorized` are * allowed) * \param thread_scope The thread scope of the thread axis to be bound, which is an invalid value if * the operation is not "bind" @@ -145,10 +146,10 @@ void CheckParallelizability(const ScheduleState& self, const For& loop, ForKind * \brief The implementation of parallelizing/vectorizing/binding a given loop * \param self The schedule state * \param loop_sref The sref of the loop to be parallelized/vectorized/bound - * \param for_kind The type of the operation (only `kParallel`, `kVectorized` and `kThreadBinding` + * \param for_kind The type of the operation (only `kParallel` and `kVectorized` * are allowed) * \param thread_axis The thread axis that the input loop is bound to, which is defined only when - * `for_kind` is `kThreadBinding` + * binding to a thread */ void ParallelizeComputation(const ScheduleState& self, const StmtSRef& loop_sref, ForKind for_kind, ffi::Optional thread_axis) { @@ -178,14 +179,12 @@ void ParallelizeComputation(const ScheduleState& self, const StmtSRef& loop_sref ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->kind = for_kind; if (thread_axis.has_value()) { - new_loop->thread_binding = IterVar(/*dom=*/Range(nullptr), // - /*var=*/ - PrimVar(thread_axis.value(), // - loop->loop_var.ty()), // - /*iter_type=*/kThreadIndex, // - /*thread_tag=*/thread_axis.value()); + new_loop->SetThreadBinding(IterVar(/*dom=*/Range(nullptr), + /*var=*/PrimVar(thread_axis.value(), loop->loop_var.ty()), + /*iter_type=*/kThreadIndex, + /*thread_tag=*/thread_axis.value())); } else { - new_loop->thread_binding = std::nullopt; + new_loop->SetThreadBinding(std::nullopt); } self->Replace(loop_sref, For(new_loop), {}); } @@ -199,14 +198,14 @@ void Vectorize(ScheduleState self, const StmtSRef& loop_sref) { } void Bind(ScheduleState self, const StmtSRef& loop_sref, const ffi::String& thread_axis) { - ParallelizeComputation(self, loop_sref, ForKind::kThreadBinding, thread_axis); + ParallelizeComputation(self, loop_sref, ForKind::kParallel, thread_axis); } void Unroll(ScheduleState self, const StmtSRef& loop_sref) { const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->kind = ForKind::kUnrolled; - new_loop->thread_binding = std::nullopt; + new_loop->SetThreadBinding(std::nullopt); self->Replace(loop_sref, For(new_loop), {}); } diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc index 57e98f5bf3cd..f386fb6af3bf 100644 --- a/src/s_tir/schedule/primitive/loop_transformation.cc +++ b/src/s_tir/schedule/primitive/loop_transformation.cc @@ -437,7 +437,7 @@ ffi::Array Split(ScheduleState self, const StmtSRef& loop_sref, // order with before. // Step 1. Check correctness const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); - if (!loop->annotations.empty() || loop->thread_binding.has_value()) { + if (!loop->annotations.empty() || loop->GetThreadBinding().has_value()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } // Currently, loops not starting with 0 are not supported @@ -712,7 +712,7 @@ const ffi::String get_sblock_name(Stmt loop_body) { ffi::Array LoopPartition(ScheduleState self, const StmtSRef& loop_sref, const ffi::Array& factors, bool preserve_unit_iters) { const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); - if (!loop->annotations.empty() || loop->thread_binding.has_value()) { + if (!loop->annotations.empty() || loop->GetThreadBinding().has_value()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } @@ -895,7 +895,7 @@ StmtSRef Merge(ScheduleState self, const ffi::Array& loop_srefs) { std::vector nest_loop_i_loops; for (auto p = sref.get(); p != lca.get(); p = p->parent) { if (auto loop = p->StmtAs()) { - if (!loop->annotations.empty() || loop->thread_binding.has_value()) { + if (!loop->annotations.empty() || loop->GetThreadBinding().has_value()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } @@ -962,7 +962,7 @@ StmtSRef Fuse(ScheduleState self, const ffi::Array& loop_srefs, // Step 1. check correctness for (const StmtSRef& sref : loop_srefs) { const ForNode* loop = TVM_SREF_TO_FOR(sref); - if (!loop->annotations.empty() || loop->thread_binding.has_value()) { + if (!loop->annotations.empty() || loop->GetThreadBinding().has_value()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } if (outer_loop_sref.defined()) { diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc index 5caa524fb978..064c5e947460 100644 --- a/src/s_tir/schedule/primitive/reduction.cc +++ b/src/s_tir/schedule/primitive/reduction.cc @@ -310,7 +310,7 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref, Var old_loop_var = old_loop->loop_var; PrimVar new_loop_var = old_loop->loop_var.CopyWithSuffix("_init"); loop_var_map[old_loop_var] = new_loop_var; - ffi::Optional opt_thread_binding = old_loop->thread_binding; + ffi::Optional opt_thread_binding = old_loop->GetThreadBinding(); if (opt_thread_binding) { auto thread_binding = opt_thread_binding.value(); auto new_var = thread_binding->var.CopyWithSuffix(""); @@ -319,7 +319,7 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref, } auto new_loop = old_loop.CopyOnWrite(); new_loop->loop_var = new_loop_var; - new_loop->thread_binding = opt_thread_binding; + new_loop->SetThreadBinding(opt_thread_binding); new_loop->body = body; body = ffi::GetRef(new_loop); } diff --git a/src/s_tir/schedule/utils.h b/src/s_tir/schedule/utils.h index cb6704c8a8e1..f68512eabea8 100644 --- a/src/s_tir/schedule/utils.h +++ b/src/s_tir/schedule/utils.h @@ -179,8 +179,8 @@ inline IterVar IterVarFromLoop(const For& loop, ffi::String name, IterVarType it * \return The thread scope bound to the loop */ inline runtime::ThreadScope GetThreadScope(const ForNode* loop) { - if (loop->kind == ForKind::kThreadBinding) { - return runtime::ThreadScope::Create(loop->thread_binding.value()->thread_tag); + if (loop->IsThreadBinding()) { + return runtime::ThreadScope::Create(loop->GetThreadBinding().value()->thread_tag); } return runtime::ThreadScope{-1, -1}; } diff --git a/src/s_tir/transform/compact_buffer_region.cc b/src/s_tir/transform/compact_buffer_region.cc index d4a947144744..a74b8ff72f5e 100644 --- a/src/s_tir/transform/compact_buffer_region.cc +++ b/src/s_tir/transform/compact_buffer_region.cc @@ -186,10 +186,9 @@ class BufferAccessRegionCollector : public StmtExprVisitor { ffi::Optional Visit_(const ForNode* op) final { Range loop_range = Range::FromMinExtent(op->min, op->extent); - IterVar iter = op->kind == ForKind::kThreadBinding - ? IterVar(Range(), op->loop_var, IterVarType::kThreadIndex, - op->thread_binding.value()->thread_tag) - : IterVar(Range(), op->loop_var, IterVarType::kDataPar); + IterVar iter = op->IsThreadBinding() ? IterVar(Range(), op->loop_var, IterVarType::kThreadIndex, + op->GetThreadBinding().value()->thread_tag) + : IterVar(Range(), op->loop_var, IterVarType::kDataPar); ancestor_iters_.push_back(iter); dom_analyzer_->Bind(op->loop_var, loop_range); dom_map_.emplace(op->loop_var.get(), sym::IntSet::FromRange(loop_range)); diff --git a/src/s_tir/transform/default_gpu_schedule.cc b/src/s_tir/transform/default_gpu_schedule.cc index 35686953f328..1642a4b931e1 100644 --- a/src/s_tir/transform/default_gpu_schedule.cc +++ b/src/s_tir/transform/default_gpu_schedule.cc @@ -40,7 +40,7 @@ void ThreadBind(s_tir::Schedule sch, const s_tir::SBlockRV& block, int64_t max_t ffi::Array loops = sch->GetLoops(block); for (const s_tir::LoopRV& loop : loops) { // skip block if already scheduled - if (sch->Get(loop)->thread_binding.has_value()) { + if (sch->Get(loop)->GetThreadBinding().has_value()) { return; } } diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc index a47af7fe022e..a08a1cee5060 100644 --- a/src/s_tir/transform/inject_software_pipeline.cc +++ b/src/s_tir/transform/inject_software_pipeline.cc @@ -1174,6 +1174,8 @@ class PipelineInjector : public StmtExprMutator { if (!HasPipelineAnnotation(op)) { return for_node; } + TVM_FFI_CHECK(!for_node->IsThreadBinding(), ValueError) + << "Software pipelining cannot replace a thread-bound loop"; // Step 2: Find the body and buffer allocations of the pipeline. The body can be direct child of // the for-loop. If the for-loop has BlockRealize as its child, the pipeline body will be the // child of the block. diff --git a/src/s_tir/transform/lift_thread_binding.cc b/src/s_tir/transform/lift_thread_binding.cc index c252563b7f8a..3ebf273831e1 100644 --- a/src/s_tir/transform/lift_thread_binding.cc +++ b/src/s_tir/transform/lift_thread_binding.cc @@ -48,7 +48,7 @@ FindLoopLCA(const Stmt& root) { ffi::Optional Visit_(const ForNode* op) final { stack.push_back(ffi::GetRef(op)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); - if (op->kind == ForKind::kThreadBinding) { + if (op->IsThreadBinding()) { UpdateLCA(op); } stack.pop_back(); @@ -56,21 +56,23 @@ FindLoopLCA(const Stmt& root) { } void UpdateLCA(const ForNode* loop) { - std::string thread_tag = loop->thread_binding.value()->thread_tag; + std::string thread_tag = loop->GetThreadBinding().value()->thread_tag; { ffi::Map* tgt = &annotations[thread_tag]; for (const auto& kv : loop->annotations) { - tgt->Set(kv.first, kv.second); + if (kv.first != tirx::attr::thread_binding) { + tgt->Set(kv.first, kv.second); + } } } IterVar& iter_var = iters[thread_tag]; if (!iter_var.defined()) { iter_var = IterVar(Range::FromMinExtent(loop->min, loop->extent), // loop->loop_var - .as_or_throw() // - .CopyWithName(thread_tag) // - .as_or_throw(), // - loop->thread_binding.value()->iter_type, // + .as_or_throw() // + .CopyWithName(thread_tag) // + .as_or_throw(), // + loop->GetThreadBinding().value()->iter_type, // thread_tag); lca[thread_tag] = stack; var_subst.Set(loop->loop_var, iter_var->var); @@ -133,7 +135,7 @@ class ThreadBindingLifter : public StmtExprMutator { UnchangedOr Mutate_(const ForNode* _op, InplaceMode inplace_mode) final { For op = ffi::GetRef(_op); bool is_kernel_root = false; - if (op->kind == ForKind::kThreadBinding) { + if (op->IsThreadBinding()) { if (iter_lca.empty()) { is_kernel_root = true; SetKernelRoot(_op); @@ -145,8 +147,8 @@ class ThreadBindingLifter : public StmtExprMutator { Stmt body = std::move(new_op.CopyOnWrite()->body); if (auto it = iter_lca.find(op); it != iter_lca.end()) { for (const auto& [iter_var, annotation] : it->second) { - body = For(iter_var->var, iter_var->dom->min, iter_var->dom->extent, - ForKind::kThreadBinding, std::move(body), + body = For(iter_var->var, iter_var->dom->min, iter_var->dom->extent, ForKind::kParallel, + std::move(body), IterVar(Range(nullptr), PrimVar(iter_var->thread_tag, iter_var->var.ty()), kThreadIndex, iter_var->thread_tag), annotation, std::nullopt); @@ -155,7 +157,7 @@ class ThreadBindingLifter : public StmtExprMutator { if (is_kernel_root) { iter_lca.clear(); } - if (op->kind == ForKind::kThreadBinding) { + if (op->IsThreadBinding()) { return body; } else { new_op.CopyOnWrite()->body = std::move(body); diff --git a/src/s_tir/transform/loop_partition.cc b/src/s_tir/transform/loop_partition.cc index 203276e4b9c8..b56a26bb1fd8 100644 --- a/src/s_tir/transform/loop_partition.cc +++ b/src/s_tir/transform/loop_partition.cc @@ -864,7 +864,7 @@ inline Stmt LoopPartitioner::MakeFor(const ffi::Object* node, PrimExpr extent, S }; return ffi::StructuralMap(body, f_substitute).as_or_throw(); } else { - TVM_FFI_ICHECK(for_node->kind != ForKind::kThreadBinding); + TVM_FFI_ICHECK(!for_node->IsThreadBinding()); auto new_loop = ffi::make_object(*for_node); new_loop->min = IntImm(for_node->min.ty(), 0); new_loop->extent = extent; diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc index fbac8846816b..2d1629a3df3d 100644 --- a/src/s_tir/transform/lower_cross_thread_reduction.cc +++ b/src/s_tir/transform/lower_cross_thread_reduction.cc @@ -65,11 +65,11 @@ struct ThreadScopeEqual { * \return True if the loop is bound to threadIdx.x/y/z */ bool IsBoundToThreadIdx(const ForNode* loop) { - if (!loop->thread_binding.has_value()) { + if (!loop->GetThreadBinding().has_value()) { return false; } runtime::ThreadScope scope = - runtime::ThreadScope::Create(loop->thread_binding.value()->thread_tag); + runtime::ThreadScope::Create(loop->GetThreadBinding().value()->thread_tag); return scope.rank == 1 && scope.dim_index >= 0; } @@ -280,7 +280,7 @@ class InThreadReducerMaker : public StmtExprMutator { .ValueOrUnchanged(ffi::GetRef(loop)) .as()) { For res = *opt_res; - if (res->thread_binding.has_value()) { + if (res->GetThreadBinding().has_value()) { if (!res->body.defined() || UnderLoopReductionBlockVarCollector::CheckHasReductionBlocks(res)) { return res->body; @@ -418,7 +418,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, } // Next arguments: all the reduction threads for (const ForNode* reduction_loop : reduction_loops) { - if (reduction_loop->thread_binding.has_value()) { + if (reduction_loop->GetThreadBinding().has_value()) { parameters.push_back(reduction_loop->loop_var); } } @@ -546,7 +546,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, ffi::StructuralWalk(realize->predicate, walk_fn); if (wb_buffers[0].scope() != "local") { for (const ForNode* loop : reduction_loops) { - if (loop->thread_binding.has_value()) { + if (loop->GetThreadBinding().has_value()) { wb_predicate = wb_predicate && (static_cast(loop->loop_var) == IntImm(loop->loop_var.ty(), 0)); } @@ -567,7 +567,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, Stmt new_stmt = SeqStmt::Flatten(std::move(stmts)); for (auto rit = reduction_loops.rbegin(); rit != reduction_loops.rend(); ++rit) { const ForNode* loop = *rit; - if (loop->thread_binding.has_value()) { + if (loop->GetThreadBinding().has_value()) { ffi::ObjectPtr n = ffi::make_object(*loop); n->body = std::move(new_stmt); new_stmt = For(n); @@ -614,7 +614,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // Step 3. Collect the loop. reduction_loops.push_back(loop); // Step 4. See whether the loop is bound to some thread axis. - if (loop->thread_binding.has_value()) { + if (loop->GetThreadBinding().has_value()) { need = true; } } @@ -656,8 +656,8 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // Erase those threads which are not free to this block. for (const ForNode* loop : loop_stack_) { - if (loop->thread_binding.has_value()) { - ThreadScope scope = ThreadScope::Create(loop->thread_binding.value()->thread_tag); + if (loop->GetThreadBinding().has_value()) { + ThreadScope scope = ThreadScope::Create(loop->GetThreadBinding().value()->thread_tag); thread2range.erase(scope); } } @@ -711,7 +711,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // bound to `threadIdx.x/y/z`. int n_bound_reduction_loops = 0; for (const ForNode* reduction_loop : reduction_loops) { - if (reduction_loop->thread_binding.has_value()) { + if (reduction_loop->GetThreadBinding().has_value()) { ++n_bound_reduction_loops; TVM_FFI_CHECK(IsBoundToThreadIdx(reduction_loop), ValueError) << "Cross-thread reduction requires all the reduction-related loops that " @@ -794,8 +794,8 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // - we are careful about thread block boundary for safety. bool is_block_idx = false; bool is_thread_idx = false; - if (loop->kind == ForKind::kThreadBinding) { - ThreadScope scope = ThreadScope::Create(loop->thread_binding.value()->thread_tag); + if (loop->IsThreadBinding()) { + ThreadScope scope = ThreadScope::Create(loop->GetThreadBinding().value()->thread_tag); if (scope.rank == 1 && scope.dim_index >= 0) { is_thread_idx = true; ++thread_idx_depth; @@ -891,9 +891,9 @@ class CrossThreadReductionTransformer : public StmtExprMutator { std::vector> reduction_threads; reduction_threads.reserve(reduction_loops.size()); for (const ForNode* loop : reduction_loops) { - if (loop->thread_binding.has_value()) { + if (loop->GetThreadBinding().has_value()) { reduction_threads.emplace_back( - ThreadScope::Create(loop->thread_binding.value()->thread_tag), + ThreadScope::Create(loop->GetThreadBinding().value()->thread_tag), Range::FromMinExtent(loop->min, loop->extent)); } } @@ -929,7 +929,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { /*loop_var=*/loop_vars[i].as_or_throw(), // /*min=*/unbound_thread2range[i].second->min, // /*extent=*/unbound_thread2range[i].second->extent, // - /*kind=*/ForKind::kThreadBinding, // + /*kind=*/ForKind::kParallel, // /*body=*/body, // /*thread_binding=*/ IterVar(Range(), PrimVar("", loop_vars[i]->ty.as_or_throw()), diff --git a/src/s_tir/transform/lower_opaque_block.cc b/src/s_tir/transform/lower_opaque_block.cc index c71e964d57f1..eb05ac08b857 100644 --- a/src/s_tir/transform/lower_opaque_block.cc +++ b/src/s_tir/transform/lower_opaque_block.cc @@ -94,7 +94,8 @@ class OpaqueBlockLower : public StmtExprMutator { // Step 1. Update unit loop info. PrimExpr min = this->Mutate(op->min, inplace_mode).ValueOrUnchanged(op->min); PrimExpr extent = this->Mutate(op->extent, inplace_mode).ValueOrUnchanged(op->extent); - if (is_one(extent) && op->annotations.empty()) { + bool has_only_thread_binding = op->IsThreadBinding() && op->annotations.size() == 1; + if (is_one(extent) && (op->annotations.empty() || has_only_thread_binding)) { // handling unit loop VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), min)); } @@ -107,10 +108,10 @@ class OpaqueBlockLower : public StmtExprMutator { ffi::Map new_annotations = HandleAnnotations(op->annotations, &pragma_attrs, /*is_block=*/false); // Step 4. Create new For loop accordingly - if (op->kind == ForKind::kThreadBinding) { + if (op->IsThreadBinding()) { // Case 1. Thread binding - TVM_FFI_ICHECK(op->thread_binding.has_value()); - ffi::String thread_tag = op->thread_binding.value()->thread_tag; + TVM_FFI_ICHECK(op->GetThreadBinding().has_value()); + ffi::String thread_tag = op->GetThreadBinding().value()->thread_tag; body = MakeLaunchThread(min, extent, op->loop_var, thread_tag, body); } else if (is_one(extent) && op->annotations.empty() && !op->annotations.count(s_tir::attr::irregular_loop_mark)) { diff --git a/src/s_tir/transform/memhammer_coalesce.cc b/src/s_tir/transform/memhammer_coalesce.cc index d141c70a0ac6..de6c0e56b77b 100644 --- a/src/s_tir/transform/memhammer_coalesce.cc +++ b/src/s_tir/transform/memhammer_coalesce.cc @@ -134,11 +134,10 @@ Stmt SplitBindVectorize(const Stmt& stmt, const ConstraintSet& constraints) { body = For(new_loop_vars.back().as_or_throw(), 0, vector_len, ForKind::kVectorized, std::move(body)); for (int i = n - 2; i >= 1; i--) { - body = - For(new_loop_vars[i].as_or_throw(), 0, factors[i], ForKind::kThreadBinding, - std::move(body), - IterVar(Range(nullptr), PrimVar(thread_axis[i - 1]), kThreadIndex, thread_axis[i - 1]), - {}, std::nullopt); + body = For( + new_loop_vars[i].as_or_throw(), 0, factors[i], ForKind::kParallel, std::move(body), + IterVar(Range(nullptr), PrimVar(thread_axis[i - 1]), kThreadIndex, thread_axis[i - 1]), {}, + std::nullopt); } return For(new_loop_vars[0].as_or_throw(), 0, factors[0], ForKind::kSerial, std::move(body)); diff --git a/src/s_tir/transform/memhammer_intermediate_stage.cc b/src/s_tir/transform/memhammer_intermediate_stage.cc index baf6082cf617..270e80798a61 100644 --- a/src/s_tir/transform/memhammer_intermediate_stage.cc +++ b/src/s_tir/transform/memhammer_intermediate_stage.cc @@ -51,7 +51,7 @@ std::pair LiftThreadBindingLoops(Stmt stmt) { std::vector thread_binding_loops; Stmt body = stmt; while (const ForNode* loop = body.as()) { - if (loop->kind == ForKind::kThreadBinding) { + if (loop->IsThreadBinding()) { thread_binding_loops.push_back(loop); } else { normal_loops.push_back(loop); @@ -274,8 +274,8 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S body = op->then_case; } for (const For& loop : outer_loops) { - if (loop->kind == ForKind::kThreadBinding) { - const ffi::String& thread_tag = loop->thread_binding.value()->thread_tag; + if (loop->IsThreadBinding()) { + const ffi::String& thread_tag = loop->GetThreadBinding().value()->thread_tag; auto thread_scope = runtime::ThreadScope::Create(thread_tag); if (CanRelaxStorageUnderThread(runtime::StorageScope::Create(storage_scope), thread_scope)) { if (is_write_cache && thread_scope.dim_index == 0) { @@ -430,7 +430,7 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S new_loop->loop_var = new_loop_vars[i]; new_loop->body = generate_body; new_loop->kind = ForKind::kSerial; - new_loop->thread_binding = std::nullopt; + new_loop->SetThreadBinding(std::nullopt); new_loop->annotations = {}; generate_body = For(new_loop); } diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc index f5c3771e58f4..482c0ec940ce 100644 --- a/src/s_tir/transform/memhammer_lower_auto_copy.cc +++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc @@ -516,11 +516,11 @@ class AutoPadder { } ffi::Optional Visit_(const ForNode* op) final { - if (op->kind != ForKind::kThreadBinding) { + if (!op->IsThreadBinding()) { substitute_map_.Set(op->loop_var, op->min); } else { int64_t extent = - warp_thread_extent_.Get(op->thread_binding.value()->thread_tag).value_or(1); + warp_thread_extent_.Get(op->GetThreadBinding().value()->thread_tag).value_or(1); var_range_.Set(op->loop_var, Range::FromMinExtent(op->min, IntImm::Int64(extent))); } if (op->kind == ForKind::kVectorized) { @@ -531,7 +531,7 @@ class AutoPadder { if (op->kind == ForKind::kVectorized) { vector_length_ = -1; } - if (op->kind != ForKind::kThreadBinding) { + if (!op->IsThreadBinding()) { substitute_map_.erase(op->loop_var); } return std::nullopt; @@ -828,9 +828,10 @@ class ThreadExtentCollector : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } ffi::Optional Visit_(const ForNode* op) final { - if (op->thread_binding.has_value() && op->thread_binding.value()->iter_type == kThreadIndex) { + if (op->GetThreadBinding().has_value() && + op->GetThreadBinding().value()->iter_type == kThreadIndex) { if (const auto* extent = op->extent.as()) { - thread_extent_.Set(op->thread_binding.value()->thread_tag, + thread_extent_.Set(op->GetThreadBinding().value()->thread_tag, static_cast(extent->value)); } } diff --git a/src/s_tir/transform/profile_instrumentation.cc b/src/s_tir/transform/profile_instrumentation.cc index 458f0e06de92..fc604accce90 100644 --- a/src/s_tir/transform/profile_instrumentation.cc +++ b/src/s_tir/transform/profile_instrumentation.cc @@ -55,7 +55,7 @@ struct LoopInfo { int32_t depth; int32_t height; bool has_siblings; - // Set to 'true' if ForKind::kParallel is set for the current loop or one of its ancestor + // Set to 'true' for CPU parallel loops and their descendants bool has_parallel; }; @@ -92,7 +92,7 @@ class LoopAnalyzer : public StmtExprVisitor { if (has_parallel) { loop_info.has_parallel = true; parent_parallel = true; - } else if (f->kind == ForKind::kParallel) { + } else if (f->IsParallel()) { // has_parallel for the current loop is being set to 'false' since the // intrinsic is added outside of the loop. The instrumentation isn't // allowed for the subsequent nested loops. @@ -127,7 +127,7 @@ class LoopAnalyzer : public StmtExprVisitor { if (has_parallel) { loop_info.has_parallel = true; parent_parallel = true; - } else if (f->kind == ForKind::kParallel) { + } else if (f->IsParallel()) { // has_parallel for the current loop is being set to 'false' since the // intrinsic is added outside of the loop. The instrumentation isn't // allowed for the subsequent nested loops. @@ -235,7 +235,7 @@ class CheckParallelLoops : public StmtExprVisitor { private: ffi::Optional Visit_(const ForNode* op) final { - if (op->kind == ForKind::kParallel) { + if (op->IsParallel()) { has_parallel = true; } else { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); diff --git a/src/s_tir/transform/unify_thread_binding.cc b/src/s_tir/transform/unify_thread_binding.cc index 62694d603e86..f0a914d1fe0b 100644 --- a/src/s_tir/transform/unify_thread_binding.cc +++ b/src/s_tir/transform/unify_thread_binding.cc @@ -71,18 +71,21 @@ class ThreadBindingUnifier : public StmtExprMutator { UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { // If this For is not thread binding attribute, return as usual. - if (op->kind != ForKind::kThreadBinding) { + if (!op->IsThreadBinding()) { return StmtExprMutator::Mutate_(op, inplace_mode); } ffi::Map annotations = op->annotations; - Stmt stmt = UnifyThreadBindingImpl(op, op->loop_var, op->thread_binding.value(), + annotations.erase(tirx::attr::thread_binding); + Stmt stmt = UnifyThreadBindingImpl(op, op->loop_var, op->GetThreadBinding().value(), Range::FromMinExtent(op->min, op->extent), inplace_mode); if (annotations.empty()) { return stmt; } if (const auto* loop = stmt.as()) { For new_loop = ffi::GetRef(loop); + ffi::Optional thread_binding = new_loop->GetThreadBinding(); new_loop.CopyOnWrite()->annotations = std::move(annotations); + new_loop.CopyOnWrite()->SetThreadBinding(thread_binding); return new_loop; } else { @@ -166,7 +169,7 @@ class ThreadBindingUnifier : public StmtExprMutator { // necessary for unit tests. result = For(thread_binding->var, thread_binding->dom->min, thread_binding->dom->extent, - ForKind::kThreadBinding, result, + ForKind::kParallel, result, IterVar(Range(), PrimVar(""), IterVarType::kThreadIndex, thread_binding->thread_tag), {}, std::nullopt); launch_threads_.pop_back(); diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc index fe43f3726df6..5011d62772e3 100644 --- a/src/target/llvm/codegen_cpu.cc +++ b/src/target/llvm/codegen_cpu.cc @@ -1192,7 +1192,7 @@ void CodeGenCPU::Dispatch_(const ForNode* op) { EmitDebugLocation(op); if (op->kind == ForKind::kSerial || op->kind == ForKind::kUnrolled) { CodeGenLLVM::Dispatch_(op); - } else if (op->kind == ForKind::kParallel) { + } else if (op->IsParallel()) { TVM_FFI_ICHECK(is_zero(op->min)) << "Parallel launch require canonical loop with zero start index"; TVM_FFI_ICHECK(op->HasTrivialStep()) diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc index 2de7d9c75a89..3497239a923f 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -465,11 +465,11 @@ UnchangedOr IndexDataTypeRewriter::Mutate_(const ForNode* op, InplaceMode n->loop_var = new_loop_var; n->min = prim::cast(new_loop_var.ty(), min); n->extent = prim::cast(new_loop_var.ty(), extent); - if (op->thread_binding.has_value()) { - auto old_thread_binding = op->thread_binding.value(); + if (op->GetThreadBinding().has_value()) { + auto old_thread_binding = op->GetThreadBinding().value(); auto* ptr = old_thread_binding.CopyOnWrite(); ptr->var = old_thread_binding->var.CopyWithDType(new_loop_var.ty()); - n->thread_binding = ffi::Optional(std::move(old_thread_binding)); + n->SetThreadBinding(ffi::Optional(std::move(old_thread_binding))); } n->body = new_body; return new_for; diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index 9ed4b8a4080d..7c5c9e7c76d3 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -210,7 +210,7 @@ TVMFFIAny AssertStmtMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, } TVMFFIAny ForVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - // skips: kind and constant annotations; unlike SBlock annotations, these carry no expressions. + // Skip the kind and auxiliary hints, but traverse the semantic thread binding. const ForNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( @@ -218,13 +218,13 @@ TVMFFIAny ForVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexc TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->min)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->extent)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->thread_binding)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->GetThreadBinding())); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->step)); return ffi::AnyView(nullptr).CopyToTVMFFIAny(); } TVMFFIAny ForMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - // skips: kind and constant annotations; unlike SBlock annotations, these carry no expressions. + // Skip the kind and auxiliary hints, but traverse the semantic thread binding. const ForNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_loop_var, @@ -238,13 +238,13 @@ TVMFFIAny ForMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noex TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_body, mutator->MutateExpected(self->body)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_thread_binding, - mutator->MutateExpected(self->thread_binding)); + mutator->MutateExpected(self->GetThreadBinding())); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_step, mutator->MutateExpected(self->step)); if (mapped_loop_var.UnchangedOrSameAs(self->loop_var) && mapped_min.UnchangedOrSameAs(self->min) && mapped_extent.UnchangedOrSameAs(self->extent) && mapped_body.UnchangedOrSameAs(self->body) && - mapped_thread_binding.UnchangedOrSameAs(self->thread_binding) && + mapped_thread_binding.UnchangedOrSameAs(self->GetThreadBinding()) && mapped_step.UnchangedOrSameAs(self->step)) { return ffi::Unchanged().CopyToTVMFFIAny(); } @@ -253,14 +253,14 @@ TVMFFIAny ForMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noex copy->min = std::move(mapped_min).ValueOrUnchanged(std::move(copy->min)); copy->extent = std::move(mapped_extent).ValueOrUnchanged(std::move(copy->extent)); copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body)); - copy->thread_binding = - std::move(mapped_thread_binding).ValueOrUnchanged(std::move(copy->thread_binding)); + copy->SetThreadBinding( + std::move(mapped_thread_binding).ValueOrUnchanged(copy->GetThreadBinding())); copy->step = std::move(mapped_step).ValueOrUnchanged(std::move(copy->step)); return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy))); } TVMFFIAny ForMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - // skips: kind and constant annotations; unlike SBlock annotations, these carry no expressions. + // Skip the kind and auxiliary hints, but traverse the semantic thread binding. ForNode* self = const_cast( ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_loop_var, @@ -277,13 +277,13 @@ TVMFFIAny ForMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView mutator->MutateExpected(self->body, ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr>, mapped_thread_binding, - mutator->MutateExpected(self->thread_binding, ffi::InplaceMode::kAllow)); + mutator->MutateExpected(self->GetThreadBinding(), ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_step, mutator->MutateExpected(self->step, ffi::InplaceMode::kAllow)); if (mapped_loop_var.UnchangedOrSameAs(self->loop_var) && mapped_min.UnchangedOrSameAs(self->min) && mapped_extent.UnchangedOrSameAs(self->extent) && mapped_body.UnchangedOrSameAs(self->body) && - mapped_thread_binding.UnchangedOrSameAs(self->thread_binding) && + mapped_thread_binding.UnchangedOrSameAs(self->GetThreadBinding()) && mapped_step.UnchangedOrSameAs(self->step)) { return ffi::Unchanged().CopyToTVMFFIAny(); } @@ -292,7 +292,7 @@ TVMFFIAny ForMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView if (!mapped_extent.IsUnchanged()) self->extent = std::move(mapped_extent).ValueUnchecked(); if (!mapped_body.IsUnchanged()) self->body = std::move(mapped_body).ValueUnchecked(); if (!mapped_thread_binding.IsUnchanged()) { - self->thread_binding = std::move(mapped_thread_binding).ValueUnchecked(); + self->SetThreadBinding(std::move(mapped_thread_binding).ValueUnchecked()); } if (!mapped_step.IsUnchanged()) self->step = std::move(mapped_step).ValueUnchecked(); return ffi::Unchanged().CopyToTVMFFIAny(); @@ -850,6 +850,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { For::For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, Stmt body, ffi::Optional thread_binding, ffi::Map annotations, ffi::Optional step, Span span) { + TVM_FFI_CHECK(kind >= ForKind::kSerial && kind <= ForKind::kUnrolled, ValueError) + << "Invalid ForKind: " << static_cast(kind); TVM_FFI_ICHECK(loop_var.defined()); TVM_FFI_ICHECK(min.defined()); TVM_FFI_ICHECK(extent.defined()); @@ -904,8 +906,16 @@ For::For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, Stmt bod node->extent = std::move(extent); node->kind = kind; node->body = std::move(body); - node->thread_binding = std::move(thread_binding); node->annotations = std::move(annotations); + if (thread_binding.has_value()) { + if (auto existing = node->GetThreadBinding()) { + TVM_FFI_CHECK(existing.value().same_as(thread_binding.value()), ValueError) + << "Conflicting thread_binding argument and annotation"; + } + node->SetThreadBinding(std::move(thread_binding)); + } else { + node->GetThreadBinding(); + } node->step = std::move(step); node->span = std::move(span); data_ = std::move(node); @@ -929,6 +939,30 @@ TVM_FFI_STATIC_INIT_BLOCK() { }); } +ffi::Optional ForNode::GetThreadBinding() const { + auto value = annotations.Get(attr::thread_binding); + if (!value.has_value()) return std::nullopt; + TVM_FFI_CHECK(kind == ForKind::kParallel, ValueError) + << "thread_binding is only valid on parallel loops"; + auto binding = value.value().as(); + TVM_FFI_CHECK(binding.has_value(), TypeError) << "thread_binding annotation must be an IterVar"; + TVM_FFI_CHECK(!binding.value()->thread_tag.empty(), ValueError) + << "thread_binding must have a nonempty thread tag"; + return binding; +} + +void ForNode::SetThreadBinding(ffi::Optional binding) { + if (binding.has_value()) { + TVM_FFI_CHECK(kind == ForKind::kParallel, ValueError) + << "thread_binding is only valid on parallel loops"; + TVM_FFI_CHECK(!binding.value()->thread_tag.empty(), ValueError) + << "thread_binding must have a nonempty thread tag"; + annotations.Set(attr::thread_binding, binding.value()); + } else { + annotations.erase(attr::thread_binding); + } +} + bool ForNode::HasTrivialStep() const { return !step.has_value() || is_one(*step); } std::ostream& operator<<(std::ostream& out, ForKind type) { // NOLINT(*) @@ -945,9 +979,6 @@ std::ostream& operator<<(std::ostream& out, ForKind type) { // NOLINT(*) case ForKind::kVectorized: out << "vectorized"; break; - case ForKind::kThreadBinding: - out << "launch_thread"; - break; } return out; } diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index f87d6d7cf57f..f7c6811afb26 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -586,7 +586,7 @@ ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, ffi::String thread, IterVar iter_var(Range(nullptr), tvm::PrimVar("iter", dtype), IterVarType::kThreadIndex, thread); return For(vars[0].as_or_throw(), doms[0]->min, doms[0]->extent, - ForKind::kThreadBinding, body, iter_var, + ForKind::kParallel, body, iter_var, annotations.value_or(ffi::Map()), std::nullopt); }; return ForFrame(n); diff --git a/src/tirx/script/printer/for_loop.cc b/src/tirx/script/printer/for_loop.cc index df1ec3a93395..afe33aa6132f 100644 --- a/src/tirx/script/printer/for_loop.cc +++ b/src/tirx/script/printer/for_loop.cc @@ -82,8 +82,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { min = d->AsDoc(loop->min, loop_p->Attr("min")); max = d->AsDoc(loop->min + loop->extent, loop_p->Attr("extent")); } - if (!loop->annotations.empty()) { - annotations = d->AsDoc(loop->annotations, loop_p->Attr("annotations")); + auto loop_annotations = loop->annotations; + loop_annotations.erase(tirx::attr::thread_binding); + if (!loop_annotations.empty()) { + annotations = d->AsDoc(loop_annotations, loop_p->Attr("annotations")); } bool use_range_sugar = false; ExprDoc prefix{ffi::UnsafeInit()}; @@ -94,16 +96,17 @@ TVM_FFI_STATIC_INIT_BLOCK() { } else { prefix = TIR(d, "serial"); } - } else if (loop->kind == tirx::ForKind::kParallel) { + } else if (loop->IsThreadBinding()) { + prefix = TIR(d, "thread_binding"); + thread = LiteralDoc::Str( + loop->GetThreadBinding().value()->thread_tag, + loop_p->Attr("annotations")->MapItem(tirx::attr::thread_binding)->Attr("thread_tag")); + } else if (loop->IsParallel()) { prefix = TIR(d, "parallel"); } else if (loop->kind == tirx::ForKind::kUnrolled) { prefix = TIR(d, "unroll"); } else if (loop->kind == tirx::ForKind::kVectorized) { prefix = TIR(d, "vectorized"); - } else if (loop->kind == tirx::ForKind::kThreadBinding) { - prefix = TIR(d, "thread_binding"); - thread = LiteralDoc::Str(loop->thread_binding.value()->thread_tag, - loop_p->Attr("thread_binding")); } else { TVM_FFI_THROW(ValueError) << "Unknown ForKind: " << tirx::ForKind2String(loop->kind); } @@ -125,14 +128,16 @@ TVM_FFI_STATIC_INIT_BLOCK() { // - annotations == {"disable_unroll": True}: print as unroll=False // - annotations == {"pragma_unroll": value}: print as unroll=value bool printed_as_unroll = false; - if (loop->annotations.size() == 1 && loop->annotations.count("disable_unroll")) { + if (!loop->IsThreadBinding() && loop_annotations.size() == 1 && + loop_annotations.count("disable_unroll")) { kwargs_keys.push_back("unroll"); kwargs_values.push_back(LiteralDoc::Boolean(false, loop_p->Attr("annotations"))); printed_as_unroll = true; - } else if (loop->annotations.size() == 1 && loop->annotations.count("pragma_unroll")) { + } else if (!loop->IsThreadBinding() && loop_annotations.size() == 1 && + loop_annotations.count("pragma_unroll")) { kwargs_keys.push_back("unroll"); kwargs_values.push_back( - d->AsDoc(loop->annotations["pragma_unroll"], loop_p->Attr("annotations"))); + d->AsDoc(loop_annotations["pragma_unroll"], loop_p->Attr("annotations"))); printed_as_unroll = true; } if (!printed_as_unroll) { diff --git a/src/tirx/transform/bind_target.cc b/src/tirx/transform/bind_target.cc index 75fd157dacb5..9ccc4ff74f63 100644 --- a/src/tirx/transform/bind_target.cc +++ b/src/tirx/transform/bind_target.cc @@ -102,7 +102,7 @@ class FunctionClassifierVisitor : public StmtExprVisitor { } ffi::Optional Visit_(const ForNode* op) final { - if (op->kind == ForKind::kThreadBinding) { + if (op->IsThreadBinding()) { // Enter GPU scope for thread binding loops bool last_is_under_gpu_scope = is_under_gpu_scope_; is_under_gpu_scope_ = true; @@ -191,7 +191,7 @@ class CallSubstitutor : public StmtExprMutator { } UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { - if (op->kind == ForKind::kThreadBinding) { + if (op->IsThreadBinding()) { // Enter GPU scope for thread binding loops bool last_is_under_gpu_scope = is_under_gpu_scope_; is_under_gpu_scope_ = true; diff --git a/src/tirx/transform/lower_tirx_opaque.cc b/src/tirx/transform/lower_tirx_opaque.cc index b91742035b67..49d4dde32531 100644 --- a/src/tirx/transform/lower_tirx_opaque.cc +++ b/src/tirx/transform/lower_tirx_opaque.cc @@ -58,7 +58,8 @@ class TIRxOpaqueLower : public StmtExprMutator { // Step 1. Update unit loop info. PrimExpr min = this->Mutate(op->min, inplace_mode).ValueOrUnchanged(op->min); PrimExpr extent = this->Mutate(op->extent, inplace_mode).ValueOrUnchanged(op->extent); - if (is_one(extent) && op->annotations.empty()) { + bool has_only_thread_binding = op->IsThreadBinding() && op->annotations.size() == 1; + if (is_one(extent) && (op->annotations.empty() || has_only_thread_binding)) { // handling unit loop VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), min)); } @@ -71,10 +72,10 @@ class TIRxOpaqueLower : public StmtExprMutator { ffi::Map new_annotations = HandleAnnotations(op->annotations, &pragma_attrs); // Step 4. Create new For loop accordingly - if (op->kind == ForKind::kThreadBinding) { + if (op->IsThreadBinding()) { // Case 1. Thread binding → AttrStmt(thread_extent) - TVM_FFI_ICHECK(op->thread_binding.has_value()); - ffi::String thread_tag = op->thread_binding.value()->thread_tag; + TVM_FFI_ICHECK(op->GetThreadBinding().has_value()); + ffi::String thread_tag = op->GetThreadBinding().value()->thread_tag; body = MakeLaunchThread(min, extent, op->loop_var, thread_tag, body); } else if (is_one(extent) && op->annotations.empty() && !op->annotations.count(s_tir::attr::irregular_loop_mark)) { diff --git a/src/tirx/transform/lower_tvm_builtin.cc b/src/tirx/transform/lower_tvm_builtin.cc index 9f14fe4d95fe..0a29fe754c5c 100644 --- a/src/tirx/transform/lower_tvm_builtin.cc +++ b/src/tirx/transform/lower_tvm_builtin.cc @@ -358,7 +358,7 @@ class BuiltinLower : public StmtExprMutator { PrimExpr extent = std::move(extent_result).ValueOrUnchanged(op->extent); Stmt body; - if (op->kind == ForKind::kParallel) { + if (op->IsParallel()) { body = this->VisitBodyAndRealizeAlloca(op->body); } else { body = scope_.WithNewScope([&]() -> Stmt { diff --git a/src/tirx/transform/make_packed_api.cc b/src/tirx/transform/make_packed_api.cc index 70cefb5602da..afb1f59b349f 100644 --- a/src/tirx/transform/make_packed_api.cc +++ b/src/tirx/transform/make_packed_api.cc @@ -57,10 +57,10 @@ class ReturnRewriter : public StmtExprMutator { explicit ReturnRewriter(Var ret_var) : ret_var_(ret_var) {} UnchangedOr Mutate_(const ForNode* node, InplaceMode inplace_mode) override { - if (node->kind == ForKind::kParallel) in_parallel_ += 1; + if (node->IsParallel()) in_parallel_ += 1; Stmt ret = StmtExprMutator::Mutate_(node, inplace_mode).ValueOrUnchanged(ffi::GetRef(node)); - if (node->kind == ForKind::kParallel) in_parallel_ -= 1; + if (node->IsParallel()) in_parallel_ -= 1; return ret; } diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc index d6a6d7bdd275..1d2875066e0e 100644 --- a/src/tirx/transform/storage_rewrite.cc +++ b/src/tirx/transform/storage_rewrite.cc @@ -694,7 +694,7 @@ class StoragePlanRewriter : public StmtExprMutator { StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); op = stmt.as(); return For(op->loop_var, op->min, op->extent, op->kind, MakeAttach(svec, op->body), - op->thread_binding, op->annotations, op->step); + op->GetThreadBinding(), op->annotations, op->step); } else { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -1096,7 +1096,7 @@ class StoragePlanRewriter : public StmtExprMutator { } } else if (s.stmt->IsInstance()) { const auto* op = static_cast(s.stmt); - if (op->kind == ForKind::kParallel) { + if (op->IsParallel()) { if (thread_scope_ == nullptr || thread_scope_ == op) { PlanNewScope(op); } diff --git a/src/tirx/transform/unroll_loop.cc b/src/tirx/transform/unroll_loop.cc index 647c8da50697..5482ee8827c7 100644 --- a/src/tirx/transform/unroll_loop.cc +++ b/src/tirx/transform/unroll_loop.cc @@ -134,6 +134,11 @@ class LoopUnroller : public StmtExprMutator { inplace_mode = InplaceMode::kDisallow; } } + // A thread-bound loop defines an execution scope even when its extent is one. + if (op->IsThreadBinding()) { + normal_loop_depth_ += 1; + return result; + } int value = GetExtent(op); // condition for auto unroll bool auto_unroll = (op->kind == ForKind::kSerial && value >= 0 && normal_loop_depth_ == 0 && diff --git a/tests/python/relax/test_pipeline.py b/tests/python/relax/test_pipeline.py index dd8a15f8194f..f81b4ca6cc90 100644 --- a/tests/python/relax/test_pipeline.py +++ b/tests/python/relax/test_pipeline.py @@ -173,7 +173,7 @@ def _has_thread_binding(func: tvm.tirx.PrimFunc) -> bool: def _visit(node): nonlocal found - if isinstance(node, tvm.tirx.For) and node.kind == tvm.tirx.ForKind.THREAD_BINDING: + if isinstance(node, tvm.tirx.For) and node.is_thread_binding(): found = True tvm_ffi.structural_walk(func.body, _visit) diff --git a/tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py b/tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py new file mode 100644 index 000000000000..67505d0fae14 --- /dev/null +++ b/tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py @@ -0,0 +1,119 @@ +# 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. +"""Schedule mutations validate the semantic thread-binding annotation atomically.""" + +import pytest + +import tvm +import tvm.testing +from tvm import tirx +from tvm.script import tirx as T + + +@T.prim_func(s_tir=True) +def copy(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")): + for i in range(16): + with T.sblock("copy"): + vi = T.axis.spatial(16, i) + B[vi] = A[vi] + + +def make_schedule(kind="parallel"): + sch = tvm.s_tir.Schedule(copy, debug_mask="all") + (loop,) = sch.get_loops(sch.get_sblock("copy")) + if kind != "serial": + getattr(sch, kind)(loop) + return sch, loop + + +def test_binding_changes_use_schedule_primitives(): + sch, loop = make_schedule() + sch.annotate(loop, "hint", 7) + sch.bind(loop, "threadIdx.x") + node = sch.get(loop) + assert node.kind == tirx.ForKind.PARALLEL + assert node.thread_binding.thread_tag == "threadIdx.x" + assert node.annotations["thread_binding"].same_as(node.thread_binding) + assert node.annotations["hint"] == 7 + + sch.parallel(loop) + node = sch.get(loop) + assert node.kind == tirx.ForKind.PARALLEL + assert node.thread_binding is None + assert "thread_binding" not in node.annotations + assert node.annotations["hint"] == 7 + + +@pytest.mark.parametrize("kind", ["parallel", "serial", "vectorize", "unroll"]) +@pytest.mark.parametrize("value", ["threadIdx.x", 1]) +def test_annotate_rejects_semantic_binding_key(kind, value): + sch, loop = make_schedule(kind) + before = sch.mod.script() + trace_before = str(sch.trace) + with pytest.raises(ValueError, match="use Schedule.bind"): + sch.annotate(loop, "thread_binding", value) + assert sch.mod.script() == before + assert str(sch.trace) == trace_before + assert sch.get(loop).thread_binding is None + + +def test_unannotate_rejects_semantic_binding_key(): + sch, loop = make_schedule() + sch.bind(loop, "threadIdx.x") + before = sch.mod.script() + trace_before = str(sch.trace) + with pytest.raises(ValueError, match="use Schedule.parallel"): + sch.unannotate(loop, "thread_binding") + assert sch.mod.script() == before + assert str(sch.trace) == trace_before + assert sch.get(loop).thread_binding.thread_tag == "threadIdx.x" + + +@T.prim_func(s_tir=True) +def reduce(A: T.Buffer((16,), "float32"), B: T.Buffer((1,), "float32")): + for i in range(16): + with T.sblock("sum"): + vi = T.axis.reduce(16, i) + with T.init(): + B[0] = T.float32(0) + B[0] = B[0] + A[vi] + + +def test_bound_reduction_cannot_bypass_parallel_legality(): + sch = tvm.s_tir.Schedule(reduce, debug_mask="all") + (loop,) = sch.get_loops(sch.get_sblock("sum")) + sch.bind(loop, "threadIdx.x") + before = sch.mod.script() + with pytest.raises(ValueError, match="use Schedule.parallel"): + sch.unannotate(loop, "thread_binding") + with pytest.raises(tvm.s_tir.ScheduleError): + sch.parallel(loop) + assert sch.mod.script() == before + assert sch.get(loop).thread_binding.thread_tag == "threadIdx.x" + + +def test_sblock_thread_binding_hint_is_not_reserved(): + sch, _ = make_schedule() + block = sch.get_sblock("copy") + sch.annotate(block, "thread_binding", "block_hint") + assert sch.get(block).annotations["thread_binding"] == "block_hint" + sch.unannotate(block, "thread_binding") + assert "thread_binding" not in sch.get(block).annotations + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/tirx-base/test_tir_for_thread_binding.py b/tests/python/tirx-base/test_tir_for_thread_binding.py new file mode 100644 index 000000000000..2e5e4e3065ca --- /dev/null +++ b/tests/python/tirx-base/test_tir_for_thread_binding.py @@ -0,0 +1,229 @@ +# 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. +"""Semantic thread binding stored in For annotations.""" + +import numpy as np +import pytest +import tvm_ffi + +import tvm +import tvm.testing +from tvm import tirx +from tvm.script import tirx as T +from tvm.testing import env + + +def make_binding(tag="threadIdx.x", dom=None, var=None): + return tirx.IterVar(dom, var if var is not None else "thread", tirx.IterVar.ThreadIndex, tag) + + +def make_loop(binding=None, annotations=None, kind=tirx.ForKind.PARALLEL, extent=8): + i = tirx.Var("i", "int32") + return tirx.For(i, 0, extent, kind, tirx.Evaluate(i), binding, annotations) + + +def test_annotation_representation_and_classification(): + binding = make_binding() + loop = make_loop(binding, {"custom_hint": 7}) + assert loop.kind == tirx.ForKind.PARALLEL + assert loop.annotations["thread_binding"].same_as(binding) + assert loop.thread_binding.same_as(binding) + assert int(loop.annotations["custom_hint"]) == 7 + assert loop.is_thread_binding() + assert not loop.is_parallel() + + annotated = make_loop(annotations={"thread_binding": binding}) + assert annotated.thread_binding.same_as(binding) + assert annotated.is_thread_binding() + parallel = make_loop() + assert parallel.thread_binding is None + assert parallel.is_parallel() + assert not parallel.is_thread_binding() + serial = make_loop(kind=tirx.ForKind.SERIAL) + assert not serial.is_parallel() + assert not serial.is_thread_binding() + + +@pytest.mark.parametrize("value", ["threadIdx.x", 1]) +def test_reject_invalid_annotation_value(value): + with pytest.raises(TypeError, match="must be an IterVar"): + make_loop(annotations={"thread_binding": value}) + + +@pytest.mark.parametrize("via_annotation", [False, True]) +def test_reject_empty_thread_tag(via_annotation): + binding = make_binding("") + with pytest.raises(ValueError, match="nonempty thread tag"): + if via_annotation: + make_loop(annotations={"thread_binding": binding}) + else: + make_loop(binding) + + +@pytest.mark.parametrize( + "kind", [tirx.ForKind.SERIAL, tirx.ForKind.UNROLLED, tirx.ForKind.VECTORIZED] +) +@pytest.mark.parametrize("via_annotation", [False, True]) +def test_reject_binding_on_nonparallel_loop(kind, via_annotation): + binding = make_binding() + with pytest.raises(ValueError, match="only valid on parallel loops"): + if via_annotation: + make_loop(annotations={"thread_binding": binding}, kind=kind) + else: + make_loop(binding, kind=kind) + + +def test_reject_removed_thread_binding_kind(): + with pytest.raises(ValueError, match="Invalid ForKind"): + make_loop(make_binding(), kind=4) + + +def test_reject_conflicting_binding_sources(): + with pytest.raises(ValueError, match="Conflicting thread_binding"): + make_loop(make_binding(), {"thread_binding": make_binding("blockIdx.x")}) + + +def test_serialization_and_structural_identity(): + binding = make_binding(dom=tvm.ir.Range.from_min_extent(2, 8)) + loop = make_loop(binding) + restored = tvm.ir.load_json(tvm.ir.save_json(loop)) + tvm.ir.assert_structural_equal(restored, loop, map_free_vars=True) + assert tvm_ffi.structural_hash(restored, map_free_vars=True) == tvm_ffi.structural_hash( + loop, map_free_vars=True + ) + assert int(restored.thread_binding.dom.min) == 2 + assert int(restored.thread_binding.dom.extent) == 8 + assert restored.thread_binding.iter_type == tirx.IterVar.ThreadIndex + assert restored.thread_binding.thread_tag == "threadIdx.x" + assert restored.thread_binding.var.name == binding.var.name + assert restored.annotations["thread_binding"].same_as(restored.thread_binding) + via_annotation = make_loop(annotations={"thread_binding": binding}) + tvm.ir.assert_structural_equal(loop, via_annotation, map_free_vars=True) + assert not tvm_ffi.structural_equal(loop, make_loop(), map_free_vars=True) + assert not tvm_ffi.structural_equal( + loop, make_loop(make_binding("blockIdx.x", binding.dom)), map_free_vars=True + ) + assert not tvm_ffi.structural_equal( + loop, make_loop(make_binding(dom=tvm.ir.Range.from_min_extent(3, 8))), map_free_vars=True + ) + + +def test_structural_walk_and_mutation_reach_binding_metadata(): + extent = tirx.Var("extent", "int32") + thread_var = tirx.Var("thread", "int32") + replacement = tirx.Var("new_thread", "int32") + binding = make_binding(dom=tvm.ir.Range.from_min_extent(2, extent), var=thread_var) + loop = make_loop(binding, {"custom_hint": extent}) + visited = [] + tvm_ffi.structural_walk(loop, visited.append) + assert any(node.same_as(binding) for node in visited) + assert any(node.same_as(thread_var) for node in visited) + assert any(node.same_as(extent) for node in visited) + + def rewrite(var): + if var.same_as(thread_var): + return replacement + if var.same_as(extent): + return tirx.IntImm("int32", 16) + return var + + rewritten = tvm_ffi.structural_map(loop, (tirx.Var, rewrite), order="post") + assert rewritten.thread_binding.var.same_as(replacement) + assert int(rewritten.thread_binding.dom.min) == 2 + assert int(rewritten.thread_binding.dom.extent) == 16 + assert rewritten.thread_binding.iter_type == binding.iter_type + assert rewritten.thread_binding.thread_tag == binding.thread_tag + # Auxiliary hints remain opaque; only the semantic annotation is traversed. + assert rewritten.annotations["custom_hint"].same_as(extent) + assert loop.thread_binding.var.same_as(thread_var) + assert loop.thread_binding.dom.extent.same_as(extent) + + +def test_script_roundtrip_keeps_hints_separate(): + @T.prim_func + def before(): + for i in T.thread_binding(8, thread="threadIdx.x", annotations={"custom_hint": 7}): + T.evaluate(i) + + script = before.script() + assert 'annotations={"thread_binding"' not in script + restored = tvm.script.from_source(script) + tvm.ir.assert_structural_equal(before, restored) + assert restored.body.is_thread_binding() + assert int(restored.body.annotations["custom_hint"]) == 7 + + +@pytest.mark.parametrize("tag", ["threadIdx.x", "blockIdx.x", "vthread.x"]) +@pytest.mark.parametrize("extent", [1, 8]) +def test_gpu_binding_lowers_to_thread_extent(tag, extent): + loop = make_loop(make_binding(tag), extent=extent) + mod = tvm.IRModule.from_expr(tirx.PrimFunc([], loop)) + lowered = tirx.transform.LowerTIRxOpaque()(mod)["main"].body + assert isinstance(lowered, tirx.AttrStmt) + assert lowered.attr_key == ("virtual_thread" if tag == "vthread.x" else "thread_extent") + assert lowered.node.thread_tag == tag + assert int(lowered.value) == extent + if extent == 1: + assert int(lowered.body.value) == 0 + else: + assert lowered.body.value.same_as(lowered.node.var) + + +@pytest.mark.parametrize( + "config", + [ + {"auto_max_extent": 1}, + {"auto_max_step": 16, "auto_max_depth": 8, "explicit_unroll": True}, + ], +) +def test_unit_thread_binding_survives_unroll(config): + loop = make_loop(make_binding(), extent=1) + mod = tvm.IRModule.from_expr(tirx.PrimFunc([], loop)) + with tvm.transform.PassContext(config={"tirx.UnrollLoop": config}): + transformed = tirx.transform.UnrollLoop()(mod)["main"].body + assert isinstance(transformed, tirx.For) + assert transformed.is_thread_binding() + assert int(transformed.extent) == 1 + tvm.ir.assert_structural_equal(transformed, loop) + + +def test_software_pipeline_cannot_discard_thread_scope(): + loop = make_loop( + make_binding(), + {"software_pipeline_stage": [0, 1], "software_pipeline_order": [0, 1]}, + ) + mod = tvm.IRModule.from_expr(tirx.PrimFunc([], loop)) + with pytest.raises(ValueError, match="cannot replace a thread-bound loop"): + tvm.s_tir.transform.InjectSoftwarePipeline()(mod) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_cpu_parallel_execution(): + @T.prim_func + def before(out: T.Buffer((16,), "int32")): + for i in T.parallel(16): + out[i] = i + 3 + + assert before.body.is_parallel() + compiled = tvm.compile(before.with_attr("global_symbol", "main"), target="llvm") + output = np.zeros(16, dtype="int32") + compiled(output) + np.testing.assert_array_equal(output, np.arange(16, dtype="int32") + 3) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py index 2932f929a146..22e57c9429c7 100644 --- a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py +++ b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py @@ -239,7 +239,7 @@ def test_ir_builder_tir_for(): loop_var=tirx.Var("", "int32"), min=0, extent=8, - kind=tirx.ForKind.THREAD_BINDING, + kind=tirx.ForKind.PARALLEL, body=tirx.Evaluate(0), thread_binding=tirx.IterVar( None, tirx.Var("", "int32"), tirx.IterVar.ThreadIndex, "threadIdx.x" diff --git a/tests/python/tvmscript/test_tvmscript_roundtrip.py b/tests/python/tvmscript/test_tvmscript_roundtrip.py index 3466246b69ba..f2088ca1288f 100644 --- a/tests/python/tvmscript/test_tvmscript_roundtrip.py +++ b/tests/python/tvmscript/test_tvmscript_roundtrip.py @@ -1834,10 +1834,10 @@ def test_for_thread_binding(): tvm.ir.assert_structural_equal(func, rt_func) assert isinstance(rt_func.body, tirx.stmt.For) - assert rt_func.body.kind == 4 + assert rt_func.body.is_thread_binding() assert rt_func.body.thread_binding.thread_tag == "threadIdx.x" assert isinstance(rt_func.body.body, tirx.stmt.For) - assert rt_func.body.body.kind == 4 + assert rt_func.body.body.is_thread_binding() assert rt_func.body.body.thread_binding.thread_tag == "threadIdx.y" assert rt_func.body.body.annotations["attr_key"] == "attr_value" From a3468bccafcbe9be3661b39c598e040c44e0dcec Mon Sep 17 00:00:00 2001 From: tqchen Date: Mon, 21 Sep 2026 21:37:21 +0000 Subject: [PATCH 2/5] [FIX][TIR] Scope semantic binding comparison to loop annotations --- src/s_tir/schedule/ir_comparator.cc | 26 ++++++++++++------- .../test_tir_schedule_binding_annotation.py | 12 +++++++++ 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc index 3263c4add139..2ee3c3195062 100644 --- a/src/s_tir/schedule/ir_comparator.cc +++ b/src/s_tir/schedule/ir_comparator.cc @@ -218,7 +218,22 @@ bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) { } return false; } - if (!CompareAnnotationMap(op->annotations, rhs->annotations)) { + // Only For annotations reserve thread_binding as an IterVar definition. + // SBlock annotations with the same spelling remain ordinary hints. + if (auto binding = op->GetThreadBinding()) { + IterVar lhs_iter = binding.value(); + IterVar rhs_iter = rhs->GetThreadBinding().value(); + if (!(CompareIterVar(lhs_iter, rhs_iter) && lhs_iter->thread_tag == rhs_iter->thread_tag && + lhs_iter->dom.defined() == rhs_iter->dom.defined() && + (!lhs_iter->dom.defined() || CompareRange(lhs_iter->dom, rhs_iter->dom)))) { + return false; + } + } + auto lhs_annotations = op->annotations; + auto rhs_annotations = rhs->annotations; + lhs_annotations.erase(tirx::attr::thread_binding); + rhs_annotations.erase(tirx::attr::thread_binding); + if (!CompareAnnotationMap(lhs_annotations, rhs_annotations)) { if (assert_mode_) { std::ostringstream os; os << "ForNode annotation maps do not match: op->annotations=" << op->annotations @@ -442,15 +457,6 @@ bool TensorizeComparator::CompareAnnotation(const std::pair(); - IterVar rhs_iter = rhs.second.as_or_throw(); - return CompareIterVar(lhs_iter, rhs_iter) && lhs_iter->thread_tag == rhs_iter->thread_tag && - lhs_iter->dom.defined() == rhs_iter->dom.defined() && - (!lhs_iter->dom.defined() || CompareRange(lhs_iter->dom, rhs_iter->dom)); - } // handle expr values if (lhs.second.as() && rhs.second.as()) { return Dispatch(lhs.second.as_or_throw(), rhs.second.as_or_throw()); diff --git a/tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py b/tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py index 67505d0fae14..3636e79ac58f 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py @@ -21,6 +21,7 @@ import tvm import tvm.testing from tvm import tirx +from tvm.s_tir.schedule.analysis import get_auto_tensorize_mapping_info from tvm.script import tirx as T @@ -115,5 +116,16 @@ def test_sblock_thread_binding_hint_is_not_reserved(): assert "thread_binding" not in sch.get(block).annotations +def test_sblock_thread_binding_hint_in_tensorize_comparison(): + sch, _ = make_schedule("serial") + block = sch.get_sblock("copy") + sch.annotate(block, "thread_binding", "block_hint") + desc = sch.mod["main"] + assert get_auto_tensorize_mapping_info(sch, block, desc) is not None + sch.unannotate(block, "thread_binding") + sch.annotate(block, "thread_binding", "different_hint") + assert get_auto_tensorize_mapping_info(sch, block, desc) is None + + if __name__ == "__main__": tvm.testing.main() From 90e2dd2535bd3d2ab3a97938f4c8d50aea6ab923 Mon Sep 17 00:00:00 2001 From: tqchen Date: Mon, 21 Sep 2026 21:37:21 +0000 Subject: [PATCH 3/5] [TEST][TIR] Cover binding mutation with transferred loop ownership --- .../python/tirx-base/test_tir_for_thread_binding.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/python/tirx-base/test_tir_for_thread_binding.py b/tests/python/tirx-base/test_tir_for_thread_binding.py index 2e5e4e3065ca..83186ebc87b7 100644 --- a/tests/python/tirx-base/test_tir_for_thread_binding.py +++ b/tests/python/tirx-base/test_tir_for_thread_binding.py @@ -122,7 +122,8 @@ def test_serialization_and_structural_identity(): ) -def test_structural_walk_and_mutation_reach_binding_metadata(): +@pytest.mark.parametrize("move", [False, True]) +def test_structural_walk_and_mutation_reach_binding_metadata(move): extent = tirx.Var("extent", "int32") thread_var = tirx.Var("thread", "int32") replacement = tirx.Var("new_thread", "int32") @@ -141,7 +142,10 @@ def rewrite(var): return tirx.IntImm("int32", 16) return var - rewritten = tvm_ffi.structural_map(loop, (tirx.Var, rewrite), order="post") + visited.clear() # Do not retain the root while exercising ownership transfer. + rewritten = tvm_ffi.structural_map( + loop._move() if move else loop, (tirx.Var, rewrite), order="post" + ) assert rewritten.thread_binding.var.same_as(replacement) assert int(rewritten.thread_binding.dom.min) == 2 assert int(rewritten.thread_binding.dom.extent) == 16 @@ -149,8 +153,9 @@ def rewrite(var): assert rewritten.thread_binding.thread_tag == binding.thread_tag # Auxiliary hints remain opaque; only the semantic annotation is traversed. assert rewritten.annotations["custom_hint"].same_as(extent) - assert loop.thread_binding.var.same_as(thread_var) - assert loop.thread_binding.dom.extent.same_as(extent) + if not move: + assert loop.thread_binding.var.same_as(thread_var) + assert loop.thread_binding.dom.extent.same_as(extent) def test_script_roundtrip_keeps_hints_separate(): From 2928fe99b2c2a85f74f11a7b17379d710ac4b56d Mon Sep 17 00:00:00 2001 From: tqchen Date: Mon, 21 Sep 2026 23:04:42 +0000 Subject: [PATCH 4/5] [REFACTOR][TIR] Store thread tags in S-TIR annotations --- include/tvm/s_tir/stmt.h | 9 + include/tvm/tirx/stmt.h | 17 +- python/tvm/s_tir/dlight/gpu/fallback.py | 7 +- python/tvm/tirx/stmt.py | 23 +- .../transform/split_call_tir_by_pattern.cc | 3 +- .../sblock_buffer_access_lca_detector.cc | 4 +- .../feature_extractor/per_store_feature.cc | 6 +- .../rewrite_parallel_vectorize_unroll.cc | 2 +- src/s_tir/schedule/analysis/analysis.cc | 4 +- src/s_tir/schedule/ir_comparator.cc | 27 +- src/s_tir/schedule/primitive/annotate.cc | 8 +- .../schedule/primitive/blockize_tensorize.cc | 2 +- .../schedule/primitive/compute_inline.cc | 2 +- .../schedule/primitive/decompose_padding.cc | 4 +- src/s_tir/schedule/primitive/for_kind.cc | 9 +- .../schedule/primitive/loop_transformation.cc | 8 +- src/s_tir/schedule/primitive/reduction.cc | 8 - src/s_tir/schedule/utils.h | 4 +- src/s_tir/stmt.cc | 19 ++ src/s_tir/transform/compact_buffer_region.cc | 6 +- src/s_tir/transform/default_gpu_schedule.cc | 2 +- .../transform/inject_software_pipeline.cc | 4 +- src/s_tir/transform/lift_thread_binding.cc | 25 +- src/s_tir/transform/loop_partition.cc | 2 +- .../transform/lower_cross_thread_reduction.cc | 37 ++- src/s_tir/transform/lower_opaque_block.cc | 10 +- src/s_tir/transform/memhammer_coalesce.cc | 7 +- .../transform/memhammer_intermediate_stage.cc | 8 +- .../transform/memhammer_lower_auto_copy.cc | 13 +- .../transform/profile_instrumentation.cc | 6 +- src/s_tir/transform/unify_thread_binding.cc | 29 +-- src/target/llvm/codegen_cpu.cc | 3 +- src/tirx/ir/data_type_rewriter.cc | 6 - src/tirx/ir/stmt.cc | 69 +----- src/tirx/script/builder/ir.cc | 19 +- src/tirx/script/printer/for_loop.cc | 17 +- src/tirx/transform/bind_target.cc | 4 +- src/tirx/transform/lower_tirx_opaque.cc | 10 +- src/tirx/transform/lower_tvm_builtin.cc | 3 +- src/tirx/transform/make_packed_api.cc | 5 +- src/tirx/transform/storage_rewrite.cc | 4 +- src/tirx/transform/unroll_loop.cc | 3 +- src/tirx/transform/vectorize_loop.cc | 6 +- tests/python/relax/test_pipeline.py | 2 +- .../test_tir_schedule_binding_annotation.py | 131 ---------- .../s_tir/schedule/test_tir_schedule_state.py | 1 - .../tirx-base/test_tir_for_thread_binding.py | 234 ------------------ .../test_tvmscript_ir_builder_tir.py | 4 +- .../tvmscript/test_tvmscript_parser_tir.py | 2 - .../tvmscript/test_tvmscript_roundtrip.py | 8 +- 50 files changed, 191 insertions(+), 655 deletions(-) delete mode 100644 tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py delete mode 100644 tests/python/tirx-base/test_tir_for_thread_binding.py diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h index eefe47d8602f..59449aac9209 100644 --- a/include/tvm/s_tir/stmt.h +++ b/include/tvm/s_tir/stmt.h @@ -32,6 +32,13 @@ namespace tvm { namespace s_tir { +/*! \brief Return the validated S-TIR thread-binding annotation, if present. */ +TVM_DLL ffi::Optional GetThreadBinding(const tirx::ForNode* loop); +/*! \brief Whether the parallel loop is bound to a device thread axis. */ +TVM_DLL bool IsThreadBinding(const tirx::ForNode* loop); +/*! \brief Whether the loop is unbound host parallel work. */ +TVM_DLL bool IsParallel(const tirx::ForNode* loop); + /*! * \brief Match introduces a constraint that the source buffer region can be remapped to the data * layout specified by the buffer field. The constraint can be checked in later part of lowering (or @@ -199,6 +206,8 @@ class SBlockRealize : public tirx::Stmt { }; namespace attr { +/*! \brief Device thread axis name attached to an S-TIR parallel loop. */ +constexpr const char* thread_binding = "thread_binding"; /*! * \brief Annotations for invoking and synchronizing asynchronous operations. diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h index c311bf7e1d1b..0e58adfc7a2b 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -595,9 +595,8 @@ class ForNode : public StmtNode { /*! * \brief Additional annotations about the loop. * - * Most annotations are auxiliary transformation hints. The reserved - * thread_binding annotation is semantic: its IterVar binds a parallel loop - * to an execution thread and must be preserved until binding is lowered. + * Annotations may carry execution semantics as well as transformation hints. + * Transformations must preserve annotations until their semantics are consumed. */ ffi::Map annotations; /*! @@ -620,15 +619,6 @@ class ForNode : public StmtNode { /*! \brief Check it is a loop without nontrivial loop step. */ bool HasTrivialStep() const; - /*! \brief Get the validated semantic thread binding, if present. */ - TVM_DLL ffi::Optional GetThreadBinding() const; - /*! \brief Set or remove the semantic binding on a parallel loop. */ - TVM_DLL void SetThreadBinding(ffi::Optional binding); - /*! \brief Whether this parallel loop is bound to an execution thread. */ - bool IsThreadBinding() const { return GetThreadBinding().has_value(); } - /*! \brief Whether this is an ordinary, unbound CPU parallel loop. */ - bool IsParallel() const { return kind == ForKind::kParallel && !IsThreadBinding(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.For", ForNode, StmtNode); }; @@ -639,7 +629,6 @@ class ForNode : public StmtNode { class For : public Stmt { public: TVM_DLL For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, Stmt body, - ffi::Optional thread_binding = std::nullopt, ffi::Map annotations = {}, ffi::Optional step = std::nullopt, Span span = Span()); @@ -794,8 +783,6 @@ class ScopeIdDefStmt : public Stmt { /*! \brief namespace of possible attributes in AttrStmt.attr_key */ namespace attr { -/*! \brief Semantic IterVar binding of a parallel For loop to an execution thread. */ -constexpr const char* thread_binding = "thread_binding"; /*! * \brief Mark the scope as when computation start to happen. * This can hint some code generator to create a new function for compute. diff --git a/python/tvm/s_tir/dlight/gpu/fallback.py b/python/tvm/s_tir/dlight/gpu/fallback.py index a194b07eb8f0..5e5d907280eb 100644 --- a/python/tvm/s_tir/dlight/gpu/fallback.py +++ b/python/tvm/s_tir/dlight/gpu/fallback.py @@ -42,7 +42,7 @@ def visit_attr(node: tirx.AttrStmt): def visit_for(node: tirx.For): nonlocal found - if node.is_thread_binding(): + if "thread_binding" in node.annotations: found = True tvm_ffi.structural_walk( @@ -85,7 +85,10 @@ def apply( # pylint: disable=too-many-locals,missing-docstring block = block.block_rv if any( - [sch.get(loop_rv).thread_binding is not None for loop_rv in sch.get_loops(block)] + [ + "thread_binding" in sch.get(loop_rv).annotations + for loop_rv in sch.get_loops(block) + ] ): continue diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 3502c03e744a..843a48972ed2 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -39,7 +39,7 @@ from . import _ffi_api from .buffer import Buffer from .exec_scope import ScopeIdDef -from .expr import IterVar, Var +from .expr import Var @tvm_ffi.register_object("tirx.Stmt") @@ -194,16 +194,12 @@ class For(Stmt): body : Stmt The body statement. - thread_binding: Optional[tirx.IterVar] - The thread this loop binds to. Only valid - if kind is PARALLEL. Stored in the semantic thread_binding annotation. - step : Expr The loop step. Default to none which represent one. annotations: Optional[Mapping[str, Object]] - Additional annotations, including the semantic thread_binding IterVar. + Additional execution annotations and transformation hints. span : Optional[Span] The location of the stmt in the source code. @@ -225,7 +221,6 @@ def __init__( extent: Expr, kind: ForKind, body: Stmt, - thread_binding: IterVar | None = None, annotations: Mapping[str, Object] | None = None, step: Expr | None = None, span: Span | None = None, @@ -238,25 +233,11 @@ def __init__( extent, kind, body, - thread_binding, annotations, step, span, ) - @property - def thread_binding(self) -> IterVar | None: - """The semantic thread binding of this parallel loop, if present.""" - return self.annotations.get("thread_binding") - - def is_thread_binding(self) -> bool: - """Whether this loop is bound to an execution thread.""" - return self.thread_binding is not None - - def is_parallel(self) -> bool: - """Whether this is an ordinary unbound CPU parallel loop.""" - return self.kind == ForKind.PARALLEL and not self.is_thread_binding() - @tvm_ffi.register_object("tirx.While") class While(Stmt): diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index cf4548ab58e5..5b514d4dfd66 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -266,7 +266,8 @@ class ForMatcher : public TensorizeComparator { if (!DefEqual(op->loop_var, rhs->loop_var)) return false; // Only handle the case where the loop start from 0 if (!is_zero(op->min) || !is_zero(rhs->min)) return false; - if (op->GetThreadBinding().has_value() || rhs->GetThreadBinding().has_value()) return false; + if (s_tir::GetThreadBinding(op).has_value() || s_tir::GetThreadBinding(rhs).has_value()) + return false; if (op->kind != ForKind::kSerial || op->kind != rhs->kind) return false; if (!op->annotations.empty() || !rhs->annotations.empty()) return false; // Match the extents of loops diff --git a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc index 1581b8e22b49..96015722da83 100644 --- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc +++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc @@ -99,9 +99,9 @@ class LCADetector : public s_tir::StmtExprVisitor { const ScopeInfo* parent_scope = ancestor_scopes_.back(); auto* current_scope = arena_.make(parent_scope, op, n); - if (op->GetThreadBinding().has_value()) { + if (s_tir::GetThreadBinding(op).has_value()) { const runtime::ThreadScope& scope = - runtime::ThreadScope::Create(op->GetThreadBinding().value()->thread_tag); + runtime::ThreadScope::Create(s_tir::GetThreadBinding(op).value()); if (scope.rank == 0) { blockidx_scopes_.push_back(current_scope); } diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc index f6fa67f4d163..0760aac85c59 100644 --- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc +++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc @@ -370,14 +370,14 @@ struct LoopNest { this->auto_unroll.push_back(*auto_unroll_attr); } ForVec* ref_loops = nullptr; - if (loop->IsParallel()) { + if (IsParallel(loop)) { ref_loops = ∥ } else if (loop->kind == ForKind::kVectorized) { ref_loops = &vectorize; } else if (loop->kind == ForKind::kUnrolled) { ref_loops = &unroll; - } else if (loop->IsThreadBinding()) { - std::string thread_tag = loop->GetThreadBinding().value()->thread_tag; + } else if (IsThreadBinding(loop)) { + std::string thread_tag = GetThreadBinding(loop).value(); if (thread_tag == "blockIdx.x") { ref_loops = &blockIdx_x; } else if (thread_tag == "blockIdx.y") { diff --git a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc index 55014ecf9de7..42d1626b6ccf 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc @@ -33,7 +33,7 @@ using namespace tvm::tirx; * \return Whether the loop has any annotation */ inline bool HasAnnOrBinding(const ForNode* loop) { - return loop->IsThreadBinding() || !loop->annotations.empty(); + return IsThreadBinding(loop) || !loop->annotations.empty(); } /*! \brief The visitor for extracting the stride of a var in a PrimExpr. */ diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index 82f21a8c98ef..c29e9a463715 100644 --- a/src/s_tir/schedule/analysis/analysis.cc +++ b/src/s_tir/schedule/analysis/analysis.cc @@ -703,8 +703,8 @@ ffi::Map LoopDomainOfSRefTreePath(const StmtSRef& low_inclusive, if (extra_relax_scope.rank != runtime::StorageRank::kGlobal) { for (; p; p = p->parent) { if (const ForNode* loop = p->StmtAs()) { - if (loop->IsThreadBinding()) { - const ffi::String& thread_tag = loop->GetThreadBinding().value()->thread_tag; + if (IsThreadBinding(loop)) { + ffi::String thread_tag = GetThreadBinding(loop).value(); if (CanRelaxStorageUnderThread(extra_relax_scope, runtime::ThreadScope::Create(thread_tag))) { result.Set(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent)); diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc index 2ee3c3195062..bd03fb36758e 100644 --- a/src/s_tir/schedule/ir_comparator.cc +++ b/src/s_tir/schedule/ir_comparator.cc @@ -200,16 +200,6 @@ bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) { } return false; } - if (op->GetThreadBinding().has_value() != rhs->GetThreadBinding().has_value()) { - if (assert_mode_) { - std::ostringstream os; - os << "ForNode thread_bindings do not match: op->GetThreadBinding().has_value()=" - << op->GetThreadBinding().has_value() - << " vs rhs->GetThreadBinding().has_value()=" << rhs->GetThreadBinding().has_value(); - EmitError(os.str()); - } - return false; - } if (op->kind != rhs->kind) { if (assert_mode_) { std::ostringstream os; @@ -218,22 +208,7 @@ bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) { } return false; } - // Only For annotations reserve thread_binding as an IterVar definition. - // SBlock annotations with the same spelling remain ordinary hints. - if (auto binding = op->GetThreadBinding()) { - IterVar lhs_iter = binding.value(); - IterVar rhs_iter = rhs->GetThreadBinding().value(); - if (!(CompareIterVar(lhs_iter, rhs_iter) && lhs_iter->thread_tag == rhs_iter->thread_tag && - lhs_iter->dom.defined() == rhs_iter->dom.defined() && - (!lhs_iter->dom.defined() || CompareRange(lhs_iter->dom, rhs_iter->dom)))) { - return false; - } - } - auto lhs_annotations = op->annotations; - auto rhs_annotations = rhs->annotations; - lhs_annotations.erase(tirx::attr::thread_binding); - rhs_annotations.erase(tirx::attr::thread_binding); - if (!CompareAnnotationMap(lhs_annotations, rhs_annotations)) { + if (!CompareAnnotationMap(op->annotations, rhs->annotations)) { if (assert_mode_) { std::ostringstream os; os << "ForNode annotation maps do not match: op->annotations=" << op->annotations diff --git a/src/s_tir/schedule/primitive/annotate.cc b/src/s_tir/schedule/primitive/annotate.cc index 29efcd0c6bf1..dc71c62a174e 100644 --- a/src/s_tir/schedule/primitive/annotate.cc +++ b/src/s_tir/schedule/primitive/annotate.cc @@ -27,7 +27,7 @@ using namespace tvm::tirx; void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_key, const Any& ann_val) { - TVM_FFI_CHECK(!sref->StmtAs() || ann_key != tirx::attr::thread_binding, ValueError) + TVM_FFI_CHECK(!sref->StmtAs() || ann_key != s_tir::attr::thread_binding, ValueError) << "thread_binding is a semantic annotation; use Schedule.bind to set it"; // Extract annotation const ffi::Map* annotations = nullptr; @@ -50,7 +50,7 @@ void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_k ffi::ObjectPtr n = ffi::make_object(*loop); n->annotations = std::move(new_ann); // Validate semantic loop annotations before installing the replacement. - n->GetThreadBinding(); + GetThreadBinding(n.get()); self->Replace(sref, For(n), {}); } else if (const auto* block = sref->StmtAs()) { ffi::ObjectPtr n = ffi::make_object(*block); @@ -64,7 +64,7 @@ void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_k } void Unannotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_key) { - TVM_FFI_CHECK(!sref->StmtAs() || ann_key != tirx::attr::thread_binding, ValueError) + TVM_FFI_CHECK(!sref->StmtAs() || ann_key != s_tir::attr::thread_binding, ValueError) << "thread_binding is a semantic annotation; use Schedule.parallel to remove it"; // Extract annotation const ffi::Map* annotations = nullptr; @@ -85,7 +85,7 @@ void Unannotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann ffi::ObjectPtr n = ffi::make_object(*loop); n->annotations = std::move(new_ann); // Validate semantic loop annotations before installing the replacement. - n->GetThreadBinding(); + GetThreadBinding(n.get()); self->Replace(sref, For(n), {}); } else if (const auto* block = sref->StmtAs()) { ffi::ObjectPtr n = ffi::make_object(*block); diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index 561fca780ec5..f67fe8a7d30f 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -759,7 +759,7 @@ class BlockizeRewriter : public StmtExprMutator { UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { if (loop == lca_->stmt) { return For(loop->loop_var, loop->min, loop->extent, loop->kind, RewriteSeq(loop->body), - loop->GetThreadBinding(), loop->annotations, loop->step, loop->span); + loop->annotations, loop->step, loop->span); } return StmtExprMutator::Mutate_(loop, inplace_mode); } diff --git a/src/s_tir/schedule/primitive/compute_inline.cc b/src/s_tir/schedule/primitive/compute_inline.cc index c8e900a917d9..3df4770103b5 100644 --- a/src/s_tir/schedule/primitive/compute_inline.cc +++ b/src/s_tir/schedule/primitive/compute_inline.cc @@ -1796,7 +1796,7 @@ class SingleBlockFusionReplacer : public StmtExprMutator { } return For(loop->loop_var, loop->min, loop->extent, loop->kind, mutated_body, - loop->GetThreadBinding(), loop->annotations); + loop->annotations); } UnchangedOr Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final { diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc b/src/s_tir/schedule/primitive/decompose_padding.cc index 4d4d6efc70bc..cdd8c136267b 100644 --- a/src/s_tir/schedule/primitive/decompose_padding.cc +++ b/src/s_tir/schedule/primitive/decompose_padding.cc @@ -373,8 +373,8 @@ static std::pair CreateInBoundBlock(const SBlockRealizeNode auto it = new_loop_ranges.find(loop->loop_var); PrimExpr min = it == new_loop_ranges.end() ? loop->min : (*it).second->min; PrimExpr extent = it == new_loop_ranges.end() ? loop->extent : (*it).second->extent; - nest_stmt_root = For(loop->loop_var, min, extent, loop->kind, nest_stmt_root, - loop->GetThreadBinding(), loop->annotations, loop->step, loop->span); + nest_stmt_root = For(loop->loop_var, min, extent, loop->kind, nest_stmt_root, loop->annotations, + loop->step, loop->span); if (loop.same_as(highest_pos_inclusive)) { break; } diff --git a/src/s_tir/schedule/primitive/for_kind.cc b/src/s_tir/schedule/primitive/for_kind.cc index cd1346b80b8b..af71d4ea8af9 100644 --- a/src/s_tir/schedule/primitive/for_kind.cc +++ b/src/s_tir/schedule/primitive/for_kind.cc @@ -179,12 +179,9 @@ void ParallelizeComputation(const ScheduleState& self, const StmtSRef& loop_sref ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->kind = for_kind; if (thread_axis.has_value()) { - new_loop->SetThreadBinding(IterVar(/*dom=*/Range(nullptr), - /*var=*/PrimVar(thread_axis.value(), loop->loop_var.ty()), - /*iter_type=*/kThreadIndex, - /*thread_tag=*/thread_axis.value())); + new_loop->annotations.Set(s_tir::attr::thread_binding, thread_axis.value()); } else { - new_loop->SetThreadBinding(std::nullopt); + new_loop->annotations.erase(s_tir::attr::thread_binding); } self->Replace(loop_sref, For(new_loop), {}); } @@ -205,7 +202,7 @@ void Unroll(ScheduleState self, const StmtSRef& loop_sref) { const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->kind = ForKind::kUnrolled; - new_loop->SetThreadBinding(std::nullopt); + new_loop->annotations.erase(s_tir::attr::thread_binding); self->Replace(loop_sref, For(new_loop), {}); } diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc index f386fb6af3bf..51f686606c84 100644 --- a/src/s_tir/schedule/primitive/loop_transformation.cc +++ b/src/s_tir/schedule/primitive/loop_transformation.cc @@ -437,7 +437,7 @@ ffi::Array Split(ScheduleState self, const StmtSRef& loop_sref, // order with before. // Step 1. Check correctness const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); - if (!loop->annotations.empty() || loop->GetThreadBinding().has_value()) { + if (!loop->annotations.empty() || GetThreadBinding(loop).has_value()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } // Currently, loops not starting with 0 are not supported @@ -712,7 +712,7 @@ const ffi::String get_sblock_name(Stmt loop_body) { ffi::Array LoopPartition(ScheduleState self, const StmtSRef& loop_sref, const ffi::Array& factors, bool preserve_unit_iters) { const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); - if (!loop->annotations.empty() || loop->GetThreadBinding().has_value()) { + if (!loop->annotations.empty() || GetThreadBinding(loop).has_value()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } @@ -895,7 +895,7 @@ StmtSRef Merge(ScheduleState self, const ffi::Array& loop_srefs) { std::vector nest_loop_i_loops; for (auto p = sref.get(); p != lca.get(); p = p->parent) { if (auto loop = p->StmtAs()) { - if (!loop->annotations.empty() || loop->GetThreadBinding().has_value()) { + if (!loop->annotations.empty() || GetThreadBinding(loop).has_value()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } @@ -962,7 +962,7 @@ StmtSRef Fuse(ScheduleState self, const ffi::Array& loop_srefs, // Step 1. check correctness for (const StmtSRef& sref : loop_srefs) { const ForNode* loop = TVM_SREF_TO_FOR(sref); - if (!loop->annotations.empty() || loop->GetThreadBinding().has_value()) { + if (!loop->annotations.empty() || GetThreadBinding(loop).has_value()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } if (outer_loop_sref.defined()) { diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc index 064c5e947460..2c9b576f78be 100644 --- a/src/s_tir/schedule/primitive/reduction.cc +++ b/src/s_tir/schedule/primitive/reduction.cc @@ -310,16 +310,8 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref, Var old_loop_var = old_loop->loop_var; PrimVar new_loop_var = old_loop->loop_var.CopyWithSuffix("_init"); loop_var_map[old_loop_var] = new_loop_var; - ffi::Optional opt_thread_binding = old_loop->GetThreadBinding(); - if (opt_thread_binding) { - auto thread_binding = opt_thread_binding.value(); - auto new_var = thread_binding->var.CopyWithSuffix(""); - thread_binding.CopyOnWrite()->var = new_var; - opt_thread_binding = thread_binding; - } auto new_loop = old_loop.CopyOnWrite(); new_loop->loop_var = new_loop_var; - new_loop->SetThreadBinding(opt_thread_binding); new_loop->body = body; body = ffi::GetRef(new_loop); } diff --git a/src/s_tir/schedule/utils.h b/src/s_tir/schedule/utils.h index f68512eabea8..8789f4cdb5a2 100644 --- a/src/s_tir/schedule/utils.h +++ b/src/s_tir/schedule/utils.h @@ -179,8 +179,8 @@ inline IterVar IterVarFromLoop(const For& loop, ffi::String name, IterVarType it * \return The thread scope bound to the loop */ inline runtime::ThreadScope GetThreadScope(const ForNode* loop) { - if (loop->IsThreadBinding()) { - return runtime::ThreadScope::Create(loop->GetThreadBinding().value()->thread_tag); + if (IsThreadBinding(loop)) { + return runtime::ThreadScope::Create(GetThreadBinding(loop).value()); } return runtime::ThreadScope{-1, -1}; } diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc index 45cbd47b0be7..24ae895dd727 100644 --- a/src/s_tir/stmt.cc +++ b/src/s_tir/stmt.cc @@ -30,6 +30,25 @@ namespace tvm { namespace s_tir { + +ffi::Optional GetThreadBinding(const tirx::ForNode* loop) { + auto value = loop->annotations.Get(attr::thread_binding); + if (!value.has_value()) return std::nullopt; + TVM_FFI_CHECK(loop->kind == tirx::ForKind::kParallel, ValueError) + << "thread_binding is only valid on parallel loops"; + auto binding = value.value().as(); + TVM_FFI_CHECK(binding.has_value(), TypeError) << "thread_binding annotation must be a String"; + TVM_FFI_CHECK(!binding.value().empty(), ValueError) + << "thread_binding must have a nonempty thread tag"; + return binding; +} + +bool IsThreadBinding(const tirx::ForNode* loop) { return GetThreadBinding(loop).has_value(); } + +bool IsParallel(const tirx::ForNode* loop) { + return loop->kind == tirx::ForKind::kParallel && !IsThreadBinding(loop); +} + using namespace tvm::tirx; using namespace tvm::prim; diff --git a/src/s_tir/transform/compact_buffer_region.cc b/src/s_tir/transform/compact_buffer_region.cc index a74b8ff72f5e..8f127bdf4c81 100644 --- a/src/s_tir/transform/compact_buffer_region.cc +++ b/src/s_tir/transform/compact_buffer_region.cc @@ -186,9 +186,9 @@ class BufferAccessRegionCollector : public StmtExprVisitor { ffi::Optional Visit_(const ForNode* op) final { Range loop_range = Range::FromMinExtent(op->min, op->extent); - IterVar iter = op->IsThreadBinding() ? IterVar(Range(), op->loop_var, IterVarType::kThreadIndex, - op->GetThreadBinding().value()->thread_tag) - : IterVar(Range(), op->loop_var, IterVarType::kDataPar); + IterVar iter = IsThreadBinding(op) ? IterVar(Range(), op->loop_var, IterVarType::kThreadIndex, + GetThreadBinding(op).value()) + : IterVar(Range(), op->loop_var, IterVarType::kDataPar); ancestor_iters_.push_back(iter); dom_analyzer_->Bind(op->loop_var, loop_range); dom_map_.emplace(op->loop_var.get(), sym::IntSet::FromRange(loop_range)); diff --git a/src/s_tir/transform/default_gpu_schedule.cc b/src/s_tir/transform/default_gpu_schedule.cc index 1642a4b931e1..fad6cb6106ef 100644 --- a/src/s_tir/transform/default_gpu_schedule.cc +++ b/src/s_tir/transform/default_gpu_schedule.cc @@ -40,7 +40,7 @@ void ThreadBind(s_tir::Schedule sch, const s_tir::SBlockRV& block, int64_t max_t ffi::Array loops = sch->GetLoops(block); for (const s_tir::LoopRV& loop : loops) { // skip block if already scheduled - if (sch->Get(loop)->GetThreadBinding().has_value()) { + if (GetThreadBinding(sch->Get(loop).get()).has_value()) { return; } } diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc index a08a1cee5060..14f5d8595fa4 100644 --- a/src/s_tir/transform/inject_software_pipeline.cc +++ b/src/s_tir/transform/inject_software_pipeline.cc @@ -1027,7 +1027,7 @@ class PipelineRewriter : public StmtExprMutator { if (!is_unit_loop) { new_loop = For(new_loop_var.as_or_throw(), pipeline_loop_->min, extent, unroll_loop ? ForKind::kUnrolled : pipeline_loop_->kind, std::move(new_loop), - std::nullopt, preserved_annotations_, std::nullopt); + preserved_annotations_, std::nullopt); } // Update producer heads in the global async states. @@ -1174,7 +1174,7 @@ class PipelineInjector : public StmtExprMutator { if (!HasPipelineAnnotation(op)) { return for_node; } - TVM_FFI_CHECK(!for_node->IsThreadBinding(), ValueError) + TVM_FFI_CHECK(!IsThreadBinding(for_node.get()), ValueError) << "Software pipelining cannot replace a thread-bound loop"; // Step 2: Find the body and buffer allocations of the pipeline. The body can be direct child of // the for-loop. If the for-loop has BlockRealize as its child, the pipeline body will be the diff --git a/src/s_tir/transform/lift_thread_binding.cc b/src/s_tir/transform/lift_thread_binding.cc index 3ebf273831e1..aef7c85926cc 100644 --- a/src/s_tir/transform/lift_thread_binding.cc +++ b/src/s_tir/transform/lift_thread_binding.cc @@ -48,7 +48,7 @@ FindLoopLCA(const Stmt& root) { ffi::Optional Visit_(const ForNode* op) final { stack.push_back(ffi::GetRef(op)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); - if (op->IsThreadBinding()) { + if (IsThreadBinding(op)) { UpdateLCA(op); } stack.pop_back(); @@ -56,11 +56,11 @@ FindLoopLCA(const Stmt& root) { } void UpdateLCA(const ForNode* loop) { - std::string thread_tag = loop->GetThreadBinding().value()->thread_tag; + std::string thread_tag = GetThreadBinding(loop).value(); { ffi::Map* tgt = &annotations[thread_tag]; for (const auto& kv : loop->annotations) { - if (kv.first != tirx::attr::thread_binding) { + if (kv.first != s_tir::attr::thread_binding) { tgt->Set(kv.first, kv.second); } } @@ -69,10 +69,10 @@ FindLoopLCA(const Stmt& root) { if (!iter_var.defined()) { iter_var = IterVar(Range::FromMinExtent(loop->min, loop->extent), // loop->loop_var - .as_or_throw() // - .CopyWithName(thread_tag) // - .as_or_throw(), // - loop->GetThreadBinding().value()->iter_type, // + .as_or_throw() // + .CopyWithName(thread_tag) // + .as_or_throw(), // + kThreadIndex, // thread_tag); lca[thread_tag] = stack; var_subst.Set(loop->loop_var, iter_var->var); @@ -135,7 +135,7 @@ class ThreadBindingLifter : public StmtExprMutator { UnchangedOr Mutate_(const ForNode* _op, InplaceMode inplace_mode) final { For op = ffi::GetRef(_op); bool is_kernel_root = false; - if (op->IsThreadBinding()) { + if (IsThreadBinding(op.get())) { if (iter_lca.empty()) { is_kernel_root = true; SetKernelRoot(_op); @@ -147,17 +147,16 @@ class ThreadBindingLifter : public StmtExprMutator { Stmt body = std::move(new_op.CopyOnWrite()->body); if (auto it = iter_lca.find(op); it != iter_lca.end()) { for (const auto& [iter_var, annotation] : it->second) { + auto annotations = annotation; + annotations.Set(s_tir::attr::thread_binding, iter_var->thread_tag); body = For(iter_var->var, iter_var->dom->min, iter_var->dom->extent, ForKind::kParallel, - std::move(body), - IterVar(Range(nullptr), PrimVar(iter_var->thread_tag, iter_var->var.ty()), - kThreadIndex, iter_var->thread_tag), - annotation, std::nullopt); + std::move(body), annotations, std::nullopt); } } if (is_kernel_root) { iter_lca.clear(); } - if (op->IsThreadBinding()) { + if (IsThreadBinding(op.get())) { return body; } else { new_op.CopyOnWrite()->body = std::move(body); diff --git a/src/s_tir/transform/loop_partition.cc b/src/s_tir/transform/loop_partition.cc index b56a26bb1fd8..efc6a9c9264b 100644 --- a/src/s_tir/transform/loop_partition.cc +++ b/src/s_tir/transform/loop_partition.cc @@ -864,7 +864,7 @@ inline Stmt LoopPartitioner::MakeFor(const ffi::Object* node, PrimExpr extent, S }; return ffi::StructuralMap(body, f_substitute).as_or_throw(); } else { - TVM_FFI_ICHECK(!for_node->IsThreadBinding()); + TVM_FFI_ICHECK(!IsThreadBinding(for_node)); auto new_loop = ffi::make_object(*for_node); new_loop->min = IntImm(for_node->min.ty(), 0); new_loop->extent = extent; diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc index 2d1629a3df3d..06a9106cc2a0 100644 --- a/src/s_tir/transform/lower_cross_thread_reduction.cc +++ b/src/s_tir/transform/lower_cross_thread_reduction.cc @@ -65,11 +65,10 @@ struct ThreadScopeEqual { * \return True if the loop is bound to threadIdx.x/y/z */ bool IsBoundToThreadIdx(const ForNode* loop) { - if (!loop->GetThreadBinding().has_value()) { + if (!GetThreadBinding(loop).has_value()) { return false; } - runtime::ThreadScope scope = - runtime::ThreadScope::Create(loop->GetThreadBinding().value()->thread_tag); + runtime::ThreadScope scope = runtime::ThreadScope::Create(GetThreadBinding(loop).value()); return scope.rank == 1 && scope.dim_index >= 0; } @@ -280,7 +279,7 @@ class InThreadReducerMaker : public StmtExprMutator { .ValueOrUnchanged(ffi::GetRef(loop)) .as()) { For res = *opt_res; - if (res->GetThreadBinding().has_value()) { + if (GetThreadBinding(res.get()).has_value()) { if (!res->body.defined() || UnderLoopReductionBlockVarCollector::CheckHasReductionBlocks(res)) { return res->body; @@ -418,7 +417,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, } // Next arguments: all the reduction threads for (const ForNode* reduction_loop : reduction_loops) { - if (reduction_loop->GetThreadBinding().has_value()) { + if (GetThreadBinding(reduction_loop).has_value()) { parameters.push_back(reduction_loop->loop_var); } } @@ -546,7 +545,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, ffi::StructuralWalk(realize->predicate, walk_fn); if (wb_buffers[0].scope() != "local") { for (const ForNode* loop : reduction_loops) { - if (loop->GetThreadBinding().has_value()) { + if (GetThreadBinding(loop).has_value()) { wb_predicate = wb_predicate && (static_cast(loop->loop_var) == IntImm(loop->loop_var.ty(), 0)); } @@ -567,7 +566,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, Stmt new_stmt = SeqStmt::Flatten(std::move(stmts)); for (auto rit = reduction_loops.rbegin(); rit != reduction_loops.rend(); ++rit) { const ForNode* loop = *rit; - if (loop->GetThreadBinding().has_value()) { + if (GetThreadBinding(loop).has_value()) { ffi::ObjectPtr n = ffi::make_object(*loop); n->body = std::move(new_stmt); new_stmt = For(n); @@ -614,7 +613,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // Step 3. Collect the loop. reduction_loops.push_back(loop); // Step 4. See whether the loop is bound to some thread axis. - if (loop->GetThreadBinding().has_value()) { + if (GetThreadBinding(loop).has_value()) { need = true; } } @@ -656,8 +655,8 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // Erase those threads which are not free to this block. for (const ForNode* loop : loop_stack_) { - if (loop->GetThreadBinding().has_value()) { - ThreadScope scope = ThreadScope::Create(loop->GetThreadBinding().value()->thread_tag); + if (GetThreadBinding(loop).has_value()) { + ThreadScope scope = ThreadScope::Create(GetThreadBinding(loop).value()); thread2range.erase(scope); } } @@ -711,7 +710,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // bound to `threadIdx.x/y/z`. int n_bound_reduction_loops = 0; for (const ForNode* reduction_loop : reduction_loops) { - if (reduction_loop->GetThreadBinding().has_value()) { + if (GetThreadBinding(reduction_loop).has_value()) { ++n_bound_reduction_loops; TVM_FFI_CHECK(IsBoundToThreadIdx(reduction_loop), ValueError) << "Cross-thread reduction requires all the reduction-related loops that " @@ -794,8 +793,8 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // - we are careful about thread block boundary for safety. bool is_block_idx = false; bool is_thread_idx = false; - if (loop->IsThreadBinding()) { - ThreadScope scope = ThreadScope::Create(loop->GetThreadBinding().value()->thread_tag); + if (IsThreadBinding(loop)) { + ThreadScope scope = ThreadScope::Create(GetThreadBinding(loop).value()); if (scope.rank == 1 && scope.dim_index >= 0) { is_thread_idx = true; ++thread_idx_depth; @@ -891,10 +890,9 @@ class CrossThreadReductionTransformer : public StmtExprMutator { std::vector> reduction_threads; reduction_threads.reserve(reduction_loops.size()); for (const ForNode* loop : reduction_loops) { - if (loop->GetThreadBinding().has_value()) { - reduction_threads.emplace_back( - ThreadScope::Create(loop->GetThreadBinding().value()->thread_tag), - Range::FromMinExtent(loop->min, loop->extent)); + if (GetThreadBinding(loop).has_value()) { + reduction_threads.emplace_back(ThreadScope::Create(GetThreadBinding(loop).value()), + Range::FromMinExtent(loop->min, loop->extent)); } } for (const BufferVar& reduction_buf : reduction_buffers) { @@ -931,10 +929,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { /*extent=*/unbound_thread2range[i].second->extent, // /*kind=*/ForKind::kParallel, // /*body=*/body, // - /*thread_binding=*/ - IterVar(Range(), PrimVar("", loop_vars[i]->ty.as_or_throw()), - IterVarType::kThreadIndex, "threadIdx." + dim_index), - /*annotations=*/{}, + /*annotations=*/{{s_tir::attr::thread_binding, ffi::String("threadIdx." + dim_index)}}, /*step=*/std::nullopt); } return body; diff --git a/src/s_tir/transform/lower_opaque_block.cc b/src/s_tir/transform/lower_opaque_block.cc index eb05ac08b857..dbbe8733d222 100644 --- a/src/s_tir/transform/lower_opaque_block.cc +++ b/src/s_tir/transform/lower_opaque_block.cc @@ -94,7 +94,7 @@ class OpaqueBlockLower : public StmtExprMutator { // Step 1. Update unit loop info. PrimExpr min = this->Mutate(op->min, inplace_mode).ValueOrUnchanged(op->min); PrimExpr extent = this->Mutate(op->extent, inplace_mode).ValueOrUnchanged(op->extent); - bool has_only_thread_binding = op->IsThreadBinding() && op->annotations.size() == 1; + bool has_only_thread_binding = IsThreadBinding(op) && op->annotations.size() == 1; if (is_one(extent) && (op->annotations.empty() || has_only_thread_binding)) { // handling unit loop VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), min)); @@ -108,10 +108,10 @@ class OpaqueBlockLower : public StmtExprMutator { ffi::Map new_annotations = HandleAnnotations(op->annotations, &pragma_attrs, /*is_block=*/false); // Step 4. Create new For loop accordingly - if (op->IsThreadBinding()) { + if (IsThreadBinding(op)) { // Case 1. Thread binding - TVM_FFI_ICHECK(op->GetThreadBinding().has_value()); - ffi::String thread_tag = op->GetThreadBinding().value()->thread_tag; + TVM_FFI_ICHECK(GetThreadBinding(op).has_value()); + ffi::String thread_tag = GetThreadBinding(op).value(); body = MakeLaunchThread(min, extent, op->loop_var, thread_tag, body); } else if (is_one(extent) && op->annotations.empty() && !op->annotations.count(s_tir::attr::irregular_loop_mark)) { @@ -120,7 +120,7 @@ class OpaqueBlockLower : public StmtExprMutator { } else { // Case 3. An ordinary loop body = For(op->loop_var, std::move(min), std::move(extent), op->kind, std::move(body), - std::nullopt, new_annotations, op->step); + new_annotations, op->step); } // Step 5. Insert nested attrs for (auto it = pragma_attrs.rbegin(); it != pragma_attrs.rend(); ++it) { diff --git a/src/s_tir/transform/memhammer_coalesce.cc b/src/s_tir/transform/memhammer_coalesce.cc index de6c0e56b77b..f1cba7c2ac27 100644 --- a/src/s_tir/transform/memhammer_coalesce.cc +++ b/src/s_tir/transform/memhammer_coalesce.cc @@ -134,10 +134,9 @@ Stmt SplitBindVectorize(const Stmt& stmt, const ConstraintSet& constraints) { body = For(new_loop_vars.back().as_or_throw(), 0, vector_len, ForKind::kVectorized, std::move(body)); for (int i = n - 2; i >= 1; i--) { - body = For( - new_loop_vars[i].as_or_throw(), 0, factors[i], ForKind::kParallel, std::move(body), - IterVar(Range(nullptr), PrimVar(thread_axis[i - 1]), kThreadIndex, thread_axis[i - 1]), {}, - std::nullopt); + body = For(new_loop_vars[i].as_or_throw(), 0, factors[i], ForKind::kParallel, + std::move(body), {{s_tir::attr::thread_binding, ffi::String(thread_axis[i - 1])}}, + std::nullopt); } return For(new_loop_vars[0].as_or_throw(), 0, factors[0], ForKind::kSerial, std::move(body)); diff --git a/src/s_tir/transform/memhammer_intermediate_stage.cc b/src/s_tir/transform/memhammer_intermediate_stage.cc index 270e80798a61..3bb9c023c614 100644 --- a/src/s_tir/transform/memhammer_intermediate_stage.cc +++ b/src/s_tir/transform/memhammer_intermediate_stage.cc @@ -51,7 +51,7 @@ std::pair LiftThreadBindingLoops(Stmt stmt) { std::vector thread_binding_loops; Stmt body = stmt; while (const ForNode* loop = body.as()) { - if (loop->IsThreadBinding()) { + if (IsThreadBinding(loop)) { thread_binding_loops.push_back(loop); } else { normal_loops.push_back(loop); @@ -274,8 +274,8 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S body = op->then_case; } for (const For& loop : outer_loops) { - if (loop->IsThreadBinding()) { - const ffi::String& thread_tag = loop->GetThreadBinding().value()->thread_tag; + if (IsThreadBinding(loop.get())) { + ffi::String thread_tag = GetThreadBinding(loop.get()).value(); auto thread_scope = runtime::ThreadScope::Create(thread_tag); if (CanRelaxStorageUnderThread(runtime::StorageScope::Create(storage_scope), thread_scope)) { if (is_write_cache && thread_scope.dim_index == 0) { @@ -430,7 +430,7 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S new_loop->loop_var = new_loop_vars[i]; new_loop->body = generate_body; new_loop->kind = ForKind::kSerial; - new_loop->SetThreadBinding(std::nullopt); + new_loop->annotations.erase(s_tir::attr::thread_binding); new_loop->annotations = {}; generate_body = For(new_loop); } diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc index 482c0ec940ce..d481bbbd83d0 100644 --- a/src/s_tir/transform/memhammer_lower_auto_copy.cc +++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc @@ -516,11 +516,10 @@ class AutoPadder { } ffi::Optional Visit_(const ForNode* op) final { - if (!op->IsThreadBinding()) { + if (!IsThreadBinding(op)) { substitute_map_.Set(op->loop_var, op->min); } else { - int64_t extent = - warp_thread_extent_.Get(op->GetThreadBinding().value()->thread_tag).value_or(1); + int64_t extent = warp_thread_extent_.Get(GetThreadBinding(op).value()).value_or(1); var_range_.Set(op->loop_var, Range::FromMinExtent(op->min, IntImm::Int64(extent))); } if (op->kind == ForKind::kVectorized) { @@ -531,7 +530,7 @@ class AutoPadder { if (op->kind == ForKind::kVectorized) { vector_length_ = -1; } - if (!op->IsThreadBinding()) { + if (!IsThreadBinding(op)) { substitute_map_.erase(op->loop_var); } return std::nullopt; @@ -828,11 +827,9 @@ class ThreadExtentCollector : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } ffi::Optional Visit_(const ForNode* op) final { - if (op->GetThreadBinding().has_value() && - op->GetThreadBinding().value()->iter_type == kThreadIndex) { + if (IsThreadBinding(op)) { if (const auto* extent = op->extent.as()) { - thread_extent_.Set(op->GetThreadBinding().value()->thread_tag, - static_cast(extent->value)); + thread_extent_.Set(GetThreadBinding(op).value(), static_cast(extent->value)); } } return StmtExprVisitor::Visit_(op); diff --git a/src/s_tir/transform/profile_instrumentation.cc b/src/s_tir/transform/profile_instrumentation.cc index fc604accce90..30c5f7818728 100644 --- a/src/s_tir/transform/profile_instrumentation.cc +++ b/src/s_tir/transform/profile_instrumentation.cc @@ -92,7 +92,7 @@ class LoopAnalyzer : public StmtExprVisitor { if (has_parallel) { loop_info.has_parallel = true; parent_parallel = true; - } else if (f->IsParallel()) { + } else if (IsParallel(f)) { // has_parallel for the current loop is being set to 'false' since the // intrinsic is added outside of the loop. The instrumentation isn't // allowed for the subsequent nested loops. @@ -127,7 +127,7 @@ class LoopAnalyzer : public StmtExprVisitor { if (has_parallel) { loop_info.has_parallel = true; parent_parallel = true; - } else if (f->IsParallel()) { + } else if (IsParallel(f)) { // has_parallel for the current loop is being set to 'false' since the // intrinsic is added outside of the loop. The instrumentation isn't // allowed for the subsequent nested loops. @@ -235,7 +235,7 @@ class CheckParallelLoops : public StmtExprVisitor { private: ffi::Optional Visit_(const ForNode* op) final { - if (op->IsParallel()) { + if (IsParallel(op)) { has_parallel = true; } else { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); diff --git a/src/s_tir/transform/unify_thread_binding.cc b/src/s_tir/transform/unify_thread_binding.cc index f0a914d1fe0b..5e48eefada23 100644 --- a/src/s_tir/transform/unify_thread_binding.cc +++ b/src/s_tir/transform/unify_thread_binding.cc @@ -71,21 +71,26 @@ class ThreadBindingUnifier : public StmtExprMutator { UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { // If this For is not thread binding attribute, return as usual. - if (!op->IsThreadBinding()) { + if (!IsThreadBinding(op)) { return StmtExprMutator::Mutate_(op, inplace_mode); } ffi::Map annotations = op->annotations; - annotations.erase(tirx::attr::thread_binding); - Stmt stmt = UnifyThreadBindingImpl(op, op->loop_var, op->GetThreadBinding().value(), - Range::FromMinExtent(op->min, op->extent), inplace_mode); + annotations.erase(s_tir::attr::thread_binding); + Stmt stmt = UnifyThreadBindingImpl( + op, op->loop_var, + IterVar(Range(), op->loop_var, kThreadIndex, GetThreadBinding(op).value()), + Range::FromMinExtent(op->min, op->extent), inplace_mode); if (annotations.empty()) { return stmt; } if (const auto* loop = stmt.as()) { For new_loop = ffi::GetRef(loop); - ffi::Optional thread_binding = new_loop->GetThreadBinding(); + ffi::Optional thread_binding = GetThreadBinding(new_loop.get()); new_loop.CopyOnWrite()->annotations = std::move(annotations); - new_loop.CopyOnWrite()->SetThreadBinding(thread_binding); + if (thread_binding) { + new_loop.CopyOnWrite()->annotations.Set(s_tir::attr::thread_binding, + thread_binding.value()); + } return new_loop; } else { @@ -95,7 +100,6 @@ class ThreadBindingUnifier : public StmtExprMutator { /*min=*/IntImm(loop_ty, 0), // /*extent=*/IntImm(loop_ty, 1), // /*kind=*/ForKind::kSerial, stmt, // - /*thread_binding=*/std::nullopt, // /*annotation=*/std::move(annotations), /*step=*/std::nullopt); } @@ -165,13 +169,10 @@ class ThreadBindingUnifier : public StmtExprMutator { Stmt result = body; while (!launch_threads_.empty()) { const IterVar& thread_binding = launch_threads_.back(); - // Recreate the IterVar as we don't duplicate `dom` in both For and IterVar. This is - // necessary for unit tests. - result = - For(thread_binding->var, thread_binding->dom->min, thread_binding->dom->extent, - ForKind::kParallel, result, - IterVar(Range(), PrimVar(""), IterVarType::kThreadIndex, thread_binding->thread_tag), - {}, std::nullopt); + // The loop carries the thread tag; its variable and bounds remain ordinary For fields. + result = For(thread_binding->var, thread_binding->dom->min, thread_binding->dom->extent, + ForKind::kParallel, result, + {{s_tir::attr::thread_binding, thread_binding->thread_tag}}, std::nullopt); launch_threads_.pop_back(); } return result; diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc index 5011d62772e3..14563f34bf60 100644 --- a/src/target/llvm/codegen_cpu.cc +++ b/src/target/llvm/codegen_cpu.cc @@ -51,6 +51,7 @@ #include #include #include +#include #include #include @@ -1192,7 +1193,7 @@ void CodeGenCPU::Dispatch_(const ForNode* op) { EmitDebugLocation(op); if (op->kind == ForKind::kSerial || op->kind == ForKind::kUnrolled) { CodeGenLLVM::Dispatch_(op); - } else if (op->IsParallel()) { + } else if (s_tir::IsParallel(op)) { TVM_FFI_ICHECK(is_zero(op->min)) << "Parallel launch require canonical loop with zero start index"; TVM_FFI_ICHECK(op->HasTrivialStep()) diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc index 3497239a923f..ff632c60711a 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -465,12 +465,6 @@ UnchangedOr IndexDataTypeRewriter::Mutate_(const ForNode* op, InplaceMode n->loop_var = new_loop_var; n->min = prim::cast(new_loop_var.ty(), min); n->extent = prim::cast(new_loop_var.ty(), extent); - if (op->GetThreadBinding().has_value()) { - auto old_thread_binding = op->GetThreadBinding().value(); - auto* ptr = old_thread_binding.CopyOnWrite(); - ptr->var = old_thread_binding->var.CopyWithDType(new_loop_var.ty()); - n->SetThreadBinding(ffi::Optional(std::move(old_thread_binding))); - } n->body = new_body; return new_for; diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index 7c5c9e7c76d3..cad181e16ef9 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -210,7 +210,7 @@ TVMFFIAny AssertStmtMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, } TVMFFIAny ForVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - // Skip the kind and auxiliary hints, but traverse the semantic thread binding. + // Skip the kind and constant annotations. const ForNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( @@ -218,13 +218,12 @@ TVMFFIAny ForVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexc TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->min)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->extent)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->GetThreadBinding())); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->step)); return ffi::AnyView(nullptr).CopyToTVMFFIAny(); } TVMFFIAny ForMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - // Skip the kind and auxiliary hints, but traverse the semantic thread binding. + // Skip the kind and constant annotations. const ForNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_loop_var, @@ -237,14 +236,11 @@ TVMFFIAny ForMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noex mutator->MutateExpected(self->extent)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_body, mutator->MutateExpected(self->body)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_thread_binding, - mutator->MutateExpected(self->GetThreadBinding())); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_step, mutator->MutateExpected(self->step)); if (mapped_loop_var.UnchangedOrSameAs(self->loop_var) && mapped_min.UnchangedOrSameAs(self->min) && mapped_extent.UnchangedOrSameAs(self->extent) && mapped_body.UnchangedOrSameAs(self->body) && - mapped_thread_binding.UnchangedOrSameAs(self->GetThreadBinding()) && mapped_step.UnchangedOrSameAs(self->step)) { return ffi::Unchanged().CopyToTVMFFIAny(); } @@ -253,14 +249,12 @@ TVMFFIAny ForMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noex copy->min = std::move(mapped_min).ValueOrUnchanged(std::move(copy->min)); copy->extent = std::move(mapped_extent).ValueOrUnchanged(std::move(copy->extent)); copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body)); - copy->SetThreadBinding( - std::move(mapped_thread_binding).ValueOrUnchanged(copy->GetThreadBinding())); copy->step = std::move(mapped_step).ValueOrUnchanged(std::move(copy->step)); return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy))); } TVMFFIAny ForMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - // Skip the kind and auxiliary hints, but traverse the semantic thread binding. + // Skip the kind and constant annotations. ForNode* self = const_cast( ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_loop_var, @@ -275,15 +269,11 @@ TVMFFIAny ForMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView mutator->MutateExpected(self->extent, ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_body, mutator->MutateExpected(self->body, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr>, mapped_thread_binding, - mutator->MutateExpected(self->GetThreadBinding(), ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_step, mutator->MutateExpected(self->step, ffi::InplaceMode::kAllow)); if (mapped_loop_var.UnchangedOrSameAs(self->loop_var) && mapped_min.UnchangedOrSameAs(self->min) && mapped_extent.UnchangedOrSameAs(self->extent) && mapped_body.UnchangedOrSameAs(self->body) && - mapped_thread_binding.UnchangedOrSameAs(self->GetThreadBinding()) && mapped_step.UnchangedOrSameAs(self->step)) { return ffi::Unchanged().CopyToTVMFFIAny(); } @@ -291,9 +281,6 @@ TVMFFIAny ForMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView if (!mapped_min.IsUnchanged()) self->min = std::move(mapped_min).ValueUnchecked(); if (!mapped_extent.IsUnchanged()) self->extent = std::move(mapped_extent).ValueUnchecked(); if (!mapped_body.IsUnchanged()) self->body = std::move(mapped_body).ValueUnchecked(); - if (!mapped_thread_binding.IsUnchanged()) { - self->SetThreadBinding(std::move(mapped_thread_binding).ValueUnchecked()); - } if (!mapped_step.IsUnchanged()) self->step = std::move(mapped_step).ValueUnchecked(); return ffi::Unchanged().CopyToTVMFFIAny(); } @@ -848,8 +835,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { // For For::For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, Stmt body, - ffi::Optional thread_binding, ffi::Map annotations, - ffi::Optional step, Span span) { + ffi::Map annotations, ffi::Optional step, Span span) { TVM_FFI_CHECK(kind >= ForKind::kSerial && kind <= ForKind::kUnrolled, ValueError) << "Invalid ForKind: " << static_cast(kind); TVM_FFI_ICHECK(loop_var.defined()); @@ -907,15 +893,6 @@ For::For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, Stmt bod node->kind = kind; node->body = std::move(body); node->annotations = std::move(annotations); - if (thread_binding.has_value()) { - if (auto existing = node->GetThreadBinding()) { - TVM_FFI_CHECK(existing.value().same_as(thread_binding.value()), ValueError) - << "Conflicting thread_binding argument and annotation"; - } - node->SetThreadBinding(std::move(thread_binding)); - } else { - node->GetThreadBinding(); - } node->step = std::move(step); node->span = std::move(span); data_ = std::move(node); @@ -930,37 +907,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { .attr(refl::type_attr::kStructuralMaybeInplaceMutate, reinterpret_cast(&ForMaybeInplaceMutate)); - refl::GlobalDef().def("tirx.For", [](PrimVar loop_var, PrimExpr min, PrimExpr extent, int kind, - Stmt body, ffi::Optional thread_binding, - ffi::Optional> annotations, - ffi::Optional step, Span span) { - return For(loop_var, min, extent, static_cast(kind), body, thread_binding, - annotations.value_or(ffi::Map()), step, span); - }); -} - -ffi::Optional ForNode::GetThreadBinding() const { - auto value = annotations.Get(attr::thread_binding); - if (!value.has_value()) return std::nullopt; - TVM_FFI_CHECK(kind == ForKind::kParallel, ValueError) - << "thread_binding is only valid on parallel loops"; - auto binding = value.value().as(); - TVM_FFI_CHECK(binding.has_value(), TypeError) << "thread_binding annotation must be an IterVar"; - TVM_FFI_CHECK(!binding.value()->thread_tag.empty(), ValueError) - << "thread_binding must have a nonempty thread tag"; - return binding; -} - -void ForNode::SetThreadBinding(ffi::Optional binding) { - if (binding.has_value()) { - TVM_FFI_CHECK(kind == ForKind::kParallel, ValueError) - << "thread_binding is only valid on parallel loops"; - TVM_FFI_CHECK(!binding.value()->thread_tag.empty(), ValueError) - << "thread_binding must have a nonempty thread tag"; - annotations.Set(attr::thread_binding, binding.value()); - } else { - annotations.erase(attr::thread_binding); - } + refl::GlobalDef().def( + "tirx.For", [](PrimVar loop_var, PrimExpr min, PrimExpr extent, int kind, Stmt body, + ffi::Optional> annotations, + ffi::Optional step, Span span) { + return For(loop_var, min, extent, static_cast(kind), body, + annotations.value_or(ffi::Map()), step, span); + }); } bool ForNode::HasTrivialStep() const { return !step.has_value() || is_one(*step); } diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index f7c6811afb26..a4ca66d69966 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -551,8 +551,8 @@ PrimExpr ConvertLoopBound(const PrimExpr& e, const PrimType& var_ty) { TVM_FFI_ICHECK_EQ(doms.size(), 1); \ TVM_FFI_ICHECK_EQ(steps.size(), 1); \ return tvm::tirx::For(vars[0].as_or_throw(), doms[0]->min, doms[0]->extent, \ - Kind, body, std::nullopt, \ - annotations.value_or(ffi::Map()), steps[0]); \ + Kind, body, annotations.value_or(ffi::Map()), \ + steps[0]); \ }; \ return ForFrame(n); \ } @@ -577,17 +577,16 @@ ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, ffi::String thread, n->vars = {Var("v", dtype)}; n->doms = {Range::FromMinExtent(min, extent)}; n->steps = {std::nullopt}; - n->f_make_for_loop = [annotations, thread, dtype](ffi::Array vars, ffi::Array doms, - ffi::Array> steps, - Stmt body) -> For { + n->f_make_for_loop = [annotations, thread](ffi::Array vars, ffi::Array doms, + ffi::Array> steps, + Stmt body) -> For { TVM_FFI_ICHECK_EQ(vars.size(), 1); TVM_FFI_ICHECK_EQ(doms.size(), 1); TVM_FFI_ICHECK(steps.size() == 1 && (!steps[0].has_value() || is_one(*steps[0]))); - IterVar iter_var(Range(nullptr), tvm::PrimVar("iter", dtype), IterVarType::kThreadIndex, - thread); + auto loop_annotations = annotations.value_or(ffi::Map()); + loop_annotations.Set(s_tir::attr::thread_binding, thread); return For(vars[0].as_or_throw(), doms[0]->min, doms[0]->extent, - ForKind::kParallel, body, iter_var, - annotations.value_or(ffi::Map()), std::nullopt); + ForKind::kParallel, body, loop_annotations, std::nullopt); }; return ForFrame(n); } @@ -628,7 +627,7 @@ ForFrame Grid(ffi::Array>> Var var = vars[i]; body = For(var.as_or_throw(), dom->min, dom->extent, ForKind::kSerial, std::move(body), - /*thread_binding=*/std::nullopt, /*annotations=*/{}, /*step=*/steps[i]); + /*annotations=*/{}, /*step=*/steps[i]); } return body; }; diff --git a/src/tirx/script/printer/for_loop.cc b/src/tirx/script/printer/for_loop.cc index afe33aa6132f..42684ab2afe4 100644 --- a/src/tirx/script/printer/for_loop.cc +++ b/src/tirx/script/printer/for_loop.cc @@ -17,6 +17,7 @@ * under the License. */ #include +#include #include "./utils.h" @@ -83,7 +84,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { max = d->AsDoc(loop->min + loop->extent, loop_p->Attr("extent")); } auto loop_annotations = loop->annotations; - loop_annotations.erase(tirx::attr::thread_binding); + loop_annotations.erase(s_tir::attr::thread_binding); if (!loop_annotations.empty()) { annotations = d->AsDoc(loop_annotations, loop_p->Attr("annotations")); } @@ -96,12 +97,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { } else { prefix = TIR(d, "serial"); } - } else if (loop->IsThreadBinding()) { + } else if (s_tir::IsThreadBinding(loop.get())) { prefix = TIR(d, "thread_binding"); - thread = LiteralDoc::Str( - loop->GetThreadBinding().value()->thread_tag, - loop_p->Attr("annotations")->MapItem(tirx::attr::thread_binding)->Attr("thread_tag")); - } else if (loop->IsParallel()) { + thread = + LiteralDoc::Str(s_tir::GetThreadBinding(loop.get()).value(), + loop_p->Attr("annotations")->MapItem(s_tir::attr::thread_binding)); + } else if (s_tir::IsParallel(loop.get())) { prefix = TIR(d, "parallel"); } else if (loop->kind == tirx::ForKind::kUnrolled) { prefix = TIR(d, "unroll"); @@ -128,12 +129,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { // - annotations == {"disable_unroll": True}: print as unroll=False // - annotations == {"pragma_unroll": value}: print as unroll=value bool printed_as_unroll = false; - if (!loop->IsThreadBinding() && loop_annotations.size() == 1 && + if (!s_tir::IsThreadBinding(loop.get()) && loop_annotations.size() == 1 && loop_annotations.count("disable_unroll")) { kwargs_keys.push_back("unroll"); kwargs_values.push_back(LiteralDoc::Boolean(false, loop_p->Attr("annotations"))); printed_as_unroll = true; - } else if (!loop->IsThreadBinding() && loop_annotations.size() == 1 && + } else if (!s_tir::IsThreadBinding(loop.get()) && loop_annotations.size() == 1 && loop_annotations.count("pragma_unroll")) { kwargs_keys.push_back("unroll"); kwargs_values.push_back( diff --git a/src/tirx/transform/bind_target.cc b/src/tirx/transform/bind_target.cc index 9ccc4ff74f63..d70f25ccb547 100644 --- a/src/tirx/transform/bind_target.cc +++ b/src/tirx/transform/bind_target.cc @@ -102,7 +102,7 @@ class FunctionClassifierVisitor : public StmtExprVisitor { } ffi::Optional Visit_(const ForNode* op) final { - if (op->IsThreadBinding()) { + if (s_tir::IsThreadBinding(op)) { // Enter GPU scope for thread binding loops bool last_is_under_gpu_scope = is_under_gpu_scope_; is_under_gpu_scope_ = true; @@ -191,7 +191,7 @@ class CallSubstitutor : public StmtExprMutator { } UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { - if (op->IsThreadBinding()) { + if (s_tir::IsThreadBinding(op)) { // Enter GPU scope for thread binding loops bool last_is_under_gpu_scope = is_under_gpu_scope_; is_under_gpu_scope_ = true; diff --git a/src/tirx/transform/lower_tirx_opaque.cc b/src/tirx/transform/lower_tirx_opaque.cc index 49d4dde32531..c3549665409b 100644 --- a/src/tirx/transform/lower_tirx_opaque.cc +++ b/src/tirx/transform/lower_tirx_opaque.cc @@ -58,7 +58,7 @@ class TIRxOpaqueLower : public StmtExprMutator { // Step 1. Update unit loop info. PrimExpr min = this->Mutate(op->min, inplace_mode).ValueOrUnchanged(op->min); PrimExpr extent = this->Mutate(op->extent, inplace_mode).ValueOrUnchanged(op->extent); - bool has_only_thread_binding = op->IsThreadBinding() && op->annotations.size() == 1; + bool has_only_thread_binding = s_tir::IsThreadBinding(op) && op->annotations.size() == 1; if (is_one(extent) && (op->annotations.empty() || has_only_thread_binding)) { // handling unit loop VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), min)); @@ -72,10 +72,10 @@ class TIRxOpaqueLower : public StmtExprMutator { ffi::Map new_annotations = HandleAnnotations(op->annotations, &pragma_attrs); // Step 4. Create new For loop accordingly - if (op->IsThreadBinding()) { + if (s_tir::IsThreadBinding(op)) { // Case 1. Thread binding → AttrStmt(thread_extent) - TVM_FFI_ICHECK(op->GetThreadBinding().has_value()); - ffi::String thread_tag = op->GetThreadBinding().value()->thread_tag; + TVM_FFI_ICHECK(s_tir::GetThreadBinding(op).has_value()); + ffi::String thread_tag = s_tir::GetThreadBinding(op).value(); body = MakeLaunchThread(min, extent, op->loop_var, thread_tag, body); } else if (is_one(extent) && op->annotations.empty() && !op->annotations.count(s_tir::attr::irregular_loop_mark)) { @@ -84,7 +84,7 @@ class TIRxOpaqueLower : public StmtExprMutator { } else { // Case 3. An ordinary loop body = For(op->loop_var, std::move(min), std::move(extent), op->kind, std::move(body), - std::nullopt, new_annotations, op->step); + new_annotations, op->step); } // Step 5. Insert nested attrs for pragma annotations for (auto it = pragma_attrs.rbegin(); it != pragma_attrs.rend(); ++it) { diff --git a/src/tirx/transform/lower_tvm_builtin.cc b/src/tirx/transform/lower_tvm_builtin.cc index 0a29fe754c5c..84f0557d9718 100644 --- a/src/tirx/transform/lower_tvm_builtin.cc +++ b/src/tirx/transform/lower_tvm_builtin.cc @@ -29,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -358,7 +359,7 @@ class BuiltinLower : public StmtExprMutator { PrimExpr extent = std::move(extent_result).ValueOrUnchanged(op->extent); Stmt body; - if (op->IsParallel()) { + if (s_tir::IsParallel(op)) { body = this->VisitBodyAndRealizeAlloca(op->body); } else { body = scope_.WithNewScope([&]() -> Stmt { diff --git a/src/tirx/transform/make_packed_api.cc b/src/tirx/transform/make_packed_api.cc index afb1f59b349f..ed4491dae19d 100644 --- a/src/tirx/transform/make_packed_api.cc +++ b/src/tirx/transform/make_packed_api.cc @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -57,10 +58,10 @@ class ReturnRewriter : public StmtExprMutator { explicit ReturnRewriter(Var ret_var) : ret_var_(ret_var) {} UnchangedOr Mutate_(const ForNode* node, InplaceMode inplace_mode) override { - if (node->IsParallel()) in_parallel_ += 1; + if (s_tir::IsParallel(node)) in_parallel_ += 1; Stmt ret = StmtExprMutator::Mutate_(node, inplace_mode).ValueOrUnchanged(ffi::GetRef(node)); - if (node->IsParallel()) in_parallel_ -= 1; + if (s_tir::IsParallel(node)) in_parallel_ -= 1; return ret; } diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc index 1d2875066e0e..1cee1f2ec3cd 100644 --- a/src/tirx/transform/storage_rewrite.cc +++ b/src/tirx/transform/storage_rewrite.cc @@ -694,7 +694,7 @@ class StoragePlanRewriter : public StmtExprMutator { StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); op = stmt.as(); return For(op->loop_var, op->min, op->extent, op->kind, MakeAttach(svec, op->body), - op->GetThreadBinding(), op->annotations, op->step); + op->annotations, op->step); } else { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -1096,7 +1096,7 @@ class StoragePlanRewriter : public StmtExprMutator { } } else if (s.stmt->IsInstance()) { const auto* op = static_cast(s.stmt); - if (op->IsParallel()) { + if (s_tir::IsParallel(op)) { if (thread_scope_ == nullptr || thread_scope_ == op) { PlanNewScope(op); } diff --git a/src/tirx/transform/unroll_loop.cc b/src/tirx/transform/unroll_loop.cc index 5482ee8827c7..f657ff3f849a 100644 --- a/src/tirx/transform/unroll_loop.cc +++ b/src/tirx/transform/unroll_loop.cc @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -135,7 +136,7 @@ class LoopUnroller : public StmtExprMutator { } } // A thread-bound loop defines an execution scope even when its extent is one. - if (op->IsThreadBinding()) { + if (s_tir::IsThreadBinding(op)) { normal_loop_depth_ += 1; return result; } diff --git a/src/tirx/transform/vectorize_loop.cc b/src/tirx/transform/vectorize_loop.cc index bed3ac8f44e6..a629a9c0cb7e 100644 --- a/src/tirx/transform/vectorize_loop.cc +++ b/src/tirx/transform/vectorize_loop.cc @@ -1370,9 +1370,9 @@ class LoopVectorizer : public StmtExprMutator { Stmt body = substituter->Mutate(op->body).ValueOrUnchanged(op->body); Stmt guarded_body = IfThenElse(index < fixed_extent, body, std::nullopt, op->span); Stmt vector_loop = For(inner, IntImm(lane_dtype, 0), scalable_lanes, ForKind::kVectorized, - guarded_body, std::nullopt, op->annotations, std::nullopt, op->span); - Stmt loop = For(outer, zero, num_chunks, ForKind::kSerial, vector_loop, std::nullopt, {}, - std::nullopt, op->span); + guarded_body, op->annotations, std::nullopt, op->span); + Stmt loop = + For(outer, zero, num_chunks, ForKind::kSerial, vector_loop, {}, std::nullopt, op->span); return this->Mutate(loop, InplaceMode::kDisallow).ValueOrUnchanged(loop); } diff --git a/tests/python/relax/test_pipeline.py b/tests/python/relax/test_pipeline.py index f81b4ca6cc90..dff25bd85b67 100644 --- a/tests/python/relax/test_pipeline.py +++ b/tests/python/relax/test_pipeline.py @@ -173,7 +173,7 @@ def _has_thread_binding(func: tvm.tirx.PrimFunc) -> bool: def _visit(node): nonlocal found - if isinstance(node, tvm.tirx.For) and node.is_thread_binding(): + if isinstance(node, tvm.tirx.For) and "thread_binding" in node.annotations: found = True tvm_ffi.structural_walk(func.body, _visit) diff --git a/tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py b/tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py deleted file mode 100644 index 3636e79ac58f..000000000000 --- a/tests/python/s_tir/schedule/test_tir_schedule_binding_annotation.py +++ /dev/null @@ -1,131 +0,0 @@ -# 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. -"""Schedule mutations validate the semantic thread-binding annotation atomically.""" - -import pytest - -import tvm -import tvm.testing -from tvm import tirx -from tvm.s_tir.schedule.analysis import get_auto_tensorize_mapping_info -from tvm.script import tirx as T - - -@T.prim_func(s_tir=True) -def copy(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")): - for i in range(16): - with T.sblock("copy"): - vi = T.axis.spatial(16, i) - B[vi] = A[vi] - - -def make_schedule(kind="parallel"): - sch = tvm.s_tir.Schedule(copy, debug_mask="all") - (loop,) = sch.get_loops(sch.get_sblock("copy")) - if kind != "serial": - getattr(sch, kind)(loop) - return sch, loop - - -def test_binding_changes_use_schedule_primitives(): - sch, loop = make_schedule() - sch.annotate(loop, "hint", 7) - sch.bind(loop, "threadIdx.x") - node = sch.get(loop) - assert node.kind == tirx.ForKind.PARALLEL - assert node.thread_binding.thread_tag == "threadIdx.x" - assert node.annotations["thread_binding"].same_as(node.thread_binding) - assert node.annotations["hint"] == 7 - - sch.parallel(loop) - node = sch.get(loop) - assert node.kind == tirx.ForKind.PARALLEL - assert node.thread_binding is None - assert "thread_binding" not in node.annotations - assert node.annotations["hint"] == 7 - - -@pytest.mark.parametrize("kind", ["parallel", "serial", "vectorize", "unroll"]) -@pytest.mark.parametrize("value", ["threadIdx.x", 1]) -def test_annotate_rejects_semantic_binding_key(kind, value): - sch, loop = make_schedule(kind) - before = sch.mod.script() - trace_before = str(sch.trace) - with pytest.raises(ValueError, match="use Schedule.bind"): - sch.annotate(loop, "thread_binding", value) - assert sch.mod.script() == before - assert str(sch.trace) == trace_before - assert sch.get(loop).thread_binding is None - - -def test_unannotate_rejects_semantic_binding_key(): - sch, loop = make_schedule() - sch.bind(loop, "threadIdx.x") - before = sch.mod.script() - trace_before = str(sch.trace) - with pytest.raises(ValueError, match="use Schedule.parallel"): - sch.unannotate(loop, "thread_binding") - assert sch.mod.script() == before - assert str(sch.trace) == trace_before - assert sch.get(loop).thread_binding.thread_tag == "threadIdx.x" - - -@T.prim_func(s_tir=True) -def reduce(A: T.Buffer((16,), "float32"), B: T.Buffer((1,), "float32")): - for i in range(16): - with T.sblock("sum"): - vi = T.axis.reduce(16, i) - with T.init(): - B[0] = T.float32(0) - B[0] = B[0] + A[vi] - - -def test_bound_reduction_cannot_bypass_parallel_legality(): - sch = tvm.s_tir.Schedule(reduce, debug_mask="all") - (loop,) = sch.get_loops(sch.get_sblock("sum")) - sch.bind(loop, "threadIdx.x") - before = sch.mod.script() - with pytest.raises(ValueError, match="use Schedule.parallel"): - sch.unannotate(loop, "thread_binding") - with pytest.raises(tvm.s_tir.ScheduleError): - sch.parallel(loop) - assert sch.mod.script() == before - assert sch.get(loop).thread_binding.thread_tag == "threadIdx.x" - - -def test_sblock_thread_binding_hint_is_not_reserved(): - sch, _ = make_schedule() - block = sch.get_sblock("copy") - sch.annotate(block, "thread_binding", "block_hint") - assert sch.get(block).annotations["thread_binding"] == "block_hint" - sch.unannotate(block, "thread_binding") - assert "thread_binding" not in sch.get(block).annotations - - -def test_sblock_thread_binding_hint_in_tensorize_comparison(): - sch, _ = make_schedule("serial") - block = sch.get_sblock("copy") - sch.annotate(block, "thread_binding", "block_hint") - desc = sch.mod["main"] - assert get_auto_tensorize_mapping_info(sch, block, desc) is not None - sch.unannotate(block, "thread_binding") - sch.annotate(block, "thread_binding", "different_hint") - assert get_auto_tensorize_mapping_info(sch, block, desc) is None - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/s_tir/schedule/test_tir_schedule_state.py b/tests/python/s_tir/schedule/test_tir_schedule_state.py index 23a55cfb81bd..4a4abb99c29f 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_state.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_state.py @@ -330,7 +330,6 @@ def test_replace_block_in_opaque_block(): extent=128, kind=tirx.ForKind.SERIAL, body=tirx.Evaluate(0), - thread_binding=None, annotations=None, ) s.replace(sref, new_for_loop) diff --git a/tests/python/tirx-base/test_tir_for_thread_binding.py b/tests/python/tirx-base/test_tir_for_thread_binding.py deleted file mode 100644 index 83186ebc87b7..000000000000 --- a/tests/python/tirx-base/test_tir_for_thread_binding.py +++ /dev/null @@ -1,234 +0,0 @@ -# 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. -"""Semantic thread binding stored in For annotations.""" - -import numpy as np -import pytest -import tvm_ffi - -import tvm -import tvm.testing -from tvm import tirx -from tvm.script import tirx as T -from tvm.testing import env - - -def make_binding(tag="threadIdx.x", dom=None, var=None): - return tirx.IterVar(dom, var if var is not None else "thread", tirx.IterVar.ThreadIndex, tag) - - -def make_loop(binding=None, annotations=None, kind=tirx.ForKind.PARALLEL, extent=8): - i = tirx.Var("i", "int32") - return tirx.For(i, 0, extent, kind, tirx.Evaluate(i), binding, annotations) - - -def test_annotation_representation_and_classification(): - binding = make_binding() - loop = make_loop(binding, {"custom_hint": 7}) - assert loop.kind == tirx.ForKind.PARALLEL - assert loop.annotations["thread_binding"].same_as(binding) - assert loop.thread_binding.same_as(binding) - assert int(loop.annotations["custom_hint"]) == 7 - assert loop.is_thread_binding() - assert not loop.is_parallel() - - annotated = make_loop(annotations={"thread_binding": binding}) - assert annotated.thread_binding.same_as(binding) - assert annotated.is_thread_binding() - parallel = make_loop() - assert parallel.thread_binding is None - assert parallel.is_parallel() - assert not parallel.is_thread_binding() - serial = make_loop(kind=tirx.ForKind.SERIAL) - assert not serial.is_parallel() - assert not serial.is_thread_binding() - - -@pytest.mark.parametrize("value", ["threadIdx.x", 1]) -def test_reject_invalid_annotation_value(value): - with pytest.raises(TypeError, match="must be an IterVar"): - make_loop(annotations={"thread_binding": value}) - - -@pytest.mark.parametrize("via_annotation", [False, True]) -def test_reject_empty_thread_tag(via_annotation): - binding = make_binding("") - with pytest.raises(ValueError, match="nonempty thread tag"): - if via_annotation: - make_loop(annotations={"thread_binding": binding}) - else: - make_loop(binding) - - -@pytest.mark.parametrize( - "kind", [tirx.ForKind.SERIAL, tirx.ForKind.UNROLLED, tirx.ForKind.VECTORIZED] -) -@pytest.mark.parametrize("via_annotation", [False, True]) -def test_reject_binding_on_nonparallel_loop(kind, via_annotation): - binding = make_binding() - with pytest.raises(ValueError, match="only valid on parallel loops"): - if via_annotation: - make_loop(annotations={"thread_binding": binding}, kind=kind) - else: - make_loop(binding, kind=kind) - - -def test_reject_removed_thread_binding_kind(): - with pytest.raises(ValueError, match="Invalid ForKind"): - make_loop(make_binding(), kind=4) - - -def test_reject_conflicting_binding_sources(): - with pytest.raises(ValueError, match="Conflicting thread_binding"): - make_loop(make_binding(), {"thread_binding": make_binding("blockIdx.x")}) - - -def test_serialization_and_structural_identity(): - binding = make_binding(dom=tvm.ir.Range.from_min_extent(2, 8)) - loop = make_loop(binding) - restored = tvm.ir.load_json(tvm.ir.save_json(loop)) - tvm.ir.assert_structural_equal(restored, loop, map_free_vars=True) - assert tvm_ffi.structural_hash(restored, map_free_vars=True) == tvm_ffi.structural_hash( - loop, map_free_vars=True - ) - assert int(restored.thread_binding.dom.min) == 2 - assert int(restored.thread_binding.dom.extent) == 8 - assert restored.thread_binding.iter_type == tirx.IterVar.ThreadIndex - assert restored.thread_binding.thread_tag == "threadIdx.x" - assert restored.thread_binding.var.name == binding.var.name - assert restored.annotations["thread_binding"].same_as(restored.thread_binding) - via_annotation = make_loop(annotations={"thread_binding": binding}) - tvm.ir.assert_structural_equal(loop, via_annotation, map_free_vars=True) - assert not tvm_ffi.structural_equal(loop, make_loop(), map_free_vars=True) - assert not tvm_ffi.structural_equal( - loop, make_loop(make_binding("blockIdx.x", binding.dom)), map_free_vars=True - ) - assert not tvm_ffi.structural_equal( - loop, make_loop(make_binding(dom=tvm.ir.Range.from_min_extent(3, 8))), map_free_vars=True - ) - - -@pytest.mark.parametrize("move", [False, True]) -def test_structural_walk_and_mutation_reach_binding_metadata(move): - extent = tirx.Var("extent", "int32") - thread_var = tirx.Var("thread", "int32") - replacement = tirx.Var("new_thread", "int32") - binding = make_binding(dom=tvm.ir.Range.from_min_extent(2, extent), var=thread_var) - loop = make_loop(binding, {"custom_hint": extent}) - visited = [] - tvm_ffi.structural_walk(loop, visited.append) - assert any(node.same_as(binding) for node in visited) - assert any(node.same_as(thread_var) for node in visited) - assert any(node.same_as(extent) for node in visited) - - def rewrite(var): - if var.same_as(thread_var): - return replacement - if var.same_as(extent): - return tirx.IntImm("int32", 16) - return var - - visited.clear() # Do not retain the root while exercising ownership transfer. - rewritten = tvm_ffi.structural_map( - loop._move() if move else loop, (tirx.Var, rewrite), order="post" - ) - assert rewritten.thread_binding.var.same_as(replacement) - assert int(rewritten.thread_binding.dom.min) == 2 - assert int(rewritten.thread_binding.dom.extent) == 16 - assert rewritten.thread_binding.iter_type == binding.iter_type - assert rewritten.thread_binding.thread_tag == binding.thread_tag - # Auxiliary hints remain opaque; only the semantic annotation is traversed. - assert rewritten.annotations["custom_hint"].same_as(extent) - if not move: - assert loop.thread_binding.var.same_as(thread_var) - assert loop.thread_binding.dom.extent.same_as(extent) - - -def test_script_roundtrip_keeps_hints_separate(): - @T.prim_func - def before(): - for i in T.thread_binding(8, thread="threadIdx.x", annotations={"custom_hint": 7}): - T.evaluate(i) - - script = before.script() - assert 'annotations={"thread_binding"' not in script - restored = tvm.script.from_source(script) - tvm.ir.assert_structural_equal(before, restored) - assert restored.body.is_thread_binding() - assert int(restored.body.annotations["custom_hint"]) == 7 - - -@pytest.mark.parametrize("tag", ["threadIdx.x", "blockIdx.x", "vthread.x"]) -@pytest.mark.parametrize("extent", [1, 8]) -def test_gpu_binding_lowers_to_thread_extent(tag, extent): - loop = make_loop(make_binding(tag), extent=extent) - mod = tvm.IRModule.from_expr(tirx.PrimFunc([], loop)) - lowered = tirx.transform.LowerTIRxOpaque()(mod)["main"].body - assert isinstance(lowered, tirx.AttrStmt) - assert lowered.attr_key == ("virtual_thread" if tag == "vthread.x" else "thread_extent") - assert lowered.node.thread_tag == tag - assert int(lowered.value) == extent - if extent == 1: - assert int(lowered.body.value) == 0 - else: - assert lowered.body.value.same_as(lowered.node.var) - - -@pytest.mark.parametrize( - "config", - [ - {"auto_max_extent": 1}, - {"auto_max_step": 16, "auto_max_depth": 8, "explicit_unroll": True}, - ], -) -def test_unit_thread_binding_survives_unroll(config): - loop = make_loop(make_binding(), extent=1) - mod = tvm.IRModule.from_expr(tirx.PrimFunc([], loop)) - with tvm.transform.PassContext(config={"tirx.UnrollLoop": config}): - transformed = tirx.transform.UnrollLoop()(mod)["main"].body - assert isinstance(transformed, tirx.For) - assert transformed.is_thread_binding() - assert int(transformed.extent) == 1 - tvm.ir.assert_structural_equal(transformed, loop) - - -def test_software_pipeline_cannot_discard_thread_scope(): - loop = make_loop( - make_binding(), - {"software_pipeline_stage": [0, 1], "software_pipeline_order": [0, 1]}, - ) - mod = tvm.IRModule.from_expr(tirx.PrimFunc([], loop)) - with pytest.raises(ValueError, match="cannot replace a thread-bound loop"): - tvm.s_tir.transform.InjectSoftwarePipeline()(mod) - - -@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") -def test_cpu_parallel_execution(): - @T.prim_func - def before(out: T.Buffer((16,), "int32")): - for i in T.parallel(16): - out[i] = i + 3 - - assert before.body.is_parallel() - compiled = tvm.compile(before.with_attr("global_symbol", "main"), target="llvm") - output = np.zeros(16, dtype="int32") - compiled(output) - np.testing.assert_array_equal(output, np.arange(16, dtype="int32") + 3) - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py index 22e57c9429c7..7cc0339100ac 100644 --- a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py +++ b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py @@ -241,9 +241,7 @@ def test_ir_builder_tir_for(): extent=8, kind=tirx.ForKind.PARALLEL, body=tirx.Evaluate(0), - thread_binding=tirx.IterVar( - None, tirx.Var("", "int32"), tirx.IterVar.ThreadIndex, "threadIdx.x" - ), + annotations={"thread_binding": "threadIdx.x"}, ) unroll_expected = tirx.For( loop_var=tirx.Var("", "int32"), diff --git a/tests/python/tvmscript/test_tvmscript_parser_tir.py b/tests/python/tvmscript/test_tvmscript_parser_tir.py index a4baa9b4568a..fbadc154d54b 100644 --- a/tests/python/tvmscript/test_tvmscript_parser_tir.py +++ b/tests/python/tvmscript/test_tvmscript_parser_tir.py @@ -533,9 +533,7 @@ def func(A: T.Buffer((128, 128)), B: T.Buffer((128, 128))): loop_i = func.body loop_j = loop_i.body assert loop_i.loop_var.ty.dtype == "int64" - assert loop_i.thread_binding.var.ty.dtype == "int64" assert loop_j.loop_var.ty.dtype == "int32" - assert loop_j.thread_binding.var.ty.dtype == "int32" def test_inferred_ty_with_prim_args(): diff --git a/tests/python/tvmscript/test_tvmscript_roundtrip.py b/tests/python/tvmscript/test_tvmscript_roundtrip.py index f2088ca1288f..fc4ff0be2f60 100644 --- a/tests/python/tvmscript/test_tvmscript_roundtrip.py +++ b/tests/python/tvmscript/test_tvmscript_roundtrip.py @@ -1834,11 +1834,11 @@ def test_for_thread_binding(): tvm.ir.assert_structural_equal(func, rt_func) assert isinstance(rt_func.body, tirx.stmt.For) - assert rt_func.body.is_thread_binding() - assert rt_func.body.thread_binding.thread_tag == "threadIdx.x" + assert rt_func.body.kind == tvm.tirx.ForKind.PARALLEL + assert rt_func.body.annotations["thread_binding"] == "threadIdx.x" assert isinstance(rt_func.body.body, tirx.stmt.For) - assert rt_func.body.body.is_thread_binding() - assert rt_func.body.body.thread_binding.thread_tag == "threadIdx.y" + assert rt_func.body.body.kind == tvm.tirx.ForKind.PARALLEL + assert rt_func.body.body.annotations["thread_binding"] == "threadIdx.y" assert rt_func.body.body.annotations["attr_key"] == "attr_value" From 2b6844b9b2e120121f63b1b4a0fba89634b88afd Mon Sep 17 00:00:00 2001 From: tqchen Date: Mon, 21 Sep 2026 23:16:41 +0000 Subject: [PATCH 5/5] [STYLE][TIR] Format For mutation conditions --- src/tirx/ir/stmt.cc | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index cad181e16ef9..9055c9d95103 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -240,8 +240,7 @@ TVMFFIAny ForMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noex mutator->MutateExpected(self->step)); if (mapped_loop_var.UnchangedOrSameAs(self->loop_var) && mapped_min.UnchangedOrSameAs(self->min) && mapped_extent.UnchangedOrSameAs(self->extent) && - mapped_body.UnchangedOrSameAs(self->body) && - mapped_step.UnchangedOrSameAs(self->step)) { + mapped_body.UnchangedOrSameAs(self->body) && mapped_step.UnchangedOrSameAs(self->step)) { return ffi::Unchanged().CopyToTVMFFIAny(); } ffi::ObjectPtr copy = ffi::make_object(*self); @@ -273,8 +272,7 @@ TVMFFIAny ForMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView mutator->MutateExpected(self->step, ffi::InplaceMode::kAllow)); if (mapped_loop_var.UnchangedOrSameAs(self->loop_var) && mapped_min.UnchangedOrSameAs(self->min) && mapped_extent.UnchangedOrSameAs(self->extent) && - mapped_body.UnchangedOrSameAs(self->body) && - mapped_step.UnchangedOrSameAs(self->step)) { + mapped_body.UnchangedOrSameAs(self->body) && mapped_step.UnchangedOrSameAs(self->step)) { return ffi::Unchanged().CopyToTVMFFIAny(); } if (!mapped_loop_var.IsUnchanged()) self->loop_var = std::move(mapped_loop_var).ValueUnchecked();