From a377b31bf5955023c5ad5bb5f13fa3e16c6eb541 Mon Sep 17 00:00:00 2001
From: Hongyi Jin <jinhongyi02@gmail.com>
Date: Wed, 29 Jul 2026 03:14:30 +0000
Subject: [PATCH] [FIX][TIRx] Make TilePrimitiveCall serializable

---
 include/tvm/tirx/tirx_stmt.h |  2 ++
 tests/python/tirx/test_op.py | 29 ++++++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/include/tvm/tirx/tirx_stmt.h b/include/tvm/tirx/tirx_stmt.h
index 41618eb1d96c..6a482165f0ae 100644
--- a/include/tvm/tirx/tirx_stmt.h
+++ b/include/tvm/tirx/tirx_stmt.h
@@ -34,6 +34,8 @@ namespace tirx {
  */
 class TilePrimitiveCallNode : public StmtNode {
  public:
+  explicit TilePrimitiveCallNode(ffi::UnsafeInit tag) : op(tag) {}
+
   TilePrimitiveCallNode(tvm::Op op, ffi::Array<ffi::Any> args,
                         ffi::Map<ffi::String, Buffer> workspace,
                         ffi::Map<ffi::String, ffi::Any> config, ffi::Optional<ffi::String> dispatch,
diff --git a/tests/python/tirx/test_op.py b/tests/python/tirx/test_op.py
index 4c417033d697..dde3662dcebf 100644
--- a/tests/python/tirx/test_op.py
+++ b/tests/python/tirx/test_op.py
@@ -14,10 +14,13 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import pickle
+
 import pytest
 
-from tvm.ir import Op
+from tvm.ir import Op, assert_structural_equal
 from tvm.tirx.buffer import decl_buffer
+from tvm.tirx.exec_scope import ExecScope
 from tvm.tirx.stmt import TilePrimitiveCall
 
 
@@ -44,6 +47,30 @@ def test_gemm():
     _test("gemm", D[:, :], A[:, :], B[:, :], C[:, :], True, False, 1.0, 0.0)
 
 
+def test_tile_primitive_call_pickle_roundtrip():
+    """TilePrimitiveCall reflection must provide a deserialization creator."""
+    A = decl_buffer((64,), "float32", scope="local")
+    workspace = decl_buffer((16,), "float32", scope="shared")
+    call = TilePrimitiveCall(
+        A[:],
+        1.0,
+        op=Op.get("tirx.tile.fill"),
+        workspace={"scratch": workspace},
+        config={"hint": "roundtrip", "stages": 2},
+        dispatch="reg",
+        scope=ExecScope("warpgroup"),
+    )
+    restored = pickle.loads(pickle.dumps(call))
+
+    assert_structural_equal(restored, call)
+    assert restored.op.same_as(call.op)
+    assert_structural_equal(restored.args, call.args)
+    assert_structural_equal(restored.workspace, call.workspace)
+    assert_structural_equal(restored.config, call.config)
+    assert restored.dispatch == call.dispatch
+    assert_structural_equal(restored.scope, call.scope)
+
+
 def test_buffer_replacer_no_shared_default():
     """Regression test for F4: BufferReplacer default dicts must not be shared."""
     from tvm.tirx.transform.common import BufferReplacer