[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
Fri Jul 24 06:53:03 PDT 2026
https://github.com/sbryngelson updated https://github.com/llvm/llvm-project/pull/211543
>From 01ed42af42ce201e33830f53d7694f5823b9f95b Mon Sep 17 00:00:00 2001
From: Spencer Bryngelson <sbryngelson at gmail.com>
Date: Thu, 23 Jul 2026 10:15:28 -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 entry-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: it may have several
blocks, and arbitrary operations may compute an assignment's operands. Rather
than reason about provenance in general, match only the canonical copy-in that
lowering emits, requiring the assignment to be in the entry block, its LHS to
be exactly the clone argument, and its RHS to be either the original argument
or a single `fir.load` of it. Anything else, including an RHS computed by a
call that may return unrelated memory, 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.
A POINTER firstprivate is excluded from the fast path: it copies the pointer
association rather than the data, so the mold and clone descriptors designate
the same storage and the copy genuinely aliases. Such a copy region is
currently lowered with fir.store rather than hlfir.assign, so this pass never
sees it today, but the guard keeps the analysis correct if that changes.
Fixes #203890.
---
.../HLFIR/Transforms/InlineHLFIRAssign.cpp | 52 +++++++++-
flang/test/HLFIR/inline-hlfir-assign.fir | 97 +++++++++++++++++++
2 files changed, 148 insertions(+), 1 deletion(-)
diff --git a/flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRAssign.cpp b/flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRAssign.cpp
index 44affd514f613..d0a7a0d6635f9 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,55 @@ static llvm::cl::opt<bool> inlineAllocatableExprAssignFlag(
"hlfir.expr (e.g., from hlfir.elemental)"),
llvm::cl::init(false));
+/// Is \p assign an `omp.private` copy region copying the mold to the clone?
+///
+/// The copy region's entry-block arguments are the mold and the freshly
+/// allocated clone, which cannot overlap, so the copy needs no aliasing check.
+/// The body is unrestricted, so match only the shape lowering emits:
+///
+/// \code
+/// ^bb0(%mold, %clone):
+/// %0 = fir.load %mold // boxed privatizer only
+/// hlfir.assign %0 to %clone
+/// \endcode
+///
+/// A POINTER privatizer is excluded: firstprivate copies the association, so
+/// mold and clone share data and the copy aliases. It is lowered with fir.store
+/// today, not hlfir.assign, so the guard only covers a future codegen change.
+static bool isOmpPrivateCopyInAssign(hlfir::AssignOp assign) {
+ auto privateOp = llvm::dyn_cast_or_null<mlir::omp::PrivateClauseOp>(
+ assign->getParentRegion()->getParentOp());
+ if (!privateOp)
+ return false;
+ mlir::Region ©Region = privateOp.getCopyRegion();
+ if (assign->getParentRegion() != ©Region)
+ return false;
+
+ // Only the entry block: another block's arguments are not the mold/clone pair.
+ mlir::Block &entry = copyRegion.front();
+ if (assign->getBlock() != &entry || entry.getNumArguments() != 2)
+ return false;
+
+ mlir::Value mold = privateOp.getCopyMoldArg();
+ if (assign.getLhs() != privateOp.getCopyPrivateArg())
+ return false;
+
+ // A POINTER firstprivate copies the association, so the data aliases.
+ auto isPointer = [](mlir::Value v) -> bool {
+ auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(fir::unwrapRefType(v.getType()));
+ return boxTy && boxTy.isPointer();
+ };
+ if (isPointer(mold))
+ return false;
+
+ mlir::Value rhs = assign.getRhs();
+ if (rhs == mold)
+ return true;
+ // A boxed privatizer loads the mold's descriptor first.
+ auto load = rhs.getDefiningOp<fir::LoadOp>();
+ return load && load.getMemref() == mold;
+}
+
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 +154,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..17ebc55602320 100644
--- a/flang/test/HLFIR/inline-hlfir-assign.fir
+++ b/flang/test/HLFIR/inline-hlfir-assign.fir
@@ -496,3 +496,100 @@ func.func @_QPtest_scalar_ref_from_lhs(%arg0: !fir.ref<!fir.array<10xf32>>) {
// CHECK: }
// CHECK: return
// CHECK: }
+
+// The copy-in writes the clone, which is distinct from the mold, 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: }
+
+// An RHS from a call may designate unrelated memory, so it keeps its aliasing
+// check.
+func.func private @get_box(!fir.ref<!fir.box<!fir.array<2xf64>>>) -> !fir.box<!fir.array<2xf64>>
+omp.private {type = firstprivate} @_QFtestEf_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.call @get_box(%arg0) : (!fir.ref<!fir.box<!fir.array<2xf64>>>) -> !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} @_QFtestEf_firstprivate : !fir.box<!fir.array<2xf64>> copy {
+// CHECK: ^bb0(%[[ORIG4:.*]]: !fir.ref<!fir.box<!fir.array<2xf64>>>, %[[CLONE4:.*]]: !fir.ref<!fir.box<!fir.array<2xf64>>>):
+// CHECK: %[[RHS4:.*]] = fir.call @get_box(%[[ORIG4]])
+// CHECK: hlfir.assign %[[RHS4]] to %[[CLONE4]] : !fir.box<!fir.array<2xf64>>, !fir.ref<!fir.box<!fir.array<2xf64>>>
+// CHECK: omp.yield(%[[CLONE4]] : !fir.ref<!fir.box<!fir.array<2xf64>>>)
+// CHECK: }
+
+// Argument 1 of a block other than the entry block is not the clone, so an
+// assignment writing it keeps its aliasing check.
+omp.private {type = firstprivate} @_QFtestEg_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 %arg0 : !fir.ref<!fir.box<!fir.array<2xf64>>>
+ cf.br ^bb1(%arg0, %arg0 : !fir.ref<!fir.box<!fir.array<2xf64>>>, !fir.ref<!fir.box<!fir.array<2xf64>>>)
+^bb1(%a: !fir.ref<!fir.box<!fir.array<2xf64>>>, %b: !fir.ref<!fir.box<!fir.array<2xf64>>>):
+ hlfir.assign %0 to %b : !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} @_QFtestEg_firstprivate : !fir.box<!fir.array<2xf64>> copy {
+// CHECK: ^bb0(%[[ORIG5:.*]]: !fir.ref<!fir.box<!fir.array<2xf64>>>, %[[CLONE5:.*]]: !fir.ref<!fir.box<!fir.array<2xf64>>>):
+// CHECK: %[[RHS5:.*]] = fir.load %[[ORIG5]]
+// CHECK: cf.br ^bb1
+// CHECK: ^bb1(%[[A:.*]]: !fir.ref<!fir.box<!fir.array<2xf64>>>, %[[B:.*]]: !fir.ref<!fir.box<!fir.array<2xf64>>>):
+// CHECK: hlfir.assign %[[RHS5]] to %[[B]] : !fir.box<!fir.array<2xf64>>, !fir.ref<!fir.box<!fir.array<2xf64>>>
+// CHECK: }
+
+// A POINTER firstprivate shares data between mold and clone, so the copy
+// aliases and keeps its check despite matching the shape. Guards a future
+// codegen change; lowering uses fir.store, not hlfir.assign, today.
+omp.private {type = firstprivate} @_QFtestEp_firstprivate : !fir.box<!fir.ptr<!fir.array<?xf64>>> copy {
+^bb0(%arg0: !fir.ref<!fir.box<!fir.ptr<!fir.array<?xf64>>>>, %arg1: !fir.ref<!fir.box<!fir.ptr<!fir.array<?xf64>>>>):
+ %0 = fir.load %arg0 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xf64>>>>
+ hlfir.assign %0 to %arg1 : !fir.box<!fir.ptr<!fir.array<?xf64>>>, !fir.ref<!fir.box<!fir.ptr<!fir.array<?xf64>>>>
+ omp.yield(%arg1 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xf64>>>>)
+}
+// CHECK-LABEL: omp.private {type = firstprivate} @_QFtestEp_firstprivate : !fir.box<!fir.ptr<!fir.array<?xf64>>> copy {
+// CHECK: ^bb0(%[[MOLD:.*]]: !fir.ref<!fir.box<!fir.ptr<!fir.array<?xf64>>>>, %[[CLONE:.*]]: !fir.ref<!fir.box<!fir.ptr<!fir.array<?xf64>>>>):
+// CHECK: %[[RHS:.*]] = fir.load %[[MOLD]]
+// CHECK: hlfir.assign %[[RHS]] to %[[CLONE]] : !fir.box<!fir.ptr<!fir.array<?xf64>>>, !fir.ref<!fir.box<!fir.ptr<!fir.array<?xf64>>>>
+// CHECK: omp.yield(%[[CLONE]] : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xf64>>>>)
+// CHECK: }
More information about the flang-commits
mailing list