[flang-commits] [flang] [mlir] [flang][OpenMP] Lower allocator-backed storage for allocate clauses (PR #211621)
Sairudra More via flang-commits
flang-commits at lists.llvm.org
Thu Jul 30 01:22:25 PDT 2026
https://github.com/Saieiei updated https://github.com/llvm/llvm-project/pull/211621
>From cf413190e2bd177803eb23ebdc6349e26bcabd97 Mon Sep 17 00:00:00 2001
From: Sairudra More <sairudra60 at gmail.com>
Date: Wed, 22 Jul 2026 02:35:15 -0500
Subject: [PATCH 1/2] [flang][OpenMP] Lower allocate clauses for scalar
parallel privates
Lower OpenMP `allocate` clauses on `parallel` constructs for fixed-size
intrinsic scalar `private` and `firstprivate` variables.
Flang records the private operand associated with each allocate item.
The OpenMP dialect verifies this mapping, and LLVM IR translation uses
the runtime-allocated pointer as the private storage and releases each
pointer with its corresponding allocator during region finalization.
This preserves one-time evaluation of allocator expressions and cleanup
on normal and cancellation paths. Unsupported types, device lowering,
other clause-bearing operations, and the `align` modifier remain
unsupported.
Tests cover source lowering, dialect verification and attribute
preservation, allocator identity, reverse-order cleanup, privatizer
deallocation, and cancellation.
Assisted-by: GitHub Copilot
---
flang/lib/Lower/OpenMP/ClauseProcessor.cpp | 6 +-
flang/lib/Lower/OpenMP/OpenMP.cpp | 59 +++++
.../Optimizer/OpenMP/LowerWorkdistribute.cpp | 32 +--
.../OpenMP/Todo/allocate-clause-allocator.f90 | 16 --
.../Todo/allocate-clause-unsupported.f90 | 65 ++++++
.../OpenMP/allocate-clause-allocator.f90 | 119 ++++++++++
.../Semantics/OpenMP/allocate-clause01.f90 | 16 ++
.../lower-workdistribute-fission-target.mlir | 19 +-
.../mlir/Dialect/OpenMP/OpenMPClauses.td | 10 +-
.../mlir/Dialect/OpenMP/OpenMPDialect.h | 4 +
.../Conversion/SCFToOpenMP/SCFToOpenMP.cpp | 1 +
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 194 +++++++++++++----
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 91 +++++++-
.../Dialect/OpenMP/host-op-filtering.mlir | 3 +-
mlir/test/Dialect/OpenMP/invalid.mlir | 82 +++++++
mlir/test/Dialect/OpenMP/ops.mlir | 35 ++-
.../Target/LLVMIR/openmp-allocate-clause.mlir | 203 ++++++++++++++++++
mlir/test/Target/LLVMIR/openmp-todo.mlir | 11 +-
18 files changed, 854 insertions(+), 112 deletions(-)
delete mode 100644 flang/test/Lower/OpenMP/Todo/allocate-clause-allocator.f90
create mode 100644 flang/test/Lower/OpenMP/Todo/allocate-clause-unsupported.f90
create mode 100644 flang/test/Lower/OpenMP/allocate-clause-allocator.f90
create mode 100644 mlir/test/Target/LLVMIR/openmp-allocate-clause.mlir
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index 5a554def351ea..f912c2af15d67 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -110,16 +110,14 @@ genAllocateClause(lower::AbstractConverter &converter,
if (std::get<std::optional<Allocate::AlignModifier>>(clause.t))
TODO(currentLocation, "OmpAllocateClause ALIGN modifier");
- // Check if allocate clause has allocator specified. If so, add it
- // to list of allocators, otherwise, add default allocator to
- // list of allocators.
+ // Use a null handle to select the binding task's default allocator.
using ComplexModifier = Allocate::AllocatorComplexModifier;
if (auto &mod = std::get<std::optional<ComplexModifier>>(clause.t)) {
mlir::Value operand = fir::getBase(converter.genExprValue(mod->v, stmtCtx));
allocatorOperands.append(objects.size(), operand);
} else {
mlir::Value operand = firOpBuilder.createIntegerConstant(
- currentLocation, firOpBuilder.getI32Type(), 1);
+ currentLocation, firOpBuilder.getI32Type(), mlir::omp::kNullAllocator);
allocatorOperands.append(objects.size(), operand);
}
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index b4fa431d42223..9fb70be2b515a 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -55,6 +55,8 @@
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Support/StateStack.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallSet.h"
@@ -3503,6 +3505,63 @@ genParallelOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
assert((!enableDelayedPrivatization || dsp) &&
"expected valid DataSharingProcessor");
+ if (!clauseOps.allocateVars.empty()) {
+ llvm::DenseMap<const semantics::Symbol *, int64_t> privateSlots;
+ for (auto [index, object] : llvm::enumerate(args.priv.objects)) {
+ const semantics::Symbol *symbol = object.sym();
+ if (!symbol)
+ fir::emitFatalError(loc, "private item without a semantic symbol");
+ const semantics::Symbol *ultimate = &symbol->GetUltimate();
+ if (!privateSlots.try_emplace(ultimate, index).second)
+ fir::emitFatalError(
+ loc, "symbol with multiple private storage slots on one construct");
+ }
+
+ llvm::DenseSet<const semantics::Symbol *> allocateSymbols;
+ for (const Clause &clause : item->clauses) {
+ if (clause.id != llvm::omp::Clause::OMPC_allocate)
+ continue;
+ const auto &allocate = std::get<clause::Allocate>(clause.u);
+ const auto &objects = std::get<ObjectList>(allocate.t);
+ for (const Object &object : objects) {
+ const semantics::Symbol *symbol = object.sym();
+ if (!symbol)
+ fir::emitFatalError(loc,
+ "ALLOCATE clause item without a semantic symbol");
+ const semantics::Symbol *ultimate = &symbol->GetUltimate();
+ if (!allocateSymbols.insert(ultimate).second)
+ TODO(loc, "ALLOCATE clause item appears more than once");
+
+ auto privateSlot = privateSlots.find(ultimate);
+ if (privateSlot == privateSlots.end())
+ fir::emitFatalError(
+ loc, "ALLOCATE clause item without private storage slot");
+
+ auto type = evaluate::DynamicType::From(*ultimate);
+ bool supportedDataSharing =
+ symbol->test(semantics::Symbol::Flag::OmpPrivate) ||
+ symbol->test(semantics::Symbol::Flag::OmpFirstPrivate);
+ bool supportedType =
+ ultimate->Rank() == 0 &&
+ !semantics::IsAllocatableOrPointer(*ultimate) && type &&
+ type->category() != common::TypeCategory::Derived &&
+ !type->RequiresDescriptor() &&
+ !type->HasDeferredOrAssumedTypeParameter();
+ if (!supportedDataSharing || !supportedType)
+ TODO(loc,
+ "ALLOCATE clause currently supports only fixed-size intrinsic "
+ "scalar PRIVATE or FIRSTPRIVATE items");
+
+ clauseOps.allocatePrivateIndices.push_back(privateSlot->second);
+ }
+ }
+
+ if (clauseOps.allocatePrivateIndices.size() !=
+ clauseOps.allocateVars.size())
+ fir::emitFatalError(loc,
+ "incomplete ALLOCATE clause private storage mapping");
+ }
+
OpWithBodyGenInfo genInfo =
OpWithBodyGenInfo(converter, symTable, semaCtx, loc, eval,
llvm::omp::Directive::OMPD_parallel)
diff --git a/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp b/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
index 3242863cd4023..5c42632d0c5fb 100644
--- a/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
+++ b/flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp
@@ -757,10 +757,10 @@ FailureOr<omp::TargetOp> splitTargetData(omp::TargetOp targetOp,
// Create the inner target op
auto newTargetOp = omp::TargetOp::create(
rewriter, targetOp.getLoc(), targetOp.getAllocateVars(),
- targetOp.getAllocatorVars(), targetOp.getDependKindsAttr(),
- targetOp.getDependVars(), targetOp.getDependIteratedKindsAttr(),
- targetOp.getDependIterated(), targetOp.getDevice(),
- targetOp.getDynGroupprivateAccessGroupAttr(),
+ targetOp.getAllocatorVars(), targetOp.getAllocatePrivateIndicesAttr(),
+ targetOp.getDependKindsAttr(), targetOp.getDependVars(),
+ targetOp.getDependIteratedKindsAttr(), targetOp.getDependIterated(),
+ targetOp.getDevice(), targetOp.getDynGroupprivateAccessGroupAttr(),
targetOp.getDynGroupprivateFallbackAttr(),
targetOp.getDynGroupprivateSize(), targetOp.getHasDeviceAddrVars(),
targetOp.getHostEvalVars(), targetOp.getIfExpr(),
@@ -1482,10 +1482,10 @@ genPreTargetOp(omp::TargetOp targetOp, SmallVector<Value> &preMapOperands,
// update the hostEvalVars of preTargetOp
omp::TargetOp preTargetOp = omp::TargetOp::create(
rewriter, targetOp.getLoc(), targetOp.getAllocateVars(),
- targetOp.getAllocatorVars(), targetOp.getDependKindsAttr(),
- targetOp.getDependVars(), targetOp.getDependIteratedKindsAttr(),
- targetOp.getDependIterated(), targetOp.getDevice(),
- targetOp.getDynGroupprivateAccessGroupAttr(),
+ targetOp.getAllocatorVars(), targetOp.getAllocatePrivateIndicesAttr(),
+ targetOp.getDependKindsAttr(), targetOp.getDependVars(),
+ targetOp.getDependIteratedKindsAttr(), targetOp.getDependIterated(),
+ targetOp.getDevice(), targetOp.getDynGroupprivateAccessGroupAttr(),
targetOp.getDynGroupprivateFallbackAttr(),
targetOp.getDynGroupprivateSize(), targetOp.getHasDeviceAddrVars(),
preHostEvalVars, targetOp.getIfExpr(), targetOp.getInReductionVars(),
@@ -1576,10 +1576,10 @@ genIsolatedTargetOp(omp::TargetOp targetOp, SmallVector<Value> &postMapOperands,
// Create the isolated target op
omp::TargetOp isolatedTargetOp = omp::TargetOp::create(
rewriter, targetOp.getLoc(), targetOp.getAllocateVars(),
- targetOp.getAllocatorVars(), targetOp.getDependKindsAttr(),
- targetOp.getDependVars(), targetOp.getDependIteratedKindsAttr(),
- targetOp.getDependIterated(), targetOp.getDevice(),
- targetOp.getDynGroupprivateAccessGroupAttr(),
+ targetOp.getAllocatorVars(), targetOp.getAllocatePrivateIndicesAttr(),
+ targetOp.getDependKindsAttr(), targetOp.getDependVars(),
+ targetOp.getDependIteratedKindsAttr(), targetOp.getDependIterated(),
+ targetOp.getDevice(), targetOp.getDynGroupprivateAccessGroupAttr(),
targetOp.getDynGroupprivateFallbackAttr(),
targetOp.getDynGroupprivateSize(), targetOp.getHasDeviceAddrVars(),
isolatedHostEvalVars, targetOp.getIfExpr(), targetOp.getInReductionVars(),
@@ -1662,10 +1662,10 @@ static omp::TargetOp genPostTargetOp(omp::TargetOp targetOp,
// Create the post target op
omp::TargetOp postTargetOp = omp::TargetOp::create(
rewriter, targetOp.getLoc(), targetOp.getAllocateVars(),
- targetOp.getAllocatorVars(), targetOp.getDependKindsAttr(),
- targetOp.getDependVars(), targetOp.getDependIteratedKindsAttr(),
- targetOp.getDependIterated(), targetOp.getDevice(),
- targetOp.getDynGroupprivateAccessGroupAttr(),
+ targetOp.getAllocatorVars(), targetOp.getAllocatePrivateIndicesAttr(),
+ targetOp.getDependKindsAttr(), targetOp.getDependVars(),
+ targetOp.getDependIteratedKindsAttr(), targetOp.getDependIterated(),
+ targetOp.getDevice(), targetOp.getDynGroupprivateAccessGroupAttr(),
targetOp.getDynGroupprivateFallbackAttr(),
targetOp.getDynGroupprivateSize(), targetOp.getHasDeviceAddrVars(),
postHostEvalVars, targetOp.getIfExpr(), targetOp.getInReductionVars(),
diff --git a/flang/test/Lower/OpenMP/Todo/allocate-clause-allocator.f90 b/flang/test/Lower/OpenMP/Todo/allocate-clause-allocator.f90
deleted file mode 100644
index 45b3aa8ee77f9..0000000000000
--- a/flang/test/Lower/OpenMP/Todo/allocate-clause-allocator.f90
+++ /dev/null
@@ -1,16 +0,0 @@
-! REQUIRES: openmp_runtime
-! RUN: %not_todo_cmd %flang_fc1 -emit-llvm %openmp_flags -fopenmp-version=51 -o - %s 2>&1 | FileCheck %s
-
-! CHECK: not yet implemented: Unhandled clause allocate in omp.parallel
-! CHECK: LLVM Translation failed for operation: omp.parallel
-program p
- use omp_lib
- integer :: x
- integer :: a
- integer :: i
- !$omp parallel private(x) allocate(allocator(omp_default_mem_alloc): x)
- do i=1,10
- a = a + i
- end do
- !$omp end parallel
-end program p
diff --git a/flang/test/Lower/OpenMP/Todo/allocate-clause-unsupported.f90 b/flang/test/Lower/OpenMP/Todo/allocate-clause-unsupported.f90
new file mode 100644
index 0000000000000..82b82fa6c33f1
--- /dev/null
+++ b/flang/test/Lower/OpenMP/Todo/allocate-clause-unsupported.f90
@@ -0,0 +1,65 @@
+! RUN: split-file %s %t
+! RUN: %not_todo_cmd %flang_fc1 -emit-hlfir %openmp_flags -fopenmp-version=51 -o - %t/duplicate.f90 2>&1 | FileCheck %s --check-prefix=DUPLICATE
+! RUN: %not_todo_cmd %flang_fc1 -emit-hlfir %openmp_flags -fopenmp-version=51 -o - %t/array.f90 2>&1 | FileCheck %s --check-prefix=ARRAY
+! RUN: %not_todo_cmd %flang_fc1 -emit-hlfir %openmp_flags -fopenmp-version=51 -o - %t/derived.f90 2>&1 | FileCheck %s --check-prefix=DERIVED
+! RUN: %not_todo_cmd %flang_fc1 -emit-hlfir %openmp_flags -fopenmp-version=51 -o - %t/pointer.f90 2>&1 | FileCheck %s --check-prefix=POINTER
+! RUN: %not_todo_cmd %flang_fc1 -emit-hlfir %openmp_flags -fopenmp-version=51 -o - %t/allocatable.f90 2>&1 | FileCheck %s --check-prefix=ALLOCATABLE
+! RUN: %not_todo_cmd %flang_fc1 -emit-hlfir %openmp_flags -fopenmp-version=51 -o - %t/assumed-length.f90 2>&1 | FileCheck %s --check-prefix=ASSUMED-LENGTH
+
+! DUPLICATE: not yet implemented: ALLOCATE clause item appears more than once
+! ARRAY: not yet implemented: ALLOCATE clause currently supports only fixed-size intrinsic scalar PRIVATE or FIRSTPRIVATE items
+! DERIVED: not yet implemented: ALLOCATE clause currently supports only fixed-size intrinsic scalar PRIVATE or FIRSTPRIVATE items
+! POINTER: not yet implemented: ALLOCATE clause currently supports only fixed-size intrinsic scalar PRIVATE or FIRSTPRIVATE items
+! ALLOCATABLE: not yet implemented: ALLOCATE clause currently supports only fixed-size intrinsic scalar PRIVATE or FIRSTPRIVATE items
+! ASSUMED-LENGTH: not yet implemented: ALLOCATE clause currently supports only fixed-size intrinsic scalar PRIVATE or FIRSTPRIVATE items
+
+!--- duplicate.f90
+subroutine duplicate(x)
+ integer :: x
+ !$omp parallel private(x) allocate(x) allocate(x)
+ x = 1
+ !$omp end parallel
+end subroutine
+
+!--- pointer.f90
+subroutine pointer(x)
+ integer, pointer :: x
+ !$omp parallel private(x) allocate(x)
+ x = 1
+ !$omp end parallel
+end subroutine
+
+!--- allocatable.f90
+subroutine allocatable(x)
+ integer, allocatable :: x
+ !$omp parallel private(x) allocate(x)
+ x = 1
+ !$omp end parallel
+end subroutine
+
+!--- assumed-length.f90
+subroutine assumed_length(x)
+ character(*) :: x
+ !$omp parallel private(x) allocate(x)
+ x = "test"
+ !$omp end parallel
+end subroutine
+
+!--- array.f90
+subroutine array(x)
+ integer :: x(4)
+ !$omp parallel private(x) allocate(x)
+ x = 1
+ !$omp end parallel
+end subroutine
+
+!--- derived.f90
+subroutine derived()
+ type t
+ integer :: value
+ end type
+ type(t) :: x
+ !$omp parallel private(x) allocate(x)
+ x%value = 1
+ !$omp end parallel
+end subroutine
diff --git a/flang/test/Lower/OpenMP/allocate-clause-allocator.f90 b/flang/test/Lower/OpenMP/allocate-clause-allocator.f90
new file mode 100644
index 0000000000000..fd64f4126f65e
--- /dev/null
+++ b/flang/test/Lower/OpenMP/allocate-clause-allocator.f90
@@ -0,0 +1,119 @@
+! REQUIRES: openmp_runtime
+! RUN: %flang_fc1 -emit-hlfir %openmp_flags -fopenmp-version=51 -o - %s 2>&1 | FileCheck %s --check-prefix=HLFIR
+
+subroutine allocator_omitted(x, y)
+ integer :: x, y
+ !$omp parallel private(x, y) allocate(x)
+ x = 1
+ y = 2
+ !$omp end parallel
+end subroutine
+
+! HLFIR-LABEL: func.func @_QPallocator_omitted
+! HLFIR: %[[NULL_ALLOC:.*]] = arith.constant 0 : i32
+! HLFIR: omp.parallel allocate(%[[NULL_ALLOC]] : i32 -> %[[X:.*]]#0 : !fir.ref<i32>)
+! HLFIR-SAME: private({{.*}} %[[X]]#0 -> %[[X_PRIVATE:.*]], {{.*}} -> %[[Y_PRIVATE:.*]] : !fir.ref<i32>, !fir.ref<i32>) {
+! HLFIR: } {allocate_private_indices = array<i64: 0>}
+
+subroutine allocator_explicit(x)
+ use omp_lib
+ integer :: x
+
+ !$omp parallel private(x) allocate(allocator(omp_null_allocator): x)
+ x = 1
+ !$omp end parallel
+
+ !$omp parallel private(x) allocate(allocator(omp_default_mem_alloc): x)
+ x = 2
+ !$omp end parallel
+end subroutine
+
+! HLFIR-LABEL: func.func @_QPallocator_explicit
+! HLFIR: omp.parallel allocate(%c0_i64 : i64 -> %[[X:.*]]#0 : !fir.ref<i32>)
+! HLFIR-SAME: private({{.*}} %[[X]]#0 -> {{.*}} : !fir.ref<i32>) {
+! HLFIR: omp.parallel allocate(%c1_i64 : i64 -> %[[X2:.*]]#0 : !fir.ref<i32>)
+! HLFIR-SAME: private({{.*}} %[[X2]]#0 -> {{.*}} : !fir.ref<i32>) {
+
+subroutine allocator_dynamic(x, allocator)
+ use omp_lib, only : omp_allocator_handle_kind
+ integer :: x
+ integer(kind=omp_allocator_handle_kind), intent(in) :: allocator
+ !$omp parallel firstprivate(x) allocate(allocator(allocator): x)
+ x = x + 1
+ !$omp end parallel
+end subroutine
+
+! HLFIR-LABEL: func.func @_QPallocator_dynamic
+! HLFIR: %[[ALLOCATOR:.*]] = fir.load
+! HLFIR: omp.parallel allocate(%[[ALLOCATOR]] : i64 -> %[[X:.*]]#0 : !fir.ref<i32>)
+! HLFIR-SAME: private({{.*}} %[[X]]#0 -> %[[X_PRIVATE:.*]] : !fir.ref<i32>) {
+! HLFIR: } {allocate_private_indices = array<i64: 0>}
+
+function allocator_value() result(allocator)
+ use omp_lib, only : omp_allocator_handle_kind, omp_default_mem_alloc
+ integer(kind=omp_allocator_handle_kind) :: allocator
+ allocator = omp_default_mem_alloc
+end function
+
+subroutine allocator_expression(x, y)
+ use omp_lib, only : omp_allocator_handle_kind
+ integer :: x, y
+ integer(kind=omp_allocator_handle_kind), external :: allocator_value
+ !$omp parallel private(x, y) allocate(allocator(allocator_value()): x, y)
+ x = 1
+ y = 2
+ !$omp end parallel
+end subroutine
+
+! HLFIR-LABEL: func.func @_QPallocator_expression
+! HLFIR-COUNT-1: fir.call @_QPallocator_value()
+! HLFIR: omp.parallel allocate(%[[ALLOCATOR:.*]] : i64 -> %{{.*}}#0 : !fir.ref<i32>, %[[ALLOCATOR]] : i64 -> %{{.*}}#0 : !fir.ref<i32>)
+! HLFIR: } {allocate_private_indices = array<i64: 0, 1>}
+
+subroutine allocator_host_associated()
+ integer :: x
+contains
+ subroutine inner()
+ !$omp parallel private(x) allocate(x)
+ x = 1
+ !$omp end parallel
+ end subroutine
+end subroutine
+
+! HLFIR-LABEL: func.func private @_QFallocator_host_associatedPinner
+! HLFIR: omp.parallel allocate({{.*}} -> %[[X:.*]]#0 : !fir.ref<i32>)
+! HLFIR-SAME: private({{.*}} %[[X]]#0 -> {{.*}} : !fir.ref<i32>) {
+! HLFIR: } {allocate_private_indices = array<i64: 0>}
+
+subroutine allocator_order(x, y, z)
+ use omp_lib
+ integer :: x, y, z
+ !$omp parallel private(x, y, z) &
+ !$omp& allocate(allocator(omp_high_bw_mem_alloc): y, x) &
+ !$omp& allocate(allocator(omp_low_lat_mem_alloc): z)
+ x = 1
+ y = 2
+ z = 3
+ !$omp end parallel
+end subroutine
+
+! HLFIR-LABEL: func.func @_QPallocator_order
+! HLFIR: omp.parallel allocate(
+! HLFIR-SAME: private(
+! HLFIR: } {allocate_private_indices = array<i64: 1, 0, 2>}
+
+subroutine allocator_scalar_types(r, c, l, s)
+ real :: r
+ complex :: c
+ logical :: l
+ character(4) :: s
+ !$omp parallel private(r, c, l, s) allocate(r, c, l, s)
+ r = 1.0
+ c = (1.0, 2.0)
+ l = .true.
+ s = "test"
+ !$omp end parallel
+end subroutine
+
+! HLFIR-LABEL: func.func @_QPallocator_scalar_types
+! HLFIR: } {allocate_private_indices = array<i64: 0, 1, 2, 3>}
diff --git a/flang/test/Semantics/OpenMP/allocate-clause01.f90 b/flang/test/Semantics/OpenMP/allocate-clause01.f90
index a014b548a8771..b84b13df5aab0 100644
--- a/flang/test/Semantics/OpenMP/allocate-clause01.f90
+++ b/flang/test/Semantics/OpenMP/allocate-clause01.f90
@@ -18,3 +18,19 @@ subroutine allocate()
!$omp allocators allocate(align(-4): c)
allocate(c)
end subroutine
+
+subroutine parallel_allocate(x, y)
+ integer :: x, y
+
+ !$omp parallel private(x) allocate(x)
+ x = 1
+ !$omp end parallel
+
+ !$omp parallel firstprivate(x) allocate(x)
+ y = x
+ !$omp end parallel
+
+ !$omp parallel private(x, y) allocate(y) allocate(x)
+ x = y
+ !$omp end parallel
+end subroutine
diff --git a/flang/test/Transforms/OpenMP/lower-workdistribute-fission-target.mlir b/flang/test/Transforms/OpenMP/lower-workdistribute-fission-target.mlir
index 95722f101efd7..4efc02e07c6d0 100644
--- a/flang/test/Transforms/OpenMP/lower-workdistribute-fission-target.mlir
+++ b/flang/test/Transforms/OpenMP/lower-workdistribute-fission-target.mlir
@@ -41,7 +41,7 @@
// CHECK: fir.store %[[VAL_25]] to %[[VAL_14]] : !fir.ref<index>
// CHECK: fir.store %[[VAL_26]] to %[[VAL_17]] : !fir.ref<index>
// CHECK: fir.store %[[VAL_30]] to %[[VAL_20]] : !fir.ref<!fir.heap<index>>
-// CHECK: omp.target kernel_type(spmd) map_entries(%[[VAL_7]] -> %[[VAL_31:.*]], %[[VAL_8]] -> %[[VAL_32:.*]], %[[VAL_9]] -> %[[VAL_33:.*]], %[[VAL_10]] -> %[[VAL_34:.*]], %[[VAL_13]] -> %[[VAL_35:.*]], %[[VAL_16]] -> %[[VAL_36:.*]], %[[VAL_19]] -> %[[VAL_37:.*]], %[[VAL_22]] -> %[[VAL_38:.*]] : !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<!fir.heap<index>>) {
+// CHECK: omp.target kernel_type(spmd) allocate(%{{.*}} : i64 -> %{{.*}} : !fir.ref<index>) map_entries(%[[VAL_7]] -> %[[VAL_31:.*]], %[[VAL_8]] -> %[[VAL_32:.*]], %[[VAL_9]] -> %[[VAL_33:.*]], %[[VAL_10]] -> %[[VAL_34:.*]], %[[VAL_13]] -> %[[VAL_35:.*]], %[[VAL_16]] -> %[[VAL_36:.*]], %[[VAL_19]] -> %[[VAL_37:.*]], %[[VAL_22]] -> %[[VAL_38:.*]] : !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<!fir.heap<index>>) private(@addr_private %{{.*}} -> %{{.*}} : !fir.ref<index>) {
// CHECK: %[[VAL_39:.*]] = fir.load %[[VAL_35]] : !fir.ref<index>
// CHECK: %[[VAL_40:.*]] = fir.load %[[VAL_36]] : !fir.ref<index>
// CHECK: %[[VAL_41:.*]] = fir.load %[[VAL_37]] : !fir.ref<index>
@@ -62,7 +62,7 @@
// CHECK: omp.terminator
// CHECK: } {omp.combined}
// CHECK: omp.terminator
-// CHECK: } {omp.combined}
+// CHECK: } {allocate_private_indices = array<i64: 0>, omp.combined}
// CHECK: %[[VAL_45:.*]] = llvm.mlir.constant(0 : i32) : i32
// CHECK: %[[VAL_46:.*]] = fir.load %[[VAL_11]] : !fir.ref<index>
// CHECK: %[[VAL_47:.*]] = fir.load %[[VAL_14]] : !fir.ref<index>
@@ -79,6 +79,7 @@
module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true} {
+omp.private {type = private} @addr_private : index
func.func @x(%lb : index, %ub : index, %step : index, %addr : !fir.ref<index>) {
%lb_ref = fir.alloca index {bindc_name = "lb"}
fir.store %lb to %lb_ref : !fir.ref<index>
@@ -91,8 +92,17 @@ func.func @x(%lb : index, %ub : index, %step : index, %addr : !fir.ref<index>) {
%ub_map = omp.map.info var_ptr(%ub_ref : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "ub"}
%step_map = omp.map.info var_ptr(%step_ref : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "step"}
%addr_map = omp.map.info var_ptr(%addr : !fir.ref<index>, index) map_clauses(tofrom) capture(ByRef) -> !fir.ref<index> {name = "addr"}
+ %allocator = arith.constant 1 : i64
- omp.target kernel_type(generic) map_entries(%lb_map -> %ARG0, %ub_map -> %ARG1, %step_map -> %ARG2, %addr_map -> %ARG3 : !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>) {
+ "omp.target"(%addr, %allocator, %lb_map, %ub_map, %step_map, %addr_map, %addr) <{
+ allocate_private_indices = array<i64: 0>,
+ kernel_type = #omp<kernel_type(generic)>,
+ operandSegmentSizes = array<i32: 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 0, 1, 0>,
+ private_syms = [@addr_private]
+ }> ({
+ ^bb0(%ARG0: !fir.ref<index>, %ARG1: !fir.ref<index>,
+ %ARG2: !fir.ref<index>, %ARG3: !fir.ref<index>,
+ %ARG4: !fir.ref<index>):
%lb_val = fir.load %ARG0 : !fir.ref<index>
%ub_val = fir.load %ARG1 : !fir.ref<index>
%step_val = fir.load %ARG2 : !fir.ref<index>
@@ -112,7 +122,8 @@ func.func @x(%lb : index, %ub : index, %step : index, %addr : !fir.ref<index>) {
omp.terminator
}
omp.terminator
- }
+ }) : (!fir.ref<index>, i64, !fir.ref<index>, !fir.ref<index>,
+ !fir.ref<index>, !fir.ref<index>, !fir.ref<index>) -> ()
return
}
}
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index cd1223dc1702c..b6637d6df90e5 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -113,7 +113,8 @@ class OpenMP_AllocateClauseSkip<
extraClassDeclaration> {
let arguments = (ins
Variadic<AnyType>:$allocate_vars,
- Variadic<AnyType>:$allocator_vars
+ Variadic<AnyType>:$allocator_vars,
+ OptionalAttr<DenseI64ArrayAttr>:$allocate_private_indices
);
let extraClassDeclaration = [{
@@ -128,9 +129,10 @@ class OpenMP_AllocateClauseSkip<
}];
let description = [{
- The `allocator_vars` and `allocate_vars` parameters are a variadic list of
- values that specify the memory allocator to be used to obtain storage for
- private values.
+ The `allocate_vars` and `allocator_vars` parameters are parallel lists that
+ pair each allocated private value with the allocator used to obtain its
+ storage. When present, `allocate_private_indices` maps each allocate variable
+ to the corresponding position in the operation's private variables.
}];
}
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h b/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
index 2507a5a36b514..d15ea19b9065c 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
@@ -50,6 +50,10 @@
mlir::omp::TargetOp
namespace mlir::omp {
+/// Value of the `omp_null_allocator` handle. Using it as an allocator operand
+/// selects the default allocator of the binding task.
+constexpr int64_t kNullAllocator = 0;
+
/// Find the omp.new_cli, generator, and consumer of a canonical loop info.
std::tuple<NewCliOp, OpOperand *, OpOperand *> decodeCli(mlir::Value cli);
diff --git a/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp b/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp
index 00945b957e188..bdc428fc289f4 100644
--- a/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp
+++ b/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp
@@ -503,6 +503,7 @@ struct ParallelOpLowering : public OpRewritePattern<scf::ParallelOp> {
rewriter, loc,
/* allocate_vars = */ llvm::SmallVector<Value>{},
/* allocator_vars = */ llvm::SmallVector<Value>{},
+ /* allocate_private_indices = */ nullptr,
/* if_expr = */ Value{},
/* num_threads_vars = */ numThreadsVars,
/* private_vars = */ ValueRange(),
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 8605b39f17c16..190014f2bab89 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -620,6 +620,74 @@ static void printAlignedClause(OpAsmPrinter &p, Operation *op,
}
}
+static LogicalResult verifyAllocateClause(
+ Operation *op, ValueRange allocateVars, ValueRange allocatorVars,
+ DenseI64ArrayAttr allocatePrivateIndices, ValueRange privateVars = {},
+ ArrayAttr privateSyms = nullptr, bool requirePrivateIndices = false) {
+ if (allocateVars.size() != allocatorVars.size())
+ return op->emitError(
+ "expected equal sizes for allocate and allocator variables");
+
+ if (allocateVars.empty()) {
+ if (allocatePrivateIndices)
+ return op->emitError(
+ "unexpected allocate private indices without allocate variables");
+ return success();
+ }
+
+ if (!allocatePrivateIndices) {
+ if (requirePrivateIndices)
+ return op->emitError(
+ "expected an allocate private index for each allocate variable");
+ return success();
+ }
+
+ ArrayRef<int64_t> indices = allocatePrivateIndices.asArrayRef();
+ if (indices.size() != allocateVars.size())
+ return op->emitError(
+ "expected as many allocate private indices as allocate variables");
+
+ DenseSet<int64_t> usedPrivateSlots;
+ for (auto [allocateVar, privateIndex] :
+ llvm::zip_equal(allocateVars, indices)) {
+ if (privateIndex < 0 ||
+ static_cast<uint64_t>(privateIndex) >= privateVars.size())
+ return op->emitError("allocate private index is out of range");
+ if (!usedPrivateSlots.insert(privateIndex).second)
+ return op->emitError(
+ "allocate private index refers to a private variable more than once");
+
+ Value privateVar = privateVars[privateIndex];
+ if (allocateVar.getType() != privateVar.getType())
+ return op->emitError()
+ << "type mismatch between allocate variable and private variable "
+ "at index "
+ << privateIndex;
+
+ if (!privateSyms ||
+ static_cast<uint64_t>(privateIndex) >= privateSyms.size())
+ return op->emitError(
+ "allocate private index does not have a privatizer symbol");
+
+ auto privateSym = dyn_cast<SymbolRefAttr>(privateSyms[privateIndex]);
+ if (!privateSym)
+ return op->emitError(
+ "allocate private index does not reference a privatizer symbol");
+ PrivateClauseOp privatizer =
+ SymbolTable::lookupNearestSymbolFrom<PrivateClauseOp>(op, privateSym);
+ if (!privatizer)
+ return op->emitError() << "failed to lookup privatizer op with symbol: '"
+ << privateSym << "'";
+ if (privatizer.getDataSharingType() != DataSharingClauseType::Private &&
+ privatizer.getDataSharingType() != DataSharingClauseType::FirstPrivate)
+ return op->emitError(
+ "allocate private index must refer to private or firstprivate "
+ "storage");
+ }
+
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// Parser, printer and verifier for Schedule Clause
//===----------------------------------------------------------------------===//
@@ -2648,9 +2716,9 @@ LogicalResult TargetUpdateOp::verify() {
void TargetOp::build(OpBuilder &builder, OperationState &state,
const TargetExtOperands &clauses) {
MLIRContext *ctx = builder.getContext();
- // TODO Store clauses in op: allocateVars, allocatorVars.
TargetOp::build(
- builder, state, /*allocate_vars=*/{}, /*allocator_vars=*/{},
+ builder, state, clauses.allocateVars, clauses.allocatorVars,
+ makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
makeArrayAttr(ctx, clauses.dependKinds), clauses.dependVars,
makeArrayAttr(ctx, clauses.dependIteratedKinds), clauses.dependIterated,
clauses.device, clauses.dynGroupprivateAccessGroup,
@@ -2715,6 +2783,12 @@ static bool targetInReductionCapturedBy(Value inReductionVar, Value mapVarPtr) {
}
LogicalResult TargetOp::verify() {
+ if (failed(verifyAllocateClause(getOperation(), getAllocateVars(),
+ getAllocatorVars(),
+ getAllocatePrivateIndicesAttr(),
+ getPrivateVars(), getPrivateSymsAttr())))
+ return failure();
+
if (getKernelType() == TargetExecMode::bare && !isCombined())
return emitOpError() << "bare kernel requires 'omp.combined'";
@@ -2860,7 +2934,8 @@ LogicalResult TargetOp::verifyRegions() {
void ParallelOp::build(OpBuilder &builder, OperationState &state,
ArrayRef<NamedAttribute> attributes) {
ParallelOp::build(builder, state, /*allocate_vars=*/ValueRange(),
- /*allocator_vars=*/ValueRange(), /*if_expr=*/nullptr,
+ /*allocator_vars=*/ValueRange(),
+ /*allocate_private_indices=*/nullptr, /*if_expr=*/nullptr,
/*num_threads_vars=*/ValueRange(),
/*private_vars=*/ValueRange(),
/*private_syms=*/nullptr, /*private_needs_barrier=*/nullptr,
@@ -2874,6 +2949,7 @@ void ParallelOp::build(OpBuilder &builder, OperationState &state,
const ParallelOperands &clauses) {
MLIRContext *ctx = builder.getContext();
ParallelOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
+ makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
clauses.ifExpr, clauses.numThreadsVars, clauses.privateVars,
makeArrayAttr(ctx, clauses.privateSyms),
clauses.privateNeedsBarrier, clauses.procBindKind,
@@ -2926,12 +3002,13 @@ static LogicalResult verifyPrivateVarList(OpType &op) {
}
LogicalResult ParallelOp::verify() {
- if (getAllocateVars().size() != getAllocatorVars().size())
- return emitError(
- "expected equal sizes for allocate and allocator variables");
-
if (failed(verifyPrivateVarList(*this)))
return failure();
+ if (failed(verifyAllocateClause(
+ getOperation(), getAllocateVars(), getAllocatorVars(),
+ getAllocatePrivateIndicesAttr(), getPrivateVars(),
+ getPrivateSymsAttr(), /*requirePrivateIndices=*/true)))
+ return failure();
return verifyReductionVarList(*this, getReductionSyms(), getReductionVars(),
getReductionByref());
@@ -2984,6 +3061,7 @@ void TeamsOp::build(OpBuilder &builder, OperationState &state,
// TODO Store clauses in op: privateVars, privateSyms, privateNeedsBarrier
TeamsOp::build(
builder, state, clauses.allocateVars, clauses.allocatorVars,
+ makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
clauses.dynGroupprivateAccessGroup, clauses.dynGroupprivateFallback,
clauses.dynGroupprivateSize, clauses.ifExpr, clauses.numTeamsLower,
clauses.numTeamsUpperVars, /*private_vars=*/{}, /*private_syms=*/nullptr,
@@ -3033,10 +3111,11 @@ LogicalResult TeamsOp::verify() {
(getNumTeamsLower() || !getNumTeamsUpperVars().empty()))
return emitOpError() << "'num_teams' not allowed in SPMD-no-loop kernels";
- // Check for allocate clause restrictions
- if (getAllocateVars().size() != getAllocatorVars().size())
- return emitError(
- "expected equal sizes for allocate and allocator variables");
+ if (failed(verifyAllocateClause(getOperation(), getAllocateVars(),
+ getAllocatorVars(),
+ getAllocatePrivateIndicesAttr(),
+ getPrivateVars(), getPrivateSymsAttr())))
+ return failure();
if (failed(verifyDynGroupprivateClause(
op, getDynGroupprivateAccessGroupAttr(),
@@ -3071,6 +3150,7 @@ void SectionsOp::build(OpBuilder &builder, OperationState &state,
MLIRContext *ctx = builder.getContext();
// TODO Store clauses in op: privateVars, privateSyms, privateNeedsBarrier
SectionsOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
+ makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
clauses.nowait, /*private_vars=*/{},
/*private_syms=*/nullptr, /*private_needs_barrier=*/nullptr,
clauses.reductionMod, clauses.reductionVars,
@@ -3082,9 +3162,11 @@ LogicalResult SectionsOp::verify() {
if (isCombined())
return emitOpError() << "cannot be a non-innermost combined construct leaf";
- if (getAllocateVars().size() != getAllocatorVars().size())
- return emitError(
- "expected equal sizes for allocate and allocator variables");
+ if (failed(verifyAllocateClause(getOperation(), getAllocateVars(),
+ getAllocatorVars(),
+ getAllocatePrivateIndicesAttr(),
+ getPrivateVars(), getPrivateSymsAttr())))
+ return failure();
return verifyReductionVarList(*this, getReductionSyms(), getReductionVars(),
getReductionByref());
@@ -3109,6 +3191,7 @@ void ScopeOp::build(OpBuilder &builder, OperationState &state,
const ScopeOperands &clauses) {
MLIRContext *ctx = builder.getContext();
ScopeOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
+ makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
clauses.nowait, clauses.privateVars,
makeArrayAttr(ctx, clauses.privateSyms),
clauses.privateNeedsBarrier, clauses.reductionMod,
@@ -3118,9 +3201,11 @@ void ScopeOp::build(OpBuilder &builder, OperationState &state,
}
LogicalResult ScopeOp::verify() {
- if (getAllocateVars().size() != getAllocatorVars().size())
- return emitError(
- "expected equal sizes for allocate and allocator variables");
+ if (failed(verifyAllocateClause(getOperation(), getAllocateVars(),
+ getAllocatorVars(),
+ getAllocatePrivateIndicesAttr(),
+ getPrivateVars(), getPrivateSymsAttr())))
+ return failure();
if (failed(verifyPrivateVarList(*this)))
return failure();
@@ -3138,6 +3223,7 @@ void SingleOp::build(OpBuilder &builder, OperationState &state,
MLIRContext *ctx = builder.getContext();
// TODO Store clauses in op: privateVars, privateSyms, privateNeedsBarrier
SingleOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
+ makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
clauses.copyprivateVars,
makeArrayAttr(ctx, clauses.copyprivateSyms), clauses.nowait,
/*private_vars=*/{}, /*private_syms=*/nullptr,
@@ -3145,10 +3231,11 @@ void SingleOp::build(OpBuilder &builder, OperationState &state,
}
LogicalResult SingleOp::verify() {
- // Check for allocate clause restrictions
- if (getAllocateVars().size() != getAllocatorVars().size())
- return emitError(
- "expected equal sizes for allocate and allocator variables");
+ if (failed(verifyAllocateClause(getOperation(), getAllocateVars(),
+ getAllocatorVars(),
+ getAllocatePrivateIndicesAttr(),
+ getPrivateVars(), getPrivateSymsAttr())))
+ return failure();
return verifyCopyprivateVarList(*this, getCopyprivateVars(),
getCopyprivateSyms());
@@ -3340,6 +3427,7 @@ LogicalResult LoopOp::verifyRegions() {
void WsloopOp::build(OpBuilder &builder, OperationState &state,
ArrayRef<NamedAttribute> attributes) {
build(builder, state, /*allocate_vars=*/{}, /*allocator_vars=*/{},
+ /*allocate_private_indices=*/nullptr,
/*linear_vars=*/ValueRange(), /*linear_step_vars=*/ValueRange(),
/*linear_var_types*/ nullptr, /*linear_modifiers=*/nullptr,
/*nowait=*/false, /*order=*/nullptr, /*order_mod=*/nullptr,
@@ -3356,20 +3444,26 @@ void WsloopOp::build(OpBuilder &builder, OperationState &state,
void WsloopOp::build(OpBuilder &builder, OperationState &state,
const WsloopOperands &clauses) {
MLIRContext *ctx = builder.getContext();
- // TODO: Store clauses in op: allocateVars, allocatorVars
WsloopOp::build(
- builder, state,
- /*allocate_vars=*/{}, /*allocator_vars=*/{}, clauses.linearVars,
- clauses.linearStepVars, clauses.linearVarTypes, clauses.linearModifiers,
- clauses.nowait, clauses.order, clauses.orderMod, clauses.ordered,
- clauses.privateVars, makeArrayAttr(ctx, clauses.privateSyms),
- clauses.privateNeedsBarrier, clauses.reductionMod, clauses.reductionVars,
+ builder, state, clauses.allocateVars, clauses.allocatorVars,
+ makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
+ clauses.linearVars, clauses.linearStepVars, clauses.linearVarTypes,
+ clauses.linearModifiers, clauses.nowait, clauses.order, clauses.orderMod,
+ clauses.ordered, clauses.privateVars,
+ makeArrayAttr(ctx, clauses.privateSyms), clauses.privateNeedsBarrier,
+ clauses.reductionMod, clauses.reductionVars,
makeDenseBoolArrayAttr(ctx, clauses.reductionByref),
makeArrayAttr(ctx, clauses.reductionSyms), clauses.scheduleKind,
clauses.scheduleChunk, clauses.scheduleMod, clauses.scheduleSimd);
}
LogicalResult WsloopOp::verify() {
+ if (failed(verifyAllocateClause(getOperation(), getAllocateVars(),
+ getAllocatorVars(),
+ getAllocatePrivateIndicesAttr(),
+ getPrivateVars(), getPrivateSymsAttr())))
+ return failure();
+
if (failed(
verifyLinearModifiers(*this, getLinearModifiers(), getLinearVars())))
return failure();
@@ -3507,9 +3601,11 @@ LogicalResult SimdOp::verifyRegions() {
void DistributeOp::build(OpBuilder &builder, OperationState &state,
const DistributeOperands &clauses) {
DistributeOp::build(builder, state, clauses.allocateVars,
- clauses.allocatorVars, clauses.distScheduleStatic,
- clauses.distScheduleChunkSize, clauses.order,
- clauses.orderMod, clauses.privateVars,
+ clauses.allocatorVars,
+ makeDenseI64ArrayAttr(builder.getContext(),
+ clauses.allocatePrivateIndices),
+ clauses.distScheduleStatic, clauses.distScheduleChunkSize,
+ clauses.order, clauses.orderMod, clauses.privateVars,
makeArrayAttr(builder.getContext(), clauses.privateSyms),
clauses.privateNeedsBarrier);
}
@@ -3519,9 +3615,11 @@ LogicalResult DistributeOp::verify() {
return emitOpError() << "chunk size set without "
"dist_schedule_static being present";
- if (getAllocateVars().size() != getAllocatorVars().size())
- return emitError(
- "expected equal sizes for allocate and allocator variables");
+ if (failed(verifyAllocateClause(getOperation(), getAllocateVars(),
+ getAllocatorVars(),
+ getAllocatePrivateIndicesAttr(),
+ getPrivateVars(), getPrivateSymsAttr())))
+ return failure();
if (failed(verifyPrivateVarList(*this)))
return failure();
@@ -3672,6 +3770,7 @@ void TaskOp::build(OpBuilder &builder, OperationState &state,
TaskOp::build(
builder, state, clauses.iterated, clauses.affinityVars,
clauses.allocateVars, clauses.allocatorVars,
+ makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
makeArrayAttr(ctx, clauses.dependKinds), clauses.dependVars,
makeArrayAttr(ctx, clauses.dependIteratedKinds), clauses.dependIterated,
clauses.final, clauses.ifExpr, clauses.inReductionVars,
@@ -3683,6 +3782,12 @@ void TaskOp::build(OpBuilder &builder, OperationState &state,
}
LogicalResult TaskOp::verify() {
+ if (failed(verifyAllocateClause(getOperation(), getAllocateVars(),
+ getAllocatorVars(),
+ getAllocatePrivateIndicesAttr(),
+ getPrivateVars(), getPrivateSymsAttr())))
+ return failure();
+
LogicalResult verifyDependVars =
verifyDependVarList(*this, getDependKinds(), getDependVars(),
getDependIteratedKinds(), getDependIterated());
@@ -3704,12 +3809,19 @@ void TaskgroupOp::build(OpBuilder &builder, OperationState &state,
const TaskgroupOperands &clauses) {
MLIRContext *ctx = builder.getContext();
TaskgroupOp::build(builder, state, clauses.allocateVars,
- clauses.allocatorVars, clauses.taskReductionVars,
+ clauses.allocatorVars,
+ makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
+ clauses.taskReductionVars,
makeDenseBoolArrayAttr(ctx, clauses.taskReductionByref),
makeArrayAttr(ctx, clauses.taskReductionSyms));
}
LogicalResult TaskgroupOp::verify() {
+ if (failed(verifyAllocateClause(getOperation(), getAllocateVars(),
+ getAllocatorVars(),
+ getAllocatePrivateIndicesAttr())))
+ return failure();
+
return verifyReductionVarList(*this, getTaskReductionSyms(),
getTaskReductionVars(),
getTaskReductionByref());
@@ -3724,7 +3836,8 @@ void TaskloopContextOp::build(OpBuilder &builder, OperationState &state,
MLIRContext *ctx = builder.getContext();
TaskloopContextOp::build(
builder, state, clauses.allocateVars, clauses.allocatorVars,
- clauses.final, clauses.grainsizeMod, clauses.grainsize, clauses.ifExpr,
+ makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices), clauses.final,
+ clauses.grainsizeMod, clauses.grainsize, clauses.ifExpr,
clauses.inReductionVars,
makeDenseBoolArrayAttr(ctx, clauses.inReductionByref),
makeArrayAttr(ctx, clauses.inReductionSyms), clauses.mergeable,
@@ -3745,12 +3858,13 @@ TaskloopWrapperOp TaskloopContextOp::getLoopOp() {
}
LogicalResult TaskloopContextOp::verify() {
- if (getAllocateVars().size() != getAllocatorVars().size())
- return emitError(
- "expected equal sizes for allocate and allocator variables");
-
if (failed(verifyPrivateVarList(*this)))
return failure();
+ if (failed(verifyAllocateClause(getOperation(), getAllocateVars(),
+ getAllocatorVars(),
+ getAllocatePrivateIndicesAttr(),
+ getPrivateVars(), getPrivateSymsAttr())))
+ return failure();
if (failed(verifyReductionVarList(*this, getReductionSyms(),
getReductionVars(), getReductionByref())) ||
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 1c50ff192c3d5..01ce1cad7237d 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -36,6 +36,7 @@
#include "llvm/IR/ReplaceConstant.h"
#include "llvm/Support/AMDGPUAddrSpace.h"
#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/MathExtras.h"
#include "llvm/Support/NVPTXAddrSpace.h"
#include "llvm/Support/VirtualFileSystem.h"
#include "llvm/TargetParser/Triple.h"
@@ -582,7 +583,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
checkReduction(op, result);
})
.Case([&](omp::ParallelOp op) {
- checkAllocate(op, result);
checkReduction(op, result);
checkNumThreads(op, result);
})
@@ -1006,6 +1006,11 @@ convertOmpCritical(Operation &opInst, llvm::IRBuilderBase &builder,
/// A util to collect info needed to convert delayed privatizers from MLIR to
/// LLVM.
struct PrivateVarsInfo {
+ struct AllocatorPrivateInfo {
+ llvm::Value *allocatedPtr;
+ llvm::Value *allocator;
+ };
+
template <typename OP>
PrivateVarsInfo(OP op)
: blockArgs(
@@ -1022,6 +1027,8 @@ struct PrivateVarsInfo {
SmallVector<mlir::Value> mlirVars;
SmallVector<llvm::Value *> llvmVars;
SmallVector<omp::PrivateClauseOp> privatizers;
+ llvm::DenseMap<Value, llvm::Value *> convertedAllocators;
+ SmallVector<AllocatorPrivateInfo> allocatorPrivates;
private:
/// Populates `privatizations` with privatization declarations used for the
@@ -1961,14 +1968,63 @@ allocatePrivateVars(T op, llvm::IRBuilderBase &builder,
->getDataLayout()
.getProgramAddressSpace();
- for (auto [privDecl, mlirPrivVar, blockArg] :
- llvm::zip_equal(privateVarsInfo.privatizers, privateVarsInfo.mlirVars,
- privateVarsInfo.blockArgs)) {
+ SmallVector<int64_t> allocateItemForPrivate(privateVarsInfo.blockArgs.size(),
+ -1);
+ ValueRange allocatorVars;
+ if constexpr (std::is_same_v<T, omp::ParallelOp>) {
+ allocatorVars = op.getAllocatorVars();
+ if (auto privateIndices = op.getAllocatePrivateIndicesAttr())
+ for (auto [allocateIndex, privateIndex] :
+ llvm::enumerate(privateIndices.asArrayRef()))
+ allocateItemForPrivate[privateIndex] = allocateIndex;
+ }
+
+ for (auto [privateIndex, tuple] : llvm::enumerate(llvm::zip_equal(
+ privateVarsInfo.privatizers, privateVarsInfo.mlirVars,
+ privateVarsInfo.blockArgs))) {
+ auto [privDecl, mlirPrivVar, blockArg] = tuple;
llvm::Type *llvmAllocType =
moduleTranslation.convertType(privDecl.getType());
builder.SetInsertPoint(allocaIP.getBlock()->getTerminator());
llvm::Value *llvmPrivateVar = nullptr;
- if (mightUseDeviceSharedMem && omp::allocaUsesRequireSharedMem(blockArg)) {
+ int64_t allocateIndex = allocateItemForPrivate[privateIndex];
+ if (allocateIndex >= 0) {
+ if (mightUseDeviceSharedMem ||
+ op->template getParentOfType<omp::TargetOp>())
+ return llvm::createStringError(
+ "allocate clause on a device parallel region is not supported");
+ if (!llvmAllocType->isSized())
+ return llvm::createStringError(
+ "allocate clause private type must have a fixed size");
+ llvm::TypeSize size = dataLayout.getTypeAllocSize(llvmAllocType);
+ if (size.isScalable())
+ return llvm::createStringError(
+ "allocate clause private type must have a fixed size");
+ llvm::IntegerType *sizeTy =
+ moduleTranslation.getLLVMModule()->getDataLayout().getIntPtrType(
+ moduleTranslation.getLLVMModule()->getContext());
+ if (!llvm::isUIntN(sizeTy->getBitWidth(), size.getFixedValue()))
+ return llvm::createStringError(
+ "OpenMP allocation size cannot be represented by the target size "
+ "type");
+ llvm::Value *sizeValue =
+ llvm::ConstantInt::get(sizeTy, size.getFixedValue());
+
+ Value allocatorVar = allocatorVars[allocateIndex];
+ auto allocator = privateVarsInfo.convertedAllocators.find(allocatorVar);
+ if (allocator == privateVarsInfo.convertedAllocators.end())
+ return llvm::createStringError(
+ "failed to find converted OpenMP allocator operand");
+ llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
+ llvmPrivateVar = ompBuilder->createOMPAlloc(
+ ompLoc, sizeValue, allocator->second, "omp.private.alloc");
+ if (!llvmPrivateVar)
+ return llvm::createStringError(
+ "failed to create OpenMP private allocation");
+ privateVarsInfo.allocatorPrivates.push_back(
+ {llvmPrivateVar, allocator->second});
+ } else if (mightUseDeviceSharedMem &&
+ omp::allocaUsesRequireSharedMem(blockArg)) {
llvmPrivateVar = ompBuilder->createOMPAllocShared(builder, llvmAllocType);
} else {
llvmPrivateVar = builder.CreateAlloca(
@@ -2108,6 +2164,7 @@ cleanupPrivateVars(T op, llvm::IRBuilderBase &builder,
/*shouldLoadCleanupRegionArg=*/false)))
return mlir::emitError(loc, "failed to inline `dealloc` region of an "
"`omp.private` op in");
+ setInsertPointForPossiblyEmptyBlock(builder);
llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
bool mightUseDeviceSharedMem = omp::opInSharedDeviceContext(*op);
@@ -2121,6 +2178,12 @@ cleanupPrivateVars(T op, llvm::IRBuilderBase &builder,
}
}
+ llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
+ for (const PrivateVarsInfo::AllocatorPrivateInfo &allocation :
+ llvm::reverse(privateVarsInfo.allocatorPrivates))
+ ompBuilder->createOMPFree(ompLoc, allocation.allocatedPtr,
+ allocation.allocator);
+
return success();
}
@@ -4780,6 +4843,24 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
return failure();
PrivateVarsInfo privateVarsInfo(opInst);
+ for (Value allocatorVar : opInst.getAllocatorVars()) {
+ if (privateVarsInfo.convertedAllocators.contains(allocatorVar))
+ continue;
+
+ llvm::Value *allocator = moduleTranslation.lookupValue(allocatorVar);
+ if (!allocator)
+ return opInst.emitError("failed to translate OpenMP allocator operand");
+ if (allocator->getType()->isIntegerTy())
+ allocator = builder.CreateIntToPtr(allocator, builder.getPtrTy());
+ else if (allocator->getType()->isPointerTy())
+ allocator = builder.CreatePointerBitCastOrAddrSpaceCast(
+ allocator, builder.getPtrTy());
+ else
+ return opInst.emitError(
+ "OpenMP allocator operand must have integer or pointer type");
+
+ privateVarsInfo.convertedAllocators.try_emplace(allocatorVar, allocator);
+ }
// Collect reduction declarations
SmallVector<omp::DeclareReductionOp> reductionDecls;
diff --git a/mlir/test/Dialect/OpenMP/host-op-filtering.mlir b/mlir/test/Dialect/OpenMP/host-op-filtering.mlir
index 8108a1af31e21..e11ef390c2147 100644
--- a/mlir/test/Dialect/OpenMP/host-op-filtering.mlir
+++ b/mlir/test/Dialect/OpenMP/host-op-filtering.mlir
@@ -353,9 +353,10 @@ module attributes {omp.is_target_device = true} {
// CHECK-NEXT: %[[MAP:.*]] = omp.map.info var_ptr(%[[ARG0]] : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr
%1 = omp.map.info var_ptr(%arg0 : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr
// CHECK-NEXT: omp.target kernel_type(generic) allocate(%[[ARG0]] : !llvm.ptr -> %[[ARG0]] : !llvm.ptr) thread_limit(%[[ARG1]] : i32) map_entries(%[[MAP]] -> %{{.*}} : !llvm.ptr) private(@privatizer %[[ARG0]] -> %{{.*}} : !llvm.ptr)
+ // CHECK: } {allocate_private_indices = array<i64: 0>}
omp.target kernel_type(generic) allocate(%arg0 : !llvm.ptr -> %arg0 : !llvm.ptr) depend(taskdependin -> %arg0 : !llvm.ptr) device(%arg1 : i32) if(%arg2) thread_limit(%arg1 : i32) in_reduction(@reduction %arg0 : !llvm.ptr) map_entries(%1 -> %arg3 : !llvm.ptr) private(@privatizer %arg0 -> %arg4 : !llvm.ptr) {
omp.terminator
- }
+ } {allocate_private_indices = array<i64: 0>}
// CHECK-NOT: omp.target_enter_data
// CHECK-NOT: omp.target_exit_data
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index a568bffe3cf99..41f2b705d31fe 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -3362,6 +3362,88 @@ func.func @omp_distribute_allocate(%data_var : memref<i32>, %lb : i32, %ub : i32
// -----
+omp.private {type = private} @allocate_private : i32
+
+func.func @omp_parallel_allocate_empty_map() {
+ // expected-error @below {{unexpected allocate private indices without allocate variables}}
+ omp.parallel {
+ omp.terminator
+ } {allocate_private_indices = array<i64>}
+ return
+}
+
+// -----
+
+omp.private {type = private} @allocate_private : i32
+
+func.func @omp_parallel_allocate_missing_map(%allocator : i64, %var : !llvm.ptr) {
+ // expected-error @below {{expected an allocate private index for each allocate variable}}
+ omp.parallel allocate(%allocator : i64 -> %var : !llvm.ptr)
+ private(@allocate_private %var -> %private : !llvm.ptr) {
+ omp.terminator
+ }
+ return
+}
+
+// -----
+
+omp.private {type = private} @allocate_private : i32
+
+func.func @omp_parallel_allocate_map_size(%allocator : i64, %var : !llvm.ptr) {
+ // expected-error @below {{expected as many allocate private indices as allocate variables}}
+ omp.parallel allocate(%allocator : i64 -> %var : !llvm.ptr)
+ private(@allocate_private %var -> %private : !llvm.ptr) {
+ omp.terminator
+ } {allocate_private_indices = array<i64: 0, 0>}
+ return
+}
+
+// -----
+
+omp.private {type = private} @allocate_private : i32
+
+func.func @omp_parallel_allocate_map_range(%allocator : i64, %var : !llvm.ptr) {
+ // expected-error @below {{allocate private index is out of range}}
+ omp.parallel allocate(%allocator : i64 -> %var : !llvm.ptr)
+ private(@allocate_private %var -> %private : !llvm.ptr) {
+ omp.terminator
+ } {allocate_private_indices = array<i64: 1>}
+ return
+}
+
+// -----
+
+omp.private {type = private} @x_private : i32
+omp.private {type = private} @y_private : i32
+
+func.func @omp_parallel_allocate_map_duplicate(
+ %allocator : i64, %x : !llvm.ptr, %y : !llvm.ptr) {
+ // expected-error @below {{allocate private index refers to a private variable more than once}}
+ omp.parallel allocate(%allocator : i64 -> %x : !llvm.ptr,
+ %allocator : i64 -> %y : !llvm.ptr)
+ private(@x_private %x -> %x_private,
+ @y_private %y -> %y_private : !llvm.ptr, !llvm.ptr) {
+ omp.terminator
+ } {allocate_private_indices = array<i64: 0, 0>}
+ return
+}
+
+// -----
+
+omp.private {type = private} @allocate_private : i32
+
+func.func @omp_parallel_allocate_type_mismatch(
+ %allocator : i64, %allocate_var : i64, %private_var : !llvm.ptr) {
+ // expected-error @below {{type mismatch between allocate variable and private variable at index 0}}
+ omp.parallel allocate(%allocator : i64 -> %allocate_var : i64)
+ private(@allocate_private %private_var -> %private : !llvm.ptr) {
+ omp.terminator
+ } {allocate_private_indices = array<i64: 0>}
+ return
+}
+
+// -----
+
func.func @omp_distribute_nested_wrapper(%lb: index, %ub: index, %step: index) -> () {
// expected-error @below {{an 'omp.wsloop' nested wrapper is only allowed when a composite 'omp.parallel' is the direct parent}}
omp.distribute {
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index ea25a0a75f646..b2fe5effd2de6 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -91,23 +91,23 @@ func.func @omp_terminator() -> () {
}
func.func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : i32, %idx : index) -> () {
- // CHECK: omp.parallel allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>) if(%{{.*}}) num_threads(%{{.*}} : i32)
- "omp.parallel" (%data_var, %data_var, %if_cond, %num_threads) ({
+ // CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : i32)
+ "omp.parallel" (%if_cond, %num_threads) ({
// test without if condition
- // CHECK: omp.parallel allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>) num_threads(%{{.*}} : i32)
- "omp.parallel"(%data_var, %data_var, %num_threads) ({
+ // CHECK: omp.parallel num_threads(%{{.*}} : i32)
+ "omp.parallel"(%num_threads) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 1,1,0,1,0,0>} : (memref<i32>, memref<i32>, i32) -> ()
+ }) {operandSegmentSizes = array<i32: 0,0,0,1,0,0>} : (i32) -> ()
// CHECK: omp.barrier
omp.barrier
// test without num_threads
- // CHECK: omp.parallel allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>) if(%{{.*}})
- "omp.parallel"(%data_var, %data_var, %if_cond) ({
+ // CHECK: omp.parallel if(%{{.*}})
+ "omp.parallel"(%if_cond) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 1,1,1,0,0,0>} : (memref<i32>, memref<i32>, i1) -> ()
+ }) {operandSegmentSizes = array<i32: 0,0,1,0,0,0>} : (i1) -> ()
// test without allocate
// CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : i32)
@@ -116,13 +116,7 @@ func.func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : i
}) {operandSegmentSizes = array<i32: 0,0,1,1,0,0>} : (i1, i32) -> ()
omp.terminator
- }) {operandSegmentSizes = array<i32: 1,1,1,1,0,0>, proc_bind_kind = #omp<procbindkind spread>} : (memref<i32>, memref<i32>, i1, i32) -> ()
-
- // test with multiple parameters for single variadic argument
- // CHECK: omp.parallel allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
- "omp.parallel" (%data_var, %data_var) ({
- omp.terminator
- }) {operandSegmentSizes = array<i32: 1,1,0,0,0,0>} : (memref<i32>, memref<i32>) -> ()
+ }) {operandSegmentSizes = array<i32: 0,0,1,1,0,0>, proc_bind_kind = #omp<procbindkind spread>} : (i1, i32) -> ()
// CHECK: omp.parallel
omp.parallel {
@@ -162,6 +156,8 @@ func.func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : i
return
}
+omp.private {type = private} @parallel_allocate_private : memref<i32>
+
func.func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads : i32, %allocator : si32) -> () {
// CHECK: omp.parallel
omp.parallel {
@@ -197,10 +193,13 @@ func.func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_thre
omp.terminator
}
- // CHECK: omp.parallel allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
- omp.parallel allocate(%data_var : memref<i32> -> %data_var : memref<i32>) {
+ // CHECK: omp.parallel allocate(
+ // CHECK-SAME: private(
+ // CHECK: } {allocate_private_indices = array<i64: 0>}
+ omp.parallel allocate(%allocator : si32 -> %data_var : memref<i32>)
+ private(@parallel_allocate_private %data_var -> %private : memref<i32>) {
omp.terminator
- }
+ } {allocate_private_indices = array<i64: 0>}
// CHECK: omp.parallel
// CHECK-NEXT: omp.parallel if(%{{.*}})
diff --git a/mlir/test/Target/LLVMIR/openmp-allocate-clause.mlir b/mlir/test/Target/LLVMIR/openmp-allocate-clause.mlir
new file mode 100644
index 0000000000000..dc2515e3f1358
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-allocate-clause.mlir
@@ -0,0 +1,203 @@
+// RUN: split-file %s %t
+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %t/valid.mlir | FileCheck %s
+// RUN: mlir-translate -mlir-to-llvmir %t/i386.mlir | FileCheck %s --check-prefix=I386
+// RUN: not mlir-translate -mlir-to-llvmir %t/i386-overflow.mlir 2>&1 | FileCheck %s --check-prefix=I386-OVERFLOW
+// RUN: not mlir-translate -mlir-to-llvmir %t/device.mlir 2>&1 | FileCheck %s --check-prefix=DEVICE
+
+//--- valid.mlir
+
+omp.private {type = firstprivate} @x.firstprivate : i32 copy {
+^bb0(%original: !llvm.ptr, %private: !llvm.ptr):
+ %value = llvm.load %original : !llvm.ptr -> i32
+ llvm.store %value, %private : i32, !llvm.ptr
+ omp.yield(%private : !llvm.ptr)
+}
+
+omp.private {type = private} @y.private : i32
+
+llvm.func @allocator_dynamic(%x: !llvm.ptr, %y: !llvm.ptr, %allocator: i64) {
+ omp.parallel allocate(%allocator : i64 -> %x : !llvm.ptr)
+ private(@x.firstprivate %x -> %x.private,
+ @y.private %y -> %y.private : !llvm.ptr, !llvm.ptr) {
+ %x.value = llvm.load %x.private : !llvm.ptr -> i32
+ %one = llvm.mlir.constant(1 : i32) : i32
+ %next = llvm.add %x.value, %one : i32
+ llvm.store %next, %x.private : i32, !llvm.ptr
+ llvm.store %one, %y.private : i32, !llvm.ptr
+ omp.terminator
+ } {allocate_private_indices = array<i64: 0>}
+ llvm.return
+}
+
+// CHECK-LABEL: define void @allocator_dynamic(
+// CHECK-COUNT-1: %[[ALLOCATOR:.*]] = inttoptr i64 %{{.*}} to ptr
+// CHECK: store ptr %[[ALLOCATOR]], ptr
+// CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_call
+// CHECK-LABEL: define internal void @allocator_dynamic..omp_par
+// CHECK: %[[CAPTURED_ALLOCATOR:.*]] = load ptr, ptr %{{.*}}, align 8
+// CHECK: %[[ALLOC:.*]] = call ptr @__kmpc_alloc({{.*}}, i64 4, ptr %[[CAPTURED_ALLOCATOR]])
+// CHECK: %[[Y_ALLOCA:.*]] = alloca i32, align 4
+// CHECK: %[[ORIGINAL:.*]] = load i32, ptr %{{.*}}, align 4
+// CHECK: store i32 %[[ORIGINAL]], ptr %[[ALLOC]], align 4
+// CHECK: %[[PRIVATE:.*]] = load i32, ptr %[[ALLOC]], align 4
+// CHECK: store i32 %{{.*}}, ptr %[[ALLOC]], align 4
+// CHECK: store i32 1, ptr %[[Y_ALLOCA]], align 4
+// CHECK: .fini:
+// CHECK: call void @__kmpc_free({{.*}}, ptr %[[ALLOC]], ptr %[[CAPTURED_ALLOCATOR]])
+
+// -----
+
+llvm.func @private_dealloc(!llvm.ptr)
+
+omp.private {type = firstprivate} @x.private : i32 copy {
+^bb0(%original: !llvm.ptr, %private: !llvm.ptr):
+ %value = llvm.load %original : !llvm.ptr -> i32
+ llvm.store %value, %private : i32, !llvm.ptr
+ omp.yield(%private : !llvm.ptr)
+} dealloc {
+^bb0(%private: !llvm.ptr):
+ llvm.call @private_dealloc(%private) : (!llvm.ptr) -> ()
+ omp.yield
+}
+omp.private {type = private} @y.private : i32
+
+llvm.func @allocator_reverse_free(%x: !llvm.ptr, %y: !llvm.ptr,
+ %allocator.x: i64, %allocator.y: i64) {
+ omp.parallel allocate(%allocator.y : i64 -> %y : !llvm.ptr,
+ %allocator.x : i64 -> %x : !llvm.ptr)
+ private(@x.private %x -> %x.private,
+ @y.private %y -> %y.private : !llvm.ptr, !llvm.ptr) {
+ omp.terminator
+ } {allocate_private_indices = array<i64: 1, 0>}
+ llvm.return
+}
+
+// CHECK-LABEL: define void @allocator_reverse_free(
+// CHECK: %[[ALLOCATOR_Y:.*]] = inttoptr i64 %{{.*}} to ptr
+// CHECK: %[[ALLOCATOR_X:.*]] = inttoptr i64 %{{.*}} to ptr
+// CHECK-LABEL: define internal void @allocator_reverse_free..omp_par
+// CHECK: %[[CAPTURED_X:.*]] = load ptr, ptr %{{.*}}, align 8
+// CHECK: %[[CAPTURED_Y:.*]] = load ptr, ptr %{{.*}}, align 8
+// CHECK: %[[X_ALLOC:.*]] = call ptr @__kmpc_alloc({{.*}}, i64 4, ptr %[[CAPTURED_X]])
+// CHECK: %[[Y_ALLOC:.*]] = call ptr @__kmpc_alloc({{.*}}, i64 4, ptr %[[CAPTURED_Y]])
+// CHECK: call void @private_dealloc(ptr %[[X_ALLOC]])
+// CHECK: call void @__kmpc_free({{.*}}, ptr %[[Y_ALLOC]], ptr %[[CAPTURED_Y]])
+// CHECK: call void @__kmpc_free({{.*}}, ptr %[[X_ALLOC]], ptr %[[CAPTURED_X]])
+// CHECK-NOT: call void @__kmpc_free
+
+// -----
+
+omp.private {type = private} @x.private : i32
+
+llvm.func @allocator_cancel(%x: !llvm.ptr) {
+ %null = llvm.mlir.constant(0 : i64) : i64
+ omp.parallel allocate(%null : i64 -> %x : !llvm.ptr)
+ private(@x.private %x -> %x.private : !llvm.ptr) {
+ omp.cancel cancellation_construct_type(parallel)
+ omp.terminator
+ } {allocate_private_indices = array<i64: 0>}
+ llvm.return
+}
+
+// CHECK-LABEL: define internal void @allocator_cancel..omp_par
+// CHECK: %[[ALLOC:.*]] = call ptr @__kmpc_alloc({{.*}}, i64 4, ptr null)
+// CHECK: {{.*}}.cncl:
+// CHECK: br label %[[FINI:.*]]
+// CHECK: .fini:
+// CHECK: call void @__kmpc_free({{.*}}, ptr %[[ALLOC]], ptr null)
+// CHECK: omp.par.pre_finalize:
+// CHECK: br label %[[FINI]]
+// CHECK-NOT: call void @__kmpc_free
+
+// -----
+
+omp.private {type = private} @x.private : i32
+
+llvm.func @allocator_cancellation_point(%x: !llvm.ptr) {
+ %null = llvm.mlir.constant(0 : i64) : i64
+ omp.parallel allocate(%null : i64 -> %x : !llvm.ptr)
+ private(@x.private %x -> %x.private : !llvm.ptr) {
+ omp.cancellation_point cancellation_construct_type(parallel)
+ omp.terminator
+ } {allocate_private_indices = array<i64: 0>}
+ llvm.return
+}
+
+// CHECK-LABEL: define internal void @allocator_cancellation_point..omp_par
+// CHECK: %[[ALLOC:.*]] = call ptr @__kmpc_alloc({{.*}}, i64 4, ptr null)
+// CHECK: {{.*}}.cncl:
+// CHECK: br label %[[FINI:.*]]
+// CHECK: .fini:
+// CHECK: call void @__kmpc_free({{.*}}, ptr %[[ALLOC]], ptr null)
+// CHECK: omp.par.pre_finalize:
+// CHECK: br label %[[FINI]]
+// CHECK-NOT: call void @__kmpc_free
+
+//--- i386.mlir
+
+module attributes {
+ llvm.data_layout = "e-m:e-p:32:32-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:32-n8:16:32-S128",
+ llvm.target_triple = "i386-unknown-linux-gnu"
+} {
+ omp.private {type = private} @x.private : i32
+
+ llvm.func @allocator_i386(%x: !llvm.ptr) {
+ %null = llvm.mlir.constant(0 : i64) : i64
+ omp.parallel allocate(%null : i64 -> %x : !llvm.ptr)
+ private(@x.private %x -> %x.private : !llvm.ptr) {
+ omp.terminator
+ } {allocate_private_indices = array<i64: 0>}
+ llvm.return
+ }
+}
+
+// I386: target datalayout = "e-m:e-p:32:32-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:32-n8:16:32-S128"
+// I386-LABEL: define internal void @allocator_i386..omp_par
+// I386: %[[ALLOC:.*]] = call ptr @__kmpc_alloc(i32 %{{.*}}, i32 4, ptr null)
+// I386: call void @__kmpc_free(i32 %{{.*}}, ptr %[[ALLOC]], ptr null)
+// I386: ret void
+// I386-LABEL: declare noalias ptr @__kmpc_alloc(i32, i32, ptr)
+
+//--- i386-overflow.mlir
+
+module attributes {
+ llvm.data_layout = "e-m:e-p:32:32-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:32-n8:16:32-S128",
+ llvm.target_triple = "i386-unknown-linux-gnu"
+} {
+ omp.private {type = private} @x.private : !llvm.array<4294967296 x i8>
+
+ llvm.func @allocator_i386_overflow(%x: !llvm.ptr) {
+ %null = llvm.mlir.constant(0 : i64) : i64
+ omp.parallel allocate(%null : i64 -> %x : !llvm.ptr)
+ private(@x.private %x -> %x.private : !llvm.ptr) {
+ omp.terminator
+ } {allocate_private_indices = array<i64: 0>}
+ llvm.return
+ }
+}
+
+// I386-OVERFLOW-NOT: __kmpc_alloc
+// I386-OVERFLOW: OpenMP allocation size cannot be represented by the target size type
+// I386-OVERFLOW-NOT: __kmpc_alloc
+// I386-OVERFLOW: LLVM Translation failed for operation: omp.parallel
+// I386-OVERFLOW-NOT: __kmpc_alloc
+
+//--- device.mlir
+
+omp.private {type = private} @device.private : i32
+
+llvm.func @allocator_device() {
+ omp.target kernel_type(generic) {
+ %allocator = llvm.mlir.constant(0 : i64) : i64
+ %x = llvm.alloca %allocator x i32 : (i64) -> !llvm.ptr
+ omp.parallel allocate(%allocator : i64 -> %x : !llvm.ptr)
+ private(@device.private %x -> %private : !llvm.ptr) {
+ omp.terminator
+ } {allocate_private_indices = array<i64: 0>}
+ omp.terminator
+ }
+ llvm.return
+}
+
+// DEVICE: allocate clause on a device parallel region is not supported
+// DEVICE: LLVM Translation failed for operation: omp.parallel
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index e29659255f690..2aedee1402798 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -52,12 +52,15 @@ llvm.func @distribute_order(%lb : i32, %ub : i32, %step : i32) {
// -----
-llvm.func @parallel_allocate(%x : !llvm.ptr) {
- // expected-error at below {{not yet implemented: Unhandled clause allocate in omp.parallel operation}}
+omp.private {type = private} @parallel_bad_allocator_private : i32
+
+llvm.func @parallel_bad_allocator(%allocator : f32, %x : !llvm.ptr) {
+ // expected-error at below {{OpenMP allocator operand must have integer or pointer type}}
// expected-error at below {{LLVM Translation failed for operation: omp.parallel}}
- omp.parallel allocate(%x : !llvm.ptr -> %x : !llvm.ptr) {
+ omp.parallel allocate(%allocator : f32 -> %x : !llvm.ptr)
+ private(@parallel_bad_allocator_private %x -> %private : !llvm.ptr) {
omp.terminator
- }
+ } {allocate_private_indices = array<i64: 0>}
llvm.return
}
>From 8e28a0b674c30eb4d222b8ac9274306bbec7d025 Mon Sep 17 00:00:00 2001
From: Sairudra More <sairudra60 at gmail.com>
Date: Thu, 30 Jul 2026 03:06:38 -0500
Subject: [PATCH 2/2] [flang][OpenMP] Fix allocate indices for common blocks
Co-authored-by: Copilot <223556219+Copilot at users.noreply.github.com>
---
flang/lib/Lower/OpenMP/OpenMP.cpp | 22 ++++++++++++++-----
.../OpenMP/allocate-clause-allocator.f90 | 18 +++++++++++++++
2 files changed, 35 insertions(+), 5 deletions(-)
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 9fb70be2b515a..109636c096804 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -3507,14 +3507,26 @@ genParallelOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
if (!clauseOps.allocateVars.empty()) {
llvm::DenseMap<const semantics::Symbol *, int64_t> privateSlots;
- for (auto [index, object] : llvm::enumerate(args.priv.objects)) {
+ int64_t privateSlot = 0;
+ auto addPrivateSlot = [&](const semantics::Symbol &symbol) {
+ if (!privateSlots.try_emplace(&symbol.GetUltimate(), privateSlot).second)
+ fir::emitFatalError(
+ loc, "symbol with multiple private storage slots on one construct");
+ ++privateSlot;
+ };
+ for (const Object &object : args.priv.objects) {
const semantics::Symbol *symbol = object.sym();
if (!symbol)
fir::emitFatalError(loc, "private item without a semantic symbol");
- const semantics::Symbol *ultimate = &symbol->GetUltimate();
- if (!privateSlots.try_emplace(ultimate, index).second)
- fir::emitFatalError(
- loc, "symbol with multiple private storage slots on one construct");
+ // A privatized common block contributes one private operand per member,
+ // so slot numbering must follow the same expansion.
+ if (const auto *commonDetails =
+ symbol->detailsIf<semantics::CommonBlockDetails>()) {
+ for (const auto &member : commonDetails->objects())
+ addPrivateSlot(*member);
+ } else {
+ addPrivateSlot(*symbol);
+ }
}
llvm::DenseSet<const semantics::Symbol *> allocateSymbols;
diff --git a/flang/test/Lower/OpenMP/allocate-clause-allocator.f90 b/flang/test/Lower/OpenMP/allocate-clause-allocator.f90
index fd64f4126f65e..79f8554dbdbda 100644
--- a/flang/test/Lower/OpenMP/allocate-clause-allocator.f90
+++ b/flang/test/Lower/OpenMP/allocate-clause-allocator.f90
@@ -117,3 +117,21 @@ subroutine allocator_scalar_types(r, c, l, s)
! HLFIR-LABEL: func.func @_QPallocator_scalar_types
! HLFIR: } {allocate_private_indices = array<i64: 0, 1, 2, 3>}
+
+subroutine allocator_common_block(y)
+ integer :: x1, x2, y
+ common /blk/ x1, x2
+ !$omp parallel private(/blk/, y) allocate(y)
+ x1 = 1
+ x2 = 2
+ y = 3
+ !$omp end parallel
+end subroutine
+
+! A privatized common block contributes one private operand per member, so the
+! recorded index of an allocated item must account for that expansion.
+! HLFIR-LABEL: func.func @_QPallocator_common_block
+! HLFIR: omp.parallel allocate({{.*}} -> %[[Y:.*]]#0 : !fir.ref<i32>)
+! HLFIR-SAME: private({{.*}} -> %{{.*}}, {{.*}} -> %{{.*}}, {{.*}} %[[Y]]#0 -> %{{.*}} :
+! HLFIR-SAME: !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>) {
+! HLFIR: } {allocate_private_indices = array<i64: 2>}
More information about the flang-commits
mailing list