Implement an scf.for range folding optimization pass.

[Anthony]

Summary

Implement an scf.for range folding optimization pass.

In cases where arithmetic (addi/muli) ops are performed on an scf.for loops induction variable with a single use, we can fold those ops directly into the scf.for loop.

For example, in the following code:

scf.for %i = %c0 to %arg1 step %c1 {
  %0 = addi %arg2, %i : index
  %1 = muli %0, %c4 : index
  %2 = memref.load %arg0[%1] : memref<?xi32>
  %3 = muli %2, %2 : i32
  memref.store %3, %arg0[%1] : memref<?xi32>
}

we can lift %0 up into the scf.for loop range, as it is the only user of %i:

%lb = addi %arg2, %c0 : index
%ub = addi %arg2, %i : index
scf.for %i = %lb to %ub step %c1 {
  %1 = muli %0, %c4 : index
  %2 = memref.load %arg0[%1] : memref<?xi32>
  %3 = muli %2, %2 : i32
  memref.store %3, %arg0[%1] : memref<?xi32>
}

[ftynse]

mlir/include/mlir/Dialect/SCF/Passes.td

49

This can be a function pass. Function passes can run in parallel on different functions.

50

This shouldn't be empty.

mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp

1

Please add the license header.

20

We mark functions local to a translation module as static.

48

There is a version of OpBuilder::clone that takes a BlockAndValueMapping. It should allow you to create a copy of use operation without casting it to the specific type, and without manually calling lookupAndDefault.

mlir/test/Dialect/SCF/loop-range.mlir

84

Please add the newline

[mravishankar]

Please use clang-format on the patch. There seems to be a lot of linter errors.

mlir/include/mlir/Dialect/SCF/Passes.h

38

Missing comment about the pass description.

mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp

35

Simpler to just use indVar.getUser().

Also naming Nit. use typically refers to the OpOperand & in the operation using the value. The Operation using the value is typically called user.

87

I think if you want to do a fixed point iteration it might be better to just redo this as a pattern on scf.for operations.

The pattern can fail if the induction variable has more than one use, or if it is used in anything other than an add or a mul. If that is not the case you can,

1. Create a new scf.for operation with different lb, ub
2. Use inlineRegionBefore to move the body of the original scf.for as a region of the new op.
3. Replace uses of the original add/mul with the induction variable.

mlir/test/Dialect/SCF/loop-range.mlir

18

Convention on checks is to use %[[ARG0:.*]] instead of [[ARG0:%.*]]. It makes some of the other checks easier to write. For example, when you want to match [%v] . With %[[ARG0:.*]], ARG0 is set to v. So you can just do [%[[ARG0]]].
OTOH when using [[ARG0:%.*]] you will need to match [[[ARG0]]]. The [[[ ]]] conflicts with the check parsing.

[Anthony]

Requested changes, minus the rewrite pattern.

Updating D104289: First crack at foor loop range folding pass.

[Hardcode84]

Is there any reason to make this separate pass instead of ForOp canonicalization? Also, diff looks incomplete

[Anthony]

I have made the changes requested, with a few inline comments to ask for further clarification.

(If I was supposed to submit as a planed revision or something else instead, my apologies)

mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp

35

I didn't see a getUser for Value, only getUsers, am I missing something?

48

Thanks for the tip. I used your suggested OpBuilder::clone method to simplify a bit. However, I still need to distinguish between AddIOp and MulIOp because the former does not fold into the for loop step, while the latter does.

87

Thanks for the comments. So I can rewrite this into a rewrite pattern, but I have one question: Does a rewrite pattern on ForOp guarantee that the innermost loop will get rewritten first? I believe that is an invariant of the transform as well.

I can come up with a test case that better demonstrates the snag I hit without the inner-most transform.

[ftynse]

It looks like the diff was updated the wrong way: now it only contains the difference between the initial version and the updated version, rather than the difference between the updated version and the mainline. If you created a separate git commit for the changes, you need to do arc diff HEAD^^ (as many ^ as commits in your branch) or arc diff <hash-of-the-commit-on-the-main-branch>.

[Anthony]

Sorry about that, I believe I have it correct now.

[ftynse]

