[OpenMP] Codegen aggregate for outlined function captures

Parallel regions are outlined as functions with capture variables explicitly generated as distinct parameters in the function's argument list. That complicates the fork_call interface in the OpenMP runtime: (1) the fork_call is variadic since there is a variable number of arguments to forward to the outlined function, (2) wrapping/unwrapping arguments happens in the OpenMP runtime, which is sub-optimal, has been a source of ABI bugs, and has a hardcoded limit (16) in the number of arguments, (3)  forwarded arguments must cast to pointer types, which complicates debugging. This patch avoids those issues by aggregating captured arguments in a struct to pass to the fork_call.

Reviewed By: jdoerfert, jhuber6

Differential Revision: https://reviews.llvm.org/D102107

Author: ggeorgakoudis

clang/lib/CodeGen/CGOpenMPRuntime.cpp

@@ -1284,7 +1284,7 @@ static llvm::Function *emitParallelOrTeamsOutlinedFunction(
   CGOpenMPOutlinedRegionInfo CGInfo(*CS, ThreadIDVar, CodeGen, InnermostKind,
                                     HasCancel, OutlinedHelperName);
   CodeGenFunction::CGCapturedStmtRAII CapInfoRAII(CGF, &CGInfo);
-  return CGF.GenerateOpenMPCapturedStmtFunction(*CS, D.getBeginLoc());
+  return CGF.GenerateOpenMPCapturedStmtFunctionAggregate(*CS, D.getBeginLoc());
 }
 
 llvm::Function *CGOpenMPRuntime::emitParallelOutlinedFunction(

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

@@ -1523,21 +1523,49 @@ void CGOpenMPRuntimeGPU::emitParallelCall(CodeGenFunction &CGF,
     // TODO: Is that needed?
     CodeGenFunction::OMPPrivateScope PrivateArgScope(CGF);
 
+    // Store addresses of global arguments to pass to the parallel call.
     Address CapturedVarsAddrs = CGF.CreateDefaultAlignTempAlloca(
         llvm::ArrayType::get(CGM.VoidPtrTy, CapturedVars.size()),
         "captured_vars_addrs");
-    // There's something to share.
+
+    // Store globalized values to push, pop through the global stack.
+    llvm::SmallDenseMap<llvm::Value *, unsigned> GlobalValuesToSizeMap;
     if (!CapturedVars.empty()) {
-      // Prepare for parallel region. Indicate the outlined function.
       ASTContext &Ctx = CGF.getContext();
       unsigned Idx = 0;
       for (llvm::Value *V : CapturedVars) {
         Address Dst = Bld.CreateConstArrayGEP(CapturedVarsAddrs, Idx);
         llvm::Value *PtrV;
         if (V->getType()->isIntegerTy())
           PtrV = Bld.CreateIntToPtr(V, CGF.VoidPtrTy);
-        else
-          PtrV = Bld.CreatePointerBitCastOrAddrSpaceCast(V, CGF.VoidPtrTy);
+        else {
+          assert(V->getType()->isPointerTy() &&
+                 "Expected Pointer Type to globalize.");
+          // Globalize and store pointer.
+          llvm::Type *PtrElemTy = V->getType()->getPointerElementType();
+          auto &DL = CGM.getDataLayout();
+          unsigned GlobalSize = DL.getTypeAllocSize(PtrElemTy);
+
+          // Use shared memory to store globalized pointer values, for now this
+          // should be the outlined args aggregate struct.
+          llvm::Value *GlobalSizeArg[] = {
+              llvm::ConstantInt::get(CGM.SizeTy, GlobalSize)};
+          llvm::Value *GlobalValue = CGF.EmitRuntimeCall(
+              OMPBuilder.getOrCreateRuntimeFunction(CGM.getModule(),
+                                                    OMPRTL___kmpc_alloc_shared),
+              GlobalSizeArg);
+          GlobalValuesToSizeMap[GlobalValue] = GlobalSize;
+
+          llvm::Value *CapturedVarVal = Bld.CreateAlignedLoad(
+              PtrElemTy, V, DL.getABITypeAlign(PtrElemTy));
+          llvm::Value *GlobalValueCast =
+              Bld.CreatePointerBitCastOrAddrSpaceCast(
+                  GlobalValue, PtrElemTy->getPointerTo());
+          Bld.CreateDefaultAlignedStore(CapturedVarVal, GlobalValueCast);
+
+          PtrV = Bld.CreatePointerBitCastOrAddrSpaceCast(GlobalValue,
+                                                         CGF.VoidPtrTy);
+        }
         CGF.EmitStoreOfScalar(PtrV, Dst, /*Volatile=*/false,
                               Ctx.getPointerType(Ctx.VoidPtrTy));
         ++Idx;
@@ -1550,8 +1578,9 @@ void CGOpenMPRuntimeGPU::emitParallelCall(CodeGenFunction &CGF,
                                     /* isSigned */ false);
     else
       IfCondVal = llvm::ConstantInt::get(CGF.Int32Ty, 1);
-
     assert(IfCondVal && "Expected a value");
+
+    // Create the parallel call.
     llvm::Value *RTLoc = emitUpdateLocation(CGF, Loc);
     llvm::Value *Args[] = {
         RTLoc,
@@ -1567,6 +1596,14 @@ void CGOpenMPRuntimeGPU::emitParallelCall(CodeGenFunction &CGF,
     CGF.EmitRuntimeCall(OMPBuilder.getOrCreateRuntimeFunction(
                             CGM.getModule(), OMPRTL___kmpc_parallel_51),
                         Args);
+
+    // Pop any globalized values from the global stack.
+    for (const auto &GV : GlobalValuesToSizeMap) {
+      CGF.EmitRuntimeCall(
+          OMPBuilder.getOrCreateRuntimeFunction(CGM.getModule(),
+                                                OMPRTL___kmpc_free_shared),
+          {GV.first, llvm::ConstantInt::get(CGM.SizeTy, GV.second)});
+    }
   };
 
   RegionCodeGenTy RCG(ParallelGen);
@@ -3477,7 +3514,6 @@ llvm::Function *CGOpenMPRuntimeGPU::createParallelDataSharingWrapper(
                     D.getBeginLoc(), D.getBeginLoc());
 
   const auto *RD = CS.getCapturedRecordDecl();
-  auto CurField = RD->field_begin();
 
   Address ZeroAddr = CGF.CreateDefaultAlignTempAlloca(CGF.Int32Ty,
                                                       /*Name=*/".zero.addr");
@@ -3489,7 +3525,6 @@ llvm::Function *CGOpenMPRuntimeGPU::createParallelDataSharingWrapper(
   Args.emplace_back(ZeroAddr.getPointer());
 
   CGBuilderTy &Bld = CGF.Builder;
-  auto CI = CS.capture_begin();
 
   // Use global memory for data sharing.
   // Handle passing of global args to workers.
@@ -3504,55 +3539,33 @@ llvm::Function *CGOpenMPRuntimeGPU::createParallelDataSharingWrapper(
   // Retrieve the shared variables from the list of references returned
   // by the runtime. Pass the variables to the outlined function.
   Address SharedArgListAddress = Address::invalid();
-  if (CS.capture_size() > 0 ||
-      isOpenMPLoopBoundSharingDirective(D.getDirectiveKind())) {
+  if (CS.capture_size() > 0) {
     SharedArgListAddress = CGF.EmitLoadOfPointer(
         GlobalArgs, CGF.getContext()
                         .getPointerType(CGF.getContext().getPointerType(
                             CGF.getContext().VoidPtrTy))
                         .castAs<PointerType>());
-  }
-  unsigned Idx = 0;
-  if (isOpenMPLoopBoundSharingDirective(D.getDirectiveKind())) {
-    Address Src = Bld.CreateConstInBoundsGEP(SharedArgListAddress, Idx);
+    const auto *CI = CS.capture_begin();
+    // Load the outlined arg aggregate struct.
+    ASTContext &CGFContext = CGF.getContext();
+    QualType RecordPointerTy =
+        CGFContext.getPointerType(CGFContext.getRecordType(RD));
+    Address Src = Bld.CreateConstInBoundsGEP(SharedArgListAddress, /*Index=*/0);
     Address TypedAddress = Bld.CreatePointerBitCastOrAddrSpaceCast(
-        Src, CGF.SizeTy->getPointerTo());
-    llvm::Value *LB = CGF.EmitLoadOfScalar(
-        TypedAddress,
-        /*Volatile=*/false,
-        CGF.getContext().getPointerType(CGF.getContext().getSizeType()),
-        cast<OMPLoopDirective>(D).getLowerBoundVariable()->getExprLoc());
-    Args.emplace_back(LB);
-    ++Idx;
-    Src = Bld.CreateConstInBoundsGEP(SharedArgListAddress, Idx);
-    TypedAddress = Bld.CreatePointerBitCastOrAddrSpaceCast(
-        Src, CGF.SizeTy->getPointerTo());
-    llvm::Value *UB = CGF.EmitLoadOfScalar(
+        Src, CGF.ConvertTypeForMem(CGFContext.getPointerType(RecordPointerTy)));
+    llvm::Value *Arg = CGF.EmitLoadOfScalar(
         TypedAddress,
-        /*Volatile=*/false,
-        CGF.getContext().getPointerType(CGF.getContext().getSizeType()),
-        cast<OMPLoopDirective>(D).getUpperBoundVariable()->getExprLoc());
-    Args.emplace_back(UB);
-    ++Idx;
-  }
-  if (CS.capture_size() > 0) {
+        /*Volatile=*/false, CGFContext.getPointerType(RecordPointerTy),
+        CI->getLocation());
+    Args.emplace_back(Arg);
+  } else {
+    // If there are no captured arguments, use nullptr.
     ASTContext &CGFContext = CGF.getContext();
-    for (unsigned I = 0, E = CS.capture_size(); I < E; ++I, ++CI, ++CurField) {
-      QualType ElemTy = CurField->getType();
-      Address Src = Bld.CreateConstInBoundsGEP(SharedArgListAddress, I + Idx);
-      Address TypedAddress = Bld.CreatePointerBitCastOrAddrSpaceCast(
-          Src, CGF.ConvertTypeForMem(CGFContext.getPointerType(ElemTy)));
-      llvm::Value *Arg = CGF.EmitLoadOfScalar(TypedAddress,
-                                              /*Volatile=*/false,
-                                              CGFContext.getPointerType(ElemTy),
-                                              CI->getLocation());
-      if (CI->capturesVariableByCopy() &&
-          !CI->getCapturedVar()->getType()->isAnyPointerType()) {
-        Arg = castValueToType(CGF, Arg, ElemTy, CGFContext.getUIntPtrType(),
-                              CI->getLocation());
-      }
-      Args.emplace_back(Arg);
-    }
+    QualType RecordPointerTy =
+        CGFContext.getPointerType(CGFContext.getRecordType(RD));
+    llvm::Value *Arg =
+        llvm::Constant::getNullValue(CGF.ConvertTypeForMem(RecordPointerTy));
+    Args.emplace_back(Arg);
   }
 
   emitOutlinedFunctionCall(CGF, D.getBeginLoc(), OutlinedParallelFn, Args);