[TIR] Fix Data Type Mismatch (int64 vs int32) in T.match_buffer when Working with Scalar Buffers in TIR (#18466)

This PR is trying to fix issues
#17392.

The issue with `T.match_buffer` for scalar elements that was causing the
int64 vs. int32 type mismatch error in TVM.

Fix:
- Safe Type Coercion: Allows automatic casting between integer types
when they have the same number of lanes
- Type Safety Preserved: Still rejects incompatible type combinations
(int vs float, different lane counts)

---------

Co-authored-by: cchung100m <[email protected]>

Author: cchung100m

src/tir/transforms/lower_match_buffer.cc

@@ -220,8 +220,15 @@ class MatchBufferLower : public StmtExprMutator {
   }
 
   void Bind(const PrimExpr& arg, PrimExpr value, const std::string& arg_name = "argument") {
-    CHECK_EQ(arg.dtype(), value.dtype())
-        << "The data type mismatched: " << arg->dtype << " vs. " << value->dtype;
+    if (arg.dtype() != value.dtype()) {
+      if (arg.dtype().is_int() && value.dtype().is_int() &&
+          arg.dtype().lanes() == value.dtype().lanes()) {
+        value = cast(arg.dtype(), value);
+      } else {
+        CHECK_EQ(arg.dtype(), value.dtype())
+            << "The data type mismatched: " << arg->dtype << " vs. " << value->dtype;
+      }
+    }
     // Handle recursive case
     value = Substitute(std::move(value), var_map_);
     if (arg->IsInstance<VarNode>()) {

tests/python/tir-transform/test_tir_transform_lower_match_buffer.py

@@ -532,5 +532,36 @@ def test_fail_match_func_param():
     _check_fail(fail_match_func_param)
 
 
+@T.prim_func
+def scalar_match_buffer_type_coercion(a: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8))
+    for i, j in T.grid(8, 8):
+        with T.block(""):
+            vi = T.axis.spatial(8, i)
+            vj = T.axis.spatial(8, j)
+            T.reads()
+            T.writes(A[vi, vj])
+            # Create scalar match buffer from single element - this triggers type coercion
+            scalar_buf = T.match_buffer(A[vi, vj], (), offset_factor=1)
+            scalar_buf[()] = T.float32(1.0)
+
+
+@T.prim_func
+def transformed_scalar_match_buffer_type_coercion(a: T.handle) -> None:
+    A = T.match_buffer(a, (8, 8))
+    for i, j in T.grid(8, 8):
+        with T.block(""):
+            vi = T.axis.spatial(8, i)
+            vj = T.axis.spatial(8, j)
+            T.reads()
+            T.writes(A[vi, vj])
+            # Scalar match_buffer eliminated, direct assignment
+            A[vi, vj] = T.float32(1.0)
+
+
+def test_scalar_match_buffer_type_coercion():
+    _check(scalar_match_buffer_type_coercion, transformed_scalar_match_buffer_type_coercion)
+
+
 if __name__ == "__main__":
     tvm.testing.main()