Added customMapper error propagation. Updated test.

Author: TIFitis

clang/lib/CodeGen/CGOpenMPRuntime.cpp

@@ -8879,17 +8879,17 @@ static void emitOffloadingArraysAndArgs(
   };
 
   auto CustomMapperCB = [&](unsigned int I) {
-    llvm::Value *MFunc = nullptr;
+    llvm::Function *MFunc = nullptr;
     if (CombinedInfo.Mappers[I]) {
       Info.HasMapper = true;
       MFunc = CGM.getOpenMPRuntime().getOrCreateUserDefinedMapperFunc(
           cast<OMPDeclareMapperDecl>(CombinedInfo.Mappers[I]));
     }
     return MFunc;
   };
-  OMPBuilder.emitOffloadingArraysAndArgs(
+  cantFail(OMPBuilder.emitOffloadingArraysAndArgs(
       AllocaIP, CodeGenIP, Info, Info.RTArgs, CombinedInfo, CustomMapperCB,
-      IsNonContiguous, ForEndCall, DeviceAddrCB);
+      IsNonContiguous, ForEndCall, DeviceAddrCB));
 }
 
 /// Check for inner distribute directive.
@@ -9082,26 +9082,25 @@ void CGOpenMPRuntime::emitUserDefinedMapper(const OMPDeclareMapperDecl *D,
     return CombinedInfo;
   };
 
