This is an archive of the discontinued LLVM Phabricator instance.

Differential D77052

[MLIR] Rename collapsePLoops -> collapseParallelLoops
ClosedPublic

Authored by tpopp on Mar 30 2020, 4:50 AM.

Details

Reviewers
bondhugula
Commits
rG90b7bbffddc6: [MLIR] Rename collapsePLoops -> collapseParallelLoops
Summary

Additionally, NFC code cleanups were done.

This is to address additional comments on
https://reviews.llvm.org/D76363

Diff Detail

Event Timeline

tpopp created this revision.Mar 30 2020, 4:50 AM
Herald added a project: Restricted Project. · View Herald TranscriptMar 30 2020, 4:50 AM
Herald added subscribers: llvm-commits, Joonsoo, liufengdb and 11 others. · View Herald Transcript
Harbormaster completed remote builds in B50939: Diff 253554.Mar 30 2020, 5:55 AM
bondhugula accepted this revision.Mar 30 2020, 9:19 AM
bondhugula added inline comments.
mlir/lib/Transforms/Utils/LoopUtils.cpp
1229

inner.getBody() to avoid confusion since this is a single block region.

This revision is now accepted and ready to land.Mar 30 2020, 9:19 AM
tpopp marked an inline comment as done.Mar 30 2020, 10:03 AM
tpopp added inline comments.
mlir/lib/Transforms/Utils/LoopUtils.cpp
1229

I'll block this on whatever happens with https://reviews.llvm.org/D77060

There's an annoying detail right now that constructing an OpBuilder on a Block vs Region create different insertion points (end() and begin()) respectively, so I'll revert this line and another below if it resolves in a different way, or do as you suggest here if it lands.

tpopp marked 2 inline comments as done.Apr 1 2020, 1:16 AM
tpopp added inline comments.
mlir/lib/Transforms/Utils/LoopUtils.cpp
1229

This is now OpBuilder::atBlockBegin() in order to be more explicit.

Herald added a subscriber: grosul1. · View Herald TranscriptApr 1 2020, 1:16 AM
Closed by commit rG90b7bbffddc6: [MLIR] Rename collapsePLoops -> collapseParallelLoops (authored by Tres Popp <tpopp@google.com>). · Explain WhyApr 1 2020, 1:37 AM
This revision was automatically updated to reflect the committed changes.

Revision Contents

PathSize
mlir/
include/
mlir/
Transforms/
4 lines
lib/
Transforms/
2 lines
Utils/
13 lines
test/
Transforms/
51 lines
40 lines

Diff 254113

mlir/include/mlir/Transforms/LoopUtils.h

Show First 20 Lines • Show All 226 Lines • ▼ Show 20 Lines
/// Replace a perfect nest of "for" loops with a single linearized loop. Assumes /// Replace a perfect nest of "for" loops with a single linearized loop. Assumes
/// `loops` contains a list of perfectly nested loops with bounds and steps /// `loops` contains a list of perfectly nested loops with bounds and steps
/// independent of any loop induction variable involved in the nest. /// independent of any loop induction variable involved in the nest.
void coalesceLoops(MutableArrayRef<loop::ForOp> loops); void coalesceLoops(MutableArrayRef<loop::ForOp> loops);
/// Take the ParallelLoop and for each set of dimension indices, combine them /// Take the ParallelLoop and for each set of dimension indices, combine them
/// into a single dimension. combinedDimensions must contain each index into /// into a single dimension. combinedDimensions must contain each index into
/// loops exactly once. /// loops exactly once.
void collapsePLoops(loop::ParallelOp loops, void collapseParallelLoops(loop::ParallelOp loops,
ArrayRef<std::vector<unsigned>> combinedDimensions); ArrayRef<std::vector<unsigned>> combinedDimensions);
/// Maps `forOp` for execution on a parallel grid of virtual `processorIds` of /// Maps `forOp` for execution on a parallel grid of virtual `processorIds` of
/// size given by `numProcessors`. This is achieved by embedding the SSA values /// size given by `numProcessors`. This is achieved by embedding the SSA values
/// corresponding to `processorIds` and `numProcessors` into the bounds and step /// corresponding to `processorIds` and `numProcessors` into the bounds and step
/// of the `forOp`. No check is performed on the legality of the rewrite, it is /// of the `forOp`. No check is performed on the legality of the rewrite, it is
/// the caller's responsibility to ensure legality. /// the caller's responsibility to ensure legality.
/// ///
/// Requires that `processorIds` and `numProcessors` have the same size and that /// Requires that `processorIds` and `numProcessors` have the same size and that
▲ Show 20 Lines • Show All 57 Lines • Show Last 20 Lines

