[clang] [llvm] [llvm][frontend][offloading] Move clang-linker-wrapper/OffloadWrapper.* to llvm/Frontend/Offloading (PR #78057)

Fabian Mora via llvm-commits llvm-commits at lists.llvm.org
Sun Jan 14 13:42:31 PST 2024


https://github.com/fabianmcg updated https://github.com/llvm/llvm-project/pull/78057

>From f56a7395b19ff634b3ac963204348db2575fdf87 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Sat, 13 Jan 2024 17:31:51 +0000
Subject: [PATCH 1/3] Move OffloadWrapper.* to llvm/Frontend/Offloading

---
 .../include/llvm/Frontend/Offloading}/OffloadWrapper.h            | 0
 .../lib/Frontend/Offloading}/OffloadWrapper.cpp                   | 0
 2 files changed, 0 insertions(+), 0 deletions(-)
 rename {clang/tools/clang-linker-wrapper => llvm/include/llvm/Frontend/Offloading}/OffloadWrapper.h (100%)
 rename {clang/tools/clang-linker-wrapper => llvm/lib/Frontend/Offloading}/OffloadWrapper.cpp (100%)

diff --git a/clang/tools/clang-linker-wrapper/OffloadWrapper.h b/llvm/include/llvm/Frontend/Offloading/OffloadWrapper.h
similarity index 100%
rename from clang/tools/clang-linker-wrapper/OffloadWrapper.h
rename to llvm/include/llvm/Frontend/Offloading/OffloadWrapper.h
diff --git a/clang/tools/clang-linker-wrapper/OffloadWrapper.cpp b/llvm/lib/Frontend/Offloading/OffloadWrapper.cpp
similarity index 100%
rename from clang/tools/clang-linker-wrapper/OffloadWrapper.cpp
rename to llvm/lib/Frontend/Offloading/OffloadWrapper.cpp

>From fad7a36c34bfed207c35cbd0e8e431a6910da792 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Sat, 13 Jan 2024 17:43:40 +0000
Subject: [PATCH 2/3] [llvm][frontend][offloading] Move
 clang-linker-wrapper/OffloadWrapper.* to llvm/Frontend/Offloading

This patch moves `clang/tools/clang-linker-wrapper/OffloadWrapper.*` to
`llvm/Frontend/Offloading` allowing them to be reutilized by other projects.

Additionally, it makes minor modifications to the API to make it more flexible.
Concretely:
 - The `wrap*` methods are moved to the `OffloadWrapper` class.
 - The `OffloadWrapper` includes `Suffix` and `EmitSurfacesAndTextures` fields
to specify some additional options.
 - The `Suffix` field is used when emitting the descriptor, registration methods,
etc, to make them more readable. It is empty by default.
 - The `EmitSurfacesAndTextures` field controls whether to emit surface and
texture registration code, as those functions were removed from `CUDART`
in CUDA 12. It is true by default.
 - The `wrap*` methods now have an optional field to specify the `EntryArray`;
this change is needed to enable JIT compilation, as ORC doesn't fully support
`__start_` and `__stop_` symbols. Thus, to JIT the code, the `EntryArray` has
to be constructed explicitly in the IR.
 - The function `getOffloadingEntryInitializer` was added to help create the
`EntryArray`, as it returns the constant initializer and not a global variable.
---
 .../tools/clang-linker-wrapper/CMakeLists.txt |   1 -
 .../ClangLinkerWrapper.cpp                    |  11 +-
 .../llvm/Frontend/Offloading/OffloadWrapper.h |  60 +++++++--
 .../llvm/Frontend/Offloading/Utility.h        |   6 +
 llvm/lib/Frontend/Offloading/CMakeLists.txt   |   2 +
 .../Frontend/Offloading/OffloadWrapper.cpp    | 123 +++++++++++-------
 llvm/lib/Frontend/Offloading/Utility.cpp      |  21 ++-
 7 files changed, 156 insertions(+), 68 deletions(-)

