[OpenMP][OMPIRBuilder] Migrate target outlined function registration to OMPIRBuilder from clang

[jsjodin]

Summary

This patch moves the outlined function registration, function attribute
configuration and function ID creation to the OpenMPIRBuilder. This will later
be used by flag as well.

[jdoerfert]

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

1467

Outlined function is too generic for the function name here and elsewhere. It's an outlined target region, or similar. We outline a lot of functions but this is specific to target regions, correct?

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3960

This is not DSO local? Maybe add a TODO to check this in the future.

3979–3984

no else after return

[jsjodin]

Changed setOutlinedFunctionAttributes to setOutlinedTargetRegionFunctionAttributes. Fixed else after return. Added TODO. Removed comment.

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

1467

Outlined function is too generic for the function name here and elsewhere. It's an outlined target region, or similar. We outline a lot of functions but this is specific to target regions, correct?

Yes, I am replacing it with setOutlinedTargetRegionFunctionAttributes.

[jsjodin]

Update against trunk.

[jsjodin]

No need to pass in IsEmbedded, it is in the config now. Perhaps the manager should also have the config available.

[jdoerfert]

LG

clang/lib/CodeGen/CGOpenMPRuntime.cpp

void CGOpenMPRuntime::emitTargetOutlinedFunctionHelper(
   OffloadEntriesInfoManager.getTargetRegionEntryFnName(EntryFnName, EntryInfo);
 
   const CapturedStmt &CS = *D.getCapturedStmt(OMPD_target);
 
   CodeGenFunction CGF(CGM, true);
   CGOpenMPTargetRegionInfo CGInfo(CS, CodeGen, EntryFnName);
   CodeGenFunction::CGCapturedStmtRAII CapInfoRAII(CGF, &CGInfo);
 
-  if (BuildOutlinedFn)
-    OutlinedFn = CGF.GenerateOpenMPCapturedStmtFunction(CS, D.getBeginLoc());
+  OutlinedFn = BuildOutlinedFn
+                   ? CGF.GenerateOpenMPCapturedStmtFunction(CS, D.getBeginLoc())
+                   : nullptr;
 
   // If this target outline function is not an offload entry, we don't need to
   // register it.
   if (!IsOffloadEntry)
     return;
 
   // The target region ID is used by the runtime library to identify the current
   // target region, so it only has to be unique and not necessarily point to
   // anything. It could be the pointer to the outlined function that implements
   // the target region, but we aren't using that so that the compiler doesn't
   // need to keep that, and could therefore inline the host function if proven
   // worthwhile during optimization. In the other hand, if emitting code for the
   // device, the ID has to be the function address so that it can retrieved from
   // the offloading entry and launched by the runtime library. We also mark the
   // outlined function to have external linkage in case we are emitting code for
   // the device, because these functions will be entry points to the device.
 
-  if (CGM.getLangOpts().OpenMPIsDevice) {
-    OutlinedFnID = llvm::ConstantExpr::getBitCast(OutlinedFn, CGM.Int8PtrTy);
-    OutlinedFn->setLinkage(llvm::GlobalValue::WeakODRLinkage);
-    OutlinedFn->setDSOLocal(false);
-    OutlinedFn->setVisibility(llvm::GlobalValue::ProtectedVisibility);
-    if (CGM.getTriple().isAMDGCN())
-      OutlinedFn->setCallingConv(llvm::CallingConv::AMDGPU_KERNEL);
-  } else {
-    std::string Name = getName({EntryFnName, "region_id"});
-    OutlinedFnID = new llvm::GlobalVariable(
-        CGM.getModule(), CGM.Int8Ty, /*isConstant=*/true,
-        llvm::GlobalValue::WeakAnyLinkage,
-        llvm::Constant::getNullValue(CGM.Int8Ty), Name);
-  }
-
-  // If we do not allow host fallback we still need a named address to use.
-  llvm::Constant *TargetRegionEntryAddr = OutlinedFn;
-  if (!BuildOutlinedFn) {
-    assert(!CGM.getModule().getGlobalVariable(EntryFnName, true) &&
-           "Named kernel already exists?");
-    TargetRegionEntryAddr = new llvm::GlobalVariable(
-        CGM.getModule(), CGM.Int8Ty, /*isConstant=*/true,
-        llvm::GlobalValue::InternalLinkage,
-        llvm::Constant::getNullValue(CGM.Int8Ty), EntryFnName);
-  }
-
-  // Register the information for the entry associated with this target region.
-  OffloadEntriesInfoManager.registerTargetRegionEntryInfo(
-      EntryInfo, TargetRegionEntryAddr, OutlinedFnID,
-      llvm::OffloadEntriesInfoManager::OMPTargetRegionEntryTargetRegion);
-
-  // Add NumTeams and ThreadLimit attributes to the outlined GPU function
+  // Get NumTeams and ThreadLimit attributes
   int32_t DefaultValTeams = -1;
-  getNumTeamsExprForTargetDirective(CGF, D, DefaultValTeams);
-  if (DefaultValTeams > 0 && OutlinedFn) {
-    OutlinedFn->addFnAttr("omp_target_num_teams",
-                          std::to_string(DefaultValTeams));
-  }
   int32_t DefaultValThreads = -1;
+  getNumTeamsExprForTargetDirective(CGF, D, DefaultValTeams);
   getNumThreadsExprForTargetDirective(CGF, D, DefaultValThreads);
-  if (DefaultValThreads > 0 && OutlinedFn) {
-    OutlinedFn->addFnAttr("omp_target_thread_limit",
-                          std::to_string(DefaultValThreads));
-  }
+
+  std::string EntryFnIDName = CGM.getLangOpts().OpenMPIsDevice
+                                  ? std::string(EntryFnName)
+                                  : getName({EntryFnName, "region_id"});
+
+  OutlinedFnID = OMPBuilder.registerTargetRegionFunction(
+      OffloadEntriesInfoManager, EntryInfo, OutlinedFn, EntryFnName,
+      EntryFnIDName, DefaultValTeams, DefaultValThreads);
 
   if (BuildOutlinedFn)
     CGM.getTargetCodeGenInfo().setTargetAttributes(nullptr, OutlinedFn, CGM);
 }
 
 /// Checks if the expression is constant or does not have non-trivial function
 /// calls.
 static bool isTrivial(ASTContext &Ctx, const Expr * E) {

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

public:
   /// \param Loc The insert and source location description.
   /// \param IsSPMD Flag to indicate if the kernel is an SPMD kernel or not.
   /// \param RequiresFullRuntime Indicate if a full device runtime is necessary.
   void createTargetDeinit(const LocationDescription &Loc, bool IsSPMD,
                           bool RequiresFullRuntime);
 
   ///}
 
+private:
+  // Sets the function attributes expected for the outlined function
+  void setOutlinedTargetRegionFunctionAttributes(Function *OutlinedFn,
+                                                 int32_t NumTeams,
+                                                 int32_t NumThreads);
+
+  // Creates the function ID/Address for the given outlined function.
+  // In the case of an embedded device function the address of the function is
+  // used, in the case of a non-offload function a constant is created.
+  Constant *createOutlinedFunctionID(Function *OutlinedFn,
+                                     StringRef EntryFnIDName);
+
+  // Creates the region entry address for the outlined function
+  Constant *createTargetRegionEntryAddr(Function *OutlinedFunction,
+                                        StringRef EntryFnName);
+
+public:
+  /// Registers the given function and sets up the attribtues of the function
+  /// Returns the FunctionID.
+  ///
+  /// \param InfoManager The info manager keeping track of the offload entries
+  /// \param EntryInfo The entry information about the function
+  /// \param OutlinedFunction Pointer to the outlined function
+  /// \param EntryFnName Name of the outlined function
+  /// \param EntryFnIDName Name of the ID o be created
+  /// \param NumTeams Number default teams
+  /// \param NumThreads Number default threads
+  Constant *registerTargetRegionFunction(OffloadEntriesInfoManager &InfoManager,
+                                         TargetRegionEntryInfo &EntryInfo,
+                                         Function *OutlinedFunction,
+                                         StringRef EntryFnName,
+                                         StringRef EntryFnIDName,
+                                         int32_t NumTeams, int32_t NumThreads);
+
   /// Declarations for LLVM-IR types (simple, array, function and structure) are
   /// generated below. Their names are defined and used in OpenMPKinds.def. Here
   /// we provide the declarations, the initializeTypes function will provide the
   /// values.
   ///
   ///{
 #define OMP_TYPE(VarName, InitValue) Type *VarName = nullptr;
 #define OMP_ARRAY_TYPE(VarName, ElemTy, ArraySize)                             \

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

ConstantInt *RequiresFullRuntimeVal =
       ConstantInt::getBool(Int32->getContext(), RequiresFullRuntime);
 
   Function *Fn = getOrCreateRuntimeFunctionPtr(
       omp::RuntimeFunction::OMPRTL___kmpc_target_deinit);
 
   Builder.CreateCall(Fn, {Ident, IsSPMDVal, RequiresFullRuntimeVal});
 }
 
