575d6892bc
This pass implements the OpenACC loop tiling transformation for acc.loop
operations that have the tile clause (OpenACC 3.4 spec, section 2.9.8).
The tile clause specifies that the iterations of the associated loops
should be divided into tiles (rectangular blocks). The pass transforms a
single or nested acc.loop with tile clauses into a structure of "tile
loops" (iterating over tiles) containing "element loops" (iterating
within tiles).
For example, tiling a 2-level nested loop with tile(T1, T2):
```
// Before tiling:
acc.loop tile(T1, T2) control(%i, %j) = ...
// After tiling:
acc.loop control(%i) step (s1*T1) { // tile loop 1
acc.loop control(%j) step (s2*T2) { // tile loop 2
acc.loop control(%ii) = (%i) to (min(ub1, %i+s1*T1)) {
acc.loop control(%jj) = (%j) to (min(ub2, %j+s2*T2)) {
// loop body using %ii, %jj
}
}
}
}
```
Key features:
- Handles constant tile sizes and wildcard tile sizes ('*') which use a
configurable default tile size
- Properly handles collapsed loops with tile counts exceeding collapse
count by uncollapsing loops before tiling
- Distributes gang/worker/vector attributes appropriately: gang -> tile
loops, vector -> element loops
- Validates that tile size types are not wider than loop IV types
- Emits optimization remarks for tiling decisions
Three test files are added:
- acc-loop-tiling.mlir: Tests single and nested loop tiling with
constant tile sizes, unknown tile sizes (*), and loops with collapse
attributes
- acc-loop-tiling-invalid.mlir: Tests error diagnostic when tile size
type is wider than the loop IV type
- acc-loop-tiling-remarks.mlir: Tests optimization remarks emitted for
tiling decisions including default tile size selection
Co-authored-by: Vijay Kandiah <vkandiah@nvidia.com>
312 lines
14 KiB
C++
312 lines
14 KiB
C++
//===- OpenACCUtilsTiling.cpp - OpenACC Loop Tiling Utilities -------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file contains utility functions for tiling OpenACC loops.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/OpenACC/OpenACCUtilsTiling.h"
|
|
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/Arith/Utils/Utils.h"
|
|
#include "mlir/Dialect/OpenACC/OpenACC.h"
|
|
#include "mlir/Dialect/Utils/StaticValueUtils.h"
|
|
#include "mlir/Transforms/RegionUtils.h"
|
|
|
|
// Resolve unknown tile sizes (represented as -1 for tile(*)) to the default.
|
|
// Returns a value with the same type as targetType.
|
|
static mlir::Value resolveAndCastTileSize(mlir::Value tileSize,
|
|
int32_t defaultTileSize,
|
|
mlir::Type targetType,
|
|
mlir::RewriterBase &rewriter,
|
|
mlir::Location loc) {
|
|
auto constVal = mlir::getConstantIntValue(tileSize);
|
|
if (constVal && *constVal < 0) {
|
|
// Create constant with the target type directly
|
|
return mlir::arith::ConstantOp::create(
|
|
rewriter, loc, targetType,
|
|
rewriter.getIntegerAttr(targetType, defaultTileSize));
|
|
}
|
|
return mlir::getValueOrCreateCastToIndexLike(rewriter, loc, targetType,
|
|
tileSize);
|
|
}
|
|
|
|
// Remove vector/worker attributes from loop
|
|
static void removeWorkerVectorFromLoop(mlir::acc::LoopOp loop) {
|
|
if (loop.hasVector() || loop.getVectorValue()) {
|
|
loop.removeVectorAttr();
|
|
loop.removeVectorOperandsDeviceTypeAttr();
|
|
} else if (loop.hasWorker() || loop.getWorkerValue()) {
|
|
loop.removeWorkerAttr();
|
|
loop.removeWorkerNumOperandsDeviceTypeAttr();
|
|
}
|
|
}
|
|
|
|
// 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) {
|
|
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(),
|
|
origLoop.getGangOperandsArgTypeAttr(),
|
|
origLoop.getGangOperandsSegmentsAttr(),
|
|
origLoop.getGangOperandsDeviceTypeAttr(), origLoop.getWorkerNumOperands(),
|
|
origLoop.getWorkerNumOperandsDeviceTypeAttr(),
|
|
origLoop.getVectorOperands(), origLoop.getVectorOperandsDeviceTypeAttr(),
|
|
origLoop.getSeqAttr(), origLoop.getIndependentAttr(),
|
|
origLoop.getAuto_Attr(), origLoop.getGangAttr(), origLoop.getWorkerAttr(),
|
|
origLoop.getVectorAttr(), mlir::ValueRange{}, mlir::DenseI32ArrayAttr{},
|
|
mlir::ArrayAttr{}, origLoop.getCacheOperands(),
|
|
origLoop.getPrivateOperands(), origLoop.getFirstprivateOperands(),
|
|
origLoop.getReductionOperands(), combinedAttr);
|
|
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();
|
|
}
|
|
if (inputLoop.hasVector() || inputLoop.getVectorValue()) {
|
|
elementLoop.removeWorkerAttr();
|
|
elementLoop.removeWorkerNumOperandsDeviceTypeAttr();
|
|
}
|
|
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,
|
|
llvm::ArrayRef<mlir::Value> newIVs,
|
|
llvm::ArrayRef<mlir::Value> origIVs,
|
|
size_t nOps, mlir::RewriterBase &rewriter) {
|
|
// Move ops from source to target loop [begin, begin + nOps - 1)
|
|
mlir::Block::iterator begin = sourceLoop.getBody().begin();
|
|
targetLoop.getBody().getOperations().splice(
|
|
targetLoop.getBody().getOperations().begin(),
|
|
sourceLoop.getBody().getOperations(), begin, std::next(begin, nOps - 1));
|
|
|
|
// Replace uses of origIV with newIV
|
|
for (auto [i, newIV] : llvm::enumerate(newIVs))
|
|
mlir::replaceAllUsesInRegionWith(origIVs[i], newIV, targetLoop.getRegion());
|
|
}
|
|
|
|
mlir::acc::LoopOp
|
|
mlir::acc::tileACCLoops(llvm::SmallVector<mlir::acc::LoopOp> &tileLoops,
|
|
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
|
|
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);
|
|
}
|
|
}
|
|
|
|
// 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);
|
|
}
|
|
}
|
|
|
|
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]);
|
|
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));
|
|
newUB = mlir::arith::SubIOp::create(rewriter, loc, stepped, c1);
|
|
}
|
|
newUBs.push_back(
|
|
mlir::arith::MinSIOp::create(rewriter, loc, origUBs[i], newUB));
|
|
}
|
|
|
|
// 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,
|
|
rewriter);
|
|
|
|
return newLoops;
|
|
}
|