[Mlir-commits] [mlir] eff269f - [MLIR][Affine][LICM] Mark users of `iter_args` variant

[Uday Bondhugula]

Author: Vinayaka Bandishti
Date: 2021-05-25T15:56:52+05:30
New Revision: eff269fc9f8b8c79e08e8295aa22da8950bcc341

URL: https://github.com/llvm/llvm-project/commit/eff269fc9f8b8c79e08e8295aa22da8950bcc341
DIFF: https://github.com/llvm/llvm-project/commit/eff269fc9f8b8c79e08e8295aa22da8950bcc341.diff

LOG: [MLIR][Affine][LICM] Mark users of `iter_args` variant

Prevent users of `iter_args` of an affine for loop from being hoisted
out of it. Otherwise, LICM leads to a violation of the SSA dominance
(as demonstrated in the added test case).

Fixes: https://bugs.llvm.org/show_bug.cgi?id=50103

Reviewed By: bondhugula, ayzhuang

Differential Revision: https://reviews.llvm.org/D102984

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 a16eac6f16a28..86fb6f83b6378 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
@@ -49,30 +49,32 @@ struct LoopInvariantCodeMotion
 } // end anonymous namespace
 
 static bool
-checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
+checkInvarianceOfNestedIfOps(Operation *op, Value indVar, ValueRange iterArgs,
                              SmallPtrSetImpl<Operation *> &opsWithUsers,
                              SmallPtrSetImpl<Operation *> &opsToHoist);
-static bool isOpLoopInvariant(Operation &op, Value indVar,
+static bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs,
                               SmallPtrSetImpl<Operation *> &opsWithUsers,
                               SmallPtrSetImpl<Operation *> &opsToHoist);
 
 static bool
 areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar,
+                                 ValueRange iterArgs,
                                  SmallPtrSetImpl<Operation *> &opsWithUsers,
                                  SmallPtrSetImpl<Operation *> &opsToHoist);
 
 // Returns true if the individual op is loop invariant.
-bool isOpLoopInvariant(Operation &op, Value indVar,
+bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs,
                        SmallPtrSetImpl<Operation *> &opsWithUsers,
                        SmallPtrSetImpl<Operation *> &opsToHoist) {
   LLVM_DEBUG(llvm::dbgs() << "iterating on op: " << op;);
 
   if (isa<AffineIfOp>(op)) {
-    if (!checkInvarianceOfNestedIfOps(&op, indVar, opsWithUsers, opsToHoist)) {
+    if (!checkInvarianceOfNestedIfOps(&op, indVar, iterArgs, opsWithUsers,
+                                      opsToHoist)) {
       return false;
     }
   } else if (auto forOp = dyn_cast<AffineForOp>(op)) {
-    if (!areAllOpsInTheBlockListInvariant(forOp.getLoopBody(), indVar,
+    if (!areAllOpsInTheBlockListInvariant(forOp.getLoopBody(), indVar, iterArgs,
                                           opsWithUsers, opsToHoist)) {
       return false;
     }
@@ -129,6 +131,12 @@ bool isOpLoopInvariant(Operation &op, Value indVar,
       return false;
     }
 
+    // If the one of the iter_args is the operand, this op isn't loop invariant.
+    if (llvm::is_contained(iterArgs, op.getOperand(i))) {
+      LLVM_DEBUG(llvm::dbgs() << "\nOne of the iter_args is the operand\n");
+      return false;
+    }
+
     if (operandSrc != nullptr) {
       LLVM_DEBUG(llvm::dbgs() << *operandSrc << "\nIterating on operand src\n");
 
@@ -148,12 +156,13 @@ 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 *> &opsWithUsers,
+    Region &blockList, Value indVar, ValueRange iterArgs,
+    SmallPtrSetImpl<Operation *> &opsWithUsers,
     SmallPtrSetImpl<Operation *> &opsToHoist) {
 
   for (auto &b : blockList) {
     for (auto &op : b) {
-      if (!isOpLoopInvariant(op, indVar, opsWithUsers, opsToHoist)) {
+      if (!isOpLoopInvariant(op, indVar, iterArgs, opsWithUsers, opsToHoist)) {
         return false;
       }
     }
@@ -164,18 +173,19 @@ bool areAllOpsInTheBlockListInvariant(
 
 // Returns true if the affine.if op can be hoisted.
 bool checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
+                                  ValueRange iterArgs,
                                   SmallPtrSetImpl<Operation *> &opsWithUsers,
                                   SmallPtrSetImpl<Operation *> &opsToHoist) {
   assert(isa<AffineIfOp>(op));
   auto ifOp = cast<AffineIfOp>(op);
 
-  if (!areAllOpsInTheBlockListInvariant(ifOp.thenRegion(), indVar, opsWithUsers,
-                                        opsToHoist)) {
+  if (!areAllOpsInTheBlockListInvariant(ifOp.thenRegion(), indVar, iterArgs,
+                                        opsWithUsers, opsToHoist)) {
     return false;
   }
 
-  if (!areAllOpsInTheBlockListInvariant(ifOp.elseRegion(), indVar, opsWithUsers,
-                                        opsToHoist)) {
+  if (!areAllOpsInTheBlockListInvariant(ifOp.elseRegion(), indVar, iterArgs,
+                                        opsWithUsers, opsToHoist)) {
     return false;
   }
 
@@ -185,6 +195,7 @@ bool checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
 void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
   auto *loopBody = forOp.getBody();
   auto indVar = forOp.getInductionVar();
+  ValueRange iterArgs = forOp.getIterOperands();
 
   // This is the place where hoisted instructions would reside.
   OpBuilder b(forOp.getOperation());
@@ -200,7 +211,7 @@ void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
     if (!op.use_empty())
       opsWithUsers.insert(&op);
     if (!isa<AffineYieldOp>(op)) {
-      if (isOpLoopInvariant(op, indVar, opsWithUsers, opsToHoist)) {
+      if (isOpLoopInvariant(op, indVar, iterArgs, 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 2203bffe3ddb8..dce934d6e5bb5 100644
--- a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
@@ -716,3 +716,20 @@ func @affine_for_not_invariant(%in : memref<30x512xf32, 1>,
 // CHECK-NEXT:  }
 // CHECK-NEXT:  mulf
 // CHECK-NEXT:  affine.store
+
+// -----
+
+// CHECK-LABEL: func @use_of_iter_args_not_invariant
+func @use_of_iter_args_not_invariant(%m : memref<10xindex>) {
+  %sum_1 = constant 0 : index
+  %v0 = affine.for %arg1 = 0 to 11 iter_args (%prevAccum = %sum_1) -> index {
+    %newAccum = addi %prevAccum, %sum_1 : index
+    affine.yield %newAccum : index
+  }
+  return
+}
+
+// CHECK:       constant
+// CHECK-NEXT:  affine.for
+// CHECK-NEXT:  addi
+// CHECK-NEXT:  affine.yield