diff --git a/mlir/include/mlir/Transforms/LoopUtils.h b/mlir/include/mlir/Transforms/LoopUtils.h --- a/mlir/include/mlir/Transforms/LoopUtils.h +++ b/mlir/include/mlir/Transforms/LoopUtils.h @@ -232,8 +232,8 @@ /// Take the ParallelLoop and for each set of dimension indices, combine them /// into a single dimension. combinedDimensions must contain each index into /// loops exactly once. -void collapsePLoops(loop::ParallelOp loops, - ArrayRef> combinedDimensions); +void collapseParallelLoops(loop::ParallelOp loops, + ArrayRef> combinedDimensions); /// Maps `forOp` for execution on a parallel grid of virtual `processorIds` of /// size given by `numProcessors`. This is achieved by embedding the SSA values diff --git a/mlir/lib/Transforms/ParallelLoopCollapsing.cpp b/mlir/lib/Transforms/ParallelLoopCollapsing.cpp --- a/mlir/lib/Transforms/ParallelLoopCollapsing.cpp +++ b/mlir/lib/Transforms/ParallelLoopCollapsing.cpp @@ -37,7 +37,7 @@ combinedLoops.push_back(clCollapsedIndices1); if (clCollapsedIndices2.size()) combinedLoops.push_back(clCollapsedIndices2); - collapsePLoops(op, combinedLoops); + collapseParallelLoops(op, combinedLoops); }); } diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp --- a/mlir/lib/Transforms/Utils/LoopUtils.cpp +++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp @@ -1226,7 +1226,7 @@ static void normalizeLoop(loop::ForOp loop, loop::ForOp outer, loop::ForOp inner) { OpBuilder builder(outer); - OpBuilder innerBuilder(inner.getBody(), inner.getBody()->begin()); + OpBuilder innerBuilder = OpBuilder::atBlockBegin(inner.getBody()); auto loopPieces = normalizeLoop(builder, innerBuilder, loop.getLoc(), loop.lowerBound(), loop.upperBound(), loop.step(), loop.getInductionVar()); @@ -1291,8 +1291,9 @@ second.erase(); } -void mlir::collapsePLoops(loop::ParallelOp loops, - ArrayRef> combinedDimensions) { +void mlir::collapseParallelLoops( + loop::ParallelOp loops, + ArrayRef> combinedDimensions) { OpBuilder outsideBuilder(loops); Location loc = loops.getLoc(); @@ -1301,7 +1302,7 @@ SmallVector normalizedSteps; SmallVector normalizedUpperBounds; 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 = normalizeLoop(outsideBuilder, insideLoopBuilder, loc, loops.lowerBound()[i], loops.upperBound()[i], @@ -1312,7 +1313,7 @@ normalizedSteps.push_back(resultBounds.step); } - // Combine iteration spaces + // Combine iteration spaces. SmallVector lowerBounds; SmallVector steps; SmallVector upperBounds; @@ -1337,7 +1338,7 @@ // that is un-normalized already by the previous logic. auto newPloop = outsideBuilder.create(loc, lowerBounds, upperBounds, steps); - OpBuilder insideBuilder(newPloop.getBody(), newPloop.getBody()->begin()); + OpBuilder insideBuilder(newPloop.region()); for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) { Value previous = newPloop.getBody()->getArgument(i); unsigned numberCombinedDimensions = combinedDimensions[i].size(); diff --git a/mlir/test/Transforms/parallel-loop-collapsing.mlir b/mlir/test/Transforms/parallel-loop-collapsing.mlir --- a/mlir/test/Transforms/parallel-loop-collapsing.mlir +++ b/mlir/test/Transforms/parallel-loop-collapsing.mlir @@ -1,17 +1,6 @@ // 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() { - // 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 %c1 = constant 1 : index %c2 = constant 2 : index @@ -28,25 +17,35 @@ %c13 = constant 13 : 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) 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 } 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 diff --git a/mlir/test/Transforms/single-parallel-loop-collapsing.mlir b/mlir/test/Transforms/single-parallel-loop-collapsing.mlir --- a/mlir/test/Transforms/single-parallel-loop-collapsing.mlir +++ b/mlir/test/Transforms/single-parallel-loop-collapsing.mlir @@ -1,36 +1,34 @@ // 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() { - // 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 %c1 = constant 7 : index %c2 = constant 11 : index %c3 = constant 29 : index %c4 = constant 3 : 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) { - // 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 } 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