[Mlir-commits] [mlir] 5d001f5 - [mlir] Fix a bug in Affine LICM.

[Amy Zhuang]

Author: Amy Zhuang
Date: 2021-10-07T15:46:43-07:00
New Revision: 5d001f58f241a0e15751feea8f087ede0ca6371f

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

LOG: [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

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 86fb6f83b6378..caab5439492c5 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
+++ b/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());

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 dce934d6e5bb5..c4146d22387a3 100644
--- a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
+++ b/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