This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 05c487d69a [FIX][TIRx] Preserve pointer expression types (#20070)
05c487d69a is described below
commit 05c487d69a959bfba2efed67751acb0fc81a2815
Author: Hongyi Jin <[email protected]>
AuthorDate: Wed Jul 29 14:01:14 2026 -0400
[FIX][TIRx] Preserve pointer expression types (#20070)
## Motivation and context
A TIRx pointer carries two pieces of information that later lowering
needs: the pointee element type and the storage scope. Both must survive
when a pointer-producing expression is assigned to a Python name and
then used as the backing storage of a buffer.
A concrete example is accessing an mbarrier in another CTA through
distributed shared memory:
```python
ptr_ty = PointerType(PrimType("uint64"), "shared")
remote_ptr = T.reinterpret(
ptr_ty,
T.ptx.map_shared_rank(mbar.ptr_to([0]), T.int32(0)),
)
remote_mbar = T.decl_buffer(
[1], "uint64", data=remote_ptr, scope="shared"
)
```
`map_shared_rank` returns the raw `uint64` address produced by PTX
`mapa`, and `reinterpret` gives that address the intended
`PointerType(uint64, shared)`. Because `decl_buffer(data=...)` requires
a pointer `Var`, assigning the expression to `remote_ptr` should create
an immutable typed pointer binding.
Before this PR, an unannotated assignment such as `remote_ptr = <pointer
expression>` followed the same parser path as a numeric assignment. That
path allocates a mutable local scalar and therefore cannot represent a
`PointerType`. The pointer expression could not be carried as a
correctly typed `Var` into `decl_buffer` and CUDA lowering.
This PR makes an unannotated pointer-valued assignment emit a TIRx
`Bind`. The bound `Var` has exactly the type of the right-hand side,
including its element type and storage scope. Pointer bindings are
immutable, so reassignment in the same scope is diagnosed; shadowing a
name supplied through `extra_vars` remains valid. Numeric assignments
keep their existing mutable-local behavior.
## Type propagation fixes
The parser fix exposed several other boundaries where pointer type
information must remain consistent:
| Boundary | Previous behavior | Behavior after this PR |
| --- | --- | --- |
| Unannotated pointer assignment | Tried to materialize the value as a
local scalar | Emits an immutable `Bind` with the RHS `PointerType` |
| `address_of(buffer)` / `buffer.ptr_to(...)` | Reused the raw backing
pointer type | Returns a pointer to `buffer.dtype` while preserving the
backing pointer storage scope |
| `tvm_access_ptr` / `ptr_byte_offset` | Accepted strings or annotation
expressions, but not a `PrimType` object directly | Accepts `PrimType`
and produces the corresponding typed pointer |
| Printed `T.ptx.mapa` call | The printer emits all intrinsic attributes
positionally, but the Python helper required keyword-only arguments |
Accepts the canonical printed form so pointer code round-trips through
TVMScript |
The `address_of` distinction matters for typed views over byte-addressed
storage. For example, if a `float32` buffer is backed by a `uint8*`
allocation in `shared.dyn`, the address of a buffer element must be
`PointerType(float32, shared.dyn)`, not `PointerType(uint8,
shared.dyn)`.
With these changes, the DSMEM example above round-trips through
TVMScript and CUDA codegen declares the remote buffer pointer as
`uint64_t*`.
## TMA dtype normalization
This PR also contains a small, separate type-representation fix in TMA
descriptor construction. `TmaPlan.elem_dtype` is a string consumed by
the host-side `runtime.cuTensorMapEncodeTiled` packed call, but
`_assemble_plan` stored `g_buf.dtype`, which is a `PrimType`. Converting
it with `str(g_buf.dtype)` ensures that the generated packed-call
argument is `StringImm("float16")` rather than an IR type object. This
does not change the pointer-binding semantics described above.
## Testing
- Verify that an unannotated pointer expression creates a `Bind` whose
`Var` type matches the RHS type.
- Verify that pointer reassignment is rejected while shadowing an
`extra_vars` name is allowed.
- Verify parser/printer structural round-tripping for the pointer
binding and canonical `T.ptx.mapa` call.
- Verify that `address_of` uses the logical buffer element type and
preserves the storage scope for byte-backed buffer views.
- Verify that `tvm_access_ptr` and `ptr_byte_offset` accept `PrimType`
inputs.
- Compile the DSMEM `map_shared_rank` example through the CUDA TIRx
pipeline and check for a typed `uint64_t*` remote buffer pointer.
- Verify that the TMA host initialization passes the descriptor dtype as
a `StringImm`.
---
docs/tirx/native_basics/cuda/buffers.rst | 10 +++-
docs/tirx/native_basics/cuda/data_types.rst | 15 ++++--
python/tvm/backend/cuda/op.py | 2 +-
.../cuda/operator/tile_primitive/copy_async/tma.py | 2 +-
python/tvm/script/parser/core/parser.py | 4 ++
python/tvm/tirx/op.py | 37 +++++++++++----
python/tvm/tirx/script/parser/parser.py | 11 ++++-
tests/python/tirx-base/test_tir_op_types.py | 10 ++++
tests/python/tirx/codegen/test_codegen_dsmem.py | 42 +++++++++++++++++
.../tile_primitive/cuda/copy_async/test_tma.py | 14 ++++++
tests/python/tirx/test_parser_printer.py | 55 ++++++++++++++++++++++
11 files changed, 183 insertions(+), 19 deletions(-)
diff --git a/docs/tirx/native_basics/cuda/buffers.rst
b/docs/tirx/native_basics/cuda/buffers.rst
index feb2763f62..18de83ea03 100644
--- a/docs/tirx/native_basics/cuda/buffers.rst
+++ b/docs/tirx/native_basics/cuda/buffers.rst
@@ -38,8 +38,9 @@ Two fundamental APIs create a buffer:
it allocates, like ``alloc_buffer``.
A buffer's ``data`` pointer is an immutable ``Var`` (``alloc_buffer`` defines
it;
-``decl_buffer`` takes one). To back a buffer with a pointer *expression*, bind
it
-first — see :doc:`data_types`.
+``decl_buffer`` takes one). To back a buffer with a pointer *expression*,
assign
+it to a name first; the parser creates an immutable pointer binding. See
+:doc:`data_types`.
Both share one descriptor; the parameters that matter most:
@@ -439,6 +440,11 @@ to an intrinsic or inline function; ``data`` is the base
pointer:
B_ptr[tx] = ld(&A_ptr[tx]); // ptr_to([tx]) -> &A_ptr[tx];
A.data -> A_ptr
+The pointer returned by ``ptr_to`` has the buffer's element type and storage
+scope. This remains true when the buffer is a typed view over a byte-addressed
+allocation pool; the pool's raw backing-pointer type does not leak through the
+element address.
+
**Vectorized access — ``vload`` / ``vstore``.** Move several elements as one
wide
transfer (see also :doc:`data_types`):
diff --git a/docs/tirx/native_basics/cuda/data_types.rst
b/docs/tirx/native_basics/cuda/data_types.rst
index 4b17543566..884cfdb07d 100644
--- a/docs/tirx/native_basics/cuda/data_types.rst
+++ b/docs/tirx/native_basics/cuda/data_types.rst
@@ -103,14 +103,19 @@ A buffer's ``data`` — its pointer — is a ``Var`` of
pointer type, and it is
- ``T.decl_buffer(..., data=ptr)`` declares a buffer over an existing pointer
``Var`` ``ptr``.
- To back a buffer with a pointer **expression** — e.g.
``T.ptx.map_shared_rank``
- (PTX ``mapa``) giving another cluster CTA's shared address — you must first
bind
- that expression to a pointer ``Var`` (``data`` must be a ``Var``, not an
- expression), using a ``T.let`` of ``PointerType``:
+ (PTX ``mapa``) giving another cluster CTA's shared address — convert the raw
+ ``uint64`` address returned by ``map_shared_rank`` to a pointer with the
+ intended element type and storage scope. Assigning that pointer expression
+ to an unannotated name automatically creates an immutable ``Bind`` (and
+ ``data`` receives the resulting pointer ``Var``):
.. code-block:: python
from tvm.ir.type import PointerType, PrimType
- ptr: T.let[T.Var(name="ptr", ty=PointerType(PrimType("uint64")))] = \
- T.reinterpret("handle", T.ptx.map_shared_rank(mbar.ptr_to([0]), 0))
+ ptr_ty = PointerType(PrimType("uint64"), "shared")
+ ptr = T.reinterpret(ptr_ty, T.ptx.map_shared_rank(mbar.ptr_to([0]), 0))
remote_mbar = T.decl_buffer([1], "uint64", data=ptr, scope="shared")
+
+ Pointer bindings cannot be reassigned; use a new name for a different
+ pointer value.
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index 13bf481543..988e48f4d0 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -4209,7 +4209,7 @@ def ptx_map_shared_rank(ptr, rank):
return ptx_mapa(ptr, rank, space="", ptx_type="u64", return_type="uint64")
-def ptx_mapa(ptr, rank, *, space="", ptx_type="u64", return_type="uint64"):
+def ptx_mapa(ptr, rank, space="", ptx_type="u64", return_type="uint64"):
"""TVM intrinsic for PTX ``mapa{.space}.type d, a, b``."""
if space not in ("", "shared::cluster"):
raise ValueError(f"Unsupported mapa space {space!r}")
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
index 8f78679de2..e83cd827c2 100644
--- a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
@@ -812,7 +812,7 @@ def _assemble_plan(
issue_axes=issue_axes,
tensor_ptr=g_buf.data,
elem_bytes=elem_bytes,
- elem_dtype=g_buf.dtype,
+ elem_dtype=str(g_buf.dtype),
)
return _merge_contig_full_box_dims(plan, analyzer)
diff --git a/python/tvm/script/parser/core/parser.py
b/python/tvm/script/parser/core/parser.py
index d402654d6e..ff09ee1177 100644
--- a/python/tvm/script/parser/core/parser.py
+++ b/python/tvm/script/parser/core/parser.py
@@ -289,6 +289,10 @@ class VarTable:
"""
return {key: values[-1] for key, values in self.name2value.items() if
values}
+ def contains_in_current_frame(self, name: str) -> bool:
+ """Check whether a variable name exists in the current frame."""
+ return bool(self.frames) and name in self.frames[-1].vars
+
def get_at_depth(self, depth: int) -> dict[str, Any]:
"""Get variables visible at the given frame depth, using current
values.
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 782fbcfe58..35fc990720 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -59,6 +59,8 @@ def _canonical_device_intrin_name(func_name: str) -> str:
def _primexpr_ty(expr):
"""Return the runtime primitive type of an expression."""
+ if isinstance(expr, tvm.ir.PrimType):
+ return expr
ty = getattr(expr, "ty", None)
if isinstance(ty, tvm.ir.PrimType):
return ty
@@ -658,6 +660,13 @@ def _is_tensormap_var(obj: Var) -> bool:
return isinstance(obj.ty, PointerType) and isinstance(obj.ty.element_type,
TensorMapType)
+def _buffer_element_pointer_type(buffer: Buffer) -> PointerType:
+ """Return a pointer to ``buffer`` elements in the buffer's storage
scope."""
+ data_ty = buffer.data.ty
+ storage_scope = data_ty.storage_scope if isinstance(data_ty, PointerType)
else "global"
+ return PointerType(buffer.dtype, storage_scope)
+
+
def address_of(obj: Buffer | BufferLoad | Var, span: Span | None = None) ->
Expr:
"""Returns the address of a buffer element or addressable variable.
@@ -677,7 +686,12 @@ def address_of(obj: Buffer | BufferLoad | Var, span: Span
| None = None) -> Expr
if isinstance(obj, Buffer):
n_dim = len(obj.shape)
buffer_load = BufferLoad(obj, [0] * n_dim)
- return Call("tirx.address_of", [buffer_load], span=span,
ret_ty=obj.data.ty)
+ return Call(
+ "tirx.address_of",
+ [buffer_load],
+ span=span,
+ ret_ty=_buffer_element_pointer_type(obj),
+ )
elif isinstance(obj, Var):
if _is_tensormap_var(obj):
return call_intrin("uint64", "tirx.address_of", obj, span=span)
@@ -685,7 +699,12 @@ def address_of(obj: Buffer | BufferLoad | Var, span: Span
| None = None) -> Expr
raise TypeError(f"address_of expects a scalar or TensorMap Var,
but got {obj.ty}")
return Call("tirx.address_of", [obj], span=span,
ret_ty=PointerType(obj.ty))
elif isinstance(obj, BufferLoad):
- return Call("tirx.address_of", [obj], span=span,
ret_ty=obj.buffer.data.ty)
+ return Call(
+ "tirx.address_of",
+ [obj],
+ span=span,
+ ret_ty=_buffer_element_pointer_type(obj.buffer),
+ )
else:
raise ValueError(f"Invalid object type: {type(obj)}")
@@ -918,11 +937,11 @@ def tvm_access_ptr(ptype, data, offset, extent, rw_mask):
Parameters
----------
- ptype : Expr or str
- The data type of pointer. If a ``str``, it is wrapped via
- :func:`type_annotation` so that the lowering rule (which reads
- ``args[0].dtype()`` for the cast type) sees the intended dtype
- instead of ``void`` from a raw StringImm.
+ ptype : Expr, PrimType, or str
+ The data type of pointer. If a ``PrimType`` or ``str``, it is wrapped
+ via :func:`type_annotation` so that the lowering rule (which reads
+ ``args[0].dtype()`` for the cast type) sees the intended dtype instead
+ of ``void`` from a raw StringImm.
data : DType*
The data of pointer.
@@ -941,7 +960,7 @@ def tvm_access_ptr(ptype, data, offset, extent, rw_mask):
call : Expr
The call expression.
"""
- if isinstance(ptype, str):
+ if isinstance(ptype, str | tvm.ir.PrimType):
ptype = type_annotation(ptype)
data_type = getattr(data, "ty", None)
storage_scope = data_type.storage_scope if isinstance(data_type,
PointerType) else "global"
@@ -962,7 +981,7 @@ def ptr_byte_offset(data, byte_offset, dtype):
``byte_offset`` is always in bytes. Use this when the source CUDA shape
needs an explicitly typed local pointer derived from a byte-addressed base.
"""
- if isinstance(dtype, str):
+ if isinstance(dtype, str | tvm.ir.PrimType):
dtype = type_annotation(dtype)
data_type = getattr(data, "ty", None)
storage_scope = data_type.storage_scope if isinstance(data_type,
PointerType) else "global"
diff --git a/python/tvm/tirx/script/parser/parser.py
b/python/tvm/tirx/script/parser/parser.py
index ddbfac29ab..40520bdb63 100644
--- a/python/tvm/tirx/script/parser/parser.py
+++ b/python/tvm/tirx/script/parser/parser.py
@@ -23,7 +23,7 @@ from functools import partial
from typing import Any
import tvm
-from tvm.ir import Expr, GlobalVar, PrimType
+from tvm.ir import Expr, GlobalVar, PointerType, PrimType
from tvm.script.ir_builder import ir as I
from tvm.script.ir_builder.base import IRBuilder
from tvm.script.ir_builder.base import IRBuilderFrame as Frame
@@ -217,6 +217,15 @@ def bind_assign_value(self: Parser, node: doc.expr,
var_name: str, value: Any) -
IRBuilder.name(var_name, value)
return value
else:
+ is_pointer_expr = isinstance(value, tvm.ir.Expr) and isinstance(
+ getattr(value, "ty", None), PointerType
+ )
+ if is_pointer_expr:
+ if self.var_table.contains_in_current_frame(var_name):
+ self.report_error(node, f"Pointer variable '{var_name}' cannot
be reassigned")
+ ann_var = T.Bind(value)
+ IRBuilder.name(var_name, ann_var)
+ return ann_var
if not tvm.ir.is_prim_expr(value):
value = tvm.tirx.const(value)
if not isinstance(value, tvm.tirx.StringImm):
diff --git a/tests/python/tirx-base/test_tir_op_types.py
b/tests/python/tirx-base/test_tir_op_types.py
index 78e379b765..e0e7bf1f42 100644
--- a/tests/python/tirx-base/test_tir_op_types.py
+++ b/tests/python/tirx-base/test_tir_op_types.py
@@ -47,6 +47,11 @@ def test_tir_op_address_of():
buffer = tirx.decl_buffer((128), "float32")
expr = tirx.address_of(buffer[0])
assert expr.op.name == "tirx.address_of"
+ storage = tirx.Var("storage", tvm.ir.PointerType(tvm.ir.PrimType("uint8"),
"shared.dyn"))
+ pooled_buffer = tirx.decl_buffer((128), "float32", data=storage,
scope="shared.dyn")
+ expected_ty = tvm.ir.PointerType(tvm.ir.PrimType("float32"), "shared.dyn")
+ assert tirx.address_of(pooled_buffer).ty == expected_ty
+ assert tirx.address_of(pooled_buffer[0]).ty == expected_ty
scalar_address = tirx.address_of(tirx.Var("value", "uint32"))
assert scalar_address.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint32"))
@@ -116,6 +121,11 @@ def test_tir_op_tvm_access_ptr():
assert expr.ty == tvm.ir.PointerType(tvm.ir.PrimType("float32"))
offset_expr = tirx.ptr_byte_offset(buffer.data, 16, "uint8")
assert offset_expr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8"))
+ prim_type = tvm.ir.PrimType("uint8")
+ typed_access_expr = tirx.tvm_access_ptr(prim_type, buffer.data, 0, 1, 2)
+ assert typed_access_expr.ty == tvm.ir.PointerType(prim_type)
+ typed_offset_expr = tirx.ptr_byte_offset(buffer.data, 16, prim_type)
+ assert typed_offset_expr.ty == tvm.ir.PointerType(prim_type)
def test_tir_op_tvm_throw_last_error():
diff --git a/tests/python/tirx/codegen/test_codegen_dsmem.py
b/tests/python/tirx/codegen/test_codegen_dsmem.py
index 052034dd0d..1a960d724e 100644
--- a/tests/python/tirx/codegen/test_codegen_dsmem.py
+++ b/tests/python/tirx/codegen/test_codegen_dsmem.py
@@ -19,6 +19,7 @@
import tvm
import tvm.testing
+from tvm.ir import PointerType, PrimType, assert_structural_equal
from tvm.script import tirx as T
@@ -87,7 +88,48 @@ def test_ptx_cp_async_bulk_s2c_codegen_address_conversion():
assert
"cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" in src
+def test_ptx_map_shared_rank_pointer_bind_codegen():
+ ptr_ty = PointerType(PrimType("uint64"), "shared")
+
+ # fmt: off
+ @T.prim_func
+ def main(A: T.Buffer((1,), "uint64")):
+ T.device_entry()
+ cta_id = T.cta_id([1])
+ tid = T.thread_id([1])
+ mbar = T.alloc_shared([2], "uint64")
+ remote_ptr = T.reinterpret(
+ ptr_ty, T.ptx.map_shared_rank(mbar.ptr_to([0]), T.int32(0))
+ )
+ remote_mbar = T.decl_buffer([1], "uint64", data=remote_ptr,
scope="shared")
+ A[0] = remote_mbar[0]
+ # fmt: on
+
+ binds = []
+ loads = []
+
+ def collect(node):
+ if isinstance(node, tvm.tirx.Bind):
+ binds.append(node)
+ elif isinstance(node, tvm.tirx.BufferLoad):
+ loads.append(node)
+
+ tvm.tirx.stmt_functor.post_order_visit(main.body, collect)
+ assert len(binds) == 1
+ assert isinstance(binds[0].var.ty, PointerType)
+ assert binds[0].var.ty.storage_scope == "shared"
+ assert binds[0].value.ty.storage_scope == "shared"
+ assert_structural_equal(binds[0].var.ty, binds[0].value.ty)
+ assert any(load.buffer.data.same_as(binds[0].var) for load in loads)
+
+ assert_structural_equal(main, tvm.script.from_source(main.script()))
+ src = _get_source(main)
+ assert "uint64_t* remote_mbar_ptr" in src
+ assert "tvm_builtin_ptx_mapa_u64" in src
+
+
if __name__ == "__main__":
test_ptx_cp_async_bulk_s2c_codegen()
test_ptx_cp_async_bulk_s2c_codegen_address_conversion()
+ test_ptx_map_shared_rank_pointer_bind_codegen()
print("All codegen tests passed!")
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
index 7cbce484ba..0189493049 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
@@ -1046,6 +1046,20 @@ def test_copy_tma_codegen(case):
tvm.ir.assert_structural_equal(host_init_stmts[0], expected_host,
map_free_vars=True)
+def test_copy_tma_host_init_dtype_is_string():
+ _, host_init_stmts = _make_tma_call(
+ g_shape=(8, 256),
+ g_region=((0, 8), (0, 256)),
+ s_shape=(8, 256),
+ s_region=((0, 8), (0, 256)),
+ gmem_layout=TileLayout(S[8, 256]),
+ smem_layout=TileLayout(S[8, 256]),
+ )
+ encode_call = host_init_stmts[0].seq[1].value
+ assert isinstance(encode_call.args[2], StringImm)
+ assert encode_call.args[2].value == "float16"
+
+
# Section 3: TMA special cases (symbolic dimension, buffer view)
# ===========================================================================
diff --git a/tests/python/tirx/test_parser_printer.py
b/tests/python/tirx/test_parser_printer.py
index e570f72200..6c5627f36b 100644
--- a/tests/python/tirx/test_parser_printer.py
+++ b/tests/python/tirx/test_parser_printer.py
@@ -1376,6 +1376,61 @@ def test_buffer_local_ir():
assert from_source(code).script() == code
+def test_pointer_expression_assignment_uses_bind():
+ # fmt: off
+ @T.prim_func
+ def func() -> None:
+ T.device_entry()
+ buf = T.alloc_buffer((4,), "uint32", scope="shared")
+ ptr = buf.ptr_to([1])
+ T.evaluate(T.reinterpret("uint64", ptr))
+ # fmt: on
+
+ binds = []
+ tvm.tirx.stmt_functor.post_order_visit(
+ func.body, lambda node: binds.append(node) if isinstance(node,
tvm.tirx.Bind) else None
+ )
+ assert len(binds) == 1
+ assert isinstance(binds[0].var.ty, PointerType)
+ assert_structural_equal(binds[0].var.ty, binds[0].value.ty)
+
+ code = func.script()
+ assert_structural_equal(func, from_source(code))
+
+
+def test_pointer_expression_assignment_rejects_reassignment():
+ with pytest.raises(tvm.error.DiagnosticError, match="cannot be
reassigned"):
+ # fmt: off
+ @T.prim_func
+ def func() -> None:
+ T.device_entry()
+ buf = T.alloc_buffer((4,), "uint32", scope="shared")
+ ptr = buf.ptr_to([0])
+ ptr = buf.ptr_to([1])
+ T.evaluate(T.reinterpret("uint64", ptr))
+ # fmt: on
+
+
+def test_pointer_expression_assignment_can_shadow_extra_var():
+ source = """
[email protected]_func
+def func() -> None:
+ T.device_entry()
+ buf = T.alloc_buffer((4,), "uint32", scope="shared")
+ ptr = buf.ptr_to([1])
+ view = T.decl_buffer((3,), "uint32", data=ptr, scope="shared")
+ view[0] = T.uint32(0)
+"""
+ func = tvm.script.from_source(source, extra_vars={"T": T, "ptr": object()})
+
+ binds = []
+ tvm.tirx.stmt_functor.post_order_visit(
+ func.body, lambda node: binds.append(node) if isinstance(node,
tvm.tirx.Bind) else None
+ )
+ assert len(binds) == 1
+ assert_structural_equal(func, from_source(func.script()))
+
+
def test_buffer_permute_ir():
"""Verify .permute(1, 0): shape swapped, layout permuted, shared data."""
