fix

Author: vinx13

src/tir/schedule/primitive/blockize_tensorize.cc

@@ -16,6 +16,8 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+#include <tvm/tir/data_type_rewriter.h>
+
 #include <functional>
 
 #include "../ir_comparator.h"
@@ -523,6 +525,19 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int
   }
   PrimFunc intrin_desc = intrin->desc;
   PrimFunc intrin_impl = DeepCopy(intrin->impl);
+
+  int index_dtype_bits = -1;
+  auto f_update_max_dtype_bits_from_region = [&](const Array<BufferRegion>& buffer_regions) {
+    for (const BufferRegion& buffer_region : buffer_regions) {
+      for (const auto& range : buffer_region->region) {
+        index_dtype_bits = std::max(index_dtype_bits, range->min.dtype().bits());
+      }
+    }
+  };
+  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);
+  intrin_impl = IndexDataTypeNormalizer(DataType::Int(index_dtype_bits)).Rewrite(intrin_impl);
   // Step 2: Structural pattern matching
   TensorizeComparator comparator(self->mod, /*assert_mode=*/true);
   comparator.VisitStmt(block_realize, intrin_desc->body);

src/tir/transforms/lower_match_buffer.cc

@@ -23,19 +23,17 @@
  */
 
 #include <tvm/arith/analyzer.h>
-#include <tvm/tir/data_type_rewriter.h>
 #include <tvm/tir/function.h>
 #include <tvm/tir/op.h>
 #include <tvm/tir/stmt_functor.h>
 #include <tvm/tir/transform.h>
 
-#include "../../printer/text_printer.h"
 #include "../ir/functor_common.h"
 #include "ir_utils.h"
 
 namespace tvm {
 namespace tir {
-class MatchBufferLower : public DataTypeLegalizer {
+class MatchBufferLower : public StmtExprMutator {
  public:
   explicit MatchBufferLower(const PrimFunc& func) {
     for (const Var& param : func->params) {
@@ -190,14 +188,14 @@ class MatchBufferLower : public DataTypeLegalizer {
       Array<PrimExpr> buffer_start_indices = source_buffer->ElemOffset(indices);
       if (buffer_start_indices.size() == 1) {
         Bind(buffer->elem_offset, buffer_start_indices[0], buffer->name + ".elem_offset");
-        CHECK(analyzer_.CanProve(truncmod(buffer_start_indices[0], buffer->offset_factor) == 0))
+        CHECK(analyzer_.CanProve(truncmod(buffer->elem_offset, buffer->offset_factor) == 0))
             << "The source elem_offset " << buffer_start_indices[0]
             << " does not satisfy the offset_factor " << buffer->offset_factor << ".";
       } else {
         // 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, make_zero(buffer->elem_offset.dtype()));
+        Bind(buffer->elem_offset, 0);
       }
     }
 
@@ -231,7 +229,7 @@ class MatchBufferLower : public DataTypeLegalizer {
   }
 
   void Bind(const PrimExpr& arg, PrimExpr value, const std::string& arg_name = "argument") {
-    CHECK_EQ(arg.dtype().code(), value.dtype().code())
+    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_);
@@ -240,7 +238,7 @@ class MatchBufferLower : public DataTypeLegalizer {
       auto it = var_map_.find(v);
       if (it == var_map_.end()) {
         var_map_.Set(v, value);
-        // analyzer_.Bind(v, value);
+        analyzer_.Bind(v, value);
       } else {
         AssertBinding((*it).second, value, arg_name);
       }
@@ -249,21 +247,10 @@ class MatchBufferLower : public DataTypeLegalizer {
     }
   }
 
-  PrimExpr LookUpArgBind(const PrimExpr& arg) {
-    if (arg->IsInstance<VarNode>()) {
-      Var v = Downcast<Var>(arg);
-      if (auto it = var_map_.find(v); it != var_map_.end()) {
-        return (*it).second;
-      }
-    }
-    return arg;
-  }
-
   void AssertBinding(const PrimExpr& lhs, const PrimExpr& rhs,
                      const std::string& arg_name = "argument") {
-    CHECK(analyzer_.CanProve(LookUpArgBind(lhs) == rhs))
-        << "The buffer match constraint for " << arg_name << " unmet: " << lhs << "==" << rhs
-        << ".";
+    CHECK(analyzer_.CanProve(lhs == rhs)) << "The buffer match constraint for " << arg_name
+                                          << " unmet: " << lhs << "==" << rhs << ".";
   }
 
  private:
@@ -277,9 +264,7 @@ class MatchBufferLower : public DataTypeLegalizer {
 
 PrimFunc LowerMatchBuffer(PrimFunc func) {
   auto fptr = func.CopyOnWrite();
-  // LOG(INFO) << "BeforeLMB:\n" << tir::AsTVMScript(func);
   fptr->body = MatchBufferLower(func)(std::move(fptr->body));
-  // LOG(INFO) << "AfterLMB:\n" << tir::AsTVMScript(func);
   return func;
 }
 
@@ -297,4 +282,4 @@ TVM_REGISTER_GLOBAL("tir.transform.LowerMatchBuffer").set_body_typed(LowerMatchB
 }  // namespace transform
 
 }  // namespace tir
-}  // namespace tvm
+}  // namespace tvm