[OFFLOAD][L0] Use device ELF as image for L0ProgramTy (#20712)

Decouple building and linking the device image from creating the
L0ProgramTy representing it. Use the GPU ELF image as the internal Image
data for L0ProgramTy. This improves compatibility with the main library
which expects to be an ELF in several places.

Author: adurang

offload/plugins-nextgen/level_zero/include/L0Device.h

@@ -307,9 +307,14 @@ class L0DeviceTy final : public GenericDeviceTy {
   }
 
   // add a new program to the device. Return a reference to the new program
-  L0ProgramTy &addProgram(int32_t ImageId,
-                          std::unique_ptr<MemoryBuffer> &&Image) {
-    Programs.emplace_back(ImageId, *this, std::move(Image));
+  Expected<L0ProgramTy &> addProgram(int32_t ImageId,
+                                     L0ProgramBuilderTy &Builder) {
+    auto ImageOrErr = Builder.getELF();
+    if (!ImageOrErr)
+      return ImageOrErr.takeError();
+    Programs.emplace_back(ImageId, *this, std::move(*ImageOrErr),
+                          Builder.getGlobalModule(),
+                          std::move(Builder.getModules()));
     return Programs.back();
   }
 

offload/plugins-nextgen/level_zero/include/L0Program.h

@@ -33,6 +33,47 @@ struct ProgramDataTy {
   int TeamsThreadLimit = 0;
 };
 
+class L0ProgramBuilderTy {
+  L0DeviceTy &Device;
+  std::unique_ptr<MemoryBuffer> Image;
+  /// Handle multiple modules within a single target image
+  llvm::SmallVector<ze_module_handle_t> Modules;
+
+  /// Module that contains global data including device RTL
+  ze_module_handle_t GlobalModule = nullptr;
+
+  /// Requires module link
+  bool RequiresModuleLink = false;
+
+  /// Build a single module with the given image, build option, and format.
+  Error addModule(const size_t Size, const uint8_t *Image,
+                  const std::string_view BuildOption,
+                  ze_module_format_t Format);
+
+  Error linkModules();
+
+public:
+  L0ProgramBuilderTy(L0DeviceTy &Device, std::unique_ptr<MemoryBuffer> &&Image)
+      : Device(Device), Image(std::move(Image)) {}
+  ~L0ProgramBuilderTy() = default;
+
+  L0DeviceTy &getL0Device() const { return Device; }
+  ze_module_handle_t getGlobalModule() const { return GlobalModule; }
+  llvm::SmallVector<ze_module_handle_t> &getModules() { return Modules; }
+
+  const void *getStart() const { return Image->getBufferStart(); }
+  size_t getSize() const { return Image->getBufferSize(); }
+
+  MemoryBufferRef getMemoryBuffer() const {
+    return MemoryBufferRef(StringRef((const char *)getStart(), getSize()),
+                           "Image");
+  }
+  Error buildModules(const std::string_view BuildOptions);
+
+  /// Retrieve the ELF binary for the program
+  Expected<std::unique_ptr<MemoryBuffer>> getELF();
+};
+
 /// Level Zero program that can contain multiple modules.
 class L0ProgramTy : public DeviceImageTy {
   /// Handle multiple modules within a single target image
@@ -49,32 +90,17 @@ class L0ProgramTy : public DeviceImageTy {
   /// Module that contains global data including device RTL
   ze_module_handle_t GlobalModule = nullptr;
 
-  /// Requires module link
-  bool RequiresModuleLink = false;
-
-  /// Is this module library
-  bool IsLibModule = false;
-
-  /// Build a single module with the given image, build option, and format.
-  Error addModule(const size_t Size, const uint8_t *Image,
-                  const std::string_view BuildOption,
-                  ze_module_format_t Format);
-  /// Read file and return the size of the binary if successful.
-  size_t readFile(const char *FileName, std::vector<uint8_t> &OutFile) const;
-  void replaceDriverOptsWithBackendOpts(const L0DeviceTy &Device,
-                                        std::string &Options) const;
-
-  /// Check if the image should be handled as a library module
-  void setLibModule();
-
   L0DeviceTy &getL0Device() const;
 
 public:
   L0ProgramTy() = delete;
 
   L0ProgramTy(int32_t ImageId, GenericDeviceTy &Device,
-              std::unique_ptr<MemoryBuffer> Image)
-      : DeviceImageTy(ImageId, Device, std::move(Image)) {}
+              std::unique_ptr<MemoryBuffer> Image,
+              ze_module_handle_t GlobalModule,
+              llvm::SmallVector<ze_module_handle_t> &&Modules)
+      : DeviceImageTy(ImageId, Device, std::move(Image)),
+        Modules(std::move(Modules)), GlobalModule(GlobalModule) {}
   ~L0ProgramTy() {}
 
   L0ProgramTy(const L0ProgramTy &other) = delete;
@@ -88,12 +114,6 @@ class L0ProgramTy : public DeviceImageTy {
     return static_cast<L0ProgramTy &>(Device);
   }
 
-  /// Build modules from the target image description
-  Error buildModules(const std::string_view BuildOptions);
-
-  /// Link modules stored in \p Modules.
-  Error linkModules();
-
   /// Loads the kernels names from all modules
   Error loadModuleKernels();
 

offload/plugins-nextgen/level_zero/src/L0Device.cpp

@@ -250,13 +250,15 @@ L0DeviceTy::loadBinaryImpl(std::unique_ptr<MemoryBuffer> &&TgtImage,
 
   CompilationOptions += " ";
   CompilationOptions += Options.InternalCompilationOptions;
-  auto &Program = addProgram(ImageId, std::move(TgtImage));
 
-  if (auto Err = Program.buildModules(CompilationOptions))
+  L0ProgramBuilderTy Builder(*this, std::move(TgtImage));
+  if (auto Err = Builder.buildModules(CompilationOptions))
     return std::move(Err);
 
-  if (auto Err = Program.linkModules())
-    return std::move(Err);
+  auto ProgramOrErr = addProgram(ImageId, Builder);
+  if (!ProgramOrErr)
+    return ProgramOrErr.takeError();
+  auto &Program = *ProgramOrErr;
 
   if (auto Err = Program.loadModuleKernels())
     return std::move(Err);

offload/plugins-nextgen/level_zero/src/L0DynWrapper.cpp

@@ -81,6 +81,7 @@ DLWRAP(zeMemGetAddressRange, 4)
 DLWRAP(zeMemGetAllocProperties, 4)
 DLWRAP(zeModuleDynamicLink, 3)
 DLWRAP(zeModuleGetGlobalPointer, 4)
+DLWRAP(zeModuleGetNativeBinary, 3)
 DLWRAP(zesDeviceEnumMemoryModules, 3)
 DLWRAP(zesMemoryGetState, 2)
 DLWRAP(zeCommandListHostSynchronize, 2)

offload/plugins-nextgen/level_zero/src/L0Program.cpp

@@ -57,45 +57,16 @@ Error L0ProgramTy::deinit() {
   return Plugin::success();
 }
 
-void L0ProgramTy::setLibModule() {
-#if _WIN32
-  return;
-#else
-  // Check if the image belongs to a dynamic library
-  Dl_info DLI{nullptr, nullptr, nullptr, nullptr};
-  if (dladdr(getStart(), &DLI) && DLI.dli_fname) {
-    std::vector<uint8_t> FileBin;
-    auto Size = readFile(DLI.dli_fname, FileBin);
-    if (Size) {
-      auto MB = MemoryBuffer::getMemBuffer(
-          StringRef(reinterpret_cast<const char *>(FileBin.data()), Size),
-          /*BufferName=*/"", /*RequiresNullTerminator=*/false);
-      auto ELF = ELFObjectFileBase::createELFObjectFile(MB->getMemBufferRef());
-      if (ELF) {
-        if (auto *Obj = dyn_cast<ELF64LEObjectFile>((*ELF).get())) {
-          const auto Header = Obj->getELFFile().getHeader();
-          if (Header.e_type == ELF::ET_DYN) {
-            DP("Processing current image as library\n");
-            IsLibModule = true;
-          }
-        }
-      }
-    }
-  }
-#endif // _WIN32
-}
-
-Error L0ProgramTy::addModule(size_t Size, const uint8_t *Image,
-                             const std::string_view CommonBuildOptions,
-                             ze_module_format_t Format) {
+Error L0ProgramBuilderTy::addModule(size_t Size, const uint8_t *Image,
+                                    const std::string_view CommonBuildOptions,
+                                    ze_module_format_t Format) {
   const ze_module_constants_t SpecConstants =
       LevelZeroPluginTy::getOptions().CommonSpecConstants.getModuleConstants();
   auto &l0Device = getL0Device();
   std::string BuildOptions(CommonBuildOptions);
 
-  // Add required flag to enable dynamic linking.
-  if (IsLibModule)
-    BuildOptions += " -library-compilation ";
+  bool IsLibModule =
+      BuildOptions.find("-library-compilation") != std::string::npos;
 
   ze_module_desc_t ModuleDesc{};
   ModuleDesc.stype = ZE_STRUCTURE_TYPE_MODULE_DESC;
@@ -128,7 +99,7 @@ Error L0ProgramTy::addModule(size_t Size, const uint8_t *Image,
   return Plugin::success();
 }
 
-Error L0ProgramTy::linkModules() {
+Error L0ProgramBuilderTy::linkModules() {
   auto &l0Device = getL0Device();
   if (!RequiresModuleLink) {
     DP("Module link is not required\n");
@@ -147,24 +118,8 @@ Error L0ProgramTy::linkModules() {
   return Plugin::success();
 }
 
-size_t L0ProgramTy::readFile(const char *FileName,
-                             std::vector<uint8_t> &OutFile) const {
-  std::ifstream IFS(FileName, std::ios::binary);
-  if (!IFS.good())
-    return 0;
-  IFS.seekg(0, IFS.end);
-  auto FileSize = static_cast<size_t>(IFS.tellg());
-  OutFile.resize(FileSize);
-  IFS.seekg(0);
-  if (!IFS.read(reinterpret_cast<char *>(OutFile.data()), FileSize)) {
-    OutFile.clear();
-    return 0;
-  }
-  return FileSize;
-}
-
-void L0ProgramTy::replaceDriverOptsWithBackendOpts(const L0DeviceTy &Device,
-                                                   std::string &Options) const {
+static void replaceDriverOptsWithBackendOpts(const L0DeviceTy &Device,
+                                             std::string &Options) {
   // Options that need to be replaced with backend-specific options
   static const struct {
     std::string Option;
@@ -256,7 +211,7 @@ bool isValidOneOmpImage(StringRef Image, uint64_t &MajorVer,
   return Res;
 }
 
-Error L0ProgramTy::buildModules(const std::string_view BuildOptions) {
+Error L0ProgramBuilderTy::buildModules(const std::string_view BuildOptions) {
   auto &l0Device = getL0Device();
   auto Image = getMemoryBuffer();
   if (identify_magic(Image.getBuffer()) == file_magic::spirv_object) {
@@ -272,8 +227,6 @@ Error L0ProgramTy::buildModules(const std::string_view BuildOptions) {
     return Plugin::error(ErrorCode::UNKNOWN, "Invalid oneAPI OpenMP image");
   }
 
-  setLibModule();
-
   // Iterate over the images and pick the first one that fits.
   uint64_t ImageCount = 0;
   struct V1ImageInfo {
@@ -473,12 +426,32 @@ Error L0ProgramTy::buildModules(const std::string_view BuildOptions) {
     }
     DP("Created module from image #%" PRIu64 ".\n", Idx);
 
+    if (RequiresModuleLink) {
+      DP("Linking modules after adding image #%" PRIu64 ".\n", Idx);
+      if (auto Err = linkModules())
+        return Err;
+    }
+
     return Plugin::success();
   }
 
   return Plugin::error(ErrorCode::UNKNOWN, "Failed to create program modules.");
 }
 
+Expected<std::unique_ptr<MemoryBuffer>> L0ProgramBuilderTy::getELF() {
+  assert(GlobalModule != nullptr && "GlobalModule is null");
+
+  size_t Size = 0;
+
+  CALL_ZE_RET_ERROR(zeModuleGetNativeBinary, GlobalModule, &Size, nullptr);
+  std::vector<uint8_t> ELFData(Size);
+  CALL_ZE_RET_ERROR(zeModuleGetNativeBinary, GlobalModule, &Size,
+                    ELFData.data());
+  return MemoryBuffer::getMemBufferCopy(
+      StringRef(reinterpret_cast<const char *>(ELFData.data()), Size),
+      /*BufferName=*/"L0Program ELF");
+}
+
 Expected<void *> L0ProgramTy::getOffloadVarDeviceAddr(const char *CName) const {
   DP("Looking up OpenMP global variable '%s'.\n", CName);