| //===- 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>(); |
| } |
