Refactor index handling in BufferStore and BufferLoad to promote 64-bit integers (#796)

- Updated index processing in `BufferStore` and `BufferLoad` to ensure that integer indices with less than 64 bits are promoted to 64-bit integers.
- Introduced a new array to store the modified indices before updating the original indices, enhancing clarity and maintainability of the code.

Author: LeiWang1999

src/transform/config_index_bitwidth.cc

@@ -123,35 +123,43 @@ class IndexLegalizer : public IRMutatorWithAnalyzer {
     auto buffer_store =
         Downcast<BufferStore>(IRMutatorWithAnalyzer::VisitStmt_(op));
     auto indices = buffer_store->indices;
+    Array<PrimExpr> new_indices;
     for (auto index : indices) {
       if (index->dtype.is_int() && index->dtype.bits() < 64) {
         auto int_bound = analyzer_->const_int_bound(index);
         if (int_bound->max_value >= (1LL << (index->dtype.bits() - 1)) - 1 ||
             int_bound->min_value < -(1LL << (index->dtype.bits() - 1))) {
           Int64Promoter promoter;
           index = promoter(index);
+          new_indices.push_back(index);
+          continue;
         }
       }
+      new_indices.push_back(index);
     }
-    buffer_store.CopyOnWrite()->indices = indices;
+    buffer_store.CopyOnWrite()->indices = new_indices;
     return std::move(buffer_store);
   }
 
   PrimExpr VisitExpr_(const BufferLoadNode *op) final {
     auto buffer_load =
         Downcast<BufferLoad>(IRMutatorWithAnalyzer::VisitExpr_(op));
     auto indices = buffer_load->indices;
+    Array<PrimExpr> new_indices;
     for (auto index : indices) {
       if (index->dtype.is_int() && index->dtype.bits() < 64) {
         auto int_bound = analyzer_->const_int_bound(index);
         if (int_bound->max_value >= (1LL << (index->dtype.bits() - 1)) - 1 ||
             int_bound->min_value < -(1LL << (index->dtype.bits() - 1))) {
           Int64Promoter promoter;
           index = promoter(index);
+          new_indices.push_back(index);
+          continue;
         }
       }
+      new_indices.push_back(index);
     }
-    buffer_load.CopyOnWrite()->indices = indices;
+    buffer_load.CopyOnWrite()->indices = new_indices;
     return std::move(buffer_load);
   }
 };