[Flang][OpenMP] Generate correct present checks for implicit maps of optional allocatables

    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).

Author: agozillon

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);

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,16 +152,23 @@ 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
     // case we cannot store or we'll segfault on the emitted memcpy.
     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;
   }

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:           }

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