[OpenMP][OMPIRBuilder] Migrate code to emit target region functions from clang to OMPIRBuilder

[jsjodin]

Summary

This patch moves some of the logic on how to emit target region functions and
adds emitTargetRegionFunction to the OpenMPIRBuilder. Also the
OpenMPOffloadMandatory flag is added to the config.

[jdoerfert]

LG

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

1506–1507

minor comment mishap

[jsjodin]

Fix comment and test.

clang/lib/CodeGen/CGOpenMPRuntime.cpp

static FieldDecl *addFieldToRecordDecl(ASTContext &C, DeclContext *DC,
   DC->addDecl(Field);
   return Field;
 }
 
 CGOpenMPRuntime::CGOpenMPRuntime(CodeGenModule &CGM)
     : CGM(CGM), OMPBuilder(CGM.getModule()), OffloadEntriesInfoManager() {
   KmpCriticalNameTy = llvm::ArrayType::get(CGM.Int32Ty, /*NumElements*/ 8);
   llvm::OpenMPIRBuilderConfig Config(CGM.getLangOpts().OpenMPIsDevice, false,
-                                     hasRequiresUnifiedSharedMemory());
+                                     hasRequiresUnifiedSharedMemory(),
+                                     CGM.getLangOpts().OpenMPOffloadMandatory);
   // Initialize Types used in OpenMPIRBuilder from OMPKinds.def
   OMPBuilder.initialize();
   OMPBuilder.setConfig(Config);
   OffloadEntriesInfoManager.setConfig(Config);
   loadOffloadInfoMetadata();
 }
 
 void CGOpenMPRuntime::clear() {
(void)CGF.EmitRuntimeCall(
                                             OMPRTL___kmpc_destroy_allocator),
       {ThreadId, AllocatorVal});
 }
 
 void CGOpenMPRuntime::emitTargetOutlinedFunctionHelper(
     const OMPExecutableDirective &D, StringRef ParentName,
     llvm::Function *&OutlinedFn, llvm::Constant *&OutlinedFnID,
     bool IsOffloadEntry, const RegionCodeGenTy &CodeGen) {
-  // Create a unique name for the entry function using the source location
-  // information of the current target region. The name will be something like:
-  //
-  // __omp_offloading_DD_FFFF_PP_lBB[_CC]
-  //
-  // where DD_FFFF is an ID unique to the file (device and file IDs), PP is the
-  // mangled name of the function that encloses the target region and BB is the
-  // line number of the target region. CC is a count added when more than one
-  // region is located at the same location.
 
-  const bool BuildOutlinedFn = CGM.getLangOpts().OpenMPIsDevice ||
-                               !CGM.getLangOpts().OpenMPOffloadMandatory;
   auto EntryInfo =
       getTargetEntryUniqueInfo(CGM.getContext(), D.getBeginLoc(), ParentName);
 
-  SmallString<64> EntryFnName;
-  OffloadEntriesInfoManager.getTargetRegionEntryFnName(EntryFnName, EntryInfo);
-
+  CodeGenFunction CGF(CGM, true);
+  llvm::OpenMPIRBuilder::FunctionGenCallback &&GenerateOutlinedFunction =
+      [this, &CGF, &D, &CodeGen](StringRef EntryFnName) {
         const CapturedStmt &CS = *D.getCapturedStmt(OMPD_target);
 
-  CodeGenFunction CGF(CGM, true);
         CGOpenMPTargetRegionInfo CGInfo(CS, CodeGen, EntryFnName);
         CodeGenFunction::CGCapturedStmtRAII CapInfoRAII(CGF, &CGInfo);
-
-  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.
+        return CGF.GenerateOpenMPCapturedStmtFunction(CS, D.getBeginLoc());
+      };
 
   // Get NumTeams and ThreadLimit attributes
   int32_t DefaultValTeams = -1;
   int32_t DefaultValThreads = -1;
   getNumTeamsExprForTargetDirective(CGF, D, DefaultValTeams);
   getNumThreadsExprForTargetDirective(CGF, D, DefaultValThreads);
 
-  std::string EntryFnIDName = CGM.getLangOpts().OpenMPIsDevice
-                                  ? std::string(EntryFnName)
-                                  : getName({EntryFnName, "region_id"});
-
-  OutlinedFnID = OMPBuilder.registerTargetRegionFunction(
-      OffloadEntriesInfoManager, EntryInfo, OutlinedFn, EntryFnName,
-      EntryFnIDName, DefaultValTeams, DefaultValThreads);
+  OMPBuilder.emitTargetRegionFunction(OffloadEntriesInfoManager, EntryInfo,
+                                      GenerateOutlinedFunction, DefaultValTeams,
+                                      DefaultValThreads, IsOffloadEntry,
+                                      OutlinedFn, OutlinedFnID);
 
-  if (BuildOutlinedFn)
+  if (OutlinedFn != nullptr)
     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) {
   // We can skip constant expressions.
   // We can skip expressions with trivial calls or simple expressions.

clang/lib/CodeGen/CGOpenMPRuntimeGPU.cpp

case EM_Unknown:
     return UndefinedMode;
   }
   llvm_unreachable("Unknown flags are requested.");
 }
 
 CGOpenMPRuntimeGPU::CGOpenMPRuntimeGPU(CodeGenModule &CGM)
     : CGOpenMPRuntime(CGM) {
   llvm::OpenMPIRBuilderConfig Config(CGM.getLangOpts().OpenMPIsDevice, true,
-                                     hasRequiresUnifiedSharedMemory());
+                                     hasRequiresUnifiedSharedMemory(),
+                                     CGM.getLangOpts().OpenMPOffloadMandatory);
   OMPBuilder.setConfig(Config);
   OffloadEntriesInfoManager.setConfig(Config);
 
   if (!CGM.getLangOpts().OpenMPIsDevice)
     llvm_unreachable("OpenMP can only handle device code.");
 
   llvm::OpenMPIRBuilder &OMPBuilder = getOMPBuilder();
   if (CGM.getLangOpts().NoGPULib || CGM.getLangOpts().OMPHostIRFile.empty())

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

