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 c612f25b0b43..0e58adfc7a2b 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,11 @@ 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. + * Annotations may carry execution semantics as well as transformation hints. + * Transformations must preserve annotations until their semantics are consumed. */ ffi::Map annotations; /*! @@ -626,7 +612,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); } @@ -644,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()); @@ -886,8 +870,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..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.kind == tirx.ForKind.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/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..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") @@ -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") @@ -195,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 ThreadBinding - step : Expr The loop step. Default to none which represent one. annotations: Optional[Mapping[str, Object]] - Additional annotation hints. + Additional execution annotations and transformation hints. span : Optional[Span] The location of the stmt in the source code. @@ -215,7 +210,6 @@ class For(Stmt): extent: Expr kind: ForKind body: Stmt - thread_binding: IterVar | None annotations: Mapping[str, Object] step: Expr | None span: Span | None @@ -227,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, @@ -240,7 +233,6 @@ def __init__( extent, kind, body, - thread_binding, annotations, step, span, diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index 53a0812001aa..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->thread_binding.has_value() || rhs->thread_binding.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 0d73fb07655e..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->thread_binding.has_value()) { + if (s_tir::GetThreadBinding(op).has_value()) { const runtime::ThreadScope& scope = - runtime::ThreadScope::Create(op->thread_binding.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 939a136eecb9..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 @@ -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 (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->kind == ForKind::kThreadBinding) { - std::string thread_tag = loop->thread_binding.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 fbfe65128c99..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->kind == ForKind::kThreadBinding || !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 8c9194958b72..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->kind == ForKind::kThreadBinding) { - const ffi::String& thread_tag = loop->thread_binding.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 fa0f4c074e34..bd03fb36758e 100644 --- a/src/s_tir/schedule/ir_comparator.cc +++ b/src/s_tir/schedule/ir_comparator.cc @@ -200,20 +200,6 @@ bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) { } return false; } - if (op->thread_binding.has_value() != rhs->thread_binding.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(); - 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; diff --git a/src/s_tir/schedule/primitive/annotate.cc b/src/s_tir/schedule/primitive/annotate.cc index 9fd1bc275246..dc71c62a174e 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 != 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; 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. + GetThreadBinding(n.get()); 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 != 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; 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. + 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 6b815827878c..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->thread_binding, 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 b93e29f5b812..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->thread_binding, 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 0ad27fc00f73..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->thread_binding, 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 558aff410d43..af71d4ea8af9 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,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->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->annotations.Set(s_tir::attr::thread_binding, thread_axis.value()); } else { - new_loop->thread_binding = std::nullopt; + new_loop->annotations.erase(s_tir::attr::thread_binding); } self->Replace(loop_sref, For(new_loop), {}); } @@ -199,14 +195,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->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 57e98f5bf3cd..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->thread_binding.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->thread_binding.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->thread_binding.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->thread_binding.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 5caa524fb978..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->thread_binding; - 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->thread_binding = 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..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->kind == ForKind::kThreadBinding) { - return runtime::ThreadScope::Create(loop->thread_binding.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 d4a947144744..8f127bdf4c81 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 = 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 35686953f328..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)->thread_binding.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 a47af7fe022e..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,6 +1174,8 @@ class PipelineInjector : public StmtExprMutator { if (!HasPipelineAnnotation(op)) { return for_node; } + 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 // 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..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->kind == ForKind::kThreadBinding) { + if (IsThreadBinding(op)) { 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 = GetThreadBinding(loop).value(); { ffi::Map* tgt = &annotations[thread_tag]; for (const auto& kv : loop->annotations) { - tgt->Set(kv.first, kv.second); + if (kv.first != s_tir::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(), // + kThreadIndex, // 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 (IsThreadBinding(op.get())) { if (iter_lca.empty()) { is_kernel_root = true; SetKernelRoot(_op); @@ -145,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) { - body = For(iter_var->var, iter_var->dom->min, iter_var->dom->extent, - ForKind::kThreadBinding, std::move(body), - IterVar(Range(nullptr), PrimVar(iter_var->thread_tag, iter_var->var.ty()), - kThreadIndex, iter_var->thread_tag), - annotation, std::nullopt); + 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), annotations, std::nullopt); } } if (is_kernel_root) { iter_lca.clear(); } - if (op->kind == ForKind::kThreadBinding) { + 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 203276e4b9c8..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->kind != ForKind::kThreadBinding); + 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 fbac8846816b..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->thread_binding.has_value()) { + if (!GetThreadBinding(loop).has_value()) { return false; } - runtime::ThreadScope scope = - runtime::ThreadScope::Create(loop->thread_binding.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->thread_binding.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->thread_binding.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->thread_binding.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->thread_binding.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->thread_binding.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->thread_binding.has_value()) { - ThreadScope scope = ThreadScope::Create(loop->thread_binding.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->thread_binding.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->kind == ForKind::kThreadBinding) { - ThreadScope scope = ThreadScope::Create(loop->thread_binding.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->thread_binding.has_value()) { - reduction_threads.emplace_back( - ThreadScope::Create(loop->thread_binding.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) { @@ -929,12 +927,9 @@ 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()), - 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 c71e964d57f1..dbbe8733d222 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 = 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)); } @@ -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 (IsThreadBinding(op)) { // 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(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)) { @@ -119,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 d141c70a0ac6..f1cba7c2ac27 100644 --- a/src/s_tir/transform/memhammer_coalesce.cc +++ b/src/s_tir/transform/memhammer_coalesce.cc @@ -134,11 +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::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), {{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 baf6082cf617..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->kind == ForKind::kThreadBinding) { + 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->kind == ForKind::kThreadBinding) { - const ffi::String& thread_tag = loop->thread_binding.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->thread_binding = 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 f5c3771e58f4..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->kind != ForKind::kThreadBinding) { + if (!IsThreadBinding(op)) { 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); + 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->kind != ForKind::kThreadBinding) { + if (!IsThreadBinding(op)) { substitute_map_.erase(op->loop_var); } return std::nullopt; @@ -828,10 +827,9 @@ 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 (IsThreadBinding(op)) { if (const auto* extent = op->extent.as()) { - thread_extent_.Set(op->thread_binding.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 458f0e06de92..30c5f7818728 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 (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->kind == ForKind::kParallel) { + } 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->kind == ForKind::kParallel) { + 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 62694d603e86..5e48eefada23 100644 --- a/src/s_tir/transform/unify_thread_binding.cc +++ b/src/s_tir/transform/unify_thread_binding.cc @@ -71,18 +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->kind != ForKind::kThreadBinding) { + if (!IsThreadBinding(op)) { return StmtExprMutator::Mutate_(op, inplace_mode); } ffi::Map annotations = op->annotations; - Stmt stmt = UnifyThreadBindingImpl(op, op->loop_var, op->thread_binding.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 = GetThreadBinding(new_loop.get()); new_loop.CopyOnWrite()->annotations = std::move(annotations); + if (thread_binding) { + new_loop.CopyOnWrite()->annotations.Set(s_tir::attr::thread_binding, + thread_binding.value()); + } return new_loop; } else { @@ -92,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); } @@ -162,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::kThreadBinding, 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 fe43f3726df6..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->kind == ForKind::kParallel) { + } 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 2de7d9c75a89..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->thread_binding.has_value()) { - auto old_thread_binding = op->thread_binding.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->body = new_body; return new_for; diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index 9ed4b8a4080d..9055c9d95103 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 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->thread_binding)); 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 constant annotations. const ForNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_loop_var, @@ -237,15 +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->thread_binding)); 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_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); @@ -253,14 +248,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->thread_binding = - std::move(mapped_thread_binding).ValueOrUnchanged(std::move(copy->thread_binding)); 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 constant annotations. ForNode* self = const_cast( ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_loop_var, @@ -275,25 +268,17 @@ 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->thread_binding, 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_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(); 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->thread_binding = std::move(mapped_thread_binding).ValueUnchecked(); - } if (!mapped_step.IsUnchanged()) self->step = std::move(mapped_step).ValueUnchecked(); return ffi::Unchanged().CopyToTVMFFIAny(); } @@ -848,8 +833,9 @@ 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()); TVM_FFI_ICHECK(min.defined()); TVM_FFI_ICHECK(extent.defined()); @@ -904,7 +890,6 @@ 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); node->step = std::move(step); node->span = std::move(span); @@ -920,13 +905,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); - }); + 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); } @@ -945,9 +930,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..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::kThreadBinding, 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 df1ec3a93395..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" @@ -82,8 +83,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(s_tir::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 +97,17 @@ TVM_FFI_STATIC_INIT_BLOCK() { } else { prefix = TIR(d, "serial"); } - } else if (loop->kind == tirx::ForKind::kParallel) { + } else if (s_tir::IsThreadBinding(loop.get())) { + prefix = TIR(d, "thread_binding"); + 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"); } 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 +129,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 (!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->annotations.size() == 1 && loop->annotations.count("pragma_unroll")) { + } 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( - 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..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->kind == ForKind::kThreadBinding) { + 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->kind == ForKind::kThreadBinding) { + 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 b91742035b67..c3549665409b 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 = 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)); } @@ -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 (s_tir::IsThreadBinding(op)) { // 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(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)) { @@ -83,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 9f14fe4d95fe..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->kind == ForKind::kParallel) { + 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 70cefb5602da..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->kind == ForKind::kParallel) in_parallel_ += 1; + if (s_tir::IsParallel(node)) in_parallel_ += 1; Stmt ret = StmtExprMutator::Mutate_(node, inplace_mode).ValueOrUnchanged(ffi::GetRef(node)); - if (node->kind == ForKind::kParallel) 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 d6a6d7bdd275..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->thread_binding, 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->kind == ForKind::kParallel) { + 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 647c8da50697..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 @@ -134,6 +135,11 @@ class LoopUnroller : public StmtExprMutator { inplace_mode = InplaceMode::kDisallow; } } + // A thread-bound loop defines an execution scope even when its extent is one. + if (s_tir::IsThreadBinding(op)) { + 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/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 dd8a15f8194f..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.kind == tvm.tirx.ForKind.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_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/tvmscript/test_tvmscript_ir_builder_tir.py b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py index 2932f929a146..7cc0339100ac 100644 --- a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py +++ b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py @@ -239,11 +239,9 @@ 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" - ), + 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 3466246b69ba..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.kind == 4 - 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.kind == 4 - 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"