From d7615fdf1121ef9190b80efac679233fc82422a5 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Mon, 3 Apr 2023 11:04:30 -0500 Subject: [PATCH] [TIR] Use same DataType of builtin::tvm_struct_set in C++ and Python Prior to this commit, the python API `tvm.tir.op.tvm_struct_set` defined the return type of `builtin::tvm_struct_set` as `"handle"`, while the C++ API `tvm::tir::TVMStructSet` defined the return type as `DataType::Int(32)`. The data type used for this builtin has no effect, because no value is returned. However, this discrepancy can cause failure to roundtrip through TVMScript. This commit updates the Python API to use `"int32"`, for consistency with the C++ API and with `CodeGenCPU`. --- python/tvm/tir/op.py | 2 +- .../unittest/test_tvmscript_roundtrip.py | 36 ++++++++++++++++++- 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/python/tvm/tir/op.py b/python/tvm/tir/op.py index 0fe460c085d7..419ab2275858 100644 --- a/python/tvm/tir/op.py +++ b/python/tvm/tir/op.py @@ -527,7 +527,7 @@ def tvm_struct_set(arr, index, field, value): call : PrimExpr The call expression. """ - return call_intrin("handle", "tir.tvm_struct_set", arr, index, field, value) + return call_intrin("int32", "tir.tvm_struct_set", arr, index, field, value) def address_of(buffer_load, span=None): diff --git a/tests/python/unittest/test_tvmscript_roundtrip.py b/tests/python/unittest/test_tvmscript_roundtrip.py index cd7f1726c9d9..bbc6dd45a83e 100644 --- a/tests/python/unittest/test_tvmscript_roundtrip.py +++ b/tests/python/unittest/test_tvmscript_roundtrip.py @@ -21,7 +21,7 @@ import tvm import tvm.testing from tvm import tir -from tvm.script import tir as T +from tvm.script import tir as T, ir as I import numpy as np @@ -3692,6 +3692,39 @@ def func( return func +def tvm_struct_set_generated_in_cpp(): + """Ensure same dtype for tvm_struct_set in Python/C++ + + The TVMStructSet method in C++, used internally by + LowerTVMBuiltin, and the Python method `T.tvm_struct_set`, used + when parsing TVMScript should use the same dtype "int32". + """ + + @I.ir_module + class Module: + @T.prim_func + def tir_packed_call(A: T.Buffer(16)): + T.attr(0, "device_id", 0) + T.attr(0, "device_type", 0) + T.evaluate( + T.tvm_call_cpacked( + "tvm_test_cpacked", + T.tvm_stack_make_array( + A.data, + T.tvm_stack_make_shape(16, dtype="handle"), + T.reinterpret(T.uint64(0), dtype="handle"), + T.uint32(1), + T.Cast("float32", 0), + 0, + dtype="handle", + ), + dtype="int32", + ) + ) + + return tvm.tir.transform.LowerTVMBuiltin()(Module) + + ir_generator = tvm.testing.parameter( launch_env_thread, opt_gemm_normalize, @@ -3757,6 +3790,7 @@ def func( merge_shape_var_def, if_then_else_var, tvm_shfl_builtins, + tvm_struct_set_generated_in_cpp, )