[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 00:55:13 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/4] [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/4] 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/4] 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 e9161514c6c86dfbd482e097cc1b067283ebbe9b 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/4] 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..faa549e14250ba 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();
});
}
More information about the Mlir-commits
mailing list