[Mlir-commits] [mlir] 8a418e5 - [mlir][Affine] Enable fusion of loops with vector loads/stores
Diego Caballero
llvmlistbot at llvm.org
Tue Jun 2 15:28:42 PDT 2020
Author: Diego Caballero
Date: 2020-06-03T01:26:22+03:00
New Revision: 8a418e5f8e89af7197e2d7dc10e39885b05fe0b8
URL: https://github.com/llvm/llvm-project/commit/8a418e5f8e89af7197e2d7dc10e39885b05fe0b8
DIFF: https://github.com/llvm/llvm-project/commit/8a418e5f8e89af7197e2d7dc10e39885b05fe0b8.diff
LOG: [mlir][Affine] Enable fusion of loops with vector loads/stores
This patch enables affine loop fusion for loops with affine vector loads
and stores. For that, we only had to use affine memory op interfaces in
LoopFusionUtils.cpp and Utils.cpp so that vector loads and stores are
also taken into account.
Reviewed By: andydavis1, ftynse
Differential Revision: https://reviews.llvm.org/D80971
Added:
Modified:
mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/lib/Transforms/Utils/LoopFusionUtils.cpp
mlir/lib/Transforms/Utils/Utils.cpp
mlir/test/Transforms/loop-fusion.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td b/mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td
index dd174da44731..a093cb9fd4be 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td
@@ -68,6 +68,19 @@ def AffineReadOpInterface : OpInterface<"AffineReadOpInterface"> {
return op.getAffineMapAttr().getValue();
}]
>,
+ InterfaceMethod<
+ /*desc=*/"Returns the AffineMapAttr associated with 'memref'.",
+ /*retTy=*/"NamedAttribute",
+ /*methodName=*/"getAffineMapAttrForMemRef",
+ /*args=*/(ins "Value":$memref),
+ /*methodBody=*/[{}],
+ /*defaultImplementation=*/[{
+ ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+ assert(memref == getMemRef());
+ return {Identifier::get(op.getMapAttrName(), op.getContext()),
+ op.getAffineMapAttr()};
+ }]
+ >,
];
}
@@ -124,6 +137,19 @@ def AffineWriteOpInterface : OpInterface<"AffineWriteOpInterface"> {
return op.getAffineMapAttr().getValue();
}]
>,
+ InterfaceMethod<
+ /*desc=*/"Returns the AffineMapAttr associated with 'memref'.",
+ /*retTy=*/"NamedAttribute",
+ /*methodName=*/"getAffineMapAttrForMemRef",
+ /*args=*/(ins "Value":$memref),
+ /*methodBody=*/[{}],
+ /*defaultImplementation=*/[{
+ ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+ assert(memref == getMemRef());
+ return {Identifier::get(op.getMapAttrName(), op.getContext()),
+ op.getAffineMapAttr()};
+ }]
+ >,
];
}
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 9329f3dbcbc8..eaf1fada9362 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -389,13 +389,6 @@ class AffineLoadOpBase<string mnemonic, list<OpTrait> traits = []> :
return getAttr(getMapAttrName()).cast<AffineMapAttr>();
}
- /// Returns the AffineMapAttr associated with 'memref'.
- NamedAttribute getAffineMapAttrForMemRef(Value memref) {
- assert(memref == getMemRef());
- return {Identifier::get(getMapAttrName(), getContext()),
- getAffineMapAttr()};
- }
-
static StringRef getMapAttrName() { return "map"; }
}];
}
diff --git a/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp b/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp
index 9ed4283101af..c40e3e26d873 100644
--- a/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp
@@ -38,10 +38,10 @@ using namespace mlir;
static void getLoadAndStoreMemRefAccesses(Operation *opA,
DenseMap<Value, bool> &values) {
opA->walk([&](Operation *op) {
- if (auto loadOp = dyn_cast<AffineLoadOp>(op)) {
+ if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
if (values.count(loadOp.getMemRef()) == 0)
values[loadOp.getMemRef()] = false;
- } else if (auto storeOp = dyn_cast<AffineStoreOp>(op)) {
+ } else if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
values[storeOp.getMemRef()] = true;
}
});
@@ -52,10 +52,10 @@ static void getLoadAndStoreMemRefAccesses(Operation *opA,
// Returns false otherwise.
static bool isDependentLoadOrStoreOp(Operation *op,
DenseMap<Value, bool> &values) {
- if (auto loadOp = dyn_cast<AffineLoadOp>(op)) {
+ if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
return values.count(loadOp.getMemRef()) > 0 &&
values[loadOp.getMemRef()] == true;
- } else if (auto storeOp = dyn_cast<AffineStoreOp>(op)) {
+ } else if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
return values.count(storeOp.getMemRef()) > 0;
}
return false;
@@ -105,7 +105,7 @@ static Operation *getLastDependentOpInRange(Operation *opA, Operation *opB) {
it != Block::reverse_iterator(opA); ++it) {
Operation *opX = &(*it);
opX->walk([&](Operation *op) {
- if (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op)) {
+ if (isa<AffineReadOpInterface>(op) || isa<AffineWriteOpInterface>(op)) {
if (isDependentLoadOrStoreOp(op, values)) {
lastDepOp = opX;
return WalkResult::interrupt();
@@ -179,7 +179,7 @@ gatherLoadsAndStores(AffineForOp forOp,
SmallVectorImpl<Operation *> &loadAndStoreOps) {
bool hasIfOp = false;
forOp.walk([&](Operation *op) {
- if (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op))
+ if (isa<AffineReadOpInterface>(op) || isa<AffineWriteOpInterface>(op))
loadAndStoreOps.push_back(op);
else if (isa<AffineIfOp>(op))
hasIfOp = true;
@@ -464,7 +464,7 @@ bool mlir::getFusionComputeCost(AffineForOp srcForOp, LoopNestStats &srcStats,
unsigned storeCount = 0;
llvm::SmallDenseSet<Value, 4> storeMemrefs;
srcForOp.walk([&](Operation *op) {
- if (auto storeOp = dyn_cast<AffineStoreOp>(op)) {
+ if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
storeMemrefs.insert(storeOp.getMemRef());
++storeCount;
}
@@ -476,7 +476,7 @@ bool mlir::getFusionComputeCost(AffineForOp srcForOp, LoopNestStats &srcStats,
// 'insertPointParent'.
for (auto value : storeMemrefs) {
for (auto *user : value.getUsers()) {
- if (auto loadOp = dyn_cast<AffineLoadOp>(user)) {
+ if (auto loadOp = dyn_cast<AffineReadOpInterface>(user)) {
SmallVector<AffineForOp, 4> loops;
// Check if any loop in loop nest surrounding 'user' is
// 'insertPointParent'.
diff --git a/mlir/lib/Transforms/Utils/Utils.cpp b/mlir/lib/Transforms/Utils/Utils.cpp
index 91ee3b4e3ce6..358d554d7df3 100644
--- a/mlir/lib/Transforms/Utils/Utils.cpp
+++ b/mlir/lib/Transforms/Utils/Utils.cpp
@@ -30,7 +30,7 @@ using namespace mlir;
// Temporary utility: will be replaced when this is modeled through
// side-effects/op traits. TODO(b/117228571)
static bool isMemRefDereferencingOp(Operation &op) {
- if (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op) ||
+ if (isa<AffineReadOpInterface>(op) || isa<AffineWriteOpInterface>(op) ||
isa<AffineDmaStartOp>(op) || isa<AffineDmaWaitOp>(op))
return true;
return false;
@@ -39,8 +39,8 @@ static bool isMemRefDereferencingOp(Operation &op) {
/// Return the AffineMapAttr associated with memory 'op' on 'memref'.
static NamedAttribute getAffineMapAttrForMemRef(Operation *op, Value memref) {
return TypeSwitch<Operation *, NamedAttribute>(op)
- .Case<AffineDmaStartOp, AffineLoadOp, AffinePrefetchOp, AffineStoreOp,
- AffineDmaWaitOp>(
+ .Case<AffineDmaStartOp, AffineReadOpInterface, AffinePrefetchOp,
+ AffineWriteOpInterface, AffineDmaWaitOp>(
[=](auto op) { return op.getAffineMapAttrForMemRef(memref); });
}
diff --git a/mlir/test/Transforms/loop-fusion.mlir b/mlir/test/Transforms/loop-fusion.mlir
index e0eee71e89d9..850071baa832 100644
--- a/mlir/test/Transforms/loop-fusion.mlir
+++ b/mlir/test/Transforms/loop-fusion.mlir
@@ -2464,3 +2464,32 @@ func @reshape_into_matmul(%lhs : memref<1024x1024xf32>,
// MAXIMAL-NEXT: affine.for
// MAXIMAL-NOT: affine.for
// MAXIMAL: return
+
+// -----
+
+// CHECK-LABEL: func @vector_loop
+func @vector_loop(%a : memref<10x20xf32>, %b : memref<10x20xf32>,
+ %c : memref<10x20xf32>) {
+ affine.for %j = 0 to 10 {
+ affine.for %i = 0 to 5 {
+ %ld0 = affine.vector_load %a[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
+ affine.vector_store %ld0, %b[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
+ }
+ }
+
+ affine.for %j = 0 to 10 {
+ affine.for %i = 0 to 5 {
+ %ld0 = affine.vector_load %b[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
+ affine.vector_store %ld0, %c[%j, %i*4] : memref<10x20xf32>, vector<4xf32>
+ }
+ }
+
+ return
+}
+// CHECK: affine.for
+// CHECK-NEXT: affine.for
+// CHECK-NEXT: affine.vector_load
+// CHECK-NEXT: affine.vector_store
+// CHECK-NEXT: affine.vector_load
+// CHECK-NEXT: affine.vector_store
+// CHECK-NOT: affine.for
More information about the Mlir-commits
mailing list