[flang-commits] [flang] [llvm] [Flang][OpenMP] Generate correct present checks for implicit maps of optional allocatables (PR #138210)
via flang-commits
flang-commits at lists.llvm.org
Thu May 1 15:44:11 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-flang-openmp
Author: None (agozillon)
<details>
<summary>Changes</summary>
Currently, we do not generate the appropriate checks to check if an optional
allocatable argument is present before accessing relevant components of it,
in particular when creating bounds, we must generate a presence check and we
must make sure we do not generate/keep an load external to the presence check
by utilising the raw address rather than the regular address of the info
data structure.
Similarly in cases for optional allocatables we must treat them like non-allocatable
arguments and generate an intermediate allocation that we can have as a location
in memory that we can access later in the lowering without causing segfaults when
we perform "mapping" on it, even if the end result is an empty allocatable
(basically, we shouldn't explode if someone tries to map a non-present optional,
similar to C++ when mapping null data).
---
Full diff: https://github.com/llvm/llvm-project/pull/138210.diff
4 Files Affected:
- (modified) flang/include/flang/Optimizer/Builder/DirectivesCommon.h (+22-3)
- (modified) flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp (+12-4)
- (added) flang/test/Lower/OpenMP/optional-argument-map-2.f90 (+46)
- (added) offload/test/offloading/fortran/optional-mapped-arguments-2.f90 (+57)
``````````diff
diff --git a/flang/include/flang/Optimizer/Builder/DirectivesCommon.h b/flang/include/flang/Optimizer/Builder/DirectivesCommon.h
index 8684299ab6792..e655c8e592364 100644
--- a/flang/include/flang/Optimizer/Builder/DirectivesCommon.h
+++ b/flang/include/flang/Optimizer/Builder/DirectivesCommon.h
@@ -156,9 +156,9 @@ genBoundsOpsFromBox(fir::FirOpBuilder &builder, mlir::Location loc,
builder.genIfOp(loc, resTypes, info.isPresent, /*withElseRegion=*/true)
.genThen([&]() {
mlir::Value box =
- !fir::isBoxAddress(info.addr.getType())
+ !fir::isBoxAddress(info.rawInput.getType())
? info.addr
- : builder.create<fir::LoadOp>(loc, info.addr);
+ : builder.create<fir::LoadOp>(loc, info.rawInput);
llvm::SmallVector<mlir::Value> boundValues =
gatherBoundsOrBoundValues<BoundsOp, BoundsType>(
builder, loc, dataExv, box,
@@ -243,6 +243,17 @@ genBaseBoundsOps(fir::FirOpBuilder &builder, mlir::Location loc,
return bounds;
}
+/// Checks if an argument is optional based on the fortran attributes
+/// that are tied to it.
+inline bool isOptionalArgument(mlir::Operation *op) {
+ if (auto declareOp = mlir::dyn_cast_or_null<hlfir::DeclareOp>(op))
+ if (declareOp.getFortranAttrs() &&
+ bitEnumContainsAny(*declareOp.getFortranAttrs(),
+ fir::FortranVariableFlagsEnum::optional))
+ return true;
+ return false;
+}
+
template <typename BoundsOp, typename BoundsType>
llvm::SmallVector<mlir::Value>
genImplicitBoundsOps(fir::FirOpBuilder &builder, AddrAndBoundsInfo &info,
@@ -251,9 +262,17 @@ genImplicitBoundsOps(fir::FirOpBuilder &builder, AddrAndBoundsInfo &info,
llvm::SmallVector<mlir::Value> bounds;
mlir::Value baseOp = info.rawInput;
- if (mlir::isa<fir::BaseBoxType>(fir::unwrapRefType(baseOp.getType())))
+ if (mlir::isa<fir::BaseBoxType>(fir::unwrapRefType(baseOp.getType()))) {
+ // if it's an optional argument, it is possible it is not present, in which
+ // case, emitting loads or stores to access bounds data will result in a
+ // runtime segfault, so we must emit guards against this.
+ if (!info.isPresent && isOptionalArgument(info.rawInput.getDefiningOp())) {
+ info.isPresent = builder.create<fir::IsPresentOp>(
+ loc, builder.getI1Type(), info.rawInput);
+ }
bounds =
genBoundsOpsFromBox<BoundsOp, BoundsType>(builder, loc, dataExv, info);
+ }
if (mlir::isa<fir::SequenceType>(fir::unwrapRefType(baseOp.getType()))) {
bounds = genBaseBoundsOps<BoundsOp, BoundsType>(builder, loc, dataExv,
dataExvIsAssumedSize);
diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
index 3fcb4b04a7b76..05d17bf71514b 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -131,7 +131,8 @@ class MapInfoFinalizationPass
boxMap.getVarPtr().getDefiningOp()))
descriptor = addrOp.getVal();
- if (!mlir::isa<fir::BaseBoxType>(descriptor.getType()))
+ if (!mlir::isa<fir::BaseBoxType>(descriptor.getType()) &&
+ !fir::factory::isOptionalArgument(descriptor.getDefiningOp()))
return descriptor;
mlir::Value &slot = localBoxAllocas[descriptor.getDefiningOp()];
@@ -151,7 +152,12 @@ class MapInfoFinalizationPass
mlir::Location loc = boxMap->getLoc();
assert(allocaBlock && "No alloca block found for this top level op");
builder.setInsertionPointToStart(allocaBlock);
- auto alloca = builder.create<fir::AllocaOp>(loc, descriptor.getType());
+
+ mlir::Type allocaType = descriptor.getType();
+ if (fir::isTypeWithDescriptor(allocaType) &&
+ !mlir::isa<fir::BaseBoxType>(descriptor.getType()))
+ allocaType = fir::unwrapRefType(allocaType);
+ auto alloca = builder.create<fir::AllocaOp>(loc, allocaType);
builder.restoreInsertionPoint(insPt);
// We should only emit a store if the passed in data is present, it is
// possible a user passes in no argument to an optional parameter, in which
@@ -159,8 +165,10 @@ class MapInfoFinalizationPass
auto isPresent =
builder.create<fir::IsPresentOp>(loc, builder.getI1Type(), descriptor);
builder.genIfOp(loc, {}, isPresent, false)
- .genThen(
- [&]() { builder.create<fir::StoreOp>(loc, descriptor, alloca); })
+ .genThen([&]() {
+ descriptor = builder.loadIfRef(loc, descriptor);
+ builder.create<fir::StoreOp>(loc, descriptor, alloca);
+ })
.end();
return slot = alloca;
}
diff --git a/flang/test/Lower/OpenMP/optional-argument-map-2.f90 b/flang/test/Lower/OpenMP/optional-argument-map-2.f90
new file mode 100644
index 0000000000000..eb89b18063f64
--- /dev/null
+++ b/flang/test/Lower/OpenMP/optional-argument-map-2.f90
@@ -0,0 +1,46 @@
+!RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
+
+module mod
+ implicit none
+contains
+ subroutine routine(a)
+ implicit none
+ real(4), allocatable, optional, intent(inout) :: a(:)
+ integer(4) :: i
+
+ !$omp target teams distribute parallel do shared(a)
+ do i=1,10
+ a(i) = i + a(i)
+ end do
+
+ end subroutine routine
+end module mod
+
+! CHECK-LABEL: func.func @_QMmodProutine(
+! CHECK-SAME: %[[ARG0:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {fir.bindc_name = "a", fir.optional}) {
+! CHECK: %[[VAL_0:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>>
+! CHECK: %[[VAL_1:.*]] = fir.dummy_scope : !fir.dscope
+! CHECK: %[[VAL_2:.*]]:2 = hlfir.declare %[[ARG0]] dummy_scope %[[VAL_1]] {fortran_attrs = #fir.var_attrs<allocatable, intent_inout, optional>, uniq_name = "_QMmodFroutineEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.dscope) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+! CHECK: %[[VAL_8:.*]] = fir.is_present %[[VAL_2]]#1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> i1
+! CHECK: %[[VAL_9:.*]]:5 = fir.if %[[VAL_8]] -> (index, index, index, index, index) {
+! CHECK: %[[VAL_10:.*]] = fir.load %[[VAL_2]]#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK: %[[VAL_11:.*]] = arith.constant 1 : index
+! CHECK: %[[VAL_12:.*]] = arith.constant 0 : index
+! CHECK: %[[VAL_13:.*]] = fir.load %[[VAL_2]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK: %[[VAL_14:.*]] = arith.constant 0 : index
+! CHECK: %[[VAL_15:.*]]:3 = fir.box_dims %[[VAL_13]], %[[VAL_14]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>, index) -> (index, index, index)
+! CHECK: %[[VAL_16:.*]]:3 = fir.box_dims %[[VAL_10]], %[[VAL_12]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>, index) -> (index, index, index)
+! CHECK: %[[VAL_17:.*]] = arith.constant 0 : index
+! CHECK: %[[VAL_18:.*]] = arith.subi %[[VAL_16]]#1, %[[VAL_11]] : index
+! CHECK: fir.result %[[VAL_17]], %[[VAL_18]], %[[VAL_16]]#1, %[[VAL_16]]#2, %[[VAL_15]]#0 : index, index, index, index, index
+! CHECK: } else {
+! CHECK: %[[VAL_19:.*]] = arith.constant 0 : index
+! CHECK: %[[VAL_20:.*]] = arith.constant -1 : index
+! CHECK: fir.result %[[VAL_19]], %[[VAL_20]], %[[VAL_19]], %[[VAL_19]], %[[VAL_19]] : index, index, index, index, index
+! CHECK: }
+! CHECK: %[[VAL_21:.*]] = omp.map.bounds lower_bound(%[[VAL_22:.*]]#0 : index) upper_bound(%[[VAL_22]]#1 : index) extent(%[[VAL_22]]#2 : index) stride(%[[VAL_22]]#3 : index) start_idx(%[[VAL_22]]#4 : index) {stride_in_bytes = true}
+! CHECK: %[[VAL_23:.*]] = fir.is_present %[[VAL_2]]#1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> i1
+! CHECK: fir.if %[[VAL_23]] {
+! CHECK: %[[VAL_24:.*]] = fir.load %[[VAL_2]]#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK: fir.store %[[VAL_24]] to %[[VAL_0]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK: }
diff --git a/offload/test/offloading/fortran/optional-mapped-arguments-2.f90 b/offload/test/offloading/fortran/optional-mapped-arguments-2.f90
new file mode 100644
index 0000000000000..0de6b7730d3a0
--- /dev/null
+++ b/offload/test/offloading/fortran/optional-mapped-arguments-2.f90
@@ -0,0 +1,57 @@
+! OpenMP offloading regression test that checks we do not cause a segfault when
+! implicitly mapping a not present optional allocatable function argument and
+! utilise it in the target region. No results requiring checking other than
+! that the program compiles and runs to completion with no error.
+! REQUIRES: flang, amdgpu
+
+! RUN: %libomptarget-compile-fortran-run-and-check-generic
+module mod
+ implicit none
+contains
+ subroutine routine(a, b)
+ implicit none
+ real(4), allocatable, optional, intent(in) :: a(:)
+ real(4), intent(out) :: b(:)
+ integer(4) :: i, ia
+ if(present(a)) then
+ ia = 1
+ write(*,*) "a is present"
+ else
+ ia=0
+ write(*,*) "a is not present"
+ end if
+
+ !$omp target teams distribute parallel do shared(a,b,ia)
+ do i=1,10
+ if (ia>0) then
+ b(i) = b(i) + a(i)
+ end if
+ end do
+
+ end subroutine routine
+
+end module mod
+
+program main
+ use mod
+ implicit none
+ real(4), allocatable :: a(:)
+ real(4), allocatable :: b(:)
+ integer(4) :: i
+ allocate(b(10))
+ do i=1,10
+ b(i)=0
+ end do
+ !$omp target data map(from: b)
+
+ call routine(b=b)
+
+ !$omp end target data
+
+ deallocate(b)
+
+ print *, "success, no segmentation fault"
+end program main
+
+!CHECK: a is not present
+!CHECK: success, no segmentation fault
``````````
</details>
https://github.com/llvm/llvm-project/pull/138210
More information about the flang-commits
mailing list