[Mlir-commits] [mlir] f56791a - [mlir] Prevent operations with users from being hoisted
Sergei Grechanik
llvmlistbot at llvm.org
Tue Apr 13 15:37:59 PDT 2021
Author: Sumesh Udayakumaran
Date: 2021-04-13T15:29:17-07:00
New Revision: f56791ae2ea081636d0fc1dd3b1b44865145c7cb
URL: https://github.com/llvm/llvm-project/commit/f56791ae2ea081636d0fc1dd3b1b44865145c7cb
DIFF: https://github.com/llvm/llvm-project/commit/f56791ae2ea081636d0fc1dd3b1b44865145c7cb.diff
LOG: [mlir] Prevent operations with users from being hoisted
This patch collects operations that have users in a for loop and uses
them when loop invariant operations are detected and hoisted.
Reviewed By: bondhugula, vinayaka-polymage
Differential Revision: https://reviews.llvm.org/D99761
Added:
Modified:
mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
index f52a7a2c5bf8a..45510703ff7fe 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
@@ -50,25 +50,25 @@ struct LoopInvariantCodeMotion
static bool
checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
- SmallPtrSetImpl<Operation *> &definedOps,
+ SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist);
static bool isOpLoopInvariant(Operation &op, Value indVar,
- SmallPtrSetImpl<Operation *> &definedOps,
+ SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist);
static bool
areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar,
- SmallPtrSetImpl<Operation *> &definedOps,
+ SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist);
// Returns true if the individual op is loop invariant.
bool isOpLoopInvariant(Operation &op, Value indVar,
- SmallPtrSetImpl<Operation *> &definedOps,
+ SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist) {
LLVM_DEBUG(llvm::dbgs() << "iterating on op: " << op;);
if (isa<AffineIfOp>(op)) {
- if (!checkInvarianceOfNestedIfOps(&op, indVar, definedOps, opsToHoist)) {
+ if (!checkInvarianceOfNestedIfOps(&op, indVar, opsWithUsers, opsToHoist)) {
return false;
}
} else if (isa<AffineForOp>(op)) {
@@ -79,11 +79,8 @@ bool isOpLoopInvariant(Operation &op, Value indVar,
// TODO: Support DMA ops.
return false;
} else if (!isa<ConstantOp>(op)) {
- // Register op in the set of ops defined inside the loop. This set is used
- // to prevent hoisting ops that depend on other ops defined inside the loop
- // which are themselves not being hoisted.
- definedOps.insert(&op);
-
+ // Register op in the set of ops that have users.
+ opsWithUsers.insert(&op);
if (isa<AffineMapAccessInterface>(op)) {
Value memref = isa<AffineReadOpInterface>(op)
? cast<AffineReadOpInterface>(op).getMemRef()
@@ -135,7 +132,8 @@ bool isOpLoopInvariant(Operation &op, Value indVar,
// If the value was defined in the loop (outside of the
// if/else region), and that operation itself wasn't meant to
// be hoisted, then mark this operation loop dependent.
- if (definedOps.count(operandSrc) && opsToHoist.count(operandSrc) == 0) {
+ if (opsWithUsers.count(operandSrc) &&
+ opsToHoist.count(operandSrc) == 0) {
return false;
}
}
@@ -149,12 +147,12 @@ bool isOpLoopInvariant(Operation &op, Value indVar,
// Checks if all ops in a region (i.e. list of blocks) are loop invariant.
bool areAllOpsInTheBlockListInvariant(
- Region &blockList, Value indVar, SmallPtrSetImpl<Operation *> &definedOps,
+ Region &blockList, Value indVar, SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist) {
for (auto &b : blockList) {
for (auto &op : b) {
- if (!isOpLoopInvariant(op, indVar, definedOps, opsToHoist)) {
+ if (!isOpLoopInvariant(op, indVar, opsWithUsers, opsToHoist)) {
return false;
}
}
@@ -165,17 +163,17 @@ bool areAllOpsInTheBlockListInvariant(
// Returns true if the affine.if op can be hoisted.
bool checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
- SmallPtrSetImpl<Operation *> &definedOps,
+ SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist) {
assert(isa<AffineIfOp>(op));
auto ifOp = cast<AffineIfOp>(op);
- if (!areAllOpsInTheBlockListInvariant(ifOp.thenRegion(), indVar, definedOps,
+ if (!areAllOpsInTheBlockListInvariant(ifOp.thenRegion(), indVar, opsWithUsers,
opsToHoist)) {
return false;
}
- if (!areAllOpsInTheBlockListInvariant(ifOp.elseRegion(), indVar, definedOps,
+ if (!areAllOpsInTheBlockListInvariant(ifOp.elseRegion(), indVar, opsWithUsers,
opsToHoist)) {
return false;
}
@@ -187,18 +185,23 @@ void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
auto *loopBody = forOp.getBody();
auto indVar = forOp.getInductionVar();
- SmallPtrSet<Operation *, 8> definedOps;
// This is the place where hoisted instructions would reside.
OpBuilder b(forOp.getOperation());
SmallPtrSet<Operation *, 8> opsToHoist;
SmallVector<Operation *, 8> opsToMove;
+ SmallPtrSet<Operation *, 8> opsWithUsers;
for (auto &op : *loopBody) {
+ // Register op in the set of ops that have users. This set is used
+ // to prevent hoisting ops that depend on these ops that are
+ // not being hoisted.
+ if (!op.use_empty())
+ opsWithUsers.insert(&op);
// We don't hoist for loops.
if (!isa<AffineForOp>(op)) {
if (!isa<AffineYieldOp>(op)) {
- if (isOpLoopInvariant(op, indVar, definedOps, opsToHoist)) {
+ if (isOpLoopInvariant(op, indVar, opsWithUsers, opsToHoist)) {
opsToMove.push_back(&op);
}
}
diff --git a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
index 0b8c036bbc294..b7bf7ac2e7c25 100644
--- a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
@@ -613,3 +613,68 @@ func @vector_loop_all_invariant() {
// CHECK-NEXT: addf
// CHECK-NEXT: affine.vector_store
// CHECK-NEXT: affine.for
+
+// -----
+
+#set = affine_set<(d0): (d0 - 10 >= 0)>
+// CHECK-LABEL: func @affine_if_not_invariant(
+func @affine_if_not_invariant(%buffer: memref<1024xf32>) -> f32 {
+ %sum_init_0 = constant 0.0 : f32
+ %sum_init_1 = constant 1.0 : f32
+ %res = affine.for %i = 0 to 10 step 2 iter_args(%sum_iter = %sum_init_0) -> f32 {
+ %t = affine.load %buffer[%i] : memref<1024xf32>
+ %sum_next = affine.if #set(%i) -> (f32) {
+ %new_sum = addf %sum_iter, %t : f32
+ affine.yield %new_sum : f32
+ } else {
+ affine.yield %sum_iter : f32
+ }
+ %modified_sum = addf %sum_next, %sum_init_1 : f32
+ affine.yield %modified_sum : f32
+ }
+ return %res : f32
+}
+
+// CHECK: constant 0.000000e+00 : f32
+// CHECK-NEXT: constant 1.000000e+00 : f32
+// CHECK-NEXT: affine.for
+// CHECK-NEXT: affine.load
+// CHECK-NEXT: affine.if
+// CHECK-NEXT: addf
+// CHECK-NEXT: affine.yield
+// CHECK-NEXT: } else {
+// CHECK-NEXT: affine.yield
+// CHECK-NEXT: }
+// CHECK-NEXT: addf
+// CHECK-NEXT: affine.yield
+// CHECK-NEXT: }
+
+// -----
+
+// CHECK-LABEL: func @affine_for_not_invariant(
+func @affine_for_not_invariant(%in : memref<30x512xf32, 1>,
+ %out : memref<30x1xf32, 1>) {
+ %sum_0 = constant 0.0 : f32
+ %cst_0 = constant 1.1 : f32
+ affine.for %j = 0 to 30 {
+ %sum = affine.for %i = 0 to 512 iter_args(%sum_iter = %sum_0) -> (f32) {
+ %t = affine.load %in[%j,%i] : memref<30x512xf32,1>
+ %sum_next = addf %sum_iter, %t : f32
+ affine.yield %sum_next : f32
+ }
+ %mod_sum = mulf %sum, %cst_0 : f32
+ affine.store %mod_sum, %out[%j, 0] : memref<30x1xf32, 1>
+ }
+ return
+}
+
+// CHECK: constant 0.000000e+00 : f32
+// CHECK-NEXT: constant 1.100000e+00 : f32
+// CHECK-NEXT: affine.for
+// CHECK-NEXT: affine.for
+// CHECK-NEXT: affine.load
+// CHECK-NEXT: addf
+// CHECK-NEXT: affine.yield
+// CHECK-NEXT: }
+// CHECK-NEXT: mulf
+// CHECK-NEXT: affine.store
More information about the Mlir-commits
mailing list