[OFFLOAD][L0] Use device ELF as image for L0ProgramTy

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,42 +57,17 @@ 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);
 
+  bool IsLibModule =
+      BuildOptions.find("-library-compilation") != std::string::npos;
+
   // Add required flag to enable dynamic linking.
   if (IsLibModule)
     BuildOptions += " -library-compilation ";
@@ -128,7 +103,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 +122,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 +215,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 +231,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 +430,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);