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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Mon Apr 15 01:13:48 PDT 2024


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

>From a37d9207ad9b286dba461f1aa53209d3cd31545e Mon Sep 17 00:00:00 2001
From: Shiva Chen <shiva.chen at imgtec.com>
Date: Wed, 3 Apr 2024 08:02:03 +0100
Subject: [PATCH 1/5] [mlir][scf] Rewrite vector.transfer_read/write after
 peeling

After peeling, the loop iteration will be multiple of step.
If the vector size of vector.transfer_read/write is equal to step in the
peeled loop, there won't be the remaining iteration smaller than the
vector size.

In this case, rewriting vector.transfer_read/write to vector.load/store
could avoid generating masks when lowering to LLVM Dialect.
---
 .../SCF/Transforms/LoopSpecialization.cpp     | 49 +++++++++++++++++++
 .../for-loop-peeling-vector-load-store.mlir   | 30 ++++++++++++
 2 files changed, 79 insertions(+)
 create mode 100644 mlir/test/Dialect/SCF/for-loop-peeling-vector-load-store.mlir

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index a30e349d49136c..4eb757d618a98d 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -20,6 +20,7 @@
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/PatternMatch.h"
@@ -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) {
+    if (!isa<vector::TransferWriteOp, vector::TransferReadOp>(affineOp))
+      return WalkResult::advance();
+    if (!hasVectorSizeEqualToStep(affineOp, stepInt))
+      return WalkResult::advance();
+    rewriteVectorReadWriteToLoadStore(rewriter, affineOp);
+    return WalkResult::advance();
+  });
+}
+
 static void rewriteAffineOpAfterPeeling(RewriterBase &rewriter, ForOp forOp,
                                         ForOp partialIteration,
                                         Value previousUb) {
@@ -200,6 +247,8 @@ LogicalResult mlir::scf::peelForLoopAndSimplifyBounds(RewriterBase &rewriter,
   if (failed(peelForLoop(rewriter, forOp, partialIteration, splitBound)))
     return failure();
 
+  rewriteVectorizedLoopAfterPeeling(rewriter, forOp);
+
   // Rewrite affine.min and affine.max ops.
   rewriteAffineOpAfterPeeling(rewriter, forOp, partialIteration, previousUb);
 
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling-vector-load-store.mlir b/mlir/test/Dialect/SCF/for-loop-peeling-vector-load-store.mlir
new file mode 100644
index 00000000000000..04991930a2c262
--- /dev/null
+++ b/mlir/test/Dialect/SCF/for-loop-peeling-vector-load-store.mlir
@@ -0,0 +1,30 @@
+// RUN: mlir-opt %s -scf-for-loop-peeling -canonicalize -verify-diagnostics | FileCheck %s
+
+func.func @vector_read_write(%a : memref<100xi32>, %b : memref<100xi32>, %ub: index) {
+// %LB to %NEW_UB will be multiple of STEP after peeling.
+// So vector.transfer_write could be transferred to vector.store to avoid
+// generating mask when lowering to LLVM.
+//
+//  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> ((s0 floordiv 64) * 64)>
+//      CHECK: func @vector_read_write(
+// CHECK-SAME:   %[[A:.*]]: memref<100xi32>, %[[B:.*]]: memref<100xi32>, %[[UB:.*]]: index
+//      CHECK:   %[[LB:.*]] = arith.constant 0 : index
+//      CHECK:   %[[STEP:.*]] = arith.constant 64 : index
+//      CHECK:   %[[NEW_UB:.*]] = affine.apply #[[MAP0]]
+//      CHECK:   scf.for %[[IV:.*]] = %[[LB]] to %[[NEW_UB]] step %[[STEP]] {
+//      CHECK:     %[[VAL:.*]] = vector.load %[[B]][%[[IV]]]
+//      CHECK:     vector.store %[[VAL]], %[[A]][%[[IV]]]
+//      CHECK:   }
+//      CHECK:   scf.for %[[IV:.*]] = %[[NEW_UB]] to %[[UB]] step %[[STEP]] {
+//      CHECK:     %[[VAL:.*]] = vector.transfer_read %[[B]][%[[IV]]]
+//      CHECK:     vector.transfer_write %[[VAL]], %[[A]][%[[IV]]]
+//      CHECK:   }
+  %c0 = arith.constant 0 : index
+  %c64 = arith.constant 64 : index
+  %pad = arith.constant 0 : i32
+  scf.for %i = %c0 to %ub step %c64 {
+    %val = vector.transfer_read %b[%i], %pad : memref<100xi32>, vector<64xi32>
+    vector.transfer_write %val, %a[%i] : vector<64xi32>, memref<100xi32>
+  }
+  return
+}

>From 90098006943f89d384f7ae81ec786c8356776b94 Mon Sep 17 00:00:00 2001
From: Shiva Chen <shiva.chen at imgtec.com>
Date: Mon, 15 Apr 2024 08:09:02 +0100
Subject: [PATCH 2/5] Set InBounds for vector read/write

---
 .../SCF/Transforms/LoopSpecialization.cpp       | 17 +++++++----------
 .../SCF/for-loop-peeling-vector-load-store.mlir |  5 +++--
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 4eb757d618a98d..399780009225a6 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -167,15 +167,12 @@ 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());
-  }
+// set InBounds attribute to lower vector read/write to vector load/store.
+static void setInBoundsForVectorReadWrite(RewriterBase &b, Operation *op) {
+  if (auto write = dyn_cast<vector::TransferWriteOp>(op))
+    write.setInBoundsAttr(b.getBoolArrayAttr({true}));
+  else if (auto read = dyn_cast<vector::TransferReadOp>(op))
+    read.setInBoundsAttr(b.getBoolArrayAttr({true}));
 }
 
 static bool hasVectorSizeEqualToStep(Operation *Op,
@@ -208,7 +205,7 @@ static void rewriteVectorizedLoopAfterPeeling(RewriterBase &rewriter,
       return WalkResult::advance();
     if (!hasVectorSizeEqualToStep(affineOp, stepInt))
       return WalkResult::advance();
-    rewriteVectorReadWriteToLoadStore(rewriter, affineOp);
+    setInBoundsForVectorReadWrite(rewriter, affineOp);
     return WalkResult::advance();
   });
 }
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling-vector-load-store.mlir b/mlir/test/Dialect/SCF/for-loop-peeling-vector-load-store.mlir
index 04991930a2c262..1dd032d9ffbe15 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling-vector-load-store.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling-vector-load-store.mlir
@@ -10,10 +10,11 @@ func.func @vector_read_write(%a : memref<100xi32>, %b : memref<100xi32>, %ub: in
 // CHECK-SAME:   %[[A:.*]]: memref<100xi32>, %[[B:.*]]: memref<100xi32>, %[[UB:.*]]: index
 //      CHECK:   %[[LB:.*]] = arith.constant 0 : index
 //      CHECK:   %[[STEP:.*]] = arith.constant 64 : index
+//      CHECK:   %[[PAD:.*]] = arith.constant 0 : i32
 //      CHECK:   %[[NEW_UB:.*]] = affine.apply #[[MAP0]]
 //      CHECK:   scf.for %[[IV:.*]] = %[[LB]] to %[[NEW_UB]] step %[[STEP]] {
-//      CHECK:     %[[VAL:.*]] = vector.load %[[B]][%[[IV]]]
-//      CHECK:     vector.store %[[VAL]], %[[A]][%[[IV]]]
+//      CHECK:     %[[VAL:.*]] = vector.transfer_read %[[B]][%[[IV]]], %[[PAD]] {in_bounds = [true]}
+//      CHECK:     vector.transfer_write %[[VAL]], %[[A]][%[[IV]]] {in_bounds = [true]}
 //      CHECK:   }
 //      CHECK:   scf.for %[[IV:.*]] = %[[NEW_UB]] to %[[UB]] step %[[STEP]] {
 //      CHECK:     %[[VAL:.*]] = vector.transfer_read %[[B]][%[[IV]]]

>From c4bde6787ae024ef5f90355e953160de67716099 Mon Sep 17 00:00:00 2001
From: Shiva Chen <shiva.chen at imgtec.com>
Date: Mon, 15 Apr 2024 08:10:45 +0100
Subject: [PATCH 3/5] change variable name to lower case

---
 mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 399780009225a6..9fd69b4ac97737 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -175,15 +175,15 @@ static void setInBoundsForVectorReadWrite(RewriterBase &b, Operation *op) {
     read.setInBoundsAttr(b.getBoolArrayAttr({true}));
 }
 
-static bool hasVectorSizeEqualToStep(Operation *Op,
+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 (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;

>From b4e5dc78e9eb2a046e1b7f8a3701a56ddd64ad47 Mon Sep 17 00:00:00 2001
From: Shiva Chen <shiva.chen at imgtec.com>
Date: Mon, 15 Apr 2024 08:17:27 +0100
Subject: [PATCH 4/5] Refine varaible name

---
 mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 9fd69b4ac97737..31fbc75304613a 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -200,12 +200,12 @@ static void rewriteVectorizedLoopAfterPeeling(RewriterBase &rewriter,
                                               ForOp forOp) {
   auto stepInt = getConstantIntValue(forOp.getStep());
 
-  forOp.walk([&](Operation *affineOp) {
-    if (!isa<vector::TransferWriteOp, vector::TransferReadOp>(affineOp))
+  forOp.walk([&](Operation *op) {
+    if (!isa<vector::TransferWriteOp, vector::TransferReadOp>(op))
       return WalkResult::advance();
-    if (!hasVectorSizeEqualToStep(affineOp, stepInt))
+    if (!hasVectorSizeEqualToStep(op, stepInt))
       return WalkResult::advance();
-    setInBoundsForVectorReadWrite(rewriter, affineOp);
+    setInBoundsForVectorReadWrite(rewriter, op);
     return WalkResult::advance();
   });
 }

>From 1f8f8111147ab764047c8dd8f952bd1c1e886d7d Mon Sep 17 00:00:00 2001
From: Shiva Chen <shiva.chen at imgtec.com>
Date: Mon, 15 Apr 2024 09:12:52 +0100
Subject: [PATCH 5/5] check the offset is induction variable

---
 mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 31fbc75304613a..9c1a8443898b65 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -175,21 +175,25 @@ static void setInBoundsForVectorReadWrite(RewriterBase &b, Operation *op) {
     read.setInBoundsAttr(b.getBoolArrayAttr({true}));
 }
 
-static bool hasVectorSizeEqualToStep(Operation *op,
+static bool hasVectorSizeEqualToStep(ForOp forOp, Operation *op,
                                      std::optional<int64_t> step) {
   if (!step)
     return false;
 
   if (isa<vector::TransferWriteOp, vector::TransferReadOp>(op)) {
+    Value IV = forOp.getInductionVar();
     auto vectorType = isa<vector::TransferWriteOp>(op)
                           ? cast<vector::TransferWriteOp>(op).getVectorType()
                           : cast<vector::TransferReadOp>(op).getVectorType();
+    auto indices = isa<vector::TransferWriteOp>(op)
+                          ? cast<vector::TransferWriteOp>(op).getIndices()
+                          : cast<vector::TransferReadOp>(op).getIndices();
 
     if (vectorType.getRank() != 1)
       return false;
 
     auto vectorSize = vectorType.getShape()[0];
-    if (vectorSize == *step)
+    if ((vectorSize == *step) && (indices[indices.size() - 1] == IV))
       return true;
   }
 
@@ -203,7 +207,7 @@ static void rewriteVectorizedLoopAfterPeeling(RewriterBase &rewriter,
   forOp.walk([&](Operation *op) {
     if (!isa<vector::TransferWriteOp, vector::TransferReadOp>(op))
       return WalkResult::advance();
-    if (!hasVectorSizeEqualToStep(op, stepInt))
+    if (!hasVectorSizeEqualToStep(forOp, op, stepInt))
       return WalkResult::advance();
     setInBoundsForVectorReadWrite(rewriter, op);
     return WalkResult::advance();



More information about the Mlir-commits mailing list