[TVMScript] Enable Safe Autocasting in BufferStore

src/script/ir_builder/tir/ir.cc

@@ -492,6 +492,34 @@ Var EnvThread(String thread_tag) {
 }
 
 void BufferStore(Buffer buffer, PrimExpr value, Array<PrimExpr> indices) {
+  runtime::DataType buffer_dtype = buffer->dtype;
+  int index_lanes = indices.size() ? indices.back().dtype().lanes() : 1;
+  runtime::DataType lhs_dtype = buffer_dtype.with_lanes(buffer_dtype.lanes() * index_lanes);
+  runtime::DataType rhs_dtype = value->dtype;
+  if (lhs_dtype != rhs_dtype) {
+    if (lhs_dtype.lanes() != rhs_dtype.lanes()) {
+      LOG(FATAL) << "TypeError: Incompatible types in BufferStore"
+                 << ": LHS is `" << lhs_dtype << "`, RHS is `" << rhs_dtype
+                 << "`, indexing lanes: " << index_lanes;
+    }
+    if (lhs_dtype.code() != rhs_dtype.code()) {
+      if (
+          // Case 1. lhs is handle, and rhs needs to be casted to handle.
+          (lhs_dtype.code() == runtime::DataType::kHandle) ||
+          // Case 2. rhs is handle, and it needs to be casted to non-handle.
+          (rhs_dtype.code() == runtime::DataType::kHandle) ||
+          // Case 3. rhs is float or bfloat, and casting to non-float can lose precision.
+          ((lhs_dtype.code() == runtime::DataType::kInt ||
+            lhs_dtype.code() == runtime::DataType::kUInt) &&
+           (rhs_dtype.code() == runtime::DataType::kFloat ||
+            rhs_dtype.code() == runtime::DataType::kBFloat))) {
+        LOG(WARNING) << "Casting in BufferStore may lose precision"
+                     << ": LHS is `" << lhs_dtype << "`, RHS is `" << rhs_dtype
+                     << "`, indexing lanes: " << index_lanes;
+      }
+    }
+    value = tvm::cast(lhs_dtype, value);
+  }
   AddToParent(tvm::tir::BufferStore(buffer, value, indices));
 }
 

src/script/ir_builder/tir/utils.h

@@ -21,6 +21,7 @@
 
 #include <tvm/script/ir_builder/tir/frame.h>
 #include <tvm/script/ir_builder/tir/ir.h>
+#include <tvm/tir/op.h>
 #include <tvm/tir/stmt.h>
 
 namespace tvm {

src/tir/ir/stmt.cc

@@ -507,6 +507,12 @@ BufferStore::BufferStore(Buffer buffer, PrimExpr value, Array<PrimExpr> indices,
       << "Cannot store value with " << value.dtype().lanes() << ", expected value with "
       << index_lanes * buffer_lanes << " (" << index_lanes << " index lanes * " << buffer_lanes
       << " buffer element lanes)";
+  if (buffer->dtype.with_lanes(buffer_lanes * index_lanes) != value.dtype()) {
+    LOG(FATAL) << "TypeError: dtype mismatch on BufferStore: "      //
+               << "buffer's dtype is `" << buffer->dtype            //
+               << "`, the lanes of indexing are: `" << index_lanes  //
+               << "`, but RHS's dtype is `" << value.dtype() << "`";
+  }
 
   ObjectPtr<BufferStoreNode> node = make_object<BufferStoreNode>();
   node->buffer = std::move(buffer);

tests/python/unittest/test_tir_constructor.py

@@ -16,7 +16,6 @@
 # under the License.
 
 import pytest
-
 import tvm
 from tvm import te
 
@@ -153,7 +152,7 @@ def test_stmt_constructor():
 
     buffer_var = tvm.tir.Var("buf", tvm.ir.PointerType(tvm.ir.PrimType("uint1")))
     buffer = tvm.tir.decl_buffer([16], "uint1", data=buffer_var)
-    x = tvm.tir.BufferStore(buffer, 1, [10])
+    x = tvm.tir.BufferStore(buffer, tvm.tir.IntImm("bool", 1), [10])
     assert isinstance(x, tvm.tir.BufferStore)
     assert x.buffer == buffer
     assert x.buffer.data == buffer_var