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 cd29ad03c1 [REFACTOR][TIRx] Consolidate bindings and emitted
allocation results (#20799)
cd29ad03c1 is described below
commit cd29ad03c1b6cac6852482c5fe2c8acf821ec395
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 19 17:12:37 2026 -0400
[REFACTOR][TIRx] Consolidate bindings and emitted allocation results
(#20799)
---
python/tvm/topi/gpu/scan.py | 14 +-
python/tvm/topi/gpu/sort.py | 36 ++---
python/tvm/topi/searchsorted.py | 4 +-
python/tvm/topi/sparse_reshape.py | 16 ++-
python/tvm/topi/vision/nms.py | 18 +--
python/tvm/topi/vision/nms_util.py | 4 +-
.../transform/test_s_tir_transform_thread_sync.py | 16 +--
tirx/include/tvm/script/ir_builder/ir.h | 15 --
.../tvm_tirx/backend/cuda/lang/alloc_pool.py | 4 +-
tirx/python/tvm_tirx/backend/cuda/lang/pipeline.py | 15 +-
.../tvm_tirx/backend/cuda/lang/tile_scheduler.py | 28 ++--
.../backend/cuda/tile/patterns/async_copy.py | 4 +-
.../tvm_tirx/backend/cuda/tile/patterns/tma.py | 2 +-
tirx/python/tvm_tirx/ir/__init__.py | 1 -
tirx/python/tvm_tirx/ir/stmt.py | 18 ---
tirx/python/tvm_tirx/script/ir_builder/__init__.py | 2 +
tirx/python/tvm_tirx/script/ir_builder/base.py | 3 +-
tirx/python/tvm_tirx/script/ir_builder/stmt.py | 52 ++-----
tirx/python/tvm_tirx/testing/bench.py | 10 +-
tirx/python/tvm_tirx/tirx/script/__init__.py | 1 -
tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py | 55 +++++---
.../tirx/script/ir_builder/parser_protocol.py | 91 ++++++------
tirx/src/script/ir_builder/ir.cc | 16 ---
.../tirx/analysis/test_check_well_defined.py | 24 ++--
tirx/tests/python/tirx/script/test_basic_usage.py | 15 +-
tirx/tests/python/tirx/script/test_ir_builder.py | 157 +++++++++++++++++++--
.../python/tirx/script/test_meta_programming.py | 6 +-
tirx/tests/python/tirx/script/test_printer.py | 39 ++---
.../python/tirx/script/test_printer_function.py | 4 +-
.../python/tirx/script/test_source_locations.py | 5 +-
tirx/tests/python/tirx/script/test_tensors.py | 10 +-
tirx/tests/python/tirx/script/test_tracing.py | 2 +-
tirx/tests/python/tirx/test_tile_namespace.py | 34 +++--
tirx/tests/python/tirx/tile/cuda/test_ewise.py | 2 +-
.../python/tirx/tile/cuda/test_tcgen05_mma.py | 2 +-
.../tirx/transform/test_common_subexpr_elim.py | 82 +++++------
.../python/tirx/transform/test_convert_ssa.py | 18 +--
.../python/tirx/transform/test_function_pass.py | 2 +-
.../tirx/transform/test_lower_tensor_layouts.py | 14 +-
39 files changed, 464 insertions(+), 377 deletions(-)
diff --git a/python/tvm/topi/gpu/scan.py b/python/tvm/topi/gpu/scan.py
index bd5bbc613b..f806e718e7 100644
--- a/python/tvm/topi/gpu/scan.py
+++ b/python/tvm/topi/gpu/scan.py
@@ -133,13 +133,13 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=operator.add, identity
width = 2 << l2_width
blocks_per_batch = cast(ceil_div(scan_axis_size,
max_threads * width), "int32")
- start_buf = T.bind(
+ start_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], "int32",
scope="local")])
)
- middle_buf = T.bind(
+ middle_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], "int32",
scope="local")])
)
- end_buf = T.bind(
+ end_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], "int32",
scope="local")])
)
with T.frame_scope(
@@ -227,16 +227,16 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=operator.add, identity
width = 2 << (lim - l2_width - 1)
blocks_per_batch = cast(ceil_div(scan_axis_size,
max_threads * width), "int32")
- start_buf = T.bind(
+ start_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], "int32",
scope="local")])
)
- middle_buf = T.bind(
+ middle_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], "int32",
scope="local")])
)
- end_buf = T.bind(
+ end_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], "int32",
scope="local")])
)
- tmp_buf = T.bind(
+ tmp_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], out_dtype,
scope="local")])
)
with T.frame_scope(
diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py
index 4c3df93f16..6d3ab862f3 100644
--- a/python/tvm/topi/gpu/sort.py
+++ b/python/tvm/topi/gpu/sort.py
@@ -112,17 +112,21 @@ def _odd_even_sort(
start = bx * block_size
# Tensor declarations without data allocate their storage.
- tmp_keys_swap = T.bind(
+ tmp_keys_swap = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([block_size], keys_swap.dtype,
scope="shared")])
)
- temp_keys = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1],
keys_swap.dtype, scope="local")]))
- temp_cond1 = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1],
keys_swap.dtype, scope="local")]))
- temp_cond2 = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1],
keys_swap.dtype, scope="local")]))
+ temp_keys = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1],
keys_swap.dtype, scope="local")]))
+ temp_cond1 = T.bind_(
+ T.alloc_tensor(ty_args=[T.Tensor([1], keys_swap.dtype,
scope="local")])
+ )
+ temp_cond2 = T.bind_(
+ T.alloc_tensor(ty_args=[T.Tensor([1], keys_swap.dtype,
scope="local")])
+ )
if values_swap is not None:
- tmp_values_swap = T.bind(
+ tmp_values_swap = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([block_size],
values_swap.dtype, scope="shared")])
)
- temp_values = T.bind(
+ temp_values = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], values_swap.dtype,
scope="local")])
)
@@ -416,10 +420,10 @@ def _sort_common(
step_count,
even,
):
- first_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")]))
- last_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")]))
- i_buf_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")]))
- j_buf_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")]))
+ first_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1],
target_dtype, scope="local")]))
+ last_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")]))
+ i_buf_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1],
target_dtype, scope="local")]))
+ j_buf_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1],
target_dtype, scope="local")]))
first = first_buf
last = last_buf
i_buf = i_buf_buf
@@ -483,16 +487,16 @@ def _sort_common(
step_count,
even,
):
- outer_first_buf = T.bind(
+ outer_first_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")])
)
- outer_last_buf = T.bind(
+ outer_last_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")])
)
- first_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")]))
- last_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")]))
- i_buf_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")]))
- j_buf_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")]))
+ first_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1],
target_dtype, scope="local")]))
+ last_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1], target_dtype,
scope="local")]))
+ i_buf_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1],
target_dtype, scope="local")]))
+ j_buf_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1],
target_dtype, scope="local")]))
outer_first = outer_first_buf
outer_last = outer_last_buf
first = first_buf
diff --git a/python/tvm/topi/searchsorted.py b/python/tvm/topi/searchsorted.py
index f0eda63934..39512653e9 100644
--- a/python/tvm/topi/searchsorted.py
+++ b/python/tvm/topi/searchsorted.py
@@ -38,8 +38,8 @@ def binary_search(sequence_offset, search_range,
sorted_sequence, value, right,
Note that we index N-D Buffer by 1-D linearlized indices.
"""
- lo_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], out_dtype,
scope="local")]))
- hi_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], out_dtype,
scope="local")]))
+ lo_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1], out_dtype,
scope="local")]))
+ hi_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1], out_dtype,
scope="local")]))
lo = lo_buf
hi = hi_buf
diff --git a/python/tvm/topi/sparse_reshape.py
b/python/tvm/topi/sparse_reshape.py
index 5542af272b..49a1d8c9dc 100644
--- a/python/tvm/topi/sparse_reshape.py
+++ b/python/tvm/topi/sparse_reshape.py
@@ -89,19 +89,19 @@ def sparse_reshape(
prev_shape_size = prev_shape_ptr.shape[0]
new_shape_size = new_shape_ptr.shape[0]
- multipliers_buf = T.bind(
+ multipliers_buf = T.bind_(
T.alloc_tensor(
ty_args=[T.Tensor([prev_shape_size], new_shape_ptr.dtype,
scope="local")]
)
)
multipliers = multipliers_buf
- dividers_buf = T.bind(
+ dividers_buf = T.bind_(
T.alloc_tensor(
ty_args=[T.Tensor([new_shape_size], new_shape_ptr.dtype,
scope="local")]
)
)
dividers = dividers_buf
- flattened_indices_buf = T.bind(
+ flattened_indices_buf = T.bind_(
T.alloc_tensor(
ty_args=[
T.Tensor([sparse_indices_ptr.shape[0]],
new_shape_ptr.dtype, scope="local")
@@ -109,15 +109,17 @@ def sparse_reshape(
)
)
flattened_indices = flattened_indices_buf
- total_ele_buf = T.bind(
+ total_ele_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], new_shape_ptr.dtype,
scope="local")])
)
total_ele = total_ele_buf
- division_total_ele_buf = T.bind(
+ division_total_ele_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor([1], new_shape_ptr.dtype,
scope="local")])
)
division_total_ele = division_total_ele_buf
- equal_shape_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1],
"bool", scope="local")]))
+ equal_shape_buf = T.bind_(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "bool", scope="local")])
+ )
equal_shape = equal_shape_buf
T.tensor_store(
@@ -243,7 +245,7 @@ def sparse_reshape(
)
with T.parallel(0, new_sparse_indices_ptr.shape[0]) as i:
- current_element_buf = T.bind(
+ current_element_buf = T.bind_(
T.alloc_tensor(
ty_args=[T.Tensor([1], new_shape_ptr.dtype,
scope="local")]
)
diff --git a/python/tvm/topi/vision/nms.py b/python/tvm/topi/vision/nms.py
index 6d0dc136cc..d72d27413f 100644
--- a/python/tvm/topi/vision/nms.py
+++ b/python/tvm/topi/vision/nms.py
@@ -227,7 +227,7 @@ def _classic_nms_ir(
with T.parallel(0, batch_size) as i:
# Step 1: Reorder data by sorted score
- nkeep_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor((1,), "int32",
scope="local")]))
+ nkeep_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor((1,),
"int32", scope="local")]))
nkeep_local = nkeep_buf
T.tensor_store(
nkeep_local,
@@ -266,20 +266,20 @@ def _classic_nms_ir(
T.tensor_store(out_box_indices, (i, j), T.int32(-1))
# Step 2: Apply NMS - greedy suppression
- num_valid_boxes_buf = T.bind(
+ num_valid_boxes_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor((1,), "int32",
scope="local")])
)
num_valid_boxes = num_valid_boxes_buf
T.tensor_store(num_valid_boxes, T.tensor_indices(num_valid_boxes,
0), T.int32(0))
- best_idx_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor((1,),
"int32", scope="local")]))
+ best_idx_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor((1,),
"int32", scope="local")]))
best_idx = best_idx_buf
- best_score_buf = T.bind(
+ best_score_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor((1,), data.dtype,
scope="local")])
)
best_score = best_score_buf
- tmp_idx_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor((1,),
"int32", scope="local")]))
+ tmp_idx_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor((1,),
"int32", scope="local")]))
tmp_idx = tmp_idx_buf
- tmp_val_buf = T.bind(
+ tmp_val_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor((1,), data.dtype,
scope="local")])
)
tmp_val = tmp_val_buf
@@ -635,7 +635,7 @@ def _classic_nms_ir(
)
if return_indices:
- valid_idx_buf = T.bind(
+ valid_idx_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor((1,), "int32",
scope="local")])
)
valid_idx = valid_idx_buf
@@ -883,7 +883,7 @@ def _rearrange_out(data, batch_size, num_anchors,
box_data_length, score_index):
out = outs[0]
with T.parallel(0, batch_size) as i:
- valid_idx_buf = T.bind(
+ valid_idx_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor((1,), "int32",
scope="local")])
)
valid_idx = valid_idx_buf
@@ -982,7 +982,7 @@ def _nms_loop(
with T.if_(tvm_tirx.tirx.all(iou_threshold > te.const(0),
valid_count[i] > te.const(0))):
with T.then_():
- num_valid_boxes_local_buf = T.bind(
+ num_valid_boxes_local_buf = T.bind_(
T.alloc_tensor(ty_args=[T.Tensor((1,), "int32",
scope="local")])
)
num_valid_boxes_local = num_valid_boxes_local_buf
diff --git a/python/tvm/topi/vision/nms_util.py
b/python/tvm/topi/vision/nms_util.py
index 26fc0b32f0..322936286a 100644
--- a/python/tvm/topi/vision/nms_util.py
+++ b/python/tvm/topi/vision/nms_util.py
@@ -68,8 +68,8 @@ def binary_search(y, num_boxes, scores, score_threshold, out):
Must be called within an IRBuilder context.
"""
- lo_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], "int32",
scope="local")]))
- hi_buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([1], "int32",
scope="local")]))
+ lo_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1], "int32",
scope="local")]))
+ hi_buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([1], "int32",
scope="local")]))
lo = lo_buf
hi = hi_buf
T.tensor_store(lo, T.tensor_indices(lo, 0), T.int32(0))
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
index 5a93835b94..a2e7a2e96b 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
@@ -191,13 +191,13 @@ def test_sync_bind():
in_thread_A_temp.data, ty_args=[T.Tensor((1,), "float32",
scope="local")]
)
in_thread_A_temp_1[0] = T.float32(0)
- A_temp_1 = T.bind(in_thread_A_temp_1[0] + A_shared_1[threadIdx_x])
+ A_temp_1 = in_thread_A_temp_1[0] + A_shared_1[threadIdx_x]
in_thread_A_temp_1[0] = A_temp_1
- A_temp_2 = T.bind(in_thread_A_temp_1[0] + A_shared_1[threadIdx_x +
128])
+ A_temp_2 = in_thread_A_temp_1[0] + A_shared_1[threadIdx_x + 128]
in_thread_A_temp_1[0] = A_temp_2
- A_temp_3 = T.bind(in_thread_A_temp_1[0] + A_shared_1[threadIdx_x +
256])
+ A_temp_3 = in_thread_A_temp_1[0] + A_shared_1[threadIdx_x + 256]
in_thread_A_temp_1[0] = A_temp_3
- A_temp_4 = T.bind(in_thread_A_temp_1[0] + A_shared_1[threadIdx_x +
384])
+ A_temp_4 = in_thread_A_temp_1[0] + A_shared_1[threadIdx_x + 384]
in_thread_A_temp_1[0] = A_temp_4
cross_thread_A_temp_1 = T.decl_tensor(
cross_thread_A_temp.data, ty_args=[T.Tensor((1,), "float32",
scope="local")]
@@ -228,13 +228,13 @@ def test_sync_bind():
)
in_thread_A_temp_1_1[0] = T.float32(0)
T.evaluate(T.call_intrin("tirx.gpu_storage_sync", "shared",
ty="int32"))
- A_temp_1 = T.bind(in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x])
+ A_temp_1 = in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x]
in_thread_A_temp_1_1[0] = A_temp_1
- A_temp_2 = T.bind(in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x +
128])
+ A_temp_2 = in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x + 128]
in_thread_A_temp_1_1[0] = A_temp_2
- A_temp_3 = T.bind(in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x +
256])
+ A_temp_3 = in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x + 256]
in_thread_A_temp_1_1[0] = A_temp_3
- A_temp_4 = T.bind(in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x +
384])
+ A_temp_4 = in_thread_A_temp_1_1[0] + A_shared_1_1[threadIdx_x + 384]
in_thread_A_temp_1_1[0] = A_temp_4
cross_thread_A_temp_1_1 = T.decl_tensor(
cross_thread_A_temp_1.data, ty_args=[T.Tensor((1,), "float32",
scope="local")]
diff --git a/tirx/include/tvm/script/ir_builder/ir.h
b/tirx/include/tvm/script/ir_builder/ir.h
index 271914c8c3..dec738f722 100644
--- a/tirx/include/tvm/script/ir_builder/ir.h
+++ b/tirx/include/tvm/script/ir_builder/ir.h
@@ -95,21 +95,6 @@ ForFrame Grid(ffi::Array<ffi::Variant<PrimExpr,
ffi::Tuple<PrimExpr, PrimExpr>>>
AssertFrame Assert(PrimExpr condition, ffi::String error_kind,
ffi::Array<ffi::String> message_parts);
-/*!
- * \brief Create a Bind (variable binding).
- *
- * Emits a flat Bind statement to the current frame and returns the bound
variable.
- *
- * \param value The value to be bound.
- * \param type_annotation The type annotation of the binding.
- * Usually it is used for fine-grained var typing,
- * particularly, PtrType.
- * \param var The variable to be bound. If not specified, a new variable will
be created.
- * \return The bound Var.
- */
-Var Bind(Expr value, ffi::Optional<Type> type_annotation = std::nullopt,
- ffi::Optional<Var> var = std::nullopt);
-
/*!
* \brief Create a while loop.
* \param condition The termination condition of the loop.
diff --git a/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py
b/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py
index 3f790f6951..dd94dc9f75 100644
--- a/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py
+++ b/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py
@@ -185,9 +185,7 @@ class SMEMPool:
def __init__(self, ptr=_POOL_UNSET):
ir = _get_ir()
if ptr is _POOL_UNSET:
- self.buf = ir._bind_tensor(
- ir.alloc_tensor(ty_args=[ir.Tensor([0], "uint8",
scope="shared.dyn")])
- )
+ self.buf = ir.alloc_tensor(ty_args=[ir.Tensor([0], "uint8",
scope="shared.dyn")]).value
self.ptr = self.buf.data
self._owns_buffer = True
else:
diff --git a/tirx/python/tvm_tirx/backend/cuda/lang/pipeline.py
b/tirx/python/tvm_tirx/backend/cuda/lang/pipeline.py
index 460b5e1381..95242e1784 100644
--- a/tirx/python/tvm_tirx/backend/cuda/lang/pipeline.py
+++ b/tirx/python/tvm_tirx/backend/cuda/lang/pipeline.py
@@ -60,8 +60,8 @@ class PipelineState:
"""
def __init__(self, depth: int, phase=None):
- self.stage = T.local_scalar("int32")
- self.phase = T.local_scalar("int32")
+ self.stage = T.bind_(T.local_scalar("int32"))
+ self.phase = T.bind_(T.local_scalar("int32"))
self.depth = depth
if phase is not None:
self.init(phase)
@@ -93,7 +93,7 @@ def _map_addr_into_cta(ptr, rank):
the enclosing frame explicitly or it is discarded, and the scratch it
writes into has to be declared in this scope to bind into the Function.
"""
- mapped = T.alloc_local([1], "uint32")
+ mapped = T.bind_(T.alloc_local([1], "uint32"))
T.evaluate(
T.ptx.mapa.shared__cluster.u32(
mapped[0], T.cuda.cvta_generic_to_shared(ptr), T.uint32(rank)
@@ -113,11 +113,14 @@ def _map_buffer_into_cta(ptr, rank, depth):
from tvm_tirx.tirx import Var as TIRVar # pylint:
disable=import-outside-toplevel
ptr_ty = PtrType(PrimType("uint64"), "shared")
- mapped = T.alloc_local([1], "uint64")
+ mapped = T.bind_(T.alloc_local([1], "uint64"))
T.evaluate(T.ptx.mapa.u64(mapped[0], ptr, T.uint32(rank)))
remote_ptr = TIRVar("remote_mbar_ptr", ptr_ty)
- T.bind(T.reinterpret(mapped[0], ty=ptr_ty), var=remote_ptr)
- return T.bind(T.decl_tensor(remote_ptr, ty_args=[T.Tensor([depth],
"uint64", scope="shared")]))
+ T.emit_(T.Bind(remote_ptr, T.reinterpret(mapped[0], ty=ptr_ty)))
+ value = T.decl_tensor(remote_ptr, ty_args=[T.Tensor([depth], "uint64",
scope="shared")])
+ remote_barrier = TIRVar("remote_mbar", value.ty)
+ T.emit_(T.Bind(remote_barrier, value))
+ return remote_barrier
def _mbarrier_arrive_remote(bar, pred=None, count=None):
diff --git a/tirx/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py
b/tirx/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py
index 24eba47b80..a99b6c41e7 100644
--- a/tirx/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py
+++ b/tirx/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py
@@ -30,9 +30,9 @@ class BaseTileScheduler:
"""Base class for tile schedulers with common state and macros."""
def __init__(self, prefix: str):
- self.m_idx = T.local_scalar("int32")
- self.n_idx = T.local_scalar("int32")
- self.linear_idx = T.local_scalar("int32")
+ self.m_idx = T.bind_(T.local_scalar("int32"))
+ self.n_idx = T.bind_(T.local_scalar("int32"))
+ self.linear_idx = T.bind_(T.local_scalar("int32"))
@T.inline
def update_current_m_n_idx(self, linear_idx):
@@ -231,7 +231,7 @@ class ClusterPersistentScheduler2D(BaseTileScheduler):
# Rename internal state for clarity
self.work_idx = self.linear_idx # alias: global work item index
- self.tile_count = T.local_scalar("int32")
+ self.tile_count = T.bind_(T.local_scalar("int32"))
is_static_m = isinstance(num_m_tiles, int)
@@ -509,9 +509,9 @@ class FlashAttentionLinearScheduler(BaseTileScheduler):
self._total_tasks = num_batches * num_heads * num_m_blocks
# Output indices
- self.batch_idx = T.local_scalar("int32")
- self.head_idx = T.local_scalar("int32")
- self.m_block_idx = T.local_scalar("int32")
+ self.batch_idx = T.bind_(T.local_scalar("int32"))
+ self.head_idx = T.bind_(T.local_scalar("int32"))
+ self.m_block_idx = T.bind_(T.local_scalar("int32"))
# fmt: off
@T.inline
@@ -596,9 +596,9 @@ class FlashAttentionLPTScheduler(BaseTileScheduler):
self._num_hb_quotient = self._num_hb // l2_swizzle
# Output indices
- self.batch_idx = T.local_scalar("int32")
- self.head_idx = T.local_scalar("int32")
- self.m_block_idx = T.local_scalar("int32")
+ self.batch_idx = T.bind_(T.local_scalar("int32"))
+ self.head_idx = T.bind_(T.local_scalar("int32"))
+ self.m_block_idx = T.bind_(T.local_scalar("int32"))
# fmt: off
@T.inline
@@ -670,8 +670,8 @@ class _CLCWorker(ClusterPersistentScheduler2D):
)
self._clc = clc
self._sa = PipelineState(1, 0)
- self._done = T.local_scalar("int32")
- self._nxt = T.local_scalar("uint32")
+ self._done = T.bind_(T.local_scalar("int32"))
+ self._nxt = T.bind_(T.local_scalar("uint32"))
@T.inline
def reset(self):
@@ -735,8 +735,8 @@ class ClusterLaunchControlScheduler:
self.sched_arr = Pipeline(pool, 1, full="tma", empty="mbar",
init_empty=1)
self.sched_fin = Pipeline(pool, 1, full="mbar", empty="mbar",
init_empty=finish_arrivals)
self.clc_handle = pool.alloc((4,), "uint32", align=16)
- self._s_done = T.local_scalar("int32")
- self._s_nxt = T.local_scalar("uint32")
+ self._s_done = T.bind_(T.local_scalar("int32"))
+ self._s_nxt = T.bind_(T.local_scalar("uint32"))
def worker(self, prefix):
return _CLCWorker(self, prefix)
diff --git a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/async_copy.py
b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/async_copy.py
index 4ae853aed5..3644da6565 100644
--- a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/async_copy.py
+++ b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/async_copy.py
@@ -31,7 +31,7 @@ before mapa. Layout strides are elements; chunk_bytes and
src_size are bytes.
import tvm_ffi
from tvm_tirx.backend.cuda.ptx import PTXNamespace
-from tvm_tirx.ir import Bind, Evaluate, Var
+from tvm_tirx.ir import Bind, Call, Evaluate, Var
from tvm_tirx.script import tirx as T
from .common import Context, scalar
@@ -66,7 +66,7 @@ def lower(opcode, entry, modifiers, args, params, sctx):
if bulk:
barrier, remote = params.pop("mbar"), params.pop("remote_cta_id")
mapped = Var("remote_addresses", T.Tensor((2,), "uint64",
scope="local", layout="default"))
- ctx.bindings.append(Bind(mapped, T.alloc_tensor(ty_args=[mapped.ty])))
+ ctx.bindings.append(Bind(mapped, Call("tirx.alloc_tensor", [],
ty_args=[mapped.ty])))
ctx.bindings.append(
Evaluate(T.ptx.mapa.u64(mapped[0], barrier, T.cast(remote,
ty="uint32")))
)
diff --git a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
index 272dd1209e..a85eeec50b 100644
--- a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
+++ b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
@@ -202,7 +202,7 @@ def _tensor_map(ctx, region, rank, box, encoding):
@T.function(check_well_formed=False)
def encode():
- T.bind(T.stack_alloca("tensormap", 1), var=tensor_map)
+ T.Bind(tensor_map, T.stack_alloca("tensormap", 1))
T.cuda.tensormap_encode_tiled(
tensor_map,
tensor.ptr_to(starts),
diff --git a/tirx/python/tvm_tirx/ir/__init__.py
b/tirx/python/tvm_tirx/ir/__init__.py
index d10809eab8..8a4d3ea4dc 100644
--- a/tirx/python/tvm_tirx/ir/__init__.py
+++ b/tirx/python/tvm_tirx/ir/__init__.py
@@ -85,7 +85,6 @@ from .stmt import (
RegionStmt,
TensorStore,
stmt_seq,
- stmt_list,
)
# Load passes after the common IR classes used by transformations are
registered.
diff --git a/tirx/python/tvm_tirx/ir/stmt.py b/tirx/python/tvm_tirx/ir/stmt.py
index b7bdddd877..d3cd1cba81 100644
--- a/tirx/python/tvm_tirx/ir/stmt.py
+++ b/tirx/python/tvm_tirx/ir/stmt.py
@@ -472,21 +472,3 @@ def stmt_seq(*args: Expr | Stmt) -> SeqStmt:
value = Evaluate(value)
ret.append(value)
return SeqStmt(ret)
-
-
-def stmt_list(stmt: Stmt) -> list[Stmt]:
- """Make list of stmt from blocks.
-
- Parameters
- ----------
- stmt : Stmt
- The input statement.
-
- Returns
- -------
- stmt_list : List[Stmt]
- The unpacked list of statements
- """
- if isinstance(stmt, SeqStmt):
- return list(stmt.seq)
- return [stmt]
diff --git a/tirx/python/tvm_tirx/script/ir_builder/__init__.py
b/tirx/python/tvm_tirx/script/ir_builder/__init__.py
index bf2b887ab5..68b9c2e1f8 100644
--- a/tirx/python/tvm_tirx/script/ir_builder/__init__.py
+++ b/tirx/python/tvm_tirx/script/ir_builder/__init__.py
@@ -17,6 +17,7 @@
"""Shared TVMScript construction APIs and lazy language variant builders."""
from tvm_tirx.ir import (
+ Bind,
Call,
DataTypeImm,
FuncType,
@@ -68,6 +69,7 @@ from .parser_protocol import (
__all__ = [
"MISSING",
"AlreadyEmitted",
+ "Bind",
"Call",
"DataTypeImm",
"FuncType",
diff --git a/tirx/python/tvm_tirx/script/ir_builder/base.py
b/tirx/python/tvm_tirx/script/ir_builder/base.py
index 4b9ee703d1..65f6ebe7da 100644
--- a/tirx/python/tvm_tirx/script/ir_builder/base.py
+++ b/tirx/python/tvm_tirx/script/ir_builder/base.py
@@ -297,8 +297,9 @@ class AlreadyEmitted(Generic[_T]):
Parameters
----------
value : T
- Native result retained without copying or emitting it. Operator APIs
+ Native result retained without copying or emitting it. Void operator
APIs
wrap their original Call; statement APIs wrap the emitted statement.
+ Allocation APIs wrap the bound Var or its scalar TensorLoad view.
Attributes
----------
diff --git a/tirx/python/tvm_tirx/script/ir_builder/stmt.py
b/tirx/python/tvm_tirx/script/ir_builder/stmt.py
index 402956631f..e8bba6641e 100644
--- a/tirx/python/tvm_tirx/script/ir_builder/stmt.py
+++ b/tirx/python/tvm_tirx/script/ir_builder/stmt.py
@@ -30,7 +30,7 @@ from typing import Any
import tvm_tirx
from tvm_tirx import ir as _ir
-from tvm_tirx.ir import Expr, IntImm, TensorRegion, Type, Var
+from tvm_tirx.ir import Expr, IntImm, TensorRegion, Var
from tvm_tirx.ir import StringImm as _StringImm
from tvm_tirx.script.ir_builder import base as _base
from tvm_tirx.script.ir_builder.base import AlreadyEmitted
@@ -53,8 +53,19 @@ def _as_expr(value):
def _enter_concise(frame: _base.IRBuilderFrame) -> Any:
- # add_callback registers on the active parent before the child enters.
- # Later statements emit into the child; parent exit closes this scope.
+ """Enter a child scope that lasts until the active parent exits.
+
+ Register the child's exit callback on the parent before entering the child.
+ Subsequent statements are emitted into the child; exiting the parent first
+ closes the child and incorporates its statements into the parent.
+
+ For example, inside an active IRBuilder::
+
+ with T.function():
+ T.func_name_("example")
+ i = _enter_concise(T.serial(4))
+ T.evaluate(i) # Emitted inside the loop, closed on function exit.
+ """
frame.add_callback(_partial(frame.__exit__, None, None, None))
return frame.__enter__()
@@ -68,41 +79,6 @@ def unpack_(value: Any) -> Any:
return value
-def bind( # pylint: disable=invalid-name
- value: Expr,
- type_annotation: Type | None = None, # pylint:
disable=redefined-outer-name
- *,
- var: Var | None = None, # pylint: disable=redefined-outer-name
-) -> Var:
- """Create a Bind (variable binding).
-
- Emits a flat Bind statement to the current frame and returns the bound
variable.
-
- Parameters
- ----------
- value : Expr
- The value to be bound.
- type_annotation : Optional[Type] = None
- The type annotation of the binding. Usually it is used for
fine-grained var typing,
- particularly, PtrType.
- var : Optional[Var] = None
- The variable to bind. If not specified, a new variable will be created.
-
- Returns
- -------
- var : Var
- The bound variable.
- """
- if type_annotation is not None:
- # Canonical Vars are callable when they denote functions. Here a Var
is
- # already a resolved type annotation, rather than a deferred
annotation factory.
- if callable(type_annotation) and not isinstance(type_annotation, Expr
| _ir.Type):
- type_annotation = type_annotation()
- if isinstance(type_annotation, _ir.Var):
- type_annotation = type_annotation.ty
- return _ffi_api.Bind(value, type_annotation, var) # type:
ignore[attr-defined] # pylint: disable=no-member
-
-
def evaluate(
value: Expr | AlreadyEmitted[tvm_tirx.ir.Call] |
AlreadyEmitted[tvm_tirx.ir.Stmt],
) -> AlreadyEmitted[tvm_tirx.ir.Call] | AlreadyEmitted[tvm_tirx.ir.Stmt]:
diff --git a/tirx/python/tvm_tirx/testing/bench.py
b/tirx/python/tvm_tirx/testing/bench.py
index 4e0cafa25c..68dd8c6786 100644
--- a/tirx/python/tvm_tirx/testing/bench.py
+++ b/tirx/python/tvm_tirx/testing/bench.py
@@ -940,11 +940,13 @@ class CudaProfiler:
# Assume Expr-like input; use as-is
self.profiler_enabled = profiler_enabled # type:
ignore[assignment]
- self.profiler_tag = T.bind(
- T.alloc_tensor(ty_args=[T.Tensor([1], "uint64", scope="local",
align=8)])
+ self.profiler_tag = T.bind_(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "uint64", scope="local",
align=8)]),
+ name="profiler_tag",
)
- self.profiler_write_offset = T.bind(
- T.alloc_tensor(ty_args=[T.Tensor([1], "uint32", scope="local",
align=8)])
+ self.profiler_write_offset = T.bind_(
+ T.alloc_tensor(ty_args=[T.Tensor([1], "uint32", scope="local",
align=8)]),
+ name="profiler_write_offset",
)
def _leader(self, leader: None | tvm_tirx.tirx.Expr | bool):
diff --git a/tirx/python/tvm_tirx/tirx/script/__init__.py
b/tirx/python/tvm_tirx/tirx/script/__init__.py
index e56b713dcd..88d0ef8ae0 100644
--- a/tirx/python/tvm_tirx/tirx/script/__init__.py
+++ b/tirx/python/tvm_tirx/tirx/script/__init__.py
@@ -41,7 +41,6 @@ def _initialize() -> None:
(name, value) for name, value in vars(builder).items() if not
name.startswith("_")
)
globals().update(
- bind=builder.bind,
ir_module=entry.ir_module,
pyfunc=entry.pyfunc,
function=entry.make_decorator(builder,
namespace_path="tirx.function"),
diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py
b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py
index a13cfd9b2f..b2784a083a 100644
--- a/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py
+++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py
@@ -194,6 +194,7 @@ def elected():
)
+@_register_mutable_decl("tirx.alloc_tensor")
def alloc_tensor(
placement=None,
*,
@@ -202,12 +203,16 @@ def alloc_tensor(
annotations: dict[str, Any] | None = None,
loc: Location = UNKNOWN_LOC,
ty=None,
-) -> ir.Call:
+) -> ir.Call | AlreadyEmitted[ir.Var]:
"""Allocate the tensor described by one explicit ``Tensor`` type argument.
Optional placement is an ordinary tuple operand. Allocation annotations
belong to ``attrs``; ``annotations`` accepts a Python dictionary whose
- scalar values are normalized to typed IR constants.
+ scalar values are normalized to typed IR constants. Under an active
builder,
+ emit one binding and return its tensor variable in an ``AlreadyEmitted``
+ receipt. Outside a builder return
+ the allocation Call. Construct ``ir.Call("tirx.alloc_tensor", ...)``
directly
+ when a pure expression is needed inside a builder.
"""
if attrs is not None and annotations is not None:
raise TypeError("alloc_tensor cannot mix attrs and annotations")
@@ -236,7 +241,7 @@ def alloc_tensor(
)
placement = ir.Tuple([normalize(value) for value in placement])
args.append(placement)
- return ir.Call(
+ value = ir.Call(
"tirx.alloc_tensor",
args,
attrs=attrs,
@@ -245,11 +250,17 @@ def alloc_tensor(
ty=ty,
)
+ return AlreadyEmitted(_bind_tensor(value)) if IRBuilder.is_in_scope() else
value
+
def _bind_tensor(value):
- from tvm_tirx.script.ir_builder.stmt import bind # pylint:
disable=import-outside-toplevel
+ from tvm_tirx.script.ir_builder.stmt import (
+ add_to_parent, # pylint: disable=import-outside-toplevel
+ )
- tensor = bind(value)
+ IRBuilder.current()._attach_loc(value)
+ tensor = ir.Var("v", value.ty)
+ add_to_parent(ir.Bind(tensor, value))
_record_meta_resource(tensor, skip_frames=2)
return tensor
@@ -263,8 +274,9 @@ def _alloc_tensor(
layout=MISSING,
allocated_addr=None,
annotations=None,
-):
- """Implement separately named allocation conveniences using the canonical
Op."""
+) -> AlreadyEmitted[ir.Var]:
+ """Emit a tensor allocation and return its bound variable in a receipt."""
+ IRBuilder.current()
scope = MemoryScope(scope)
if allocated_addr is not None:
if scope in (_SCOPE_GLOBAL, _SCOPE_SHARED, _SCOPE_SHARED_DYN,
_SCOPE_LOCAL):
@@ -281,9 +293,7 @@ def _alloc_tensor(
layout,
data=data,
)
- return _bind_tensor(
- alloc_tensor(allocated_addr, ty_args=[tensor_type],
annotations=annotations)
- )
+ return alloc_tensor(allocated_addr, ty_args=[tensor_type],
annotations=annotations)
class LetAnnotation:
@@ -363,16 +373,18 @@ def alloc_scalar(
scope: str | MemoryScope = _SCOPE_GLOBAL,
*,
annotations: dict[str, Any] | None = None,
-) -> _ir.TensorLoad:
+) -> AlreadyEmitted[_ir.TensorLoad]:
"""Allocate a zero-dimensional tensor (scalar), with optional allocation
annotations."""
- buf = _alloc_tensor(shape=(1,), dtype=dtype, scope=scope, layout=None,
annotations=annotations)
+ buf = _alloc_tensor(
+ shape=(1,), dtype=dtype, scope=scope, layout=None,
annotations=annotations
+ ).value
assert is_tensor_var(buf)
scalar = buf[0]
- return scalar
+ return AlreadyEmitted(scalar)
@_register_mutable_decl("tirx.decl_scalar")
-def decl_scalar(dtype, data, scope, layout=None) -> _ir.TensorLoad:
+def decl_scalar(dtype, data, scope, layout=None) ->
AlreadyEmitted[_ir.TensorLoad]:
"""Declare a zero-dimensional tensor (scalar) from a pointer."""
buf = _bind_tensor(
decl_tensor(
@@ -390,13 +402,13 @@ def decl_scalar(dtype, data, scope, layout=None) ->
_ir.TensorLoad:
)
assert is_tensor_var(buf)
scalar = buf[0]
- return scalar
+ return AlreadyEmitted(scalar)
@_register_mutable_decl("tirx.shared_scalar")
def shared_scalar(
dtype: str = "float32", *, annotations: dict[str, Any] | None = None
-) -> _ir.TensorLoad:
+) -> AlreadyEmitted[_ir.TensorLoad]:
"""Allocate a zero-dimensional tensor in shared memory."""
return alloc_scalar(dtype=dtype, scope="shared", annotations=annotations)
@@ -408,7 +420,7 @@ def local_scalar(
*,
name: str | None = None,
annotations: dict[str, Any] | None = None,
-) -> _ir.TensorLoad:
+) -> AlreadyEmitted[_ir.TensorLoad]:
"""Allocate scalar local storage, optionally naming and initializing it.
Parameters
@@ -425,15 +437,16 @@ def local_scalar(
Returns
-------
- scalar : TensorLoad
- A load from index zero of the one-element local tensor, with no layout.
+ scalar : AlreadyEmitted[TensorLoad]
+ A receipt for the load from index zero of the one-element local tensor.
"""
- scalar = alloc_scalar(dtype=dtype, scope="local", annotations=annotations)
+ result = alloc_scalar(dtype=dtype, scope="local", annotations=annotations)
+ scalar = result.value
if name is not None:
IRBuilder.name(name, scalar.source)
if init is not None:
assign(scalar, init)
- return scalar
+ return result
def assign(dst: _ir.TensorLoad | _ir.Var, value: Any) ->
AlreadyEmitted[_ir.Stmt]:
diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py
b/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py
index 42387cc3a2..87e14747a8 100644
--- a/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py
+++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py
@@ -31,9 +31,6 @@ from typing import Any
import tvm_ffi
import tvm_tirx
-
-# isort: off
-# isort: on
from tvm_tirx import ir as _ir
from tvm_tirx import tirx as _tir
from tvm_tirx.ir import IntImm, is_prim_expr
@@ -43,51 +40,26 @@ from tvm_tirx.script.ir_builder import base as _base
from tvm_tirx.script.ir_builder import frame as _shared_frame
from tvm_tirx.script.ir_builder.base import AlreadyEmitted
from tvm_tirx.script.ir_builder.base import IRBuilder as _IRBuilder
+
+# isort: off
from tvm_tirx.script.ir_builder.stmt import (
_enter_concise as _enter_concise,
-)
-from tvm_tirx.script.ir_builder.stmt import (
add_to_parent as add_to_parent,
-)
-from tvm_tirx.script.ir_builder.stmt import (
assert_ as assert_,
-)
-from tvm_tirx.script.ir_builder.stmt import (
- bind as _bind,
-)
-from tvm_tirx.script.ir_builder.stmt import (
break_ as break_,
-)
-from tvm_tirx.script.ir_builder.stmt import (
continue_ as continue_,
-)
-from tvm_tirx.script.ir_builder.stmt import (
else_ as else_,
-)
-from tvm_tirx.script.ir_builder.stmt import (
evaluate as evaluate,
-)
-from tvm_tirx.script.ir_builder.stmt import (
grid as grid,
-)
-from tvm_tirx.script.ir_builder.stmt import (
if_ as if_,
-)
-from tvm_tirx.script.ir_builder.stmt import (
region as region,
-)
-from tvm_tirx.script.ir_builder.stmt import (
return_ as return_,
-)
-from tvm_tirx.script.ir_builder.stmt import (
then_ as then_,
-)
-from tvm_tirx.script.ir_builder.stmt import (
unpack_ as unpack_,
-)
-from tvm_tirx.script.ir_builder.stmt import (
while_ as while_,
)
+
+# isort: on
from tvm_tirx.tirx import Expr
from . import _ffi_api, frame
@@ -226,14 +198,6 @@ def _name(value: Any, name: str | None, loc: _Loc) -> Any:
return _base.at_(loc, value)
-def bind(value, type_annotation=None, *, var=None):
- """Bind an expression and track tensor resources created in meta
classes."""
- result = _bind(value, type_annotation, var=var)
- if _tir.is_tensor_var(result):
- _native._record_meta_resource(result)
- return result
-
-
def bind_(
value: Any = _base.MISSING,
*,
@@ -248,9 +212,26 @@ def bind_(
Returned Vars, including tensors, and metadata retain identity, names and
locs.
Other expressions create native Bind nodes; value_loc belongs to the RHS.
- Explicit typed bindings and frame targets retain their separate contracts.
+ Emitted allocation receipts unwrap to their original Var or TensorLoad,
+ checking annotations without creating another binding. Other receipts
remain
+ intact. Explicit typed bindings and frame targets retain their separate
contracts.
"""
name_loc = loc if name_loc is None else name_loc
+ if isinstance(value, _base.AlreadyEmitted):
+ result = _base.at_(value_loc, value).value
+ if not isinstance(result, _ir.Var | _ir.TensorLoad):
+ return value
+ if ty is not None:
+ if isinstance(ty, _native.LetAnnotation):
+ annotation = ty.as_var(rhs_dtype=result.ty).ty
+ else:
+ annotation = (
+ ty() if callable(ty) and not isinstance(ty, _ir.Expr |
_ir.Type) else ty
+ )
+ annotation = annotation.ty if isinstance(annotation, _ir.Expr)
else annotation
+ _ir.assert_structural_equal(result.ty, annotation)
+ _name(result.source if isinstance(result, _ir.TensorLoad) else result,
name, name_loc)
+ return result
if frame_value:
if isinstance(value, _python.list | _python.tuple | tvm_ffi.Array):
for index, item in enumerate(value):
@@ -273,7 +254,7 @@ def bind_(
if not isinstance(value, _ir.Var):
_base.at_(value_loc, value)
variable = _name(ty.as_var(rhs_dtype=value.ty), name, name_loc)
- _base.with_at_group_(loc, lambda: bind(value, var=variable))
+ _base.with_at_group_(loc, lambda: emit_(_ir.Bind(variable, value)))
return variable
if ty is not None:
annotation = ty() if callable(ty) and not isinstance(ty, _ir.Expr |
_ir.Type) else ty
@@ -282,11 +263,10 @@ def bind_(
if not isinstance(value, _ir.Var):
_base.at_(value_loc, value)
variable = _ir.Var(name or "", annotation)
- return _name(_base.with_at_group_(loc, lambda: bind(value,
var=variable)), name, name_loc)
+ _base.with_at_group_(loc, lambda: emit_(_ir.Bind(variable, value)))
+ return _name(variable, name, name_loc)
if value is _base.MISSING:
raise ValueError("An uninitialized binding requires a scalar type
annotation")
- if isinstance(value, _base.AlreadyEmitted):
- return _base.at_(value_loc, value)
if isinstance(value, _base.IRBuilderFrame):
frame_loc = value_loc if value_loc is not None else loc
return _name(_enter_concise(_base.at_(frame_loc, value)), name,
name_loc)
@@ -299,7 +279,9 @@ def bind_(
if not isinstance(value, _ir.Expr | _python.int | _python.float |
_python.bool | str):
return value
value = _base.at_(value_loc, _op._as_expr(value))
- return _name(_base.with_at_group_(loc, lambda: bind(value)), name,
name_loc)
+ variable = _ir.Var("v", value.ty)
+ _base.with_at_group_(loc, lambda: emit_(_ir.Bind(variable, value)))
+ return _name(variable, name, name_loc)
def decl_mutable_cell_(
@@ -316,11 +298,21 @@ def decl_mutable_cell_(
allocate their declared shape. Var declaration producers retain their own
effects.
"""
name_loc = loc if name_loc is None else name_loc
+ if isinstance(value, _base.AlreadyEmitted):
+ return bind_(value, ty=ty, name=name, loc=loc, name_loc=name_loc)
+ if _tir.is_tensor_var(value) and ty is not None:
+ if isinstance(ty, _native.LetAnnotation):
+ annotation = ty.as_var(rhs_dtype=value.ty).ty
+ else:
+ annotation = ty() if callable(ty) and not isinstance(ty, _ir.Expr
| _ir.Type) else ty
+ annotation = annotation.ty if isinstance(annotation, _ir.Expr)
else annotation
+ _ir.assert_structural_equal(value.ty, annotation)
+ return _name(value, name, name_loc)
if isinstance(ty, _PrimTypeArrayAnnotation):
if value is not _base.MISSING:
raise ValueError("Vector annotation does not support an
initializer")
return _name(
- _base.with_at_group_(loc, lambda: _native.alloc_local(ty.shape,
ty.dtype)),
+ _base.with_at_group_(loc, lambda: _native.alloc_local(ty.shape,
ty.dtype)).value,
name,
name_loc,
)
@@ -329,7 +321,7 @@ def decl_mutable_cell_(
annotation = annotation.ty if isinstance(annotation, _ir.Expr) else
annotation
if not isinstance(annotation, _ir.PrimType) or str(annotation) ==
"handle":
raise TypeError("Mutable scalar annotations require a primitive
scalar type")
- storage = _base.with_at_group_(loc, lambda:
_native.local_scalar(str(annotation)))
+ storage = _base.with_at_group_(loc, lambda:
_native.local_scalar(str(annotation))).value
if value is not _base.MISSING:
set_mutable_cell_(storage, value, loc=loc)
else:
@@ -391,6 +383,8 @@ def emit_(value: Any, *, loc: _Loc = None) -> None:
_enter_concise(_base.at_(loc, frame))
elif isinstance(value, tvm_tirx.ir.Stmt):
add_to_parent(_base.at_(loc, value))
+ if isinstance(value, _ir.Bind) and _tir.is_tensor_var(value.var):
+ _native._record_meta_resource(value.var)
else:
# Native conversion owns Python literals; annotate the exact expression
# it stored, as well as the statement, without converting or emitting
twice.
@@ -788,7 +782,6 @@ __all__ = [
"and_",
"arg_",
"assert_",
- "bind",
"bind_",
"break_",
"call_global_var_",
diff --git a/tirx/src/script/ir_builder/ir.cc b/tirx/src/script/ir_builder/ir.cc
index 9bb4e99059..c855dd9f4a 100644
--- a/tirx/src/script/ir_builder/ir.cc
+++ b/tirx/src/script/ir_builder/ir.cc
@@ -290,21 +290,6 @@ AssertFrame Assert(PrimExpr condition, ffi::String
error_kind,
return AssertFrame(n);
}
-Var Bind(Expr value, ffi::Optional<Type> type_annotation, ffi::Optional<Var>
var) {
- Expr value_expr = value;
- Var bind_var = [&]() {
- if (var.has_value()) {
- return var.value();
- } else if (type_annotation.has_value()) {
- return Var("v", type_annotation.value());
- } else {
- return Var("v", value_expr->ty);
- }
- }();
- AddToParent(tvm::Bind(bind_var, value_expr));
- return bind_var;
-}
-
RegionFrame Region(Op op, ffi::Array<Expr> args,
ffi::Optional<ffi::Array<Var>> body_params,
DictAttrs attrs) {
TVM_FFI_CHECK(tvm::IsRegionOp(op), ValueError)
@@ -385,7 +370,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("script.ir_builder.For", For)
.def("script.ir_builder.Grid", Grid)
.def("script.ir_builder.Assert", Assert)
- .def("script.ir_builder.Bind", Bind)
.def("script.ir_builder.While", While)
.def("script.ir_builder.Return", Return)
.def("script.ir_builder.Break", Break)
diff --git a/tirx/tests/python/tirx/analysis/test_check_well_defined.py
b/tirx/tests/python/tirx/analysis/test_check_well_defined.py
index c781d95e04..f5806fbc27 100644
--- a/tirx/tests/python/tirx/analysis/test_check_well_defined.py
+++ b/tirx/tests/python/tirx/analysis/test_check_well_defined.py
@@ -80,8 +80,8 @@ def test_error_for_nested_rebind_usage():
@T.function(check_well_formed=False)
def func():
- T.bind(42, var=i)
- T.bind(42, var=i)
+ T.Bind(i, 42)
+ T.Bind(i, 42)
T.evaluate(i)
with pytest.raises(
@@ -103,9 +103,9 @@ def test_error_for_repeated_binding():
@T.function(check_well_formed=False)
def func():
- T.bind(42, var=i)
+ T.Bind(i, 42)
T.evaluate(i)
- T.bind(17, var=i)
+ T.Bind(i, 17)
T.evaluate(i)
with pytest.raises(
@@ -124,12 +124,12 @@ def test_error_for_cross_function_reuse():
class mod:
@T.function
def func1():
- T.bind(42, var=i)
+ T.Bind(i, 42)
T.evaluate(i)
@T.function
def func2():
- T.bind(42, var=i)
+ T.Bind(i, 42)
T.evaluate(i)
with pytest.raises((ValueError, RuntimeError), match="multiple definitions
of variable i"):
@@ -215,10 +215,10 @@ def
test_error_message_without_previous_definition_location():
@T.function(check_well_formed=False)
def func():
- T.bind(42, var=x)
+ T.Bind(x, 42)
T.evaluate(x)
- T.bind(99, var=x) # This should trigger the error
+ T.Bind(x, 99) # This should trigger the error
T.evaluate(x)
with pytest.raises((ValueError, RuntimeError)) as exc_info:
@@ -242,8 +242,8 @@ def test_error_message_with_previous_definition_location():
@T.function(check_well_formed=False)
def func():
- T.bind(42, var=x)
- T.bind(99, var=x) # This should trigger the error
+ T.Bind(x, 42)
+ T.Bind(x, 99) # This should trigger the error
T.evaluate(x)
with pytest.raises((ValueError, RuntimeError)) as exc_info:
@@ -271,10 +271,10 @@ def test_sequential_redefinition_with_location():
@T.function(check_well_formed=False)
def func():
- T.bind(1, var=x)
+ T.Bind(x, 1)
T.evaluate(x)
- T.bind(2, var=x) # This should trigger the error
+ T.Bind(x, 2) # This should trigger the error
T.evaluate(x)
with pytest.raises((ValueError, RuntimeError)) as exc_info:
diff --git a/tirx/tests/python/tirx/script/test_basic_usage.py
b/tirx/tests/python/tirx/script/test_basic_usage.py
index 67199b35f9..9834831fab 100644
--- a/tirx/tests/python/tirx/script/test_basic_usage.py
+++ b/tirx/tests/python/tirx/script/test_basic_usage.py
@@ -1340,8 +1340,8 @@ def elif_chain_with_else():
def bind_var():
@T.function
def func():
- x = T.bind(0)
- y = T.bind(0)
+ x = 0
+ y = 0
T.evaluate(0)
T.evaluate(0)
@@ -1599,18 +1599,17 @@ def test_implicit_evaluate_call_extern():
def test_preserve_trivial_let_binding():
- """Trivial `T.let[...]` annotations survive the parser as LetStmt and are
not inlined.
+ """Trivial `T.let[...]` annotations survive the parser as Bind and are not
inlined.
- In fork, bare `j = i` lowers to a local_scalar (AllocTensor +
TensorStore); the
- LetStmt form is opt-in via `T.let[T.dtype]`. Both the explicit
`T.bind(..., var=j)`
- builder API and the `j: T.let[T.dtype]` annotation produce the same
LetStmt IR.
+ Both an explicit `T.Bind(j, i)` statement and the `j: T.let[T.dtype]`
+ annotation produce the same Bind IR, preserving a distinct variable for
`j`.
"""
j = T.dynamic("j", "int32")
@T.function
def explicit(i: T.int32):
- T.bind(i, var=j)
+ T.Bind(j, i)
T.evaluate(j)
@T.function
@@ -1628,7 +1627,7 @@ def test_preserve_trivial_let_binding_of_value():
@T.function
def explicit(i: T.int32):
- T.bind(42, var=j)
+ T.Bind(j, 42)
T.evaluate(j)
@T.function
diff --git a/tirx/tests/python/tirx/script/test_ir_builder.py
b/tirx/tests/python/tirx/script/test_ir_builder.py
index 3b414f091e..aa7da77c0e 100644
--- a/tirx/tests/python/tirx/script/test_ir_builder.py
+++ b/tirx/tests/python/tirx/script/test_ir_builder.py
@@ -25,7 +25,7 @@ import tvm_tirx.testing
from tvm_tirx import tirx
from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.ir.location import SourceLoc, SourceName
-from tvm_tirx.script.ir_builder import IRBuilder
+from tvm_tirx.script.ir_builder import AlreadyEmitted, IRBuilder
from tvm_tirx.tirx.layout import S, TileLayout
from tvm_tirx.tirx.script import ir_builder as T
@@ -39,6 +39,116 @@ def _is_tensor_binding(node, *op_names):
)
+def test_alloc_tensor_builder_boundary():
+ tensor_type = T.Tensor((4,), "int32", scope="local")
+ outside = T.alloc_tensor(ty_args=[tensor_type])
+ assert isinstance(outside, tvm_tirx.ir.Call)
+ loc = SourceLoc(SourceName("allocation.py"), 3, 1, 3, 20)
+ with IRBuilder() as ib:
+ with ib.with_loc(loc):
+ raw = tvm_tirx.ir.Call("tirx.alloc_tensor", [],
ty_args=[tensor_type])
+ receipt = T.alloc_tensor(ty_args=[tensor_type],
annotations={"alignment": 16})
+ assert isinstance(receipt, AlreadyEmitted)
+ tensor = T.bind_(receipt, name="tensor", ty=tensor_type,
value_loc=loc)
+ assert tensor.same_as(receipt.value)
+ assert T.bind_(tensor, name="alias").same_as(tensor)
+ assert tensor.name == "tensor"
+ T.emit_(receipt)
+ with pytest.raises(ValueError):
+ T.decl_mutable_cell_(receipt, ty=T.Tensor((8,), "int32",
scope="local"))
+ binding = ib.get()
+ assert isinstance(binding, tvm_tirx.ir.Bind)
+ assert binding.var.same_as(tensor)
+ assert binding.var.ty.same_as(tensor_type)
+ assert binding.var.loc.same_as(loc)
+ assert not binding.value.same_as(raw)
+ assert binding.value.attrs["alignment"].value == 16
+ assert binding.loc.same_as(loc)
+ assert binding.value.loc.same_as(loc)
+
+
[email protected](
+ "producer,scalar",
+ [
+ (lambda: T.alloc_local((4,), "int32"), False),
+ (lambda: T.alloc_shared((4,), "int32"), False),
+ (lambda: T.smem((4,), "int32"), False),
+ (lambda: T.tmem((4,), "int32"), False),
+ (lambda: T.alloc_scalar("int32", "local"), True),
+ (lambda: T.local_scalar("int32"), True),
+ (lambda: T.shared_scalar("int32"), True),
+ (
+ lambda: T.decl_scalar(
+ "int32", tvm_tirx.ir.Var("data", tvm_tirx.ir.PtrType("int32",
"local")), "local"
+ ),
+ True,
+ ),
+ ],
+)
+def test_allocation_convenience_receipts(producer, scalar):
+ with IRBuilder() as ib:
+ receipt = producer()
+ assert isinstance(receipt, AlreadyEmitted)
+ value = T.bind_(receipt, name="storage")
+ assert value.same_as(receipt.value)
+ assert isinstance(value, tvm_tirx.ir.TensorLoad if scalar else
tvm_tirx.ir.Var)
+ assert T.decl_mutable_cell_(receipt, ty=value.ty,
name="storage").same_as(value)
+ T.emit_(receipt)
+ binding = ib.get()
+ assert isinstance(binding, tvm_tirx.ir.Bind)
+ assert binding.var.same_as(value.source if scalar else value)
+ assert binding.var.name == "storage"
+
+
+def test_binding_preserves_void_and_statement_receipts():
+ with IRBuilder() as ib:
+ void = T.call_extern("side_effect", ty="void")
+ statement = T.evaluate(1)
+ for receipt in (void, statement):
+ assert isinstance(receipt, AlreadyEmitted)
+ assert T.bind_(receipt) is receipt
+ assert T.bind_(receipt, ty=T.int32) is receipt
+ T.emit_(receipt)
+ assert len(ib.get().seq) == 2
+
+
+def test_alloc_tensor_parsed_identity():
+ from tvm_tirx.script import tirx as TS
+
+ @TS.function(private=True)
+ def function():
+ tensor = TS.alloc_tensor(ty_args=[TS.Tensor((4,), "int32",
scope="local")])
+ alias = tensor
+ alias[0] = 1
+
+ binding, store = function.body.seq
+ assert _is_tensor_binding(binding, "tirx.alloc_tensor")
+ assert binding.var.name == "tensor"
+ assert store.dest.same_as(binding.var)
+ assert isinstance(binding.var.loc, tvm_tirx.ir.CallSiteLoc)
+ assert binding.var.loc.caller.end_column -
binding.var.loc.caller.start_column == len("tensor")
+ assert isinstance(binding.value.loc, SourceLoc)
+
+
[email protected]("use_let", [False, True])
+def test_alloc_tensor_annotated_binding(use_let):
+ from tvm_tirx.script import tirx as TS
+
+ tensor_type = TS.Tensor((4,), "int32", scope="local")
+ annotation = TS.let[tensor_type] if use_let else tensor_type
+
+ @TS.function(private=True)
+ def function():
+ tensor: annotation = TS.alloc_tensor(ty_args=[TS.Tensor((4,), "int32",
scope="local")])
+ tensor[0] = 1
+
+ binding, store = function.body.seq
+ assert _is_tensor_binding(binding, "tirx.alloc_tensor")
+ assert binding.var.name == "tensor"
+ assert store.dest.same_as(binding.var)
+ assert_structural_equal(binding.var.ty, TS.Tensor((4,), "int32",
scope="local"))
+
+
def test_ir_builder_tir_for():
with IRBuilder() as ib:
with T.serial(128) as a:
@@ -144,10 +254,12 @@ def test_ir_builder_tir_assert():
def test_ir_builder_tir_bind():
- # Test that T.bind emits a flat Bind statement and returns the Var.
+ # Construct and emit a flat Bind statement with an explicit variable.
+ assert T.Bind is tvm_tirx.ir.Bind
with IRBuilder() as ib:
- v = T.bind(tvm_tirx.ir.IntImm("int32", 2))
- # the let binding generated by IRBuilder
+ v = tirx.Var("v", "int32")
+ T.emit_(T.Bind(v, tvm_tirx.ir.IntImm("int32", 2)))
+ # the binding generated by IRBuilder
let_actual = ib.get()
# Bind is now flat (no body), so a single Bind stmt is emitted.
@@ -155,8 +267,19 @@ def test_ir_builder_tir_bind():
# Check if the generated ir is expected
assert_structural_equal(let_actual, let_expected, map_free_vars=True)
- # Check that the returned value is a Var
- assert isinstance(v, tirx.Var)
+ assert let_actual.var.same_as(v)
+
+
+def test_removed_binding_helpers_are_not_exported():
+ from tvm_tirx.ir import stmt
+ from tvm_tirx.script import tirx as TS
+ from tvm_tirx.script.ir_builder import stmt as builder_stmt
+
+ assert not hasattr(T, "bind")
+ assert not hasattr(TS, "bind")
+ assert not hasattr(builder_stmt, "bind")
+ assert not hasattr(tvm_tirx.ir, "stmt_list")
+ assert not hasattr(stmt, "stmt_list")
def test_ir_builder_tir_while():
@@ -331,12 +454,14 @@ def test_concrete_mutable_scalar(declare):
with IRBuilder() as ib:
with T.function():
if declare:
- owner = T.bind(T.alloc_tensor(ty_args=[T.Tensor((4,), "int32",
scope="local")]))
- scalar = T.decl_scalar(
- "int32", owner.data, "local",
layout=T.TileLayout(T.S[(1,)] + (2))
+ owner = T.bind_(T.alloc_tensor(ty_args=[T.Tensor((4,),
"int32", scope="local")]))
+ scalar = T.bind_(
+ T.decl_scalar(
+ "int32", owner.data, "local",
layout=T.TileLayout(T.S[(1,)] + (2))
+ )
)
else:
- scalar = T.alloc_scalar("int32", "local")
+ scalar = T.bind_(T.alloc_scalar("int32", "local"))
assert type(scalar) is tvm_tirx.ir.TensorLoad
named = T.decl_mutable_cell_(scalar, name="counter")
assert named.same_as(scalar)
@@ -369,13 +494,15 @@ def test_local_scalar_initialization_and_assign():
with IRBuilder() as ib:
with T.function():
with ib.with_loc(loc):
- scalar = T.local_scalar("int32", 2, name="counter",
annotations=annotations)
- uninitialized = T.local_scalar()
+ scalar = T.bind_(
+ T.local_scalar("int32", 2, name="counter",
annotations=annotations)
+ )
+ uninitialized = T.bind_(T.local_scalar())
receipt = T.assign(scalar, 3)
assert isinstance(receipt, AlreadyEmitted)
T.emit_(receipt)
source_receipt = T.assign(scalar.source, 4)
- global_tensor = T.bind(T.alloc_tensor(ty_args=[T.Tensor((2,),
"int32")]))
+ global_tensor = T.bind_(T.alloc_tensor(ty_args=[T.Tensor((2,),
"int32")]))
indexed_receipt = T.assign(global_tensor[1], 5)
for invalid in (tirx.Var("immutable", "int32"), scalar + 1,
global_tensor):
with pytest.raises(TypeError, match="scalar storage"):
@@ -465,7 +592,7 @@ def test_ir_builder_tir_allocate():
with IRBuilder() as ib:
with T.function():
T.func_name_("test")
- buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([10], "float32",
scope="local")]))
+ buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([10], "float32",
scope="local")]))
T.evaluate(1)
# the allocate generated by IRBuilder
@@ -484,7 +611,7 @@ def test_ir_builder_tir_decl_tensor():
with IRBuilder() as ib:
with T.function():
T.func_name_("test")
- buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor([128, 128],
"float32")]))
+ buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor([128, 128],
"float32")]))
T.evaluate(1)
# the decl_tensor generated by IRBuilder
diff --git a/tirx/tests/python/tirx/script/test_meta_programming.py
b/tirx/tests/python/tirx/script/test_meta_programming.py
index 5b3e83aa63..f095c44feb 100644
--- a/tirx/tests/python/tirx/script/test_meta_programming.py
+++ b/tirx/tests/python/tirx/script/test_meta_programming.py
@@ -33,7 +33,7 @@ def test_meta_class_constructor_rejects_unowned_resource():
@T.meta_class
class Bad:
def __init__(self):
- tmp = T.bind(T.alloc_tensor(ty_args=[T.Tensor((1,), "int32",
scope="local")]))
+ tmp = T.bind_(T.alloc_tensor(ty_args=[T.Tensor((1,), "int32",
scope="local")]))
with pytest.raises(ValueError):
@@ -50,8 +50,8 @@ def
test_meta_class_multiple_instances_preserve_owned_resources():
class Holder:
def __init__(self, external):
self.external = external
- self.buf = T.bind(T.alloc_tensor(ty_args=[T.Tensor((2,), "int32",
scope="local")]))
- self.scalar = T.local_scalar("int32")
+ self.buf = T.bind_(T.alloc_tensor(ty_args=[T.Tensor((2,), "int32",
scope="local")]))
+ self.scalar = T.bind_(T.local_scalar("int32"))
instances.append(self)
@T.function(private=True)
diff --git a/tirx/tests/python/tirx/script/test_printer.py
b/tirx/tests/python/tirx/script/test_printer.py
index fb0d49d62f..5f68b5811e 100644
--- a/tirx/tests/python/tirx/script/test_printer.py
+++ b/tirx/tests/python/tirx/script/test_printer.py
@@ -256,9 +256,11 @@ def test_allocate(declare):
TB.func_name_("test")
if declare is TB.decl_tensor:
data = TB.ptr("float32")
- buf = TB.bind(declare(data, ty_args=[TB.Tensor([128, 128],
"float32")]))
+ buf_value = declare(data, ty_args=[TB.Tensor([128, 128],
"float32")])
+ buf = tvm_tirx.ir.Var("", buf_value.ty)
+ TB.emit_(TB.Bind(buf, buf_value))
else:
- buf = TB.bind(declare(ty_args=[TB.Tensor([128, 128],
"float32")]))
+ buf = TB.bind_(declare(ty_args=[TB.Tensor([128, 128],
"float32")]))
TB.evaluate(1)
obj = ib.get()
_assert_print(
@@ -281,10 +283,12 @@ def test_allocate_with_decl_tensor_sugar():
with IRBuilder() as ib:
with TB.function():
TB.func_name_("test")
- buf = TB.bind(TB.alloc_tensor(ty_args=[TB.Tensor([128, 128],
"float32")]))
- buf2 = TB.bind(
- TB.decl_tensor(buf.data, ty_args=[TB.Tensor([128, 128],
"float32", data=buf.data)])
+ buf = TB.bind_(TB.alloc_tensor(ty_args=[TB.Tensor([128, 128],
"float32")]))
+ buf2_value = TB.decl_tensor(
+ buf.data, ty_args=[TB.Tensor([128, 128], "float32",
data=buf.data)]
)
+ buf2 = tvm_tirx.ir.Var("", buf2_value.ty)
+ TB.emit_(TB.Bind(buf2, buf2_value))
TB.evaluate(1)
obj = ib.get()
_assert_print(
@@ -302,10 +306,12 @@ def test_allocate_with_decl_tensor_sugar_multi_usage():
with IRBuilder() as ib:
with TB.function():
TB.func_name_("test")
- buf = TB.bind(TB.alloc_tensor(ty_args=[TB.Tensor([128, 128],
"float32")]))
- buf2 = TB.bind(
- TB.decl_tensor(buf.data, ty_args=[TB.Tensor([128, 128],
"float32", data=buf.data)])
+ buf = TB.bind_(TB.alloc_tensor(ty_args=[TB.Tensor([128, 128],
"float32")]))
+ buf2_value = TB.decl_tensor(
+ buf.data, ty_args=[TB.Tensor([128, 128], "float32",
data=buf.data)]
)
+ buf2 = tvm_tirx.ir.Var("", buf2_value.ty)
+ TB.emit_(TB.Bind(buf2, buf2_value))
TB.evaluate(buf.data)
obj = ib.get()
_assert_print(
@@ -322,10 +328,12 @@ def test_allocate_with_decl_tensor_no_sugar_mismatch():
with IRBuilder() as ib:
with TB.function():
TB.func_name_("test")
- buf = TB.bind(TB.alloc_tensor(ty_args=[TB.Tensor([128, 128],
"float32")]))
- buf2 = TB.bind(
- TB.decl_tensor(buf.data, ty_args=[TB.Tensor([256, 256],
"float32", data=buf.data)])
+ buf = TB.bind_(TB.alloc_tensor(ty_args=[TB.Tensor([128, 128],
"float32")]))
+ buf2_value = TB.decl_tensor(
+ buf.data, ty_args=[TB.Tensor([256, 256], "float32",
data=buf.data)]
)
+ buf2 = tvm_tirx.ir.Var("", buf2_value.ty)
+ TB.emit_(TB.Bind(buf2, buf2_value))
TB.evaluate(buf.data)
obj = ib.get()
_assert_print(
@@ -343,12 +351,11 @@ def test_decl_tensor():
with IRBuilder() as ib:
with TB.function():
TB.func_name_("test")
- buf = TB.bind(
- TB.decl_tensor(
- TB.ptr("float32"),
- ty_args=[TB.Tensor((10, 10), "float32",
data=TB.ptr("float32"))],
- )
+ buf_value = TB.decl_tensor(
+ TB.ptr("float32"), ty_args=[TB.Tensor((10, 10), "float32",
data=TB.ptr("float32"))]
)
+ buf = tvm_tirx.ir.Var("", buf_value.ty)
+ TB.emit_(TB.Bind(buf, buf_value))
TB.evaluate(1)
obj = ib.get()
# Print only the body (skip Function wrapper)
diff --git a/tirx/tests/python/tirx/script/test_printer_function.py
b/tirx/tests/python/tirx/script/test_printer_function.py
index 07a0c8d58f..cae67a165a 100644
--- a/tirx/tests/python/tirx/script/test_printer_function.py
+++ b/tirx/tests/python/tirx/script/test_printer_function.py
@@ -179,7 +179,9 @@ def main(A: T.Tensor((128, 128), T.float32), B:
T.Tensor((256, 256), T.float32))
def test_bind():
with IRBuilder() as ib:
with TB.function():
- v = TB.bind(TB.float32(10))
+ v_value = TB.float32(10)
+ v = tvm_tirx.ir.Var("", v_value.ty)
+ TB.emit_(TB.Bind(v, v_value))
ib.name("v", v)
TB.evaluate(1)
obj = ib.get()
diff --git a/tirx/tests/python/tirx/script/test_source_locations.py
b/tirx/tests/python/tirx/script/test_source_locations.py
index 0c8b8f00dc..0f818fdfcc 100644
--- a/tirx/tests/python/tirx/script/test_source_locations.py
+++ b/tirx/tests/python/tirx/script/test_source_locations.py
@@ -228,7 +228,7 @@ def test_explicit_loc_on_global_var_lambda_and_min_max():
assert isinstance(T.max(a, b), prim.Max)
-def test_native_bind_keeps_returned_and_stored_variable_identity():
+def test_explicit_bind_keeps_stored_and_aliased_variable_identity():
# Aliases of an explicitly produced native variable must not rename, reloc
or bind it again.
from tvm_tirx import ir
@@ -247,7 +247,8 @@ def
test_native_bind_keeps_returned_and_stored_variable_identity():
@T.function
def main(x: T.int32):
- renamed = T.bind(argument(x), var=produced)
+ T.Bind(produced, argument(x))
+ renamed = produced
alias = renamed
called = argument(alias)
observe(renamed)
diff --git a/tirx/tests/python/tirx/script/test_tensors.py
b/tirx/tests/python/tirx/script/test_tensors.py
index c383f22add..08753fce8b 100644
--- a/tirx/tests/python/tirx/script/test_tensors.py
+++ b/tirx/tests/python/tirx/script/test_tensors.py
@@ -243,12 +243,12 @@ def test_alloc_apis():
def __init__(self, Ta, inner_pool):
self.Ta = Ta
self.inner_pool = inner_pool
- self.Tb = T.shared_scalar("float16")
- self.idx = T.local_scalar("int32")
- self.inner_pool2 = T.decl_scalar(
+ self.Tb = T.bind_(T.shared_scalar("float16"))
+ self.idx = T.bind_(T.local_scalar("int32"))
+ self.inner_pool2 = T.bind_(T.decl_scalar(
"float16", self.inner_pool.data, "shared.dyn",
layout=T.TileLayout(T.S[(1,)] + 5),
- )
+ ))
@T.inline
def init(self):
@@ -1212,7 +1212,7 @@ def test_tensor_factored_views():
"""Splitting and selecting logical factors preserves the physical
mapping."""
with IRBuilder():
with Tx_builder.function():
- A = Tx_builder.bind(
+ A = Tx_builder.bind_(
Tx_builder.alloc_tensor(
ty_args=[T.Tensor((3, 64, 512), "float16", scope="shared",
layout="default")]
)
diff --git a/tirx/tests/python/tirx/script/test_tracing.py
b/tirx/tests/python/tirx/script/test_tracing.py
index bea151d4ac..e446960bce 100644
--- a/tirx/tests/python/tirx/script/test_tracing.py
+++ b/tirx/tests/python/tirx/script/test_tracing.py
@@ -93,7 +93,7 @@ def test_backend_tile_namespace_emits_once(trace):
with IRBuilder() as builder:
with I.function_():
source = I.arg_("source", T.Tensor((1, 4), T.float32))
- dest = I.bind(T.alloc_tensor(ty_args=[T.Tensor((1, 4), T.float32,
scope="local")]))
+ dest = T.bind_(T.alloc_tensor(ty_args=[T.Tensor((1, 4), T.float32,
scope="local")]))
with builder.line_trace_scope() if trace else nullcontext():
receipt = T.cuda.tile.apply(T.ptx["ld.global.v4.f32"], dest,
source)
assert isinstance(receipt, AlreadyEmitted)
diff --git a/tirx/tests/python/tirx/test_tile_namespace.py
b/tirx/tests/python/tirx/test_tile_namespace.py
index 48997db995..803a944f25 100644
--- a/tirx/tests/python/tirx/test_tile_namespace.py
+++ b/tirx/tests/python/tirx/test_tile_namespace.py
@@ -104,8 +104,10 @@ def test_omitted_and_explicit_none_layout(scope):
with IRBuilder():
with I.function():
tensor_type = I.Tensor((4, 8), "float32", scope=scope)
- allocated = I.bind(I.alloc_tensor(ty_args=[tensor_type]))
- declared = I.bind(I.decl_tensor(allocated.data,
ty_args=[tensor_type]))
+ allocated = I.bind_(I.alloc_tensor(ty_args=[tensor_type]))
+ declared_value = I.decl_tensor(allocated.data,
ty_args=[tensor_type])
+ declared = tvm_tirx.ir.Var("declared", declared_value.ty)
+ I.emit_(I.Bind(declared, declared_value))
assert allocated.layout is None and declared.layout is None
@@ -149,12 +151,12 @@ def test_none_layout_views_keep_offsets_and_strides():
builder = IRBuilder()
with builder:
with I.function():
- source = I.bind(
- I.decl_tensor(
- I.ptr("float32"),
- ty_args=[I.Tensor((4, 8), "float32",
layout=T.TileLayout(T.S[(4, 8)] + (5)))],
- )
+ source_value = I.decl_tensor(
+ I.ptr("float32"),
+ ty_args=[I.Tensor((4, 8), "float32",
layout=T.TileLayout(T.S[(4, 8)] + (5)))],
)
+ source = tvm_tirx.ir.Var("source", source_value.ty)
+ I.emit_(I.Bind(source, source_value))
reshaped = source.view(2, 2, 8)
permuted = source.permute(1, 0)
sliced = source.sub[1:3, 2:6]
@@ -204,7 +206,9 @@ def test_none_layout_dtype_views(dtype, new_dtype,
old_shape, new_shape):
builder = IRBuilder()
with builder:
with I.function():
- source = I.bind(I.decl_tensor(I.ptr(dtype),
ty_args=[I.Tensor(old_shape, dtype)]))
+ source_value = I.decl_tensor(I.ptr(dtype),
ty_args=[I.Tensor(old_shape, dtype)])
+ source = tvm_tirx.ir.Var("source", source_value.ty)
+ I.emit_(I.Bind(source, source_value))
viewed = source.view(new_dtype)
assert viewed.layout is None
assert list(map(int, viewed.shape)) == list(new_shape)
@@ -215,16 +219,18 @@ def test_none_layout_dtype_view_accepts_symbolic_extent():
with IRBuilder():
with I.function():
size = tvm_tirx.tirx.Var("size", "int32")
- source = I.bind(I.decl_tensor(I.ptr("float32"),
ty_args=[I.Tensor((size,), "float32")]))
+ source_value = I.decl_tensor(I.ptr("float32"),
ty_args=[I.Tensor((size,), "float32")])
+ source = tvm_tirx.ir.Var("source", source_value.ty)
+ I.emit_(I.Bind(source, source_value))
viewed = source.view("uint64")
assert viewed.layout is None
assert_structural_equal(viewed.shape[0], size // 2)
- misaligned = I.bind(
- I.decl_tensor(
- I.ptr("float32"),
- ty_args=[I.Tensor((8,), "float32",
layout=T.TileLayout(T.S[(8,)] + (1)))],
- )
+ misaligned_value = I.decl_tensor(
+ I.ptr("float32"),
+ ty_args=[I.Tensor((8,), "float32",
layout=T.TileLayout(T.S[(8,)] + (1)))],
)
+ misaligned = tvm_tirx.ir.Var("misaligned", misaligned_value.ty)
+ I.emit_(I.Bind(misaligned, misaligned_value))
with pytest.raises(ValueError, match="align"):
misaligned.view("uint64")
diff --git a/tirx/tests/python/tirx/tile/cuda/test_ewise.py
b/tirx/tests/python/tirx/tile/cuda/test_ewise.py
index 4661b26d8e..fc5ad16f08 100644
--- a/tirx/tests/python/tirx/tile/cuda/test_ewise.py
+++ b/tirx/tests/python/tirx/tile/cuda/test_ewise.py
@@ -95,7 +95,7 @@ def _shared_case(input, operation, src_dtype, dst_dtype,
bias_type=None, target=
def allocate_output(source):
if in_place:
return source
- return T.alloc_shared(shape, dst_dtype, layout=layout)
+ return T.bind_(T.alloc_shared(shape, dst_dtype, layout=layout))
@T.function
def kernel(
diff --git a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma.py
b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma.py
index 0355aac865..61648c0e88 100644
--- a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma.py
+++ b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma.py
@@ -430,7 +430,7 @@ def test_gemm_tcgen05_cta_group_1_layout_f_m64():
def _mapa(ptr, rank):
"""`mapa.u64` into a declared register, returned as a value."""
- mapped = T.alloc_local([1], "uint64", layout="default")
+ mapped = T.bind_(T.alloc_local([1], "uint64", layout="default"))
T.evaluate(T.ptx.mapa.u64(mapped[0], ptr, T.uint32(rank)))
return mapped[0]
diff --git a/tirx/tests/python/tirx/transform/test_common_subexpr_elim.py
b/tirx/tests/python/tirx/transform/test_common_subexpr_elim.py
index ecf25aef86..3d33e08a51 100644
--- a/tirx/tests/python/tirx/transform/test_common_subexpr_elim.py
+++ b/tirx/tests/python/tirx/transform/test_common_subexpr_elim.py
@@ -27,28 +27,28 @@ def test_basic():
class Before:
@T.function
def main(B: T.Tensor((50,), "int32"), i1: T.int32, i2: T.int32, z3:
T.int32):
- z1 = T.bind(1)
- z2 = T.bind(2)
+ z1 = 1
+ z2 = 2
B[i1] = z1 + z2
- x = T.bind(1)
- y = T.bind(1)
- a = T.bind((x + y) + (z1 + z2))
- b = T.bind((x + y) + z3)
+ x = 1
+ y = 1
+ a = (x + y) + (z1 + z2)
+ b = (x + y) + z3
B[i2] = a + b
@tvm_tirx.script.ir_module
class Expected:
@T.function
def main(B: T.Tensor((50,), "int32"), i1: T.int32, i2: T.int32, z3:
T.int32):
- z1 = T.bind(1)
- z2 = T.bind(2)
- cse_v1 = T.bind(z1 + z2)
+ z1 = 1
+ z2 = 2
+ cse_v1 = z1 + z2
B[i1] = cse_v1
- x = T.bind(1)
- y = T.bind(1)
- cse_v2 = T.bind(x + y)
- a = T.bind(cse_v2 + cse_v1)
- b = T.bind(cse_v2 + z3)
+ x = 1
+ y = 1
+ cse_v2 = x + y
+ a = cse_v2 + cse_v1
+ b = cse_v2 + z3
B[i2] = a + b
after = tvm_tirx.tirx.transform.CommonSubexprElim()(Before)
@@ -67,7 +67,7 @@ def test_if_single_branch():
y: T.int32,
z: T.int32,
):
- b = T.bind(1)
+ b = 1
if b:
B[i1] = y + z
B[i2] = y + z
@@ -85,9 +85,9 @@ def test_if_single_branch():
y: T.int32,
z: T.int32,
):
- b = T.bind(1)
+ b = 1
if b:
- cse_v1 = T.bind(y + z)
+ cse_v1 = y + z
B[i1] = cse_v1
B[i2] = cse_v1
else:
@@ -109,7 +109,7 @@ def test_if_both_branches():
y: T.int32,
z: T.int32,
):
- b = T.bind(1)
+ b = 1
if b:
B[i1] = y + z
B[i2] = y
@@ -127,8 +127,8 @@ def test_if_both_branches():
y: T.int32,
z: T.int32,
):
- b = T.bind(1)
- cse_v1 = T.bind(y + z)
+ b = 1
+ cse_v1 = y + z
if b:
B[i1] = cse_v1
B[i2] = y
@@ -168,8 +168,8 @@ def test_cascade():
y: T.int32,
z: T.int32,
):
- cse_v2 = T.bind(x + y)
- cse_v1 = T.bind(cse_v2 + z)
+ cse_v2 = x + y
+ cse_v1 = cse_v2 + z
B[i1] = cse_v1
B[i2] = cse_v1
B[i3] = cse_v2
@@ -183,7 +183,7 @@ def test_no_duplication():
class Before:
@T.function
def main(x: T.int32, y: T.int32, z: T.int32):
- a = T.bind(x + (y + z))
+ a = x + (y + z)
T.evaluate(a)
Expected = Before
@@ -206,9 +206,9 @@ def test_deterministic():
class Expected:
@T.function
def main(B: T.Tensor((4,), "int32"), x: T.int32):
- cse_v1 = T.bind(x + 1)
+ cse_v1 = x + 1
B[0] = cse_v1
- cse_v2 = T.bind(x + 2)
+ cse_v2 = x + 2
B[1] = cse_v2
B[2] = cse_v1
B[3] = cse_v2
@@ -235,7 +235,7 @@ def test_for_loop():
@T.function
def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
for i in range(10):
- cse_v1 = T.bind(y + z)
+ cse_v1 = y + z
B[i] = cse_v1
B[i + 10] = cse_v1
@@ -256,7 +256,7 @@ def test_for_hoist():
class Expected:
@T.function
def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- cse_v1 = T.bind(y + z)
+ cse_v1 = y + z
B[0] = cse_v1
for i in range(10):
B[i + 1] = cse_v1
@@ -309,7 +309,7 @@ def test_nested_if():
z: T.int32,
):
if c1:
- cse_v1 = T.bind(y + z)
+ cse_v1 = y + z
if c2:
B[0] = cse_v1
else:
@@ -347,9 +347,9 @@ def test_multi_independent():
c: T.int32,
d: T.int32,
):
- cse_v1 = T.bind(a + b)
+ cse_v1 = a + b
B[0] = cse_v1
- cse_v2 = T.bind(c + d)
+ cse_v2 = c + d
B[1] = cse_v2
B[2] = cse_v1
B[3] = cse_v2
@@ -370,7 +370,7 @@ def test_if_condition():
class Expected:
@T.function
def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- cse_v1 = T.bind(y + z)
+ cse_v1 = y + z
if cse_v1 > 0:
B[0] = cse_v1
@@ -414,7 +414,7 @@ def test_no_single_use_binding():
y: T.int32,
z: T.int32,
):
- cse_v1 = T.bind((x + y) + z)
+ cse_v1 = (x + y) + z
B[0] = cse_v1
B[1] = cse_v1
@@ -434,7 +434,7 @@ def test_for_extent_lift():
class Expected:
@T.function
def main(B: T.Tensor((50,), "int32"), y: T.int32, z: T.int32):
- cse_v1 = T.bind(y + z)
+ cse_v1 = y + z
for i in range(cse_v1):
B[i] = cse_v1
@@ -461,7 +461,7 @@ def test_loop_var_expr_stays_inside():
B: T.Tensor((50,), "int32"),
):
for i in range(10):
- cse_v1 = T.bind(i * 4)
+ cse_v1 = i * 4
A[cse_v1] = B[cse_v1]
after = tvm_tirx.tirx.transform.CommonSubexprElim()(Before)
@@ -497,7 +497,7 @@ def test_let_value_cse():
@T.function
def main(B: T.Tensor((10,), "int32"), y: T.int32, z: T.int32):
for i in range(10):
- cse_v1 = T.bind(y + z)
+ cse_v1 = y + z
B[i] = cse_v1 + T.Let(x + 1, where={x: cse_v1})
After = tvm_tirx.tirx.transform.CommonSubexprElim()(Before)
@@ -598,15 +598,15 @@ def test_shared_subtree_stays_ssa():
class Expected:
@T.function
def main(B: T.Tensor((50,), "int32"), i1: T.int32, i2: T.int32):
- value = T.bind(i1 + i2)
- twice = T.bind(value * 2)
+ value = i1 + i2
+ twice = value * 2
B[i1] = twice
- thrice = T.bind(value * 3)
+ thrice = value * 3
B[i2] = thrice
- value2 = T.bind(i1 + i2)
- twice2 = T.bind(value2 * 2)
+ value2 = i1 + i2
+ twice2 = value2 * 2
B[i1] = twice2
- thrice2 = T.bind(value * 3)
+ thrice2 = value * 3
B[i2] = thrice2
unrolled = tvm_tirx.tirx.transform.UnrollLoop()(Before)
diff --git a/tirx/tests/python/tirx/transform/test_convert_ssa.py
b/tirx/tests/python/tirx/transform/test_convert_ssa.py
index 632adbb8ad..0ef8c45acf 100644
--- a/tirx/tests/python/tirx/transform/test_convert_ssa.py
+++ b/tirx/tests/python/tirx/transform/test_convert_ssa.py
@@ -32,18 +32,18 @@ def test_reuse_in_sequential_bind():
class Before:
@T.function(private=True)
def main():
- T.bind(16, var=var)
+ T.Bind(var, 16)
T.evaluate(var)
- T.bind(32, var=var)
+ T.Bind(var, 32)
T.evaluate(var)
assert Before["main"].body[0].var.same_as(Before["main"].body[2].var)
@T.function(private=True)
def expected():
- var1 = T.bind(T.int32(16))
+ var1 = T.int32(16)
T.evaluate(var1)
- var2 = T.bind(T.int32(32))
+ var2 = T.int32(32)
T.evaluate(var2)
After = tvm_tirx.tirx.transform.ConvertSSA()(Before)
@@ -58,9 +58,9 @@ def test_reuse_in_nested_bind():
class Before:
@T.function(private=True)
def main():
- T.bind(32, var=var)
+ T.Bind(var, 32)
T.evaluate(var)
- T.bind(16, var=var)
+ T.Bind(var, 16)
T.evaluate(var)
T.evaluate(var)
@@ -68,9 +68,9 @@ def test_reuse_in_nested_bind():
class Expected:
@T.function(private=True)
def main():
- var1 = T.bind(32)
+ var1 = 32
T.evaluate(var1)
- var2 = T.bind(16)
+ var2 = 16
T.evaluate(var2)
T.evaluate(var2)
@@ -84,7 +84,7 @@ def test_reused_var_across_module():
@T.function(private=True)
def func():
- var = T.bind(10)
+ var = 10
T.evaluate(var)
@I.ir_module(check_well_formed=False)
diff --git a/tirx/tests/python/tirx/transform/test_function_pass.py
b/tirx/tests/python/tirx/transform/test_function_pass.py
index 6a14f8c483..c2fffb88ba 100644
--- a/tirx/tests/python/tirx/transform/test_function_pass.py
+++ b/tirx/tests/python/tirx/transform/test_function_pass.py
@@ -26,7 +26,7 @@ def test_function_pass(capsys):
class Before:
@T.function(private=True)
def main(x: T.int32, y: T.int32, B: T.Tensor((16,), "float32")):
- value = T.bind(10)
+ value = 10
T.evaluate(value + 1)
@I.ir_module
diff --git a/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
b/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
index edae515ca5..91b4708449 100644
--- a/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
+++ b/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
@@ -1275,15 +1275,17 @@ def test_lower_alloc_decl_tensor_outside_of_parser():
@T.meta_class
class State:
def __init__(self, smem):
- self.A = T.alloc_local([1], "float16")
- self.B = T.alloc_local([1], "float16")
- self.C = T.bind(
- T.decl_tensor(smem, ty_args=[T.Tensor([1], "float16",
scope="shared.dyn")])
+ self.A = T.bind_(T.alloc_local([1], "float16"))
+ self.B = T.bind_(T.alloc_local([1], "float16"))
+ declaration = T.decl_tensor(
+ smem, ty_args=[T.Tensor([1], "float16", scope="shared.dyn")]
)
+ self.C = tvm_tirx.ir.Var("C", declaration.ty)
+ T.emit_(T.Bind(self.C, declaration))
@register_mutable_decl("TestMutableCells.int_var1")
def int_var1(val):
- buf = T.local_scalar("int32")
+ buf = T.bind_(T.local_scalar("int32"))
if val is not None:
T.tensor_store(buf.source, 0, val)
return buf
@@ -1292,7 +1294,7 @@ def test_lower_alloc_decl_tensor_outside_of_parser():
register_namespace("TestMutableCells", TestMutableCells)
def int_var2(val):
- buf = T.alloc_local([1], "int32")
+ buf = T.bind_(T.alloc_local([1], "int32"))
if val is not None:
T.tensor_store(buf, 0, val)
return buf
