25fcc8702d
Fixes a problem encountered with enabling coalesceLoops when bounds were constructed inside expanded loops. Additionally, ensures that all loop utilities use rewriter instead of their own builders for proper tracking.
326 lines
13 KiB
C++
326 lines
13 KiB
C++
//===- OpenACCUtilsLoop.cpp - OpenACC Loop 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 converting OpenACC loops to SCF.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/OpenACC/OpenACCUtilsLoop.h"
|
|
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/Arith/Utils/Utils.h"
|
|
#include "mlir/Dialect/OpenACC/OpenACC.h"
|
|
#include "mlir/Dialect/SCF/IR/SCF.h"
|
|
#include "mlir/Dialect/SCF/Utils/Utils.h"
|
|
#include "mlir/IR/IRMapping.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
|
|
/// Calculate trip count for a loop: (ub - lb + step) / step
|
|
/// If inclusiveUpperbound is false, subtracts 1 from ub first.
|
|
static Value calculateTripCount(OpBuilder &b, Location loc, Value lb, Value ub,
|
|
Value step, bool inclusiveUpperbound) {
|
|
Type type = b.getIndexType();
|
|
|
|
// Convert original loop arguments to index type
|
|
lb = getValueOrCreateCastToIndexLike(b, loc, type, lb);
|
|
ub = getValueOrCreateCastToIndexLike(b, loc, type, ub);
|
|
step = getValueOrCreateCastToIndexLike(b, loc, type, step);
|
|
|
|
if (!inclusiveUpperbound) {
|
|
Value one = arith::ConstantIndexOp::create(b, loc, 1);
|
|
ub = b.createOrFold<arith::SubIOp>(loc, ub, one,
|
|
arith::IntegerOverflowFlags::nsw);
|
|
}
|
|
|
|
Value sub = b.createOrFold<arith::SubIOp>(loc, ub, lb,
|
|
arith::IntegerOverflowFlags::nsw);
|
|
Value add = b.createOrFold<arith::AddIOp>(loc, sub, step,
|
|
arith::IntegerOverflowFlags::nsw);
|
|
return b.createOrFold<arith::DivSIOp>(loc, add, step);
|
|
}
|
|
|
|
/// Get exclusive upper bound from acc.loop (add 1 if inclusive).
|
|
/// The result is always in index type.
|
|
static Value getExclusiveUpperBoundAsIndex(acc::LoopOp loopOp, size_t ivPos,
|
|
OpBuilder &b) {
|
|
bool isInclusive = false;
|
|
if (loopOp.getInclusiveUpperbound().has_value())
|
|
isInclusive = loopOp.getInclusiveUpperboundAttr().asArrayRef()[ivPos];
|
|
|
|
Value origUB = loopOp.getUpperbound()[ivPos];
|
|
Location loc = origUB.getLoc();
|
|
Type indexType = b.getIndexType();
|
|
|
|
// Cast to index first, then add if inclusive
|
|
Value ub = getValueOrCreateCastToIndexLike(b, loc, indexType, origUB);
|
|
if (isInclusive) {
|
|
Value one = arith::ConstantIndexOp::create(b, loc, 1);
|
|
ub = b.createOrFold<arith::AddIOp>(loc, ub, one,
|
|
arith::IntegerOverflowFlags::nsw);
|
|
}
|
|
return ub;
|
|
}
|
|
|
|
/// Handle differing types between SCF (index) and ACC loops.
|
|
/// Creates casts from the new SCF IVs to the original ACC IV types and updates
|
|
/// the mapping. The newIVs should correspond 1:1 with the ACC loop's IVs.
|
|
static void mapACCLoopIVsToSCFIVs(acc::LoopOp accLoop, ValueRange newIVs,
|
|
OpBuilder &b, IRMapping &mapping) {
|
|
for (auto [origIV, newIV] :
|
|
llvm::zip(accLoop.getBody().getArguments(), newIVs)) {
|
|
Value replacementIV = getValueOrCreateCastToIndexLike(
|
|
b, accLoop->getLoc(), origIV.getType(), newIV);
|
|
mapping.map(origIV, replacementIV);
|
|
}
|
|
}
|
|
|
|
/// Normalize IV uses after converting to normalized loop form.
|
|
/// For normalized loops (lb=0, step=1), we need to denormalize the IV:
|
|
/// original_iv = new_iv * orig_step + orig_lb
|
|
static void normalizeIVUses(OpBuilder &b, Location loc, Value iv, Value origLB,
|
|
Value origStep) {
|
|
Type indexType = b.getIndexType();
|
|
Value lb = getValueOrCreateCastToIndexLike(b, loc, indexType, origLB);
|
|
Value step = getValueOrCreateCastToIndexLike(b, loc, indexType, origStep);
|
|
|
|
// new_iv * step + lb
|
|
Value scaled =
|
|
arith::MulIOp::create(b, loc, iv, step, arith::IntegerOverflowFlags::nsw);
|
|
Value denormalized = arith::AddIOp::create(b, loc, scaled, lb,
|
|
arith::IntegerOverflowFlags::nsw);
|
|
|
|
// Replace uses of iv with denormalized value, except for the ops that
|
|
// compute the denormalized value itself (muli and addi)
|
|
llvm::SmallPtrSet<Operation *, 2> exceptions;
|
|
exceptions.insert(scaled.getDefiningOp());
|
|
exceptions.insert(denormalized.getDefiningOp());
|
|
iv.replaceAllUsesExcept(denormalized, exceptions);
|
|
}
|
|
|
|
/// Clone an ACC region into a destination block, handling the ACC terminators.
|
|
/// Returns the insertion point after the cloned operations.
|
|
static Block::iterator cloneACCRegionInto(Region *src, Block *dest,
|
|
Block::iterator insertionPoint,
|
|
IRMapping &mapping,
|
|
RewriterBase &rewriter) {
|
|
assert(src->hasOneBlock() && "expected single-block region");
|
|
|
|
Region *insertRegion = dest->getParent();
|
|
Block *postInsertBlock = rewriter.splitBlock(dest, insertionPoint);
|
|
rewriter.cloneRegionBefore(*src, *insertRegion,
|
|
postInsertBlock->getIterator(), mapping);
|
|
|
|
auto lastNewBlock = std::prev(postInsertBlock->getIterator());
|
|
|
|
Block::iterator newInsertionPoint;
|
|
Operation *terminator = lastNewBlock->getTerminator();
|
|
|
|
if (auto yieldOp = dyn_cast<acc::YieldOp>(terminator)) {
|
|
newInsertionPoint = std::prev(yieldOp->getIterator());
|
|
rewriter.eraseOp(yieldOp);
|
|
} else if (auto terminatorOp = dyn_cast<acc::TerminatorOp>(terminator)) {
|
|
newInsertionPoint = std::prev(terminatorOp->getIterator());
|
|
rewriter.eraseOp(terminatorOp);
|
|
} else {
|
|
llvm_unreachable("unexpected terminator in ACC region");
|
|
}
|
|
|
|
// Merge last block with the postInsertBlock
|
|
rewriter.mergeBlocks(postInsertBlock, &*lastNewBlock);
|
|
|
|
// Merge first block with original dest block
|
|
Block *firstNewBlock = &*std::next(dest->getIterator());
|
|
rewriter.mergeBlocks(firstNewBlock, dest);
|
|
|
|
return newInsertionPoint;
|
|
}
|
|
|
|
/// Wrap a multi-block region with scf.execute_region.
|
|
static scf::ExecuteRegionOp
|
|
wrapMultiBlockRegionWithSCFExecuteRegion(Region ®ion, IRMapping &mapping,
|
|
Location loc, RewriterBase &rewriter) {
|
|
auto exeRegionOp = scf::ExecuteRegionOp::create(rewriter, loc, TypeRange{});
|
|
|
|
rewriter.cloneRegionBefore(region, exeRegionOp.getRegion(),
|
|
exeRegionOp.getRegion().end(), mapping);
|
|
|
|
// Find and replace the ACC terminator with scf.yield
|
|
Operation *terminator = exeRegionOp.getRegion().back().getTerminator();
|
|
if (auto yieldOp = dyn_cast<acc::YieldOp>(terminator)) {
|
|
if (yieldOp.getNumOperands() > 0) {
|
|
region.getParentOp()->emitError(
|
|
"acc.loop with results not yet supported");
|
|
return nullptr;
|
|
}
|
|
} else if (!isa<acc::TerminatorOp>(terminator)) {
|
|
llvm_unreachable("unexpected terminator in ACC region");
|
|
}
|
|
|
|
rewriter.eraseOp(terminator);
|
|
rewriter.setInsertionPointToEnd(&exeRegionOp.getRegion().back());
|
|
scf::YieldOp::create(rewriter, loc);
|
|
return exeRegionOp;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
namespace mlir {
|
|
namespace acc {
|
|
|
|
scf::ForOp convertACCLoopToSCFFor(LoopOp loopOp, RewriterBase &rewriter,
|
|
bool enableCollapse) {
|
|
assert(!loopOp.getUnstructured() &&
|
|
"use convertUnstructuredACCLoopToSCFExecuteRegion for unstructured "
|
|
"loops");
|
|
|
|
Location loc = loopOp->getLoc();
|
|
Type indexType = rewriter.getIndexType();
|
|
|
|
// Create nested scf.for loops and build IR mapping for IVs
|
|
IRMapping mapping;
|
|
SmallVector<scf::ForOp> forOps;
|
|
|
|
// Save the original insertion point
|
|
OpBuilder::InsertionGuard guard(rewriter);
|
|
rewriter.setInsertionPoint(loopOp);
|
|
|
|
// First, compute ALL loop bounds at the current insertion point (before
|
|
// any ForOp). This ensures all bounds are defined in the outer scope,
|
|
// which is required for coalesceLoops to work correctly.
|
|
SmallVector<Value> lowerBounds, upperBounds, steps;
|
|
for (BlockArgument iv : loopOp.getBody().getArguments()) {
|
|
size_t idx = iv.getArgNumber();
|
|
Value newLowerBound = getValueOrCreateCastToIndexLike(
|
|
rewriter, loc, indexType, loopOp.getLowerbound()[idx]);
|
|
Value newUpperBound = getExclusiveUpperBoundAsIndex(loopOp, idx, rewriter);
|
|
Value newStep = getValueOrCreateCastToIndexLike(rewriter, loc, indexType,
|
|
loopOp.getStep()[idx]);
|
|
lowerBounds.push_back(newLowerBound);
|
|
upperBounds.push_back(newUpperBound);
|
|
steps.push_back(newStep);
|
|
}
|
|
|
|
// Now create the nested ForOps using the pre-computed bounds
|
|
for (BlockArgument iv : loopOp.getBody().getArguments()) {
|
|
size_t idx = iv.getArgNumber();
|
|
|
|
// For nested loops, insert inside the previous loop's body
|
|
if (idx > 0)
|
|
rewriter.setInsertionPointToStart(forOps.back().getBody());
|
|
|
|
scf::ForOp forOp = scf::ForOp::create(rewriter, loc, lowerBounds[idx],
|
|
upperBounds[idx], steps[idx]);
|
|
forOps.push_back(forOp);
|
|
mapping.map(iv, forOp.getInductionVar());
|
|
}
|
|
|
|
// Set insertion point inside the innermost loop for IV casts and body cloning
|
|
rewriter.setInsertionPointToStart(forOps.back().getBody());
|
|
|
|
// Handle IV type conversion (index -> original type)
|
|
SmallVector<Value> scfIVs;
|
|
for (scf::ForOp forOp : forOps)
|
|
scfIVs.push_back(forOp.getInductionVar());
|
|
mapACCLoopIVsToSCFIVs(loopOp, scfIVs, rewriter, mapping);
|
|
|
|
// Clone the loop body into the innermost scf.for
|
|
cloneACCRegionInto(&loopOp.getRegion(), forOps.back().getBody(),
|
|
rewriter.getInsertionPoint(), mapping, rewriter);
|
|
|
|
// Optionally collapse nested loops
|
|
if (enableCollapse && forOps.size() > 1)
|
|
if (failed(coalesceLoops(rewriter, forOps)))
|
|
loopOp.emitError("failed to collapse acc.loop");
|
|
|
|
return forOps.front();
|
|
}
|
|
|
|
scf::ParallelOp convertACCLoopToSCFParallel(LoopOp loopOp,
|
|
RewriterBase &rewriter) {
|
|
assert(!loopOp.getUnstructured() &&
|
|
"use convertUnstructuredACCLoopToSCFExecuteRegion for unstructured "
|
|
"loops");
|
|
assert(
|
|
rewriter.getInsertionBlock() &&
|
|
!loopOp->isProperAncestor(rewriter.getInsertionBlock()->getParentOp()) &&
|
|
"builder insertion point must not be inside the loop being converted");
|
|
|
|
Location loc = loopOp->getLoc();
|
|
|
|
SmallVector<Value> lowerBounds, upperBounds, steps;
|
|
|
|
// Normalize all loops: lb=0, step=1, ub=tripCount
|
|
Value lb = arith::ConstantIndexOp::create(rewriter, loc, 0);
|
|
Value step = arith::ConstantIndexOp::create(rewriter, loc, 1);
|
|
|
|
for (auto [idx, iv] : llvm::enumerate(loopOp.getBody().getArguments())) {
|
|
bool inclusiveUpperbound = false;
|
|
if (loopOp.getInclusiveUpperbound().has_value())
|
|
inclusiveUpperbound = loopOp.getInclusiveUpperbound().value()[idx];
|
|
|
|
Value ub = calculateTripCount(rewriter, loc, loopOp.getLowerbound()[idx],
|
|
loopOp.getUpperbound()[idx],
|
|
loopOp.getStep()[idx], inclusiveUpperbound);
|
|
|
|
lowerBounds.push_back(lb);
|
|
upperBounds.push_back(ub);
|
|
steps.push_back(step);
|
|
}
|
|
|
|
auto parallelOp =
|
|
scf::ParallelOp::create(rewriter, loc, lowerBounds, upperBounds, steps);
|
|
|
|
// Create IV type conversions
|
|
IRMapping mapping;
|
|
rewriter.setInsertionPointToStart(parallelOp.getBody());
|
|
mapACCLoopIVsToSCFIVs(loopOp, parallelOp.getInductionVars(), rewriter,
|
|
mapping);
|
|
|
|
if (!loopOp.getRegion().hasOneBlock()) {
|
|
auto exeRegion = wrapMultiBlockRegionWithSCFExecuteRegion(
|
|
loopOp.getRegion(), mapping, loc, rewriter);
|
|
if (!exeRegion) {
|
|
rewriter.eraseOp(parallelOp);
|
|
return nullptr;
|
|
}
|
|
} else {
|
|
cloneACCRegionInto(&loopOp.getRegion(), parallelOp.getBody(),
|
|
rewriter.getInsertionPoint(), mapping, rewriter);
|
|
}
|
|
|
|
// Denormalize IV uses
|
|
rewriter.setInsertionPointToStart(parallelOp.getBody());
|
|
for (auto [idx, iv] : llvm::enumerate(parallelOp.getBody()->getArguments()))
|
|
if (!iv.use_empty())
|
|
normalizeIVUses(rewriter, loc, iv, loopOp.getLowerbound()[idx],
|
|
loopOp.getStep()[idx]);
|
|
|
|
return parallelOp;
|
|
}
|
|
|
|
scf::ExecuteRegionOp
|
|
convertUnstructuredACCLoopToSCFExecuteRegion(LoopOp loopOp,
|
|
RewriterBase &rewriter) {
|
|
assert(loopOp.getUnstructured() &&
|
|
"use convertACCLoopToSCFFor for structured loops");
|
|
assert(
|
|
rewriter.getInsertionBlock() &&
|
|
!loopOp->isProperAncestor(rewriter.getInsertionBlock()->getParentOp()) &&
|
|
"builder insertion point must not be inside the loop being converted");
|
|
|
|
IRMapping mapping;
|
|
return wrapMultiBlockRegionWithSCFExecuteRegion(loopOp.getRegion(), mapping,
|
|
loopOp->getLoc(), rewriter);
|
|
}
|
|
|
|
} // namespace acc
|
|
} // namespace mlir
|