[mlir] Fix a bug in Affine LICM.

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

Reviewed By: bondhugula

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

Author: ayzhuang

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

@@ -195,7 +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();
+  ValueRange iterArgs = forOp.getRegionIterArgs();
 
   // This is the place where hoisted instructions would reside.
   OpBuilder b(forOp.getOperation());

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

@@ -719,6 +719,25 @@ func @affine_for_not_invariant(%in : memref<30x512xf32, 1>,
 
 // -----
 
+// 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