diff --git a/clang/tools/clang-linker-wrapper/CMakeLists.txt b/clang/tools/clang-linker-wrapper/CMakeLists.txt
index 744026a37b22c0..5556869affaa62 100644
--- a/clang/tools/clang-linker-wrapper/CMakeLists.txt
+++ b/clang/tools/clang-linker-wrapper/CMakeLists.txt
@@ -28,7 +28,6 @@ endif()
 
 add_clang_tool(clang-linker-wrapper
   ClangLinkerWrapper.cpp
-  OffloadWrapper.cpp
 
   DEPENDS
   ${tablegen_deps}
diff --git a/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp b/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
index 122ba1998eb83f..ebe8b634c7ae73 100644
--- a/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
+++ b/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
@@ -14,11 +14,11 @@
 //
 //===---------------------------------------------------------------------===//
 
-#include "OffloadWrapper.h"
 #include "clang/Basic/Version.h"
 #include "llvm/BinaryFormat/Magic.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/CodeGen/CommandFlags.h"
+#include "llvm/Frontend/Offloading/OffloadWrapper.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/IR/Module.h"
@@ -906,15 +906,18 @@ wrapDeviceImages(ArrayRef<std::unique_ptr<MemoryBuffer>> Buffers,
 
   switch (Kind) {
   case OFK_OpenMP:
-    if (Error Err = wrapOpenMPBinaries(M, BuffersToWrap))
+    if (Error Err =
+            offloading::OffloadWrapper().wrapOpenMPBinaries(M, BuffersToWrap))
       return std::move(Err);
     break;
   case OFK_Cuda:
-    if (Error Err = wrapCudaBinary(M, BuffersToWrap.front()))
+    if (Error Err = offloading::OffloadWrapper().wrapCudaBinary(
+            M, BuffersToWrap.front()))
       return std::move(Err);
     break;
   case OFK_HIP:
-    if (Error Err = wrapHIPBinary(M, BuffersToWrap.front()))
+    if (Error Err = offloading::OffloadWrapper().wrapHIPBinary(
+            M, BuffersToWrap.front()))
       return std::move(Err);
     break;
   default:
diff --git a/llvm/include/llvm/Frontend/Offloading/OffloadWrapper.h b/llvm/include/llvm/Frontend/Offloading/OffloadWrapper.h
index 679333975b2120..6b23f875a8f15f 100644
--- a/llvm/include/llvm/Frontend/Offloading/OffloadWrapper.h
+++ b/llvm/include/llvm/Frontend/Offloading/OffloadWrapper.h
@@ -6,23 +6,57 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef LLVM_CLANG_TOOLS_CLANG_LINKER_WRAPPER_OFFLOAD_WRAPPER_H
-#define LLVM_CLANG_TOOLS_CLANG_LINKER_WRAPPER_OFFLOAD_WRAPPER_H
+#ifndef LLVM_FRONTEND_OFFLOADING_OFFLOADWRAPPER_H
+#define LLVM_FRONTEND_OFFLOADING_OFFLOADWRAPPER_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/IR/Module.h"
 
-/// Wraps the input device images into the module \p M as global symbols and
-/// registers the images with the OpenMP Offloading runtime libomptarget.
-llvm::Error wrapOpenMPBinaries(llvm::Module &M,
-                               llvm::ArrayRef<llvm::ArrayRef<char>> Images);
+namespace llvm {
+namespace offloading {
+/// Class for embedding and registering offloading images and related objects in
+/// a Module.
+class OffloadWrapper {
+public:
+  using EntryArrayTy = std::pair<GlobalVariable *, GlobalVariable *>;
 
-/// Wraps the input fatbinary image into the module \p M as global symbols and
-/// registers the images with the CUDA runtime.
-llvm::Error wrapCudaBinary(llvm::Module &M, llvm::ArrayRef<char> Images);
+  OffloadWrapper(const Twine &Suffix = "", bool EmitSurfacesAndTextures = true)
+      : Suffix(Suffix.str()), EmitSurfacesAndTextures(EmitSurfacesAndTextures) {
+  }
 
-/// Wraps the input bundled image into the module \p M as global symbols and
-/// registers the images with the HIP runtime.
-llvm::Error wrapHIPBinary(llvm::Module &M, llvm::ArrayRef<char> Images);
+  /// Wraps the input device images into the module \p M as global symbols and
+  /// registers the images with the OpenMP Offloading runtime libomptarget.
+  /// \param EntryArray Optional pair pointing to the `__start` and `__stop`
+  /// symbols holding the `__tgt_offload_entry` array.
+  llvm::Error wrapOpenMPBinaries(
+      llvm::Module &M, llvm::ArrayRef<llvm::ArrayRef<char>> Images,
+      std::optional<EntryArrayTy> EntryArray = std::nullopt) const;
 
-#endif
+  /// Wraps the input fatbinary image into the module \p M as global symbols and
+  /// registers the images with the CUDA runtime.
+  /// \param EntryArray Optional pair pointing to the `__start` and `__stop`
+  /// symbols holding the `__tgt_offload_entry` array.
+  llvm::Error
+  wrapCudaBinary(llvm::Module &M, llvm::ArrayRef<char> Images,
+                 std::optional<EntryArrayTy> EntryArray = std::nullopt) const;
+
+  /// Wraps the input bundled image into the module \p M as global symbols and
+  /// registers the images with the HIP runtime.
+  /// \param EntryArray Optional pair pointing to the `__start` and `__stop`
+  /// symbols holding the `__tgt_offload_entry` array.
+  llvm::Error
+  wrapHIPBinary(llvm::Module &M, llvm::ArrayRef<char> Images,
+                std::optional<EntryArrayTy> EntryArray = std::nullopt) const;
+
+protected:
+  /// Suffix used when emitting symbols. It defaults to the empty string.
+  std::string Suffix;
+
+  /// Whether to emit surface and textures registration code. It defaults to
+  /// false.
+  bool EmitSurfacesAndTextures;
+};
+} // namespace offloading
+} // namespace llvm
+
+#endif // LLVM_FRONTEND_OFFLOADING_OFFLOADWRAPPER_H
diff --git a/llvm/include/llvm/Frontend/Offloading/Utility.h b/llvm/include/llvm/Frontend/Offloading/Utility.h
index 520c192996a066..f54dd7ba7ab45f 100644
--- a/llvm/include/llvm/Frontend/Offloading/Utility.h
+++ b/llvm/include/llvm/Frontend/Offloading/Utility.h
@@ -61,6 +61,12 @@ StructType *getEntryTy(Module &M);
 void emitOffloadingEntry(Module &M, Constant *Addr, StringRef Name,
                          uint64_t Size, int32_t Flags, int32_t Data,
                          StringRef SectionName);
+/// Create a constant struct initializer used to register this global at
+/// runtime.
+/// \return the constant struct and the global variable holding the symbol name.
+std::pair<Constant *, GlobalVariable *>
+getOffloadingEntryInitializer(Module &M, Constant *Addr, StringRef Name,
+                              uint64_t Size, int32_t Flags, int32_t Data);
 
 /// Creates a pair of globals used to iterate the array of offloading entries by
 /// accessing the section variables provided by the linker.
diff --git a/llvm/lib/Frontend/Offloading/CMakeLists.txt b/llvm/lib/Frontend/Offloading/CMakeLists.txt
index 2d0117c9e10059..16e0dcfa0e90d6 100644
--- a/llvm/lib/Frontend/Offloading/CMakeLists.txt
+++ b/llvm/lib/Frontend/Offloading/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_llvm_component_library(LLVMFrontendOffloading
   Utility.cpp
+  OffloadWrapper.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Frontend
@@ -9,6 +10,7 @@ add_llvm_component_library(LLVMFrontendOffloading
 
   LINK_COMPONENTS
   Core
+  BinaryFormat
   Support
   TransformUtils
   TargetParser
diff --git a/llvm/lib/Frontend/Offloading/OffloadWrapper.cpp b/llvm/lib/Frontend/Offloading/OffloadWrapper.cpp
index 161374ae555233..f34a879b99dd02 100644
--- a/llvm/lib/Frontend/Offloading/OffloadWrapper.cpp
+++ b/llvm/lib/Frontend/Offloading/OffloadWrapper.cpp
@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "OffloadWrapper.h"
+#include "llvm/Frontend/Offloading/OffloadWrapper.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/BinaryFormat/Magic.h"
 #include "llvm/Frontend/Offloading/Utility.h"
@@ -21,8 +21,11 @@
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 using namespace llvm;
+using namespace llvm::offloading;
 
 namespace {
+using EntryArrayTy = OffloadWrapper::EntryArrayTy;
+
 /// Magic number that begins the section containing the CUDA fatbinary.
 constexpr unsigned CudaFatMagic = 0x466243b1;
 constexpr unsigned HIPFatMagic = 0x48495046;
@@ -110,10 +113,10 @@ PointerType *getBinDescPtrTy(Module &M) {
 /// };
 ///
 /// Global variable that represents BinDesc is returned.
-GlobalVariable *createBinDesc(Module &M, ArrayRef<ArrayRef<char>> Bufs) {
+GlobalVariable *createBinDesc(Module &M, ArrayRef<ArrayRef<char>> Bufs,
+                              EntryArrayTy EntryArray, StringRef Suffix) {
   LLVMContext &C = M.getContext();
-  auto [EntriesB, EntriesE] =
-      offloading::getOffloadEntryArray(M, "omp_offloading_entries");
+  auto [EntriesB, EntriesE] = EntryArray;
 
   auto *Zero = ConstantInt::get(getSizeTTy(M), 0u);
   Constant *ZeroZero[] = {Zero, Zero};
@@ -126,7 +129,7 @@ GlobalVariable *createBinDesc(Module &M, ArrayRef<ArrayRef<char>> Bufs) {
     auto *Data = ConstantDataArray::get(C, Buf);
     auto *Image = new GlobalVariable(M, Data->getType(), /*isConstant=*/true,
                                      GlobalVariable::InternalLinkage, Data,
-                                     ".omp_offloading.device_image");
+                                     ".omp_offloading.device_image" + Suffix);
     Image->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
     Image->setSection(".llvm.offloading");
     Image->setAlignment(Align(object::OffloadBinary::getAlignment()));
@@ -166,7 +169,7 @@ GlobalVariable *createBinDesc(Module &M, ArrayRef<ArrayRef<char>> Bufs) {
   auto *Images =
       new GlobalVariable(M, ImagesData->getType(), /*isConstant*/ true,
                          GlobalValue::InternalLinkage, ImagesData,
-                         ".omp_offloading.device_images");
+                         ".omp_offloading.device_images" + Suffix);
   Images->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
 
   auto *ImagesB =
@@ -180,14 +183,15 @@ GlobalVariable *createBinDesc(Module &M, ArrayRef<ArrayRef<char>> Bufs) {
 
   return new GlobalVariable(M, DescInit->getType(), /*isConstant*/ true,
                             GlobalValue::InternalLinkage, DescInit,
-                            ".omp_offloading.descriptor");
+                            ".omp_offloading.descriptor" + Suffix);
 }
 
-void createRegisterFunction(Module &M, GlobalVariable *BinDesc) {
+void createRegisterFunction(Module &M, GlobalVariable *BinDesc,
+                            StringRef Suffix) {
   LLVMContext &C = M.getContext();
   auto *FuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
   auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage,
-                                ".omp_offloading.descriptor_reg", &M);
+                                ".omp_offloading.descriptor_reg" + Suffix, &M);
   Func->setSection(".text.startup");
 
   // Get __tgt_register_lib function declaration.
@@ -210,11 +214,13 @@ void createRegisterFunction(Module &M, GlobalVariable *BinDesc) {
   appendToGlobalCtors(M, Func, /*Priority*/ 1);
 }
 
-void createUnregisterFunction(Module &M, GlobalVariable *BinDesc) {
+void createUnregisterFunction(Module &M, GlobalVariable *BinDesc,
+                              StringRef Suffix) {
   LLVMContext &C = M.getContext();
   auto *FuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
-  auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage,
-                                ".omp_offloading.descriptor_unreg", &M);
+  auto *Func =
+      Function::Create(FuncTy, GlobalValue::InternalLinkage,
+                       ".omp_offloading.descriptor_unreg" + Suffix, &M);
   Func->setSection(".text.startup");
 
   // Get __tgt_unregister_lib function declaration.
@@ -251,7 +257,8 @@ StructType *getFatbinWrapperTy(Module &M) {
 
 /// Embed the image \p Image into the module \p M so it can be found by the
 /// runtime.
-GlobalVariable *createFatbinDesc(Module &M, ArrayRef<char> Image, bool IsHIP) {
+GlobalVariable *createFatbinDesc(Module &M, ArrayRef<char> Image, bool IsHIP,
+                                 StringRef Suffix) {
   LLVMContext &C = M.getContext();
   llvm::Type *Int8PtrTy = PointerType::getUnqual(C);
   llvm::Triple Triple = llvm::Triple(M.getTargetTriple());
@@ -263,7 +270,7 @@ GlobalVariable *createFatbinDesc(Module &M, ArrayRef<char> Image, bool IsHIP) {
   auto *Data = ConstantDataArray::get(C, Image);
   auto *Fatbin = new GlobalVariable(M, Data->getType(), /*isConstant*/ true,
                                     GlobalVariable::InternalLinkage, Data,
-                                    ".fatbin_image");
+                                    ".fatbin_image" + Suffix);
   Fatbin->setSection(FatbinConstantSection);
 
   // Create the fatbinary wrapper
@@ -282,7 +289,7 @@ GlobalVariable *createFatbinDesc(Module &M, ArrayRef<char> Image, bool IsHIP) {
   auto *FatbinDesc =
       new GlobalVariable(M, getFatbinWrapperTy(M),
                          /*isConstant*/ true, GlobalValue::InternalLinkage,
-                         FatbinInitializer, ".fatbin_wrapper");
+                         FatbinInitializer, ".fatbin_wrapper" + Suffix);
   FatbinDesc->setSection(FatbinWrapperSection);
   FatbinDesc->setAlignment(Align(8));
 
@@ -312,10 +319,12 @@ GlobalVariable *createFatbinDesc(Module &M, ArrayRef<char> Image, bool IsHIP) {
 ///                         0, entry->size, 0, 0);
 ///   }
 /// }
-Function *createRegisterGlobalsFunction(Module &M, bool IsHIP) {
+Function *createRegisterGlobalsFunction(Module &M, bool IsHIP,
+                                        EntryArrayTy EntryArray,
+                                        StringRef Suffix,
+                                        bool EmitSurfacesAndTextures) {
   LLVMContext &C = M.getContext();
-  auto [EntriesB, EntriesE] = offloading::getOffloadEntryArray(
-      M, IsHIP ? "hip_offloading_entries" : "cuda_offloading_entries");
+  auto [EntriesB, EntriesE] = EntryArray;
 
   // Get the __cudaRegisterFunction function declaration.
   PointerType *Int8PtrTy = PointerType::get(C, 0);
@@ -339,7 +348,7 @@ Function *createRegisterGlobalsFunction(Module &M, bool IsHIP) {
       IsHIP ? "__hipRegisterVar" : "__cudaRegisterVar", RegVarTy);
 
   // Get the __cudaRegisterSurface function declaration.
-  auto *RegSurfaceTy =
+  FunctionType *RegSurfaceTy =
       FunctionType::get(Type::getVoidTy(C),
                         {Int8PtrPtrTy, Int8PtrTy, Int8PtrTy, Int8PtrTy,
                          Type::getInt32Ty(C), Type::getInt32Ty(C)},
@@ -348,7 +357,7 @@ Function *createRegisterGlobalsFunction(Module &M, bool IsHIP) {
       IsHIP ? "__hipRegisterSurface" : "__cudaRegisterSurface", RegSurfaceTy);
 
   // Get the __cudaRegisterTexture function declaration.
-  auto *RegTextureTy = FunctionType::get(
+  FunctionType *RegTextureTy = FunctionType::get(
       Type::getVoidTy(C),
       {Int8PtrPtrTy, Int8PtrTy, Int8PtrTy, Int8PtrTy, Type::getInt32Ty(C),
        Type::getInt32Ty(C), Type::getInt32Ty(C)},
@@ -454,19 +463,20 @@ Function *createRegisterGlobalsFunction(Module &M, bool IsHIP) {
   Builder.CreateBr(IfEndBB);
   Switch->addCase(Builder.getInt32(llvm::offloading::OffloadGlobalManagedEntry),
                   SwManagedBB);
-
   // Create surface variable registration code.
   Builder.SetInsertPoint(SwSurfaceBB);
-  Builder.CreateCall(
-      RegSurface, {RegGlobalsFn->arg_begin(), Addr, Name, Name, Data, Extern});
+  if (EmitSurfacesAndTextures)
+    Builder.CreateCall(RegSurface, {RegGlobalsFn->arg_begin(), Addr, Name, Name,
+                                    Data, Extern});
   Builder.CreateBr(IfEndBB);
   Switch->addCase(Builder.getInt32(llvm::offloading::OffloadGlobalSurfaceEntry),
                   SwSurfaceBB);
 
   // Create texture variable registration code.
   Builder.SetInsertPoint(SwTextureBB);
-  Builder.CreateCall(RegTexture, {RegGlobalsFn->arg_begin(), Addr, Name, Name,
-                                  Data, Normalized, Extern});
+  if (EmitSurfacesAndTextures)
+    Builder.CreateCall(RegTexture, {RegGlobalsFn->arg_begin(), Addr, Name, Name,
+                                    Data, Normalized, Extern});
   Builder.CreateBr(IfEndBB);
   Switch->addCase(Builder.getInt32(llvm::offloading::OffloadGlobalTextureEntry),
                   SwTextureBB);
@@ -497,18 +507,21 @@ Function *createRegisterGlobalsFunction(Module &M, bool IsHIP) {
 // Create the constructor and destructor to register the fatbinary with the CUDA
 // runtime.
 void createRegisterFatbinFunction(Module &M, GlobalVariable *FatbinDesc,
-                                  bool IsHIP) {
+                                  bool IsHIP,
+                                  std::optional<EntryArrayTy> EntryArrayOpt,
+                                  StringRef Suffix,
+                                  bool EmitSurfacesAndTextures) {
   LLVMContext &C = M.getContext();
   auto *CtorFuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
-  auto *CtorFunc =
-      Function::Create(CtorFuncTy, GlobalValue::InternalLinkage,
-                       IsHIP ? ".hip.fatbin_reg" : ".cuda.fatbin_reg", &M);
+  auto *CtorFunc = Function::Create(
+      CtorFuncTy, GlobalValue::InternalLinkage,
+      (IsHIP ? ".hip.fatbin_reg" : ".cuda.fatbin_reg") + Suffix, &M);
   CtorFunc->setSection(".text.startup");
 
   auto *DtorFuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
-  auto *DtorFunc =
-      Function::Create(DtorFuncTy, GlobalValue::InternalLinkage,
-                       IsHIP ? ".hip.fatbin_unreg" : ".cuda.fatbin_unreg", &M);
+  auto *DtorFunc = Function::Create(
+      DtorFuncTy, GlobalValue::InternalLinkage,
+      (IsHIP ? ".hip.fatbin_unreg" : ".cuda.fatbin_unreg") + Suffix, &M);
   DtorFunc->setSection(".text.startup");
 
   auto *PtrTy = PointerType::getUnqual(C);
@@ -536,7 +549,7 @@ void createRegisterFatbinFunction(Module &M, GlobalVariable *FatbinDesc,
   auto *BinaryHandleGlobal = new llvm::GlobalVariable(
       M, PtrTy, false, llvm::GlobalValue::InternalLinkage,
       llvm::ConstantPointerNull::get(PtrTy),
-      IsHIP ? ".hip.binary_handle" : ".cuda.binary_handle");
+      (IsHIP ? ".hip.binary_handle" : ".cuda.binary_handle") + Suffix);
 
   // Create the constructor to register this image with the runtime.
   IRBuilder<> CtorBuilder(BasicBlock::Create(C, "entry", CtorFunc));
@@ -546,7 +559,16 @@ void createRegisterFatbinFunction(Module &M, GlobalVariable *FatbinDesc,
   CtorBuilder.CreateAlignedStore(
       Handle, BinaryHandleGlobal,
       Align(M.getDataLayout().getPointerTypeSize(PtrTy)));
-  CtorBuilder.CreateCall(createRegisterGlobalsFunction(M, IsHIP), Handle);
+  EntryArrayTy EntryArray =
+      (EntryArrayOpt ? *EntryArrayOpt
+                     : (IsHIP ? offloading::getOffloadEntryArray(
+                                    M, "hip_offloading_entries")
+                              : offloading::getOffloadEntryArray(
+                                    M, "cuda_offloading_entries")));
+  CtorBuilder.CreateCall(createRegisterGlobalsFunction(M, IsHIP, EntryArray,
+                                                       Suffix,
+                                                       EmitSurfacesAndTextures),
+                         Handle);
   if (!IsHIP)
     CtorBuilder.CreateCall(RegFatbinEnd, Handle);
   CtorBuilder.CreateCall(AtExit, DtorFunc);
@@ -568,32 +590,45 @@ void createRegisterFatbinFunction(Module &M, GlobalVariable *FatbinDesc,
 
 } // namespace
 
-Error wrapOpenMPBinaries(Module &M, ArrayRef<ArrayRef<char>> Images) {
-  GlobalVariable *Desc = createBinDesc(M, Images);
+Error OffloadWrapper::wrapOpenMPBinaries(
+    Module &M, ArrayRef<ArrayRef<char>> Images,
+    std::optional<EntryArrayTy> EntryArray) const {
+  GlobalVariable *Desc = createBinDesc(
+      M, Images,
+      EntryArray
+          ? *EntryArray
+          : offloading::getOffloadEntryArray(M, "omp_offloading_entries"),
+      Suffix);
   if (!Desc)
     return createStringError(inconvertibleErrorCode(),
                              "No binary descriptors created.");
-  createRegisterFunction(M, Desc);
-  createUnregisterFunction(M, Desc);
+  createRegisterFunction(M, Desc, Suffix);
+  createUnregisterFunction(M, Desc, Suffix);
   return Error::success();
 }
 
-Error wrapCudaBinary(Module &M, ArrayRef<char> Image) {
-  GlobalVariable *Desc = createFatbinDesc(M, Image, /* IsHIP */ false);
+Error OffloadWrapper::wrapCudaBinary(
+    Module &M, ArrayRef<char> Image,
+    std::optional<EntryArrayTy> EntryArray) const {
+  GlobalVariable *Desc = createFatbinDesc(M, Image, /* IsHIP */ false, Suffix);
   if (!Desc)
     return createStringError(inconvertibleErrorCode(),
                              "No fatinbary section created.");
 
-  createRegisterFatbinFunction(M, Desc, /* IsHIP */ false);
+  createRegisterFatbinFunction(M, Desc, /* IsHIP */ false, EntryArray, Suffix,
+                               EmitSurfacesAndTextures);
   return Error::success();
 }
 
-Error wrapHIPBinary(Module &M, ArrayRef<char> Image) {
-  GlobalVariable *Desc = createFatbinDesc(M, Image, /* IsHIP */ true);
+Error OffloadWrapper::wrapHIPBinary(
+    Module &M, ArrayRef<char> Image,
+    std::optional<EntryArrayTy> EntryArray) const {
+  GlobalVariable *Desc = createFatbinDesc(M, Image, /* IsHIP */ true, Suffix);
   if (!Desc)
     return createStringError(inconvertibleErrorCode(),
                              "No fatinbary section created.");
 
-  createRegisterFatbinFunction(M, Desc, /* IsHIP */ true);
+  createRegisterFatbinFunction(M, Desc, /* IsHIP */ true, EntryArray, Suffix,
+                               EmitSurfacesAndTextures);
   return Error::success();
 }
diff --git a/llvm/lib/Frontend/Offloading/Utility.cpp b/llvm/lib/Frontend/Offloading/Utility.cpp
index 25f609517ebeb7..531919bccb94e3 100644
--- a/llvm/lib/Frontend/Offloading/Utility.cpp
+++ b/llvm/lib/Frontend/Offloading/Utility.cpp
@@ -1,4 +1,4 @@
-//===- Utility.cpp ------ Collection of geneirc offloading utilities ------===//
+//===- Utility.cpp ------ Collection of generic offloading utilities ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -28,11 +28,10 @@ StructType *offloading::getEntryTy(Module &M) {
 }
 
 // TODO: Rework this interface to be more generic.
-void offloading::emitOffloadingEntry(Module &M, Constant *Addr, StringRef Name,
-                                     uint64_t Size, int32_t Flags, int32_t Data,
-                                     StringRef SectionName) {
-  llvm::Triple Triple(M.getTargetTriple());
-
+std::pair<Constant *, GlobalVariable *>
+offloading::getOffloadingEntryInitializer(Module &M, Constant *Addr,
+                                          StringRef Name, uint64_t Size,
+                                          int32_t Flags, int32_t Data) {
   Type *Int8PtrTy = PointerType::getUnqual(M.getContext());
   Type *Int32Ty = Type::getInt32Ty(M.getContext());
   Type *SizeTy = M.getDataLayout().getIntPtrType(M.getContext());
@@ -54,6 +53,16 @@ void offloading::emitOffloadingEntry(Module &M, Constant *Addr, StringRef Name,
       ConstantInt::get(Int32Ty, Data),
   };
   Constant *EntryInitializer = ConstantStruct::get(getEntryTy(M), EntryData);
+  return {EntryInitializer, Str};
+}
+
+void offloading::emitOffloadingEntry(Module &M, Constant *Addr, StringRef Name,
+                                     uint64_t Size, int32_t Flags, int32_t Data,
+                                     StringRef SectionName) {
+  llvm::Triple Triple(M.getTargetTriple());
+
+  auto [EntryInitializer, NameGV] =
+      getOffloadingEntryInitializer(M, Addr, Name, Size, Flags, Data);
 
   auto *Entry = new GlobalVariable(
       M, getEntryTy(M),

>From 543c58cb03a9b506c4a4e5b65fe0118b9c984dc0 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Sun, 14 Jan 2024 21:42:18 +0000
Subject: [PATCH 3/3] Address comments

---
 .../ClangLinkerWrapper.cpp                    |  9 +--
 .../llvm/Frontend/Offloading/OffloadWrapper.h | 77 +++++++++----------
 .../Frontend/Offloading/OffloadWrapper.cpp    | 31 ++++----
 3 files changed, 53 insertions(+), 64 deletions(-)

diff --git a/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp b/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
index ebe8b634c7ae73..c30d66821dae4e 100644
--- a/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
+++ b/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
@@ -906,18 +906,15 @@ wrapDeviceImages(ArrayRef<std::unique_ptr<MemoryBuffer>> Buffers,
 
   switch (Kind) {
   case OFK_OpenMP:
-    if (Error Err =
-            offloading::OffloadWrapper().wrapOpenMPBinaries(M, BuffersToWrap))
+    if (Error Err = offloading::wrapOpenMPBinaries(M, BuffersToWrap))
       return std::move(Err);
     break;
   case OFK_Cuda:
-    if (Error Err = offloading::OffloadWrapper().wrapCudaBinary(
-            M, BuffersToWrap.front()))
+    if (Error Err = offloading::wrapCudaBinary(M, BuffersToWrap.front()))
       return std::move(Err);
     break;
   case OFK_HIP:
-    if (Error Err = offloading::OffloadWrapper().wrapHIPBinary(
-            M, BuffersToWrap.front()))
+    if (Error Err = offloading::wrapHIPBinary(M, BuffersToWrap.front()))
       return std::move(Err);
     break;
   default:
diff --git a/llvm/include/llvm/Frontend/Offloading/OffloadWrapper.h b/llvm/include/llvm/Frontend/Offloading/OffloadWrapper.h
index 6b23f875a8f15f..f6ab1f475cdb90 100644
--- a/llvm/include/llvm/Frontend/Offloading/OffloadWrapper.h
+++ b/llvm/include/llvm/Frontend/Offloading/OffloadWrapper.h
@@ -14,48 +14,41 @@
 
 namespace llvm {
 namespace offloading {
-/// Class for embedding and registering offloading images and related objects in
-/// a Module.
-class OffloadWrapper {
-public:
-  using EntryArrayTy = std::pair<GlobalVariable *, GlobalVariable *>;
-
-  OffloadWrapper(const Twine &Suffix = "", bool EmitSurfacesAndTextures = true)
-      : Suffix(Suffix.str()), EmitSurfacesAndTextures(EmitSurfacesAndTextures) {
-  }
-
-  /// Wraps the input device images into the module \p M as global symbols and
-  /// registers the images with the OpenMP Offloading runtime libomptarget.
-  /// \param EntryArray Optional pair pointing to the `__start` and `__stop`
-  /// symbols holding the `__tgt_offload_entry` array.
-  llvm::Error wrapOpenMPBinaries(
-      llvm::Module &M, llvm::ArrayRef<llvm::ArrayRef<char>> Images,
-      std::optional<EntryArrayTy> EntryArray = std::nullopt) const;
-
-  /// Wraps the input fatbinary image into the module \p M as global symbols and
-  /// registers the images with the CUDA runtime.
-  /// \param EntryArray Optional pair pointing to the `__start` and `__stop`
-  /// symbols holding the `__tgt_offload_entry` array.
-  llvm::Error
-  wrapCudaBinary(llvm::Module &M, llvm::ArrayRef<char> Images,
-                 std::optional<EntryArrayTy> EntryArray = std::nullopt) const;
-
-  /// Wraps the input bundled image into the module \p M as global symbols and
-  /// registers the images with the HIP runtime.
-  /// \param EntryArray Optional pair pointing to the `__start` and `__stop`
-  /// symbols holding the `__tgt_offload_entry` array.
-  llvm::Error
-  wrapHIPBinary(llvm::Module &M, llvm::ArrayRef<char> Images,
-                std::optional<EntryArrayTy> EntryArray = std::nullopt) const;
-
-protected:
-  /// Suffix used when emitting symbols. It defaults to the empty string.
-  std::string Suffix;
-
-  /// Whether to emit surface and textures registration code. It defaults to
-  /// false.
-  bool EmitSurfacesAndTextures;
-};
+using EntryArrayTy = std::pair<GlobalVariable *, GlobalVariable *>;
+/// Wraps the input device images into the module \p M as global symbols and
+/// registers the images with the OpenMP Offloading runtime libomptarget.
+/// \param EntryArray Optional pair pointing to the `__start` and `__stop`
+/// symbols holding the `__tgt_offload_entry` array.
+/// \param Suffix An optional suffix appended to the emitted symbols.
+llvm::Error
+wrapOpenMPBinaries(llvm::Module &M, llvm::ArrayRef<llvm::ArrayRef<char>> Images,
+                   std::optional<EntryArrayTy> EntryArray = std::nullopt,
+                   llvm::StringRef Suffix = "");
+
+/// Wraps the input fatbinary image into the module \p M as global symbols and
+/// registers the images with the CUDA runtime.
+/// \param EntryArray Optional pair pointing to the `__start` and `__stop`
+/// symbols holding the `__tgt_offload_entry` array.
+/// \param Suffix An optional suffix appended to the emitted symbols.
+/// \param EmitSurfacesAndTextures Whether to emit surface and textures
+/// registration code. It defaults to false.
+llvm::Error
+wrapCudaBinary(llvm::Module &M, llvm::ArrayRef<char> Images,
+               std::optional<EntryArrayTy> EntryArray = std::nullopt,
+               llvm::StringRef Suffix = "",
+               bool EmitSurfacesAndTextures = true);
+
+/// Wraps the input bundled image into the module \p M as global symbols and
+/// registers the images with the HIP runtime.
+/// \param EntryArray Optional pair pointing to the `__start` and `__stop`
+/// symbols holding the `__tgt_offload_entry` array.
+/// \param Suffix An optional suffix appended to the emitted symbols.
+/// \param EmitSurfacesAndTextures Whether to emit surface and textures
+/// registration code. It defaults to false.
+llvm::Error wrapHIPBinary(llvm::Module &M, llvm::ArrayRef<char> Images,
+                          std::optional<EntryArrayTy> EntryArray = std::nullopt,
+                          llvm::StringRef Suffix = "",
+                          bool EmitSurfacesAndTextures = true);
 } // namespace offloading
 } // namespace llvm
 
diff --git a/llvm/lib/Frontend/Offloading/OffloadWrapper.cpp b/llvm/lib/Frontend/Offloading/OffloadWrapper.cpp
index f34a879b99dd02..2cc5e110510abd 100644
--- a/llvm/lib/Frontend/Offloading/OffloadWrapper.cpp
+++ b/llvm/lib/Frontend/Offloading/OffloadWrapper.cpp
@@ -24,8 +24,6 @@ using namespace llvm;
 using namespace llvm::offloading;
 
 namespace {
-using EntryArrayTy = OffloadWrapper::EntryArrayTy;
-
 /// Magic number that begins the section containing the CUDA fatbinary.
 constexpr unsigned CudaFatMagic = 0x466243b1;
 constexpr unsigned HIPFatMagic = 0x48495046;
@@ -587,12 +585,11 @@ void createRegisterFatbinFunction(Module &M, GlobalVariable *FatbinDesc,
   // Add this function to constructors.
   appendToGlobalCtors(M, CtorFunc, /*Priority*/ 1);
 }
-
 } // namespace
 
-Error OffloadWrapper::wrapOpenMPBinaries(
-    Module &M, ArrayRef<ArrayRef<char>> Images,
-    std::optional<EntryArrayTy> EntryArray) const {
+Error offloading::wrapOpenMPBinaries(Module &M, ArrayRef<ArrayRef<char>> Images,
+                                     std::optional<EntryArrayTy> EntryArray,
+                                     llvm::StringRef Suffix) {
   GlobalVariable *Desc = createBinDesc(
       M, Images,
       EntryArray
@@ -607,28 +604,30 @@ Error OffloadWrapper::wrapOpenMPBinaries(
   return Error::success();
 }
 
-Error OffloadWrapper::wrapCudaBinary(
-    Module &M, ArrayRef<char> Image,
-    std::optional<EntryArrayTy> EntryArray) const {
-  GlobalVariable *Desc = createFatbinDesc(M, Image, /* IsHIP */ false, Suffix);
+Error offloading::wrapCudaBinary(Module &M, ArrayRef<char> Image,
+                                 std::optional<EntryArrayTy> EntryArray,
+                                 llvm::StringRef Suffix,
+                                 bool EmitSurfacesAndTextures) {
+  GlobalVariable *Desc = createFatbinDesc(M, Image, /*IsHip=*/false, Suffix);
   if (!Desc)
     return createStringError(inconvertibleErrorCode(),
                              "No fatinbary section created.");
 
-  createRegisterFatbinFunction(M, Desc, /* IsHIP */ false, EntryArray, Suffix,
+  createRegisterFatbinFunction(M, Desc, /*IsHip=*/false, EntryArray, Suffix,
                                EmitSurfacesAndTextures);
   return Error::success();
 }
 
-Error OffloadWrapper::wrapHIPBinary(
-    Module &M, ArrayRef<char> Image,
-    std::optional<EntryArrayTy> EntryArray) const {
-  GlobalVariable *Desc = createFatbinDesc(M, Image, /* IsHIP */ true, Suffix);
+Error offloading::wrapHIPBinary(Module &M, ArrayRef<char> Image,
+                                std::optional<EntryArrayTy> EntryArray,
+                                llvm::StringRef Suffix,
+                                bool EmitSurfacesAndTextures) {
+  GlobalVariable *Desc = createFatbinDesc(M, Image, /*IsHip=*/true, Suffix);
   if (!Desc)
     return createStringError(inconvertibleErrorCode(),
                              "No fatinbary section created.");
 
-  createRegisterFatbinFunction(M, Desc, /* IsHIP */ true, EntryArray, Suffix,
+  createRegisterFatbinFunction(M, Desc, /*IsHip=*/true, EntryArray, Suffix,
                                EmitSurfacesAndTextures);
   return Error::success();
 }



More information about the llvm-commits mailing list