[flang-commits] [flang] [flang][OpenMP] Inline the firstprivate array copy instead of calling Assign() (PR #211543)
Spencer Bryngelson via flang-commits
flang-commits at lists.llvm.org
Thu Jul 23 07:22:44 PDT 2026
https://github.com/sbryngelson updated https://github.com/llvm/llvm-project/pull/211543
>From 97d61c5d6ab0937175c27d87907b9e6acebd933d Mon Sep 17 00:00:00 2001
From: Spencer Bryngelson <sbryngelson at gmail.com>
Date: Thu, 23 Jul 2026 07:58:52 -0500
Subject: [PATCH] [flang][OpenMP] Inline the firstprivate array copy instead of
calling Assign()
The copy region of an `omp.private` firstprivate privatizer holds an
`hlfir.assign` from the original host variable to the privatized clone.
`InlineHLFIRAssign` declines to expand it because both sides are block
arguments, so alias analysis cannot prove the two do not overlap, and the
assignment falls back to a `_FortranAAssign` runtime call.
That call is unusable in device code. Compiling a `target` region with a
`firstprivate` fixed-size array for `amdgcn-amd-amdhsa` fails to link:
ld.lld: error: undefined symbol: _FortranAAssign
The operation already defines the two copy-region block arguments to be the
original host variable and the memory allocated for the clone, so they cannot
overlap. Skip the aliasing check for an assignment that writes the clone from
the original, and the existing inlining produces a plain element loop.
The body of a copy region is otherwise unrestricted, so this requires both that
the assignment writes the clone and that its RHS is rooted in the original
block argument. An RHS that reaches the clone, an unrecognised value, or an
operation carrying regions keeps the usual aliasing check.
For a one-element `real(8)` array on gfx90a this removes both
`_FortranAAssign` calls from the device IR, drops the device compile's
ScratchSize from 208 to 112 bytes/lane, and the program links and returns the
correct result on an MI210.
Fixes #203890.
---
.../HLFIR/Transforms/InlineHLFIRAssign.cpp | 60 ++++++++++++++++++-
flang/test/HLFIR/inline-hlfir-assign.fir | 49 +++++++++++++++
2 files changed, 108 insertions(+), 1 deletion(-)
diff --git a/flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRAssign.cpp b/flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRAssign.cpp
index 44affd514f613..7947b7889f3dd 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRAssign.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRAssign.cpp
@@ -19,6 +19,7 @@
#include "flang/Optimizer/HLFIR/HLFIROps.h"
#include "flang/Optimizer/HLFIR/Passes.h"
#include "flang/Optimizer/OpenMP/Passes.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
@@ -39,6 +40,63 @@ static llvm::cl::opt<bool> inlineAllocatableExprAssignFlag(
"hlfir.expr (e.g., from hlfir.elemental)"),
llvm::cl::init(false));
+/// Does \p v derive only from \p origArg, the copy region's original-variable
+/// block argument? Returns false if anything else is reached, including
+/// \p cloneArg or a value with no defining operation, so an unrecognised RHS
+/// is treated as potentially overlapping.
+static bool isRootedInOriginalArg(mlir::Value v, mlir::BlockArgument origArg,
+ mlir::BlockArgument cloneArg) {
+ llvm::SmallVector<mlir::Value> worklist{v};
+ llvm::SmallPtrSet<mlir::Value, 8> seen;
+ bool reachedOrig = false;
+ while (!worklist.empty()) {
+ mlir::Value cur = worklist.pop_back_val();
+ if (!seen.insert(cur).second)
+ continue;
+ if (cur == cloneArg)
+ return false;
+ if (cur == origArg) {
+ reachedOrig = true;
+ continue;
+ }
+ mlir::Operation *def = cur.getDefiningOp();
+ if (!def)
+ return false;
+ // An operation carrying regions may capture values that are not among its
+ // operands, so do not try to reason about it.
+ if (def->getNumRegions() != 0)
+ return false;
+ worklist.append(def->getOperands().begin(), def->getOperands().end());
+ }
+ return reachedOrig;
+}
+
+/// Is \p assign the copy-in of an `omp.private` copy region, writing the
+/// privatized clone from the original host variable? The operation defines its
+/// first copy-region block argument to be the original and the second to be the
+/// memory allocated for the clone, so the two cannot overlap and the copy can
+/// be inlined without an aliasing check.
+///
+/// The body of a copy region is otherwise unrestricted, so this requires both
+/// that the assignment writes the clone and that its RHS is rooted in the
+/// original. Anything else keeps the usual aliasing check.
+static bool isOmpPrivateCopyInAssign(hlfir::AssignOp assign) {
+ mlir::Region *region = assign->getParentRegion();
+ auto privateOp =
+ llvm::dyn_cast_or_null<mlir::omp::PrivateClauseOp>(region->getParentOp());
+ if (!privateOp || region != &privateOp.getCopyRegion())
+ return false;
+ auto blockArg = llvm::dyn_cast<mlir::BlockArgument>(assign.getLhs());
+ if (!blockArg || blockArg.getOwner()->getParent() != region ||
+ blockArg.getArgNumber() != 1)
+ return false;
+ mlir::Block *entry = ®ion->front();
+ if (entry->getNumArguments() != 2)
+ return false;
+ return isRootedInOriginalArg(assign.getRhs(), entry->getArgument(0),
+ entry->getArgument(1));
+}
+
namespace {
/// Expand hlfir.assign of array RHS to array LHS into a loop nest
/// of element-by-element assignments. Also handles scalar RHS broadcast
@@ -104,7 +162,7 @@ class InlineHLFIRAssignConversion
bool rhsNeedsTemporary = false;
if (rhs.isArray() && !mlir::isa<hlfir::ExprType>(rhs.getType()) &&
- !assign.getTemporaryLhs()) {
+ !assign.getTemporaryLhs() && !isOmpPrivateCopyInAssign(assign)) {
fir::AliasAnalysis aliasAnalysis;
mlir::AliasResult aliasRes = aliasAnalysis.alias(lhs, rhs);
if (!aliasRes.isNo()) {
diff --git a/flang/test/HLFIR/inline-hlfir-assign.fir b/flang/test/HLFIR/inline-hlfir-assign.fir
index 6cf6fcedb0af1..3c7a5406114c3 100644
--- a/flang/test/HLFIR/inline-hlfir-assign.fir
+++ b/flang/test/HLFIR/inline-hlfir-assign.fir
@@ -496,3 +496,52 @@ func.func @_QPtest_scalar_ref_from_lhs(%arg0: !fir.ref<!fir.array<10xf32>>) {
// CHECK: }
// CHECK: return
// CHECK: }
+
+// The copy-in of an omp.private copy region writes the privatized clone,
+// which the operation defines to be distinct from the original host variable,
+// so it inlines even though alias analysis cannot prove it from the block
+// arguments alone.
+omp.private {type = firstprivate} @_QFtestEc_firstprivate : !fir.box<!fir.array<1xf64>> copy {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.array<1xf64>>>, %arg1: !fir.ref<!fir.box<!fir.array<1xf64>>>):
+ %0 = fir.load %arg0 : !fir.ref<!fir.box<!fir.array<1xf64>>>
+ hlfir.assign %0 to %arg1 : !fir.box<!fir.array<1xf64>>, !fir.ref<!fir.box<!fir.array<1xf64>>>
+ omp.yield(%arg1 : !fir.ref<!fir.box<!fir.array<1xf64>>>)
+}
+// CHECK-LABEL: omp.private {type = firstprivate} @_QFtestEc_firstprivate : !fir.box<!fir.array<1xf64>> copy {
+// CHECK: ^bb0(%[[ORIG:.*]]: !fir.ref<!fir.box<!fir.array<1xf64>>>, %[[CLONE:.*]]: !fir.ref<!fir.box<!fir.array<1xf64>>>):
+// CHECK: %[[RHS:.*]] = fir.load %[[ORIG]] : !fir.ref<!fir.box<!fir.array<1xf64>>>
+// CHECK: %[[LHS:.*]] = fir.load %[[CLONE]] : !fir.ref<!fir.box<!fir.array<1xf64>>>
+// CHECK: fir.do_loop
+// CHECK-NOT: hlfir.assign %[[RHS]] to %[[CLONE]]
+// CHECK: omp.yield(%[[CLONE]] : !fir.ref<!fir.box<!fir.array<1xf64>>>)
+// CHECK: }
+
+// An assignment in the copy region that does not write the clone keeps its
+// aliasing check.
+omp.private {type = firstprivate} @_QFtestEd_firstprivate : !fir.box<!fir.array<1xf64>> copy {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.array<1xf64>>>, %arg1: !fir.ref<!fir.box<!fir.array<1xf64>>>):
+ %0 = fir.load %arg1 : !fir.ref<!fir.box<!fir.array<1xf64>>>
+ hlfir.assign %0 to %arg0 : !fir.box<!fir.array<1xf64>>, !fir.ref<!fir.box<!fir.array<1xf64>>>
+ omp.yield(%arg1 : !fir.ref<!fir.box<!fir.array<1xf64>>>)
+}
+// CHECK-LABEL: omp.private {type = firstprivate} @_QFtestEd_firstprivate : !fir.box<!fir.array<1xf64>> copy {
+// CHECK: ^bb0(%[[ORIG2:.*]]: !fir.ref<!fir.box<!fir.array<1xf64>>>, %[[CLONE2:.*]]: !fir.ref<!fir.box<!fir.array<1xf64>>>):
+// CHECK: %[[RHS2:.*]] = fir.load %[[CLONE2]] : !fir.ref<!fir.box<!fir.array<1xf64>>>
+// CHECK: hlfir.assign %[[RHS2]] to %[[ORIG2]] : !fir.box<!fir.array<1xf64>>, !fir.ref<!fir.box<!fir.array<1xf64>>>
+// CHECK: omp.yield(%[[CLONE2]] : !fir.ref<!fir.box<!fir.array<1xf64>>>)
+// CHECK: }
+
+// A copy-region assignment whose RHS is rooted in the clone rather than the
+// original can overlap the LHS, so it keeps its aliasing check.
+omp.private {type = firstprivate} @_QFtestEe_firstprivate : !fir.box<!fir.array<2xf64>> copy {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.array<2xf64>>>, %arg1: !fir.ref<!fir.box<!fir.array<2xf64>>>):
+ %0 = fir.load %arg1 : !fir.ref<!fir.box<!fir.array<2xf64>>>
+ hlfir.assign %0 to %arg1 : !fir.box<!fir.array<2xf64>>, !fir.ref<!fir.box<!fir.array<2xf64>>>
+ omp.yield(%arg1 : !fir.ref<!fir.box<!fir.array<2xf64>>>)
+}
+// CHECK-LABEL: omp.private {type = firstprivate} @_QFtestEe_firstprivate : !fir.box<!fir.array<2xf64>> copy {
+// CHECK: ^bb0(%[[ORIG3:.*]]: !fir.ref<!fir.box<!fir.array<2xf64>>>, %[[CLONE3:.*]]: !fir.ref<!fir.box<!fir.array<2xf64>>>):
+// CHECK: %[[RHS3:.*]] = fir.load %[[CLONE3]] : !fir.ref<!fir.box<!fir.array<2xf64>>>
+// CHECK: hlfir.assign %[[RHS3]] to %[[CLONE3]] : !fir.box<!fir.array<2xf64>>, !fir.ref<!fir.box<!fir.array<2xf64>>>
+// CHECK: omp.yield(%[[CLONE3]] : !fir.ref<!fir.box<!fir.array<2xf64>>>)
+// CHECK: }
More information about the flang-commits
mailing list