Skip to content

Commit b291cfc

Browse files
agozillonagozillon
authored
[Flang][OpenMP] Generate correct present checks for implicit maps of optional allocatables (#138210)
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).
1 parent dd42112 commit b291cfc

File tree

5 files changed

+127
-5
lines changed

5 files changed

+127
-5
lines changed

‎flang/include/flang/Optimizer/Builder/DirectivesCommon.h

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,17 @@ genBaseBoundsOps(fir::FirOpBuilder &builder, mlir::Location loc,
243243
return bounds;
244244
}
245245

246+
/// Checks if an argument is optional based on the fortran attributes
247+
/// that are tied to it.
248+
inline bool isOptionalArgument(mlir::Operation *op) {
249+
if (auto declareOp = mlir::dyn_cast_or_null<hlfir::DeclareOp>(op))
250+
if (declareOp.getFortranAttrs() &&
251+
bitEnumContainsAny(*declareOp.getFortranAttrs(),
252+
fir::FortranVariableFlagsEnum::optional))
253+
return true;
254+
return false;
255+
}
256+
246257
template <typename BoundsOp, typename BoundsType>
247258
llvm::SmallVector<mlir::Value>
248259
genImplicitBoundsOps(fir::FirOpBuilder &builder, AddrAndBoundsInfo &info,

‎flang/lib/Lower/OpenMP/OpenMP.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2322,7 +2322,8 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
23222322

23232323
fir::factory::AddrAndBoundsInfo info =
23242324
Fortran::lower::getDataOperandBaseAddr(
2325-
converter, firOpBuilder, sym, converter.getCurrentLocation());
2325+
converter, firOpBuilder, sym.GetUltimate(),
2326+
converter.getCurrentLocation());
23262327
llvm::SmallVector<mlir::Value> bounds =
23272328
fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
23282329
mlir::omp::MapBoundsType>(

‎flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,8 @@ class MapInfoFinalizationPass
131131
boxMap.getVarPtr().getDefiningOp()))
132132
descriptor = addrOp.getVal();
133133

134-
if (!mlir::isa<fir::BaseBoxType>(descriptor.getType()))
134+
if (!mlir::isa<fir::BaseBoxType>(descriptor.getType()) &&
135+
!fir::factory::isOptionalArgument(descriptor.getDefiningOp()))
135136
return descriptor;
136137

137138
mlir::Value &slot = localBoxAllocas[descriptor.getDefiningOp()];
@@ -151,16 +152,22 @@ class MapInfoFinalizationPass
151152
mlir::Location loc = boxMap->getLoc();
152153
assert(allocaBlock && "No alloca block found for this top level op");
153154
builder.setInsertionPointToStart(allocaBlock);
154-
auto alloca = builder.create<fir::AllocaOp>(loc, descriptor.getType());
155+
156+
mlir::Type allocaType = descriptor.getType();
157+
if (fir::isBoxAddress(allocaType))
158+
allocaType = fir::unwrapRefType(allocaType);
159+
auto alloca = builder.create<fir::AllocaOp>(loc, allocaType);
155160
builder.restoreInsertionPoint(insPt);
156161
// We should only emit a store if the passed in data is present, it is
157162
// possible a user passes in no argument to an optional parameter, in which
158163
// case we cannot store or we'll segfault on the emitted memcpy.
159164
auto isPresent =
160165
builder.create<fir::IsPresentOp>(loc, builder.getI1Type(), descriptor);
161166
builder.genIfOp(loc, {}, isPresent, false)
162-
.genThen(
163-
[&]() { builder.create<fir::StoreOp>(loc, descriptor, alloca); })
167+
.genThen([&]() {
168+
descriptor = builder.loadIfRef(loc, descriptor);
169+
builder.create<fir::StoreOp>(loc, descriptor, alloca);
170+
})
164171
.end();
165172
return slot = alloca;
166173
}

