(tvm) 09/13: [FIX][TIR] Cast opaque workspace allocations

[tq...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch tvm-further-cleanup-python-tests-followup
in repository https://gitbox.apache.org/repos/asf/tvm.git

commit d310cbb34767cd79c1c59a4f627ab0c6fa1a111f
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Sun Jul 5 18:54:32 2026 +0000

    [FIX][TIR] Cast opaque workspace allocations
---
 src/tirx/transform/lower_tvm_builtin.cc            |  9 +++----
 tests/python/codegen/test_target_codegen_c_host.py | 16 +++++++++++++
 .../test_tir_transform_lower_tvm_builtin.py        | 28 ++++++++++++++++++++++
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/src/tirx/transform/lower_tvm_builtin.cc b/src/tirx/transform/lower_tvm_builtin.cc
index 55790e51a4..8f940fbd4d 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -294,12 +294,13 @@ class BuiltinLower : public StmtExprMutator {
     // Push free to enclosing scope's pending_frees (LIFO ordering preserved).
     scope_.Current().pending_frees.push_back(free_stmt);
 
-    Stmt alloc_bind = Bind(
-        op->buffer->data,
-        Call(op->buffer->data->ty, alloc_workspace_op,
+    Expr opaque_alloc =
+        Call(PointerType::VoidPointerTy(), alloc_workspace_op,
              {cast(PrimType::Int(32), device_type_.value()),
               cast(PrimType::Int(32), device_id_.value()), total_bytes,
-              IntImm::Int32(op->buffer->dtype.code()), IntImm::Int32(op->buffer->dtype.bits())}));
+              IntImm::Int32(op->buffer->dtype.code()), IntImm::Int32(op->buffer->dtype.bits())});
+    Stmt alloc_bind =
+        Bind(op->buffer->data, reinterpret(op->buffer->data->ty, std::move(opaque_alloc)));
 
     return SeqStmt({alloc_bind, alloc_nullptr_check});
   }
diff --git a/tests/python/codegen/test_target_codegen_c_host.py b/tests/python/codegen/test_target_codegen_c_host.py
index 5dac50d48e..5ccf6aba17 100644
--- a/tests/python/codegen/test_target_codegen_c_host.py
+++ b/tests/python/codegen/test_target_codegen_c_host.py
@@ -227,5 +227,21 @@ def test_subroutine_call():
     )
 
 
+def test_workspace_allocation_cast():
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((16,), "float32")):
+            workspace = T.alloc_buffer((16,), "float32", scope="local")
+            workspace[0] = A[0]
+            A[0] = workspace[0]
+
+    built = tvm.tirx.build(Module, target="c")
+    assert "((float*)TVMBackendAllocWorkspace(" in built.inspect_source()
+
+    temp = utils.tempdir()
+    built.export_library(temp.relpath("workspace.so"))
+
+
 if __name__ == "__main__":
     tvm.testing.main()
diff --git a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
index aeab2eb9ed..8f1cf4afcf 100644
--- a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
+++ b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
@@ -234,6 +234,34 @@ def test_lower_device_allocate():
     # DeclBuffer should appear as a flat statement
     assert "T.decl_buffer" in script_output
 
+    workspace_binds = []
+
+    def collect_workspace_bind(node):
+        if not isinstance(node, tvm.tirx.Bind):
+            return
+        if not isinstance(node.value, tvm.ir.Call):
+            return
+        if node.value.op != tvm.ir.Op.get("tirx.reinterpret"):
+            return
+        opaque_alloc = node.value.args[0]
+        if isinstance(opaque_alloc, tvm.ir.Call) and opaque_alloc.op == tvm.ir.Op.get(
+            "tirx.TVMBackendAllocWorkspace"
+        ):
+            workspace_binds.append(node)
+
+    tvm.tirx.stmt_functor.post_order_visit(After["main"].body, collect_workspace_bind)
+    assert len(workspace_binds) == 1
+    workspace_bind = workspace_binds[0]
+    tvm.ir.assert_structural_equal(workspace_bind.var.ty, workspace_bind.value.ty)
+    tvm.ir.assert_structural_equal(
+        workspace_bind.value.args[0].ty, tvm.ir.PointerType(tvm.ir.PrimType("void"))
+    )
+    roundtrip = tvm.script.from_source(After.script())
+    workspace_binds.clear()
+    tvm.tirx.stmt_functor.post_order_visit(roundtrip["main"].body, collect_workspace_bind)
+    assert len(workspace_binds) == 1
+    tvm.ir.assert_structural_equal(workspace_bind.value, workspace_binds[0].value)
+
 
 def test_lower_cpu_allocation():
     """CPU allocations can be handled at codegen time"""