[Affine][LICM] Mark users of `iter_args` variant

[vinayaka-polymage]

Summary

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

[bondhugula]

LGTM - some minor comments on compactness.

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

53

Operation::operand_range aren't normally used on methods/API - just ValueRange should do - more compact/cleaner.

57

Likewise.

138

auto -> Value

138–142

You can use llvm::is_contained to encode this in one line.

140

Nit: terminate with a full stop.

203

auto -> ValueRange should work I think.

[vinayaka-polymage]

Thanks for the suggestions for compactness, all of them worked well.

[vinayaka-polymage]

Incorporated all the suggestions, thanks @bondhugula !

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

140

Not adding a full stop to be consistent with other places. Sounds good ?

[bondhugula]

LGTM.

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

140

Sounds fine.

[ayzhuang]

Looks great! Thanks.

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

struct LoopInvariantCodeMotion
     : public AffineLoopInvariantCodeMotionBase<LoopInvariantCodeMotion> {
   void runOnFunction() override;
   void runOnAffineForOp(AffineForOp forOp);
 };
 } // end anonymous namespace
 
 static bool
 checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
+                             Operation::operand_range iterArgs,
                              SmallPtrSetImpl<Operation *> &opsWithUsers,
                              SmallPtrSetImpl<Operation *> &opsToHoist);
 static bool isOpLoopInvariant(Operation &op, Value indVar,
+                              Operation::operand_range iterArgs,
                               SmallPtrSetImpl<Operation *> &opsWithUsers,
                               SmallPtrSetImpl<Operation *> &opsToHoist);
 
 static bool
 areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar,
+                                 Operation::operand_range iterArgs,
                                  SmallPtrSetImpl<Operation *> &opsWithUsers,
                                  SmallPtrSetImpl<Operation *> &opsToHoist);
 
 // Returns true if the individual op is loop invariant.
 bool isOpLoopInvariant(Operation &op, Value indVar,
+                       Operation::operand_range 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;
     }
   } else if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) {
     // TODO: Support DMA ops.
     return false;
   } else if (!isa<ConstantOp>(op)) {
     // Register op in the set of ops that have users.
LLVM_DEBUG(
         op.getOperand(i).print(llvm::dbgs() << "\nIterating on operand\n"));
 
     // If the loop IV is the operand, this op isn't loop invariant.
     if (indVar == op.getOperand(i)) {
       LLVM_DEBUG(llvm::dbgs() << "\nLoop IV is the operand\n");
       return false;
     }
 
+    // If the one of the iter_args is the operand, this op isn't loop invariant.
+    for (auto iterArg : iterArgs) {
+      if (iterArg == 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");
 
       // 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 (opsWithUsers.count(operandSrc) && opsToHoist.count(operandSrc) == 0) {
         return false;
       }
     }
   }
 
   // If no operand was loop variant, mark this op for motion.
   opsToHoist.insert(&op);
   return true;
 }
 
 // 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, Operation::operand_range 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;
       }
     }
   }
 
   return true;
 }
 
 // Returns true if the affine.if op can be hoisted.
 bool checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
+                                  Operation::operand_range 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;
   }
 
   return true;
 }
 
 void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
   auto *loopBody = forOp.getBody();
   auto indVar = forOp.getInductionVar();
+  auto iterArgs = forOp.getIterOperands();
 
   // 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);
     if (!isa<AffineYieldOp>(op)) {
-      if (isOpLoopInvariant(op, indVar, opsWithUsers, opsToHoist)) {
+      if (isOpLoopInvariant(op, indVar, iterArgs, opsWithUsers, opsToHoist)) {
         opsToMove.push_back(&op);
       }
     }
   }
 
   // For all instructions that we found to be invariant, place sequentially
   // right before the for loop.
   for (auto *op : opsToMove) {

mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir

 // 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
+
+// -----
+
+// 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