(tvm) branch unity updated: [Unity] Fix IndexDataTypeNormalizer so that it correctly handles corner case (#16245)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new a2e19d21eb [Unity] Fix IndexDataTypeNormalizer so that it correctly 
handles corner case (#16245)
a2e19d21eb is described below

commit a2e19d21eb89329df03c7e5ffcd14ddb5e09dd86
Author: Hongyi Jin <jinhongy...@gmail.com>
AuthorDate: Thu Dec 14 20:30:15 2023 -0500

    [Unity] Fix IndexDataTypeNormalizer so that it correctly handles corner 
case (#16245)
    
    fix
---
 src/tir/ir/data_type_rewriter.cc | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/tir/ir/data_type_rewriter.cc b/src/tir/ir/data_type_rewriter.cc
index aa8f2f3f6f..e68d085b44 100644
--- a/src/tir/ir/data_type_rewriter.cc
+++ b/src/tir/ir/data_type_rewriter.cc
@@ -582,6 +582,11 @@ IndexDataTypeNormalizer::IndexDataTypeNormalizer(DataType 
target_data_type)
     : target_data_type_(std::move(target_data_type)) {}
 
 PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) {
+  // collect var remap
+  VisitStmt(std::move(func->body));
+  buffer_remap_.clear();
+  ivmap_.clear();
+  // start rewrite
   Map<Var, Buffer> new_buffer_map = func->buffer_map;
   for (const auto& [var, buffer] : func->buffer_map) {
     new_buffer_map.Set(var, VisitBuffer(buffer));