[REFACTOR][TIR] Represent thread binding as a parallel-loop annotation

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<ffi::String> 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.

include/tvm/tirx/stmt.h

@@ -559,22 +559,15 @@ 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.
    *  The loop body will be vectorized.
    */
   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<IterVar> 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<ffi::String, ffi::Any> 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<IterVar> thread_binding = std::nullopt,
               ffi::Map<ffi::String, ffi::Any> annotations = {},
               ffi::Optional<PrimExpr> 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();

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
 

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

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,

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

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<ScopeInfo>(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);
       }

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") {

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. */

src/s_tir/schedule/analysis/analysis.cc

@@ -703,8 +703,8 @@ ffi::Map<Var, Range> LoopDomainOfSRefTreePath(const StmtSRef& low_inclusive,
   if (extra_relax_scope.rank != runtime::StorageRank::kGlobal) {
     for (; p; p = p->parent) {
       if (const ForNode* loop = p->StmtAs<ForNode>()) {
-        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));

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;

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<ForNode>() || 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<ffi::String, ffi::Any>* annotations = nullptr;
   if (const auto* loop = sref->StmtAs<ForNode>()) {
@@ -47,6 +49,8 @@ void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_k
   if (const auto* loop = sref->StmtAs<ForNode>()) {
     ffi::ObjectPtr<ForNode> n = ffi::make_object<ForNode>(*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<SBlockNode>()) {
     ffi::ObjectPtr<SBlockNode> n = ffi::make_object<SBlockNode>(*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<ForNode>() || 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<ffi::String, ffi::Any>* annotations = nullptr;
   if (const auto* loop = sref->StmtAs<ForNode>()) {
@@ -78,6 +84,8 @@ void Unannotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann
   if (const auto* loop = sref->StmtAs<ForNode>()) {
     ffi::ObjectPtr<ForNode> n = ffi::make_object<ForNode>(*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<SBlockNode>()) {
     ffi::ObjectPtr<SBlockNode> n = ffi::make_object<SBlockNode>(*block);

src/s_tir/schedule/primitive/blockize_tensorize.cc

@@ -759,7 +759,7 @@ class BlockizeRewriter : public StmtExprMutator {
   UnchangedOr<Stmt> 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);
   }

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<Stmt> Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final {

src/s_tir/schedule/primitive/decompose_padding.cc

@@ -373,8 +373,8 @@ static std::pair<Stmt, SBlockRealize> 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;
     }