mlir/test/Dialect/SCF/loop-range.mlir

1

We prefer to only test one pass unless strictly necessary otherwise. Helps debugging.

[mehdi_amini]

In D104289#2827120, @Hardcode84 wrote:

Is there any reason to make this separate pass instead of ForOp canonicalization?

Was this answered?

Also please update the title to be descriptive, and update the description to reflect what the commit message will be.

[Anthony]

In D104289#2831392, @mehdi_amini wrote:

In D104289#2827120, @Hardcode84 wrote:

Is there any reason to make this separate pass instead of ForOp canonicalization?

Was this answered?

Also please update the title to be descriptive, and update the description to reflect what the commit message will be.

Sorry, I missed this. I followed the pattern for the loop invariant code motion pass, which is not a canonicalization on for. For this choice, I actually am not sure whether this would be better off as a canonicalization, or as its own pass. Can I have some guidance here?

(Will update description and commit messages).

[mehdi_amini]

Thinking about canonicalization, in the past I've been working with a canonical form for loops that would try to form the iteration space so that it starts at 0 and has a step increment of one, which may be the opposite of this transformation.

So having this as a separate pass makes sense, but I'd be cautious and keep in mind that this may just be as a lowering step which could be undone by canonicalization.

mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp

46

https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/Value.h#L209-L214

62–63

(keep auto for when it improves readability and/or it is obvious from the context, same below in this patch)

[bondhugula]

Can you fix the foor -> for typo in the revision title? Also "scf for" in the title will be useful.

[Anthony]

Implement an scf.for range folding optimization pass.

In cases where arithmetic (addi/muli) ops are performed on an scf.for
loops induction variable with a single use, we can fold those ops
directly into the scf.for loop.

For example, in the following code:

scf.for %i = %c0 to %arg1 step %c1 {
  %0 = addi %arg2, %i : index
  %1 = muli %0, %c4 : index
  %2 = memref.load %arg0[%1] : memref<?xi32>
  %3 = muli %2, %2 : i32
  memref.store %3, %arg0[%1] : memref<?xi32>
}

we can lift %0 up into the scf.for loop range, as it is the only user
of %i:

%lb = addi %arg2, %c0 : index
%ub = addi %arg2, %i : index
scf.for %i = %lb to %ub step %c1 {
  %1 = muli %0, %c4 : index
  %2 = memref.load %arg0[%1] : memref<?xi32>
  %3 = muli %2, %2 : i32
  memref.store %3, %arg0[%1] : memref<?xi32>
}

[Anthony]

In D104289#2831536, @mehdi_amini wrote:

Thinking about canonicalization, in the past I've been working with a canonical form for loops that would try to form the iteration space so that it starts at 0 and has a step increment of one, which may be the opposite of this transformation.

So having this as a separate pass makes sense, but I'd be cautious and keep in mind that this may just be as a lowering step which could be undone by canonicalization.

Thanks, got it. I'm happy to change the pass in the future if it turns we need it as a canonicalization.

I've also made the requested changes, and added another test case (see description in loop-range.mlir).

I believe there is still the question of whether or not this is better off as a rewrite pattern per @mravishankar comments.

[mehdi_amini]

LGTM as-is, with a few comments.

mlir/include/mlir/Dialect/SCF/Passes.td

49

Can you actually make it an Operation pass: there is no need to anchor this to a function I believe.

mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp

85

