This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 758553133c5824e6ef0c264134b8de6b547c9f8e Author: Tianqi Chen <[email protected]> AuthorDate: Mon Oct 19 17:12:37 2026 -0400 [REFACTOR][TIRx] Consolidate bindings and emitted allocation results (#20799) --- 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 +- 32 files changed, 407 insertions(+), 326 deletions(-) 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