+void OpenMPIRBuilder::setOutlinedTargetRegionFunctionAttributes(
+    Function *OutlinedFn, int32_t NumTeams, int32_t NumThreads) {
+  if (Config.isEmbedded()) {
+    OutlinedFn->setLinkage(GlobalValue::WeakODRLinkage);
+    // TODO: Determine if DSO local can be set to true.
+    OutlinedFn->setDSOLocal(false);
+    OutlinedFn->setVisibility(GlobalValue::ProtectedVisibility);
+    if (Triple(M.getTargetTriple()).isAMDGCN())
+      OutlinedFn->setCallingConv(CallingConv::AMDGPU_KERNEL);
+  }
+
+  if (NumTeams > 0)
+    OutlinedFn->addFnAttr("omp_target_num_teams", std::to_string(NumTeams));
+  if (NumThreads > 0)
+    OutlinedFn->addFnAttr("omp_target_thread_limit",
+                          std::to_string(NumThreads));
+}
+
+Constant *OpenMPIRBuilder::createOutlinedFunctionID(Function *OutlinedFn,
+                                                    StringRef EntryFnIDName) {
+  if (Config.isEmbedded()) {
+    assert(OutlinedFn && "The outlined function must exist if embedded");
+    return ConstantExpr::getBitCast(OutlinedFn, Builder.getInt8PtrTy());
+  }
+
+  return new GlobalVariable(
+      M, Builder.getInt8Ty(), /*isConstant=*/true, GlobalValue::WeakAnyLinkage,
+      Constant::getNullValue(Builder.getInt8Ty()), EntryFnIDName);
+}
+
+Constant *OpenMPIRBuilder::createTargetRegionEntryAddr(Function *OutlinedFn,
+                                                       StringRef EntryFnName) {
+  if (OutlinedFn)
+    return OutlinedFn;
+
+  assert(!M.getGlobalVariable(EntryFnName, true) &&
+         "Named kernel already exists?");
+  return new GlobalVariable(
+      M, Builder.getInt8Ty(), /*isConstant=*/true, GlobalValue::InternalLinkage,
+      Constant::getNullValue(Builder.getInt8Ty()), EntryFnName);
+}
+
+Constant *OpenMPIRBuilder::registerTargetRegionFunction(
+    OffloadEntriesInfoManager &InfoManager, TargetRegionEntryInfo &EntryInfo,
+    Function *OutlinedFn, StringRef EntryFnName, StringRef EntryFnIDName,
+    int32_t NumTeams, int32_t NumThreads) {
+  if (OutlinedFn)
+    setOutlinedTargetRegionFunctionAttributes(OutlinedFn, NumTeams, NumThreads);
+  auto OutlinedFnID = createOutlinedFunctionID(OutlinedFn, EntryFnIDName);
+  auto EntryAddr = createTargetRegionEntryAddr(OutlinedFn, EntryFnName);
+  InfoManager.registerTargetRegionEntryInfo(
+      EntryInfo, EntryAddr, OutlinedFnID,
+      OffloadEntriesInfoManager::OMPTargetRegionEntryTargetRegion);
+  return OutlinedFnID;
+}
+
 std::string OpenMPIRBuilder::getNameWithSeparators(ArrayRef<StringRef> Parts,
                                                    StringRef FirstSeparator,
                                                    StringRef Separator) {
   SmallString<128> Buffer;
   llvm::raw_svector_ostream OS(Buffer);
   StringRef Sep = FirstSeparator;
   for (StringRef Part : Parts) {
     OS << Sep << Part;