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 <[email protected]> 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"""
