Index: mlir/include/mlir/Analysis/AffineAnalysis.h =================================================================== --- mlir/include/mlir/Analysis/AffineAnalysis.h +++ mlir/include/mlir/Analysis/AffineAnalysis.h @@ -40,6 +40,10 @@ Value value; }; +/// Populate `supportedReductions` with descriptors of the supported reductions. +void getSupportedReductions( + AffineForOp forOp, SmallVectorImpl *supportedReductions); + /// Returns true if `forOp' is a parallel loop. If `parallelReductions` is /// provided, populates it with descriptors of the parallelizable reductions and /// treats them as not preventing parallelization. Index: mlir/lib/Analysis/AffineAnalysis.cpp =================================================================== --- mlir/lib/Analysis/AffineAnalysis.cpp +++ mlir/lib/Analysis/AffineAnalysis.cpp @@ -98,6 +98,20 @@ return nullptr; } +/// Populate `supportedReductions` with descriptors of the supported reductions. +void mlir::getSupportedReductions( + AffineForOp forOp, SmallVectorImpl *supportedReductions) { + unsigned numIterArgs = forOp.getNumIterOperands(); + if (numIterArgs == 0 || !supportedReductions) + return; + supportedReductions->reserve(numIterArgs); + for (unsigned i = 0; i < numIterArgs; ++i) { + AtomicRMWKind kind; + if (Value value = getSupportedReduction(forOp, i, kind)) + supportedReductions->emplace_back(LoopReduction{kind, i, value}); + } +} + /// Returns true if `forOp' is a parallel loop. If `parallelReductions` is /// provided, populates it with descriptors of the parallelizable reductions and /// treats them as not preventing parallelization. @@ -112,13 +126,7 @@ // Find supported reductions of requested. if (parallelReductions) { - parallelReductions->reserve(forOp.getNumIterOperands()); - for (unsigned i = 0; i < numIterArgs; ++i) { - AtomicRMWKind kind; - if (Value value = getSupportedReduction(forOp, i, kind)) - parallelReductions->emplace_back(LoopReduction{kind, i, value}); - } - + getSupportedReductions(forOp, parallelReductions); // Return later to allow for identifying all parallel reductions even if the // loop is not parallel. if (parallelReductions->size() != numIterArgs) Index: mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp =================================================================== --- mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp +++ mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp @@ -34,6 +34,7 @@ //===----------------------------------------------------------------------===// #include "PassDetail.h" +#include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/LoopAnalysis.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/Passes.h" @@ -73,6 +74,9 @@ // unroll-and-jammed by this pass. However, runOnAffineForOp can be called on // any for operation. auto &entryBlock = getFunction().front(); - if (auto forOp = dyn_cast(entryBlock.front())) - (void)loopUnrollJamByFactor(forOp, unrollJamFactor); + if (auto forOp = dyn_cast(entryBlock.front())) { + SmallVector reductions; + if (forOp.getNumIterOperands() == 0 || isLoopParallel(forOp, &reductions)) + (void)loopUnrollJamByFactor(forOp, unrollJamFactor); + } } Index: mlir/lib/Transforms/Utils/LoopUtils.cpp =================================================================== --- mlir/lib/Transforms/Utils/LoopUtils.cpp +++ mlir/lib/Transforms/Utils/LoopUtils.cpp @@ -1387,6 +1387,27 @@ return failure(); } + // 1.Check if any control operand of any loop is defined within `forOp` and + // return failure if so. 2.Collect loops with iter_args. + SmallVector loopsWithIterArgs; + auto walkResult = forOp.walk([&](AffineForOp aForOp) { + for (auto controlOperand : + aForOp.getOperands().take_front(aForOp.getNumControlOperands())) { + if (!forOp.isDefinedOutsideOfLoop(controlOperand)) + return WalkResult::interrupt(); + } + if (aForOp.getNumIterOperands() > 0) + loopsWithIterArgs.push_back(aForOp); + return WalkResult::advance(); + }); + if (walkResult.wasInterrupted()) + return failure(); + + // Get supported reductions to be used for creating reduction ops at the end. + SmallVector reductions; + if (forOp.getNumIterOperands() > 0) + getSupportedReductions(forOp, &reductions); + // Gather all sub-blocks to jam upon the loop being unrolled. JamBlockGatherer jbg; jbg.walk(forOp); @@ -1406,6 +1427,19 @@ cleanupOperands); cleanupAffineForOp.setLowerBound(cleanupOperands, cleanupMap); + if (forOp.getNumResults() > 0) { + // Update uses of `forOp` results. `cleanupAffineForOp` should use `forOp` + // results and produce results for the original users of `forOp` results. + ValueRange results = forOp.getResults(); + ValueRange cleanupResults = cleanupAffineForOp.getResults(); + ValueRange cleanupIterOperands = cleanupAffineForOp.getIterOperands(); + + for (auto it : llvm::zip(results, cleanupResults, cleanupIterOperands)) { + std::get<0>(it).replaceAllUsesWith(std::get<1>(it)); + cleanupAffineForOp->replaceUsesOfWith(std::get<2>(it), std::get<0>(it)); + } + } + // Promote the cleanup loop if it has turned into a single iteration loop. (void)promoteIfSingleIteration(cleanupAffineForOp); @@ -1414,6 +1448,55 @@ forOp.setUpperBound(cleanupOperands, cleanupMap); } + // `operandMaps[i]` carries old->new operand mapping for the ith unrolled + // iteration. There are (`unrollJamFactor` - 1) iterations. + SmallVector operandMaps(unrollJamFactor - 1); + + // For any loop with iter_args, replace it with a new loop that has + // `unrollJamFactor` copies of its iterOperands, iter_args and yield + // operands. + SmallVector newLoopsWithIterArgs; + OpBuilder builder(forOp.getContext()); + for (AffineForOp oldForOp : loopsWithIterArgs) { + SmallVector dupIterOperands, dupIterArgs, dupYieldOperands; + ValueRange oldIterOperands = oldForOp.getIterOperands(); + ValueRange oldIterArgs = oldForOp.getRegionIterArgs(); + ValueRange oldYieldOperands = + cast(oldForOp.getBody()->getTerminator()).getOperands(); + // Get additional iterOperands, iterArgs, and yield operands. We will + // fix iterOperands and yield operands after cloning of sub-blocks. + for (unsigned i = unrollJamFactor - 1; i >= 1; --i) { + dupIterOperands.append(oldIterOperands.begin(), oldIterOperands.end()); + dupIterArgs.append(oldIterArgs.begin(), oldIterArgs.end()); + dupYieldOperands.append(oldYieldOperands.begin(), oldYieldOperands.end()); + } + // Create a new loop with additional iterOperands, iter_args and yield + // operands. This new loop will take the loop body of the original loop. + AffineForOp newForOp = mlir::replaceForOpWithNewYields( + builder, oldForOp, dupIterOperands, dupYieldOperands, dupIterArgs); + newLoopsWithIterArgs.push_back(newForOp); + // `forOp` has been replaced with a new loop. + if (oldForOp == forOp) + forOp = newForOp; + assert(oldForOp.use_empty() && "old for op should not have any user"); + oldForOp.erase(); + // Update `operandMaps` for `newForOp` iterArgs and results. + ValueRange newIterArgs = newForOp.getRegionIterArgs(); + unsigned oldNumIterArgs = oldIterArgs.size(); + ValueRange newResults = newForOp.getResults(); + unsigned oldNumResults = newResults.size() / unrollJamFactor; + assert(oldNumIterArgs == oldNumResults && + "oldNumIterArgs must be the same as oldNumResults"); + for (unsigned i = unrollJamFactor - 1; i >= 1; --i) { + for (unsigned j = 0; j < oldNumIterArgs; ++j) { + operandMaps[i - 1].map(newIterArgs[j], + newIterArgs[i * oldNumIterArgs + j]); + operandMaps[i - 1].map(newResults[j], + newResults[i * oldNumResults + j]); + } + } + } + // Scale the step of loop being unroll-jammed by the unroll-jam factor. int64_t step = forOp.getStep(); forOp.setStep(step * unrollJamFactor); @@ -1436,11 +1519,58 @@ auto bumpMap = AffineMap::get(1, 0, d0 + i * step); auto ivUnroll = builder.create(forOp.getLoc(), bumpMap, forOpIV); - operandMap.map(forOpIV, ivUnroll); + operandMaps[i - 1].map(forOpIV, ivUnroll); } // Clone the sub-block being unroll-jammed. for (auto it = subBlock.first; it != std::next(subBlock.second); ++it) - builder.clone(*it, operandMap); + builder.clone(*it, operandMaps[i - 1]); + } + // Fix iterOperands and yield op operands of newly created loops. + for (auto newForOp : newLoopsWithIterArgs) { + unsigned oldNumIterOperands = + newForOp.getNumIterOperands() / unrollJamFactor; + unsigned numControlOperands = newForOp.getNumControlOperands(); + auto yieldOp = cast(newForOp.getBody()->getTerminator()); + unsigned oldNumYieldOperands = yieldOp.getNumOperands() / unrollJamFactor; + assert(oldNumIterOperands == oldNumYieldOperands && + "oldNumIterOperands must be the same as oldNumYieldOperands"); + for (unsigned j = 0; j < oldNumIterOperands; ++j) { + newForOp.setOperand(numControlOperands + i * oldNumIterOperands + j, + operandMaps[i - 1].lookupOrDefault( + newForOp.getOperand(numControlOperands + j))); + yieldOp.setOperand( + i * oldNumYieldOperands + j, + operandMaps[i - 1].lookupOrDefault(yieldOp.getOperand(j))); + } + } + } + + if (forOp.getNumResults() > 0) { + // Add reduction ops after `forOp` since the number of the new results is + // `unrollJamFactor` times of the number of the old results. + // For example, if `unrollJamFactor` is 2, reduction kind is addf, + // and we get %0:2 = affine.for ..., we need to add + // %1 = addf %0#0, %0#1 + // and replace the following uses of %0#0 with %1. + builder.setInsertionPointAfter(forOp); + auto loc = forOp.getLoc(); + unsigned oldNumResults = forOp.getNumResults() / unrollJamFactor; + for (LoopReduction &reduction : reductions) { + unsigned pos = reduction.iterArgPosition; + Value lhs = forOp.getResult(pos); + Value rhs; + SmallPtrSet newOps; + for (unsigned i = unrollJamFactor - 1; i >= 1; --i) { + rhs = forOp.getResult(i * oldNumResults + pos); + // Create ops based on reduction type. + lhs = getReductionOp(reduction.kind, builder, loc, lhs, rhs); + if (!lhs) + return failure(); + Operation *op = lhs.getDefiningOp(); + newOps.insert(op); + } + // Do not replace the use in newly created reduction ops. + forOp.getResult(pos).replaceAllUsesExcept(lhs, newOps); } } Index: mlir/test/Dialect/Affine/unroll-jam.mlir =================================================================== --- mlir/test/Dialect/Affine/unroll-jam.mlir +++ mlir/test/Dialect/Affine/unroll-jam.mlir @@ -122,3 +122,414 @@ // CHECK-NEXT: } // CHECK-NEXT: } // CHECK-NEXT: return + +// The inner loop trip count changes each iteration of outer loop. +// Do no unroll-and-jam. +// CHECK-LABEL: func @no_unroll_jam_dependent_ubound +func @no_unroll_jam_dependent_ubound(%in0: memref) { + affine.for %i = 0 to 100 { + affine.for %k = 0 to affine_map<(d0) -> (d0 + 1)>(%i) { + %y = "addi32"(%k, %k) : (index, index) -> i32 + } + } + return +} +// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 { +// CHECK-NEXT: affine.for [[IV1:%arg[0-9]+]] = 0 to [[$MAP_PLUS_1]]([[IV0]]) { +// CHECK-NEXT: "addi32"([[IV1]], [[IV1]]) +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: return + +// Inner loop with one iter_arg. +// CHECK-LABEL: func @unroll_jam_one_iter_arg +func @unroll_jam_one_iter_arg() { + affine.for %i = 0 to 101 { + %cst = constant 1 : i32 + %x = "addi32"(%i, %i) : (index, index) -> i32 + %red = affine.for %j = 0 to 17 iter_args(%acc = %cst) -> (i32) { + %y = "bar"(%i, %j, %acc) : (index, index, i32) -> i32 + affine.yield %y : i32 + } + %w = "foo"(%i, %x, %red) : (index, i32, i32) -> i32 + } + return +} +// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 { +// CHECK-NEXT: [[CONST1:%[a-zA-Z0-9_]*]] = constant 1 : i32 +// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]]) +// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[CONST2:%[a-zA-Z0-9_]*]] = constant 1 : i32 +// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]]) +// CHECK-NEXT: [[RES3:%[0-9]+]]:2 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[CONST1]], [[ACC2:%arg[0-9]+]] = [[CONST2]]) -> (i32, i32) { +// CHECK-NEXT: [[RES4:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[ACC1]]) +// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[ACC2]]) +// CHECK-NEXT: affine.yield [[RES4]], [[RES5]] +// CHECK-NEXT: } +// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0) +// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#1) +// CHECK: } +// Cleanup loop (single iteration). +// CHECK: constant 1 : i32 +// CHECK-NEXT: "addi32"(%c100, %c100) +// CHECK-NEXT: [[RES6:%[0-9]+]] = affine.for +// CHECK-NEXT: [[RES7:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}) +// CHECK-NEXT: affine.yield [[RES7]] : i32 +// CHECK-NEXT: } +// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES6]]) +// CHECK-NEXT: return + +// Inner loop with multiple iter_args. +// CHECK-LABEL: func @unroll_jam_iter_args +func @unroll_jam_iter_args() { + affine.for %i = 0 to 101 { + %cst = constant 0 : i32 + %cst1 = constant 1 : i32 + %x = "addi32"(%i, %i) : (index, index) -> i32 + %red:2 = affine.for %j = 0 to 17 iter_args(%acc = %cst, %acc1 = %cst1) -> (i32, i32) { + %y = "bar"(%i, %j, %acc) : (index, index, i32) -> i32 + %z = "bar1"(%i, %j, %acc1) : (index, index, i32) -> i32 + affine.yield %y, %z : i32, i32 + } + %w = "foo"(%i, %x, %red#0, %red#1) : (index, i32, i32, i32) -> i32 + } + return +} +// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 { +// CHECK-NEXT: [[CONST0:%[a-zA-Z0-9_]*]] = constant 0 : i32 +// CHECK-NEXT: [[CONST1:%[a-zA-Z0-9_]*]] = constant 1 : i32 +// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]]) +// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[CONST2:%[a-zA-Z0-9_]*]] = constant 0 : i32 +// CHECK-NEXT: [[CONST3:%[a-zA-Z0-9_]*]] = constant 1 : i32 +// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]]) +// CHECK-NEXT: [[RES3:%[0-9]+]]:4 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC0:%arg[0-9]+]] = [[CONST0]], [[ACC1:%arg[0-9]+]] = [[CONST1]], +// CHECK-SAME: [[ACC2:%arg[0-9]+]] = [[CONST2]], [[ACC3:%arg[0-9]+]] = [[CONST3]]) -> (i32, i32, i32, i32) { +// CHECK-NEXT: [[RES4:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[ACC0]]) +// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar1"([[IV0]], [[IV1]], [[ACC1]]) +// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[RES6:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[ACC2]]) +// CHECK-NEXT: [[RES7:%[0-9]+]] = "bar1"([[INC1]], [[IV1]], [[ACC3]]) +// CHECK-NEXT: affine.yield [[RES4]], [[RES5]], [[RES6]], [[RES7]] +// CHECK-NEXT: } +// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0, [[RES3]]#1) +// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#2, [[RES3]]#3) +// CHECK: } +// Cleanup loop (single iteration). +// CHECK: constant 0 : i32 +// CHECK-NEXT: constant 1 : i32 +// CHECK-NEXT: "addi32"(%c100, %c100) +// CHECK-NEXT: [[RES8:%[0-9]+]]:2 = affine.for +// CHECK-NEXT: [[RES9:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}) +// CHECK-NEXT: [[RES10:%[0-9]+]] = "bar1"(%c100, {{.*}}, {{.*}}) +// CHECK-NEXT: affine.yield [[RES9]], [[RES10]] : i32, i32 +// CHECK-NEXT: } +// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES8]]#0, [[RES8]]#1) +// CHECK-NEXT: return + +// When an iter operand is a function argument, do not replace any use of the +// operand . +// CHECK-LABEL: func @unroll_jam_iter_args_func_arg +// CHECK-SAME: [[INIT:%arg[0-9]+]]: i32 +func @unroll_jam_iter_args_func_arg(%in: i32) { + affine.for %i = 0 to 101 { + %x = "addi32"(%i, %i) : (index, index) -> i32 + %red = affine.for %j = 0 to 17 iter_args(%acc = %in) -> (i32) { + %y = "bar"(%i, %j, %acc) : (index, index, i32) -> i32 + affine.yield %y : i32 + } + %w = "foo"(%i, %x, %red) : (index, i32, i32) -> i32 + } + return +} +// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 { +// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]]) +// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]]) +// CHECK-NEXT: [[RES3:%[0-9]+]]:2 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[INIT]], [[ACC2:%arg[0-9]+]] = [[INIT]]) -> (i32, i32) { +// CHECK-NEXT: [[RES4:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[ACC1]]) +// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[ACC2]]) +// CHECK-NEXT: affine.yield [[RES4]], [[RES5]] +// CHECK-NEXT: } +// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0) +// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#1) +// CHECK: } +// Cleanup loop (single iteration). +// CHECK: "addi32"(%c100, %c100) +// CHECK-NEXT: [[RES6:%[0-9]+]] = affine.for +// CHECK-NEXT: [[RES7:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}) +// CHECK-NEXT: affine.yield [[RES7]] : i32 +// CHECK-NEXT: } +// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES6]]) +// CHECK-NEXT: return + +// Nested inner loops, each with one iter_arg. The inner most loop uses its +// outer loop's iter_arg as its iter operand. +// CHECK-LABEL: func @unroll_jam_iter_args_nested +func @unroll_jam_iter_args_nested() { + affine.for %i = 0 to 101 { + %cst = constant 1 : i32 + %x = "addi32"(%i, %i) : (index, index) -> i32 + %red = affine.for %j = 0 to 17 iter_args(%acc = %cst) -> (i32) { + %red1 = affine.for %k = 0 to 35 iter_args(%acc1 = %acc) -> (i32) { + %y = "bar"(%i, %j, %k, %acc1) : (index, index, index, i32) -> i32 + affine.yield %y : i32 + } + affine.yield %red1 : i32 + } + %w = "foo"(%i, %x, %red) : (index, i32, i32) -> i32 + } + return +} +// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 { +// CHECK-NEXT: [[CONST1:%[a-zA-Z0-9_]*]] = constant 1 : i32 +// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]]) +// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[CONST2:%[a-zA-Z0-9_]*]] = constant 1 : i32 +// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]]) +// CHECK-NEXT: [[RES3:%[0-9]+]]:2 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[CONST1]], [[ACC2:%arg[0-9]+]] = [[CONST2]]) -> (i32, i32) { +// CHECK-NEXT: [[RES4:%[0-9]+]]:2 = affine.for [[IV2:%arg[0-9]+]] = 0 to 35 iter_args([[ACC3:%arg[0-9]+]] = [[ACC1]], [[ACC4:%arg[0-9]+]] = [[ACC2]]) -> (i32, i32) { +// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[IV2]], [[ACC3]]) +// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[RES6:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[IV2]], [[ACC4]]) +// CHECK-NEXT: affine.yield [[RES5]], [[RES6]] +// CHECK-NEXT: } +// CHECK-NEXT: affine.yield [[RES4]]#0, [[RES4]]#1 +// CHECK-NEXT: } +// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0) +// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#1) +// CHECK: } +// Cleanup loop (single iteration). +// CHECK: constant 1 : i32 +// CHECK-NEXT: "addi32"(%c100, %c100) +// CHECK-NEXT: [[RES6:%[0-9]+]] = affine.for +// CHECK-NEXT: [[RES7:%[0-9]+]] = affine.for +// CHECK-NEXT: [[RES8:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}, {{.*}}) +// CHECK-NEXT: affine.yield [[RES8]] : i32 +// CHECK-NEXT: } +// CHECK-NEXT: affine.yield [[RES7]] : i32 +// CHECK-NEXT: } +// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES6]]) +// CHECK-NEXT: return + +// Nested inner loops, each with one iter_arg. One loop uses its sibling loop's +// result as its iter operand. +// CHECK-LABEL: func @unroll_jam_iter_args_nested_affine_for_result +func @unroll_jam_iter_args_nested_affine_for_result() { + affine.for %i = 0 to 101 { + %cst = constant 1 : i32 + %x = "addi32"(%i, %i) : (index, index) -> i32 + %red = affine.for %j = 0 to 17 iter_args(%acc = %cst) -> (i32) { + %red1 = affine.for %k = 0 to 35 iter_args(%acc1 = %acc) -> (i32) { + %y = "bar"(%i, %j, %k, %acc1) : (index, index, index, i32) -> i32 + affine.yield %acc : i32 + } + %red2 = affine.for %l = 0 to 36 iter_args(%acc2 = %red1) -> (i32) { + %y = "bar"(%i, %j, %l, %acc2) : (index, index, index, i32) -> i32 + affine.yield %y : i32 + } + affine.yield %red2 : i32 + } + %w = "foo"(%i, %x, %red) : (index, i32, i32) -> i32 + } + return +} +// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 { +// CHECK-NEXT: [[CONST1:%[a-zA-Z0-9_]*]] = constant 1 : i32 +// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]]) +// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[CONST2:%[a-zA-Z0-9_]*]] = constant 1 : i32 +// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]]) +// CHECK-NEXT: [[RES3:%[0-9]+]]:2 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[CONST1]], [[ACC2:%arg[0-9]+]] = [[CONST2]]) -> (i32, i32) { +// CHECK-NEXT: [[RES4:%[0-9]+]]:2 = affine.for [[IV2:%arg[0-9]+]] = 0 to 35 iter_args([[ACC3:%arg[0-9]+]] = [[ACC1]], [[ACC4:%arg[0-9]+]] = [[ACC2]]) -> (i32, i32) { +// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[IV2]], [[ACC3]]) +// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[RES6:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[IV2]], [[ACC4]]) +// CHECK-NEXT: affine.yield [[ACC1]], [[ACC2]] +// CHECK-NEXT: } +// CHECK-NEXT: [[RES14:%[0-9]+]]:2 = affine.for [[IV3:%arg[0-9]+]] = 0 to 36 iter_args([[ACC13:%arg[0-9]+]] = [[RES4]]#0, [[ACC14:%arg[0-9]+]] = [[RES4]]#1) -> (i32, i32) { +// CHECK-NEXT: [[RES15:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[IV3]], [[ACC13]]) +// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[RES16:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[IV3]], [[ACC14]]) +// CHECK-NEXT: affine.yield [[RES15]], [[RES16]] +// CHECK-NEXT: } +// CHECK-NEXT: affine.yield [[RES14]]#0, [[RES14]]#1 +// CHECK-NEXT: } +// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0) +// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#1) +// CHECK: } +// Cleanup loop (single iteration). +// CHECK: constant 1 : i32 +// CHECK-NEXT: "addi32"(%c100, %c100) +// CHECK-NEXT: [[RES6:%[0-9]+]] = affine.for +// CHECK-NEXT: [[RES7:%[0-9]+]] = affine.for +// CHECK-NEXT: [[RES8:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}, {{.*}}) +// CHECK-NEXT: affine.yield +// CHECK-NEXT: } +// CHECK-NEXT: [[RES17:%[0-9]+]] = affine.for +// CHECK-NEXT: [[RES18:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}, {{.*}}) +// CHECK-NEXT: affine.yield [[RES18]] : i32 +// CHECK-NEXT: } +// CHECK-NEXT: affine.yield [[RES17]] : i32 +// CHECK-NEXT: } +// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES6]]) +// CHECK-NEXT: return + +// Nested inner loops, each with one or more iter_args. Yeild the same value +// multiple times. +// CHECK-LABEL: func @unroll_jam_iter_args_nested_yield +func @unroll_jam_iter_args_nested_yield() { + affine.for %i = 0 to 101 { + %cst = constant 1 : i32 + %x = "addi32"(%i, %i) : (index, index) -> i32 + %red:3 = affine.for %j = 0 to 17 iter_args(%acc = %cst, %acc1 = %cst, %acc2 = %cst) -> (i32, i32, i32) { + %red1 = affine.for %k = 0 to 35 iter_args(%acc3 = %acc) -> (i32) { + %y = "bar"(%i, %j, %k, %acc3) : (index, index, index, i32) -> i32 + affine.yield %y : i32 + } + %red2:2 = affine.for %l = 0 to 36 iter_args(%acc4 = %acc1, %acc5 = %acc2) -> (i32, i32) { + %y = "bar1"(%i, %j, %l, %acc4, %acc5) : (index, index, index, i32, i32) -> i32 + affine.yield %y, %y : i32, i32 + } + affine.yield %red1, %red1, %red2#1 : i32, i32, i32 + } + %w = "foo"(%i, %x, %red#0, %red#2) : (index, i32, i32, i32) -> i32 + } + return +} +// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 { +// CHECK-NEXT: [[CONST1:%[a-zA-Z0-9_]*]] = constant 1 : i32 +// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]]) +// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[CONST2:%[a-zA-Z0-9_]*]] = constant 1 : i32 +// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]]) +// CHECK-NEXT: [[RES3:%[0-9]+]]:6 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[CONST1]], [[ACC2:%arg[0-9]+]] = [[CONST1]], +// CHECK-SAME: [[ACC3:%arg[0-9]+]] = [[CONST1]], [[ACC4:%arg[0-9]+]] = [[CONST2]], [[ACC5:%arg[0-9]+]] = [[CONST2]], [[ACC6:%arg[0-9]+]] = [[CONST2]]) -> (i32, i32, i32, i32, i32, i32) { +// CHECK-NEXT: [[RES4:%[0-9]+]]:2 = affine.for [[IV2:%arg[0-9]+]] = 0 to 35 iter_args([[ACC7:%arg[0-9]+]] = [[ACC1]], [[ACC8:%arg[0-9]+]] = [[ACC4]]) -> (i32, i32) { +// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[IV2]], [[ACC7]]) +// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[RES6:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[IV2]], [[ACC8]]) +// CHECK-NEXT: affine.yield [[RES5]], [[RES6]] +// CHECK-NEXT: } +// CHECK-NEXT: [[RES14:%[0-9]+]]:4 = affine.for [[IV3:%arg[0-9]+]] = 0 to 36 iter_args([[ACC13:%arg[0-9]+]] = [[ACC2]], [[ACC14:%arg[0-9]+]] = [[ACC3]], +// CHECK-SAME: [[ACC15:%arg[0-9]+]] = [[ACC5]], [[ACC16:%arg[0-9]+]] = [[ACC6]]) -> (i32, i32, i32, i32) { +// CHECK-NEXT: [[RES15:%[0-9]+]] = "bar1"([[IV0]], [[IV1]], [[IV3]], [[ACC13]], [[ACC14]]) +// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: [[RES16:%[0-9]+]] = "bar1"([[INC1]], [[IV1]], [[IV3]], [[ACC15]], [[ACC16]]) +// CHECK-NEXT: affine.yield [[RES15]], [[RES15]], [[RES16]], [[RES16]] +// CHECK-NEXT: } +// CHECK-NEXT: affine.yield [[RES4]]#0, [[RES4]]#0, [[RES14]]#1, [[RES4]]#1, [[RES4]]#1, [[RES14]]#3 +// CHECK-NEXT: } +// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0, [[RES3]]#2) +// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]]) +// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#3, [[RES3]]#5) +// CHECK: } +// Cleanup loop (single iteration). +// CHECK: constant 1 : i32 +// CHECK-NEXT: "addi32"(%c100, %c100) +// CHECK-NEXT: [[RES6:%[0-9]+]]:3 = affine.for +// CHECK-NEXT: [[RES7:%[0-9]+]] = affine.for +// CHECK-NEXT: [[RES8:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}, {{.*}}) +// CHECK-NEXT: affine.yield [[RES8]] : i32 +// CHECK-NEXT: } +// CHECK-NEXT: [[RES17:%[0-9]+]]:2 = affine.for +// CHECK-NEXT: [[RES18:%[0-9]+]] = "bar1"(%c100, {{.*}}, {{.*}}, {{.*}}, {{.*}}) +// CHECK-NEXT: affine.yield [[RES18]], [[RES18]] : i32, i32 +// CHECK-NEXT: } +// CHECK-NEXT: affine.yield [[RES7]], [[RES7]], [[RES17]]#1 : i32, i32, i32 +// CHECK-NEXT: } +// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES6]]#0, [[RES6]]#2) +// CHECK-NEXT: return + +// Do not unroll-jam loops with iter_args that are not parallel. +// CHECK-LABEL: func @unroll_jam_not_supported +// CHECK-SAME: [[INIT:%arg[0-9]+]]: i32 +func @unroll_jam_not_supported(%in : i32) { + %red = affine.for %i = 0 to 17 iter_args(%acc = %in) -> (i32) { + %y = "bar"(%i, %acc) : (index, i32) -> i32 + affine.yield %y : i32 + } + %w = "foo"(%red) : (i32) -> i32 + return +} +// CHECK: [[RES:%[0-9]+]] = affine.for [[IV0:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[INIT]]) -> (i32) { +// CHECK-NEXT: [[RES1:%[0-9]+]] = "bar"([[IV0]], [[ACC1]]) +// CHECK-NEXT: affine.yield [[RES1]] : i32 +// CHECK-NEXT: } +// CHECK-NEXT: "foo"([[RES]]) +// CHECK-NEXT: return + +// CHECK-LABEL: func @unroll_jam_nested_iter_args_mulf +// CHECK-SAME: [[INIT0:%arg[0-9]+]]: f32, [[INIT1:%arg[0-9]+]]: f32 +func @unroll_jam_nested_iter_args_mulf(%arg0: memref<21x30xf32, 1>, %init : f32, %init1 : f32) { + %0 = affine.for %arg3 = 0 to 21 iter_args(%arg4 = %init) -> (f32) { + %1 = affine.for %arg5 = 0 to 30 iter_args(%arg6 = %init1) -> (f32) { + %3 = affine.load %arg0[%arg3, %arg5] : memref<21x30xf32, 1> + %4 = addf %arg6, %3 : f32 + affine.yield %4 : f32 + } + %2 = mulf %arg4, %1 : f32 + affine.yield %2 : f32 + } + return +} + +// CHECK: %[[CONST0:[a-zA-Z0-9_]*]] = constant 20 : index +// CHECK-NEXT: [[RES:%[0-9]+]]:2 = affine.for %[[IV0:arg[0-9]+]] = 0 to 20 step 2 iter_args([[ACC0:%arg[0-9]+]] = [[INIT0]], [[ACC1:%arg[0-9]+]] = [[INIT0]]) -> (f32, f32) { +// CHECK-NEXT: [[RES1:%[0-9]+]]:2 = affine.for %[[IV1:arg[0-9]+]] = 0 to 30 iter_args([[ACC2:%arg[0-9]+]] = [[INIT1]], [[ACC3:%arg[0-9]+]] = [[INIT1]]) -> (f32, f32) { +// CHECK-NEXT: [[LOAD1:%[0-9]+]] = affine.load {{.*}}[%[[IV0]], %[[IV1]]] +// CHECK-NEXT: [[ADD1:%[0-9]+]] = addf [[ACC2]], [[LOAD1]] : f32 +// CHECK-NEXT: %[[INC1:[0-9]+]] = affine.apply [[$MAP_PLUS_1]](%[[IV0]]) +// CHECK-NEXT: [[LOAD2:%[0-9]+]] = affine.load {{.*}}[%[[INC1]], %[[IV1]]] +// CHECK-NEXT: [[ADD2:%[0-9]+]] = addf [[ACC3]], [[LOAD2]] : f32 +// CHECK-NEXT: affine.yield [[ADD1]], [[ADD2]] +// CHECK-NEXT: } +// CHECK-NEXT: [[MUL1:%[0-9]+]] = mulf [[ACC0]], [[RES1]]#0 : f32 +// CHECK-NEXT: affine.apply +// CHECK-NEXT: [[MUL2:%[0-9]+]] = mulf [[ACC1]], [[RES1]]#1 : f32 +// CHECK-NEXT: affine.yield [[MUL1]], [[MUL2]] +// CHECK-NEXT: } +// Reduction op. +// CHECK-NEXT: [[MUL3:%[0-9]+]] = mulf [[RES]]#0, [[RES]]#1 : f32 +// Cleanup loop (single iteration). +// CHECK-NEXT: [[RES2:%[0-9]+]] = affine.for %[[IV2:arg[0-9]+]] = 0 to 30 iter_args([[ACC4:%arg[0-9]+]] = [[INIT1]]) -> (f32) { +// CHECK-NEXT: [[LOAD3:%[0-9]+]] = affine.load {{.*}}[%[[CONST0]], %[[IV2]]] +// CHECK-NEXT: [[ADD3:%[0-9]+]] = addf [[ACC4]], [[LOAD3]] : f32 +// CHECK-NEXT: affine.yield [[ADD3]] : f32 +// CHECK-NEXT: } +// CHECK-NEXT: [[MUL4:%[0-9]+]] = mulf [[MUL3]], [[RES2]] : f32 +// CHECK-NEXT: return + +// CHECK-LABEL: func @unroll_jam_iter_args_addi +// CHECK-SAME: [[INIT0:%arg[0-9]+]]: i32 +func @unroll_jam_iter_args_addi(%arg0: memref<21xi32, 1>, %init : i32) { + %0 = affine.for %arg3 = 0 to 21 iter_args(%arg4 = %init) -> (i32) { + %1 = affine.load %arg0[%arg3] : memref<21xi32, 1> + %2 = addi %arg4, %1 : i32 + affine.yield %2 : i32 + } + return +} + +// CHECK: %[[CONST0:[a-zA-Z0-9_]*]] = constant 20 : index +// CHECK-NEXT: [[RES:%[0-9]+]]:2 = affine.for %[[IV0:arg[0-9]+]] = 0 to 20 step 2 iter_args([[ACC0:%arg[0-9]+]] = [[INIT0]], [[ACC1:%arg[0-9]+]] = [[INIT0]]) -> (i32, i32) { +// CHECK-NEXT: [[LOAD1:%[0-9]+]] = affine.load {{.*}}[%[[IV0]]] +// CHECK-NEXT: [[ADD1:%[0-9]+]] = addi [[ACC0]], [[LOAD1]] : i32 +// CHECK-NEXT: %[[INC1:[0-9]+]] = affine.apply [[$MAP_PLUS_1]](%[[IV0]]) +// CHECK-NEXT: [[LOAD2:%[0-9]+]] = affine.load {{.*}}[%[[INC1]]] +// CHECK-NEXT: [[ADD2:%[0-9]+]] = addi [[ACC1]], [[LOAD2]] : i32 +// CHECK-NEXT: affine.yield [[ADD1]], [[ADD2]] +// CHECK-NEXT: } +// Reduction op. +// CHECK-NEXT: [[ADD3:%[0-9]+]] = addi [[RES]]#0, [[RES]]#1 : i32 +// Cleanup loop (single iteration). +// CHECK-NEXT: [[LOAD3:%[0-9]+]] = affine.load {{.*}}[%[[CONST0]]] +// CHECK-NEXT: [[ADD4:%[0-9]+]] = addi [[ADD3]], [[LOAD3]] : i32 +// CHECK-NEXT: return