[Unity] Fix IndexDataTypeNormalizer so that it correctly handles corn…

…er case (#16245) fix
apache · Dec 15, 2023 · a2e19d2 · a2e19d2
1 parent cd9445d
commit a2e19d2
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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));