(After you make it an Operation pass, this will be just getOperation()->walk()

86

foldRanges can't fail, can you make it return void?
At this point I would inline the body here as well.

[mehdi_amini]

Oh and if you can update the description here to match what will be the commit description please.

[Anthony]

In D104289#2834278, @mehdi_amini wrote:

LGTM as-is, with a few comments.

I can change it to operation pass---originally I had it as one, but changed it to a function pass at the request of another reviewer @ftynse so that it can be run in parallel on functions. Do we still want to go back to operation pass?

(I'll update description as well).

[rriddle]

In D104289#2834495, @Anthony wrote:

In D104289#2834278, @mehdi_amini wrote:

LGTM as-is, with a few comments.

I can change it to operation pass---originally I had it as one, but changed it to a function pass at the request of another reviewer @ftynse so that it can be run in parallel on functions. Do we still want to go back to operation pass?

(I'll update description as well).

An OperationPass is strictly better, because it can run on different operation types (think different kinds of FuncOp). Being an OperationPass doesn't change how the parallelization works, because the pass should still be scheduled per FuncOp. The only potential reason why parallelization would be hurt, is if a user accidentally schedules it on a ModuleOp (which is allowed for generic OperationPass').

[mehdi_amini]

To extend on what River describes with an example:

-pass-pipeline='func(for-loop-range-folding)' -> runs in parallel on function
`-pass-pipeline='module(for-loop-range-folding)' -> won't run in parallel.

The point is that making it an operation pass also allows to do:

-pass-pipeline='gpu.func(for-loop-range-folding)' -> run in parallel on GPU functions!

[Anthony]

Implement an scf.for range folding optimization pass.

In cases where arithmetic (addi/muli) ops are performed on an scf.for
loops induction variable with a single use, we can fold those ops
directly into the scf.for loop.

For example, in the following code:

scf.for %i = %c0 to %arg1 step %c1 {
  %0 = addi %arg2, %i : index
  %1 = muli %0, %c4 : index
  %2 = memref.load %arg0[%1] : memref<?xi32>
  %3 = muli %2, %2 : i32
  memref.store %3, %arg0[%1] : memref<?xi32>
}

we can lift %0 up into the scf.for loop range, as it is the only user
of %i:

%lb = addi %arg2, %c0 : index
%ub = addi %arg2, %i : index
scf.for %i = %lb to %ub step %c1 {
  %1 = muli %0, %c4 : index
  %2 = memref.load %arg0[%1] : memref<?xi32>
  %3 = muli %2, %2 : i32
  memref.store %3, %arg0[%1] : memref<?xi32>
}

[Anthony]

Thanks for the clarifications guys. I've gone ahead and switched it to an operation pass, and inlined the foldRanges function into the operation walk.

[mehdi_amini]

In D104289#2835402, @Anthony wrote:

Thanks for the clarifications guys.

As a nit you likely intend to use "guys" in a neutral way, but nowadays it is quite controversial. It may be better to avoid it.

I've gone ahead and switched it to an operation pass, and inlined the foldRanges function into the operation walk.

Thanks! LG, do you need someone to land this for you?

[Anthony]

In D104289#2837490, @mehdi_amini wrote:

In D104289#2835402, @Anthony wrote:

Thanks for the clarifications guys.

As a nit you likely intend to use "guys" in a neutral way, but nowadays it is quite controversial. It may be better to avoid it.

Thanks for pointing that out---I'll be more conscientious about my wording in the future.

I've gone ahead and switched it to an operation pass, and inlined the foldRanges function into the operation walk.

Thanks! LG, do you need someone to land this for you?

Yes please! This is my first commit.

[stellaraccident]

I think we need to revisit this patch as it is broken for some very common cases. Left notes here: https://github.com/llvm/llvm-project/issues/56235

mlir/include/mlir/Dialect/SCF/Passes.h

 /// Creates a pass that specializes parallel loop for unrolling and
 /// vectorization.
 std::unique_ptr<Pass> createParallelLoopSpecializationPass();
 
 /// Creates a pass which tiles innermost parallel loops.
 std::unique_ptr<Pass>
 createParallelLoopTilingPass(llvm::ArrayRef<int64_t> tileSize = {});
 
+/// Creates a pass which folds arith ops on induction variable into
+/// loop range.
 std::unique_ptr<Pass> createForLoopRangeFoldingPass();
 
-
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/SCF/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_PASSES_H_

mlir/include/mlir/Dialect/SCF/Passes.td

let options = [
     ListOption<"tileSizes", "parallel-loop-tile-sizes", "int64_t",
                "Factors to tile parallel loops by",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
   ];
   let dependentDialects = ["AffineDialect"];
 }
 
 def SCFForLoopRangeFolding
-    : Pass<"for-loop-range-folding"> {
-  let summary = "";
+    : FunctionPass<"for-loop-range-folding"> {
+  let summary = "Fold add/mul ops into loop range";
   let constructor = "mlir::createForLoopRangeFoldingPass()";
 }
 
 #endif // MLIR_DIALECT_SCF_PASSES

mlir/lib/Dialect/SCF/Transforms/LoopRangeFolding.cpp

+//===- LoopRangeFolding.cpp - Code to perform loop range folding-----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements loop range folding.
+//
+//===----------------------------------------------------------------------===//
+
 #include "PassDetail.h"
 #include "mlir/Dialect/SCF/Passes.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/SCF/Utils.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 
-
 using namespace mlir;
 using namespace mlir::scf;
 
 namespace {
 struct ForLoopRangeFolding
     : public SCFForLoopRangeFoldingBase<ForLoopRangeFolding> {
-  void runOnOperation() override;
+  void runOnFunction() override;
 };
-}
+} // namespace
 
-LogicalResult foldRanges(ForOp op) {
+static LogicalResult foldRanges(ForOp op) {
   // Fold until a fixed point is reached
   Value indVar = op.getInductionVar();
 
   auto canBeFolded = [&](Value value) {
     return op.isDefinedOutsideOfLoop(value) || value == indVar;
   };
 
   while (true) {
 
     // If the induction variable is used more than once, we can't fold its arith
     // ops into the loop range
     if (!indVar.hasOneUse())
       break;
 
-    Operation *use = indVar.getUses().begin().getUser();
-    if (!isa<AddIOp, MulIOp>(use))
+    Operation *user = indVar.getUses().begin().getUser();
+    if (!isa<AddIOp, MulIOp>(user))
       break;
 
-    if (!llvm::all_of(use->getOperands(), canBeFolded))
+    if (!llvm::all_of(user->getOperands(), canBeFolded))
       break;
 
     OpBuilder b(op);
-    BlockAndValueMapping lbMap; lbMap.map(indVar, op.lowerBound());
-    BlockAndValueMapping ubMap; ubMap.map(indVar, op.upperBound());
-    BlockAndValueMapping stepMap; stepMap.map(indVar, op.step());
-
-    if (auto addOp = dyn_cast<AddIOp>(use)) {
-      auto lbFold = b.create<AddIOp>(
-          op.getLoc(),
-          lbMap.lookupOrDefault(addOp.getOperand(0)),
-          lbMap.lookupOrDefault(addOp.getOperand(1)));
-
-      auto ubFold = b.create<AddIOp>(
-          op.getLoc(),
-          ubMap.lookupOrDefault(addOp.getOperand(0)),
-          ubMap.lookupOrDefault(addOp.getOperand(1)));
-
-      op.setLowerBound(lbFold);
-      op.setUpperBound(ubFold);
-
-      addOp.replaceAllUsesWith(indVar);
-      addOp.erase();
-
-    } else if (auto mulOp = dyn_cast<MulIOp>(use)) {
-      auto ubFold = b.create<MulIOp>(
-          op.getLoc(),
-          ubMap.lookupOrDefault(mulOp.getOperand(0)),
-          ubMap.lookupOrDefault(mulOp.getOperand(1)));
-
-      auto stepFold = b.create<MulIOp>(
-          op.getLoc(),
-          stepMap.lookupOrDefault(mulOp.getOperand(0)),
-          stepMap.lookupOrDefault(mulOp.getOperand(1)));
-
-      op.setUpperBound(ubFold);
-      op.setStep(stepFold);
-
-      mulOp.replaceAllUsesWith(indVar);
-      mulOp.erase();
-    } 
+    BlockAndValueMapping lbMap;
+    lbMap.map(indVar, op.lowerBound());
+    BlockAndValueMapping ubMap;
+    ubMap.map(indVar, op.upperBound());
+    BlockAndValueMapping stepMap;
+    stepMap.map(indVar, op.step());
+
+    if (isa<AddIOp>(user)) {
+      auto lbFold = b.clone(*user, lbMap);
+      auto ubFold = b.clone(*user, ubMap);
+
+      op.setLowerBound(lbFold->getResult(0));
+      op.setUpperBound(ubFold->getResult(0));
+
+    } else if (isa<MulIOp>(user)) {
+      auto ubFold = b.clone(*user, ubMap);
+      auto stepFold = b.clone(*user, stepMap);
+
+      op.setUpperBound(ubFold->getResult(0));
+      op.setStep(stepFold->getResult(0));
+    }
+
+    ValueRange wrapIndvar(indVar);
+    user->replaceAllUsesWith(wrapIndvar);
+    user->erase();
   }
 
   return success();
 }
 
-void ForLoopRangeFolding::runOnOperation() {
-  getOperation()->walk([&](ForOp forOp) {
+void ForLoopRangeFolding::runOnFunction() {
+  getFunction().getOperation()->walk([&](ForOp forOp) {
     if (failed(foldRanges(forOp)))
       signalPassFailure();
   });
 }
 
 std::unique_ptr<Pass> mlir::createForLoopRangeFoldingPass() {
   return std::make_unique<ForLoopRangeFolding>();
 }

mlir/test/Dialect/SCF/loop-range.mlir

 // RUN: mlir-opt %s -pass-pipeline='func(for-loop-range-folding,canonicalize)' -split-input-file | FileCheck %s
 
 func @fold_one_loop(%arg0: memref<?xi32>, %arg1: index, %arg2: index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c4 = constant 4 : index
   scf.for %i = %c0 to %arg1 step %c1 {
     %0 = addi %arg2, %i : index
     %1 = muli %0, %c4 : index
     %2 = memref.load %arg0[%1] : memref<?xi32>
     %3 = muli %2, %2 : i32
     memref.store %3, %arg0[%1] : memref<?xi32>
   }
   return
 }
 
 // CHECK-LABEL: func @fold_one_loop
-// CHECK-SAME:   ([[ARG0:%.*]]: {{.*}}, [[ARG1:%.*]]: {{.*}}, [[ARG2:%.*]]: {{.*}}
-// CHECK:       [[C4:%.*]] = constant 4 : index
-// CHECK:       [[I0:%.*]] = addi [[ARG2]], [[ARG1]] : index
-// CHECK:       [[I1:%.*]] = muli [[I0]], [[C4]] : index
-// CHECK:       scf.for [[I:%.*]] = [[ARG2]] to [[I1]] step [[C4]] {
-// CHECK:         [[I2:%.*]] = memref.load [[ARG0]]{{\[}}[[I]]
-// CHECK:         [[I3:%.*]] = muli [[I2]], [[I2]] : i32
-// CHECK:         memref.store [[I3]], [[ARG0]]{{\[}}[[I]]
+// CHECK-SAME:   (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}
+// CHECK:       %[[C4:.*]] = constant 4 : index
+// CHECK:       %[[I0:.*]] = addi %[[ARG2]], %[[ARG1]] : index
+// CHECK:       %[[I1:.*]] = muli %[[I0]], %[[C4]] : index
+// CHECK:       scf.for %[[I:.*]] = %[[ARG2]] to %[[I1]] step %[[C4]] {
+// CHECK:         %[[I2:.*]] = memref.load %[[ARG0]]{{\[}}%[[I]]
+// CHECK:         %[[I3:.*]] = muli %[[I2]], %[[I2]] : i32
+// CHECK:         memref.store %[[I3]], %[[ARG0]]{{\[}}%[[I]]
 
 func @fold_one_loop2(%arg0: memref<?xi32>, %arg1: index, %arg2: index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c4 = constant 4 : index
   %c10 = constant 10 : index
   scf.for %j = %c0 to %c10 step %c1 {
     scf.for %i = %c0 to %arg1 step %c1 {
       %0 = addi %arg2, %i : index
       %1 = muli %0, %c4 : index
       %2 = memref.load %arg0[%1] : memref<?xi32>
       %3 = muli %2, %2 : i32
       memref.store %3, %arg0[%1] : memref<?xi32>
     }
   }
   return
 }
 
 // CHECK-LABEL: func @fold_one_loop2
-// CHECK-SAME:   ([[ARG0:%.*]]: {{.*}}, [[ARG1:%.*]]: {{.*}}, [[ARG2:%.*]]: {{.*}}
-// CHECK:       [[C4:%.*]] = constant 4 : index
-// CHECK:       [[I0:%.*]] = addi [[ARG2]], [[ARG1]] : index
-// CHECK:       [[I1:%.*]] = muli [[I0]], [[C4]] : index
-// CHECK:       scf.for [[I:%.*]] = [[ARG2]] to [[I1]] step [[C4]] {
-// CHECK:         [[I2:%.*]] = memref.load [[ARG0]]{{\[}}[[I]]
-// CHECK:         [[I3:%.*]] = muli [[I2]], [[I2]] : i32
-// CHECK:         memref.store [[I3]], [[ARG0]]{{\[}}[[I]]
+// CHECK-SAME:   (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}
+// CHECK:       %[[C4:.*]] = constant 4 : index
+// CHECK:       %[[I0:.*]] = addi %[[ARG2]], %[[ARG1]] : index
+// CHECK:       %[[I1:.*]] = muli %[[I0]], %[[C4]] : index
+// CHECK:       scf.for %[[I:.*]] = %[[ARG2]] to %[[I1]] step %[[C4]] {
+// CHECK:         %[[I2:.*]] = memref.load %[[ARG0]]{{\[}}%[[I]]
+// CHECK:         %[[I3:.*]] = muli %[[I2]], %[[I2]] : i32
+// CHECK:         memref.store %[[I3]], %[[ARG0]]{{\[}}%[[I]]
 
 func @fold_two_loops(%arg0: memref<?xi32>, %arg1: index, %arg2: index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c4 = constant 4 : index
   %c10 = constant 10 : index
   scf.for %j = %c0 to %c10 step %c1 {
     scf.for %i = %j to %arg1 step %c1 {
       %0 = addi %arg2, %i : index
       %1 = muli %0, %c4 : index
       %2 = memref.load %arg0[%1] : memref<?xi32>
       %3 = muli %2, %2 : i32
       memref.store %3, %arg0[%1] : memref<?xi32>
     }
   }
   return
 }
 
 // CHECK-LABEL: func @fold_two_loops
-// CHECK-SAME:   ([[ARG0:%.*]]: {{.*}}, [[ARG1:%.*]]: {{.*}}, [[ARG2:%.*]]: {{.*}}
-// CHECK:       [[C10:%.*]] = constant 10 : index
-// CHECK:       [[C4:%.*]] = constant 4 : index
-// CHECK:       [[C1:%.*]] = constant 1 : index
-// CHECK:       [[I0:%.*]] = addi [[ARG2]], [[C10]] : index
-// CHECK:       scf.for [[J:%.*]] = [[ARG2]] to [[I0]] step [[C1]] {
-// CHECK:         [[I1:%.*]] = addi [[ARG2]], [[ARG1]] : index
-// CHECK:         [[I2:%.*]] = muli [[I1]], [[C4]] : index
-// CHECK:         scf.for [[I:%.*]] = [[J]] to [[I2]] step [[C4]] {
-// CHECK:           [[I3:%.*]] = memref.load [[ARG0]]{{\[}}[[I]]
-// CHECK:           [[I4:%.*]] = muli [[I3]], [[I3]] : i32
-// CHECK:           memref.store [[I4]], [[ARG0]]{{\[}}[[I]]
+// CHECK-SAME:   (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}
+// CHECK:       %[[C10:.*]] = constant 10 : index
+// CHECK:       %[[C4:.*]] = constant 4 : index
+// CHECK:       %[[C1:.*]] = constant 1 : index
+// CHECK:       %[[I0:.*]] = addi %[[ARG2]], %[[C10]] : index
+// CHECK:       scf.for %[[J:.*]] = %[[ARG2]] to %[[I0]] step %[[C1]] {
+// CHECK:         %[[I1:.*]] = addi %[[ARG2]], %[[ARG1]] : index
+// CHECK:         %[[I2:.*]] = muli %[[I1]], %[[C4]] : index
+// CHECK:         scf.for %[[I:.*]] = %[[J]] to %[[I2]] step %[[C4]] {
+// CHECK:           %[[I3:.*]] = memref.load %[[ARG0]]{{\[}}%[[I]]
+// CHECK:           %[[I4:.*]] = muli %[[I3]], %[[I3]] : i32
+// CHECK:           memref.store %[[I4]], %[[ARG0]]{{\[}}%[[I]]