mlir/lib/Transforms/ParallelLoopCollapsing.cpp

Show All 31 Lines module->walk([&](loop::ParallelOp op) {
// arguments, so we do that here to simplify things. // arguments, so we do that here to simplify things.
llvm::SmallVector<std::vector<unsigned>, 3> combinedLoops; llvm::SmallVector<std::vector<unsigned>, 3> combinedLoops;
if (clCollapsedIndices0.size()) if (clCollapsedIndices0.size())
combinedLoops.push_back(clCollapsedIndices0); combinedLoops.push_back(clCollapsedIndices0);
if (clCollapsedIndices1.size()) if (clCollapsedIndices1.size())
combinedLoops.push_back(clCollapsedIndices1); combinedLoops.push_back(clCollapsedIndices1);
if (clCollapsedIndices2.size()) if (clCollapsedIndices2.size())
combinedLoops.push_back(clCollapsedIndices2); combinedLoops.push_back(clCollapsedIndices2);
collapsePLoops(op, combinedLoops); collapseParallelLoops(op, combinedLoops);
}); });
} }
ListOption<unsigned> clCollapsedIndices0{ ListOption<unsigned> clCollapsedIndices0{
*this, "collapsed-indices-0", *this, "collapsed-indices-0",
llvm::cl::desc("Which loop indices to combine 0th loop index"), llvm::cl::desc("Which loop indices to combine 0th loop index"),
llvm::cl::MiscFlags::CommaSeparated}; llvm::cl::MiscFlags::CommaSeparated};
Show All 21 Lines

mlir/lib/Transforms/Utils/LoopUtils.cpp