public:
   /// Flag for specifying if the compilation is done for an offloading target,
   /// like GPU.
   std::optional<bool> IsTargetCodegen;
 
   /// Flag for specifying weather a requires unified_shared_memory
   /// directive is present or not.
   std::optional<bool> HasRequiresUnifiedSharedMemory;
 
+  // Flag for specifying if offloading is mandatory.
+  Optional<bool> OpenMPOffloadMandatory;
+
   /// First separator used between the initial two parts of a name.
   std::optional<StringRef> FirstSeparator;
   /// Separator used between all of the rest consecutive parts of s name
   std::optional<StringRef> Separator;
 
   OpenMPIRBuilderConfig() {}
   OpenMPIRBuilderConfig(bool IsEmbedded, bool IsTargetCodegen,
-                        bool HasRequiresUnifiedSharedMemory)
+                        bool HasRequiresUnifiedSharedMemory,
+                        bool OpenMPOffloadMandatory)
       : IsEmbedded(IsEmbedded), IsTargetCodegen(IsTargetCodegen),
-        HasRequiresUnifiedSharedMemory(HasRequiresUnifiedSharedMemory) {}
+        HasRequiresUnifiedSharedMemory(HasRequiresUnifiedSharedMemory),
+        OpenMPOffloadMandatory(OpenMPOffloadMandatory) {}
 
   // Getters functions that assert if the required values are not present.
   bool isEmbedded() const {
     assert(IsEmbedded.has_value() && "IsEmbedded is not set");
     return IsEmbedded.value();
   }
 
   bool isTargetCodegen() const {
     assert(IsTargetCodegen.has_value() && "IsTargetCodegen is not set");
     return IsTargetCodegen.value();
   }
 
   bool hasRequiresUnifiedSharedMemory() const {
     assert(HasRequiresUnifiedSharedMemory.has_value() &&
            "HasUnifiedSharedMemory is not set");
     return HasRequiresUnifiedSharedMemory.value();
   }
 
