[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 08:15:55 PDT 2026
https://github.com/sbryngelson updated https://github.com/llvm/llvm-project/pull/211543
>From 630f75a70f185ace0cc71f0c1747eaf442a5d7a5 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.
Fixes #203890.
---
.../HLFIR/Transforms/InlineHLFIRAssign.cpp | 52 +++++++++++-
flang/test/HLFIR/inline-hlfir-assign.fir | 84 +++++++++++++++++++
2 files changed, 135 insertions(+), 1 deletion(-)
diff --git a/flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRAssign.cpp b/flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRAssign.cpp
index 44affd514f613..c82e21a42dd44 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 the copy-in of an `omp.private` copy region, writing the
+/// privatized clone from the original host variable?
+///
+/// `omp.private` defines the copy region's first *entry-block* argument to be
+/// the original host variable and the second to be the memory allocated for the
+/// clone, so those two cannot overlap and the copy needs no aliasing check.
+///
+/// The body of a copy region is otherwise unrestricted: it may have several
+/// blocks, and arbitrary operations may compute the operands of an assignment.
+/// So rather than reason about provenance in general, this matches only the
+/// canonical shape lowering emits for the copy-in,
+///
+/// \code
+/// ^bb0(%orig, %clone):
+/// %0 = fir.load %orig // boxed privatizer only
+/// hlfir.assign %0 to %clone
+/// \endcode
+///
+/// and requires the assignment to be in the entry block itself. Anything else,
+/// including an RHS computed by a call that might return unrelated memory,
+/// keeps the usual aliasing check; being conservative here costs only the
+/// inlining.
+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 original/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;
+
+ if (assign.getLhs() != entry.getArgument(1))
+ return false;
+
+ mlir::Value orig = entry.getArgument(0);
+ mlir::Value rhs = assign.getRhs();
+ if (rhs == orig)
+ return true;
+ // A boxed privatizer loads the original's descriptor first.
+ auto load = rhs.getDefiningOp<fir::LoadOp>();
+ return load && load.getMemref() == orig;
+}
+
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..a594bc3e9f13e 100644
--- a/flang/test/HLFIR/inline-hlfir-assign.fir
+++ b/flang/test/HLFIR/inline-hlfir-assign.fir
@@ -496,3 +496,87 @@ 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: }
More information about the flang-commits
mailing list