Repository: apache/tvm
Pull request: [TIR] Shuffle in PointerValueTypeRewrite for scalar reads 
Added an option `rewrite_scalar_read_to_vector_shuffle` in `PointerValueTypeRewrite` (currently
only enabled for Vulkan). When enabled, when a buffer has both scalar
and vector reads, the buffer will be vectorized if possible and scalar
reads will be achieved via T.Shuffle.

Close https://github.com/apache/tvm/issues/15463.

cc @tqchen @Lunderberg @sunggg 

Problem statement:
[TensorIR] Enhance PointerTypeRewrite
Our current GPU backends sometimes can generate the following code:

```c++
void kernel(float* global_data, float* res) {
    float data[4];
    ((float4*)data)[0] = ((float4*)global_data)[0];
     res[0] = data[0] + data[1] + data[2] + data[3];
}
```

Here we are doing vectorized store into a local memory then do a scalarized access from the local memory. While such pattern can be supported by compilers like cuda, it is not always the most robust path. Additionally, backends like vulkan, webgpu comes with the restriction that we can only have one kind of access to the data, resulting the code to be scalarized.

A better version of the code is as follows, when the following conditions are met:
- We detect we have aligned vector store into data
- For all read access to the data, they are scalar access, and can be represented by extract_element(Shuffle), this can be done by running modular analysis and check that all access have coeff that is multiple of the vector lane 

When such condition are met, we can still rewrite data into the vector type, and rewrite all read access of the data to extract_element(represented via shuffle).

```c++
void new_kernel(float4* global_data, float* res) {
    float4 vec_data[1];
    vec_data[0] = global_data[0];
    // access the elements of the vector.
    res[0] = vec_data[0].x + vec_data[0].y + vec_data[0].z + vec_data[0].w;
}
```

This is a pretty common case in GPU programs so it would be desirable to upgrade  `PointerTypeRewrite` to support such cases. This would likely help to robustify our overall support as the new_kernel is also what most programmer will write by handle 

## Action Items

- [ ] Enhance the var analysis to decouple buffer read and write/opaque access, 
     - The read part needs to contain a coeff entry via Modular analysis that summarizes whether rewrite to extract_element is possible
     - For all read access, do Modular analysis to get coeff, so we know that index = x * coeff + offset
     - If coeff % vector_lane == 0, then it means we can rewrite all scalar access to vector access `vec_data[(index - offset % veclane) // vec_lane].extract_element(offset % veclane)`. We can keep GCD of all coeff access to trigger the condition
     - To keep backward compatibility, we can put it behind a flag `rewrite_scalar_read_to_vector_shuffle` 
- [ ] Rewrite scalar read access into extract_element when we find such cases.
- [ ] Support extract_element shuffle pattern (more general shuffle pattern is OK but extract_element is the main one) in codegen , then turn `rewrite_scalar_read_to_vector_shuffle` on for the cases


cc @Hzfengsy @junrushao @quic-sanirudh @shingjan

Hints from issue discussion:
I'm interested in working on this
I really like this idea, and it will simplify a lot of the codegen-time analysis that occurs to efficiently handle Ramp nodes.

Benefits:

