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

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

@@ -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,

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -2322,7 +2322,8 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
 
       fir::factory::AddrAndBoundsInfo info =
           Fortran::lower::getDataOperandBaseAddr(
-              converter, firOpBuilder, sym, converter.getCurrentLocation());
+              converter, firOpBuilder, sym.GetUltimate(),
+              converter.getCurrentLocation());
       llvm::SmallVector<mlir::Value> bounds =
           fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
                                              mlir::omp::MapBoundsType>(

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,11 @@ 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::isBoxAddress(allocaType))
+      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 +164,10 @@ class MapInfoFinalizationPass
     auto isPresent =
         builder.create<fir::IsPresentOp>(loc, builder.getI1Type(), descriptor);
     builder.genIfOp(loc, {}, isPresent, false)
-        .genThen(
+        .genThen([&]() {
-            [&]() { builder.create<fir::StoreOp>(loc, descriptor, alloca); })
+          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]]#0 : !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_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}
+! 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