[Mlir-commits] [mlir] [mlir][OpenACC] Emit multi-IV tile and element loops from ACCLoopTiling (PR #211651)
Vijay Kandiah
llvmlistbot at llvm.org
Thu Jul 23 12:47:52 PDT 2026
https://github.com/VijayKandiah created https://github.com/llvm/llvm-project/pull/211651
`ACCLoopTiling` used to lower an N-dimensional `tile()` clause by *uncollapsing* the fused loop into a deep nest of single-IV `acc.loop`s (`tile_1 → … → tile_N → elem_1 → … → elem_N`). This discards the natural grouping of the tile clause: the 2N single-IV loops no longer express "these are the tile iterations" and "these are the in-tile iterations" as collapsible units, and gang/vector end up spread across a deep nest where only the outermost loop of each group carries the attribute. This could potentially lead to poor parallelism assignment.
With this MR:
`tileACCLoops` now rewrites the single fused `acc.loop` in-place into exactly two multi-IV loops, each carrying all the tiled induction variables:
- a **tile-group** loop (gang) whose steps are the original steps scaled by the tile sizes, and
- a nested **element-group** loop (vector) that walks the iterations within one tile, with upper bounds clamped to `min(origUB, tileStart + tileExtent)`.
Keeping each group as a single multi-IV `acc.loop` preserves the tile/element structure as two collapsible units instead of a 2N-deep single-IV nest, which is a simpler canonical form. The gang/worker/vector distribution across the two groups is unchanged from before.
This change also diagnoses `tile` + `collapse` on the same loop because these clauses give conflicting loop-association counts. The pass now emits a NYI diagnostic instead of silently dropping the `collapse` clause. Added a new invalid-input test to verify this.
With the pass no longer uncollapsing, `uncollapseLoops` and the vector-of-loops `tileACCLoops(SmallVector<LoopOp>&, …)` overload (and the now-orphaned `createInnerLoop` helper) are deleted, along with their unit tests. New unit tests are added to cover the surviving fused-loop overload.
>From 3cb4151c7f60f5549dce86110583ed8cbb5f85e0 Mon Sep 17 00:00:00 2001
From: Vijay Kandiah <vkandiah at nvidia.com>
Date: Thu, 23 Jul 2026 12:38:02 -0700
Subject: [PATCH] [mlir][OpenACC] Emit multi-IV tile/element loops from
ACCLoopTiling
---
.../mlir/Dialect/OpenACC/OpenACCUtilsTiling.h | 53 ++-
.../OpenACC/Transforms/ACCLoopTiling.cpp | 32 +-
.../OpenACC/Utils/OpenACCUtilsTiling.cpp | 345 +++++++-----------
.../OpenACC/acc-loop-tiling-invalid.mlir | 23 ++
.../test/Dialect/OpenACC/acc-loop-tiling.mlir | 85 ++---
.../OpenACC/OpenACCUtilsTilingTest.cpp | 225 ++----------
6 files changed, 253 insertions(+), 510 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsTiling.h b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsTiling.h
index 6fcb706aa3488..05ac37d1106ec 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsTiling.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsTiling.h
@@ -20,23 +20,20 @@
namespace mlir {
namespace acc {
-/// Uncollapse tile loops with multiple IVs and collapseCount < tileCount.
-/// This is used to prepare loops for tiling when the collapse count is less
-/// than the tile count.
+/// Tile a single fused acc.loop that carries all associated induction
+/// variables (one IV per tile dimension).
///
-/// \param origLoop The original loop operation to uncollapse.
-/// \param tileCount The number of tile dimensions.
-/// \param collapseCount The collapse count from the original loop.
-/// \param rewriter The rewriter to use for modifications.
-/// \return A vector of uncollapsed loop operations.
-llvm::SmallVector<mlir::acc::LoopOp>
-uncollapseLoops(mlir::acc::LoopOp origLoop, unsigned tileCount,
- unsigned collapseCount, mlir::RewriterBase &rewriter);
-
-/// Tile ACC loops according to the given tile sizes.
+/// This produces exactly two multi-IV loops, each carrying all of the tiled
+/// induction variables:
///
-/// Tiling a 2-level nested loop will create two 'tile' loops containing two
-/// 'element' loops. The transformation looks like:
+/// - a "tile group" loop that steps over tiles: each step is the
+/// original step multiplied by the tile size. It keeps the original loop's
+/// gang attribute.
+/// - an "element group" loop that walks the iterations inside one tile: it
+/// keeps the original step, its lower bound is the current tile's starting
+/// index, and its upper bound is clamped to min(original upper bound,
+/// tile start + tile extent). It keeps the original loop's vector (or
+/// worker) attribute.
///
/// Before Tiling:
/// \code
@@ -48,17 +45,15 @@ uncollapseLoops(mlir::acc::LoopOp origLoop, unsigned tileCount,
/// }
/// \endcode
///
-/// After Tiling:
+/// After Tiling (each group is one multi-IV loop over all tiled IVs):
/// \code
-/// for (i = lb1; i < ub1; i += (step1 * tile_size1)) { // tile loop 1
-/// for (j = lb2; j < ub2; j += (step2 * tile_size2)) { // tile loop 2
-/// for (ii = i; ii < min(ub1, (step1 * tile_size1) + i); ii += step1) {
-/// // element loop 1
-/// for (jj = j; jj < min(ub2, (step2 * tile_size2) + j); jj += step2)
-/// { // element loop 2
-/// a[ii,jj] = i + j;
-/// }
-/// }
+/// // tile group
+/// for (i = lb1; i < ub1; i += (step1 * tile_size1),
+/// j = lb2; j < ub2; j += (step2 * tile_size2)) {
+/// // element group
+/// for (ii = i; ii < min(ub1, (step1 * tile_size1) + i); ii += step1,
+/// jj = j; jj < min(ub2, (step2 * tile_size2) + j); jj += step2) {
+/// a[ii,jj] = i + j;
/// }
/// }
/// \endcode
@@ -66,13 +61,13 @@ uncollapseLoops(mlir::acc::LoopOp origLoop, unsigned tileCount,
/// Unknown tile sizes (represented as -1 in acc dialect for `tile(*)`) are
/// resolved to the provided default tile size.
///
-/// \param tileLoops The loops to tile (outermost first).
-/// \param tileSizes The tile sizes for each dimension. Values of -1 are
+/// \param tileLoop The fused loop to tile.
+/// \param tileSizes The tile sizes for each tiled dimension. Values of -1 are
/// treated as unknown and resolved to defaultTileSize.
/// \param defaultTileSize The default tile size to use for unknown (*) tiles.
/// \param rewriter The rewriter to use for modifications.
-/// \return The outermost loop after tiling.
-mlir::acc::LoopOp tileACCLoops(llvm::SmallVector<mlir::acc::LoopOp> &tileLoops,
+/// \return The tile group loop that is modified in place.
+mlir::acc::LoopOp tileACCLoops(mlir::acc::LoopOp tileLoop,
const llvm::SmallVector<mlir::Value> &tileSizes,
int32_t defaultTileSize,
mlir::RewriterBase &rewriter);
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCLoopTiling.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCLoopTiling.cpp
index 6bc95ca896f37..d5865bc98727e 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCLoopTiling.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCLoopTiling.cpp
@@ -162,8 +162,16 @@ struct ACCLoopTilingImpl : public OpRewritePattern<acc::LoopOp> {
SmallVector<Value> tileSizes(origLoop.getTileValues().begin(),
origLoop.getTileValues().end());
- unsigned tileCount = tileSizes.size();
- unsigned collapseCount = origLoop.getCollapseValue().value_or(1);
+
+ // A tile clause and a collapse clause on the same loop are not supported:
+ // they give conflicting descriptions of how many loops the construct
+ // associates.
+ if (origLoop.getCollapseAttr()) {
+ accSupport.emitNYI(origLoop.getLoc(),
+ "a tile clause combined with a collapse clause on the "
+ "same loop");
+ return failure();
+ }
// Sanity check tile size types
if (failed(checkTileSizeTypes(origLoop, tileSizes)))
@@ -182,24 +190,12 @@ struct ACCLoopTilingImpl : public OpRewritePattern<acc::LoopOp> {
origLoop.removeTileOperandsDeviceTypeAttr();
rewriter.finalizeOpModification(origLoop);
- SmallVector<acc::LoopOp> loopsToTile;
- if (collapseCount < tileCount) {
- // Uncollapse tile loops before tiling if necessary
- loopsToTile =
- acc::uncollapseLoops(origLoop, tileCount, collapseCount, rewriter);
- rewriter.replaceOp(origLoop, loopsToTile[0]);
- LLVM_DEBUG(llvm::dbgs() << "\nAfter uncollapsing:\n"
- << *loopsToTile[0] << "\n");
- } else {
- loopsToTile.push_back(origLoop);
- }
-
- // loopsToTile is a vector of perfectly nested loops. The outermost loop
- // may have multiple IVs but inner loops can only have one IV.
+ // Tile the single fused loop in place: the tile iterations become one
+ // multi-IV "tile group" loop wrapping one multi-IV "element group" loop.
// The utility handles unknown tile sizes (*) by using `defaultTileSize`.
- acc::tileACCLoops(loopsToTile, tileSizes, defaultTileSize, rewriter);
+ acc::tileACCLoops(origLoop, tileSizes, defaultTileSize, rewriter);
- LLVM_DEBUG(llvm::dbgs() << "\nAfter tiling:\n " << *loopsToTile[0] << "\n");
+ LLVM_DEBUG(llvm::dbgs() << "\nAfter tiling:\n " << *origLoop << "\n");
return success();
}
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsTiling.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsTiling.cpp
index 09e83af2bc686..32fd3957894fb 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsTiling.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsTiling.cpp
@@ -17,6 +17,7 @@
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Transforms/RegionUtils.h"
+#include "llvm/ADT/STLExtras.h"
// Resolve unknown tile sizes (represented as -1 for tile(*)) to the default.
// Returns a value with the same type as targetType.
@@ -48,19 +49,13 @@ static void removeWorkerVectorFromLoop(mlir::acc::LoopOp loop) {
}
// Create a new ACC loop with new steps, lb, ub from original loop
-static mlir::acc::LoopOp
-createACCLoopFromOriginal(mlir::acc::LoopOp origLoop,
- mlir::RewriterBase &rewriter, mlir::ValueRange lb,
- mlir::ValueRange ub, mlir::ValueRange step,
- mlir::DenseBoolArrayAttr inclusiveUBAttr,
- mlir::acc::CombinedConstructsTypeAttr combinedAttr,
- mlir::Location loc, bool preserveCollapse) {
+static mlir::acc::LoopOp createACCLoopFromOriginal(
+ mlir::acc::LoopOp origLoop, mlir::RewriterBase &rewriter,
+ mlir::ValueRange lb, mlir::ValueRange ub, mlir::ValueRange step,
+ mlir::DenseBoolArrayAttr inclusiveUBAttr,
+ mlir::acc::CombinedConstructsTypeAttr combinedAttr, mlir::Location loc) {
mlir::ArrayAttr collapseAttr = mlir::ArrayAttr{};
mlir::ArrayAttr collapseDeviceTypeAttr = mlir::ArrayAttr{};
- if (preserveCollapse) {
- collapseAttr = origLoop.getCollapseAttr();
- collapseDeviceTypeAttr = origLoop.getCollapseDeviceTypeAttr();
- }
auto newLoop = mlir::acc::LoopOp::create(
rewriter, loc, origLoop->getResultTypes(), lb, ub, step, inclusiveUBAttr,
collapseAttr, collapseDeviceTypeAttr, origLoop.getGangOperands(),
@@ -78,53 +73,6 @@ createACCLoopFromOriginal(mlir::acc::LoopOp origLoop,
return newLoop;
}
-// Create inner loop inside input loop
-static mlir::acc::LoopOp
-createInnerLoop(mlir::acc::LoopOp inputLoop, mlir::RewriterBase &rewriter,
- mlir::ValueRange lb, mlir::ValueRange ub, mlir::ValueRange step,
- mlir::DenseBoolArrayAttr inclusiveUBAttr, mlir::Location loc) {
- mlir::acc::LoopOp elementLoop = createACCLoopFromOriginal(
- inputLoop, rewriter, lb, ub, step, inclusiveUBAttr,
- mlir::acc::CombinedConstructsTypeAttr{}, loc, /*preserveCollapse*/ false);
-
- // Remove gang/worker attributes from inner loops
- rewriter.startOpModification(elementLoop);
- if (inputLoop.hasGang() ||
- inputLoop.getGangValue(mlir::acc::GangArgType::Num) ||
- inputLoop.getGangValue(mlir::acc::GangArgType::Dim) ||
- inputLoop.getGangValue(mlir::acc::GangArgType::Static)) {
- elementLoop.removeGangAttr();
- elementLoop.removeGangOperandsArgTypeAttr();
- elementLoop.removeGangOperandsSegmentsAttr();
- elementLoop.removeGangOperandsDeviceTypeAttr();
- // Also drop the operand values themselves so that elementLoop does not
- // end up with a non-empty gang operand list but no corresponding
- // device-type/segment/arg-type attributes. Leaving stale operands behind
- // makes elementLoop look like it still has gang operands to later
- // queries (e.g. LoopOp::getGangValue), which then dereference the
- // now-missing device-type attribute.
- elementLoop.getGangOperandsMutable().clear();
- }
- if (inputLoop.hasVector() || inputLoop.getVectorValue()) {
- elementLoop.removeWorkerAttr();
- elementLoop.removeWorkerNumOperandsDeviceTypeAttr();
- // As above for gang, also drop the worker operand values so elementLoop
- // does not keep a dangling worker operand with no device-type attribute.
- elementLoop.getWorkerNumOperandsMutable().clear();
- }
- rewriter.finalizeOpModification(elementLoop);
-
- // Create empty block in elementLoop and add IV argument
- mlir::Block *blk = rewriter.createBlock(&elementLoop.getRegion(),
- elementLoop.getRegion().begin());
- rewriter.setInsertionPointToEnd(blk);
- mlir::acc::YieldOp::create(rewriter, loc);
- elementLoop.getBody().addArgument(
- inputLoop.getBody().getArgument(0).getType(), loc);
-
- return elementLoop;
-}
-
// Move ops from source to target Loop and replace uses of IVs
static void moveOpsAndReplaceIVs(mlir::acc::LoopOp sourceLoop,
mlir::acc::LoopOp targetLoop,
@@ -159,180 +107,145 @@ static void moveOpsAndReplaceIVs(mlir::acc::LoopOp sourceLoop,
rewriter.finalizeOpModification(op);
}
+// Create a single "element group" loop nested in `tileLoop`, carrying
+// `ivTypes.size()` induction variables so the whole element space is one
+// multi-IV loop. The element group carries vector or worker but not gang.
+static mlir::acc::LoopOp
+createElementGroupLoop(mlir::acc::LoopOp tileLoop, mlir::RewriterBase &rewriter,
+ mlir::ValueRange lbs, mlir::ValueRange ubs,
+ mlir::ValueRange steps,
+ mlir::DenseBoolArrayAttr inclusiveUBAttr,
+ llvm::ArrayRef<mlir::Type> ivTypes, mlir::Location loc) {
+ mlir::acc::LoopOp elementLoop = createACCLoopFromOriginal(
+ tileLoop, rewriter, lbs, ubs, steps, inclusiveUBAttr,
+ mlir::acc::CombinedConstructsTypeAttr{}, loc);
+
+ // Drop gang from the element group, keeping vector/worker. The operand
+ // values must be cleared too, not just the attributes.
+ rewriter.startOpModification(elementLoop);
+ if (tileLoop.hasGang() ||
+ tileLoop.getGangValue(mlir::acc::GangArgType::Num) ||
+ tileLoop.getGangValue(mlir::acc::GangArgType::Dim) ||
+ tileLoop.getGangValue(mlir::acc::GangArgType::Static)) {
+ elementLoop.removeGangAttr();
+ elementLoop.removeGangOperandsArgTypeAttr();
+ elementLoop.removeGangOperandsSegmentsAttr();
+ elementLoop.removeGangOperandsDeviceTypeAttr();
+ elementLoop.getGangOperandsMutable().clear();
+ }
+ if (tileLoop.hasVector() || tileLoop.getVectorValue()) {
+ elementLoop.removeWorkerAttr();
+ elementLoop.removeWorkerNumOperandsDeviceTypeAttr();
+ elementLoop.getWorkerNumOperandsMutable().clear();
+ }
+ rewriter.finalizeOpModification(elementLoop);
+
+ // Create the element loop body: one block argument per IV plus a terminator.
+ mlir::Block *blk = rewriter.createBlock(&elementLoop.getRegion(),
+ elementLoop.getRegion().begin());
+ rewriter.setInsertionPointToEnd(blk);
+ mlir::acc::YieldOp::create(rewriter, loc);
+ for (mlir::Type ivType : ivTypes)
+ elementLoop.getBody().addArgument(ivType, loc);
+
+ return elementLoop;
+}
+
mlir::acc::LoopOp
-mlir::acc::tileACCLoops(llvm::SmallVector<mlir::acc::LoopOp> &tileLoops,
+mlir::acc::tileACCLoops(mlir::acc::LoopOp tileLoop,
const llvm::SmallVector<mlir::Value> &tileSizes,
int32_t defaultTileSize, mlir::RewriterBase &rewriter) {
- // Tile collapsed and/or nested loops
- mlir::acc::LoopOp outerLoop = tileLoops[0];
- const mlir::Location loc = outerLoop.getLoc();
-
- mlir::acc::LoopOp innerLoop = tileLoops[tileLoops.size() - 1];
- llvm::SmallVector<mlir::Value, 3> origIVs;
- llvm::SmallVector<mlir::Value, 3> origSteps;
- llvm::SmallVector<mlir::Value, 3> origUBs;
- llvm::SmallVector<mlir::Value, 3> newSteps;
- llvm::SmallVector<mlir::Value, 3> newUBs;
- llvm::SmallVector<mlir::Value, 3> newIVs;
- size_t nOps = innerLoop.getBody().getOperations().size();
-
- // Extract original inclusiveUBs
+ // Tile a single fused acc.loop that carries all associated induction
+ // variables. This keeps the tile iterations as one multi-IV "tile group"
+ // loop and the in-tile iterations as one multi-IV "element group" loop, each
+ // spanning all of its induction variables.
+ const mlir::Location loc = tileLoop.getLoc();
+ const unsigned tileCount = tileSizes.size();
+
+ llvm::SmallVector<mlir::Value, 3> origIVs(tileLoop.getBody().getArguments());
+ llvm::SmallVector<mlir::Value, 3> origUBs(tileLoop.getUpperbound());
+ llvm::SmallVector<mlir::Value, 3> origSteps(tileLoop.getStep());
+ const unsigned numIVs = origIVs.size();
+ const size_t nOps = tileLoop.getBody().getOperations().size();
+
+ // Original inclusive-UB flags (default false when the attribute is absent).
llvm::SmallVector<bool> inclusiveUBs;
- for (auto tileLoop : tileLoops) {
- for (auto [j, step] : llvm::enumerate(tileLoop.getStep())) {
- // inclusiveUBs are present on the IR from Fortran frontend for DO loops
- // but might not be present from other frontends (python)
- // So check if it exists
- if (tileLoop.getInclusiveUpperboundAttr())
- inclusiveUBs.push_back(
- tileLoop.getInclusiveUpperboundAttr().asArrayRef()[j]);
- else
- inclusiveUBs.push_back(false);
- }
+ for (unsigned i = 0; i < numIVs; ++i) {
+ if (tileLoop.getInclusiveUpperboundAttr())
+ inclusiveUBs.push_back(
+ tileLoop.getInclusiveUpperboundAttr().asArrayRef()[i]);
+ else
+ inclusiveUBs.push_back(false);
}
- // Extract original ivs, UBs, steps, and calculate new steps
- rewriter.setInsertionPoint(outerLoop);
- for (auto [i, tileLoop] : llvm::enumerate(tileLoops)) {
- for (auto arg : tileLoop.getBody().getArguments())
- origIVs.push_back(arg);
- for (auto ub : tileLoop.getUpperbound())
- origUBs.push_back(ub);
-
- llvm::SmallVector<mlir::Value, 3> currentLoopSteps;
- for (auto [j, step] : llvm::enumerate(tileLoop.getStep())) {
- origSteps.push_back(step);
- if (i + j >= tileSizes.size()) {
- currentLoopSteps.push_back(step);
- } else {
- mlir::Value tileSize = resolveAndCastTileSize(
- tileSizes[i + j], defaultTileSize, step.getType(), rewriter, loc);
- auto newLoopStep =
- mlir::arith::MulIOp::create(rewriter, loc, step, tileSize);
- currentLoopSteps.push_back(newLoopStep);
- newSteps.push_back(newLoopStep);
- }
+ // Scale each tiled dimension's step by its tile size to form the tile group
+ // loop steps.
+ rewriter.setInsertionPoint(tileLoop);
+ llvm::SmallVector<mlir::Value, 3> scaledSteps;
+ llvm::SmallVector<mlir::Value, 3> tileLoopSteps;
+ for (unsigned i = 0; i < numIVs; ++i) {
+ if (i < tileCount) {
+ mlir::Value tileSize = resolveAndCastTileSize(
+ tileSizes[i], defaultTileSize, origSteps[i].getType(), rewriter, loc);
+ mlir::Value scaled =
+ mlir::arith::MulIOp::create(rewriter, loc, origSteps[i], tileSize);
+ scaledSteps.push_back(scaled);
+ tileLoopSteps.push_back(scaled);
+ } else {
+ tileLoopSteps.push_back(origSteps[i]);
}
-
- rewriter.startOpModification(tileLoop);
- tileLoop.getStepMutable().clear();
- tileLoop.getStepMutable().append(currentLoopSteps);
- rewriter.finalizeOpModification(tileLoop);
}
- // Calculate new upper bounds for element loops
- for (size_t i = 0; i < newSteps.size(); i++) {
- rewriter.setInsertionPoint(innerLoop.getBody().getTerminator());
- // UpperBound: min(origUB, origIV+(originalStep*tile_size))
- auto stepped =
- mlir::arith::AddIOp::create(rewriter, loc, origIVs[i], newSteps[i]);
+ // Compute the element-loop upper bounds min(origUB, origIV + scaledStep).
+ rewriter.setInsertionPoint(tileLoop.getBody().getTerminator());
+ llvm::SmallVector<mlir::Value, 3> elemLBs, elemUBs, elemSteps;
+ llvm::SmallVector<mlir::Type, 3> elemIVTypes;
+ llvm::SmallVector<bool> elemInclusiveUBs;
+ for (unsigned i = 0; i < tileCount; ++i) {
+ mlir::Value stepped =
+ mlir::arith::AddIOp::create(rewriter, loc, origIVs[i], scaledSteps[i]);
mlir::Value newUB = stepped;
if (inclusiveUBs[i]) {
- // Handle InclusiveUB
- // UpperBound: min(origUB, origIV+(originalStep*tile_size - 1))
- auto c1 = mlir::arith::ConstantOp::create(
- rewriter, loc, newSteps[i].getType(),
- rewriter.getIntegerAttr(newSteps[i].getType(), 1));
+ // Inclusive UB: min(origUB, origIV + (scaledStep - 1)).
+ mlir::Value c1 = mlir::arith::ConstantOp::create(
+ rewriter, loc, scaledSteps[i].getType(),
+ rewriter.getIntegerAttr(scaledSteps[i].getType(), 1));
newUB = mlir::arith::SubIOp::create(rewriter, loc, stepped, c1);
}
- newUBs.push_back(
+ elemUBs.push_back(
mlir::arith::MinSIOp::create(rewriter, loc, origUBs[i], newUB));
+ elemLBs.push_back(origIVs[i]);
+ elemSteps.push_back(origSteps[i]);
+ elemIVTypes.push_back(origIVs[i].getType());
+ elemInclusiveUBs.push_back(inclusiveUBs[i]);
}
- // Create and insert nested elementLoopOps before terminator of outer loopOp
- mlir::acc::LoopOp currentLoop = innerLoop;
- for (size_t i = 0; i < tileSizes.size(); i++) {
- rewriter.setInsertionPoint(currentLoop.getBody().getTerminator());
- mlir::DenseBoolArrayAttr inclusiveUBAttr = mlir::DenseBoolArrayAttr{};
- if (inclusiveUBs[i])
- inclusiveUBAttr = rewriter.getDenseBoolArrayAttr({true});
-
- mlir::acc::LoopOp elementLoop =
- createInnerLoop(innerLoop, rewriter, mlir::ValueRange{origIVs[i]},
- mlir::ValueRange{newUBs[i]},
- mlir::ValueRange{origSteps[i]}, inclusiveUBAttr, loc);
-
- // Remove vector/worker attributes from inner element loops except
- // outermost element loop
- if (i > 0) {
- rewriter.startOpModification(elementLoop);
- removeWorkerVectorFromLoop(elementLoop);
- rewriter.finalizeOpModification(elementLoop);
- }
- newIVs.push_back(elementLoop.getBody().getArgument(0));
- currentLoop = elementLoop;
- }
-
- // Remove vector/worker attributes from outer tile loops
- for (auto tileLoop : tileLoops) {
- rewriter.startOpModification(tileLoop);
- removeWorkerVectorFromLoop(tileLoop);
- rewriter.finalizeOpModification(tileLoop);
- }
-
- // Move ops from inner tile loop to inner element loop and replace IV uses
- moveOpsAndReplaceIVs(innerLoop, currentLoop, newIVs, origIVs, nOps, rewriter);
-
- return outerLoop;
-}
-
-llvm::SmallVector<mlir::acc::LoopOp>
-mlir::acc::uncollapseLoops(mlir::acc::LoopOp origLoop, unsigned tileCount,
- unsigned collapseCount,
- mlir::RewriterBase &rewriter) {
- llvm::SmallVector<mlir::acc::LoopOp> newLoops;
- llvm::SmallVector<mlir::Value, 3> newIVs;
- mlir::Location loc = origLoop.getLoc();
- llvm::SmallVector<bool> newInclusiveUBs;
- llvm::SmallVector<mlir::Value, 3> lbs, ubs, steps;
- for (unsigned i = 0; i < collapseCount; i++) {
- // inclusiveUpperbound attribute might not be set, default to false
- bool inclusiveUB = false;
- if (origLoop.getInclusiveUpperboundAttr())
- inclusiveUB = origLoop.getInclusiveUpperboundAttr().asArrayRef()[i];
- newInclusiveUBs.push_back(inclusiveUB);
- lbs.push_back(origLoop.getLowerbound()[i]);
- ubs.push_back(origLoop.getUpperbound()[i]);
- steps.push_back(origLoop.getStep()[i]);
- }
- mlir::acc::LoopOp outerLoop = createACCLoopFromOriginal(
- origLoop, rewriter, lbs, ubs, steps,
- rewriter.getDenseBoolArrayAttr(newInclusiveUBs),
- origLoop.getCombinedAttr(), loc, /*preserveCollapse*/ true);
- mlir::Block *blk = rewriter.createBlock(&outerLoop.getRegion(),
- outerLoop.getRegion().begin());
- rewriter.setInsertionPointToEnd(blk);
- mlir::acc::YieldOp::create(rewriter, loc);
- for (unsigned i = 0; i < collapseCount; i++) {
- outerLoop.getBody().addArgument(origLoop.getBody().getArgument(i).getType(),
- loc);
- newIVs.push_back(outerLoop.getBody().getArgument(i));
- }
- newLoops.push_back(outerLoop);
-
- mlir::acc::LoopOp currentLoopOp = outerLoop;
- for (unsigned i = collapseCount; i < tileCount; i++) {
- rewriter.setInsertionPoint(currentLoopOp.getBody().getTerminator());
- bool inclusiveUB = false;
- if (origLoop.getInclusiveUpperboundAttr())
- inclusiveUB = origLoop.getInclusiveUpperboundAttr().asArrayRef()[i];
- mlir::DenseBoolArrayAttr inclusiveUBAttr =
- rewriter.getDenseBoolArrayAttr({inclusiveUB});
- mlir::acc::LoopOp innerLoop = createInnerLoop(
- origLoop, rewriter, mlir::ValueRange{origLoop.getLowerbound()[i]},
- mlir::ValueRange{origLoop.getUpperbound()[i]},
- mlir::ValueRange{origLoop.getStep()[i]}, inclusiveUBAttr, loc);
- newIVs.push_back(innerLoop.getBody().getArgument(0));
- newLoops.push_back(innerLoop);
- currentLoopOp = innerLoop;
- }
- // Move ops from origLoop to innermost loop and replace uses of IVs
- size_t nOps = origLoop.getBody().getOperations().size();
- llvm::SmallVector<mlir::Value, 3> origIVs;
- for (auto arg : origLoop.getBody().getArguments())
- origIVs.push_back(arg);
- moveOpsAndReplaceIVs(origLoop, currentLoopOp, newIVs, origIVs, nOps,
+ // Only attach an inclusiveUpperbound attribute if at least one element
+ // dimension is inclusive.
+ mlir::DenseBoolArrayAttr elemInclAttr = mlir::DenseBoolArrayAttr{};
+ if (llvm::is_contained(elemInclusiveUBs, true))
+ elemInclAttr = rewriter.getDenseBoolArrayAttr(elemInclusiveUBs);
+
+ // Create the element group loop from the unmodified tile loop.
+ mlir::acc::LoopOp elementLoop =
+ createElementGroupLoop(tileLoop, rewriter, elemLBs, elemUBs, elemSteps,
+ elemInclAttr, elemIVTypes, loc);
+
+ // Move the original body into the element loop and remap the tiled IVs to the
+ // element IVs.
+ llvm::SmallVector<mlir::Value, 3> newIVs(
+ elementLoop.getBody().getArguments());
+ llvm::SmallVector<mlir::Value, 3> tiledOrigIVs(origIVs.begin(),
+ origIVs.begin() + tileCount);
+ moveOpsAndReplaceIVs(tileLoop, elementLoop, newIVs, tiledOrigIVs, nOps,
rewriter);
- return newLoops;
+ // Turn the fused loop into the tile group: scaled steps, gang only.
+ rewriter.startOpModification(tileLoop);
+ tileLoop.getStepMutable().clear();
+ tileLoop.getStepMutable().append(tileLoopSteps);
+ removeWorkerVectorFromLoop(tileLoop);
+ rewriter.finalizeOpModification(tileLoop);
+
+ return tileLoop;
}
diff --git a/mlir/test/Dialect/OpenACC/acc-loop-tiling-invalid.mlir b/mlir/test/Dialect/OpenACC/acc-loop-tiling-invalid.mlir
index 6ef1884345f82..53b0467976658 100644
--- a/mlir/test/Dialect/OpenACC/acc-loop-tiling-invalid.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-loop-tiling-invalid.mlir
@@ -13,3 +13,26 @@ func.func @tile_wider_than_iv(%arg0: memref<100xf32>) {
} attributes {independent = [#acc.device_type<none>]}
return
}
+
+// -----
+
+// A tile clause combined with a collapse clause on the same loop is not
+// supported: the two clauses describe conflicting loop-association counts. The
+// pass must diagnose it rather than silently drop the collapse clause.
+
+func.func @tile_with_collapse(%arg0: memref<100x50xf32>) {
+ %c0 = arith.constant 0 : index
+ %c100 = arith.constant 100 : index
+ %c50 = arith.constant 50 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ // expected-error @below {{not yet implemented: a tile clause combined with a collapse clause on the same loop}}
+ acc.loop tile({%c4 : index, %c8 : index}) control(%i : index, %j : index) = (%c0, %c0 : index, index) to (%c100, %c50 : index, index) step (%c1, %c1 : index, index) {
+ %val = arith.index_castui %i : index to i32
+ %fval = arith.sitofp %val : i32 to f32
+ memref.store %fval, %arg0[%i, %j] : memref<100x50xf32>
+ acc.yield
+ } attributes {collapse = [2], collapseDeviceType = [#acc.device_type<none>], independent = [#acc.device_type<none>]}
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/acc-loop-tiling.mlir b/mlir/test/Dialect/OpenACC/acc-loop-tiling.mlir
index a1f80ad9defdf..724fba48b7595 100644
--- a/mlir/test/Dialect/OpenACC/acc-loop-tiling.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-loop-tiling.mlir
@@ -32,7 +32,8 @@ func.func @single_loop_tile(%arg0: memref<10xf32>) {
}
// Test 2-level nested loop tiling with tile(4, 8)
-// Creates: tile_loop_1 -> tile_loop_2 -> element_loop_1 -> element_loop_2
+// Produces one multi-IV tile-group loop (gang) wrapping one multi-IV
+// element-group loop (vector), each carrying both tiled induction variables.
// CHECK-LABEL: func.func @nested_loop_tile
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
@@ -41,22 +42,16 @@ func.func @single_loop_tile(%arg0: memref<10xf32>) {
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
-// Outer tile loop with gang
-// CHECK: acc.loop gang control(%[[I:.*]] : index) = (%[[C0]] : index) to (%[[C100]] : index) step (%[[C4]] : index) {
-// Inner tile loop
-// CHECK: acc.loop control(%[[J:.*]] : index) = (%[[C0]] : index) to (%[[C50]] : index) step (%[[C8]] : index) {
-// Outer element loop with vector
-// CHECK: acc.loop vector control({{.*}} : index) = (%[[I]] : index) to ({{.*}} : index) step (%[[C1]] : index) {
-// Inner element loop
-// CHECK: acc.loop control({{.*}} : index) = (%[[J]] : index) to ({{.*}} : index) step (%[[C1]] : index) {
-// CHECK: acc.yield
-// CHECK: }
-// CHECK: acc.yield
-// CHECK: }
+// Tile-group loop with gang; steps scaled by the tile sizes.
+// CHECK: acc.loop gang control(%[[I:.*]] : index, %[[J:.*]] : index) = (%[[C0]], %[[C0]] : index, index) to (%[[C100]], %[[C50]] : index, index) step (%[[C4]], %[[C8]] : index, index) {
+// CHECK: %[[MUB0:.*]] = arith.minsi
+// CHECK: %[[MUB1:.*]] = arith.minsi
+// Element-group loop with vector.
+// CHECK: acc.loop vector control(%{{.*}} : index, %{{.*}} : index) = (%[[I]], %[[J]] : index, index) to (%[[MUB0]], %[[MUB1]] : index, index) step (%[[C1]], %[[C1]] : index, index) {
// CHECK: acc.yield
-// CHECK: }
+// CHECK: } attributes {independent = [#acc.device_type<none>]}
// CHECK: acc.yield
-// CHECK: }
+// CHECK: } attributes {independent = [#acc.device_type<none>]}
func.func @nested_loop_tile(%arg0: memref<100x50xf32>) {
%c0 = arith.constant 0 : index
%c100 = arith.constant 100 : index
@@ -76,29 +71,19 @@ func.func @nested_loop_tile(%arg0: memref<100x50xf32>) {
// Regression test: a loop with GANG(STATIC: N) combined with a multi-dim
// TILE clause used to crash with an assertion failure inside
-// LoopOp::getGangValue() because uncollapseLoops() (needed here since the
-// tile count exceeds the implicit collapse count of 1) left one of the
-// generated inner loops with a leftover gang operand but no corresponding
-// gang device-type attribute. Check that the pass completes and that gang
-// is only preserved on the outermost tile loop.
+// LoopOp::getGangValue() when a leftover gang operand was copied onto a
+// generated loop without a corresponding gang device-type attribute. Check
+// that the pass completes and that gang (with its static operand) is only
+// preserved on the tile-group loop, never on the element group.
// CHECK-LABEL: func.func @gang_static_with_multi_dim_tile
-// CHECK: acc.loop gang({static=%{{.*}} : i32}) control(%[[I:.*]] : index) = ({{.*}}) to ({{.*}}) step ({{.*}}) {
-// CHECK-NOT: gang
-// CHECK: acc.loop control(%[[J:.*]] : index) = ({{.*}}) to ({{.*}}) step ({{.*}}) {
+// CHECK: acc.loop gang({static=%{{.*}} : i32}) control(%[[I:.*]] : index, %[[J:.*]] : index) = ({{.*}}) to ({{.*}}) step ({{.*}}) {
+// CHECK: acc.loop control(%{{.*}} : index, %{{.*}} : index) = (%[[I]], %[[J]] : index, index) to ({{.*}}) step ({{.*}}) {
// CHECK-NOT: gang
-// CHECK: acc.loop control({{.*}} : index) = (%[[I]] : index) to ({{.*}}) step ({{.*}}) {
-// CHECK-NOT: gang
-// CHECK: acc.loop control({{.*}} : index) = (%[[J]] : index) to ({{.*}}) step ({{.*}}) {
-// CHECK-NOT: gang
-// CHECK: acc.yield
-// CHECK: }
-// CHECK: acc.yield
-// CHECK: }
// CHECK: acc.yield
-// CHECK: }
+// CHECK: } attributes {independent = [#acc.device_type<none>]}
// CHECK: acc.yield
-// CHECK: }
+// CHECK: } attributes {independent = [#acc.device_type<none>]}
func.func @gang_static_with_multi_dim_tile(%arg0: memref<100x50xf32>) {
%c0 = arith.constant 0 : index
%c100 = arith.constant 100 : index
@@ -118,36 +103,22 @@ func.func @gang_static_with_multi_dim_tile(%arg0: memref<100x50xf32>) {
}
// Regression test: a loop with WORKER(N) combined with VECTOR and a
-// multi-dim TILE clause used to produce invalid IR because createInnerLoop()
-// (via uncollapseLoops()) and removeWorkerVectorFromLoop() removed the
-// worker attributes from generated inner/tile loops without also clearing
-// the worker operand value copied onto them. This left loops with a
-// non-empty worker operand list but no worker device-type attribute,
-// which the verifier rejects ('worker operands count must match worker
-// device_type count'). Check that the pass produces valid IR, with worker
-// only preserved on the outermost tile loop and vector only on the
-// outermost element loop.
+// multi-dim TILE clause used to produce invalid IR because a worker attribute
+// was removed from a generated loop without also clearing the worker operand
+// value copied onto it, leaving a non-empty worker operand list with no worker
+// device-type attribute (rejected by the verifier). Check that the pass
+// produces valid IR, with worker only on the tile-group loop and vector only
+// on the element-group loop.
// CHECK-LABEL: func.func @worker_num_with_multi_dim_tile
-// CHECK: acc.loop worker(%{{.*}} : i32) control(%[[I:.*]] : index) = ({{.*}}) to ({{.*}}) step ({{.*}}) {
-// CHECK-NOT: worker
+// CHECK: acc.loop worker(%{{.*}} : i32) control(%[[I:.*]] : index, %[[J:.*]] : index) = ({{.*}}) to ({{.*}}) step ({{.*}}) {
// CHECK-NOT: vector
-// CHECK: acc.loop control(%[[J:.*]] : index) = ({{.*}}) to ({{.*}}) step ({{.*}}) {
+// CHECK: acc.loop vector control(%{{.*}} : index, %{{.*}} : index) = (%[[I]], %[[J]] : index, index) to ({{.*}}) step ({{.*}}) {
// CHECK-NOT: worker
-// CHECK-NOT: vector
-// CHECK: acc.loop vector control({{.*}} : index) = (%[[I]] : index) to ({{.*}}) step ({{.*}}) {
-// CHECK-NOT: worker
-// CHECK: acc.loop control({{.*}} : index) = (%[[J]] : index) to ({{.*}}) step ({{.*}}) {
-// CHECK-NOT: worker
-// CHECK-NOT: vector
-// CHECK: acc.yield
-// CHECK: }
-// CHECK: acc.yield
-// CHECK: }
// CHECK: acc.yield
-// CHECK: }
+// CHECK: } attributes {independent = [#acc.device_type<none>]}
// CHECK: acc.yield
-// CHECK: }
+// CHECK: } attributes {independent = [#acc.device_type<none>]}
func.func @worker_num_with_multi_dim_tile(%arg0: memref<100x50xf32>) {
%c0 = arith.constant 0 : index
%c100 = arith.constant 100 : index
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsTilingTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsTilingTest.cpp
index 95bc1eab7d3fe..302e60cca168d 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsTilingTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsTilingTest.cpp
@@ -98,41 +98,36 @@ TEST_F(OpenACCUtilsTilingTest, tileACCLoopsSingleLoop) {
Value tileSize =
arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(4));
- // Create the loop
+ // Create the loop (single IV)
acc::LoopOp loopOp = createLoopOp(b, {lb}, {ub}, {step});
// Tile the loop using IRRewriter
IRRewriter rewriter(&context);
rewriter.setInsertionPoint(loopOp);
- SmallVector<acc::LoopOp> loopsToTile = {loopOp};
SmallVector<Value> tileSizes = {tileSize};
- acc::LoopOp tiledLoop =
- tileACCLoops(loopsToTile, tileSizes, /*defaultTileSize=*/128, rewriter);
+ acc::LoopOp tileGroup =
+ tileACCLoops(loopOp, tileSizes, /*defaultTileSize=*/128, rewriter);
- // Verify the tiled loop was created
- EXPECT_TRUE(tiledLoop != nullptr);
- EXPECT_FALSE(tiledLoop.getBody().empty());
+ // Verify the tile-group loop was created
+ EXPECT_TRUE(tileGroup != nullptr);
+ EXPECT_FALSE(tileGroup.getBody().empty());
- // After tiling a single loop with tile(4), we should have:
- // - 1 tile loop (the outer loop)
- // - 1 element loop nested inside
- // Total: 1 nested loop inside the tile loop
- EXPECT_EQ(countNestedLoops(tiledLoop), 1u);
+ // A single-IV tile(4) produces one tile-group loop wrapping one element-group
+ // loop: exactly one nested loop.
+ EXPECT_EQ(countNestedLoops(tileGroup), 1u);
- // The tile loop (outer) should have 1 IV
- EXPECT_EQ(tiledLoop.getBody().getNumArguments(), 1u);
+ // The tile-group loop carries its single IV.
+ EXPECT_EQ(tileGroup.getBody().getNumArguments(), 1u);
- // Collect nested loops and verify
- auto nestedLoops = collectNestedLoops(tiledLoop);
- EXPECT_EQ(nestedLoops.size(), 1u);
- // The element loop should have 1 IV
- if (!nestedLoops.empty())
- EXPECT_EQ(nestedLoops[0].getBody().getNumArguments(), 1u);
+ auto nestedLoops = collectNestedLoops(tileGroup);
+ ASSERT_EQ(nestedLoops.size(), 1u);
+ // The element-group loop should carry the single IV as well.
+ EXPECT_EQ(nestedLoops[0].getBody().getNumArguments(), 1u);
}
-TEST_F(OpenACCUtilsTilingTest, tileACCLoopsNestedLoops) {
+TEST_F(OpenACCUtilsTilingTest, tileACCLoopsFusedTwoDim) {
// Create a module to hold the function
OwningOpRef<ModuleOp> module = ModuleOp::create(loc);
Block *moduleBlock = module->getBody();
@@ -148,15 +143,13 @@ TEST_F(OpenACCUtilsTilingTest, tileACCLoopsNestedLoops) {
b.setInsertionPointToStart(funcBlock);
- // Create loop bounds for outer loop
+ // Create bounds for a single fused loop carrying two IVs.
Value lb1 =
arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(0));
Value ub1 =
arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(100));
Value step1 =
arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(1));
-
- // Create loop bounds for inner loop
Value lb2 =
arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(0));
Value ub2 =
@@ -170,179 +163,31 @@ TEST_F(OpenACCUtilsTilingTest, tileACCLoopsNestedLoops) {
Value tileSize2 =
arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(8));
- // Create outer loop
- acc::LoopOp outerLoop = createLoopOp(b, {lb1}, {ub1}, {step1});
-
- // Create inner loop inside outer loop
- b.setInsertionPoint(outerLoop.getBody().getTerminator());
- acc::LoopOp innerLoop = createLoopOp(b, {lb2}, {ub2}, {step2});
-
- // Tile the loops
- IRRewriter rewriter(&context);
- rewriter.setInsertionPoint(outerLoop);
-
- SmallVector<acc::LoopOp> loopsToTile = {outerLoop, innerLoop};
- SmallVector<Value> tileSizes = {tileSize1, tileSize2};
-
- acc::LoopOp tiledLoop =
- tileACCLoops(loopsToTile, tileSizes, /*defaultTileSize=*/128, rewriter);
-
- // Verify the tiled loop nest was created
- EXPECT_TRUE(tiledLoop != nullptr);
- EXPECT_FALSE(tiledLoop.getBody().empty());
-
- // After tiling a 2-level nested loop with tile(4,8), we should have:
- // tile_loop_1 -> tile_loop_2 -> element_loop_1 -> element_loop_2
- // Total: 3 nested loops inside the outermost tile loop
- unsigned nestedCount = countNestedLoops(tiledLoop);
- EXPECT_EQ(nestedCount, 3u);
-
- // The outermost tile loop should have 1 IV
- EXPECT_EQ(tiledLoop.getBody().getNumArguments(), 1u);
-
- // Collect all nested loops and verify each has 1 IV
- auto nestedLoops = collectNestedLoops(tiledLoop);
- EXPECT_EQ(nestedLoops.size(), 3u);
- for (auto loop : nestedLoops)
- EXPECT_EQ(loop.getBody().getNumArguments(), 1u);
-}
-
-//===----------------------------------------------------------------------===//
-// uncollapseLoops Tests
-//===----------------------------------------------------------------------===//
-
-TEST_F(OpenACCUtilsTilingTest, uncollapseLoopsBasic) {
- // Create a module to hold the function
- OwningOpRef<ModuleOp> module = ModuleOp::create(loc);
- Block *moduleBlock = module->getBody();
-
- OpBuilder::InsertionGuard guard(b);
- b.setInsertionPointToStart(moduleBlock);
-
- // Create a function
- auto funcType = b.getFunctionType({}, {});
- OwningOpRef<func::FuncOp> funcOp =
- func::FuncOp::create(b, loc, "test_func", funcType);
- Block *funcBlock = funcOp->addEntryBlock();
-
- b.setInsertionPointToStart(funcBlock);
-
- // Create loop bounds for a collapsed 2-level loop
- Value lb1 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(0));
- Value ub1 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(10));
- Value step1 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(1));
- Value lb2 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(0));
- Value ub2 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(20));
- Value step2 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(1));
-
- // Create a collapsed loop with 2 IVs
- acc::LoopOp collapsedLoop =
+ // Create a single fused loop with two IVs.
+ acc::LoopOp fusedLoop =
createLoopOp(b, {lb1, lb2}, {ub1, ub2}, {step1, step2});
- // Set the collapse attribute
- collapsedLoop.setCollapseForDeviceTypes(&context, {acc::DeviceType::None},
- llvm::APInt(64, 1));
-
- // Uncollapse the loop: tileCount=2, collapseCount=1
IRRewriter rewriter(&context);
- rewriter.setInsertionPoint(collapsedLoop);
-
- SmallVector<acc::LoopOp> uncollapsedLoops = uncollapseLoops(
- collapsedLoop, /*tileCount=*/2, /*collapseCount=*/1, rewriter);
-
- // Should produce 2 loops (one outer with collapse=1, one inner)
- EXPECT_EQ(uncollapsedLoops.size(), 2u);
-
- if (uncollapsedLoops.size() >= 2) {
- // Verify the outer loop has 1 IV (collapseCount=1)
- acc::LoopOp outerLoop = uncollapsedLoops[0];
- EXPECT_EQ(outerLoop.getBody().getNumArguments(), 1u);
- EXPECT_EQ(outerLoop.getLowerbound().size(), 1u);
- EXPECT_EQ(outerLoop.getUpperbound().size(), 1u);
- EXPECT_EQ(outerLoop.getStep().size(), 1u);
-
- // Verify the inner loop has 1 IV
- acc::LoopOp innerLoop = uncollapsedLoops[1];
- EXPECT_EQ(innerLoop.getBody().getNumArguments(), 1u);
- EXPECT_EQ(innerLoop.getLowerbound().size(), 1u);
- EXPECT_EQ(innerLoop.getUpperbound().size(), 1u);
- EXPECT_EQ(innerLoop.getStep().size(), 1u);
-
- // Verify nesting: inner loop should be inside outer loop
- unsigned nestedCount = countNestedLoops(outerLoop);
- EXPECT_EQ(nestedCount, 1u);
- }
-}
-
-TEST_F(OpenACCUtilsTilingTest, uncollapseLoopsThreeLevels) {
- // Test uncollapsing with 3 levels: collapse(2) with tile(3)
- OwningOpRef<ModuleOp> module = ModuleOp::create(loc);
- Block *moduleBlock = module->getBody();
+ rewriter.setInsertionPoint(fusedLoop);
- OpBuilder::InsertionGuard guard(b);
- b.setInsertionPointToStart(moduleBlock);
-
- auto funcType = b.getFunctionType({}, {});
- OwningOpRef<func::FuncOp> funcOp =
- func::FuncOp::create(b, loc, "test_func", funcType);
- Block *funcBlock = funcOp->addEntryBlock();
-
- b.setInsertionPointToStart(funcBlock);
-
- // Create 3 sets of bounds
- Value lb1 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(0));
- Value ub1 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(10));
- Value step1 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(1));
- Value lb2 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(0));
- Value ub2 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(20));
- Value step2 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(1));
- Value lb3 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(0));
- Value ub3 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(30));
- Value step3 =
- arith::ConstantOp::create(b, loc, b.getIndexType(), b.getIndexAttr(1));
-
- // Create a collapsed loop with 3 IVs
- acc::LoopOp collapsedLoop =
- createLoopOp(b, {lb1, lb2, lb3}, {ub1, ub2, ub3}, {step1, step2, step3});
-
- // Set collapse(2)
- collapsedLoop.setCollapseForDeviceTypes(&context, {acc::DeviceType::None},
- llvm::APInt(64, 2));
+ SmallVector<Value> tileSizes = {tileSize1, tileSize2};
- // Uncollapse: tileCount=3, collapseCount=2
- // This should create: outer loop with 2 IVs, then 1 inner loop
- IRRewriter rewriter(&context);
- rewriter.setInsertionPoint(collapsedLoop);
+ acc::LoopOp tileGroup =
+ tileACCLoops(fusedLoop, tileSizes, /*defaultTileSize=*/128, rewriter);
- SmallVector<acc::LoopOp> uncollapsedLoops = uncollapseLoops(
- collapsedLoop, /*tileCount=*/3, /*collapseCount=*/2, rewriter);
+ // Verify the tile-group loop nest was created.
+ EXPECT_TRUE(tileGroup != nullptr);
+ EXPECT_FALSE(tileGroup.getBody().empty());
- // Should produce 2 loops
- EXPECT_EQ(uncollapsedLoops.size(), 2u);
+ // tile(4,8) on a fused 2-IV loop produces exactly two multi-IV loops: a
+ // tile-group loop wrapping a single element-group loop (one nested loop),
+ // rather than a 4-deep nest of single-IV loops.
+ EXPECT_EQ(countNestedLoops(tileGroup), 1u);
- if (uncollapsedLoops.size() >= 2) {
- // Outer loop should have 2 IVs (from collapse=2)
- acc::LoopOp outerLoop = uncollapsedLoops[0];
- EXPECT_EQ(outerLoop.getBody().getNumArguments(), 2u);
- EXPECT_EQ(outerLoop.getLowerbound().size(), 2u);
+ // Both groups carry all (two) IVs.
+ EXPECT_EQ(tileGroup.getBody().getNumArguments(), 2u);
- // Inner loop should have 1 IV (the 3rd dimension)
- acc::LoopOp innerLoop = uncollapsedLoops[1];
- EXPECT_EQ(innerLoop.getBody().getNumArguments(), 1u);
- EXPECT_EQ(innerLoop.getLowerbound().size(), 1u);
- }
+ auto nestedLoops = collectNestedLoops(tileGroup);
+ ASSERT_EQ(nestedLoops.size(), 1u);
+ EXPECT_EQ(nestedLoops[0].getBody().getNumArguments(), 2u);
}
More information about the Mlir-commits
mailing list