‎flang/test/Lower/OpenMP/optional-argument-map-2.f90

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
!RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
2+
3+
module mod
4+
implicit none
5+
contains
6+
subroutine routine(a)
7+
implicit none
8+
real(4), allocatable, optional, intent(inout) :: a(:)
9+
integer(4) :: i
10+
11+
!$omp target teams distribute parallel do shared(a)
12+
do i=1,10
13+
a(i) = i + a(i)
14+
end do
15+
16+
end subroutine routine
17+
end module mod
18+
19+
! CHECK-LABEL: func.func @_QMmodProutine(
20+
! CHECK-SAME: %[[ARG0:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {fir.bindc_name = "a", fir.optional}) {
21+
! CHECK: %[[VAL_0:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>>
22+
! CHECK: %[[VAL_1:.*]] = fir.dummy_scope : !fir.dscope
23+
! 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>>>>)
24+
! CHECK: %[[VAL_8:.*]] = fir.is_present %[[VAL_2]]#1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> i1
25+
! CHECK: %[[VAL_9:.*]]:5 = fir.if %[[VAL_8]] -> (index, index, index, index, index) {
26+
! CHECK: %[[VAL_10:.*]] = fir.load %[[VAL_2]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
27+
! CHECK: %[[VAL_11:.*]] = arith.constant 1 : index
28+
! CHECK: %[[VAL_12:.*]] = arith.constant 0 : index
29+
! CHECK: %[[VAL_13:.*]] = fir.load %[[VAL_2]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
30+
! CHECK: %[[VAL_14:.*]] = arith.constant 0 : index
31+
! CHECK: %[[VAL_15:.*]]:3 = fir.box_dims %[[VAL_13]], %[[VAL_14]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>, index) -> (index, index, index)
32+
! CHECK: %[[VAL_16:.*]]:3 = fir.box_dims %[[VAL_10]], %[[VAL_12]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>, index) -> (index, index, index)
33+
! CHECK: %[[VAL_17:.*]] = arith.constant 0 : index
34+
! CHECK: %[[VAL_18:.*]] = arith.subi %[[VAL_16]]#1, %[[VAL_11]] : index
35+
! CHECK: fir.result %[[VAL_17]], %[[VAL_18]], %[[VAL_16]]#1, %[[VAL_16]]#2, %[[VAL_15]]#0 : index, index, index, index, index
36+
! CHECK: } else {
37+
! CHECK: %[[VAL_19:.*]] = arith.constant 0 : index
38+
! CHECK: %[[VAL_20:.*]] = arith.constant -1 : index
39+
! CHECK: fir.result %[[VAL_19]], %[[VAL_20]], %[[VAL_19]], %[[VAL_19]], %[[VAL_19]] : index, index, index, index, index
40+
! CHECK: }
41+
! CHECK: %[[VAL_21:.*]] = omp.map.bounds lower_bound(%[[VAL_9]]#0 : index) upper_bound(%[[VAL_9]]#1 : index) extent(%[[VAL_9]]#2 : index) stride(%[[VAL_9]]#3 : index) start_idx(%[[VAL_9]]#4 : index) {stride_in_bytes = true}
42+
! CHECK: %[[VAL_23:.*]] = fir.is_present %[[VAL_2]]#1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> i1
43+
! CHECK: fir.if %[[VAL_23]] {
44+
! CHECK: %[[VAL_24:.*]] = fir.load %[[VAL_2]]#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
45+
! CHECK: fir.store %[[VAL_24]] to %[[VAL_0]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
46+
! CHECK: }

‎offload/test/offloading/fortran/optional-mapped-arguments-2.f90

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
! OpenMP offloading regression test that checks we do not cause a segfault when
2+
! implicitly mapping a not present optional allocatable function argument and
3+
! utilise it in the target region. No results requiring checking other than
4+
! that the program compiles and runs to completion with no error.
5+
! REQUIRES: flang, amdgpu
6+
7+
! RUN: %libomptarget-compile-fortran-run-and-check-generic
8+
module mod
9+
implicit none
10+
contains
11+
subroutine routine(a, b)
12+
implicit none
13+
real(4), allocatable, optional, intent(in) :: a(:)
14+
real(4), intent(out) :: b(:)
15+
integer(4) :: i, ia
16+
if(present(a)) then
17+
ia = 1
18+
write(*,*) "a is present"
19+
else
20+
ia=0
21+
write(*,*) "a is not present"
22+
end if
23+
24+
!$omp target teams distribute parallel do shared(a,b,ia)
25+
do i=1,10
26+
if (ia>0) then
27+
b(i) = b(i) + a(i)
28+
end if
29+
end do
30+
31+
end subroutine routine
32+
33+
end module mod
34+
35+
program main
36+
use mod
37+
implicit none
38+
real(4), allocatable :: a(:)
39+
real(4), allocatable :: b(:)
40+
integer(4) :: i
41+
allocate(b(10))
42+
do i=1,10
43+
b(i)=0
44+
end do
45+
!$omp target data map(from: b)
46+
47+
call routine(b=b)
48+
49+
!$omp end target data
50+
51+
deallocate(b)
52+
53+
print *, "success, no segmentation fault"
54+
end program main
55+
56+
!CHECK: a is not present
57+
!CHECK: success, no segmentation fault

0 commit comments

Comments
 (0)