[flang-commits] [flang] [flang][acc] Add ACCUseDeviceCanonicalizer pass (PR #175228)
Razvan Lupusoru via flang-commits
flang-commits at lists.llvm.org
Fri Jan 9 11:41:05 PST 2026
https://github.com/razvanlupusoru created https://github.com/llvm/llvm-project/pull/175228
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
}
```
>From 5a055959fa432375da142faaf47b960e77576a0a Mon Sep 17 00:00:00 2001
From: nvptm <pmathew at nvidia.com>
Date: Fri, 9 Jan 2026 11:26:35 -0800
Subject: [PATCH] [flang][acc] Add ACCUseDeviceCanonicalizer pass
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
}
```
---
.../include/flang/Optimizer/OpenACC/Passes.h | 1 +
.../include/flang/Optimizer/OpenACC/Passes.td | 21 +
.../Transforms/ACCUseDeviceCanonicalizer.cpp | 428 ++++++++++++++++++
.../OpenACC/Transforms/CMakeLists.txt | 3 +
.../Fir/OpenACC/use-device-canonicalizer.mlir | 91 ++++
5 files changed, 544 insertions(+)
create mode 100644 flang/lib/Optimizer/OpenACC/Transforms/ACCUseDeviceCanonicalizer.cpp
create mode 100644 flang/test/Fir/OpenACC/use-device-canonicalizer.mlir
diff --git a/flang/include/flang/Optimizer/OpenACC/Passes.h b/flang/include/flang/Optimizer/OpenACC/Passes.h
index c27c7ebc3b06f..64ddb84e63c3e 100644
--- a/flang/include/flang/Optimizer/OpenACC/Passes.h
+++ b/flang/include/flang/Optimizer/OpenACC/Passes.h
@@ -30,6 +30,7 @@ namespace acc {
std::unique_ptr<mlir::Pass> createACCInitializeFIRAnalysesPass();
std::unique_ptr<mlir::Pass> createACCRecipeBufferizationPass();
+std::unique_ptr<mlir::Pass> createACCUseDeviceCanonicalizerPass();
} // namespace acc
} // namespace fir
diff --git a/flang/include/flang/Optimizer/OpenACC/Passes.td b/flang/include/flang/Optimizer/OpenACC/Passes.td
index d947aa470494a..8579a471d9a56 100644
--- a/flang/include/flang/Optimizer/OpenACC/Passes.td
+++ b/flang/include/flang/Optimizer/OpenACC/Passes.td
@@ -49,4 +49,25 @@ def ACCRecipeBufferization
}];
}
+def ACCUseDeviceCanonicalizer
+ : Pass<"acc-use-device-canonicalizer", "mlir::func::FuncOp"> {
+ let summary = "Canonicalize acc.use_device operations for FIR box types";
+ let description = [{
+ 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.
+ }];
+ let dependentDialects = ["mlir::acc::OpenACCDialect", "fir::FIROpsDialect"];
+}
+
#endif // FORTRAN_OPTIMIZER_OPENACC_PASSES
diff --git a/flang/lib/Optimizer/OpenACC/Transforms/ACCUseDeviceCanonicalizer.cpp b/flang/lib/Optimizer/OpenACC/Transforms/ACCUseDeviceCanonicalizer.cpp
new file mode 100644
index 0000000000000..27a5b2316cbfe
--- /dev/null
+++ b/flang/lib/Optimizer/OpenACC/Transforms/ACCUseDeviceCanonicalizer.cpp
@@ -0,0 +1,428 @@
+//===- 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 {
+
+static bool isFirRefToBox(Type type) {
+ // First check if it's a FIR reference type
+ if (!fir::isa_ref_type(type))
+ return false;
+
+ // Unwrap the reference to get the inner type
+ Type innerType = fir::unwrapRefType(type);
+
+ // Check if the inner type is a box type
+ return isa<fir::BoxType>(innerType);
+}
+
+struct UseDeviceHostDataHoisting : public OpRewritePattern<acc::HostDataOp> {
+ using OpRewritePattern<acc::HostDataOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(acc::HostDataOp op,
+ PatternRewriter &rewriter) const override {
+
+ // Check if any of the data operands are acc.use_device results
+ bool hasUseDeviceOperand = false;
+ 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>()) {
+ hasUseDeviceOperand = true;
+ if (isFirRefToBox(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 {
+ // use_device operands that don't need hoisting, they are passed as
+ // is to __tgt_acc_get_deviceptr.
+ 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();
+ }
+
+ bool patternModified = false;
+ if (!refToBoxUseDeviceOps.empty()) {
+ // Handle pointer to box types
+ for (acc::UseDeviceOp useDeviceOp : refToBoxUseDeviceOps) {
+ patternModified =
+ hoistRefToBox(rewriter, useDeviceOp.getResult(), useDeviceOp, op);
+ break; // Only handle one at a time to avoid iterator invalidation
+ }
+ return patternModified ? success() : failure();
+ }
+
+ if (!boxUseDeviceOps.empty()) {
+ // Handle box types
+ for (acc::UseDeviceOp useDeviceOp : boxUseDeviceOps) {
+ patternModified =
+ hoistBox(rewriter, useDeviceOp.getResult(), useDeviceOp, op);
+ break; // Only handle one at a time to avoid iterator invalidation
+ }
+ return patternModified ? success() : failure();
+ }
+
+ if (!hasUseDeviceOperand)
+ return failure();
+
+ return patternModified ? success() : failure();
+ }
+
+private:
+ /// Create new use_device operation with the given box address as operand.
+ 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());
+
+ // Replace the old use_device operand in the host_data with the new one
+ SmallVector<Value> newOperands;
+ for (Value operand : hostDataOp.getDataClauseOperands()) {
+ if (operand == useDeviceOp.getResult())
+ newOperands.push_back(newUseDeviceOp.getResult());
+ else
+ newOperands.push_back(operand);
+ }
+
+ LLVM_DEBUG(llvm::dbgs() << "Created new hoisted pattern for box access:\n"
+ << " box_addr: " << *boxAddr << "\n"
+ << " new use_device: " << *newUseDeviceOp << "\n");
+
+ // Update the host_data operation
+ rewriter.modifyOpInPlace(hostDataOp, [&]() {
+ hostDataOp.getDataClauseOperandsMutable().assign(newOperands);
+ });
+
+ SmallVector<Operation *> usersOfUseDeviceOp;
+ for (mlir::Operation *user : useDeviceOp->getUsers())
+ if (hostDataOp.getRegion().isAncestor(user->getParentRegion()))
+ usersOfUseDeviceOp.push_back(user);
+
+ for (mlir::Operation *user : usersOfUseDeviceOp)
+ user->replaceUsesOfWith(useDeviceOp.getResult(),
+ newUseDeviceOp.getResult());
+
+ // Set insertion point to the first op inside the host_data region
+ rewriter.setInsertionPoint(&hostDataOp.getRegion().front().front());
+
+ 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>();
+
+ 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");
+ // 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");
+ SmallVector<Operation *> usersOfNewUseDeviceOp;
+ for (mlir::Operation *user : newUseDeviceOp->getUsers())
+ if (hostDataOp.getRegion().isAncestor(user->getParentRegion()))
+ usersOfNewUseDeviceOp.push_back(user);
+
+ // replace all uses of the newUseDeviceOp.getResult() strictly inside the
+ // host_data region with the newBoxwithDevicePtr
+ for (mlir::Operation *user : usersOfNewUseDeviceOp)
+ if (hostDataOp.getRegion().isAncestor(user->getParentRegion()))
+ if (user != hostDataOp && user != newBoxwithDevicePtr.getDefiningOp())
+ user->replaceUsesOfWith(newUseDeviceOp.getResult(), newMemLoc);
+ 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;
+ }
+
+ // 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);
+
+ // 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");
+
+ SmallVector<Operation *> usersOfNewUseDeviceOp;
+ for (mlir::Operation *user : newUseDeviceOp->getUsers())
+ if (hostDataOp.getRegion().isAncestor(user->getParentRegion()))
+ usersOfNewUseDeviceOp.push_back(user);
+
+ // Replace all uses of the newUseDeviceOp.getResult() strictly inside the
+ // host_data region with the newBoxWithDevicePtr
+ for (mlir::Operation *user : usersOfNewUseDeviceOp)
+ if (hostDataOp.getRegion().isAncestor(user->getParentRegion()))
+ if (user != hostDataOp && user != newBoxWithDevicePtr.getDefiningOp())
+ user->replaceUsesOfWith(newUseDeviceOp.getResult(),
+ newBoxWithDevicePtr);
+
+ 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;
+ config.setUseTopDownTraversal(true);
+ config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Normal);
+ config.setMaxIterations(10);
+ config.setMaxNumRewrites(-1);
+
+ (void)applyPatternsGreedily(getOperation(), std::move(patterns), config);
+ }
+};
+
+} // namespace
+
+std::unique_ptr<mlir::Pass> fir::acc::createACCUseDeviceCanonicalizerPass() {
+ return std::make_unique<ACCUseDeviceCanonicalizer>();
+}
diff --git a/flang/lib/Optimizer/OpenACC/Transforms/CMakeLists.txt b/flang/lib/Optimizer/OpenACC/Transforms/CMakeLists.txt
index d41e99a6c0679..9bd56c1618544 100644
--- a/flang/lib/Optimizer/OpenACC/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/OpenACC/Transforms/CMakeLists.txt
@@ -1,4 +1,5 @@
add_flang_library(FIROpenACCTransforms
+ ACCUseDeviceCanonicalizer.cpp
ACCInitializeFIRAnalyses.cpp
ACCRecipeBufferization.cpp
@@ -7,6 +8,7 @@ add_flang_library(FIROpenACCTransforms
LINK_LIBS
FIRAnalysis
+ FIRBuilder
FIRDialect
FIROpenACCAnalysis
HLFIRDialect
@@ -16,4 +18,5 @@ add_flang_library(FIROpenACCTransforms
MLIRPass
MLIROpenACCDialect
MLIROpenACCUtils
+ MLIRTransformUtils
)
diff --git a/flang/test/Fir/OpenACC/use-device-canonicalizer.mlir b/flang/test/Fir/OpenACC/use-device-canonicalizer.mlir
new file mode 100644
index 0000000000000..8776113f6eebd
--- /dev/null
+++ b/flang/test/Fir/OpenACC/use-device-canonicalizer.mlir
@@ -0,0 +1,91 @@
+// RUN: fir-opt %s --acc-use-device-canonicalizer -split-input-file | FileCheck %s
+
+// -----
+
+// Test hoisting of load/box_addr/convert pattern out of acc.host_data with function call
+func.func @test_host_data_hoisting_function_call(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) {
+ // CHECK: %[[LOADED:.*]] = fir.load %arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ // CHECK: %[[ADDR:.*]] = fir.box_addr %[[LOADED]] : (!fir.box<!fir.heap<!fir.array<?xf64>>>) -> !fir.heap<!fir.array<?xf64>>
+ // CHECK: %[[DEV_PTR:.*]] = acc.use_device varPtr(%[[ADDR]] : !fir.heap<!fir.array<?xf64>>) varType(!fir.box<!fir.heap<!fir.array<?xf64>>>) -> !fir.heap<!fir.array<?xf64>>
+ // CHECK: acc.host_data dataOperands(%[[DEV_PTR]]
+ // CHECK: %[[EMBOX:.*]] = fir.embox %[[DEV_PTR]]
+ // CHECK: %[[ALLOCA:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf64>>>
+ // CHECK: fir.store %[[EMBOX]] to %[[ALLOCA]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ // CHECK: %[[LOAD2:.*]] = fir.load %[[ALLOCA]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ // CHECK: %[[DEV_PTRADDR:.*]] = fir.box_addr %[[LOAD2]] : (!fir.box<!fir.heap<!fir.array<?xf64>>>
+ // CHECK: %[[CONV:.*]] = fir.convert %[[DEV_PTRADDR]] : (!fir.heap<!fir.array<?xf64>>) -> !fir.ref<!fir.array<?xf64>>
+ // CHECK: fir.call @_QMmPvadd(%[[CONV]]
+ %dev_ptr = acc.use_device varPtr(%arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>> {name = "a"}
+ acc.host_data dataOperands(%dev_ptr : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) {
+ %loaded = fir.load %dev_ptr : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ %addr = fir.box_addr %loaded : (!fir.box<!fir.heap<!fir.array<?xf64>>>) -> !fir.heap<!fir.array<?xf64>>
+ %conv = fir.convert %addr : (!fir.heap<!fir.array<?xf64>>) -> !fir.ref<!fir.array<?xf64>>
+ fir.call @_QMmPvadd(%conv, %conv) : (!fir.ref<!fir.array<?xf64>>, !fir.ref<!fir.array<?xf64>>) -> ()
+ acc.terminator
+ }
+ return
+}
+
+// -----
+
+// Test hoisting of load/box_addr/convert pattern out of acc.host_data with load operation
+func.func @test_host_data_hoisting_load(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) {
+ // CHECK: %[[LOADED:.*]] = fir.load %arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ // CHECK: %[[ADDR:.*]] = fir.box_addr %[[LOADED]] : (!fir.box<!fir.heap<!fir.array<?xf64>>>) -> !fir.heap<!fir.array<?xf64>>
+ // CHECK: %[[DEV_PTR:.*]] = acc.use_device varPtr(%[[ADDR]] :
+ // CHECK: acc.host_data dataOperands(%[[DEV_PTR]]
+ // CHECK: %[[EMBOX:.*]] = fir.embox %[[DEV_PTR]]
+ // CHECK: %[[ALLOCA:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf64>>>
+ // CHECK: fir.store %[[EMBOX]] to %[[ALLOCA]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ // CHECK: %[[LOAD2:.*]] = fir.load %[[ALLOCA]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ // CHECK: %[[DEV_PTRADDR:.*]] = fir.box_addr %[[LOAD2]] : (!fir.box<!fir.heap<!fir.array<?xf64>>>
+ // CHECK: %[[CONV:.*]] = fir.convert %[[DEV_PTRADDR]] : (!fir.heap<!fir.array<?xf64>>) -> !fir.ref<!fir.array<?xf64>>
+ // CHECK: %[[VAL:.*]] = fir.load %[[CONV]]
+ %dev_ptr = acc.use_device varPtr(%arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>> {name = "a"}
+ acc.host_data dataOperands(%dev_ptr : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>) {
+ %loaded = fir.load %dev_ptr : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>
+ %addr = fir.box_addr %loaded : (!fir.box<!fir.heap<!fir.array<?xf64>>>) -> !fir.heap<!fir.array<?xf64>>
+ %conv = fir.convert %addr : (!fir.heap<!fir.array<?xf64>>) -> !fir.ref<!fir.array<?xf64>>
+ %val = fir.load %conv : !fir.ref<!fir.array<?xf64>>
+ fir.call @foo(%val) : (!fir.array<?xf64>) -> ()
+ acc.terminator
+ }
+ return
+}
+
+// -----
+
+// Test hoisting for pointer attributes: load/box_addr hoisted, remove additional
+// unused use_device clause
+func.func @test_host_data_hoisting_ref_to_box() {
+ %1 = fir.alloca !fir.box<!fir.ptr<i32>> {bindc_name = "ptr", uniq_name = "_QFEptr"}
+ // CHECK: %[[ALLOCA:.*]] = fir.alloca !fir.box<!fir.ptr<i32>> {bindc_name = "ptr", uniq_name = "_QFEptr"}
+ %4 = fir.declare %1 {fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QFEptr"} : (!fir.ref<!fir.box<!fir.ptr<i32>>>) -> !fir.ref<!fir.box<!fir.ptr<i32>>>
+ // CHECK: %[[DECLARE:.*]] = fir.declare %[[ALLOCA]] {fortran_attrs = #fir.var_attrs<pointer>, uniq_name = "_QFEptr"} : (!fir.ref<!fir.box<!fir.ptr<i32>>>) -> !fir.ref<!fir.box<!fir.ptr<i32>>>
+ %5 = fir.address_of(@_QFEtgt) : !fir.ref<i32>
+ %6 = fir.declare %5 {fortran_attrs = #fir.var_attrs<target>, uniq_name = "_QFEtgt"} : (!fir.ref<i32>) -> !fir.ref<i32>
+ %8 = fir.embox %6 : (!fir.ref<i32>) -> !fir.box<!fir.ptr<i32>>
+ fir.store %8 to %4 : !fir.ref<!fir.box<!fir.ptr<i32>>>
+ // CHECK: %[[LOAD:.*]] = fir.load %[[DECLARE]]
+ // CHECK: %[[BOXADDR:.*]] = fir.box_addr %[[LOAD]]
+ // CHECK: %[[DEV_PTR:.*]] = acc.use_device varPtr(%[[BOXADDR]] : !fir.ptr<i32>) varType(!fir.box<!fir.ptr<i32>>) -> !fir.ptr<i32> {name = "ptr"}
+ // CHECK: acc.host_data dataOperands(%[[DEV_PTR]] : !fir.ptr<i32>) {
+ // CHECK: %[[EMBOX:.*]] = fir.embox %[[DEV_PTR]] : (!fir.ptr<i32>) -> !fir.box<!fir.ptr<i32>>
+ // CHECK: %[[ALLOCA:.*]] = fir.alloca !fir.box<!fir.ptr<i32>>
+ // CHECK: fir.store %[[EMBOX]] to %[[ALLOCA]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
+ // CHECK: %[[LOAD2:.*]] = fir.load %[[ALLOCA]] : !fir.ref<!fir.box<!fir.ptr<i32>>>
+ // CHECK: %[[DEV_PTRADDR:.*]] = fir.box_addr %[[LOAD2]] : (!fir.box<!fir.ptr<i32>>) -> !fir.ptr<i32>
+ // CHECK: %[[CONV:.*]] = fir.convert %[[DEV_PTRADDR]] : (!fir.ptr<i32>) -> i64
+ // CHECK: fir.call @foo(%[[CONV]]) : (i64) -> ()
+ %9 = acc.use_device varPtr(%4 : !fir.ref<!fir.box<!fir.ptr<i32>>>) -> !fir.ref<!fir.box<!fir.ptr<i32>>> {name = "ptr"}
+ %12 = acc.use_device varPtr(%4 : !fir.ref<!fir.box<!fir.ptr<i32>>>) -> !fir.ref<!fir.box<!fir.ptr<i32>>> {name = "ptr"}
+ acc.host_data dataOperands(%9, %12 : !fir.ref<!fir.box<!fir.ptr<i32>>>, !fir.ref<!fir.box<!fir.ptr<i32>>>) {
+ %14 = fir.load %9 : !fir.ref<!fir.box<!fir.ptr<i32>>>
+ %15 = fir.box_addr %14 : (!fir.box<!fir.ptr<i32>>) -> !fir.ptr<i32>
+ %16 = fir.convert %15 : (!fir.ptr<i32>) -> i64
+ fir.call @foo(%16) : (i64) -> ()
+ acc.terminator
+ }
+ return
+}
+
More information about the flang-commits
mailing list