[Mlir-commits] [mlir] [mlir][scf] Rewrite vector.transfer_read/write after peeling (PR #88684)

[llvmlistbot at llvm.org]

================
@@ -166,6 +167,52 @@ static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp,
   return success();
 }
 
+static void rewriteVectorReadWriteToLoadStore(RewriterBase &b, Operation *op) {
+  b.setInsertionPoint(op);
+  if (auto write = dyn_cast<vector::TransferWriteOp>(op)) {
+    b.replaceOpWithNewOp<vector::StoreOp>(
+        op, write.getVector(), write.getSource(), write.getIndices());
+  } else if (auto read = dyn_cast<vector::TransferReadOp>(op)) {
+    b.replaceOpWithNewOp<vector::LoadOp>(op, read.getVectorType(),
+                                         read.getSource(), read.getIndices());
+  }
+}
+
+static bool hasVectorSizeEqualToStep(Operation *Op,
+                                     std::optional<int64_t> step) {
+  if (!step)
+    return false;
+
+  if (isa<vector::TransferWriteOp, vector::TransferReadOp>(Op)) {
+    auto vectorType = isa<vector::TransferWriteOp>(Op)
+                          ? cast<vector::TransferWriteOp>(Op).getVectorType()
+                          : cast<vector::TransferReadOp>(Op).getVectorType();
+
+    if (vectorType.getRank() != 1)
+      return false;
+
+    auto vectorSize = vectorType.getShape()[0];
+    if (vectorSize == *step)
+      return true;
+  }
+
+  return false;
+}
+
+static void rewriteVectorizedLoopAfterPeeling(RewriterBase &rewriter,
+                                              ForOp forOp) {
+  auto stepInt = getConstantIntValue(forOp.getStep());
+
+  forOp.walk([&](Operation *affineOp) {
----------------
ShivaChen wrote:

Is it ok to change to op?

https://github.com/llvm/llvm-project/pull/88684