Show First 20 Lines • Show All 1,220 Lines • ▼ Show 20 Lines
/// %i = %ii * %s + %lb /// %i = %ii * %s + %lb
/// Insert the induction variable remapping in the body of `inner`, which is /// Insert the induction variable remapping in the body of `inner`, which is
/// expected to be either `loop` or another loop perfectly nested under `loop`. /// expected to be either `loop` or another loop perfectly nested under `loop`.
/// Insert the definition of new bounds immediate before `outer`, which is /// Insert the definition of new bounds immediate before `outer`, which is
/// expected to be either `loop` or its parent in the loop nest. /// expected to be either `loop` or its parent in the loop nest.
static void normalizeLoop(loop::ForOp loop, loop::ForOp outer, static void normalizeLoop(loop::ForOp loop, loop::ForOp outer,
loop::ForOp inner) { loop::ForOp inner) {
OpBuilder builder(outer); OpBuilder builder(outer);
OpBuilder innerBuilder(inner.getBody(), inner.getBody()->begin()); OpBuilder innerBuilder = OpBuilder::atBlockBegin(inner.getBody());
bondhugulaUnsubmitted
Done
ReplyInline Actions

inner.getBody() to avoid confusion since this is a single block region.

bondhugula: `inner.getBody()` to avoid confusion since this is a single block region.
tpoppAuthorUnsubmitted
Done
ReplyInline Actions

I'll block this on whatever happens with https://reviews.llvm.org/D77060

There's an annoying detail right now that constructing an OpBuilder on a Block vs Region create different insertion points (end() and begin()) respectively, so I'll revert this line and another below if it resolves in a different way, or do as you suggest here if it lands.

tpopp: I'll block this on whatever happens with https://reviews.llvm.org/D77060 There's an annoying…
tpoppAuthorUnsubmitted
Done
ReplyInline Actions

This is now OpBuilder::atBlockBegin() in order to be more explicit.

tpopp: This is now `OpBuilder::atBlockBegin()` in order to be more explicit.
auto loopPieces = auto loopPieces =
normalizeLoop(builder, innerBuilder, loop.getLoc(), loop.lowerBound(), normalizeLoop(builder, innerBuilder, loop.getLoc(), loop.lowerBound(),
loop.upperBound(), loop.step(), loop.getInductionVar()); loop.upperBound(), loop.step(), loop.getInductionVar());
loop.setLowerBound(loopPieces.lowerBound); loop.setLowerBound(loopPieces.lowerBound);
loop.setUpperBound(loopPieces.upperBound); loop.setUpperBound(loopPieces.upperBound);
loop.setStep(loopPieces.step); loop.setStep(loopPieces.step);
} }
▲ Show 20 Lines • Show All 48 Lines • ▼ Show 20 Linesvoid mlir::coalesceLoops(MutableArrayRef<loop::ForOp> loops) {
loop::ForOp second = loops[1]; loop::ForOp second = loops[1];
innermost.getBody()->back().erase(); innermost.getBody()->back().erase();
outermost.getBody()->getOperations().splice( outermost.getBody()->getOperations().splice(
Block::iterator(second.getOperation()), Block::iterator(second.getOperation()),
innermost.getBody()->getOperations()); innermost.getBody()->getOperations());
second.erase(); second.erase();
} }
void mlir::collapsePLoops(loop::ParallelOp loops, void mlir::collapseParallelLoops(
loop::ParallelOp loops,
ArrayRef<std::vector<unsigned>> combinedDimensions) { ArrayRef<std::vector<unsigned>> combinedDimensions) {
OpBuilder outsideBuilder(loops); OpBuilder outsideBuilder(loops);
Location loc = loops.getLoc(); Location loc = loops.getLoc();
// Normalize ParallelOp's iteration pattern. // Normalize ParallelOp's iteration pattern.
SmallVector<Value, 3> normalizedLowerBounds; SmallVector<Value, 3> normalizedLowerBounds;
SmallVector<Value, 3> normalizedSteps; SmallVector<Value, 3> normalizedSteps;
SmallVector<Value, 3> normalizedUpperBounds; SmallVector<Value, 3> normalizedUpperBounds;
for (unsigned i = 0, e = loops.getNumLoops(); i < e; ++i) { for (unsigned i = 0, e = loops.getNumLoops(); i < e; ++i) {
OpBuilder insideLoopBuilder(loops.getBody(), loops.getBody()->begin()); OpBuilder insideLoopBuilder = OpBuilder::atBlockBegin(loops.getBody());
auto resultBounds = auto resultBounds =
normalizeLoop(outsideBuilder, insideLoopBuilder, loc, normalizeLoop(outsideBuilder, insideLoopBuilder, loc,
loops.lowerBound()[i], loops.upperBound()[i], loops.lowerBound()[i], loops.upperBound()[i],
loops.step()[i], loops.getBody()->getArgument(i)); loops.step()[i], loops.getBody()->getArgument(i));
normalizedLowerBounds.push_back(resultBounds.lowerBound); normalizedLowerBounds.push_back(resultBounds.lowerBound);
normalizedUpperBounds.push_back(resultBounds.upperBound); normalizedUpperBounds.push_back(resultBounds.upperBound);
normalizedSteps.push_back(resultBounds.step); normalizedSteps.push_back(resultBounds.step);
} }
// Combine iteration spaces // Combine iteration spaces.
SmallVector<Value, 3> lowerBounds; SmallVector<Value, 3> lowerBounds;
SmallVector<Value, 3> steps; SmallVector<Value, 3> steps;
SmallVector<Value, 3> upperBounds; SmallVector<Value, 3> upperBounds;
auto cst0 = outsideBuilder.create<ConstantIndexOp>(loc, 0); auto cst0 = outsideBuilder.create<ConstantIndexOp>(loc, 0);
auto cst1 = outsideBuilder.create<ConstantIndexOp>(loc, 1); auto cst1 = outsideBuilder.create<ConstantIndexOp>(loc, 1);
for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) { for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) {
Value newUpperBound = outsideBuilder.create<ConstantIndexOp>(loc, 1); Value newUpperBound = outsideBuilder.create<ConstantIndexOp>(loc, 1);
for (auto idx : combinedDimensions[i]) { for (auto idx : combinedDimensions[i]) {
newUpperBound = outsideBuilder.create<MulIOp>(loc, newUpperBound, newUpperBound = outsideBuilder.create<MulIOp>(loc, newUpperBound,
normalizedUpperBounds[idx]); normalizedUpperBounds[idx]);
} }
lowerBounds.push_back(cst0); lowerBounds.push_back(cst0);
steps.push_back(cst1); steps.push_back(cst1);
upperBounds.push_back(newUpperBound); upperBounds.push_back(newUpperBound);
} }
// Create new ParallelLoop with conversions to the original induction values. // Create new ParallelLoop with conversions to the original induction values.
// The loop below uses divisions to get the relevant range of values in the // The loop below uses divisions to get the relevant range of values in the
// new induction value that represent each range of the original induction // new induction value that represent each range of the original induction
// value. The remainders then determine based on that range, which iteration // value. The remainders then determine based on that range, which iteration
// of the original induction value this represents. This is a normalized value // of the original induction value this represents. This is a normalized value
// that is un-normalized already by the previous logic. // that is un-normalized already by the previous logic.
auto newPloop = outsideBuilder.create<loop::ParallelOp>(loc, lowerBounds, auto newPloop = outsideBuilder.create<loop::ParallelOp>(loc, lowerBounds,
upperBounds, steps); upperBounds, steps);
OpBuilder insideBuilder(newPloop.getBody(), newPloop.getBody()->begin()); OpBuilder insideBuilder(newPloop.region());
for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) { for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) {
Value previous = newPloop.getBody()->getArgument(i); Value previous = newPloop.getBody()->getArgument(i);
unsigned numberCombinedDimensions = combinedDimensions[i].size(); unsigned numberCombinedDimensions = combinedDimensions[i].size();
// Iterate over all except the last induction value. // Iterate over all except the last induction value.
for (unsigned j = 0, e = numberCombinedDimensions - 1; j < e; ++j) { for (unsigned j = 0, e = numberCombinedDimensions - 1; j < e; ++j) {
unsigned idx = combinedDimensions[i][j]; unsigned idx = combinedDimensions[i][j];
// Determine the current induction value's current loop iteration // Determine the current induction value's current loop iteration
▲ Show 20 Lines • Show All 991 Lines • Show Last 20 Lines

mlir/test/Transforms/parallel-loop-collapsing.mlir

// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='func(parallel-loop-collapsing{collapsed-indices-0=0,3 collapsed-indices-1=1,4 collapsed-indices-2=2}, canonicalize)' | FileCheck %s // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='func(parallel-loop-collapsing{collapsed-indices-0=0,3 collapsed-indices-1=1,4 collapsed-indices-2=2}, canonicalize)' | FileCheck %s
// CHECK-LABEL: func @parallel_many_dims() {
func @parallel_many_dims() { func @parallel_many_dims() {
// CHECK: [[VAL_0:%.*]] = constant 6 : index
// CHECK: [[VAL_1:%.*]] = constant 7 : index
// CHECK: [[VAL_2:%.*]] = constant 9 : index
// CHECK: [[VAL_3:%.*]] = constant 10 : index
// CHECK: [[VAL_4:%.*]] = constant 12 : index
// CHECK: [[VAL_5:%.*]] = constant 13 : index
// CHECK: [[VAL_6:%.*]] = constant 3 : index
// CHECK: [[VAL_7:%.*]] = constant 0 : index
// CHECK: [[VAL_8:%.*]] = constant 1 : index
// CHECK: [[VAL_9:%.*]] = constant 2 : index
%c0 = constant 0 : index %c0 = constant 0 : index
%c1 = constant 1 : index %c1 = constant 1 : index
%c2 = constant 2 : index %c2 = constant 2 : index
%c3 = constant 3 : index %c3 = constant 3 : index
%c4 = constant 4 : index %c4 = constant 4 : index
%c5 = constant 5 : index %c5 = constant 5 : index
%c6 = constant 6 : index %c6 = constant 6 : index
%c7 = constant 7 : index %c7 = constant 7 : index
%c8 = constant 8 : index %c8 = constant 8 : index
%c9 = constant 9 : index %c9 = constant 9 : index
%c10 = constant 10 : index %c10 = constant 10 : index
%c11 = constant 11 : index %c11 = constant 11 : index
%c12 = constant 12 : index %c12 = constant 12 : index
%c13 = constant 13 : index %c13 = constant 13 : index
%c14 = constant 14 : index %c14 = constant 14 : index
// CHECK: loop.parallel ([[VAL_10:%.*]], [[VAL_11:%.*]], [[VAL_12:%.*]]) = ([[VAL_7]], [[VAL_7]], [[VAL_7]]) to ([[VAL_9]], [[VAL_8]], [[VAL_8]]) step ([[VAL_8]], [[VAL_8]], [[VAL_8]]) {
loop.parallel (%i0, %i1, %i2, %i3, %i4) = (%c0, %c3, %c6, %c9, %c12) to (%c2, %c5, %c8, %c11, %c14) loop.parallel (%i0, %i1, %i2, %i3, %i4) = (%c0, %c3, %c6, %c9, %c12) to (%c2, %c5, %c8, %c11, %c14)
step (%c1, %c4, %c7, %c10, %c13) { step (%c1, %c4, %c7, %c10, %c13) {
// CHECK: [[VAL_13:%.*]] = remi_signed [[VAL_10]], [[VAL_9]] : index
// CHECK: [[VAL_14:%.*]] = divi_signed [[VAL_10]], [[VAL_8]] : index
// CHECK: [[VAL_15:%.*]] = divi_signed [[VAL_11]], [[VAL_8]] : index
// CHECK: [[VAL_16:%.*]] = muli [[VAL_15]], [[VAL_5]] : index
// CHECK: [[VAL_17:%.*]] = addi [[VAL_16]], [[VAL_4]] : index
// CHECK: [[VAL_18:%.*]] = muli [[VAL_14]], [[VAL_3]] : index
// CHECK: [[VAL_19:%.*]] = addi [[VAL_18]], [[VAL_2]] : index
// CHECK: [[VAL_20:%.*]] = muli [[VAL_12]], [[VAL_1]] : index
// CHECK: [[VAL_21:%.*]] = addi [[VAL_20]], [[VAL_0]] : index
// CHECK: [[VAL_22:%.*]] = "magic.op"([[VAL_13]], [[VAL_6]], [[VAL_21]], [[VAL_19]], [[VAL_17]]) : (index, index, index, index, index) -> index
%result = "magic.op"(%i0, %i1, %i2, %i3, %i4): (index, index, index, index, index) -> index %result = "magic.op"(%i0, %i1, %i2, %i3, %i4): (index, index, index, index, index) -> index
} }
return return
} }
// CHECK: loop.yield
// CHECK: }
// CHECK: return
// CHECK: }
// CHECK-LABEL: func @parallel_many_dims() {
// CHECK: [[C6:%.*]] = constant 6 : index
// CHECK: [[C7:%.*]] = constant 7 : index
// CHECK: [[C9:%.*]] = constant 9 : index
// CHECK: [[C10:%.*]] = constant 10 : index
// CHECK: [[C12:%.*]] = constant 12 : index
// CHECK: [[C13:%.*]] = constant 13 : index
// CHECK: [[C3:%.*]] = constant 3 : index
// CHECK: [[C0:%.*]] = constant 0 : index
// CHECK: [[C1:%.*]] = constant 1 : index
// CHECK: [[C2:%.*]] = constant 2 : index
// CHECK: loop.parallel ([[NEW_I0:%.*]], [[NEW_I1:%.*]], [[NEW_I2:%.*]]) = ([[C0]], [[C0]], [[C0]]) to ([[C2]], [[C1]], [[C1]]) step ([[C1]], [[C1]], [[C1]]) {
// CHECK: [[I0:%.*]] = remi_signed [[NEW_I0]], [[C2]] : index
// CHECK: [[I3_COUNT:%.*]] = divi_signed [[NEW_I0]], [[C1]] : index
// CHECK: [[I4_COUNT:%.*]] = divi_signed [[NEW_I1]], [[C1]] : index
// CHECK: [[VAL_16:%.*]] = muli [[I4_COUNT]], [[C13]] : index
// CHECK: [[I4:%.*]] = addi [[VAL_16]], [[C12]] : index
// CHECK: [[VAL_18:%.*]] = muli [[I3_COUNT]], [[C10]] : index
// CHECK: [[I3:%.*]] = addi [[VAL_18]], [[C9]] : index
// CHECK: [[VAL_20:%.*]] = muli [[NEW_I2]], [[C7]] : index
// CHECK: [[I2:%.*]] = addi [[VAL_20]], [[C6]] : index
// CHECK: "magic.op"([[I0]], [[C3]], [[I2]], [[I3]], [[I4]]) : (index, index, index, index, index) -> index
// CHECK: loop.yield
// CHECK-NEXT: }
// CHECK-NEXT: return

mlir/test/Transforms/single-parallel-loop-collapsing.mlir

// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='func(parallel-loop-collapsing{collapsed-indices-0=0,1}, canonicalize)' | FileCheck %s // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='func(parallel-loop-collapsing{collapsed-indices-0=0,1}, canonicalize)' | FileCheck %s
// CHECK-LABEL: func @collapse_to_single() {
func @collapse_to_single() { func @collapse_to_single() {
// CHECK: [[VAL_0:%.*]] = constant 7 : index
// CHECK: [[VAL_1:%.*]] = constant 4 : index
// CHECK: [[VAL_2:%.*]] = constant 18 : index
// CHECK: [[VAL_3:%.*]] = constant 3 : index
// CHECK: [[VAL_4:%.*]] = constant 6 : index
// CHECK: [[VAL_5:%.*]] = constant 0 : index
// CHECK: [[VAL_6:%.*]] = constant 1 : index
%c0 = constant 3 : index %c0 = constant 3 : index
%c1 = constant 7 : index %c1 = constant 7 : index
%c2 = constant 11 : index %c2 = constant 11 : index
%c3 = constant 29 : index %c3 = constant 29 : index
%c4 = constant 3 : index %c4 = constant 3 : index
%c5 = constant 4 : index %c5 = constant 4 : index
// CHECK: loop.parallel ([[VAL_7:%.*]]) = ([[VAL_5]]) to ([[VAL_2]]) step ([[VAL_6]]) {
loop.parallel (%i0, %i1) = (%c0, %c1) to (%c2, %c3) step (%c4, %c5) { loop.parallel (%i0, %i1) = (%c0, %c1) to (%c2, %c3) step (%c4, %c5) {
// CHECK: [[VAL_8:%.*]] = remi_signed [[VAL_7]], [[VAL_3]] : index
// CHECK: [[VAL_9:%.*]] = divi_signed [[VAL_7]], [[VAL_4]] : index
// CHECK: [[VAL_10:%.*]] = muli [[VAL_9]], [[VAL_1]] : index
// CHECK: [[VAL_11:%.*]] = addi [[VAL_10]], [[VAL_0]] : index
// CHECK: [[VAL_12:%.*]] = muli [[VAL_8]], [[VAL_3]] : index
// CHECK: [[VAL_13:%.*]] = addi [[VAL_12]], [[VAL_3]] : index
// CHECK: [[VAL_14:%.*]] = "magic.op"([[VAL_13]], [[VAL_11]]) : (index, index) -> index
%result = "magic.op"(%i0, %i1): (index, index) -> index %result = "magic.op"(%i0, %i1): (index, index) -> index
} }
return return
} }
// CHECK: loop.yield
// CHECK: }
// CHECK: return
// CHECK: }
// CHECK-LABEL: func @collapse_to_single() {
// CHECK: [[C7:%.*]] = constant 7 : index
// CHECK: [[C4:%.*]] = constant 4 : index
// CHECK: [[C18:%.*]] = constant 18 : index
// CHECK: [[C3:%.*]] = constant 3 : index
// CHECK: [[C6:%.*]] = constant 6 : index
// CHECK: [[C0:%.*]] = constant 0 : index
// CHECK: [[C1:%.*]] = constant 1 : index
// CHECK: loop.parallel ([[NEW_I:%.*]]) = ([[C0]]) to ([[C18]]) step ([[C1]]) {
// CHECK: [[I0_COUNT:%.*]] = remi_signed [[NEW_I]], [[C3]] : index
// CHECK: [[I1_COUNT:%.*]] = divi_signed [[NEW_I]], [[C6]] : index
// CHECK: [[VAL_10:%.*]] = muli [[I1_COUNT]], [[C4]] : index
// CHECK: [[I1:%.*]] = addi [[VAL_10]], [[C7]] : index
// CHECK: [[VAL_12:%.*]] = muli [[I0_COUNT]], [[C3]] : index
// CHECK: [[I0:%.*]] = addi [[VAL_12]], [[C3]] : index
// CHECK: "magic.op"([[I0]], [[I1]]) : (index, index) -> index
// CHECK: loop.yield
// CHECK-NEXT: }
// CHECK-NEXT: return