+  bool openMPOffloadMandatory() const {
+    assert(OpenMPOffloadMandatory.has_value() &&
+           "OpenMPOffloadMandatory is not set");
+    return OpenMPOffloadMandatory.value();
+  }
   // Returns the FirstSeparator if set, otherwise use the default
   // separator depending on isTargetCodegen
   StringRef firstSeparator() const {
     if (FirstSeparator.has_value())
       return FirstSeparator.value();
     if (isTargetCodegen())
       return "_";
     return ".";
private:
   Constant *createOutlinedFunctionID(Function *OutlinedFn,
                                      StringRef EntryFnIDName);
 
   // Creates the region entry address for the outlined function
   Constant *createTargetRegionEntryAddr(Function *OutlinedFunction,
                                         StringRef EntryFnName);
 
 public:
+  /// Functions used to generate a function with the given name.
+  using FunctionGenCallback = std::function<Function *(StringRef FunctionName)>;
+
+  /// Create a unique name for the entry function using the source location
+  /// information of the current target region. The name will be something like:
+  ///
+  /// __omp_offloading_DD_FFFF_PP_lBB[_CC]
+  ///
+  /// where DD_FFFF is an ID unique to the file (device and file IDs), PP is the
+  /// mangled name of the function that encloses the target region and BB is the
+  /// line number of the target region. CC is a count added when more than one
+  /// region is located at the same location.
+  ///
+  /// If this target outline function is not an offload entry, we don't need to
+  /// register it. This may happen if it is guarded by an if clause that is
+  /// false at compile time, or no target archs have been specified.
+  ///
+  /// The created 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.
+  ///
+  /// \param InfoManager The info manager keeping track of the offload entries
+  /// \param EntryInfo The entry information about the function
+  /// \param GenerateFunctionCallback The callback function to generate the code
+  /// \param NumTeams Number default teams
+  /// \param NumThreads Number default threads
+  /// \param OutlinedFunction Pointer to the outlined function
+  /// \param EntryFnIDName Name of the ID o be created
+  void emitTargetRegionFunction(OffloadEntriesInfoManager &InfoManager,
+                                TargetRegionEntryInfo &EntryInfo,
+                                FunctionGenCallback &GenerateFunctionCallback,
+                                int32_t NumTeams, int32_t NumThreads,
+                                bool IsOffloadEntry, Function *&OutlinedFn,
+                                Constant *&OutlinedFnID);
+
   /// 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

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Constant *OpenMPIRBuilder::createTargetRegionEntryAddr(Function *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);
 }
 
+void OpenMPIRBuilder::emitTargetRegionFunction(
+    OffloadEntriesInfoManager &InfoManager, TargetRegionEntryInfo &EntryInfo,
+    FunctionGenCallback &GenerateFunctionCallback, int32_t NumTeams,
+    int32_t NumThreads, bool IsOffloadEntry, Function *&OutlinedFn,
+    Constant *&OutlinedFnID) {
+
+  SmallString<64> EntryFnName;
+  InfoManager.getTargetRegionEntryFnName(EntryFnName, EntryInfo);
+
+  OutlinedFn = Config.isEmbedded() || !Config.openMPOffloadMandatory()
+                   ? GenerateFunctionCallback(EntryFnName)
+                   : nullptr;
+
+  // If this target outline function is not an offload entry, we don't need to
+  // register it. This may be in the case of a false if clause, or if there are
+  // no OpenMP targets.
+  if (!IsOffloadEntry)
+    return;
+
+  std::string EntryFnIDName =
+      Config.isEmbedded()
+          ? std::string(EntryFnName)
+          : createPlatformSpecificName({EntryFnName, "region_id"});
+
+  OutlinedFnID = registerTargetRegionFunction(
+      InfoManager, EntryInfo, OutlinedFn, EntryFnName, EntryFnIDName, NumTeams,
+      NumThreads);
+}
+
 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);

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

TEST_F(OpenMPIRBuilderTest, EmitOffloadingArraysArguments) {
   EXPECT_EQ(RTArgs.SizesArray->getType(), Int64PtrTy);
   EXPECT_EQ(RTArgs.MapTypesArray->getType(), Int64PtrTy);
   EXPECT_EQ(RTArgs.MappersArray->getType(), VoidPtrPtrTy);
   EXPECT_EQ(RTArgs.MapNamesArray->getType(), VoidPtrPtrTy);
 }
 
 TEST_F(OpenMPIRBuilderTest, OffloadEntriesInfoManager) {
   OffloadEntriesInfoManager InfoManager;
-  InfoManager.setConfig(OpenMPIRBuilderConfig(true, false, false));
+  InfoManager.setConfig(OpenMPIRBuilderConfig(true, false, false, false));
   TargetRegionEntryInfo EntryInfo("parent", 1, 2, 4, 0);
   InfoManager.initializeTargetRegionEntryInfo(EntryInfo, 0);
   EXPECT_TRUE(InfoManager.hasTargetRegionEntryInfo(EntryInfo));
   InfoManager.initializeDeviceGlobalVarEntryInfo(
       "gvar", OffloadEntriesInfoManager::OMPTargetGlobalVarEntryTo, 0);
   InfoManager.registerTargetRegionEntryInfo(
       EntryInfo, nullptr, nullptr,
       OffloadEntriesInfoManager::OMPTargetRegionEntryTargetRegion);
   InfoManager.registerDeviceGlobalVarEntryInfo(
       "gvar", 0x0, 8, OffloadEntriesInfoManager::OMPTargetGlobalVarEntryTo,
       GlobalValue::WeakAnyLinkage);
   EXPECT_TRUE(InfoManager.hasDeviceGlobalVarEntryInfo("gvar"));
 }
 } // namespace