This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 3dc160cdff [REFACTOR][TIR] Store loop thread binding in annotations
(#20606)
3dc160cdff is described below
commit 3dc160cdffd2d156a90e2d839fa717b4b5afd864
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 8 11:17:59 2026 -0400
[REFACTOR][TIR] Store loop thread binding in annotations (#20606)
Thread-bound loops now use the parallel kind and store their thread tag
in `annotations["thread_binding"]`. The dedicated field and loop kind
are removed, while `kDefault` / `ForKind.DEFAULT` retain ordinary
sequential iteration semantics.
Update native and Python construction, scheduling, analysis, lowering,
and script conversion to use the annotation consistently. Preserve
`T.thread_binding` syntax and distinguish explicit thread placement from
unbound CPU parallel loops.
---
include/tvm/tirx/stmt.h | 39 +++++--------
python/tvm/s_tir/dlight/gpu/fallback.py | 5 +-
python/tvm/tirx/stmt.py | 21 +++----
src/relax/transform/split_call_tir_by_pattern.cc | 4 +-
.../analysis/sblock_buffer_access_lca_detector.cc | 5 +-
.../feature_extractor/per_store_feature.cc | 33 +++++------
.../postproc/rewrite_parallel_vectorize_unroll.cc | 2 +-
.../meta_schedule/schedule/cuda/thread_bind.cc | 2 +-
src/s_tir/schedule/analysis/analysis.cc | 4 +-
src/s_tir/schedule/ir_comparator.cc | 13 -----
src/s_tir/schedule/primitive/annotate.cc | 4 ++
src/s_tir/schedule/primitive/blockize_tensorize.cc | 2 +-
src/s_tir/schedule/primitive/cache_index.cc | 2 +-
src/s_tir/schedule/primitive/cache_read_write.cc | 6 +-
src/s_tir/schedule/primitive/compute_at.cc | 2 +-
src/s_tir/schedule/primitive/compute_inline.cc | 2 +-
src/s_tir/schedule/primitive/decompose_padding.cc | 6 +-
src/s_tir/schedule/primitive/for_kind.cc | 25 +++++----
.../schedule/primitive/layout_transformation.cc | 10 ++--
.../schedule/primitive/loop_transformation.cc | 24 ++++----
src/s_tir/schedule/primitive/pad_einsum.cc | 2 +-
src/s_tir/schedule/primitive/read_write_at.cc | 2 +-
src/s_tir/schedule/primitive/reduction.cc | 4 +-
src/s_tir/schedule/utils.h | 4 +-
src/s_tir/transform/annotate_irregular_loop.cc | 2 +-
src/s_tir/transform/compact_buffer_region.cc | 8 +--
src/s_tir/transform/default_gpu_schedule.cc | 4 +-
src/s_tir/transform/inject_software_pipeline.cc | 2 +-
src/s_tir/transform/inject_virtual_thread.cc | 4 +-
src/s_tir/transform/lift_thread_binding.cc | 15 ++---
src/s_tir/transform/loop_partition.cc | 17 +++---
.../transform/lower_cross_thread_reduction.cc | 34 ++++++-----
src/s_tir/transform/lower_opaque_block.cc | 6 +-
src/s_tir/transform/lower_thread_binding.cc | 11 ++--
.../manifest_shared_memory_local_stage.cc | 2 +-
src/s_tir/transform/memhammer_coalesce.cc | 11 ++--
.../transform/memhammer_intermediate_stage.cc | 9 ++-
src/s_tir/transform/memhammer_lower_auto_copy.cc | 10 ++--
.../transform/memhammer_tensorcore_rewrite.cc | 16 +++---
src/target/llvm/codegen_cpu.cc | 4 +-
src/target/llvm/codegen_llvm.cc | 2 +-
src/te/operation/create_function.cc | 2 +-
src/tirx/ir/stmt.cc | 51 +++++++----------
src/tirx/script/ir_builder/ir.cc | 65 ++++++++++++----------
src/tirx/script/printer/for_loop.cc | 29 +++++-----
src/tirx/transform/bind_target.cc | 4 +-
src/tirx/transform/lower_tirx_opaque.cc | 15 ++---
src/tirx/transform/lower_tvm_builtin.cc | 2 +-
src/tirx/transform/make_packed_api.cc | 6 +-
src/tirx/transform/storage_rewrite.cc | 4 +-
src/tirx/transform/unroll_loop.cc | 7 ++-
src/tirx/transform/vectorize_loop.cc | 10 ++--
tests/python/ir/test_nonnullable_ir.py | 2 +-
tests/python/relax/test_pipeline.py | 2 +-
tests/python/runtime/test_runtime_module_load.py | 2 +-
.../test_s_tir_analysis_verify_block_scope.py | 2 +-
.../s_tir/analysis/test_sblock_access_region.py | 2 +-
.../s_tir/schedule/test_tir_schedule_analysis.py | 2 +-
.../s_tir/schedule/test_tir_schedule_state.py | 3 +-
.../s_tir/script/test_s_tir_script_basic_usage.py | 8 +--
.../test_tir_analysis_verify_well_formed.py | 2 +-
tests/python/tirx-base/test_tir_constructor.py | 2 +-
tests/python/tirx-base/test_tir_nodes.py | 2 +-
.../test_tir_transform_common_subexpr_elim.py | 8 +--
.../test_tir_transform_convert_ssa.py | 2 +-
.../test_tir_transform_lower_intrin.py | 2 +-
.../test_tir_transform_remove_no_op.py | 8 +--
.../tirx/script/test_tirx_script_basic_usage.py | 4 +-
.../tirx/script/test_tirx_script_ir_builder.py | 8 +--
69 files changed, 306 insertions(+), 330 deletions(-)
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 1e90a815d3..2a2df56114 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -431,9 +431,9 @@ class IfThenElse : public Stmt {
* in all TIR passes.
*/
enum class ForKind : int {
- /*! \brief default semantics -- serial execution. */
- kSerial = 0,
- /*! \brief Parallel execution on CPU. */
+ /*! \brief Ordinary loop with sequential iteration semantics. */
+ kDefault = 0,
+ /*! \brief Parallel execution, optionally placed by the thread_binding
annotation. */
kParallel = 1,
/*!
* \brief Vector SIMD loop.
@@ -442,13 +442,6 @@ 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
};
/*!
@@ -481,18 +474,12 @@ class ForNode : public StmtNode {
ForKind kind;
/*! \brief The body of the for loop. */
SeqStmt body;
- /*!
- * \brief Only valid when kind == ForKind::kThreadBinding
- * The hardware thread tag to which this loop variable is bound.
- */
- ffi::Optional<ffi::String> 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.
+ * The reserved thread_binding entry is a string specifying the thread tag
+ * for a parallel loop. It determines execution placement and must be
+ * preserved by transformations. Other entries provide auxiliary hints.
*/
ffi::Map<ffi::String, ffi::Any> annotations;
/*!
@@ -508,11 +495,18 @@ 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);
}
+ /*! \brief Return the explicit thread placement, if present. */
+ ffi::Optional<ffi::String> GetThreadBinding() const {
+ if (auto value = annotations.Get("thread_binding")) {
+ return value->cast<ffi::String>();
+ }
+ return std::nullopt;
+ }
+
/*! \brief Check it is a loop without nontrivial loop step. */
bool HasTrivialStep() const;
@@ -526,7 +520,6 @@ class ForNode : public StmtNode {
class For : public Stmt {
public:
TVM_DLL For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind,
SeqStmt body,
- ffi::Optional<ffi::String> thread_binding = std::nullopt,
ffi::Map<ffi::String, ffi::Any> annotations = {},
ffi::Optional<PrimExpr> step = std::nullopt, Span span = Span());
@@ -741,7 +734,7 @@ TVM_DLL std::ostream& operator<<(std::ostream& os, ForKind
kind);
// inline implementations
inline const char* ForKind2String(ForKind t) {
switch (t) {
- case ForKind::kSerial:
+ case ForKind::kDefault:
return "serial";
case ForKind::kParallel:
return "parallel";
@@ -749,8 +742,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 c5dd78b63f..8fd2c87755 100644
--- a/python/tvm/s_tir/dlight/gpu/fallback.py
+++ b/python/tvm/s_tir/dlight/gpu/fallback.py
@@ -42,7 +42,7 @@ def _has_internal_thread_env(stmt: tirx.Stmt) -> bool:
def visit_for(node: tirx.For):
nonlocal found
- if node.kind == tirx.ForKind.THREAD_BINDING:
+ if node.annotations.get("thread_binding") is not None:
found = True
tvm_ffi.structural_walk(
@@ -100,7 +100,8 @@ class Fallback(GPUScheduleRule):
block = block.block_rv
if any(
- [sch.get(loop_rv).thread_binding is not None for loop_rv in
sch.get_loops(block)]
+ sch.get(loop_rv).annotations.get("thread_binding") is not None
+ for loop_rv in sch.get_loops(block)
):
continue
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index 48c0e03d79..9026d76411 100644
--- a/python/tvm/tirx/stmt.py
+++ b/python/tvm/tirx/stmt.py
@@ -130,11 +130,10 @@ class ForKind(IntEnum):
of the loop and need to be considered in all TIR passes.
"""
- SERIAL = 0
- PARALLEL = 1
+ DEFAULT = 0 # Ordinary loop with sequential iteration semantics.
+ PARALLEL = 1 # Parallel execution, optionally with explicit thread
placement.
VECTORIZED = 2
UNROLLED = 3
- THREAD_BINDING = 4 # pylint: disable=invalid-name
@tvm_ffi.register_object("tirx.For")
@@ -158,16 +157,13 @@ class For(Stmt):
body : Stmt | Sequence[Stmt]
The body statement.
- thread_binding: Optional[str]
- The thread this loop binds to. Only valid
- if kind is ThreadBinding
+ annotations: Optional[Mapping[str, Object]]
+ Additional loop annotations. For parallel loops, the optional
+ ``thread_binding`` entry is a string naming the bound thread.
+ Parallel loops without this entry use CPU parallel execution.
step : Expr
- The loop step. Default to none which
- represent one.
-
- annotations: Optional[Mapping[str, Object]]
- Additional annotation hints.
+ The loop step. Defaults to None, which represents one.
span : Optional[Span]
The location of the stmt in the source code.
@@ -178,7 +174,6 @@ class For(Stmt):
extent: Expr
kind: ForKind
body: "SeqStmt"
- thread_binding: str | None
annotations: Mapping[str, Object]
step: Expr | None
span: Span | None
@@ -190,7 +185,6 @@ class For(Stmt):
extent: Expr,
kind: ForKind,
body: Stmt | Sequence[Stmt],
- thread_binding: str | None = None,
annotations: Mapping[str, Object] | None = None,
step: Expr | None = None,
span: Span | None = None,
@@ -202,7 +196,6 @@ class For(Stmt):
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 7be76c2c46..e7a192a870 100644
--- a/src/relax/transform/split_call_tir_by_pattern.cc
+++ b/src/relax/transform/split_call_tir_by_pattern.cc
@@ -274,8 +274,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 (!IsZero(op->min) || !IsZero(rhs->min)) return false;
- if (op->thread_binding.has_value() || rhs->thread_binding.has_value())
return false;
- if (op->kind != ForKind::kSerial || op->kind != rhs->kind) return false;
+ if (op->GetThreadBinding().has_value() ||
rhs->GetThreadBinding().has_value()) return false;
+ if (op->kind != ForKind::kDefault || op->kind != rhs->kind) return false;
if (!op->annotations.empty() || !rhs->annotations.empty()) return false;
// Match the extents of loops
if (!Dispatch(op->extent, rhs->extent)) return false;
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 3568e948ca..4c52ec2126 100644
--- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
+++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
@@ -100,8 +100,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()) {
- const runtime::ThreadScope& scope =
runtime::ThreadScope::Create(op->thread_binding.value());
+ if (op->GetThreadBinding().has_value()) {
+ const runtime::ThreadScope& scope =
+ runtime::ThreadScope::Create(op->GetThreadBinding().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 c4c1a1c4ed..6499f1b304 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
@@ -311,8 +311,8 @@ Pass SimplifyForFeatureExtraction(bool
normalize_thread_bindings = false) {
if (!normalize_thread_bindings_ && IsZero(loop->extent)) {
return Evaluate(0);
}
- if (normalize_thread_bindings_ && loop->kind == ForKind::kThreadBinding)
{
- const auto& tag = loop->thread_binding.value();
+ if (normalize_thread_bindings_ && loop->GetThreadBinding().has_value()) {
+ const auto tag = loop->GetThreadBinding().value();
if (support::StartsWith(tag, "vthread")) {
// Virtual axes are independent iterations and isolate hardware
aliases.
thread_bindings_.push_back(nullptr);
@@ -322,7 +322,7 @@ Pass SimplifyForFeatureExtraction(bool
normalize_thread_bindings = false) {
}
for (auto it = thread_bindings_.rbegin(); it !=
thread_bindings_.rend() && *it; ++it) {
const ForNode* outer = *it;
- if (outer->thread_binding.value() != tag) continue;
+ if (outer->GetThreadBinding().value() != tag) continue;
// Feature counts describe hardware work, not repeated lexical names
// for one axis. Normalize only nested aliases in this analysis
input.
sym::Analyzer analyzer;
@@ -337,10 +337,11 @@ Pass SimplifyForFeatureExtraction(bool
normalize_thread_bindings = false) {
.ValueOrUnchanged(loop->annotations);
Stmt body = Mutate(loop->body,
inplace_mode).ValueOrUnchanged(loop->body);
VarRemapSet(loop->loop_var, previous_remap);
+ annotations.erase("thread_binding");
if (!annotations.empty()) {
PrimType ty = loop->loop_var.ty();
- body = For(PrimVar("annotation", ty), IntImm(ty, 0), IntImm(ty,
1), ForKind::kSerial,
- std::move(body), std::nullopt, std::move(annotations),
std::nullopt);
+ body = For(PrimVar("annotation", ty), IntImm(ty, 0), IntImm(ty,
1), ForKind::kDefault,
+ std::move(body), std::move(annotations), std::nullopt);
}
return body;
}
@@ -350,7 +351,7 @@ Pass SimplifyForFeatureExtraction(bool
normalize_thread_bindings = false) {
return result;
}
if (!normalize_thread_bindings_ && IsZero(loop->min) &&
IsOne(loop->extent) &&
- loop->kind == ForKind::kSerial && loop->annotations.empty()) {
+ loop->kind == ForKind::kDefault && loop->annotations.empty()) {
VarRemapSet(loop->loop_var, MakeConst(loop->loop_var.ty(), 0.0));
return Mutate(loop->body, inplace_mode).ValueOrUnchanged(loop->body);
} else {
@@ -420,13 +421,13 @@ struct LoopNest {
ForVec parallel; // The loops whose ForKind are kParallel
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 explicitly bound to blockIdx.x
+ ForVec blockIdx_y; // The loops explicitly bound to blockIdx.y
+ ForVec blockIdx_z; // The loops explicitly bound to blockIdx.z
+ ForVec threadIdx_x; // The loops explicitly bound to threadIdx.x
+ ForVec threadIdx_y; // The loops explicitly bound to threadIdx.y
+ ForVec threadIdx_z; // The loops explicitly bound to threadIdx.z
+ ForVec vthread; // The loops explicitly bound to vthread.*
/*!
* \brief Push a new loop into the loop nest
@@ -445,14 +446,14 @@ struct LoopNest {
this->auto_unroll.push_back(*auto_unroll_attr);
}
ForVec* ref_loops = nullptr;
- if (loop->kind == ForKind::kParallel) {
+ if (loop->kind == ForKind::kParallel &&
!loop->GetThreadBinding().has_value()) {
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();
+ } else if (loop->GetThreadBinding().has_value()) {
+ std::string thread_tag = loop->GetThreadBinding().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 74be92bc71..02a3feaa03 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->GetThreadBinding().has_value() || !loop->annotations.empty();
}
/*! \brief The visitor for extracting the stride of a var in a PrimExpr. */
diff --git a/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc
b/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc
index 843cbe5a7e..2ac35735ca 100644
--- a/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc
+++ b/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc
@@ -127,7 +127,7 @@ void BindBlockThreadIdx(Schedule sch, SBlockRV block_rv, //
i_thread_idx = i;
}
}
- if (loop->kind != ForKind::kSerial) {
+ if (loop->kind != ForKind::kDefault) {
if (i_multi_child == -1) {
i_multi_child = i;
}
diff --git a/src/s_tir/schedule/analysis/analysis.cc
b/src/s_tir/schedule/analysis/analysis.cc
index 1707a65a54..84fefce476 100644
--- a/src/s_tir/schedule/analysis/analysis.cc
+++ b/src/s_tir/schedule/analysis/analysis.cc
@@ -705,8 +705,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();
+ if (loop->GetThreadBinding().has_value()) {
+ const ffi::String thread_tag = loop->GetThreadBinding().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 bfacc7b061..4e4f00fef4 100644
--- a/src/s_tir/schedule/ir_comparator.cc
+++ b/src/s_tir/schedule/ir_comparator.cc
@@ -207,19 +207,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() && 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 170c308ea9..7cd8eb8db7 100644
--- a/src/s_tir/schedule/primitive/annotate.cc
+++ b/src/s_tir/schedule/primitive/annotate.cc
@@ -30,6 +30,8 @@ void Annotate(ScheduleState self, const StmtSRef& sref, const
ffi::String& ann_k
// Extract annotation
const ffi::Map<ffi::String, ffi::Any>* annotations = nullptr;
if (const auto* loop = sref->StmtAs<ForNode>()) {
+ TVM_FFI_CHECK(ann_key != "thread_binding", ValueError)
+ << "Use loop scheduling primitives to change thread binding";
annotations = &loop->annotations;
} else if (const auto* block = sref->StmtAs<SBlockNode>()) {
annotations = &block->annotations;
@@ -63,6 +65,8 @@ void Unannotate(ScheduleState self, const StmtSRef& sref,
const ffi::String& ann
// Extract annotation
const ffi::Map<ffi::String, ffi::Any>* annotations = nullptr;
if (const auto* loop = sref->StmtAs<ForNode>()) {
+ TVM_FFI_CHECK(ann_key != "thread_binding", ValueError)
+ << "Use loop scheduling primitives to change thread binding";
annotations = &loop->annotations;
} else if (const auto* block = sref->StmtAs<SBlockNode>()) {
annotations = &block->annotations;
diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc
b/src/s_tir/schedule/primitive/blockize_tensorize.cc
index cfeaa2ff3b..94e7b6440e 100644
--- a/src/s_tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc
@@ -760,7 +760,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);
}
diff --git a/src/s_tir/schedule/primitive/cache_index.cc
b/src/s_tir/schedule/primitive/cache_index.cc
index d925841f85..eef296fc51 100644
--- a/src/s_tir/schedule/primitive/cache_index.cc
+++ b/src/s_tir/schedule/primitive/cache_index.cc
@@ -359,7 +359,7 @@ ffi::Array<SBlock> MakeIndexCacheStage(IndexInfo* info,
const ffi::String& stora
body = For(/*loop_var=*/loop_vars[i - 1],
/*min=*/0,
/*extent=*/info->range_map.at(iter_vars[i - 1])->extent,
- /*kind=*/ForKind::kSerial,
+ /*kind=*/ForKind::kDefault,
/*body=*/body);
}
bodies.push_back(body);
diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc
b/src/s_tir/schedule/primitive/cache_read_write.cc
index 70b2eb9c9b..7598bcef06 100644
--- a/src/s_tir/schedule/primitive/cache_read_write.cc
+++ b/src/s_tir/schedule/primitive/cache_read_write.cc
@@ -236,7 +236,7 @@ SBlock MakeReindexCacheStage(const TensorRegion&
cache_region, ReindexCacheStage
body = For(/*loop_var=*/loop_vars[i - 1],
/*min=*/info->loop_ranges[i - 1]->min,
/*extent=*/info->loop_ranges[i - 1]->extent,
- /*kind=*/ForKind::kSerial,
+ /*kind=*/ForKind::kDefault,
/*body=*/body);
}
info->cache_stage = std::move(body);
@@ -339,7 +339,7 @@ SBlock MakeCacheStage(const TensorRegion& cache_region,
CacheStageInfo* info,
body = For(/*loop_var=*/loop_vars[i - 1],
/*min=*/0,
/*extent=*/cache_region->region[i - 1]->extent,
- /*kind=*/ForKind::kSerial,
+ /*kind=*/ForKind::kDefault,
/*body=*/body);
}
info->cache_stage = std::move(body);
@@ -457,7 +457,7 @@ SBlock MakeReIndexStage(const SBlock& block,
CacheStageInfo* info,
body = For(/*loop_var=*/loop_vars[i],
/*min=*/new_block_iters[i]->dom->min,
/*extent=*/new_block_iters[i]->dom->extent,
- /*kind=*/ForKind::kSerial,
+ /*kind=*/ForKind::kDefault,
/*body=*/std::move(body));
}
// Update cache info, which will be used in the later rewriting.
diff --git a/src/s_tir/schedule/primitive/compute_at.cc
b/src/s_tir/schedule/primitive/compute_at.cc
index 8fac5f6301..d5e41006e6 100644
--- a/src/s_tir/schedule/primitive/compute_at.cc
+++ b/src/s_tir/schedule/primitive/compute_at.cc
@@ -311,7 +311,7 @@ class ScopeReconstructor : public StmtExprMutator {
new_subtree = For(/*loop_var=*/loop_var.as_or_throw<PrimVar>(),
/*min=*/IntImm::Int32(0),
/*extent=*/loop_extent,
- /*ForKind=*/ForKind::kSerial,
+ /*ForKind=*/ForKind::kDefault,
/*body=*/std::move(new_subtree));
}
ffi::Array<Stmt> subtrees = AsArray(loop_->body);
diff --git a/src/s_tir/schedule/primitive/compute_inline.cc
b/src/s_tir/schedule/primitive/compute_inline.cc
index f3a2b2d4de..369920ef58 100644
--- a/src/s_tir/schedule/primitive/compute_inline.cc
+++ b/src/s_tir/schedule/primitive/compute_inline.cc
@@ -1812,7 +1812,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 {
diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc
b/src/s_tir/schedule/primitive/decompose_padding.cc
index 348652efd7..61b85fba06 100644
--- a/src/s_tir/schedule/primitive/decompose_padding.cc
+++ b/src/s_tir/schedule/primitive/decompose_padding.cc
@@ -278,7 +278,7 @@ static std::pair<Stmt, SBlockRealize>
CreateConstBlock(const SBlockRealizeNode*
for (size_t i = 0; i < new_loop_vars.size(); ++i) {
For loop = loops[i];
nest_stmt_root = For(new_loop_vars[i].as_or_throw<PrimVar>(), loop->min,
loop->extent,
- ForKind::kSerial, nest_stmt_root);
+ ForKind::kDefault, nest_stmt_root);
}
return {nest_stmt_root, new_realize};
@@ -374,8 +374,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;
}
diff --git a/src/s_tir/schedule/primitive/for_kind.cc
b/src/s_tir/schedule/primitive/for_kind.cc
index a5d35fc4f0..a18bdd7ae4 100644
--- a/src/s_tir/schedule/primitive/for_kind.cc
+++ b/src/s_tir/schedule/primitive/for_kind.cc
@@ -28,9 +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, Var
loop_var, SBlock block,
+ bool is_binding)
: mod_(std::move(mod)), loop_var_(std::move(loop_var)),
block_(std::move(block)) {
- op_str_ = for_kind == ForKind::kParallel
+ op_str_ = is_binding ? "bind"
+ : for_kind == ForKind::kParallel
? "parallel"
: (for_kind == ForKind::kVectorized ? "vectorize" : "bind");
}
@@ -73,7 +75,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 +114,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<WrongBlockIterTypeError>(self->mod, for_kind,
loop_var, block);
+ throw MakeScheduleError<WrongBlockIterTypeError>(self->mod, for_kind,
loop_var, block,
+ thread_scope.rank !=
-1);
}
}
}
@@ -122,7 +125,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 +148,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 a parallel loop to an explicit thread
*/
void ParallelizeComputation(const ScheduleState& self, const StmtSRef&
loop_sref, ForKind for_kind,
ffi::Optional<ffi::String> thread_axis) {
@@ -178,9 +181,9 @@ void ParallelizeComputation(const ScheduleState& self,
const StmtSRef& loop_sref
ffi::ObjectPtr<ForNode> new_loop = ffi::make_object<ForNode>(*loop);
new_loop->kind = for_kind;
if (thread_axis.has_value()) {
- new_loop->thread_binding = thread_axis.value();
+ new_loop->annotations.Set("thread_binding", thread_axis.value());
} else {
- new_loop->thread_binding = std::nullopt;
+ new_loop->annotations.erase("thread_binding");
}
self->Replace(loop_sref, For(new_loop), {});
}
@@ -194,14 +197,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<ForNode> new_loop = ffi::make_object<ForNode>(*loop);
new_loop->kind = ForKind::kUnrolled;
- new_loop->thread_binding = std::nullopt;
+ new_loop->annotations.erase("thread_binding");
self->Replace(loop_sref, For(new_loop), {});
}
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc
b/src/s_tir/schedule/primitive/layout_transformation.cc
index 4e9754bdce..0dc6232c73 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -204,7 +204,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
.as_or_throw<PrimExpr>();
bool is_loop_over_axis = index.same_as(loop->loop_var) &&
IsConstInt(loop->min, 0) &&
prim::ExprDeepEqual()(loop->extent,
buffer_dim) &&
- loop->kind == ForKind::kSerial;
+ loop->kind == ForKind::kDefault;
if (!is_loop_over_axis) {
return false;
}
@@ -540,7 +540,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
size_t i = (inverse->initial_indices.size() - 1) - rev_i;
Var loop_var = inverse->initial_indices[i];
PrimExpr extent = new_buffer->shape[i];
- stmt = For(loop_var.as_or_throw<PrimVar>(), 0, extent, ForKind::kSerial,
stmt);
+ stmt = For(loop_var.as_or_throw<PrimVar>(), 0, extent,
ForKind::kDefault, stmt);
}
return ProloguePlan{stmt};
}
@@ -574,7 +574,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
size_t i = (inverse->initial_indices.size() - 1) - rev_i;
Var loop_var = inverse->initial_indices[i];
PrimExpr extent = new_buffer->shape[i];
- stmt = For(loop_var.as_or_throw<PrimVar>(), 0, extent,
ForKind::kSerial, stmt);
+ stmt = For(loop_var.as_or_throw<PrimVar>(), 0, extent,
ForKind::kDefault, stmt);
}
return stmt;
@@ -634,7 +634,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
size_t i = (inverse->initial_indices.size() - 1) - rev_i;
Var loop_var = inverse->initial_indices[i];
PrimExpr extent = new_buffer->shape[i];
- stmt = For(loop_var.as_or_throw<PrimVar>(), 0, extent, ForKind::kSerial,
stmt);
+ stmt = For(loop_var.as_or_throw<PrimVar>(), 0, extent,
ForKind::kDefault, stmt);
}
const auto& info = write_info_.back();
@@ -1572,7 +1572,7 @@ void TransformBlockLayout(ScheduleState self, const
StmtSRef& block_sref,
Stmt body = ffi::GetRef<Stmt>(new_block_realize);
for (int i = static_cast<int>(new_loop_vars.size()) - 1; i >= 0; --i) {
body = For(new_loop_vars[i].as_or_throw<PrimVar>(), 0,
new_block_iter_range[i],
- ForKind::kSerial, std::move(body));
+ ForKind::kDefault, std::move(body));
}
// Step 6: Do the actual replacement
diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc
b/src/s_tir/schedule/primitive/loop_transformation.cc
index 40b0c1f5f0..22a553d1b8 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<StmtSRef> 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()) {
throw MakeScheduleError<HasAnnotationOrThreadBindingError>(self->mod,
ffi::GetRef<For>(loop));
}
// Currently, loops not starting with 0 are not supported
@@ -480,7 +480,7 @@ ffi::Array<StmtSRef> Split(ScheduleState self, const
StmtSRef& loop_sref,
// Step 4. Generate nested loops to replace the original loop and simplify
the binding
for (int i = n - 1; i >= 0; i--) {
new_stmt =
- For(new_loop_vars[i].as_or_throw<PrimVar>(), 0, factors[i],
ForKind::kSerial, new_stmt);
+ For(new_loop_vars[i].as_or_throw<PrimVar>(), 0, factors[i],
ForKind::kDefault, new_stmt);
}
new_stmt =
IterMapSimplifyBlockBinding::SimplifyBindings(std::move(new_stmt),
GetLoops(loop_sref),
opaque_block_reuse.CopyOnWrite(),
@@ -716,7 +716,7 @@ const ffi::String get_sblock_name(SeqStmt loop_body) {
ffi::Array<StmtSRef> LoopPartition(ScheduleState self, const StmtSRef&
loop_sref,
const ffi::Array<PrimExpr>& 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()) {
throw MakeScheduleError<HasAnnotationOrThreadBindingError>(self->mod,
ffi::GetRef<For>(loop));
}
@@ -759,7 +759,7 @@ ffi::Array<StmtSRef> LoopPartition(ScheduleState self,
const StmtSRef& loop_sref
.ValueOrUnchanged(std::move(loop_body));
// Create new for loop with appropriate range
auto for_node = For(new_loop_var.as_or_throw<PrimVar>(), min_value,
extent_value - min_value,
- ForKind::kSerial, loop_body);
+ ForKind::kDefault, loop_body);
const auto& partition_block_name = block_name + std::to_string(i) +
"_partition";
// Create partition_block for the partitioned for loop
@@ -838,13 +838,13 @@ class LoopReconstructor : public StmtExprMutator {
this->need_remove_loop_.push_back(loops_[i].back());
}
auto new_loop = For(new_loop_vars[0].as_or_throw<PrimVar>(),
IntImm::Int32(0),
- new_loop_extents[0], ForKind::kSerial,
SeqStmt(std::move(new_stmts)));
+ new_loop_extents[0], ForKind::kDefault,
SeqStmt(std::move(new_stmts)));
this->new_inner_loop_ = new_loop;
for (size_t i = 1; i < new_loop_vars.size(); ++i) {
const Var& loop_var = new_loop_vars[i];
const PrimExpr& loop_extent = new_loop_extents[i];
new_loop = For(loop_var.as_or_throw<PrimVar>(), IntImm::Int32(0),
loop_extent,
- ForKind::kSerial, new_loop);
+ ForKind::kDefault, new_loop);
}
this->new_outer_loop_ = new_loop;
}
@@ -899,7 +899,7 @@ StmtSRef Merge(ScheduleState self, const
ffi::Array<StmtSRef>& loop_srefs) {
std::vector<For> nest_loop_i_loops;
for (auto p = sref.get(); p != lca.get(); p = p->parent) {
if (auto loop = p->StmtAs<ForNode>()) {
- if (!loop->annotations.empty() || loop->thread_binding.has_value()) {
+ if (!loop->annotations.empty()) {
throw MakeScheduleError<HasAnnotationOrThreadBindingError>(self->mod,
ffi::GetRef<For>(loop));
}
@@ -967,7 +967,7 @@ StmtSRef Fuse(ScheduleState self, const
ffi::Array<StmtSRef>& 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()) {
throw MakeScheduleError<HasAnnotationOrThreadBindingError>(self->mod,
ffi::GetRef<For>(loop));
}
if (outer_loop_sref.defined()) {
@@ -1036,7 +1036,7 @@ StmtSRef Fuse(ScheduleState self, const
ffi::Array<StmtSRef>& loop_srefs,
fused_extent *= loops[i]->extent;
}
fused_extent = analyzer->Simplify(fused_extent);
- new_stmt = For(fused_var.as_or_throw<PrimVar>(), 0, fused_extent,
ForKind::kSerial, new_stmt);
+ new_stmt = For(fused_var.as_or_throw<PrimVar>(), 0, fused_extent,
ForKind::kDefault, new_stmt);
new_stmt = IterMapSimplifyBlockBinding::SimplifyBindings(
std::move(new_stmt), GetLoops(loop_srefs[0]),
opaque_block_reuse.CopyOnWrite(),
preserve_unit_iters);
@@ -1225,8 +1225,8 @@ void Reorder(ScheduleState self, const
ffi::Array<StmtSRef>& ordered_loop_srefs)
StmtSRef AddUnitLoop(ScheduleState self, StmtSRef sref) {
if (sref->stmt->IsInstance<ForNode>()) {
- For new_loop =
- For(PrimVar("u", PrimType::Int(32)), 0, 1, ForKind::kSerial,
ffi::GetRef<Stmt>(sref->stmt));
+ For new_loop = For(PrimVar("u", PrimType::Int(32)), 0, 1,
ForKind::kDefault,
+ ffi::GetRef<Stmt>(sref->stmt));
self->Replace(sref, new_loop, {});
return self->stmt2ref.at(new_loop.get());
}
@@ -1243,7 +1243,7 @@ StmtSRef AddUnitLoop(ScheduleState self, StmtSRef sref) {
UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* realize, InplaceMode
inplace_mode) final {
if (realize->block.get() == src_block_) {
- new_loop_ = For(PrimVar("u", PrimType::Int(32)), 0, 1,
ForKind::kSerial,
+ new_loop_ = For(PrimVar("u", PrimType::Int(32)), 0, 1,
ForKind::kDefault,
ffi::GetRef<SBlockRealize>(realize));
return new_loop_.value();
}
diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc
b/src/s_tir/schedule/primitive/pad_einsum.cc
index 912f78d141..05e6b682e5 100644
--- a/src/s_tir/schedule/primitive/pad_einsum.cc
+++ b/src/s_tir/schedule/primitive/pad_einsum.cc
@@ -204,7 +204,7 @@ struct BufferPadding {
body = SBlockRealize(prim_loop_vars, IntImm::Bool(true), new_block);
for (int i = ndim - 1; i >= 0; --i) {
body = For(loop_vars[i].as_or_throw<PrimVar>(), loop_doms[i]->min,
loop_doms[i]->extent,
- ForKind::kSerial, std::move(body));
+ ForKind::kDefault, std::move(body));
}
return body;
}
diff --git a/src/s_tir/schedule/primitive/read_write_at.cc
b/src/s_tir/schedule/primitive/read_write_at.cc
index e14179881b..b9a46e3354 100644
--- a/src/s_tir/schedule/primitive/read_write_at.cc
+++ b/src/s_tir/schedule/primitive/read_write_at.cc
@@ -337,7 +337,7 @@ struct ReadWriteAtImpl {
TensorStore(copy_to, /*value=*/MakeTensorLoad(copy_from, indices),
/*indices=*/indices);
for (int i = n - 1; i >= 0; --i) {
stmt = For(loop_vars[i].as_or_throw<PrimVar>(), IntImm::Int32(0),
domain[i]->extent,
- ForKind::kSerial, stmt);
+ ForKind::kDefault, stmt);
}
return SBlockRealize(
/*values=*/iter_values,
diff --git a/src/s_tir/schedule/primitive/reduction.cc
b/src/s_tir/schedule/primitive/reduction.cc
index a2f595d6df..cf2c5f36e2 100644
--- a/src/s_tir/schedule/primitive/reduction.cc
+++ b/src/s_tir/schedule/primitive/reduction.cc
@@ -311,10 +311,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.insert_or_assign(old_loop_var, new_loop_var);
- ffi::Optional<ffi::String> opt_thread_binding = old_loop->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<For>(new_loop);
}
@@ -1377,7 +1375,7 @@ StmtSRef RFactor(ScheduleState self, const StmtSRef&
rf_loop_sref, int factor_ax
CheckReductionBlock(self, block_sref, scope_root);
}
const ForNode* rf_loop = TVM_SREF_TO_FOR(rf_loop_sref);
- if (rf_loop->kind != ForKind::kSerial) {
+ if (rf_loop->kind != ForKind::kDefault) {
throw MakeScheduleError<NotSerialLoopKindError>(self->mod,
ffi::GetRef<For>(rf_loop));
}
diff --git a/src/s_tir/schedule/utils.h b/src/s_tir/schedule/utils.h
index 773f22cc6e..fbec8ac634 100644
--- a/src/s_tir/schedule/utils.h
+++ b/src/s_tir/schedule/utils.h
@@ -166,8 +166,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());
+ if (loop->GetThreadBinding().has_value()) {
+ return runtime::ThreadScope::Create(loop->GetThreadBinding().value());
}
return runtime::ThreadScope{-1, -1};
}
diff --git a/src/s_tir/transform/annotate_irregular_loop.cc
b/src/s_tir/transform/annotate_irregular_loop.cc
index e0d0a0b6a2..7161f0a0cf 100644
--- a/src/s_tir/transform/annotate_irregular_loop.cc
+++ b/src/s_tir/transform/annotate_irregular_loop.cc
@@ -55,7 +55,7 @@ class IrregularLoopAnnotator : public StmtExprMutator {
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
.as_or_throw<For>();
if (has_jump_) {
- TVM_FFI_ICHECK(op->kind == ForKind::kSerial)
+ TVM_FFI_ICHECK(op->kind == ForKind::kDefault)
<< "Loop kind " << op->kind << " is invalid for irregular loop " <<
op->loop_var;
for (const char* key :
{tirx::attr::auto_unroll_max_step, tirx::attr::unroll_explicit,
diff --git a/src/s_tir/transform/compact_buffer_region.cc
b/src/s_tir/transform/compact_buffer_region.cc
index d862066cf1..1112ba1966 100644
--- a/src/s_tir/transform/compact_buffer_region.cc
+++ b/src/s_tir/transform/compact_buffer_region.cc
@@ -184,10 +184,10 @@ class BufferAccessRegionCollector : public
StmtExprVisitor {
ffi::Optional<VisitInterrupt> 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())
- : IterVar(Range(), op->loop_var, IterVarType::kDataPar);
+ IterVar iter = op->GetThreadBinding().has_value()
+ ? IterVar(Range(), op->loop_var,
IterVarType::kThreadIndex,
+ op->GetThreadBinding().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 a430675125..efd769cc25 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<s_tir::LoopRV> 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;
}
}
@@ -147,7 +147,7 @@ tirx::Function WrapBareSBlockBody(const tirx::Function&
func) {
/*iter_values=*/ffi::Array<tvm::PrimExpr>{loop_var.as_or_throw<tvm::PrimExpr>()},
/*predicate=*/realize->predicate, inner_block);
tirx::Stmt for_stmt =
- tirx::For(loop_var.as_or_throw<PrimVar>(), zero, one,
tirx::ForKind::kSerial, inner_realize);
+ tirx::For(loop_var.as_or_throw<PrimVar>(), zero, one,
tirx::ForKind::kDefault, inner_realize);
s_tir::SBlock root_block(/*iter_vars=*/ffi::Array<s_tir::IterVar>{},
/*reads=*/ffi::Array<tvm::TensorRegion>{},
/*writes=*/ffi::Array<tvm::TensorRegion>{},
diff --git a/src/s_tir/transform/inject_software_pipeline.cc
b/src/s_tir/transform/inject_software_pipeline.cc
index ac3107478e..2c7e77ae72 100644
--- a/src/s_tir/transform/inject_software_pipeline.cc
+++ b/src/s_tir/transform/inject_software_pipeline.cc
@@ -1010,7 +1010,7 @@ class PipelineRewriter : public StmtExprMutator {
if (!is_unit_loop) {
new_loop = For(new_loop_var.as_or_throw<PrimVar>(), 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.
diff --git a/src/s_tir/transform/inject_virtual_thread.cc
b/src/s_tir/transform/inject_virtual_thread.cc
index 23c4a16c78..4357be5ecc 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -678,7 +678,7 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
stmt = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(stmt,
f_substitute).as_or_throw<Stmt>();
PrimType idx_dtype = idx->ty.as_or_throw<PrimType>();
return For(idx.as_or_throw<PrimVar>(), IntImm(idx_dtype, 0),
- prim::MakeConst(idx_dtype, num_threads_), ForKind::kSerial,
stmt);
+ prim::MakeConst(idx_dtype, num_threads_), ForKind::kDefault,
stmt);
}
}
@@ -737,7 +737,7 @@ class VirtualThreadInjector : public
s_tir::IRMutatorWithAnalyzer {
if (vs->has_opaque_region_) {
// Keep unknown operations and their result definitions inside one
// lexical loop; they have no sharing or distribution semantics.
- return For(var, IntImm(var.ty(), 0),
op->args[1].as_or_throw<PrimExpr>(), ForKind::kSerial,
+ return For(var, IntImm(var.ty(), 0),
op->args[1].as_or_throw<PrimExpr>(), ForKind::kDefault,
op->body);
}
auto injector =
diff --git a/src/s_tir/transform/lift_thread_binding.cc
b/src/s_tir/transform/lift_thread_binding.cc
index 476108beba..f4ec631bbd 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<VisitInterrupt> Visit_(const ForNode* op) final {
stack.push_back(ffi::GetRef<Stmt>(op));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
- if (op->kind == ForKind::kThreadBinding) {
+ if (op->GetThreadBinding().has_value()) {
UpdateLCA(op);
}
stack.pop_back();
@@ -56,7 +56,7 @@ FindLoopLCA(const Stmt& root) {
}
void UpdateLCA(const ForNode* loop) {
- std::string thread_tag = loop->thread_binding.value();
+ std::string thread_tag = loop->GetThreadBinding().value();
{
ffi::Map<ffi::String, ffi::Any>* tgt = &annotations[thread_tag];
for (const auto& kv : loop->annotations) {
@@ -134,7 +134,7 @@ class ThreadBindingLifter : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_(const ForNode* _op, InplaceMode inplace_mode)
final {
For op = ffi::GetRef<For>(_op);
bool is_kernel_root = false;
- if (op->kind == ForKind::kThreadBinding) {
+ if (op->GetThreadBinding().has_value()) {
if (iter_lca.empty()) {
is_kernel_root = true;
SetKernelRoot(_op);
@@ -146,15 +146,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), iter_var->thread_tag, annotation,
std::nullopt);
+ auto annotations = annotation;
+ annotations.Set("thread_binding", iter_var->thread_tag);
+ body = For(iter_var->var, iter_var->dom->min, iter_var->dom->extent,
ForKind::kParallel,
+ std::move(body), std::move(annotations), std::nullopt);
}
}
if (is_kernel_root) {
iter_lca.clear();
}
- if (op->kind == ForKind::kThreadBinding) {
+ if (op->GetThreadBinding().has_value()) {
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 26b8536f79..fe7c2bd0be 100644
--- a/src/s_tir/transform/loop_partition.cc
+++ b/src/s_tir/transform/loop_partition.cc
@@ -109,8 +109,8 @@ using ExpressionSet = std::unordered_set<PrimExpr,
ffi::ObjectPtrHash, ffi::Obje
// Virtual threads are lowered separately and are not hardware partition
scopes.
static bool IsVirtualThread(const ForNode* op) {
- if (op->kind != ForKind::kThreadBinding) return false;
- const auto& tag = op->thread_binding.value();
+ if (!op->GetThreadBinding().has_value()) return false;
+ const auto tag = op->GetThreadBinding().value();
return std::string(tag).rfind("vthread", 0) == 0;
}
@@ -165,8 +165,8 @@ class CandidateSelector final : public StmtExprVisitor {
candidates.insert(ffi::GetRef<Stmt>(op));
return StmtExprVisitor::Visit_(op);
}
- if (op->kind == ForKind::kThreadBinding &&
- runtime::ThreadScope::Create(op->thread_binding.value()).rank != 0) {
+ if (op->GetThreadBinding().has_value() &&
+ runtime::ThreadScope::Create(op->GetThreadBinding().value()).rank !=
0) {
return StmtExprVisitor::Visit_(op);
}
// partition const loop when sets partition_const_loop_
@@ -451,7 +451,7 @@ class ThreadPartitionInserter : public StmtExprMutator {
: ps_(ps), cond_(cond), innermost_thread_scope_(false) {}
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
- if (op->kind != ForKind::kThreadBinding || IsVirtualThread(op)) {
+ if (!op->GetThreadBinding().has_value() || IsVirtualThread(op)) {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
innermost_thread_scope_ = true;
@@ -521,7 +521,7 @@ class LoopPartitioner : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
if (IsVirtualThread(op)) return StmtExprMutator::Mutate_(op, inplace_mode);
- bool thread_scope = op->kind == ForKind::kThreadBinding;
+ bool thread_scope = op->GetThreadBinding().has_value();
analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent),
true);
auto fs = ffi::GetRef<Stmt>(op);
if (selector->candidates.count(fs)) {
@@ -531,7 +531,8 @@ class LoopPartitioner : public StmtExprMutator {
}
// Relax threadIdx ranges to avoid introducing divergent partition
branches.
- bool relax = thread_scope &&
runtime::ThreadScope::Create(op->thread_binding.value()).rank == 1;
+ bool relax =
+ thread_scope &&
runtime::ThreadScope::Create(op->GetThreadBinding().value()).rank == 1;
auto& ranges = relax ? relax_map_ : hint_map_;
ranges.insert({op->loop_var.get(), IntSet::Interval(op->min, op->min +
op->extent - 1)});
Stmt res = StmtExprMutator::Mutate_(op, InplaceMode::kDisallow)
@@ -923,7 +924,7 @@ inline Stmt LoopPartitioner::MakeFor(const ffi::Object*
node, PrimExpr extent, S
};
return ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(body,
f_substitute).as_or_throw<Stmt>();
} else {
- TVM_FFI_ICHECK(for_node->kind != ForKind::kThreadBinding);
+ TVM_FFI_ICHECK(!for_node->GetThreadBinding().has_value());
auto new_loop = ffi::make_object<ForNode>(*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 e60bafb1c0..7796f78e5f 100644
--- a/src/s_tir/transform/lower_cross_thread_reduction.cc
+++ b/src/s_tir/transform/lower_cross_thread_reduction.cc
@@ -66,10 +66,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 (!loop->GetThreadBinding().has_value()) {
return false;
}
- runtime::ThreadScope scope =
runtime::ThreadScope::Create(loop->thread_binding.value());
+ runtime::ThreadScope scope =
runtime::ThreadScope::Create(loop->GetThreadBinding().value());
return scope.rank == 1 && scope.dim_index >= 0;
}
@@ -296,7 +296,7 @@ class InThreadReducerMaker : public StmtExprMutator {
if (!body.value().same_as(loop->body)) {
res.CopyOnWrite()->body = body.value();
}
- if (res->thread_binding.has_value() &&
+ if (res->GetThreadBinding().has_value() &&
UnderLoopReductionBlockVarCollector::CheckHasReductionBlocks(res)) {
return res->body;
}
@@ -429,7 +429,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode*
realize,
}
ffi::Array<PrimExpr> thread_axes;
for (const ForNode* reduction_loop : reduction_loops) {
- if (reduction_loop->thread_binding.has_value()) {
+ if (reduction_loop->GetThreadBinding().has_value()) {
thread_axes.push_back(reduction_loop->loop_var);
}
}
@@ -559,7 +559,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode*
realize,
ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(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<PrimExpr>(loop->loop_var) ==
IntImm(loop->loop_var.ty(), 0));
}
@@ -580,7 +580,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode*
realize,
Stmt new_stmt = SeqStmt(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<ForNode> n = ffi::make_object<ForNode>(*loop);
n->body = std::move(new_stmt);
new_stmt = For(n);
@@ -627,7 +627,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;
}
}
@@ -669,8 +669,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());
+ if (loop->GetThreadBinding().has_value()) {
+ ThreadScope scope =
ThreadScope::Create(loop->GetThreadBinding().value());
thread2range.erase(scope);
}
}
@@ -724,7 +724,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 "
@@ -807,8 +807,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());
+ if (loop->GetThreadBinding().has_value()) {
+ ThreadScope scope =
ThreadScope::Create(loop->GetThreadBinding().value());
if (scope.rank == 1 && scope.dim_index >= 0) {
is_thread_idx = true;
++thread_idx_depth;
@@ -905,8 +905,8 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
std::vector<std::pair<ThreadScope, Range>> 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()),
+ if (loop->GetThreadBinding().has_value()) {
+
reduction_threads.emplace_back(ThreadScope::Create(loop->GetThreadBinding().value()),
Range::FromMinExtent(loop->min,
loop->extent));
}
}
@@ -942,11 +942,9 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
/*loop_var=*/loop_vars[i].as_or_throw<PrimVar>(), //
/*min=*/unbound_thread2range[i].second->min, //
/*extent=*/unbound_thread2range[i].second->extent, //
- /*kind=*/ForKind::kThreadBinding, //
+ /*kind=*/ForKind::kParallel, //
/*body=*/body, //
- /*thread_binding=*/
- ffi::String("threadIdx." + dim_index),
- /*annotations=*/{},
+ /*annotations=*/{{"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 5599d824a1..7ea2d3b321 100644
--- a/src/s_tir/transform/lower_opaque_block.cc
+++ b/src/s_tir/transform/lower_opaque_block.cc
@@ -111,7 +111,7 @@ class OpaqueBlockLower : public StmtExprMutator {
.as_or_throw<UnchangedOr<ffi::Optional<PrimExpr>>>()
.ValueOrUnchanged(op->step);
auto previous_remap = VarRemapGet(op->loop_var);
- if (op->kind != ForKind::kThreadBinding && IsOne(extent) &&
op->annotations.empty()) {
+ if (!op->GetThreadBinding().has_value() && IsOne(extent) &&
op->annotations.empty()) {
// handling unit loop
VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), min));
}
@@ -137,12 +137,12 @@ class OpaqueBlockLower : public StmtExprMutator {
VarRemapSet(op->loop_var, previous_remap);
// Step 2. Keep thread-binding loops until LowerThreadBinding.
- if (op->kind != ForKind::kThreadBinding && IsOne(extent) &&
op->annotations.empty() &&
+ if (!op->GetThreadBinding().has_value() && IsOne(extent) &&
op->annotations.empty() &&
!op->annotations.count(s_tir::attr::irregular_loop_mark)) {
return body;
}
return For(op->loop_var, std::move(min), std::move(extent), op->kind,
std::move(body),
- op->thread_binding, std::move(new_annotations),
std::move(step), op->span);
+ std::move(new_annotations), std::move(step), op->span);
}
void UpdateUnrollPolicy(const ffi::Map<ffi::String, ffi::Any>& annotations) {
diff --git a/src/s_tir/transform/lower_thread_binding.cc
b/src/s_tir/transform/lower_thread_binding.cc
index 1c3b46d9b2..427f16cd66 100644
--- a/src/s_tir/transform/lower_thread_binding.cc
+++ b/src/s_tir/transform/lower_thread_binding.cc
@@ -45,7 +45,7 @@ class ThreadBindingLowerer : public StmtExprMutator {
}
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
- if (op->kind != ForKind::kThreadBinding) {
+ if (!op->GetThreadBinding().has_value()) {
auto annotations = Mutate(op->annotations, inplace_mode)
.as_or_throw<UnchangedOr<ffi::Map<ffi::String,
ffi::Any>>>();
auto result = StmtExprMutator::Mutate_(op, inplace_mode);
@@ -57,7 +57,7 @@ class ThreadBindingLowerer : public StmtExprMutator {
PrimExpr min = Mutate(op->min, inplace_mode).ValueOrUnchanged(op->min);
PrimExpr extent = Mutate(op->extent,
inplace_mode).ValueOrUnchanged(op->extent);
TVM_FFI_ICHECK(IsZero(min)) << "Thread binding loops must start at zero";
- TVM_FFI_ICHECK(op->thread_binding.has_value());
+ TVM_FFI_ICHECK(op->GetThreadBinding().has_value());
TVM_FFI_ICHECK(!op->annotations.count("loop_partition_hint") ||
op->annotations.at("loop_partition_hint") == nullptr)
<< "Run LoopPartition before LowerThreadBinding";
@@ -69,12 +69,13 @@ class ThreadBindingLowerer : public StmtExprMutator {
.ValueOrUnchanged(op->annotations);
Stmt body = Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
VarRemapSet(op->loop_var, previous_remap);
+ annotations.erase("thread_binding");
if (!annotations.empty()) {
PrimType ty = op->loop_var.ty();
- body = For(PrimVar("annotation", ty), IntImm(ty, 0), IntImm(ty, 1),
ForKind::kSerial,
- std::move(body), std::nullopt, std::move(annotations),
std::nullopt);
+ body = For(PrimVar("annotation", ty), IntImm(ty, 0), IntImm(ty, 1),
ForKind::kDefault,
+ std::move(body), std::move(annotations), std::nullopt);
}
- return RegionStmt(tirx::launch_thread_op(),
{StringImm(op->thread_binding.value()), extent},
+ return RegionStmt(tirx::launch_thread_op(),
{StringImm(op->GetThreadBinding().value()), extent},
{launch_var}, DictAttrs(), std::move(body), {},
op->span);
}
};
diff --git a/src/s_tir/transform/manifest_shared_memory_local_stage.cc
b/src/s_tir/transform/manifest_shared_memory_local_stage.cc
index 251734e7ad..eecd626a98 100644
--- a/src/s_tir/transform/manifest_shared_memory_local_stage.cc
+++ b/src/s_tir/transform/manifest_shared_memory_local_stage.cc
@@ -92,7 +92,7 @@ class IntermediateStageRewriter {
const Stmt& ancestor = ancestor_loop_or_blocks_[i];
if (const ForNode* ancestor_loop = ancestor.as<ForNode>()) {
TVM_FFI_CHECK(
- ancestor_loop->kind == ForKind::kSerial || ancestor_loop->kind ==
ForKind::kVectorized,
+ ancestor_loop->kind == ForKind::kDefault || ancestor_loop->kind ==
ForKind::kVectorized,
ValueError)
<< "Expect the ancestor loops to be serial or vectorized, got " <<
ancestor_loop->kind;
relaxed_loops.push_back(ancestor.as<ForNode>());
diff --git a/src/s_tir/transform/memhammer_coalesce.cc
b/src/s_tir/transform/memhammer_coalesce.cc
index 0e010d7beb..e810fa4d0e 100644
--- a/src/s_tir/transform/memhammer_coalesce.cc
+++ b/src/s_tir/transform/memhammer_coalesce.cc
@@ -60,7 +60,7 @@ Stmt FuseNestLoops(Stmt body) {
}
body = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(std::move(body),
f_substitute)
.as_or_throw<Stmt>();
- return For(fused_var, 0, fused_extent, ForKind::kSerial, std::move(body));
+ return For(fused_var, 0, fused_extent, ForKind::kDefault, std::move(body));
}
/*!
@@ -134,10 +134,11 @@ Stmt SplitBindVectorize(const Stmt& stmt, const
ConstraintSet& constraints) {
body = For(new_loop_vars.back().as_or_throw<PrimVar>(), 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<PrimVar>(), 0, factors[i],
ForKind::kThreadBinding,
- std::move(body), ffi::String(thread_axis[i - 1]), {},
std::nullopt);
+ body =
+ For(new_loop_vars[i].as_or_throw<PrimVar>(), 0, factors[i],
ForKind::kParallel,
+ std::move(body), {{"thread_binding", ffi::String(thread_axis[i -
1])}}, std::nullopt);
}
- return For(new_loop_vars[0].as_or_throw<PrimVar>(), 0, factors[0],
ForKind::kSerial,
+ return For(new_loop_vars[0].as_or_throw<PrimVar>(), 0, factors[0],
ForKind::kDefault,
std::move(body));
}
@@ -242,7 +243,7 @@ Stmt InverseMapping::Rewrite(const Stmt& stmt, const
ConstraintSet& constraints,
// Step 3.3 construct loop body
for (int i = static_cast<int>(new_loop_vars.size()) - 1; i >= 0; i--) {
PrimExpr extent = write_region->region[i]->extent;
- ret = For(new_loop_vars[i], 0, extent, ForKind::kSerial, std::move(ret));
+ ret = For(new_loop_vars[i], 0, extent, ForKind::kDefault, std::move(ret));
}
return ret;
}
diff --git a/src/s_tir/transform/memhammer_intermediate_stage.cc
b/src/s_tir/transform/memhammer_intermediate_stage.cc
index 95d484404a..d26a78a0d1 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<Stmt, ffi::Optional<For>>
LiftThreadBindingLoops(Stmt stmt) {
std::vector<const ForNode*> thread_binding_loops;
Stmt body = stmt;
while (const ForNode* loop = body.as<ForNode>()) {
- if (loop->kind == ForKind::kThreadBinding) {
+ if (loop->GetThreadBinding().has_value()) {
thread_binding_loops.push_back(loop);
} else {
normal_loops.push_back(loop);
@@ -275,8 +275,8 @@ std::pair<Stmt, SeqStmt> InsertCacheStage(Stmt stmt, bool
is_write_cache, ffi::S
body = op->then_case->seq[0];
}
for (const For& loop : outer_loops) {
- if (loop->kind == ForKind::kThreadBinding) {
- const ffi::String& thread_tag = loop->thread_binding.value();
+ if (loop->GetThreadBinding().has_value()) {
+ const ffi::String thread_tag = loop->GetThreadBinding().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,8 +430,7 @@ std::pair<Stmt, SeqStmt> InsertCacheStage(Stmt stmt, bool
is_write_cache, ffi::S
ffi::ObjectPtr<ForNode> new_loop = ffi::make_object<ForNode>(*orig_loop);
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->kind = ForKind::kDefault;
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 c64e87d4e2..c95536cb0f 100644
--- a/src/s_tir/transform/memhammer_lower_auto_copy.cc
+++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc
@@ -516,10 +516,10 @@ class AutoPadder {
}
ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
- if (op->kind != ForKind::kThreadBinding) {
+ if (!op->GetThreadBinding().has_value()) {
substitute_map_.Set(op->loop_var, op->min);
} else {
- int64_t extent =
warp_thread_extent_.Get(op->thread_binding.value()).value_or(1);
+ int64_t extent =
warp_thread_extent_.Get(op->GetThreadBinding().value()).value_or(1);
var_range_.Set(op->loop_var, Range::FromMinExtent(op->min,
IntImm::Int64(extent)));
}
if (op->kind == ForKind::kVectorized) {
@@ -530,7 +530,7 @@ class AutoPadder {
if (op->kind == ForKind::kVectorized) {
vector_length_ = -1;
}
- if (op->kind != ForKind::kThreadBinding) {
+ if (!op->GetThreadBinding().has_value()) {
substitute_map_.erase(op->loop_var);
}
return std::nullopt;
@@ -831,9 +831,9 @@ class ThreadExtentCollector : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
- if (op->thread_binding.has_value()) {
+ if (op->GetThreadBinding().has_value()) {
if (const auto* extent = op->extent.as<IntImmNode>()) {
- thread_extent_.Set(op->thread_binding.value(),
static_cast<int64_t>(extent->value));
+ thread_extent_.Set(op->GetThreadBinding().value(),
static_cast<int64_t>(extent->value));
}
}
return StmtExprVisitor::Visit_(op);
diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
index e69f23a916..121bc1e509 100644
--- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
+++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
@@ -78,13 +78,13 @@ std::pair<Stmt, ffi::Optional<For>> TileWmmaBlock(Stmt
stmt) {
/*3:*/ 16, //
/*4:*/ 16, //
};
- body = For(new_loop_vars[3].as_or_throw<PrimVar>(), 0, factor[3],
ForKind::kSerial,
+ body = For(new_loop_vars[3].as_or_throw<PrimVar>(), 0, factor[3],
ForKind::kDefault,
std::move(body));
- body = For(new_loop_vars[2].as_or_throw<PrimVar>(), 0, factor[2],
ForKind::kSerial,
+ body = For(new_loop_vars[2].as_or_throw<PrimVar>(), 0, factor[2],
ForKind::kDefault,
std::move(body));
- body = For(new_loop_vars[1].as_or_throw<PrimVar>(), 0, factor[1],
ForKind::kSerial,
+ body = For(new_loop_vars[1].as_or_throw<PrimVar>(), 0, factor[1],
ForKind::kDefault,
std::move(body));
- body = For(new_loop_vars[0].as_or_throw<PrimVar>(), 0, factor[0],
ForKind::kSerial,
+ body = For(new_loop_vars[0].as_or_throw<PrimVar>(), 0, factor[0],
ForKind::kDefault,
std::move(body));
}
For compute_location = body.as_or_throw<For>();
@@ -409,13 +409,13 @@ std::pair<Stmt, ffi::Optional<For>>
TileMmaToGlobalBlock(Stmt stmt) {
/*3:*/ 8, //
/*4:*/ 8, //
};
- body = For(new_loop_vars[3].as_or_throw<PrimVar>(), 0, factor[3],
ForKind::kSerial,
+ body = For(new_loop_vars[3].as_or_throw<PrimVar>(), 0, factor[3],
ForKind::kDefault,
std::move(body));
- body = For(new_loop_vars[2].as_or_throw<PrimVar>(), 0, factor[2],
ForKind::kSerial,
+ body = For(new_loop_vars[2].as_or_throw<PrimVar>(), 0, factor[2],
ForKind::kDefault,
std::move(body));
- body = For(new_loop_vars[1].as_or_throw<PrimVar>(), 0, factor[1],
ForKind::kSerial,
+ body = For(new_loop_vars[1].as_or_throw<PrimVar>(), 0, factor[1],
ForKind::kDefault,
std::move(body));
- body = For(new_loop_vars[0].as_or_throw<PrimVar>(), 0, factor[0],
ForKind::kSerial,
+ body = For(new_loop_vars[0].as_or_throw<PrimVar>(), 0, factor[0],
ForKind::kDefault,
std::move(body));
}
For compute_location = body.as_or_throw<For>();
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index f2d75a5e72..82ad33fffc 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -1147,9 +1147,9 @@ void CodeGenCPU::Dispatch_(const RegionStmtNode* op) {
void CodeGenCPU::Dispatch_(const ForNode* op) {
EmitDebugLocation(op);
- if (op->kind == ForKind::kSerial || op->kind == ForKind::kUnrolled) {
+ if (op->kind == ForKind::kDefault || op->kind == ForKind::kUnrolled) {
CodeGenLLVM::Dispatch_(op);
- } else if (op->kind == ForKind::kParallel) {
+ } else if (op->kind == ForKind::kParallel &&
!op->GetThreadBinding().has_value()) {
TVM_FFI_ICHECK(IsZero(op->min))
<< "Parallel launch require canonical loop with zero start index";
TVM_FFI_ICHECK(op->HasTrivialStep())
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 972c8d1199..afcb507c8d 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -2100,7 +2100,7 @@ void CodeGenLLVM::Dispatch_(const ForNode* op) {
LOG(WARNING) << "Unroll hint get ignore at CodeGenLLVM backend, "
<< " consider set unroll_explicit=True";
} else {
- TVM_FFI_ICHECK(op->kind == ForKind::kSerial);
+ TVM_FFI_ICHECK(op->kind == ForKind::kDefault);
}
PrimExpr step = op->step.value_or(IntImm(op->extent.ty(), 1));
PrimExpr end = IsZero(op->min) ? op->extent : analyzer_->Simplify(op->min +
op->extent);
diff --git a/src/te/operation/create_function.cc
b/src/te/operation/create_function.cc
index 1f2b80fb46..5ac18fd6cc 100644
--- a/src/te/operation/create_function.cc
+++ b/src/te/operation/create_function.cc
@@ -743,7 +743,7 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp&
compute_op, CreateFuncInfo* in
}
for (size_t j = cur.loop_vars.size(); j > 0; --j) {
const auto& [loop_var, dom] = cur.loop_vars[j - 1];
- body = For(loop_var.as_or_throw<PrimVar>(), dom->min, dom->extent,
ForKind::kSerial, body);
+ body = For(loop_var.as_or_throw<PrimVar>(), dom->min, dom->extent,
ForKind::kDefault, body);
}
}
return body;
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index 377957433f..270ffd5674 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -283,7 +283,6 @@ TVM_FFI_INLINE
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> ForVisit(
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 std::nullopt;
}
@@ -303,16 +302,11 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
ForMutate(
mutator->MutateExpected(self->extent));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body,
mutator->MutateExpected(self->body));
-
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<ffi::String>>,
- mapped_thread_binding,
-
mutator->MutateExpected(self->thread_binding));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<PrimExpr>>,
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();
}
ffi::ObjectPtr<ForNode> copy = ffi::make_object<ForNode>(*self);
@@ -320,8 +314,6 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
ForMutate(
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::Any(std::move(copy));
}
@@ -343,25 +335,17 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
ForMaybeInplaceMutate(
mutator->MutateExpected(self->extent, ffi::InplaceMode::kAllow));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body,
mutator->MutateExpected(self->body,
ffi::InplaceMode::kAllow));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
- ffi::UnchangedOr<ffi::Optional<ffi::String>>, mapped_thread_binding,
- mutator->MutateExpected(self->thread_binding, ffi::InplaceMode::kAllow));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<PrimExpr>>,
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();
}
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 =
SeqStmt(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();
}
@@ -855,8 +839,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// For
For::For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind,
SeqStmt body,
- ffi::Optional<ffi::String> thread_binding, ffi::Map<ffi::String, Any>
annotations,
- ffi::Optional<PrimExpr> step, Span span)
+ ffi::Map<ffi::String, Any> annotations, ffi::Optional<PrimExpr> step,
Span span)
: Stmt(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(loop_var.defined());
TVM_FFI_ICHECK(min.defined());
@@ -908,8 +891,15 @@ For::For(PrimVar loop_var, PrimExpr min, PrimExpr extent,
ForKind kind, SeqStmt
ffi::ObjectPtr<ForNode> node =
ffi::make_object<ForNode>(std::move(loop_var), std::move(min),
std::move(extent),
std::move(body));
+ if (auto tag = annotations.Get("thread_binding")) {
+ TVM_FFI_CHECK(kind == ForKind::kParallel, ValueError)
+ << "thread_binding requires a parallel loop";
+ TVM_FFI_CHECK(tag->as<ffi::String>().has_value(), TypeError)
+ << "thread_binding annotation must be a string";
+ TVM_FFI_CHECK(!step.has_value() || IsOne(*step), ValueError)
+ << "Thread binding loops require a unit step";
+ }
node->kind = kind;
- node->thread_binding = std::move(thread_binding);
node->annotations = std::move(annotations);
node->step = std::move(step);
node->span = std::move(span);
@@ -925,20 +915,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
ffi::FStructuralMutate::FromNative<&ForMaybeInplaceMutate>());
- refl::GlobalDef().def("tirx.For", [](PrimVar loop_var, PrimExpr min,
PrimExpr extent, int kind,
- SeqStmt body,
ffi::Optional<ffi::String> thread_binding,
- ffi::Optional<ffi::Map<ffi::String,
Any>> annotations,
- ffi::Optional<PrimExpr> step, Span
span) {
- return For(loop_var, min, extent, static_cast<ForKind>(kind), body,
thread_binding,
- annotations.value_or(ffi::Map<ffi::String, Any>()), step, span);
- });
+ refl::GlobalDef().def(
+ "tirx.For", [](PrimVar loop_var, PrimExpr min, PrimExpr extent, int
kind, SeqStmt body,
+ ffi::Optional<ffi::Map<ffi::String, Any>> annotations,
+ ffi::Optional<PrimExpr> step, Span span) {
+ return For(loop_var, min, extent, static_cast<ForKind>(kind), body,
+ annotations.value_or(ffi::Map<ffi::String, Any>()), step,
span);
+ });
}
bool ForNode::HasTrivialStep() const { return !step.has_value() ||
IsOne(*step); }
std::ostream& operator<<(std::ostream& out, ForKind type) { // NOLINT(*)
switch (type) {
- case ForKind::kSerial:
+ case ForKind::kDefault:
out << "for";
break;
case ForKind::kParallel:
@@ -950,9 +940,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/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index 8ca0a71a61..2a8af91843 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -281,34 +281,34 @@ PrimExpr ConvertLoopBound(const PrimExpr& e, const
PrimType& var_ty) {
return tvm::prim::Cast(var_ty, e);
}
-#define TVM_TIRX_IR_BUILDER_FOR_FRAME(Method, Kind)
\
- ForFrame Method(PrimExpr start, PrimExpr stop,
\
- ffi::Optional<ffi::Map<ffi::String, Any>> annotations,
\
- ffi::Optional<PrimExpr> step, ffi::Optional<PrimType> dtype)
{ \
- PrimType var_ty = InferLoopVarDtype(start, stop, dtype);
\
- PrimExpr min = ConvertLoopBound(start, var_ty);
\
- PrimExpr extent = sym::Analyzer()->Simplify(ConvertLoopBound(stop, var_ty)
- min); \
- if (step.has_value()) {
\
- step = ConvertLoopBound(step.value(), var_ty);
\
- }
\
- ffi::ObjectPtr<ForFrameNode> n = ffi::make_object<ForFrameNode>();
\
- n->vars = {Var("v", var_ty)};
\
- n->doms = {Range::FromMinExtent(min, extent)};
\
- n->steps = {step};
\
- n->f_make_for_loop = [annotations](ffi::Array<Var> vars, ffi::Array<Range>
doms, \
- ffi::Array<ffi::Optional<PrimExpr>>
steps, \
- tvm::tirx::SeqStmt body, Span span) {
\
- TVM_FFI_ICHECK_EQ(vars.size(), 1);
\
- TVM_FFI_ICHECK_EQ(doms.size(), 1);
\
- TVM_FFI_ICHECK_EQ(steps.size(), 1);
\
- return tvm::tirx::For(vars[0].as_or_throw<tvm::PrimVar>(), doms[0]->min,
doms[0]->extent, \
- Kind, body, std::nullopt,
\
- annotations.value_or(ffi::Map<ffi::String,
Any>()), steps[0], span); \
- };
\
- return ForFrame(n);
\
+#define TVM_TIRX_IR_BUILDER_FOR_FRAME(Method, Kind)
\
+ ForFrame Method(PrimExpr start, PrimExpr stop,
\
+ ffi::Optional<ffi::Map<ffi::String, Any>> annotations,
\
+ ffi::Optional<PrimExpr> step, ffi::Optional<PrimType> dtype)
{ \
+ PrimType var_ty = InferLoopVarDtype(start, stop, dtype);
\
+ PrimExpr min = ConvertLoopBound(start, var_ty);
\
+ PrimExpr extent = sym::Analyzer()->Simplify(ConvertLoopBound(stop, var_ty)
- min); \
+ if (step.has_value()) {
\
+ step = ConvertLoopBound(step.value(), var_ty);
\
+ }
\
+ ffi::ObjectPtr<ForFrameNode> n = ffi::make_object<ForFrameNode>();
\
+ n->vars = {Var("v", var_ty)};
\
+ n->doms = {Range::FromMinExtent(min, extent)};
\
+ n->steps = {step};
\
+ n->f_make_for_loop = [annotations](ffi::Array<Var> vars, ffi::Array<Range>
doms, \
+ ffi::Array<ffi::Optional<PrimExpr>>
steps, \
+ tvm::tirx::SeqStmt body, Span span) {
\
+ TVM_FFI_ICHECK_EQ(vars.size(), 1);
\
+ TVM_FFI_ICHECK_EQ(doms.size(), 1);
\
+ TVM_FFI_ICHECK_EQ(steps.size(), 1);
\
+ return tvm::tirx::For(vars[0].as_or_throw<tvm::PrimVar>(), doms[0]->min,
doms[0]->extent, \
+ Kind, body,
annotations.value_or(ffi::Map<ffi::String, Any>()), \
+ steps[0], span);
\
+ };
\
+ return ForFrame(n);
\
}
-TVM_TIRX_IR_BUILDER_FOR_FRAME(Serial, tvm::tirx::ForKind::kSerial);
+TVM_TIRX_IR_BUILDER_FOR_FRAME(Serial, tvm::tirx::ForKind::kDefault);
TVM_TIRX_IR_BUILDER_FOR_FRAME(Parallel, tvm::tirx::ForKind::kParallel);
TVM_TIRX_IR_BUILDER_FOR_FRAME(Vectorized, tvm::tirx::ForKind::kVectorized);
TVM_TIRX_IR_BUILDER_FOR_FRAME(Unroll, tvm::tirx::ForKind::kUnrolled);
@@ -334,9 +334,14 @@ ForFrame ThreadBinding(PrimExpr start, PrimExpr stop,
ffi::String thread,
TVM_FFI_ICHECK_EQ(vars.size(), 1);
TVM_FFI_ICHECK_EQ(doms.size(), 1);
TVM_FFI_ICHECK(steps.size() == 1 && (!steps[0].has_value() ||
IsOne(*steps[0])));
+ auto loop_annotations = annotations.value_or(ffi::Map<ffi::String,
ffi::Any>());
+ if (auto existing = loop_annotations.Get("thread_binding")) {
+ TVM_FFI_CHECK(existing->cast<ffi::String>() == thread, ValueError)
+ << "Conflicting thread_binding annotation and thread argument";
+ }
+ loop_annotations.Set("thread_binding", thread);
return For(vars[0].as_or_throw<tvm::PrimVar>(), doms[0]->min,
doms[0]->extent,
- ForKind::kThreadBinding, body, thread,
- annotations.value_or(ffi::Map<ffi::String, ffi::Any>()),
std::nullopt, span);
+ ForKind::kParallel, body, std::move(loop_annotations),
std::nullopt, span);
};
return ForFrame(n);
}
@@ -377,9 +382,9 @@ ForFrame Grid(ffi::Array<ffi::Variant<PrimExpr,
ffi::Tuple<PrimExpr, PrimExpr>>>
for (int i = n - 1; i >= 0; --i) {
Range dom = doms[i];
Var var = vars[i];
- result = For(var.as_or_throw<tvm::PrimVar>(), dom->min, dom->extent,
ForKind::kSerial,
+ result = For(var.as_or_throw<tvm::PrimVar>(), dom->min, dom->extent,
ForKind::kDefault,
SeqStmt(std::move(result)),
- /*thread_binding=*/std::nullopt, /*annotations=*/{},
/*step=*/steps[i], span);
+ /*annotations=*/{}, /*step=*/steps[i], span);
}
return result;
};
diff --git a/src/tirx/script/printer/for_loop.cc
b/src/tirx/script/printer/for_loop.cc
index 48aa0e2357..df89d500da 100644
--- a/src/tirx/script/printer/for_loop.cc
+++ b/src/tirx/script/printer/for_loop.cc
@@ -40,13 +40,14 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
tirx::ForNode>(input);
TVM_FFI_CHECK(destination == nullptr, TypeError)
<< "printer statement-only node cannot fulfill a destination";
+ bool use_thread_binding = loop->GetThreadBinding().has_value() &&
!loop->step.has_value();
ffi::String method;
switch (loop->kind) {
- case tirx::ForKind::kSerial:
+ case tirx::ForKind::kDefault:
method = "serial";
break;
case tirx::ForKind::kParallel:
- method = "parallel";
+ method = use_thread_binding ? "thread_binding" : "parallel";
break;
case tirx::ForKind::kVectorized:
method = "vectorized";
@@ -54,9 +55,6 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
case tirx::ForKind::kUnrolled:
method = "unroll";
break;
- case tirx::ForKind::kThreadBinding:
- method = "thread_binding";
- break;
default:
TVM_FFI_THROW(TypeError) << "printer unknown loop kind";
}
@@ -72,14 +70,14 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
}
ffi::Array<ffi::String> keys;
ffi::Array<ExprDoc> values;
- if (loop->kind == tirx::ForKind::kThreadBinding) {
- TVM_FFI_CHECK(loop->thread_binding.has_value(), TypeError)
+ if (use_thread_binding) {
+ TVM_FFI_CHECK(loop->GetThreadBinding().has_value(), TypeError)
<< "printer thread-binding loop lacks thread tag";
keys.push_back("thread");
- values.push_back(LiteralDoc::Str(loop->thread_binding.value(),
std::nullopt));
+ values.push_back(LiteralDoc::Str(loop->GetThreadBinding().value(),
std::nullopt));
}
bool unroll_option = false;
- if (loop->kind == tirx::ForKind::kSerial && loop->annotations.size() == 1) {
+ if (loop->kind == tirx::ForKind::kDefault && loop->annotations.size() == 1) {
const auto& [key, value] = *loop->annotations.begin();
auto boolean = value.as<bool>();
auto integer = value.as<int64_t>();
@@ -94,9 +92,11 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
unroll_option = true;
}
}
- if (!unroll_option && !loop->annotations.empty()) {
+ auto annotations = loop->annotations;
+ if (use_thread_binding) annotations.erase("thread_binding");
+ if (!unroll_option && !annotations.empty()) {
std::vector<std::pair<ffi::String, ffi::Any>> sorted;
- for (const auto& [key, value] : loop->annotations)
sorted.emplace_back(key, value);
+ for (const auto& [key, value] : annotations) sorted.emplace_back(key,
value);
std::sort(sorted.begin(), sorted.end(),
[](const auto& a, const auto& b) { return a.first < b.first; });
ffi::Array<ExprDoc> annotation_keys;
@@ -109,8 +109,6 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
values.push_back(DictDoc(annotation_keys, annotation_values));
}
if (loop->step.has_value()) {
- TVM_FFI_CHECK(loop->kind != tirx::ForKind::kThreadBinding, TypeError)
- << "printer thread-binding loop cannot have a step";
keys.push_back("step");
values.push_back(d->Translate(loop->step.value()).value());
}
@@ -120,12 +118,11 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj*
d, ffi::AnyView input,
{
IdDoc var = VarDoc(d, loop->loop_var);
ffi::Array<ExprDoc> bounds = {min, end};
- if (loop->kind == tirx::ForKind::kThreadBinding && prim::IsZero(loop->min)
&&
- loop->min.ty() == loop->extent.ty()) {
+ if (use_thread_binding && prim::IsZero(loop->min) && loop->min.ty() ==
loop->extent.ty()) {
bounds = {end};
}
ExprDoc callee = NamespaceDoc("tirx")->Attr(method);
- if (loop->kind == tirx::ForKind::kSerial && loop->annotations.empty()) {
+ if (loop->kind == tirx::ForKind::kDefault && loop->annotations.empty()) {
callee = IdDoc("range");
// range's step is positional. Retain even an explicit unit step because
// it is part of the source For node, unlike an absent step.
diff --git a/src/tirx/transform/bind_target.cc
b/src/tirx/transform/bind_target.cc
index b36a32e8f1..278d9b303d 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<VisitInterrupt> Visit_(const ForNode* op) final {
- if (op->kind == ForKind::kThreadBinding) {
+ if (op->GetThreadBinding().has_value()) {
// Enter GPU scope for thread binding loops
bool last_is_under_gpu_scope = is_under_gpu_scope_;
is_under_gpu_scope_ = true;
@@ -188,7 +188,7 @@ class CallSubstitutor : public StmtExprMutator {
}
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
- if (op->kind == ForKind::kThreadBinding) {
+ if (op->GetThreadBinding().has_value()) {
// 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 1d3cc54d94..903e451483 100644
--- a/src/tirx/transform/lower_tirx_opaque.cc
+++ b/src/tirx/transform/lower_tirx_opaque.cc
@@ -71,7 +71,7 @@ class TIRxOpaqueLower : public StmtExprMutator {
.ValueOrUnchanged(op->step);
ffi::Any previous_remap = VarRemapGet(op->loop_var);
PrimVar launch_var(ffi::UnsafeInit{});
- if (op->kind == ForKind::kThreadBinding) {
+ if (op->GetThreadBinding().has_value()) {
TVM_FFI_ICHECK(IsZero(min)) << "Thread binding must have zero minimum";
launch_var = PrimVar(op->loop_var->name, extent.ty());
VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), launch_var));
@@ -101,24 +101,25 @@ class TIRxOpaqueLower : public StmtExprMutator {
VarRemapSet(op->loop_var, previous_remap);
// Step 2. Create the lowered loop or launch region.
- if (op->kind == ForKind::kThreadBinding) {
+ if (op->GetThreadBinding().has_value()) {
// Case 1. Thread binding → RegionStmt(launch_thread)
TVM_FFI_ICHECK(!op->annotations.count("loop_partition_hint") ||
op->annotations.at("loop_partition_hint") == nullptr)
<< "Run LoopPartition before opaque lowering of a thread-binding
loop with "
"loop_partition_hint";
- TVM_FFI_ICHECK(op->thread_binding.has_value());
+ TVM_FFI_ICHECK(op->GetThreadBinding().has_value());
} else if (IsOne(extent) && op->annotations.empty()) {
// Case 2. Unit loop elimination
return body;
} else {
// Case 3. An ordinary loop
body = For(op->loop_var, std::move(min), std::move(extent), op->kind,
std::move(body),
- std::nullopt, FilterAnnotations(annotations), step);
+ FilterAnnotations(annotations), step);
}
- if (op->kind == ForKind::kThreadBinding) {
- return RegionStmt(tirx::launch_thread_op(),
{StringImm(op->thread_binding.value()), extent},
- {launch_var}, DictAttrs(), body, {}, op->span);
+ if (op->GetThreadBinding().has_value()) {
+ return RegionStmt(tirx::launch_thread_op(),
+ {StringImm(op->GetThreadBinding().value()), extent},
{launch_var},
+ DictAttrs(), body, {}, op->span);
}
return body;
}
diff --git a/src/tirx/transform/lower_tvm_builtin.cc
b/src/tirx/transform/lower_tvm_builtin.cc
index a4197e7788..03ce57b729 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -374,7 +374,7 @@ class BuiltinLower : public StmtExprMutator {
PrimExpr extent = std::move(extent_result).ValueOrUnchanged(op->extent);
Stmt body = op->body;
- if (op->kind == ForKind::kParallel) {
+ if (op->kind == ForKind::kParallel && !op->GetThreadBinding().has_value())
{
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 ae433c2f89..d6acde6604 100644
--- a/src/tirx/transform/make_packed_api.cc
+++ b/src/tirx/transform/make_packed_api.cc
@@ -57,10 +57,12 @@ class ReturnRewriter : public StmtExprMutator {
explicit ReturnRewriter(Var ret_var) : ret_var_(ret_var) {}
UnchangedOr<Stmt> Mutate_(const ForNode* node, InplaceMode inplace_mode)
override {
- if (node->kind == ForKind::kParallel) in_parallel_ += 1;
+ if (node->kind == ForKind::kParallel &&
!node->GetThreadBinding().has_value())
+ in_parallel_ += 1;
Stmt ret =
StmtExprMutator::Mutate_(node,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(node));
- if (node->kind == ForKind::kParallel) in_parallel_ -= 1;
+ if (node->kind == ForKind::kParallel &&
!node->GetThreadBinding().has_value())
+ in_parallel_ -= 1;
return ret;
}
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index 7c0ad7a6d1..a4023ec519 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -688,7 +688,7 @@ class StoragePlanRewriter : public StmtExprMutator {
StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
op = stmt.as<ForNode>();
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);
}
@@ -1161,7 +1161,7 @@ class StoragePlanRewriter : public StmtExprMutator {
PlanNewScope(s.stmt);
} else if (s.stmt->IsInstance<ForNode>()) {
const auto* op = static_cast<const ForNode*>(s.stmt);
- if (op->kind == ForKind::kParallel) {
+ if (op->kind == ForKind::kParallel &&
!op->GetThreadBinding().has_value()) {
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 044b229ec6..0465e64aa3 100644
--- a/src/tirx/transform/unroll_loop.cc
+++ b/src/tirx/transform/unroll_loop.cc
@@ -148,9 +148,14 @@ class LoopUnroller : public StmtExprMutator {
node->annotations.erase(attr::unroll_explicit);
op = node;
}
+ if (op->GetThreadBinding().has_value()) {
+ // Explicit thread placement must survive until thread-binding lowering.
+ 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 &&
+ bool auto_unroll = (op->kind == ForKind::kDefault && value >= 0 &&
normal_loop_depth_ == 0 &&
unroll_depth_ <= auto_max_depth_);
const auto& policy = unroll_policy_.Current();
diff --git a/src/tirx/transform/vectorize_loop.cc
b/src/tirx/transform/vectorize_loop.cc
index 46ba771923..933c6a2b29 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -1166,7 +1166,7 @@ class Vectorizer : public StmtExprMutator {
auto substituter = ffi::make_object<StmtExprMutator>();
substituter->VarRemapSet(var_, idx);
stmt = substituter->Mutate(stmt).ValueOrUnchanged(stmt);
- return For(idx.as_or_throw<PrimVar>(), IntImm(var_ty, 0), var_lanes_,
ForKind::kSerial, stmt);
+ return For(idx.as_or_throw<PrimVar>(), IntImm(var_ty, 0), var_lanes_,
ForKind::kDefault, stmt);
}
private:
@@ -1371,9 +1371,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::kDefault, vector_loop, {},
std::nullopt, op->span);
return this->Mutate(loop, InplaceMode::kDisallow).ValueOrUnchanged(loop);
}
@@ -1393,7 +1393,7 @@ class VectorizeSkipper : public StmtExprMutator {
Stmt stmt = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
op = stmt.as<ForNode>();
if (op->kind == ForKind::kVectorized) {
- return For(op->loop_var, op->min, op->extent, ForKind::kSerial,
op->body);
+ return For(op->loop_var, op->min, op->extent, ForKind::kDefault,
op->body);
} else {
return stmt;
}
diff --git a/tests/python/ir/test_nonnullable_ir.py
b/tests/python/ir/test_nonnullable_ir.py
index e159c45697..435570b867 100644
--- a/tests/python/ir/test_nonnullable_ir.py
+++ b/tests/python/ir/test_nonnullable_ir.py
@@ -45,7 +45,7 @@ def test_optional_statement_fields_and_roundtrip():
body = tvm.tirx.Evaluate(0)
conditional = tvm.tirx.IfThenElse(True, body, None)
assert conditional.else_case is None
- loop = tvm.tirx.For(tvm.tirx.Var("i", "int32"), 0, 4,
tvm.tirx.ForKind.SERIAL, body, step=None)
+ loop = tvm.tirx.For(tvm.tirx.Var("i", "int32"), 0, 4,
tvm.tirx.ForKind.DEFAULT, body, step=None)
assert loop.step is None
declaration = tvm.tirx.Function([], None)
assert declaration.body is None
diff --git a/tests/python/relax/test_pipeline.py
b/tests/python/relax/test_pipeline.py
index 69eda6256b..a27caa23e7 100644
--- a/tests/python/relax/test_pipeline.py
+++ b/tests/python/relax/test_pipeline.py
@@ -174,7 +174,7 @@ def _has_thread_binding(func: tvm.tirx.Function) -> 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.annotations.get("thread_binding") is not None:
found = True
tvm_ffi.structural_walk(func.body, _visit)
diff --git a/tests/python/runtime/test_runtime_module_load.py
b/tests/python/runtime/test_runtime_module_load.py
index 6b3f85e839..4daabf11da 100644
--- a/tests/python/runtime/test_runtime_module_load.py
+++ b/tests/python/runtime/test_runtime_module_load.py
@@ -58,7 +58,7 @@ def test_dso_module_load():
i,
0,
n - 1,
- tvm.tirx.ForKind.SERIAL,
+ tvm.tirx.ForKind.DEFAULT,
tvm.tirx.TensorStore(Ab, tvm.tirx.TensorLoad(Ab, [i]) + 1, [i +
1]),
)
mod = tvm.IRModule.from_expr(
diff --git
a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
index 0d0d7daabf..0272df8a3f 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
@@ -142,7 +142,7 @@ def test_error_undeclared_buffer_in_schedulable_tir():
function = tvm.tirx.Function(
params=[A, B_data],
- body=tvm.tirx.For(i, 0, n, tvm.tirx.ForKind.SERIAL, block_realize),
+ body=tvm.tirx.For(i, 0, n, tvm.tirx.ForKind.DEFAULT, block_realize),
# Note: B is NOT a function parameter, so its declaration scope is only
# within a DeclTensor node (which we intentionally omit here).
)
diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py
b/tests/python/s_tir/analysis/test_sblock_access_region.py
index 5e7a9d6e77..9668845868 100644
--- a/tests/python/s_tir/analysis/test_sblock_access_region.py
+++ b/tests/python/s_tir/analysis/test_sblock_access_region.py
@@ -506,7 +506,7 @@ def test_conditional_inequality_access_regions(case):
]
)
for var, (minimum, extent) in reversed(list(zip(variables, domains))):
- body = tirx.For(var, minimum, extent, tirx.ForKind.SERIAL, body)
+ body = tirx.For(var, minimum, extent, tirx.ForKind.DEFAULT, body)
block = s_tir.SBlock([], [], [], "conditional", body)
# Unbounded access sets conservatively cover the whole buffer.
outside_expected = [(0, 256)] if case == "unbounded" else domains
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
index 4be08f067a..c13a481b89 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
@@ -63,7 +63,7 @@ def _make_loops(loop_vars: list[Var], extents: list[int]) ->
list[For]:
loop_var=loop_var,
min=0,
extent=extent,
- kind=ForKind.SERIAL,
+ kind=ForKind.DEFAULT,
body=Evaluate(0),
)
for loop_var, extent in zip(loop_vars, extents)
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 f74340afdb..9ed2c6ce04 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_state.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_state.py
@@ -338,9 +338,8 @@ def test_replace_block_in_opaque_block():
loop_var=for_loop.loop_var,
min=0,
extent=128,
- kind=tirx.ForKind.SERIAL,
+ kind=tirx.ForKind.DEFAULT,
body=tirx.Evaluate(0),
- thread_binding=None,
annotations=None,
)
s.replace(sref, new_for_loop)
diff --git a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
index 7d8c32346b..207dddf504 100644
--- a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
+++ b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
@@ -470,11 +470,11 @@ def test_for_thread_binding():
tvm.ir.assert_structural_equal(func, rt_func)
assert isinstance(rt_func.body[0], tirx.stmt.For)
- assert rt_func.body[0].kind == 4
- assert rt_func.body[0].thread_binding == "threadIdx.x"
+ assert rt_func.body[0].kind == tvm.tirx.ForKind.PARALLEL
+ assert rt_func.body[0].annotations.get("thread_binding") == "threadIdx.x"
assert isinstance(rt_func.body[0].body[0], tirx.stmt.For)
- assert rt_func.body[0].body[0].kind == 4
- assert rt_func.body[0].body[0].thread_binding == "threadIdx.y"
+ assert rt_func.body[0].body[0].kind == tvm.tirx.ForKind.PARALLEL
+ assert rt_func.body[0].body[0].annotations.get("thread_binding") ==
"threadIdx.y"
assert rt_func.body[0].body[0].annotations["attr_key"] == "attr_value"
diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
index bbd149088e..45b78d6746 100644
--- a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
+++ b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py
@@ -58,7 +58,7 @@ def test_error_for_out_of_scope_usage():
tvm.tirx.Var("j", "int32"),
0,
1,
- tvm.tirx.ForKind.SERIAL,
+ tvm.tirx.ForKind.DEFAULT,
tvm.tirx.SeqStmt([tvm.tirx.Bind(i, 42), tvm.tirx.Evaluate(i)]),
)
# Use i outside the For loop — this is out of scope
diff --git a/tests/python/tirx-base/test_tir_constructor.py
b/tests/python/tirx-base/test_tir_constructor.py
index b9be658e28..20d45c6056 100644
--- a/tests/python/tirx-base/test_tir_constructor.py
+++ b/tests/python/tirx-base/test_tir_constructor.py
@@ -306,7 +306,7 @@ def test_stmt_constructor():
assert len(x.message_parts) == 1
assert x.message_parts[0].value == "hellow"
- x = tvm.tirx.For(tvm.tirx.Var("x", "int32"), 0, 10,
tvm.tirx.ForKind.SERIAL, nop)
+ x = tvm.tirx.For(tvm.tirx.Var("x", "int32"), 0, 10,
tvm.tirx.ForKind.DEFAULT, nop)
assert isinstance(x, tvm.tirx.For)
assert x.min.value == 0
assert x.extent.value == 10
diff --git a/tests/python/tirx-base/test_tir_nodes.py
b/tests/python/tirx-base/test_tir_nodes.py
index 5c212e1ff1..0a49ae7ddf 100644
--- a/tests/python/tirx-base/test_tir_nodes.py
+++ b/tests/python/tirx-base/test_tir_nodes.py
@@ -129,7 +129,7 @@ def test_basic():
def test_stmt():
x = tvm.tirx.Evaluate(0)
- tvm.tirx.For(tvm.tirx.Var("i", "int32"), 0, 1, tvm.tirx.ForKind.SERIAL, x)
+ tvm.tirx.For(tvm.tirx.Var("i", "int32"), 0, 1, tvm.tirx.ForKind.DEFAULT, x)
tvm.tirx.For(tvm.tirx.Var("i", "int32"), 0, 1, tvm.tirx.ForKind.UNROLLED,
x, step=2)
diff --git
a/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
b/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
index 010fc533fe..e0b2895d82 100644
--- a/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
+++ b/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py
@@ -596,7 +596,7 @@ def test_let_body_no_extraction():
i,
tvm.tirx.const(0, "int32"),
tvm.tirx.const(10, "int32"),
- tvm.tirx.ForKind.SERIAL,
+ tvm.tirx.ForKind.DEFAULT,
store,
)
func = tvm.tirx.Function([buf, y], loop)
@@ -626,7 +626,7 @@ def test_let_value_cse():
i,
tvm.tirx.const(0, "int32"),
tvm.tirx.const(10, "int32"),
- tvm.tirx.ForKind.SERIAL,
+ tvm.tirx.ForKind.DEFAULT,
store,
)
func = tvm.tirx.Function([buf, y, z], loop)
@@ -659,7 +659,7 @@ def test_nested_let_no_extraction():
i,
tvm.tirx.const(0, "int32"),
tvm.tirx.const(10, "int32"),
- tvm.tirx.ForKind.SERIAL,
+ tvm.tirx.ForKind.DEFAULT,
store,
)
func = tvm.tirx.Function([buf, z], loop)
@@ -702,7 +702,7 @@ def test_let_floordiv_pattern():
i,
tvm.tirx.const(0, "int32"),
tvm.tirx.const(10, "int32"),
- tvm.tirx.ForKind.SERIAL,
+ tvm.tirx.ForKind.DEFAULT,
store,
)
func = tvm.tirx.Function([buf_a, buf_b, buf_c], loop)
diff --git a/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
b/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
index 37b2d64237..adbb983b61 100644
--- a/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
+++ b/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
@@ -510,7 +510,7 @@ def test_reused_loop_var_in_decl_buffer_elem_offset():
loop_var,
0,
128,
- tirx.ForKind.SERIAL,
+ tirx.ForKind.DEFAULT,
tirx.SeqStmt(
[
tirx.Bind(
diff --git a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
index 064d2c3999..2b451a768b 100644
--- a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
+++ b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
@@ -78,7 +78,7 @@ def check_value(expr, variables, data, fref):
loop_var,
tvm.tirx.const(0, "int32"),
tvm.tirx.const(n, "int32"),
- tvm.tirx.ForKind.SERIAL,
+ tvm.tirx.ForKind.DEFAULT,
make_store(loop_var),
)
diff --git a/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
b/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
index ccdf8df88a..6f3409ca85 100644
--- a/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
+++ b/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
@@ -38,17 +38,17 @@ def test_remove_no_op():
i,
0,
4,
- tvm.tirx.ForKind.SERIAL,
+ tvm.tirx.ForKind.DEFAULT,
tvm.tirx.For(
j,
0,
n,
- tvm.tirx.ForKind.SERIAL,
+ tvm.tirx.ForKind.DEFAULT,
tvm.tirx.For(
k,
0,
m,
- tvm.tirx.ForKind.SERIAL,
+ tvm.tirx.ForKind.DEFAULT,
tvm.tirx.IfThenElse(
(i * m + j + k < n), tvm.tirx.Evaluate(m),
tvm.tirx.Evaluate(n)
),
@@ -68,7 +68,7 @@ def test_remove_no_op():
assert len(ret) == 1 and ret[0] == store
# remove zero extent loop
- stmt3 = tvm.tirx.For(i, 0, 0, tvm.tirx.ForKind.SERIAL, store)
+ stmt3 = tvm.tirx.For(i, 0, 0, tvm.tirx.ForKind.DEFAULT, store)
mod = tvm.IRModule.from_expr(tvm.tirx.Function([Ab], stmt3))
ret = tvm.tirx.transform.RemoveNoOp()(mod)["main"].body
assert isinstance(ret, tvm.tirx.SeqStmt) and len(ret) == 0
diff --git a/tests/python/tirx/script/test_tirx_script_basic_usage.py
b/tests/python/tirx/script/test_tirx_script_basic_usage.py
index cd8c85a166..7ac1c50d7e 100644
--- a/tests/python/tirx/script/test_tirx_script_basic_usage.py
+++ b/tests/python/tirx/script/test_tirx_script_basic_usage.py
@@ -929,7 +929,7 @@ def test_thread_binding_has_no_dtype_parameter():
def test_hand_built_for_promotes_int_literal_bounds_to_uint32():
"""The For constructor retypes literal bounds to the loop var's dtype."""
loop_var = tvm.tirx.Var("i", "uint32")
- loop = tvm.tirx.For(loop_var, 0, 128, tvm.tirx.ForKind.SERIAL,
tvm.tirx.Evaluate(0))
+ loop = tvm.tirx.For(loop_var, 0, 128, tvm.tirx.ForKind.DEFAULT,
tvm.tirx.Evaluate(0))
assert loop.min.ty == PrimType("uint32")
assert loop.extent.ty == PrimType("uint32")
@@ -937,4 +937,4 @@ def
test_hand_built_for_promotes_int_literal_bounds_to_uint32():
def test_hand_built_for_rejects_negative_literal_for_uint32():
loop_var = tvm.tirx.Var("i", "uint32")
with pytest.raises(Exception, match="not representable"):
- tvm.tirx.For(loop_var, -1, 128, tvm.tirx.ForKind.SERIAL,
tvm.tirx.Evaluate(0))
+ tvm.tirx.For(loop_var, -1, 128, tvm.tirx.ForKind.DEFAULT,
tvm.tirx.Evaluate(0))
diff --git a/tests/python/tirx/script/test_tirx_script_ir_builder.py
b/tests/python/tirx/script/test_tirx_script_ir_builder.py
index 219fe8cc0c..ff8f189197 100644
--- a/tests/python/tirx/script/test_tirx_script_ir_builder.py
+++ b/tests/python/tirx/script/test_tirx_script_ir_builder.py
@@ -57,9 +57,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="threadIdx.x",
+ annotations={"thread_binding": "threadIdx.x"},
)
unroll_expected = tirx.For(
loop_var=tirx.Var("", "int32"),
@@ -86,7 +86,7 @@ def test_ir_builder_tir_for():
loop_var=tirx.Var("", "int32"),
min=0,
extent=128,
- kind=tirx.ForKind.SERIAL,
+ kind=tirx.ForKind.DEFAULT,
body=parallel_expected,
)
@@ -106,7 +106,7 @@ def test_ir_builder_tir_for_uint():
loop_var=tirx.Var("", "uint32"),
min=tirx.const(0, "uint32"),
extent=tirx.const(128, "uint32"),
- kind=tirx.ForKind.SERIAL,
+ kind=tirx.ForKind.DEFAULT,
body=tirx.Evaluate(0),
)
