[Feature Request] Lower Vectorized Loop Pass should be enhanced to adapt layout inference

[LeiWang1999]

When utilize TileLang, some layout transformation like swizzling or padding will implicitly apply layout transformation, though this approach is efficient and powerful, but sometimes will lead to a crash for vectorization.

Considering dequantize gemm on volta:

for v in T.vectorized(0, 8):
        index = i * threads * local_size + tx * local_size + v
        vi = index // block_K
        vj = index % block_K
        B_dequantize_shared[vi, vj] = B_dequantize_local[v]

On Volta, applying a swizzle operation will adjust the memory layout to align with groups of 4 elements instead of 8 elements. This optimization enhances memory coalescing and data locality for efficient GPU execution.

We should enhance lower vectorize pass to automatically convert the vectorize stage into:

for ov in T.serial(0, local_size // 4):
      for iv in T.vectorized(0, 4):
          index = (
              i * threads * local_size
              + tx * local_size
              + ov * 4
              + iv
          )
          vi = index // block_K
          vj = index % block_K
          B_dequantize_shared[vi, vj] = B_dequantize_local[ov * 4 + iv]

[LeiWang1999]

PR #255 also introduced a new pass LoopVectorizedLegalizer that we allow smart rewrite to a given vectorize loop:

class LoopVectorizedLegalizer : IRMutatorWithAnalyzer {
 public:
  static PrimFunc Substitute(PrimFunc f) {
    arith::Analyzer analyzer;
    LoopVectorizedLegalizer substituter(&analyzer);
    PrimFuncNode* fptr = f.CopyOnWrite();
    fptr->body = substituter.VisitStmt(f->body);
    return f;
  }

 private:
  LoopVectorizedLegalizer(arith::Analyzer* analyzer) : arith::IRMutatorWithAnalyzer(analyzer) {};

  Stmt VisitStmt_(const ForNode* op) final {
    For for_node = Downcast<For>(IRMutatorWithAnalyzer::VisitStmt_(op));
    if (for_node->kind != ForKind::kVectorized) {
      return IRMutatorWithAnalyzer::VisitStmt_(op);
    }
    for_node.CopyOnWrite()->kind = ForKind::kSerial;
    return VectorizeLoop(for_node);
  }
};

Think the pass name should be replaced with VectorizedLoopSmartRewrite.

[LeiWang1999]

closed as has been resolved.