ce6dd8b02d
This pass canonicalizes the use_device clause on acc.host_data
constructs to enable simpler runtime lowering. For use_device operands
that are box types or references to boxes, the pass:
1. Extracts the host base address for mapping to a device address using
acc.use_device
2. Creates a new boxed descriptor with the device address as the base
address for use inside the host_data region
The pass also removes unused use_device clauses to reduce runtime calls.
This canonicalization hoists load/box_addr patterns out of the host_data
region so they are applied to the host variable before acc.use_device,
ensuring the device pointer is used directly inside the region.
Example transformation for a reference to a box (!fir.ref<!fir.box<>>):
Before:
```
%ptr = acc.use_device varPtr(%ref : !fir.ref<!fir.box<!fir.ptr<i32>>>)
acc.host_data dataOperands(%ptr) {
%box = fir.load %ptr
%addr = fir.box_addr %box
// use %addr
}
```
After:
```
%box = fir.load %ref
%addr = fir.box_addr %box
%dev_ptr = acc.use_device varPtr(%addr : !fir.ptr<i32>)
acc.host_data dataOperands(%dev_ptr) {
%new_box = fir.embox %dev_ptr
// use device pointer through new descriptor
}
```
---------
Co-authored-by: nvptm <pmathew@nvidia.com>
401 lines
17 KiB
C++
401 lines
17 KiB
C++
//===- ACCUseDeviceCanonicalizer.cpp --------------------------------------===//
|
|
//
|
|
// 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 pass canonicalizes the use_device clause on a host_data construct such
|
|
// that use_device(x) can be lowered to a simple runtime call that takes the
|
|
// actual host pointer as argument.
|
|
//
|
|
// For a use_device operand that is a box type or a reference to a box, the
|
|
// pass:
|
|
// 1. Extracts the host base address for mapping to a device address using
|
|
// acc.use_device.
|
|
// 2. Creates a new boxed descriptor with the device address as the base
|
|
// address for use inside the host_data region.
|
|
//
|
|
// The pass also removes unused use_device clauses, reducing the number of
|
|
// runtime calls.
|
|
//
|
|
// Supported use_device operand types:
|
|
//
|
|
// Scalars:
|
|
// - !fir.ref<i32>, !fir.ref<f64>, etc.
|
|
//
|
|
// Arrays:
|
|
// - Explicit shape (no descriptor): !fir.ref<!fir.array<100xi32>>
|
|
// - Adjustable size: !fir.ref<!fir.array<?xi32>>
|
|
// - Assumed shape (handled by hoistBox): !fir.box<!fir.array<?xi32>>
|
|
// - Assumed size: !fir.ref<!fir.array<?xi32>>
|
|
// - Deferred shape (handled by hoistRefToBox):
|
|
// - Allocatable: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
|
|
// - Pointer: !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>
|
|
// - Subarray specification (handled by hoistBox):
|
|
// !fir.box<!fir.array<?xi32>>
|
|
//
|
|
// Not yet supported:
|
|
// - Assumed rank arrays
|
|
// - Composite variables: !fir.ref<!fir.type<...>>
|
|
// - Array elements (device pointer arithmetic in host_data region)
|
|
// - Composite variable members
|
|
// - Fortran common blocks: use_device(/cm_block/)
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "flang/Optimizer/Builder/BoxValue.h"
|
|
#include "flang/Optimizer/Builder/FIRBuilder.h"
|
|
#include "flang/Optimizer/Dialect/FIROps.h"
|
|
#include "flang/Optimizer/Dialect/FIRType.h"
|
|
#include "flang/Optimizer/OpenACC/Passes.h"
|
|
#include "mlir/Dialect/OpenACC/OpenACC.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/IR/Value.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
#include "llvm/Support/Debug.h"
|
|
#include <cassert>
|
|
|
|
namespace fir::acc {
|
|
#define GEN_PASS_DEF_ACCUSEDEVICECANONICALIZER
|
|
#include "flang/Optimizer/OpenACC/Passes.h.inc"
|
|
} // namespace fir::acc
|
|
|
|
#define DEBUG_TYPE "acc-use-device-canonicalizer"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
|
|
struct UseDeviceHostDataHoisting : public OpRewritePattern<acc::HostDataOp> {
|
|
using OpRewritePattern<acc::HostDataOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(acc::HostDataOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
SmallVector<Value> usedOperands;
|
|
SmallVector<Value> unusedUseDeviceOperands;
|
|
SmallVector<acc::UseDeviceOp> refToBoxUseDeviceOps;
|
|
SmallVector<acc::UseDeviceOp> boxUseDeviceOps;
|
|
|
|
for (Value operand : op.getDataClauseOperands()) {
|
|
if (acc::UseDeviceOp useDeviceOp =
|
|
operand.getDefiningOp<acc::UseDeviceOp>()) {
|
|
if (fir::isBoxAddress(useDeviceOp.getVar().getType())) {
|
|
if (!llvm::hasSingleElement(useDeviceOp->getUsers()))
|
|
refToBoxUseDeviceOps.push_back(useDeviceOp);
|
|
} else if (isa<fir::BoxType>(useDeviceOp.getVar().getType())) {
|
|
if (!llvm::hasSingleElement(useDeviceOp->getUsers()))
|
|
boxUseDeviceOps.push_back(useDeviceOp);
|
|
}
|
|
|
|
// host_data is the only user of this use_device operand - mark for
|
|
// removal
|
|
if (llvm::hasSingleElement(useDeviceOp->getUsers()))
|
|
unusedUseDeviceOperands.push_back(useDeviceOp.getResult());
|
|
else
|
|
usedOperands.push_back(useDeviceOp.getResult());
|
|
} else {
|
|
// Operand is not an `acc.use_device` result, keep it as is.
|
|
usedOperands.push_back(operand);
|
|
}
|
|
}
|
|
|
|
assert(!usedOperands.empty() && "Host_data operation has no used operands");
|
|
|
|
if (!unusedUseDeviceOperands.empty()) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "ACCUseDeviceCanonicalizer: Removing "
|
|
<< unusedUseDeviceOperands.size()
|
|
<< " unused use_device operands from host_data operation\n");
|
|
|
|
// Update the host_data operation to have only used operands
|
|
rewriter.modifyOpInPlace(op, [&]() {
|
|
op.getDataClauseOperandsMutable().assign(usedOperands);
|
|
});
|
|
|
|
// Remove unused use_device operations
|
|
for (Value operand : unusedUseDeviceOperands) {
|
|
acc::UseDeviceOp useDeviceOp =
|
|
operand.getDefiningOp<acc::UseDeviceOp>();
|
|
LLVM_DEBUG(llvm::dbgs() << "ACCUseDeviceCanonicalizer: Erasing: "
|
|
<< *useDeviceOp << "\n");
|
|
rewriter.eraseOp(useDeviceOp);
|
|
}
|
|
return success();
|
|
}
|
|
|
|
// Handle references to box types
|
|
bool modified = false;
|
|
for (acc::UseDeviceOp useDeviceOp : refToBoxUseDeviceOps)
|
|
modified |=
|
|
hoistRefToBox(rewriter, useDeviceOp.getResult(), useDeviceOp, op);
|
|
|
|
// Handle box types
|
|
for (acc::UseDeviceOp useDeviceOp : boxUseDeviceOps)
|
|
modified |= hoistBox(rewriter, useDeviceOp.getResult(), useDeviceOp, op);
|
|
|
|
return modified ? success() : failure();
|
|
}
|
|
|
|
private:
|
|
/// Collect users of `acc.use_device` operation inside the `acc.host_data`
|
|
/// region that need to be updated with the final replacement value.
|
|
void collectUseDeviceUsersToUpdate(
|
|
acc::UseDeviceOp useDeviceOp, acc::HostDataOp hostDataOp,
|
|
SmallVectorImpl<Operation *> &usersToUpdate) const {
|
|
for (mlir::Operation *user : useDeviceOp->getUsers())
|
|
if (hostDataOp.getRegion().isAncestor(user->getParentRegion()))
|
|
usersToUpdate.push_back(user);
|
|
}
|
|
|
|
/// Create new `acc.use_device` operation with the given box address as
|
|
/// operand. Updates the `acc.host_data` operation to use the new
|
|
/// `acc.use_device` result.
|
|
acc::UseDeviceOp createNewUseDeviceOp(PatternRewriter &rewriter,
|
|
acc::UseDeviceOp useDeviceOp,
|
|
acc::HostDataOp hostDataOp,
|
|
fir::BoxAddrOp boxAddr) const {
|
|
// Create use_device on the raw pointer
|
|
acc::UseDeviceOp newUseDeviceOp = acc::UseDeviceOp::create(
|
|
rewriter, useDeviceOp.getLoc(), boxAddr.getType(), boxAddr.getResult(),
|
|
useDeviceOp.getVarTypeAttr(), useDeviceOp.getVarPtrPtr(),
|
|
useDeviceOp.getBounds(), useDeviceOp.getAsyncOperands(),
|
|
useDeviceOp.getAsyncOperandsDeviceTypeAttr(),
|
|
useDeviceOp.getAsyncOnlyAttr(), useDeviceOp.getDataClauseAttr(),
|
|
useDeviceOp.getStructuredAttr(), useDeviceOp.getImplicitAttr(),
|
|
useDeviceOp.getModifiersAttr(), useDeviceOp.getNameAttr(),
|
|
useDeviceOp.getRecipeAttr());
|
|
|
|
LLVM_DEBUG(llvm::dbgs() << "Created new hoisted pattern for box access:\n"
|
|
<< " box_addr: " << *boxAddr << "\n"
|
|
<< " new use_device: " << *newUseDeviceOp << "\n");
|
|
|
|
// Replace the old `acc.use_device` operand in the `acc.host_data` operation
|
|
// with the new one
|
|
rewriter.modifyOpInPlace(hostDataOp, [&]() {
|
|
hostDataOp->replaceUsesOfWith(useDeviceOp.getResult(),
|
|
newUseDeviceOp.getResult());
|
|
});
|
|
|
|
return newUseDeviceOp;
|
|
}
|
|
|
|
/// Canonicalize use_device operand that is a reference to a box.
|
|
/// Transforms:
|
|
/// %3 = fir.address_of(@_QFEtgt) : !fir.ref<i32>
|
|
/// %5 = fir.embox %3 : (!fir.ref<i32>) -> !fir.box<!fir.ptr<i32>>
|
|
/// fir.store %5 to %0 : !fir.ref<!fir.box<!fir.ptr<i32>>>
|
|
/// %9 = acc.use_device varPtr(%0 : !fir.ref<!fir.box<!fir.ptr<i32>>>)
|
|
/// -> !fir.ref<!fir.box<!fir.ptr<i32>>> {name = "ptr"}
|
|
/// acc.host_data dataOperands(%9 : !fir.ref<!fir.box<!fir.ptr<i32>>>) {
|
|
/// %loaded = fir.load %9 : !fir.ref<!fir.box<!fir.ptr<i32>>>
|
|
/// %addr = fir.box_addr %loaded : (!fir.box<!fir.ptr<i32>>) ->
|
|
/// !fir.ptr<i32> %conv = fir.convert %addr : (!fir.ptr<i32>) -> i64
|
|
/// fir.call @foo(%conv) : (i64) -> ()
|
|
/// acc.terminator
|
|
/// }
|
|
/// into:
|
|
/// %loaded = fir.load %0 : !fir.ref<!fir.box<!fir.ptr<i32>>>
|
|
/// %addr = fir.box_addr %loaded : (!fir.box<!fir.ptr<i32>>) ->
|
|
/// !fir.ptr<i32>
|
|
/// %dev_ptr = acc.use_device varPtr(%addr : !fir.ptr<i32>) ->
|
|
/// !fir.ptr<i32>
|
|
/// -> !fir.ref<!fir.box<!fir.ptr<i32>>> {name = "ptr"}
|
|
/// acc.host_data dataOperands(%dev_ptr : !fir.ref<!fir.box<!fir.ptr<i32>>>)
|
|
/// {
|
|
/// %embox = fir.embox %dev_ptr : (!fir.ptr<i32>) ->
|
|
/// !fir.box<!fir.ptr<i32>> %alloca = fir.alloca !fir.box<!fir.ptr<i32>>
|
|
/// fir.store %embox to %alloca : !fir.ref<!fir.box<!fir.ptr<i32>>>
|
|
/// %loaded2 = fir.load %alloca : !fir.ref<!fir.box<!fir.ptr<i32>>>
|
|
/// %addr2 = fir.box_addr %loaded2 : (!fir.box<!fir.ptr<i32>>) ->
|
|
/// !fir.ptr<i32> %conv = fir.convert %addr2 : (!fir.ptr<i32>) -> i64
|
|
/// fir.call @foo(%conv) : (i64) -> ()
|
|
/// acc.terminator
|
|
/// }
|
|
bool hoistRefToBox(PatternRewriter &rewriter, Value operand,
|
|
acc::UseDeviceOp useDeviceOp,
|
|
acc::HostDataOp hostDataOp) const {
|
|
|
|
// Safety check: if the use_device operation is already using a box_addr
|
|
// result, it means it has already been processed, so skip to avoid infinite
|
|
// loop
|
|
if (useDeviceOp.getVar().getDefiningOp<fir::BoxAddrOp>()) {
|
|
LLVM_DEBUG(llvm::dbgs() << "ACCUseDeviceCanonicalizer: Skipping "
|
|
"already processed use_device operation\n");
|
|
return false;
|
|
}
|
|
// Get the ModuleOp before we erase useDeviceOp to avoid invalid reference
|
|
ModuleOp mod = useDeviceOp->getParentOfType<ModuleOp>();
|
|
|
|
// Collect users of the original `acc.use_device` operation that need to be
|
|
// updated
|
|
SmallVector<Operation *> usersToUpdate;
|
|
collectUseDeviceUsersToUpdate(useDeviceOp, hostDataOp, usersToUpdate);
|
|
|
|
rewriter.setInsertionPoint(useDeviceOp);
|
|
// Create a load operation to get the box from the variable
|
|
fir::LoadOp box = fir::LoadOp::create(rewriter, useDeviceOp.getLoc(),
|
|
useDeviceOp.getVar());
|
|
// Create a box_addr operation to get the address from the box
|
|
fir::BoxAddrOp boxAddr =
|
|
fir::BoxAddrOp::create(rewriter, useDeviceOp.getLoc(), box);
|
|
|
|
acc::UseDeviceOp newUseDeviceOp =
|
|
createNewUseDeviceOp(rewriter, useDeviceOp, hostDataOp, boxAddr);
|
|
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Created new hoisted pattern for pointer access:\n"
|
|
<< " load box: " << *box << "\n"
|
|
<< " box_addr: " << *boxAddr << "\n"
|
|
<< " new use_device: " << *newUseDeviceOp << "\n");
|
|
|
|
// Set insertion point to the first op inside the host_data region
|
|
rewriter.setInsertionPoint(&hostDataOp.getRegion().front().front());
|
|
|
|
// Create a FirOpBuilder from the PatternRewriter using the module we got
|
|
// earlier
|
|
fir::FirOpBuilder builder(rewriter, mod);
|
|
Value newBoxwithDevicePtr = fir::factory::getDescriptorWithNewBaseAddress(
|
|
builder, useDeviceOp.getLoc(), box.getResult(),
|
|
newUseDeviceOp.getResult());
|
|
|
|
// Create new memory location and store the newBoxwithDevicePtr into new
|
|
// memory location
|
|
fir::AllocaOp newMemLoc = fir::AllocaOp::create(
|
|
rewriter, useDeviceOp.getLoc(), newBoxwithDevicePtr.getType());
|
|
[[maybe_unused]] fir::StoreOp newStoreOp = fir::StoreOp::create(
|
|
rewriter, useDeviceOp.getLoc(), newBoxwithDevicePtr, newMemLoc);
|
|
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "host_data region updated with new host descriptor "
|
|
"containing device pointer:\n"
|
|
<< " box with device pointer: "
|
|
<< *newBoxwithDevicePtr.getDefiningOp() << "\n"
|
|
<< " mem loc: " << *newMemLoc << "\n"
|
|
<< " store op: " << *newStoreOp << "\n");
|
|
|
|
// Replace all uses of the original `acc.use_device` operation inside the
|
|
// `acc.host_data` region with the new memory location containing the box
|
|
// with device pointer
|
|
for (mlir::Operation *user : usersToUpdate)
|
|
user->replaceUsesOfWith(useDeviceOp.getResult(), newMemLoc);
|
|
|
|
assert(useDeviceOp.getResult().use_empty() &&
|
|
"expected all uses of use_device to be replaced");
|
|
rewriter.eraseOp(useDeviceOp);
|
|
return true;
|
|
}
|
|
|
|
/// Canonicalize use_device operand that is a box type.
|
|
/// Transforms:
|
|
/// %box = ... : !fir.box<!fir.array<?xi32>>
|
|
/// %dev_box = acc.use_device varPtr(%box : !fir.box<!fir.array<?xi32>>)
|
|
/// -> !fir.box<!fir.array<?xi32>>
|
|
/// acc.host_data dataOperands(%dev_box : !fir.box<!fir.array<?xi32>>) {
|
|
/// %addr = fir.box_addr %dev_box : (!fir.box<!fir.array<?xi32>>) ->
|
|
/// !fir.heap<!fir.array<?xi32>>
|
|
/// // use %addr
|
|
/// }
|
|
/// into:
|
|
/// %box = ... : !fir.box<!fir.array<?xi32>>
|
|
/// %addr = fir.box_addr %box : (!fir.box<!fir.array<?xi32>>) ->
|
|
/// !fir.heap<!fir.array<?xi32>>
|
|
/// %dev_ptr = acc.use_device varPtr(%addr : !fir.heap<!fir.array<?xi32>>)
|
|
/// -> !fir.heap<!fir.array<?xi32>>
|
|
/// acc.host_data dataOperands(%dev_ptr : !fir.heap<!fir.array<?xi32>>) {
|
|
/// %new_box = fir.embox %dev_ptr ... : !fir.box<!fir.array<?xi32>>
|
|
/// %new_addr = fir.box_addr %new_box : (!fir.box<!fir.array<?xi32>>) ->
|
|
/// !fir.heap<!fir.array<?xi32>>
|
|
/// // use %new_addr instead of %addr
|
|
/// }
|
|
bool hoistBox(PatternRewriter &rewriter, Value operand,
|
|
acc::UseDeviceOp useDeviceOp,
|
|
acc::HostDataOp hostDataOp) const {
|
|
|
|
// Safety check: if the use_device operation is already using a box_addr
|
|
// result, it means it has already been processed, so skip to avoid infinite
|
|
// loop
|
|
if (useDeviceOp.getVar().getDefiningOp<fir::BoxAddrOp>()) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "ACCUseDeviceCanonicalizer: Skipping "
|
|
"already processed box use_device operation\n");
|
|
return false;
|
|
}
|
|
|
|
// Collect users of the original `acc.use_device` operation that need to be
|
|
// updated
|
|
SmallVector<Operation *> usersToUpdate;
|
|
collectUseDeviceUsersToUpdate(useDeviceOp, hostDataOp, usersToUpdate);
|
|
|
|
// Get the ModuleOp before we erase useDeviceOp to avoid invalid reference
|
|
ModuleOp mod = useDeviceOp->getParentOfType<ModuleOp>();
|
|
|
|
rewriter.setInsertionPoint(useDeviceOp);
|
|
// Extract the raw pointer from the box descriptor
|
|
fir::BoxAddrOp boxAddr = fir::BoxAddrOp::create(
|
|
rewriter, useDeviceOp.getLoc(), useDeviceOp.getVar());
|
|
|
|
acc::UseDeviceOp newUseDeviceOp =
|
|
createNewUseDeviceOp(rewriter, useDeviceOp, hostDataOp, boxAddr);
|
|
|
|
// Set insertion point to the first op inside the host_data region
|
|
rewriter.setInsertionPoint(&hostDataOp.getRegion().front().front());
|
|
|
|
// Create a FirOpBuilder from the PatternRewriter using the module we got
|
|
// earlier
|
|
fir::FirOpBuilder builder(rewriter, mod);
|
|
|
|
// Create a new host descriptor at the start of the host_data region
|
|
// with the device pointer as the base address
|
|
Value newBoxWithDevicePtr = fir::factory::getDescriptorWithNewBaseAddress(
|
|
builder, useDeviceOp.getLoc(), useDeviceOp.getVar(),
|
|
newUseDeviceOp.getResult());
|
|
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "host_data region updated with new host descriptor "
|
|
"containing device pointer:\n"
|
|
<< " box with device pointer: "
|
|
<< *newBoxWithDevicePtr.getDefiningOp() << "\n");
|
|
|
|
// Replace all uses of the original `acc.use_device` operation inside the
|
|
// `acc.host_data` region with the new box containing device pointer
|
|
for (mlir::Operation *user : usersToUpdate)
|
|
user->replaceUsesOfWith(useDeviceOp.getResult(), newBoxWithDevicePtr);
|
|
|
|
assert(useDeviceOp.getResult().use_empty() &&
|
|
"expected all uses of use_device to be replaced");
|
|
rewriter.eraseOp(useDeviceOp);
|
|
return true;
|
|
}
|
|
};
|
|
|
|
class ACCUseDeviceCanonicalizer
|
|
: public fir::acc::impl::ACCUseDeviceCanonicalizerBase<
|
|
ACCUseDeviceCanonicalizer> {
|
|
public:
|
|
void runOnOperation() override {
|
|
MLIRContext *context = getOperation()->getContext();
|
|
|
|
RewritePatternSet patterns(context);
|
|
|
|
// Add the custom use_device canonicalization patterns
|
|
patterns.insert<UseDeviceHostDataHoisting>(context);
|
|
|
|
// Apply patterns greedily
|
|
GreedyRewriteConfig config;
|
|
// Prevent the pattern driver from merging blocks.
|
|
config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled);
|
|
config.setUseTopDownTraversal(true);
|
|
|
|
(void)applyPatternsGreedily(getOperation(), std::move(patterns), config);
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
std::unique_ptr<mlir::Pass> fir::acc::createACCUseDeviceCanonicalizerPass() {
|
|
return std::make_unique<ACCUseDeviceCanonicalizer>();
|
|
}
|