[OpenMP] Add callback annotation to `__kmpc_parallel_51`

[jdoerfert]

Summary

Regardless of the linking of the device runtime we want to ensure the
__kmpc_parallel_51 will have callback annotations.

Depends on D102107.

[jdoerfert]

Update tests, though, need to manually apply the effect of D102107 once that landed instead.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

if (!Fn) {
     switch (FnID) {
 #define OMP_RTL(Enum, Str, ...)                                                \
   case Enum:                                                                   \
     Fn = Function::Create(FnTy, GlobalValue::ExternalLinkage, Str, M);         \
     break;
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
     }
 
-    // Add information if the runtime function takes a callback function
-    if (FnID == OMPRTL___kmpc_fork_call || FnID == OMPRTL___kmpc_fork_teams) {
-      if (!Fn->hasMetadata(LLVMContext::MD_callback)) {
-        LLVMContext &Ctx = Fn->getContext();
-        MDBuilder MDB(Ctx);
-        // Annotate the callback behavior of the runtime function:
-        //  - The callback callee is argument number 2 (microtask).
-        //  - The first two arguments of the callback callee are unknown (-1).
-        //  - All variadic arguments to the runtime function are passed to the
-        //    callback callee.
-        Fn->addMetadata(
-            LLVMContext::MD_callback,
-            *MDNode::get(Ctx, {MDB.createCallbackEncoding(
-                                  2, {-1, -1}, /* VarArgsArePassed */ true)}));
-      }
+    LLVM_DEBUG(dbgs() << "Created OpenMP runtime function " << Fn->getName()
+                      << " with type " << *Fn->getFunctionType() << "\n");
+    addAttributes(FnID, *Fn);
+
+  } else {
+    LLVM_DEBUG(dbgs() << "Found OpenMP runtime function " << Fn->getName()
+                      << " with type " << *Fn->getFunctionType() << "\n");
+  }
+
+  } else if (FnID == OMPRTL___kmpc_parallel_51) {
+    if (!Fn->hasMetadata(LLVMContext::MD_callback)) {
+      LLVMContext &Ctx = Fn->getContext();
+      MDBuilder MDB(Ctx);
+      // Annotate the callback behavior of the runtime function:
+      //  - The callback callee is argument number 5 (outlined function).
+      //  - The first two arguments of the callback callee are unknown (-1).
+      //  - Argument 7, (args) is passed next, no variadic arguments are used.
+      Fn->addMetadata(LLVMContext::MD_callback,
+                      *MDNode::get(Ctx, {MDB.createCallbackEncoding(
+                                            5, {-1, -1, 7},
+                                            /* VarArgsArePassed */ false)}));
     }
-
-    LLVM_DEBUG(dbgs() << "Created OpenMP runtime function " << Fn->getName()
-                      << " with type " << *Fn->getFunctionType() << "\n");
-    addAttributes(FnID, *Fn);
-
-  } else {
-    LLVM_DEBUG(dbgs() << "Found OpenMP runtime function " << Fn->getName()
-                      << " with type " << *Fn->getFunctionType() << "\n");
   }
 
   assert(Fn && "Failed to create OpenMP runtime function");
 
   // Cast the function to the expected type if necessary
   Constant *C = ConstantExpr::getBitCast(Fn, FnTy->getPointerTo());
   return {FnTy, C};
 }

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 
 struct AAICVTracker;
 
 /// OpenMP specific information. For now, stores RFIs and ICVs also needed for
 /// Attributor runs.
 struct OMPInformationCache : public InformationCache {
   OMPInformationCache(Module &M, AnalysisGetter &AG,
                       BumpPtrAllocator &Allocator, SetVector<Function *> &CGSCC,
-                      SmallPtrSetImpl<Kernel> &Kernels)
-      : InformationCache(M, AG, Allocator, &CGSCC), OMPBuilder(M),
+                      SmallPtrSetImpl<Kernel> &Kernels,
+                      OpenMPIRBuilder &OMPBuilder)
+      : InformationCache(M, AG, Allocator, &CGSCC), OMPBuilder(OMPBuilder),
         Kernels(Kernels) {
 
-    OMPBuilder.initialize();
     initializeRuntimeFunctions();
     initializeInternalControlVars();
   }
 
   /// Generic information that describes an internal control variable.
   struct InternalControlVarInfo {
     /// The kind, as described by InternalControlVar enum.
     InternalControlVar Kind;
struct OMPInformationCache : public InformationCache {
 
   public:
     /// Iterators for the uses of this runtime function.
     decltype(UsesMap)::iterator begin() { return UsesMap.begin(); }
     decltype(UsesMap)::iterator end() { return UsesMap.end(); }
   };
 
   /// An OpenMP-IR-Builder instance
-  OpenMPIRBuilder OMPBuilder;
+  OpenMPIRBuilder &OMPBuilder;
 
   /// Map from runtime function kind to the runtime function description.
   EnumeratedArray<RuntimeFunctionInfo, RuntimeFunction,
                   RuntimeFunction::OMPRTL___last>
       RFIs;
 
   /// Map from function declarations/definitions to their runtime enum type.
   DenseMap<Function *, RuntimeFunction> RuntimeFunctionIDMap;
if (!containsOpenMP(M))
     return PreservedAnalyses::all();
   if (DisableOpenMPOptimizations)
     return PreservedAnalyses::all();
 
   FunctionAnalysisManager &FAM =
       AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   KernelSet Kernels = getDeviceKernels(M);
 
+  OpenMPIRBuilder OMPBuilder(M);
+  OMPBuilder.initialize();
+  // Make sure we have the callback annotations attached to the
+  // __kmpc_parallel_51 function declaration/definition.
+  OMPBuilder.getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_parallel_51);
+
   auto IsCalled = [&](Function &F) {
     if (Kernels.contains(&F))
       return true;
     for (const User *U : F.users())
       if (!isa<BlockAddress>(U))
         return true;
     return false;
   };
