[TIR] Use same DataType of builtin::tvm_struct_set in C++ and Python (#…

…14489)

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

@@ -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):

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,
 )