-  auto CustomMapperCB = [&](unsigned I, llvm::Function **MapperFunc) {
+  auto CustomMapperCB = [&](unsigned I) {
+    llvm::Function *MapperFunc = nullptr;
     if (CombinedInfo.Mappers[I]) {
       // Call the corresponding mapper function.
-      *MapperFunc = getOrCreateUserDefinedMapperFunc(
+      MapperFunc = getOrCreateUserDefinedMapperFunc(
           cast<OMPDeclareMapperDecl>(CombinedInfo.Mappers[I]));
-      assert(*MapperFunc && "Expect a valid mapper function is available.");
-      return true;
+      assert(MapperFunc && "Expect a valid mapper function is available.");
     }
-    return false;
+    return MapperFunc;
   };
 
   SmallString<64> TyStr;
   llvm::raw_svector_ostream Out(TyStr);
   CGM.getCXXABI().getMangleContext().mangleCanonicalTypeName(Ty, Out);
   std::string Name = getName({"omp_mapper", TyStr, D->getName()});
 
-  llvm::Expected<llvm::Function *> NewFn = OMPBuilder.emitUserDefinedMapper(
-      PrivatizeAndGenMapInfoCB, ElemTy, Name, CustomMapperCB);
-  assert(NewFn && "Unexpected error in emitUserDefinedMapper");
-  UDMMap.try_emplace(D, *NewFn);
+  llvm::Function *NewFn = cantFail(OMPBuilder.emitUserDefinedMapper(
+      PrivatizeAndGenMapInfoCB, ElemTy, Name, CustomMapperCB));
+  UDMMap.try_emplace(D, NewFn);
   if (CGF)
     FunctionUDMMap[CGF->CurFn].push_back(D);
 }
@@ -10074,7 +10073,7 @@ void CGOpenMPRuntime::emitTargetDataCalls(
   };
 
   auto CustomMapperCB = [&](unsigned int I) {
-    llvm::Value *MFunc = nullptr;
+    llvm::Function *MFunc = nullptr;
     if (CombinedInfo.Mappers[I]) {
       Info.HasMapper = true;
       MFunc = CGF.CGM.getOpenMPRuntime().getOrCreateUserDefinedMapperFunc(

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -2408,6 +2408,11 @@ class OpenMPIRBuilder {
   using EmitFallbackCallbackTy =
       function_ref<InsertPointOrErrorTy(InsertPointTy)>;
 
+  // Callback function type for emitting and fetching user defined custom
+  // mappers.
+  using CustomMapperCallbackTy =
+      function_ref<Expected<Function *>(unsigned int)>;
+
   /// Generate a target region entry call and host fallback call.
   ///
   /// \param Loc The location at which the request originated and is fulfilled.
@@ -2474,9 +2479,9 @@ class OpenMPIRBuilder {
   /// return nullptr by reference. Accepts a reference to a MapInfosTy object
   /// that contains information generated for mappable clauses,
   /// including base pointers, pointers, sizes, map types, user-defined mappers.
-  void emitOffloadingArrays(
+  Error emitOffloadingArrays(
       InsertPointTy AllocaIP, InsertPointTy CodeGenIP, MapInfosTy &CombinedInfo,
-      TargetDataInfo &Info, function_ref<Value *(unsigned int)> CustomMapperCB,
+      TargetDataInfo &Info, CustomMapperCallbackTy CustomMapperCB,
       bool IsNonContiguous = false,
       function_ref<void(unsigned int, Value *)> DeviceAddrCB = nullptr);
 
@@ -2486,11 +2491,11 @@ class OpenMPIRBuilder {
   /// library. In essence, this function is a combination of
   /// emitOffloadingArrays and emitOffloadingArraysArgument and should arguably
   /// be preferred by clients of OpenMPIRBuilder.
-  void emitOffloadingArraysAndArgs(
+  Error emitOffloadingArraysAndArgs(
       InsertPointTy AllocaIP, InsertPointTy CodeGenIP, TargetDataInfo &Info,
       TargetDataRTArgs &RTArgs, MapInfosTy &CombinedInfo,
-      function_ref<Value *(unsigned int)> CustomMapperCB,
-      bool IsNonContiguous = false, bool ForEndCall = false,
+      CustomMapperCallbackTy CustomMapperCB, bool IsNonContiguous = false,
+      bool ForEndCall = false,
       function_ref<void(unsigned int, Value *)> DeviceAddrCB = nullptr);
 
   /// Creates offloading entry for the provided entry ID \a ID, address \a
@@ -2956,7 +2961,7 @@ class OpenMPIRBuilder {
           InsertPointTy CodeGenIP, llvm::Value *PtrPHI, llvm::Value *BeginArg)>
           PrivAndGenMapInfoCB,
       llvm::Type *ElemTy, StringRef FuncName,
-      function_ref<bool(unsigned int, Function **)> CustomMapperCB);
+      CustomMapperCallbackTy CustomMapperCB);
 
   /// Generator for '#omp target data'
   ///
@@ -2979,7 +2984,7 @@ class OpenMPIRBuilder {
       const LocationDescription &Loc, InsertPointTy AllocaIP,
       InsertPointTy CodeGenIP, Value *DeviceID, Value *IfCond,
       TargetDataInfo &Info, GenMapInfoCallbackTy GenMapInfoCB,
-      function_ref<Value *(unsigned int)> CustomMapperCB,
+      CustomMapperCallbackTy CustomMapperCB,
       omp::RuntimeFunction *MapperFunc = nullptr,
       function_ref<InsertPointOrErrorTy(InsertPointTy CodeGenIP,
                                         BodyGenTy BodyGenType)>
@@ -3028,7 +3033,7 @@ class OpenMPIRBuilder {
       SmallVectorImpl<Value *> &Inputs, GenMapInfoCallbackTy GenMapInfoCB,
       TargetBodyGenCallbackTy BodyGenCB,
       TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB,
-      function_ref<Value *(unsigned int)> CustomMapperCB,
+      CustomMapperCallbackTy CustomMapperCB,
       SmallVector<DependData> Dependencies = {}, bool HasNowait = false);
 
   /// Returns __kmpc_for_static_init_* runtime function for the specified

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -6555,8 +6555,7 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTargetData(
     const LocationDescription &Loc, InsertPointTy AllocaIP,
     InsertPointTy CodeGenIP, Value *DeviceID, Value *IfCond,
     TargetDataInfo &Info, GenMapInfoCallbackTy GenMapInfoCB,
-    function_ref<Value *(unsigned int)> CustomMapperCB,
-    omp::RuntimeFunction *MapperFunc,
+    CustomMapperCallbackTy CustomMapperCB, omp::RuntimeFunction *MapperFunc,
     function_ref<InsertPointOrErrorTy(InsertPointTy CodeGenIP,
                                       BodyGenTy BodyGenType)>
         BodyGenCB,
@@ -6585,9 +6584,10 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTargetData(
   auto BeginThenGen = [&](InsertPointTy AllocaIP,
                           InsertPointTy CodeGenIP) -> Error {
     MapInfo = &GenMapInfoCB(Builder.saveIP());
-    emitOffloadingArrays(AllocaIP, Builder.saveIP(), *MapInfo, Info,
-                         CustomMapperCB,
-                         /*IsNonContiguous=*/true, DeviceAddrCB);
+    if (Error Err = emitOffloadingArrays(
+            AllocaIP, Builder.saveIP(), *MapInfo, Info, CustomMapperCB,
+            /*IsNonContiguous=*/true, DeviceAddrCB))
+      return Err;
 
     TargetDataRTArgs RTArgs;
     emitOffloadingArraysArgument(Builder, RTArgs, Info);
@@ -7486,14 +7486,17 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::emitTargetTask(
   return Builder.saveIP();
 }
 
-void OpenMPIRBuilder::emitOffloadingArraysAndArgs(
+Error OpenMPIRBuilder::emitOffloadingArraysAndArgs(
     InsertPointTy AllocaIP, InsertPointTy CodeGenIP, TargetDataInfo &Info,
     TargetDataRTArgs &RTArgs, MapInfosTy &CombinedInfo,
-    function_ref<Value *(unsigned int)> CustomMapperCB, bool IsNonContiguous,
+    CustomMapperCallbackTy CustomMapperCB, bool IsNonContiguous,
     bool ForEndCall, function_ref<void(unsigned int, Value *)> DeviceAddrCB) {
-  emitOffloadingArrays(AllocaIP, CodeGenIP, CombinedInfo, Info, CustomMapperCB,
-                       IsNonContiguous, DeviceAddrCB);
+  if (Error Err =
+          emitOffloadingArrays(AllocaIP, CodeGenIP, CombinedInfo, Info,
+                               CustomMapperCB, IsNonContiguous, DeviceAddrCB))
+    return Err;
   emitOffloadingArraysArgument(Builder, RTArgs, Info, ForEndCall);
+  return Error::success();
 }
 
 static void
@@ -7505,7 +7508,7 @@ emitTargetCall(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
                Value *IfCond, Function *OutlinedFn, Constant *OutlinedFnID,
                SmallVectorImpl<Value *> &Args,
                OpenMPIRBuilder::GenMapInfoCallbackTy GenMapInfoCB,
-               function_ref<Value *(unsigned int)> CustomMapperCB,
+               OpenMPIRBuilder::CustomMapperCallbackTy CustomMapperCB,
                SmallVector<llvm::OpenMPIRBuilder::DependData> Dependencies,
                bool HasNoWait) {
   // Generate a function call to the host fallback implementation of the target
@@ -7580,10 +7583,11 @@ emitTargetCall(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
           OpenMPIRBuilder::InsertPointTy CodeGenIP) -> Error {
     OpenMPIRBuilder::MapInfosTy &MapInfo = GenMapInfoCB(Builder.saveIP());
     OpenMPIRBuilder::TargetDataRTArgs RTArgs;
-    OMPBuilder.emitOffloadingArraysAndArgs(AllocaIP, Builder.saveIP(), Info,
-                                           RTArgs, MapInfo, CustomMapperCB,
-                                           /*IsNonContiguous=*/true,
-                                           /*ForEndCall=*/false);
+    if (Error Err = OMPBuilder.emitOffloadingArraysAndArgs(
+            AllocaIP, Builder.saveIP(), Info, RTArgs, MapInfo, CustomMapperCB,
+            /*IsNonContiguous=*/true,
+            /*ForEndCall=*/false))
+      return Err;
 
     SmallVector<Value *, 3> NumTeamsC;
     for (auto [DefaultVal, RuntimeVal] :
@@ -7692,8 +7696,8 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTarget(
     SmallVectorImpl<Value *> &Inputs, GenMapInfoCallbackTy GenMapInfoCB,
     OpenMPIRBuilder::TargetBodyGenCallbackTy CBFunc,
     OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB,
-    function_ref<Value *(unsigned int)> CustomMapperCB,
-    SmallVector<DependData> Dependencies, bool HasNowait) {
+    CustomMapperCallbackTy CustomMapperCB, SmallVector<DependData> Dependencies,
+    bool HasNowait) {
 
   if (!updateToLocation(Loc))
     return InsertPointTy();
@@ -8045,8 +8049,7 @@ Expected<Function *> OpenMPIRBuilder::emitUserDefinedMapper(
     function_ref<MapInfosOrErrorTy(InsertPointTy CodeGenIP, llvm::Value *PtrPHI,
                                    llvm::Value *BeginArg)>
         GenMapInfoCB,
-    Type *ElemTy, StringRef FuncName,
-    function_ref<bool(unsigned int, Function **)> CustomMapperCB) {
+    Type *ElemTy, StringRef FuncName, CustomMapperCallbackTy CustomMapperCB) {
   SmallVector<Type *> Params;
   Params.emplace_back(Builder.getPtrTy());
   Params.emplace_back(Builder.getPtrTy());
@@ -8226,17 +8229,19 @@ Expected<Function *> OpenMPIRBuilder::emitUserDefinedMapper(
 
     Value *OffloadingArgs[] = {MapperHandle, CurBaseArg, CurBeginArg,
                                CurSizeArg,   CurMapType, CurNameArg};
-    Function *ChildMapperFn = nullptr;
-    if (CustomMapperCB && CustomMapperCB(I, &ChildMapperFn)) {
+
+    auto ChildMapperFn = CustomMapperCB(I);
+    if (!ChildMapperFn)
+      return ChildMapperFn.takeError();
+    if (*ChildMapperFn)
       // Call the corresponding mapper function.
-      Builder.CreateCall(ChildMapperFn, OffloadingArgs)->setDoesNotThrow();
-    } else {
+      Builder.CreateCall(*ChildMapperFn, OffloadingArgs)->setDoesNotThrow();
+    else
       // Call the runtime API __tgt_push_mapper_component to fill up the runtime
       // data structure.
       Builder.CreateCall(
           getOrCreateRuntimeFunction(M, OMPRTL___tgt_push_mapper_component),
           OffloadingArgs);
-    }
   }
 
   // Update the pointer to point to the next element that needs to be mapped,
@@ -8263,9 +8268,9 @@ Expected<Function *> OpenMPIRBuilder::emitUserDefinedMapper(
   return MapperFn;
 }
 
-void OpenMPIRBuilder::emitOffloadingArrays(
+Error OpenMPIRBuilder::emitOffloadingArrays(
     InsertPointTy AllocaIP, InsertPointTy CodeGenIP, MapInfosTy &CombinedInfo,
-    TargetDataInfo &Info, function_ref<Value *(unsigned int)> CustomMapperCB,
+    TargetDataInfo &Info, CustomMapperCallbackTy CustomMapperCB,
     bool IsNonContiguous,
     function_ref<void(unsigned int, Value *)> DeviceAddrCB) {
 
@@ -8274,7 +8279,7 @@ void OpenMPIRBuilder::emitOffloadingArrays(
   Info.NumberOfPtrs = CombinedInfo.BasePointers.size();
 
   if (Info.NumberOfPtrs == 0)
-    return;
+    return Error::success();
 
   Builder.restoreIP(AllocaIP);
   // Detect if we have any capture size requiring runtime evaluation of the
@@ -8438,9 +8443,13 @@ void OpenMPIRBuilder::emitOffloadingArrays(
     // Fill up the mapper array.
     unsigned IndexSize = M.getDataLayout().getIndexSizeInBits(0);
     Value *MFunc = ConstantPointerNull::get(PtrTy);
-    if (CustomMapperCB)
-      if (Value *CustomMFunc = CustomMapperCB(I))
-        MFunc = Builder.CreatePointerCast(CustomMFunc, PtrTy);
+
+    auto CustomMFunc = CustomMapperCB(I);
+    if (!CustomMFunc)
+      return CustomMFunc.takeError();
+    if (*CustomMFunc)
+      MFunc = Builder.CreatePointerCast(*CustomMFunc, PtrTy);
+
     Value *MAddr = Builder.CreateInBoundsGEP(
         MappersArray->getAllocatedType(), MappersArray,
         {Builder.getIntN(IndexSize, 0), Builder.getIntN(IndexSize, I)});
@@ -8450,8 +8459,9 @@ void OpenMPIRBuilder::emitOffloadingArrays(
 
   if (!IsNonContiguous || CombinedInfo.NonContigInfo.Offsets.empty() ||
       Info.NumberOfPtrs == 0)
-    return;
+    return Error::success();
   emitNonContiguousDescriptor(AllocaIP, CodeGenIP, CombinedInfo, Info);
+  return Error::success();
 }
 
 void OpenMPIRBuilder::emitBranch(BasicBlock *Target) {

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -3608,16 +3608,17 @@ emitUserDefinedMapper(Operation *op, llvm::IRBuilderBase &builder,
     return combinedInfo;
   };
 
-  auto customMapperCB = [&](unsigned i, llvm::Function **mapperFunc) {
+  auto customMapperCB = [&](unsigned i) -> llvm::Expected<llvm::Function *> {
+    llvm::Function *mapperFunc = nullptr;
     if (combinedInfo.Mappers[i]) {
       // Call the corresponding mapper function.
       llvm::Expected<llvm::Function *> newFn = getOrCreateUserDefinedMapperFunc(
           combinedInfo.Mappers[i], builder, moduleTranslation);
-      assert(newFn && "Expect a valid mapper function is available");
-      *mapperFunc = *newFn;
-      return true;
+      if (!newFn)
+        return newFn.takeError();
+      mapperFunc = *newFn;
     }
-    return false;
+    return mapperFunc;
   };
 
   llvm::Expected<llvm::Function *> newFn = ompBuilder->emitUserDefinedMapper(
@@ -3840,13 +3841,15 @@ convertOmpTargetData(Operation *op, llvm::IRBuilderBase &builder,
     return builder.saveIP();
   };
 
-  auto customMapperCB = [&](unsigned int i) {
+  auto customMapperCB =
+      [&](unsigned int i) -> llvm::Expected<llvm::Function *> {
     llvm::Function *mapperFunc = nullptr;
     if (combinedInfo.Mappers[i]) {
       info.HasMapper = true;
       llvm::Expected<llvm::Function *> newFn = getOrCreateUserDefinedMapperFunc(
           combinedInfo.Mappers[i], builder, moduleTranslation);
-      assert(newFn && "Expect a valid mapper function is available");
+      if (!newFn)
+        return newFn.takeError();
       mapperFunc = *newFn;
     }
     return mapperFunc;
@@ -4551,13 +4554,15 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
       /*RequiresDevicePointerInfo=*/false,
       /*SeparateBeginEndCalls=*/true);
 
-  auto customMapperCB = [&](unsigned int i) {
-    llvm::Value *mapperFunc = nullptr;
+  auto customMapperCB =
+      [&](unsigned int i) -> llvm::Expected<llvm::Function *> {
+    llvm::Function *mapperFunc = nullptr;
     if (combinedInfos.Mappers[i]) {
       info.HasMapper = true;
       llvm::Expected<llvm::Function *> newFn = getOrCreateUserDefinedMapperFunc(
           combinedInfos.Mappers[i], builder, moduleTranslation);
-      assert(newFn && "Expect a valid mapper function is available");
+      if (!newFn)
+        return newFn.takeError();
       mapperFunc = *newFn;
     }
     return mapperFunc;

offload/test/offloading/fortran/target-custom-mapper.f90

@@ -10,26 +10,32 @@ program test_openmp_mapper
       integer :: data(n)
    end type mytype
 
+   type :: mytype2
+      type(mytype) :: my_data
+   end type mytype2
+
    ! Declare custom mappers for the derived type `mytype`
-   !$omp declare mapper(my_mapper1 : mytype :: t) map(to: t%data)
-   !$omp declare mapper(my_mapper2 : mytype :: t) map(mapper(my_mapper1): t%data)
+   !$omp declare mapper(my_mapper1 : mytype :: t) map(to: t%data(1 : n))
+
+   ! Declare custom mappers for the derived type `mytype2`
+   !$omp declare mapper(my_mapper2 : mytype2 :: t) map(mapper(my_mapper1): t%my_data)
 
-   type(mytype) :: obj
+   type(mytype2) :: obj
    integer :: i, sum_host, sum_device
 
    ! Initialize the host data
    do i = 1, n
-      obj%data(i) = 1
+      obj%my_data%data(i) = 1
    end do
 
    ! Compute the sum on the host for verification
-   sum_host = sum(obj%data)
+   sum_host = sum(obj%my_data%data)
 
    ! Offload computation to the device using the named mapper `my_mapper2`
    sum_device = 0
    !$omp target map(tofrom: sum_device) map(mapper(my_mapper2) : obj)
    do i = 1, n
-      sum_device = sum_device + obj%data(i)
+      sum_device = sum_device + obj%my_data%data(i)
    end do
    !$omp end target