update tensorize

Author: vinx13

src/tir/schedule/primitive/blockize_tensorize.cc

@@ -537,6 +537,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int
   f_update_max_dtype_bits_from_region(block_realize->block->reads);
   f_update_max_dtype_bits_from_region(block_realize->block->writes);
   ICHECK(index_dtype_bits > 0);
+  LOG(INFO) << "normalize to " << index_dtype_bits << " bits";
   intrin_impl = IndexDataTypeNormalizer(DataType::Int(index_dtype_bits)).Rewrite(intrin_impl);
   // Step 2: Structural pattern matching
   TensorizeComparator comparator(self->mod, /*assert_mode=*/true);

src/tir/transforms/lower_match_buffer.cc

@@ -195,7 +195,7 @@ class MatchBufferLower : public StmtExprMutator {
         // Non-zero elem_offset is ill-defined for non-flat memory.
         // If needed in the future, will require `Array<PrimExpr>
         // elem_offsets`, with one offset for each flattened index.
-        Bind(buffer->elem_offset, 0);
+        Bind(buffer->elem_offset, make_const(buffer->elem_offset.dtype(), 0));
       }
     }
 
@@ -206,7 +206,7 @@ class MatchBufferLower : public StmtExprMutator {
     if (!buffer->strides.empty()) {
       ICHECK_EQ(buffer->strides.size(), buffer->shape.size());
       if (source_buffer->strides.empty()) {
-        PrimExpr stride = make_const(DataType::Int(32), 1);
+        PrimExpr stride = make_const(buffer->strides.back().dtype(), 1);
         for (size_t i = buffer->shape.size(); i > 0; --i) {
           const PrimExpr& shape = source_buffer->shape[i - 1 + offset];
           Bind(buffer->strides[i - 1], stride, buffer->name + ".strides_" + std::to_string(i - 1));
@@ -230,7 +230,8 @@ 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;
+        << "The data type mismatched: " << arg->dtype << " vs. " << value->dtype
+        << " when binding " << arg_name << " to " << value;
     // Handle recursive case
     value = Substitute(std::move(value), var_map_);
     if (arg->IsInstance<VarNode>()) {
@@ -282,4 +283,4 @@ TVM_REGISTER_GLOBAL("tir.transform.LowerMatchBuffer").set_body_typed(LowerMatchB
 }  // namespace transform
 
 }  // namespace tir
-}  // namespace tvm
+}  // namespace tvm

tests/python/unittest/test_tir_schedule_tensorize.py

@@ -697,34 +697,34 @@ def tensorized_matmul_int64_shape(
                         ]
                     )
                     T.writes(C[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vj * T.int64(16) : vj * T.int64(16) + T.int64(16)])
-                    A_elem_offset = T.var("int32")
-                    B_elem_offset = T.var("int32")
-                    C_elem_offset = T.var("int32")
+                    A_elem_offset = T.var("int64")
+                    B_elem_offset = T.var("int64")
+                    C_elem_offset = T.var("int64")
                     A_sub = T.match_buffer(
                         A[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vk * T.int64(16) : vk * T.int64(16) + T.int64(16)],
-                        [16, 16],
+                        [T.int64(16), T.int64(16)],
                         elem_offset=A_elem_offset,
                     )
                     B_sub = T.match_buffer(
                         B[vj * T.int64(16) : vj * T.int64(16) + T.int64(16), vk * T.int64(16) : vk * T.int64(16) + T.int64(16)],
-                        [16, 16],
+                        [T.int64(16), T.int64(16)],
                         elem_offset=B_elem_offset,
                     )
                     C_sub = T.match_buffer(
                         C[vi * T.int64(16) : vi * T.int64(16) + T.int64(16), vj * T.int64(16) : vj * T.int64(16) + T.int64(16)],
-                        [16, 16],
+                        [T.int64(16), T.int64(16)],
                         elem_offset=C_elem_offset,
                     )
                     T.evaluate(
                         T.tvm_mma_sync(
                             C_sub.data,
-                            T.floordiv(C_sub.elem_offset, 256),
+                            T.floordiv(C_sub.elem_offset, T.int64(256)),
                             A_sub.data,
-                            T.floordiv(A_sub.elem_offset, 256),
+                            T.floordiv(A_sub.elem_offset, T.int64(256)),
                             B_sub.data,
-                            T.floordiv(B_sub.elem_offset, 256),
+                            T.floordiv(B_sub.elem_offset, T.int64(256)),
                             C_sub.data,
-                            T.floordiv(C_sub.elem_offset, 256),
+                            T.floordiv(C_sub.elem_offset, T.int64(256)),
                             dtype="handle",
                         )
                     )