[mlir] Fix a bug in Affine LICM.

[ayzhuang]

Summary

Currently Affine LICM checks iterOperands and does not hoist out any instruction containing iterOperands. We should check iterArgs instead.

[bondhugula]

Thanks for fixing this! Minor comment below.

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

734–738

Can you please indent to maintain the input nesting for readability?

[ayzhuang]

Address review comment to indent test.

[ayzhuang]

@bondhugula Thank you for the review. We'll merge it tomorrow if there are no more comments.

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

bool checkInvarianceOfNestedIfOps(Operation *op, Value indVar,
   }
 
   return true;
 }
 
 void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
   auto *loopBody = forOp.getBody();
   auto indVar = forOp.getInductionVar();
-  ValueRange iterArgs = forOp.getIterOperands();
+  ValueRange iterArgs = forOp.getRegionIterArgs();
 
   // 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;
 

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

 // CHECK-NEXT:  addf
 // CHECK-NEXT:  affine.yield
 // CHECK-NEXT:  }
 // CHECK-NEXT:  mulf
 // CHECK-NEXT:  affine.store
 
 // -----
 
+// CHECK-LABEL: func @use_of_iter_operands_invariant
+func @use_of_iter_operands_invariant(%m : memref<10xindex>) {
+  %sum_1 = constant 0 : index
+  %v0 = affine.for %arg1 = 0 to 11 iter_args (%prevAccum = %sum_1) -> index {
+    %prod = muli %sum_1, %sum_1 : index
+    %newAccum = addi %prevAccum, %prod : index
+    affine.yield %newAccum : index
+  }
+  return
+}
+
+// CHECK:       constant
+// CHECK-NEXT:  muli
+// CHECK-NEXT:  affine.for
+// CHECK-NEXT:    addi
+// CHECK-NEXT:    affine.yield
+
+// -----
+
 // 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