[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.

Author: fabianmcg

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}

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:

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

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.

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