[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:39:56 PDT 2026
https://github.com/sbryngelson updated https://github.com/llvm/llvm-project/pull/211543
>From 42d5e475f65e252381f4dbdaa63fe61fc6159e8c 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 | 57 +++++++++-
flang/test/HLFIR/inline-hlfir-assign.fir | 101 ++++++++++++++++++
2 files changed, 157 insertions(+), 1 deletion(-)
diff --git a/flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRAssign.cpp b/flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRAssign.cpp
index 44affd514f613..a92cca20e0270 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,60 @@ static llvm::cl::opt<bool> inlineAllocatableExprAssignFlag(
"hlfir.expr (e.g., from hlfir.elemental)"),
llvm::cl::init(false));
+/// Is \p assign the copy-in of an `omp.private` copy region, writing the
+/// privatized clone from the mold?
+///
+/// The copy region's entry-block arguments are the mold (original variable) and
+/// the freshly allocated clone, which cannot overlap, so the copy needs no
+/// aliasing check. The region body is otherwise unrestricted, so rather than
+/// reason about provenance this matches only the shape lowering emits:
+///
+/// \code
+/// ^bb0(%mold, %clone):
+/// %0 = fir.load %mold // boxed privatizer only
+/// hlfir.assign %0 to %clone
+/// \endcode
+///
+/// The assign must be in the entry block; anything else keeps the aliasing
+/// check. A POINTER privatizer is also excluded: firstprivate copies the
+/// association, so the mold and clone descriptors share data and the copy
+/// aliases. That copy region is lowered with fir.store today, not hlfir.assign,
+/// so this pass never sees it; the guard 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: the arguments of any other block are unrelated to
+ // the mold/clone pair, even where they happen to have the same numbers.
+ 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 +159,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..cf4852924caec 100644
--- a/flang/test/HLFIR/inline-hlfir-assign.fir
+++ b/flang/test/HLFIR/inline-hlfir-assign.fir
@@ -496,3 +496,104 @@ 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: }
+
+// An RHS computed by a call is not known to designate the original: the call
+// may return unrelated memory reached through global state. 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 copies the association, so mold and clone share data:
+// the copy aliases and keeps its aliasing check despite matching the canonical
+// shape. Lowering uses fir.store today, not hlfir.assign, so this does not
+// arise yet; the test guards a future codegen change.
+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