* This mimics what is already done as part of the LLVM codegen ([link](https://github.com/apache/tvm/blob/main/src/target/llvm/codegen_llvm.cc#L1364)), in case of a
* This would remove the need for the `CodeGenC` generation of vector loads from Ramp indices([link](https://github.com/apache/tvm/blob/main/src/target/source/codegen_c.cc#L691)), since any vectorizable loads would already be expressed with scalar indices.  After the enhanced PointerTypeRewrite, `bool can_vector_load` would always be false in `CodeGenC::VisitExpr_(const BufferLoadNode*)`.
* The special handling that Vulkan has ([link](https://github.com/apache/tvm/blob/main/src/target/spirv/spirv_utils.cc#L123) could be removed entirely, since the lowered module would already satisfy the constraint `buffer_load.dtype == buffer_load.buffer.dtype`.

Potential breakages

* Vector types that whose struct representation contains padding.  For example, LLVM uses 32-bits to store `int8x3`, with padding in bits 24-31.  For these types, the address `((int8x3*)ptr)[N]` is not the same as `((int8*)ptr)[3*N]`.
* Padded vectorization of `AllocateConst`.  When a buffer's dtype is vectorized, the underlying data in `AllocateConst` is kept as-is.
* DLTensor arguments.  The size of a DLTensor is computed ([link](https://github.com/apache/tvm/blob/main/include/tvm/runtime/ndarray.h#L341)) assuming that elements are padded out to the nearest byte (e.g. padding a 1-bit `DataType::Bool()` out to 8 bits), but does not introduce any padding between vectorized elements (e.g. padding a 24-bit `DataType::Int(8, 3)` to 32 bits).

Likely follow-ups:

* Any remaining scalar read access would be updated from `((float*)ptr)[index]` to `extract_element(((float4*)ptr)[index/4], index%4)`, and the `((float4*)ptr)[index/4]` would likely benefit from CSE.  Currrently, the `tir.transform.CommonSubexprElimTIR` pass forbids removal of `BufferLoad` statements entirely ([link](https://github.com/apache/tvm/blob/main/src/tir/transforms/common_subexpr_elim.cc#L69)), as their contents may change.  We will probably want to loosen up this restriction, to allow CSE elimination of input buffers.

Commit message: [TIR] Shuffle in PointerValueTypeRewrite for scalar reads

Added an option `rewrite_scalar_read_to_vector_shuffle` in `PointerValueTypeRewrite` (currently
only enabled for Vulkan). When enabled, when a buffer has both scalar
and vector reads, the buffer will be vectorized if possible and scalar
reads will be achieved via T.Shuffle.

Close https://github.com/apache/tvm/issues/15463.

Patch to review:
diff --git a/python/tvm/tir/transform/transform.py b/python/tvm/tir/transform/transform.py
index 0cd54064a7b5..2f564e4ec8e6 100644
--- a/python/tvm/tir/transform/transform.py
+++ b/python/tvm/tir/transform/transform.py
@@ -230,6 +230,19 @@ def StorageRewrite():
     return _ffi_api.StorageRewrite()  # type: ignore
 
 
+def PointerValueTypeRewrite():
+    """
+    Rewrite the pointer content type of arguments, as well as Alloc internal to the function to use
+    the most frequently accessed type for load/store to avoid pointer casting in backend when possible.
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.PointerValueTypeRewrite()  # type: ignore
+
+
 def UnrollLoop():
     """Unroll the constant loop marked by unroll.
 
diff --git a/src/target/spirv/codegen_spirv.cc b/src/target/spirv/codegen_spirv.cc
index ab9aec077542..5cc3f8f8ddc0 100644
--- a/src/target/spirv/codegen_spirv.cc
+++ b/src/target/spirv/codegen_spirv.cc
@@ -610,6 +610,17 @@ void CodeGenSPIRV::Scalarize(const PrimExpr& e, std::function<void(int i, spirv:
   }
 }
 
+spirv::Value CodeGenSPIRV::VisitExpr_(const ShuffleNode* op) {
+  ICHECK(op->vectors.size() == 1 && op->indices.size() == 1)
+      << "SPIR-V codegen only supports shuffle "
+      << "of one vector with one index";
+  spirv::Value vector = MakeValue(op->vectors[0]);
+  int index = Downcast<Integer>(op->indices[0])->value;
+  spirv::SType etype = builder_->GetSType(op->dtype);
+  spirv::Value element = builder_->MakeValue(spv::OpCompositeExtract, etype, vector, index);
+  return element;
+}
+
 void CodeGenSPIRV::VisitStmt_(const BufferStoreNode* op) {
   ICHECK_EQ(op->indices.size(), 1) << "SPIR-V codegen expects flat memory buffers";
   Var buffer_var = op->buffer->data;
diff --git a/src/target/spirv/codegen_spirv.h b/src/target/spirv/codegen_spirv.h
index 1e7b53558508..8ea90a9c4b80 100644
--- a/src/target/spirv/codegen_spirv.h
+++ b/src/target/spirv/codegen_spirv.h
@@ -102,6 +102,7 @@ class CodeGenSPIRV : public ExprFunctor<spirv::Value(const PrimExpr&)>,
   spirv::Value VisitExpr_(const RampNode* op) override;
   spirv::Value VisitExpr_(const BroadcastNode* op) override;
   spirv::Value VisitExpr_(const BufferLoadNode* op) override;
+  spirv::Value VisitExpr_(const ShuffleNode* op) override;
   // stmt
   void VisitStmt_(const BufferStoreNode* op) override;
   void VisitStmt_(const ForNode* op) override;
diff --git a/src/tir/transforms/storage_rewrite.cc b/src/tir/transforms/storage_rewrite.cc
index 3ecd0f64bb44..44315da9414a 100644
--- a/src/tir/transforms/storage_rewrite.cc
+++ b/src/tir/transforms/storage_rewrite.cc
@@ -33,6 +33,7 @@
 #include <tvm/tir/transform.h>
 
 #include <map>
+#include <numeric>
 #include <unordered_map>
 #include <unordered_set>
 
@@ -1083,17 +1084,14 @@ struct BufferVarInfo {
 
     DataType preferred_base_type = *base_access_dtype.begin();
 
-    // If there is only one vectorizable size used to access the
-    // buffer, and if that access size is compatible with the array
-    // size, then the buffer is vectorizable.  In the future, this
-    // could be improved to allow vectorized buffer access of size
-    // GCD(*lanes_used), if necessary.
     int preferred_lanes = element_dtype.lanes();
-    if ((element_dtype.lanes() == 1) && (access_dtype.size() == 1)) {
+    if (element_dtype.lanes() == 1) {
+      int lanes = access_dtype.begin()->lanes();
+      for (auto dtype : access_dtype) {
+        lanes = std::gcd(lanes, dtype.lanes());
+      }
       arith::Analyzer analyzer_;
       arith::ModularSet me = analyzer_.modular_set(extent);
-
-      int lanes = access_dtype.begin()->lanes();
       if ((me->coeff % lanes == 0) && (me->base % lanes == 0)) {
         preferred_lanes = lanes;
       }
@@ -1120,8 +1118,10 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
    * type as it is later accessed, with scalar element types.
    */
   VectorTypeAccessChecker(const Array<tir::Var>& params, const Map<Var, Buffer>& buffer_map,
-                          bool allow_untyped_pointers = false)
-      : allow_untyped_pointers_(allow_untyped_pointers) {
+                          bool allow_untyped_pointers = false,
+                          bool detect_scalar_read_patterns = true)
+      : allow_untyped_pointers_(allow_untyped_pointers),
+        detect_scalar_read_patterns_(detect_scalar_read_patterns) {
     // If a parameter is in the buffer map, we want to track the
     // version in the map.
     for (auto it : buffer_map) {
@@ -1145,12 +1145,12 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
   }
 
   void VisitExpr_(const BufferLoadNode* op) final {
-    OnArrayAccess(op->dtype, op->buffer->data.get(), op->indices);
+    OnArrayAccess(op->dtype, op->buffer->data.get(), op->indices, /*is_buffer_load=*/true);
     StmtExprVisitor::VisitExpr_(op);
   }
 
   void VisitStmt_(const BufferStoreNode* op) final {
-    OnArrayAccess(op->value.dtype(), op->buffer->data.get(), op->indices);
+    OnArrayAccess(op->value.dtype(), op->buffer->data.get(), op->indices, /*is_buffer_load=*/false);
     StmtExprVisitor::VisitStmt_(op);
   }
 
@@ -1159,7 +1159,10 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
       DataType dtype = op->args[0].dtype();
       const VarNode* buffer = op->args[1].as<VarNode>();
       PrimExpr index = op->args[2];
-      OnArrayAccess(dtype, buffer, {index});
+      OnArrayAccess(dtype, buffer, {index}, false);
+    } else if (op->op.same_as(builtin::address_of())) {
+      BufferLoad load = Downcast<BufferLoad>(op->args[0]);
+      OnArrayAccess(load->dtype, load->buffer->data.get(), load->indices, /*is_buffer_load=*/false);
     }
     StmtExprVisitor::VisitExpr_(op);
   }
@@ -1226,8 +1229,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
     if (element_dtype == DataType::Bool()) {
       element_dtype = DataType::Int(8).with_lanes(element_dtype.lanes());
     }
-
-    info_map_[buffer.get()] = {buffer, element_dtype, extent, declaration_location};
+    info_map_[buffer.get()] = BufferVarInfo{buffer, element_dtype, extent, declaration_location};
   }
 
   /* Update the type map for a buffer based on its usage
@@ -1237,11 +1239,12 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
    *
    * @param buffer The VarNode representing the buffer.
    *
-   * @param index The index at which the value is being stored/loaded.
+   * @param indices The index at which the value is being stored/loaded.
    *
-   * @param predicate The predicate used for the store/load.
+   * @param is_buffer_load Whether the access is BufferLoad
    */
-  void OnArrayAccess(DataType value_dtype, const VarNode* buffer, const Array<PrimExpr>& indices) {
+  void OnArrayAccess(DataType value_dtype, const VarNode* buffer, const Array<PrimExpr>& indices,
+                     bool is_buffer_load) {
     auto it = info_map_.find(buffer);
     ICHECK(it != info_map_.end()) << "Load/Store of buffer " << buffer->name_hint << " (" << buffer
                                   << ") occurred before its declaration.";
@@ -1304,6 +1307,18 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
       }
     }
 
+    if (detect_scalar_read_patterns_ && is_buffer_load && indices.size()) {
+      const PrimExpr last_dim_index = indices[indices.size() - 1];
+      if (last_dim_index.dtype().lanes() == 1) {
+        arith::ModularSet me = analyzer_.modular_set(last_dim_index);
+        if (me->coeff > 0) {
+          // When coeff == 0, the index is constant and doesn't need to be recorded since it can
+          // always be rewritten to shuffle.
+          var_info.access_dtype.insert(access_dtype.with_lanes(me->coeff));
+        }
+        return;
+      }
+    }
     var_info.access_dtype.insert(access_dtype.with_lanes(lanes_used));
   }
 
@@ -1312,6 +1327,8 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
 
   //
   bool allow_untyped_pointers_{false};
+  // Whether to detect scalar read patterns for rewriting to vector shuffle
+  bool detect_scalar_read_patterns_{true};
 
   // internal analyzer
   arith::Analyzer analyzer_;
@@ -1366,7 +1383,8 @@ class VectorTypeRewriter : public StmtExprMutator {
   VectorTypeRewriter(const std::unordered_map<const VarNode*, BufferVarInfo>& info_map,
                      bool rewrite_params = true, bool rewrite_buffer_map = true,
                      bool rewrite_allocate_node = true, bool rewrite_indices = true,
-                     bool rewrite_let_node = true, bool rewrite_allocate_const_node = true)
+                     bool rewrite_let_node = true, bool rewrite_allocate_const_node = true,
+                     bool rewrite_scalar_read_to_vector_shuffle = true)
       : rewrite_indices_(rewrite_indices) {
     int rewrite_mask = 0;
     if (rewrite_params) {
@@ -1401,42 +1419,52 @@ class VectorTypeRewriter : public StmtExprMutator {
     }
   }
 
+  /*!
+   * \brief Mutator for BufferLoad or BufferStore.
+   * \return The rewritten node and the shuffle index. (Only for BufferLoad) When the shuffle index
+   * is non-negative, the caller should generate Shuffle to extract the element from the vector.
+   */
   template <typename Node>
-  Node VisitBufferAccess(Node node) {
+  std::pair<Node, int> VisitBufferAccess(Node node) {
+    int shuffle_index = -1;
     if (!rewrite_indices_) {
-      return node;
+      return {node, shuffle_index};
     }
 
     auto it = rewrite_map_.find(node->buffer->data.get());
     if (it == rewrite_map_.end()) {
-      return node;
+      return {node, shuffle_index};
     }
     const auto& info = it->second;
 
     Array<PrimExpr> indices = node->indices;
-
-    const RampNode* ramp_index = indices[indices.size() - 1].as<RampNode>();
-    if (ramp_index && is_one(ramp_index->stride)) {
+    const PrimExpr& last_dim_index = indices[indices.size() - 1];
+    if (const RampNode* ramp_index = last_dim_index.as<RampNode>();
+        ramp_index && is_one(ramp_index->stride)) {
       PrimExpr new_index =
           ramp_index->base / make_const(ramp_index->base.dtype(), ramp_index->lanes);
       if (ramp_index->lanes != info.factor()) {
         new_index = Ramp(new_index, ramp_index->stride, ramp_index->lanes / info.factor(),
                          ramp_index->span);
       }
-
+      indices.Set(indices.size() - 1, new_index);
+    } else if (last_dim_index.dtype().lanes() == 1 && info.factor() > 1) {
+      arith::ModularSet me = analyzer_.modular_set(last_dim_index);
+      ICHECK(me->coeff == 0 || info.factor() % me->coeff == 0);
+      PrimExpr new_index = last_dim_index / make_const(last_dim_index.dtype(), info.factor());
+      shuffle_index = me->base;
       indices.Set(indices.size() - 1, new_index);
     }
 
     auto writer = node.CopyOnWrite();
     writer->buffer = RemapBuffer(node->buffer);
     writer->indices = indices;
-
-    return node;
+    return {node, shuffle_index};
   }
 
   PrimExpr VisitExpr_(const BufferLoadNode* op) final {
     auto node = Downcast<BufferLoad>(StmtExprMutator::VisitExpr_(op));
-    auto modified = VisitBufferAccess(node);
+    auto [modified, shuffle_index] = VisitBufferAccess(node);
 
     // Not needed for BufferStoreNode, so we can't just call
     // LegalizeDtype() in VisitBufferAccess.
@@ -1445,13 +1473,18 @@ class VectorTypeRewriter : public StmtExprMutator {
     } else {
       auto writer = modified.CopyOnWrite();
       writer->LegalizeDType();
+      if (shuffle_index >= 0) {
+        return Shuffle::ExtractElement(std::move(modified), shuffle_index);
+      }
       return std::move(modified);
     }
   }
 
   Stmt VisitStmt_(const BufferStoreNode* op) final {
     auto node = Downcast<BufferStore>(StmtExprMutator::VisitStmt_(op));
-    return VisitBufferAccess(std::move(node));
+    auto [modified, shuffle_index] = VisitBufferAccess(std::move(node));
+    ICHECK(shuffle_index < 0);
+    return std::move(modified);
   }
 
   Stmt VisitStmt_(const LetStmtNode* op) final {
@@ -1627,6 +1660,7 @@ class VectorTypeRewriter : public StmtExprMutator {
   bool rewrite_indices_{true};
   std::unordered_map<const VarNode*, RewriteInfo> rewrite_map_;
   std::unordered_map<const BufferNode*, Buffer> buffer_map_;
+  arith::Analyzer analyzer_;
 };
 
 // Rewrite allocates, pointer parameters, and buffer map into vectorized versions
@@ -1635,13 +1669,15 @@ PrimFunc PointerValueTypeRewrite(PrimFunc f, bool allow_untyped_pointers = false
                                  bool rewrite_params = true, bool rewrite_buffer_map = true,
                                  bool rewrite_allocate_node = true, bool rewrite_indices = true,
                                  bool rewrite_let_node = true,
-                                 bool rewrite_allocate_const_node = true) {
-  VectorTypeAccessChecker checker(f->params, f->buffer_map, allow_untyped_pointers);
+                                 bool rewrite_allocate_const_node = true,
+                                 bool rewrite_scalar_read_to_vector_shuffle = true) {
+  VectorTypeAccessChecker checker(f->params, f->buffer_map, allow_untyped_pointers,
+                                  rewrite_scalar_read_to_vector_shuffle);
   checker(f->body);
 
   VectorTypeRewriter rewriter(checker.info_map_, rewrite_params, rewrite_buffer_map,
                               rewrite_allocate_node, rewrite_indices, rewrite_let_node,
-                              rewrite_allocate_const_node);
+                              rewrite_allocate_const_node, rewrite_scalar_read_to_vector_shuffle);
   PrimFuncNode* n = f.CopyOnWrite();
   n->body = rewriter(std::move(n->body));
   rewriter.Finalize(&f);
@@ -1661,7 +1697,8 @@ Pass StorageRewrite() {
     // padded out to 32 bits) would require either rewriting
     // AllocateConst::data, or would require the code generators to
     // handle vectorized constants.
-    return PointerValueTypeRewrite(std::move(f), true, false, false, true, true, true, false);
+    return PointerValueTypeRewrite(std::move(f), true, false, false, true, true, true, false,
+                                   false);
   };
   return CreatePrimFuncPass(pass_func, 0, "tir.StorageRewrite", {});
 }
diff --git a/tests/python/unittest/test_tir_transform_pointer_value_type_rewrite.py b/tests/python/unittest/test_tir_transform_pointer_value_type_rewrite.py
new file mode 100644
index 000000000000..b7e948170a13
--- /dev/null
+++ b/tests/python/unittest/test_tir_transform_pointer_value_type_rewrite.py
@@ -0,0 +1,64 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import tvm
+import tvm.testing
+from tvm import te
+from tvm.driver.build_module import schedule_to_module
+from tvm.script import tir as T
+
+
+class BaseCompare(tvm.testing.CompareBeforeAfter):
+    transform = tvm.tir.transform.PointerValueTypeRewrite()
+
+
+class TestRewriteToShuffle(BaseCompare):
+    @T.prim_func
+    def before(A: T.Buffer((16,), "float32"), B: T.Buffer((4,), "float32")):
+        A_local_data = T.allocate([16], "float32", scope="local")
+        A_local = T.Buffer((16,), "float32", data=A_local_data, scope="local")
+        for i in range(4):
+            A_local[i * 4 : i * 4 + 4] = A[i * 4 : i * 4 + 4]
+        for i in range(4):
+            B[i] = A_local[i * 4] + A_local[i * 4 + 1] + A_local[i * 4 + 2] + A_local[i * 4 + 3]
+
+    @T.prim_func
+    def expected(A: T.Buffer((4,), "float32x4"), B: T.Buffer((4,), "float32")):
+        A_local_data = T.allocate([4], "float32x4", scope="local")
+        A_local = T.Buffer((4,), "float32x4", data=A_local_data, scope="local")
+        for i in range(4):
+            A_local[T.Div(i * 4, 4)] = A[T.Div(i * 4, 4)]
+        for i in range(4):
+            B[i] = (
+                T.Shuffle([A_local[T.Div(i * 4, 4)]], [0])
+                + T.Shuffle([A_local[T.Div(i * 4 + 1, 4)]], [1])
+                + T.Shuffle([A_local[T.Div(i * 4 + 2, 4)]], [2])
+                + T.Shuffle([A_local[T.Div(i * 4 + 3, 4)]], [3])
+            )
+
+
+class TestAddressOf(BaseCompare):
+    @T.prim_func
+    def before(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
+        for i in range(4):
+            T.evaluate(T.address_of(A[i * 4]))
+            B[i * 4 : i * 4 + 4] = A[i * 4 : i * 4 + 4]
+
+    @T.prim_func
+    def expected(A: T.Buffer((16,), "float32"), B: T.Buffer((4,), "float32x4")):
+        for i in range(4):
+            T.evaluate(T.address_of(A[i * 4]))
+            B[T.Div(i * 4, 4)] = A[i * 4 : i * 4 + 4]


Code review comments:
diff_hunk: @@ -1304,6 +1307,18 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
       }
     }
 
+    if (detect_scalar_read_patterns_ && is_buffer_load && indices.size()) {
+      const PrimExpr last_dim_index = indices[indices.size() - 1];
+      if (last_dim_index.dtype().lanes() == 1) {
+        arith::ModularSet me = analyzer_.modular_set(last_dim_index);
+        if (me->coeff > 0) {
+          // When coeff == 0, the index is constant and doesn't need to be recorded since it can
+          // always be rewritten to shuffle.
+          var_info.access_dtype.insert(access_dtype.with_lanes(me->coeff));
path: src/tir/transforms/storage_rewrite.cc
text: @user1:
i think we need a data structure that captures two things: 

```c++
struct VarReadInfo {
   DataType access_dtype;
   // maintained as GCD of all coef of access index
   int64_t simd_coeff;
};
```
This is mainly to ensure that we don't end up use a very large vector here

@author:
access_dtype also contains dtype for write, so the result dtype will be eventually bound by the vectorization lanes of writes

@user1:
the main corner case would be when we do not have writes and only reads. 

@author:
updated

diff_hunk: @@ -1083,17 +1084,14 @@ struct BufferVarInfo {
 
     DataType preferred_base_type = *base_access_dtype.begin();
 
-    // If there is only one vectorizable size used to access the
-    // buffer, and if that access size is compatible with the array
-    // size, then the buffer is vectorizable.  In the future, this
-    // could be improved to allow vectorized buffer access of size
-    // GCD(*lanes_used), if necessary.
     int preferred_lanes = element_dtype.lanes();
-    if ((element_dtype.lanes() == 1) && (access_dtype.size() == 1)) {
+    if (element_dtype.lanes() == 1) {
+      int lanes = access_dtype.begin()->lanes();
+      for (auto dtype : access_dtype) {
+        lanes = std::gcd(lanes, dtype.lanes());
path: src/tir/transforms/storage_rewrite.cc
text: @user1:
use https://github.com/apache/tvm/blob/main/src/arith/int_operator.h#L171