[mlir] Prevent operations with users from being hoisted

[sumesh13]

Summary

This patch collects operations that have users in a for loop and uses them when loop invariant operations are detected and hoisted.

[bondhugula]

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

187–188

This choice of name definedOps appears to be confusing: there is no such thing as ops being defined. It would be good rename this and add a comment while on this (since you are adding to definedOps below) to make this readable.

197

Nit: op.getNumResults() > 0 if you prefer.

[sumesh13]

1. Address code review comments including renaming definedOps to opsWithusers

[sumesh13]

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

187–188

I have renamed it to opsWithUsers

[sumesh13]

Gentle ping

[bondhugula]

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

197

Just because the op has results doesn't mean those results have users. You can do a check in O(1) time to determine if the results have no uses (use_empty()). This will make the check more powerful and you'll still be able to hoist those that have results but are ultimately unused.

[vinayaka-polymage]

Thank you for bringing up this issue in LICM, it is an important one. I have a few comments on the fix:

1. Inherent issue is definedOps ( going to be renamed to opsWithUsers) is being populated only under some conditions, and does not actually do the job of collecting all ops that define an SSA value in that block.
2. Even after this fix, only AffineForOps will be added to the list, but there are other cases like AffineDma... ops, AffineIfOps that are actually NOT code invariant etc. which will still cause similar problems.

So, I suggest that you can implement a generic fix for this inherent issue, where any op that defines an SSA value can be added to the above list, and this is decoupled from the checks like 'isOpLoopInvariant`. This will make the check more robust, and relevant test cases can also be added.

[vinayaka-polymage]

Minor suggestion.

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

620

Minor suggestion: It is better to create a minimal test case that captures the intended check. A similar test case without vector_load and vector.reduction can be created to improve readability.

[sumesh13]

Rework based on comments

• Separate the part that identifies operations that have users and use it in the later phase.

[sumesh13]

In D99761#2675903, @vinayaka-polymage wrote:

Thank you for bringing up this issue in LICM, it is an important one. I have a few comments on the fix:

1. Inherent issue is definedOps ( going to be renamed to opsWithUsers) is being populated only under some conditions, and does not actually do the job of collecting all ops that define an SSA value in that block.
2. Even after this fix, only AffineForOps will be added to the list, but there are other cases like AffineDma... ops, AffineIfOps that are actually NOT code invariant etc. which will still cause similar problems.

So, I suggest that you can implement a generic fix for this inherent issue, where any op that defines an SSA value can be added to the above list, and this is decoupled from the checks like 'isOpLoopInvariant`. This will make the check more robust, and relevant test cases can also be added.

Thanks @vinayaka-polymage for the very useful feedback. Hadnt realized the issue might affect other operations such as Affineif etc. I have reworked the code and also addes two tests.

[sumesh13]

I have addresses the comments. Pls take a look.

[vinayaka-polymage]

@sumesh13 This new generic fix looks good, it handles more cases that were previously not handled. Thanks for this. I have added some comments related to code organization, please check them. Logic looks good to me.

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

48

It is better to keep this local like before, see comment below.

190

This is a per AffineForOp var, and can be kept local here like before.

201

Populating into definedOps (opsWithUsers) can happen here. No need to walk the affine for, because source of an operand for any op, can only come before it. So,

if (op->getNumResults() > 0 && !op->use_empty())
        opsWithUsers.insert(op);

can be placed here.

[sumesh13]

Address review comments on code reorg

[sumesh13]

@vinayaka-polymage Thanks for the comments. Pls take a look

[sumesh13]

Undo a change w.r.t to the base that had crept in

[sumesh13]

Mark review comments as done and undo a typo that had crept in

[vinayaka-polymage]

@sumesh13 Thanks for addressing all the comments, the change looks good now. Can you please check these last two comments before I approve this ?

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

81–82

Don't need this now, right ? This line can be deleted, and comment be moved to the location where it is getting populated.

226–227

This section now does not have any change, so please make sure that there is no diff here.

[vinayaka-polymage]

Please change the commit title to a more relevant one, I am approving this.

@bondhugula Could you please take a final pass at this?

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

81–82

Actually, I take my above comment back. This function is being used in ifOp hoisting as well, and with the current implementation it will be needed. Refactoring that can be a separate change in itself.

[sumesh13]

Address comments

[sumesh13]

@vinayaka-polymage Just noticed that you accepted a previous revision with comment on handling of ifops. Considering ifops, would you think it might be helpful to just collect all the ops like I was doing in a previous revision ?

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

81–82

Thanks ...that was one reason I was collecting ops separately. Would you think that might be preferable ?

[sumesh13]

Revert to earlier revision that was approved and some minor cleanup of comments.

[sumesh13]

Adding back a change from previous revision that was approved and some minor cleanup.

[sumesh13]

Gentle reminder..also pls suggest if it would be better to collect all the opsWithUsers separately as i was doing in a previous revision. For nested ifs, looks like otherwise you may not be collecting all the opsWithUsers by just collecting at the top level and leaves.

[bondhugula]

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

199

I think just !op.use_empty() should be sufficient.

[sumesh13]

Simplify condition to collect ops with users.

[sumesh13]

Address comment on simplifying condition to add ops with users.

[bondhugula]

Looks great.

[sumesh13]

@bondhugula , @vinayaka-polymage thanks for the review. @sgrechanik Can you pls help land this

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

 /// TODO: The pass is missing zero-trip tests.
 /// TODO: Check for the presence of side effects before hoisting.
 /// TODO: This code should be removed once the new LICM pass can handle its
 ///       uses.
 struct LoopInvariantCodeMotion
     : public AffineLoopInvariantCodeMotionBase<LoopInvariantCodeMotion> {
   void runOnFunction() override;
   void runOnAffineForOp(AffineForOp forOp);
 };
 } // end anonymous namespace
 
 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)) {
     // If the body of a predicated region has a for loop, we don't hoist the
     // 'affine.if'.
     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 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()
                          : cast<AffineWriteOpInterface>(op).getMemRef();
       for (auto *user : memref.getUsers()) {
         // If this memref has a user that is a DMA, give up because these
         // operations write to this memref.
         if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) {
for (unsigned int i = 0; i < op.getNumOperands(); ++i) {
 
       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 (definedOps.count(operandSrc) && opsToHoist.count(operandSrc) == 0) {
+        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 *> &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;
       }
     }
   }
 
   return true;
 }
 
 // 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;
   }
 
   return true;
 }
 
 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);
         }
       }
     }
   }
 
   // For all instructions that we found to be invariant, place sequentially
   // right before the for loop.
   for (auto *op : opsToMove) {
     op->moveBefore(forOp);
   }
 
   LLVM_DEBUG(forOp->print(llvm::dbgs() << "Modified loop\n"));
 }
 
 void LoopInvariantCodeMotion::runOnFunction() {
   // Walk through all loops in a function in innermost-loop-first order.  This
   // way, we first LICM from the inner loop, and place the ops in
   // the outer loop, which in turn can be further LICM'ed.
   getFunction().walk([&](AffineForOp op) {
     LLVM_DEBUG(op->print(llvm::dbgs() << "\nOriginal loop\n"));
     runOnAffineForOp(op);
   });
 }
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createAffineLoopInvariantCodeMotionPass() {
   return std::make_unique<LoopInvariantCodeMotion>();
 }

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

 // CHECK:       memref.alloc()
 // CHECK-NEXT:  memref.alloc()
 // CHECK-NEXT:  memref.alloc()
 // CHECK-NEXT:  affine.vector_load
 // CHECK-NEXT:  affine.vector_load
 // 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