PreservedAnalyses OpenMPOptPass::run(Module &M, ModuleAnalysisManager &AM) {
   AnalysisGetter AG(FAM);
 
   auto OREGetter = [&FAM](Function *F) -> OptimizationRemarkEmitter & {
     return FAM.getResult<OptimizationRemarkEmitterAnalysis>(*F);
   };
 
   BumpPtrAllocator Allocator;
   CallGraphUpdater CGUpdater;
-
   SetVector<Function *> Functions(SCC.begin(), SCC.end());
-  OMPInformationCache InfoCache(M, AG, Allocator, /*CGSCC*/ Functions, Kernels);
+  OMPInformationCache InfoCache(M, AG, Allocator, /*CGSCC*/ Functions, Kernels,
+                                OMPBuilder);
 
   unsigned MaxFixpointIterations = (isOpenMPDevice(M)) ? 128 : 32;
   Attributor A(Functions, InfoCache, CGUpdater, nullptr, true, false,
                MaxFixpointIterations, OREGetter, DEBUG_TYPE);
 
   OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
   bool Changed = OMPOpt.run(true);
 
PreservedAnalyses OpenMPOptCGSCCPass::run(LazyCallGraph::SCC &C,
   auto OREGetter = [&FAM](Function *F) -> OptimizationRemarkEmitter & {
     return FAM.getResult<OptimizationRemarkEmitterAnalysis>(*F);
   };
 
   BumpPtrAllocator Allocator;
   CallGraphUpdater CGUpdater;
   CGUpdater.initialize(CG, C, AM, UR);
 
+  OpenMPIRBuilder OMPBuilder(M);
+  OMPBuilder.initialize();
   SetVector<Function *> Functions(SCC.begin(), SCC.end());
-  OMPInformationCache InfoCache(*(Functions.back()->getParent()), AG, Allocator,
-                                /*CGSCC*/ Functions, Kernels);
+  OMPInformationCache InfoCache(M, AG, Allocator,
+                                /*CGSCC*/ Functions, Kernels, OMPBuilder);
 
   unsigned MaxFixpointIterations = (isOpenMPDevice(M)) ? 128 : 32;
   Attributor A(Functions, InfoCache, CGUpdater, nullptr, false, true,
                MaxFixpointIterations, OREGetter, DEBUG_TYPE);
 
   OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
   bool Changed = OMPOpt.run(false);
 
auto OREGetter = [&OREMap](Function *F) -> OptimizationRemarkEmitter & {
       if (!ORE)
         ORE = std::make_unique<OptimizationRemarkEmitter>(F);
       return *ORE;
     };
 
     AnalysisGetter AG;
     SetVector<Function *> Functions(SCC.begin(), SCC.end());
     BumpPtrAllocator Allocator;
-    OMPInformationCache InfoCache(*(Functions.back()->getParent()), AG,
-                                  Allocator,
-                                  /*CGSCC*/ Functions, Kernels);
+    OpenMPIRBuilder OMPBuilder(M);
+    OMPBuilder.initialize();
+    OMPInformationCache InfoCache(M, AG, Allocator,
+                                  /*CGSCC*/ Functions, Kernels, OMPBuilder);
 
     unsigned MaxFixpointIterations = (isOpenMPDevice(M)) ? 128 : 32;
     Attributor A(Functions, InfoCache, CGUpdater, nullptr, false, true,
                  MaxFixpointIterations, OREGetter, DEBUG_TYPE);
 
     OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
     bool Result = OMPOpt.run(false);
 

llvm/test/Transforms/OpenMP/parallel_level_fold.ll

 ;
   %level = call i8 @__kmpc_parallel_level()
   store i8 %level, i8* @G
   ret void
 }
 
 define internal void @__kmpc_parallel_51(%struct.ident_t*, i32, i32, i32, i32, i8*, i8*, i8**, i64) {
 ; CHECK-LABEL: define {{[^@]+}}@__kmpc_parallel_51
-; CHECK-SAME: (%struct.ident_t* noalias nocapture nofree readnone align 1073741824 [[TMP0:%.*]], i32 [[TMP1:%.*]], i32 [[TMP2:%.*]], i32 [[TMP3:%.*]], i32 [[TMP4:%.*]], i8* noalias nocapture nofree readnone align 1073741824 [[TMP5:%.*]], i8* noalias nocapture nofree readnone align 1073741824 [[TMP6:%.*]], i8** noalias nocapture nofree readnone align 1073741824 [[TMP7:%.*]], i64 [[TMP8:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-SAME: (%struct.ident_t* noalias nocapture nofree readnone align 1073741824 [[TMP0:%.*]], i32 [[TMP1:%.*]], i32 [[TMP2:%.*]], i32 [[TMP3:%.*]], i32 [[TMP4:%.*]], i8* noalias nocapture nofree readnone align 1073741824 [[TMP5:%.*]], i8* noalias nocapture nofree readnone align 1073741824 [[TMP6:%.*]], i8** noalias nocapture nofree readnone align 1073741824 [[TMP7:%.*]], i64 [[TMP8:%.*]]) #[[ATTR0:[0-9]+]] !callback !5 {
 ; CHECK-NEXT:    call void @parallel_helper()
 ; CHECK-NEXT:    ret void
 ;
   call void @parallel_helper()
   ret void
 }
 
 define internal void @parallel_helper() {

 ;.
 ; CHECK: attributes #[[ATTR0]] = { alwaysinline }
 ;.
 ; CHECK: [[META0:![0-9]+]] = !{i32 7, !"openmp", i32 50}
 ; CHECK: [[META1:![0-9]+]] = !{i32 7, !"openmp-device", i32 50}
 ; CHECK: [[META2:![0-9]+]] = !{void ()* @none_spmd, !"kernel", i32 1}
 ; CHECK: [[META3:![0-9]+]] = !{void ()* @spmd, !"kernel", i32 1}
 ; CHECK: [[META4:![0-9]+]] = !{void ()* @parallel, !"kernel", i32 1}
+; CHECK: [[META5:![0-9]+]] = !{!6}
+; CHECK: [[META6:![0-9]+]] = !{i64 5, i64 -1, i64 -1, i64 7, i1 false}
 ;.