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 366561c3bb [REFACTOR][TIRx] Simplify argument and unroll policy
plumbing (#20579)
366561c3bb is described below
commit 366561c3bbcfc4ab2240acf2dff2d976f51c0931
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 21:44:25 2026 -0400
[REFACTOR][TIRx] Simplify argument and unroll policy plumbing (#20579)
Share function argument registration through the ordinary Var path while
preserving the typed TensorVar overload and variable identity. Store
initialized LoopUnroller policy directly while retaining configuration,
inheritance and restoration.
Clarify the one-element, index-zero allreduce result convention, remove
a duplicate sibling-launch fixture, and update lexical-launch and
LowerThreadBinding terminology and small documentation remnants.
---
docs/reference/api/python/script/script.rst | 2 +-
include/tvm/tirx/builtin.h | 6 ++--
python/tvm/tirx/op.py | 3 +-
python/tvm/tirx/transform/common.py | 4 +--
src/tirx/script/ir_builder/ir.cc | 7 +----
src/tirx/transform/unroll_loop.cc | 11 ++++----
...> test_s_tir_transform_lower_thread_binding.py} | 32 +++++++++++-----------
.../test_tir_analysis_verify_well_formed.py | 21 ++------------
tests/python/tirx-base/test_tir_op_types.py | 3 +-
tests/python/tirx/test_compilation_pipeline.py | 2 +-
10 files changed, 36 insertions(+), 55 deletions(-)
diff --git a/docs/reference/api/python/script/script.rst
b/docs/reference/api/python/script/script.rst
index 8826052473..ff43fa04d1 100644
--- a/docs/reference/api/python/script/script.rst
+++ b/docs/reference/api/python/script/script.rst
@@ -39,7 +39,7 @@ tvm.script.tirx
***************
.. automodule:: tvm.script.tirx
:members:
- :exclude-members: Range, meta_var, Var, Call, CommReducer, Reduce,
SMEMPool, TMEMPool, FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv,
FloorMod, LShift, RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min,
Max, EQ, NE, LT, LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle,
CallEffectKind, IterVar, ComposeLayout, DtypeConstructor, ExecScope, Iter,
Layout, LetAnnotation, LocalVectorAnnotation, ScopeIdDef, TileLayout, Buffer,
buffer, prim_func, jit, inline, macro
+ :exclude-members: Range, meta_var, Var, Call, CommReducer, Reduce,
SMEMPool, TMEMPool, FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv,
FloorMod, LShift, RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min,
Max, EQ, NE, LT, LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle,
CallEffectKind, IterVar, ComposeLayout, DtypeConstructor, ExecScope, Iter,
Layout, LetAnnotation, LocalVectorAnnotation, ScopeIdDef, TileLayout,
prim_func, jit, inline, macro
.. autofunction:: tvm.script.tirx.prim_func
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index 688198c60c..d7a48f0592 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -519,8 +519,10 @@ TVM_DLL const Op& tvm_warp_activemask();
* and thread_axes may each be a scalar or an explicit Tuple of fields.
* Each value, identity, pair of parameters and result have
* the same primitive type. Inactive inputs are replaced by their identities.
- * Destinations are N tensor loads (optionally cast for boolean storage), and
- * thread_axes are reduction thread variables or zero for simplified unit axes.
+ * Destinations are N tensor loads at index zero of one-element result
temporaries
+ * (optionally cast for boolean storage). Each result temporary must be
accessed
+ * only at index zero. Thread axes are reduction thread variables or zero for
+ * simplified unit axes.
* Other thread indices remain fixed. The operation writes the reduced values
* to the destination tensors and returns void.
*/
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 9014497141..5d656c95ea 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -837,7 +837,8 @@ def tvm_thread_allreduce(combine, identity, values,
predicate, destinations, thr
predicate : PrimExpr
Boolean participation predicate. Inactive threads contribute
identities.
destinations : Expr or Sequence[Expr]
- Tensor loads identifying the destinations of the reduction results.
+ Tensor loads at index zero of one-element result temporaries,
optionally
+ cast for boolean storage. Each temporary must be accessed only at
index zero.
thread_axes : Expr or Sequence[Expr]
Thread variables participating in the reduction.
diff --git a/python/tvm/tirx/transform/common.py
b/python/tvm/tirx/transform/common.py
index cd3a6c79d1..111d6bac86 100644
--- a/python/tvm/tirx/transform/common.py
+++ b/python/tvm/tirx/transform/common.py
@@ -25,8 +25,8 @@ from tvm.tirx.layout import Iter, TileLayout
class BufferReplacer:
"""
- Replace buffer with another buffer.
- Var values are ordinary Vars, so the same mapping also rewrites
+ Replace tensor variables with other tensor variables.
+ Tensor variables are ordinary Vars, so the same mapping also rewrites
``buffer_data`` projections.
"""
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index 7d4d468df9..f4d614e113 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -112,9 +112,7 @@ Var Arg(ffi::String name, Var var) {
}
TensorVar Arg(ffi::String name, TensorVar buffer) {
- PrimFuncFrame frame = FindPrimFuncFrame("T.Arg");
- details::Namer::Name(buffer, name);
- frame->args.push_back(buffer.var());
+ Arg(std::move(name), buffer.var());
return buffer;
}
@@ -650,9 +648,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("script.ir_builder.tirx.Arg",
[](ffi::String name, ffi::ObjectRef obj) -> ffi::ObjectRef {
using namespace tvm::tirx;
- if (auto buffer = obj.as<TensorVar>()) {
- return Arg(name, buffer.value());
- }
if (auto var = obj.as<Var>()) {
return Arg(name, var.value());
}
diff --git a/src/tirx/transform/unroll_loop.cc
b/src/tirx/transform/unroll_loop.cc
index 7f77a383b2..8535def0aa 100644
--- a/src/tirx/transform/unroll_loop.cc
+++ b/src/tirx/transform/unroll_loop.cc
@@ -33,7 +33,6 @@
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
-#include <optional>
#include <unordered_set>
#include "../../runtime/thread_storage_scope.h"
@@ -155,8 +154,8 @@ class LoopUnroller : public StmtExprMutator {
unroll_depth_ <= auto_max_depth_);
const auto& policy = unroll_policy_.Current();
- auto_unroll = auto_unroll && (value * step_count_ <=
policy.auto_unroll_max_step.value() ||
- value <= auto_max_extent_);
+ auto_unroll = auto_unroll &&
+ (value * step_count_ <= policy.auto_unroll_max_step || value
<= auto_max_extent_);
if (op->kind == ForKind::kUnrolled) {
TVM_FFI_ICHECK_GE(value, 0) << "Cannot unroll non-constant loop";
@@ -176,7 +175,7 @@ class LoopUnroller : public StmtExprMutator {
normal_loop_depth_ += 1;
}
- if ((auto_unroll && policy.unroll_explicit.value()) ||
+ if ((auto_unroll && policy.unroll_explicit) ||
// unroll loops with extent = 1, no matter how many steps in body
(0 <= value && value <= auto_max_extent_ && auto_max_extent_ == 1)) {
return Unroll(op);
@@ -292,8 +291,8 @@ class LoopUnroller : public StmtExprMutator {
}
struct UnrollPolicy {
- std::optional<int> auto_unroll_max_step;
- std::optional<bool> unroll_explicit;
+ int auto_unroll_max_step{0};
+ bool unroll_explicit{true};
};
ScopeStack<UnrollPolicy> unroll_policy_;
int auto_max_depth_;
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_unify_thread_binding.py
b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py
similarity index 92%
rename from
tests/python/s_tir/transform/test_s_tir_transform_unify_thread_binding.py
rename to
tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py
index a2a9a85184..e1b8223f41 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_unify_thread_binding.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py
@@ -61,7 +61,7 @@ def element_wise_thread_x(
@Ts.prim_func
-def unified_element_wise_thread_x(
+def lowered_element_wise_thread_x(
A: T.Tensor([128, 128]), B: T.Tensor([128, 128]), C: T.Tensor([128, 128])
) -> None:
with T.launch_thread("blockIdx.x", 128) as i:
@@ -93,7 +93,7 @@ def element_wise_thread_x_different_dtype(
@Ts.prim_func
-def unified_element_wise_thread_x_different_dtype(
+def lowered_element_wise_thread_x_different_dtype(
A: T.Tensor((128, 128), "float32"),
B: T.Tensor((128, 128), "float32"),
C: T.Tensor((128, 128), "float32"),
@@ -110,7 +110,7 @@ def unified_element_wise_thread_x_different_dtype(
@Ts.prim_func
-def element_wise_env_thread_x(
+def element_wise_launch_thread_x(
A: T.Tensor([128, 128]), B: T.Tensor([128, 128]), C: T.Tensor([128, 128])
) -> None:
i = T.launch_thread("blockIdx.x", 128)
@@ -126,7 +126,7 @@ def element_wise_env_thread_x(
@Ts.prim_func
-def unified_element_wise_env_thread_x(
+def lowered_element_wise_launch_thread_x(
A: T.Tensor([128, 128]), B: T.Tensor([128, 128]), C: T.Tensor([128, 128])
) -> None:
i = T.launch_thread("blockIdx.x", 128)
@@ -152,7 +152,7 @@ def element_wise_vthread_x(A: T.Tensor([128, 128]), B:
T.Tensor([128, 128])) ->
@Ts.prim_func
-def unified_element_wise_vthread_x(A: T.Tensor([128, 128]), B: T.Tensor([128,
128])) -> None:
+def lowered_element_wise_vthread_x(A: T.Tensor([128, 128]), B: T.Tensor([128,
128])) -> None:
with T.launch_thread("vthread.x", 2) as i_0:
with T.launch_thread("threadIdx.x", 64) as i_1:
with T.launch_thread("vthread.x", 2) as j_0:
@@ -188,7 +188,7 @@ def element_wise_kernels_with_different_size(
@Ts.prim_func
-def unified_element_wise_kernels_with_different_size(
+def lowered_element_wise_kernels_with_different_size(
A: T.Tensor([128, 128]),
B: T.Tensor([128, 128]),
C: T.Tensor([256, 256]),
@@ -218,7 +218,7 @@ def element_wise_implicit_block(
@Ts.prim_func
-def unified_element_wise_implicit_block(
+def lowered_element_wise_implicit_block(
A: T.Tensor([128, 128]), B: T.Tensor([128, 128]), C: T.Tensor([128, 128])
) -> None:
with T.launch_thread("threadIdx.y", 128) as i:
@@ -233,19 +233,19 @@ def unified_element_wise_implicit_block(
def test_thread_x():
- _check(element_wise_thread_x, unified_element_wise_thread_x)
+ _check(element_wise_thread_x, lowered_element_wise_thread_x)
def test_thread_x_different_dtype():
- _check(element_wise_thread_x_different_dtype,
unified_element_wise_thread_x_different_dtype)
+ _check(element_wise_thread_x_different_dtype,
lowered_element_wise_thread_x_different_dtype)
-def test_env_thread_x():
- _check(element_wise_env_thread_x, unified_element_wise_env_thread_x)
+def test_launch_thread_x():
+ _check(element_wise_launch_thread_x, lowered_element_wise_launch_thread_x)
def test_vthread_x():
- _check(element_wise_vthread_x, unified_element_wise_vthread_x)
+ _check(element_wise_vthread_x, lowered_element_wise_vthread_x)
def test_two_thread_x_in_same_kernel_not_equal():
@@ -254,12 +254,12 @@ def test_two_thread_x_in_same_kernel_not_equal():
def test_kernels_with_different_size():
_check(
- element_wise_kernels_with_different_size,
unified_element_wise_kernels_with_different_size
+ element_wise_kernels_with_different_size,
lowered_element_wise_kernels_with_different_size
)
def test_implicit_block():
- _check(element_wise_implicit_block, unified_element_wise_implicit_block)
+ _check(element_wise_implicit_block, lowered_element_wise_implicit_block)
def test_inner_binding_with_annotation():
@@ -272,7 +272,7 @@ def test_inner_binding_with_annotation():
B[v] = A[v]
@Ts.prim_func
- def unified_inner_binding_with_annotation(
+ def lowered_inner_binding_with_annotation(
A: T.Tensor((64,), "float32"), B: T.Tensor((64,), "float32")
):
with T.launch_thread("blockIdx.x", 32) as blockIdx_x:
@@ -284,7 +284,7 @@ def test_inner_binding_with_annotation():
Ts.writes(B[v])
B[v] = A[v]
- _check(inner_binding_with_annotation,
unified_inner_binding_with_annotation)
+ _check(inner_binding_with_annotation,
lowered_inner_binding_with_annotation)
if __name__ == "__main__":
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 893afc5787..af80aec89b 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
@@ -136,7 +136,7 @@ def test_error_for_cross_function_reuse():
tvm.tirx.analysis.verify_well_formed(mod)
-def test_reuse_of_env_thread_in_function_is_well_formed():
+def test_sibling_launch_bindings_are_well_formed():
"""Sibling launch regions each introduce a fresh lexical binding."""
@T.prim_func
@@ -150,24 +150,7 @@ def test_reuse_of_env_thread_in_function_is_well_formed():
tvm.tirx.analysis.verify_well_formed(func)
-def test_reuse_of_env_thread_in_function_is_mandatory():
- """Separate launch bindings for the same hardware axis are well formed.
-
- The former shared-identity requirement is retired; tag equality does not
- make sibling region parameters the same definition."""
-
- @T.prim_func
- def func(A: T.Tensor([256], "float32")):
- with T.launch_thread("threadIdx.x", 256) as threadIdx_x:
- A[threadIdx_x] = A[threadIdx_x] + 1.0
-
- with T.launch_thread("threadIdx.x", 256) as threadIdx_x:
- A[threadIdx_x] = A[threadIdx_x] + 2.0
-
- tvm.tirx.analysis.verify_well_formed(func)
-
-
-def test_reuse_of_env_thread_across_functions_is_ill_formed():
+def test_reuse_of_region_parameter_across_functions_is_ill_formed():
"""An explicit region parameter may not be defined across multiple
functions."""
threadIdx_x = T.dynamic("threadIdx_x", "int32")
diff --git a/tests/python/tirx-base/test_tir_op_types.py
b/tests/python/tirx-base/test_tir_op_types.py
index 83a3944b81..550bf47384 100644
--- a/tests/python/tirx-base/test_tir_op_types.py
+++ b/tests/python/tirx-base/test_tir_op_types.py
@@ -94,10 +94,11 @@ def test_tir_op_call_likely():
def test_tir_op_tvm_thread_allreduce():
tensor = tirx.decl_tensor((128,), "float32")
+ result = tirx.decl_tensor((1,), "float32")
axis = tirx.Var("thread", "int32")
combine = tvm.ir.LambdaExpr(["float32", "float32"], lambda lhs, rhs: (lhs
+ rhs,))
expr = tirx.tvm_thread_allreduce(
- combine, (tirx.const(0, "float32"),), (tensor[0],), True,
(tensor[1],), (axis,)
+ combine, (tirx.const(0, "float32"),), (tensor[0],), True,
(result[0],), (axis,)
)
assert expr.op.name == "tirx.tvm_thread_allreduce"
diff --git a/tests/python/tirx/test_compilation_pipeline.py
b/tests/python/tirx/test_compilation_pipeline.py
index cac37ad982..32bffecfc3 100644
--- a/tests/python/tirx/test_compilation_pipeline.py
+++ b/tests/python/tirx/test_compilation_pipeline.py
@@ -44,7 +44,7 @@ def test_default_pipeline_allocations(pipeline):
np.testing.assert_equal(b.numpy(), source + 1)
-def test_unify_thread_binding():
+def test_lower_thread_binding():
@T.prim_func
def before(A: T.Tensor((32,), "int32")):
for bx in T.thread_binding(1, thread="blockIdx.x"):
