[Offload] Implement 'olIsValidBinary' in offload and clean up (llvm#159658)

Summary:
This exposes the 'isDeviceCompatible' routine for checking if a binary
*can* be loaded. This is useful if people don't want to consume errors
everywhere when figuring out which image to put to what device.

I don't know if this is a good name, I was thining like `olIsCompatible`
or whatever. Let me know what you think.

Long term I'd like to be able to do something similar to what OpenMP
does where we can conditionally only initialize devices if we need them.
That's going to be support needed if we want this to be more
generic.

Author: jhuber6

offload/liboffload/API/Program.td

@@ -24,6 +24,18 @@ def olCreateProgram : Function {
     let returns = [];
 }
 
+def olIsValidBinary : Function {
+    let desc = "Validate if the binary image pointed to by `ProgData` is compatible with the device.";
+    let details = ["The provided `ProgData` will not be loaded onto the device"];
+    let params = [
+        Param<"ol_device_handle_t", "Device", "handle of the device", PARAM_IN>,
+        Param<"const void*", "ProgData", "pointer to the program binary data", PARAM_IN>,
+        Param<"size_t", "ProgDataSize", "size of the program binary in bytes", PARAM_IN>,
+        Param<"bool*", "Valid", "output is true if the image is compatible", PARAM_OUT>
+    ];
+    let returns = [];
+}
+
 def olDestroyProgram : Function {
     let desc = "Destroy the program and free all underlying resources.";
     let details = [];

offload/liboffload/src/OffloadImpl.cpp

@@ -990,7 +990,6 @@ Error olMemFill_impl(ol_queue_handle_t Queue, void *Ptr, size_t PatternSize,
 
 Error olCreateProgram_impl(ol_device_handle_t Device, const void *ProgData,
                            size_t ProgDataSize, ol_program_handle_t *Program) {
-  // Make a copy of the program binary in case it is released by the caller.
   StringRef Buffer(reinterpret_cast<const char *>(ProgData), ProgDataSize);
   Expected<plugin::DeviceImageTy *> Res =
       Device->Device->loadBinary(Device->Device->Plugin, Buffer);
@@ -1002,6 +1001,14 @@ Error olCreateProgram_impl(ol_device_handle_t Device, const void *ProgData,
   return Error::success();
 }
 
+Error olIsValidBinary_impl(ol_device_handle_t Device, const void *ProgData,
+                           size_t ProgDataSize, bool *IsValid) {
+  StringRef Buffer(reinterpret_cast<const char *>(ProgData), ProgDataSize);
+  *IsValid = Device->Device->Plugin.isDeviceCompatible(
+      Device->Device->getDeviceId(), Buffer);
+  return Error::success();
+}
+
 Error olDestroyProgram_impl(ol_program_handle_t Program) {
   auto &Device = Program->Image->getDevice();
   if (auto Err = Device.unloadBinary(Program->Image))

offload/libomptarget/PluginManager.cpp

@@ -243,7 +243,10 @@ void PluginManager::registerLib(__tgt_bin_desc *Desc) {
     // Scan the RTLs that have associated images until we find one that supports
     // the current image.
     for (auto &R : plugins()) {
-      if (!R.is_plugin_compatible(Img))
+      StringRef Buffer(reinterpret_cast<const char *>(Img->ImageStart),
+                       utils::getPtrDiff(Img->ImageEnd, Img->ImageStart));
+
+      if (!R.isPluginCompatible(Buffer))
         continue;
 
       if (!initializePlugin(R))
@@ -266,7 +269,7 @@ void PluginManager::registerLib(__tgt_bin_desc *Desc) {
           continue;
         }
 
-        if (!R.is_device_compatible(DeviceId, Img))
+        if (!R.isDeviceCompatible(DeviceId, Buffer))
           continue;
 
         DP("Image " DPxMOD " is compatible with RTL %s device %d!\n",

offload/plugins-nextgen/common/include/PluginInterface.h

@@ -1714,10 +1714,10 @@ struct GenericPluginTy {
 
   /// Returns non-zero if the \p Image is compatible with the plugin. This
   /// function does not require the plugin to be initialized before use.
-  int32_t is_plugin_compatible(__tgt_device_image *Image);
+  int32_t isPluginCompatible(StringRef Image);
 
   /// Returns non-zero if the \p Image is compatible with the device.
-  int32_t is_device_compatible(int32_t DeviceId, __tgt_device_image *Image);
+  int32_t isDeviceCompatible(int32_t DeviceId, StringRef Image);
 
   /// Returns non-zero if the plugin device has been initialized.
   int32_t is_device_initialized(int32_t DeviceId) const;

offload/plugins-nextgen/common/src/PluginInterface.cpp

@@ -1981,85 +1981,66 @@ int32_t GenericPluginTy::supports_empty_images() {
   return supportsEmptyImages();
 }
 
-int32_t GenericPluginTy::is_plugin_compatible(__tgt_device_image *Image) {
-  auto T = logger::log<int32_t>(__func__, Image);
-  auto R = [&]() {
-    StringRef Buffer(reinterpret_cast<const char *>(Image->ImageStart),
-                     utils::getPtrDiff(Image->ImageEnd, Image->ImageStart));
-
-    auto HandleError = [&](Error Err) -> bool {
-      [[maybe_unused]] std::string ErrStr = toString(std::move(Err));
-      DP("Failure to check validity of image %p: %s", Image, ErrStr.c_str());
-      return false;
-    };
-    switch (identify_magic(Buffer)) {
-    case file_magic::elf:
-    case file_magic::elf_relocatable:
-    case file_magic::elf_executable:
-    case file_magic::elf_shared_object:
-    case file_magic::elf_core: {
-      auto MatchOrErr = checkELFImage(Buffer);
-      if (Error Err = MatchOrErr.takeError())
-        return HandleError(std::move(Err));
-      return *MatchOrErr;
-    }
-    case file_magic::bitcode: {
-      auto MatchOrErr = checkBitcodeImage(Buffer);
-      if (Error Err = MatchOrErr.takeError())
-        return HandleError(std::move(Err));
-      return *MatchOrErr;
-    }
-    default:
-      return false;
-    }
-  }();
-  T.res(R);
-  return R;
+int32_t GenericPluginTy::isPluginCompatible(StringRef Image) {
+  auto HandleError = [&](Error Err) -> bool {
+    [[maybe_unused]] std::string ErrStr = toString(std::move(Err));
+    DP("Failure to check validity of image %p: %s", Image, ErrStr.c_str());
+    return false;
+  };
+  switch (identify_magic(Image)) {
+  case file_magic::elf:
+  case file_magic::elf_relocatable:
+  case file_magic::elf_executable:
+  case file_magic::elf_shared_object:
+  case file_magic::elf_core: {
+    auto MatchOrErr = checkELFImage(Image);
+    if (Error Err = MatchOrErr.takeError())
+      return HandleError(std::move(Err));
+    return *MatchOrErr;
+  }
+  case file_magic::bitcode: {
+    auto MatchOrErr = checkBitcodeImage(Image);
+    if (Error Err = MatchOrErr.takeError())
+      return HandleError(std::move(Err));
+    return *MatchOrErr;
+  }
+  default:
+    return false;
+  }
 }
 
-int32_t GenericPluginTy::is_device_compatible(int32_t DeviceId,
-                                              __tgt_device_image *Image) {
-  auto T = logger::log<int32_t>(__func__, DeviceId, Image);
-  auto R = [&]() {
-    StringRef Buffer(reinterpret_cast<const char *>(Image->ImageStart),
-                     utils::getPtrDiff(Image->ImageEnd, Image->ImageStart));
-
-    auto HandleError = [&](Error Err) -> bool {
-      [[maybe_unused]] std::string ErrStr = toString(std::move(Err));
-      DP("Failure to check validity of image %p: %s", Image, ErrStr.c_str());
-      return false;
-    };
-    switch (identify_magic(Buffer)) {
-    case file_magic::elf:
-    case file_magic::elf_relocatable:
-    case file_magic::elf_executable:
-    case file_magic::elf_shared_object:
-    case file_magic::elf_core: {
-      auto MatchOrErr = checkELFImage(Buffer);
-      if (Error Err = MatchOrErr.takeError())
-        return HandleError(std::move(Err));
-      if (!*MatchOrErr)
-        return false;
-
-      // Perform plugin-dependent checks for the specific architecture if
-      // needed.
-      auto CompatibleOrErr = isELFCompatible(DeviceId, Buffer);
-      if (Error Err = CompatibleOrErr.takeError())
-        return HandleError(std::move(Err));
-      return *CompatibleOrErr;
-    }
-    case file_magic::bitcode: {
-      auto MatchOrErr = checkBitcodeImage(Buffer);
-      if (Error Err = MatchOrErr.takeError())
-        return HandleError(std::move(Err));
-      return *MatchOrErr;
-    }
-    default:
+int32_t GenericPluginTy::isDeviceCompatible(int32_t DeviceId, StringRef Image) {
+  auto HandleError = [&](Error Err) -> bool {
+    [[maybe_unused]] std::string ErrStr = toString(std::move(Err));
+    DP("Failure to check validity of image %p: %s", Image, ErrStr.c_str());
+    return false;
+  };
+  switch (identify_magic(Image)) {
+  case file_magic::elf:
+  case file_magic::elf_relocatable:
+  case file_magic::elf_executable:
+  case file_magic::elf_shared_object:
+  case file_magic::elf_core: {
+    auto MatchOrErr = checkELFImage(Image);
+    if (Error Err = MatchOrErr.takeError())
+      return HandleError(std::move(Err));
+    if (!*MatchOrErr)
       return false;
-    }
-  }();
-  T.res(R);
-  return R;
+    // Perform plugin-dependent checks for the specific architecture if needed.
+    auto CompatibleOrErr = isELFCompatible(DeviceId, Image);
+    if (Error Err = CompatibleOrErr.takeError())
+      return HandleError(std::move(Err));
+    return *CompatibleOrErr;
+  }
+  case file_magic::bitcode: {
+    auto MatchOrErr = checkBitcodeImage(Image);
+    if (Error Err = MatchOrErr.takeError())
+      return HandleError(std::move(Err));
+    return *MatchOrErr;
+  }
+  default:
+    return false;
+  }
 }
 
 int32_t GenericPluginTy::is_device_initialized(int32_t DeviceId) const {

offload/unittests/OffloadAPI/CMakeLists.txt

@@ -37,6 +37,7 @@ add_offload_unittest("platform"
 
 add_offload_unittest("program"
     program/olCreateProgram.cpp
+    program/olIsValidBinary.cpp
     program/olDestroyProgram.cpp)
 
 add_offload_unittest("queue"

offload/unittests/OffloadAPI/program/olIsValidBinary.cpp

@@ -0,0 +1,49 @@
+//===------- Offload API tests - olIsValidBinary --------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "../common/Fixtures.hpp"
+#include <OffloadAPI.h>
+#include <gtest/gtest.h>
+
+using olIsValidBinaryTest = OffloadDeviceTest;
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olIsValidBinaryTest);
+
+TEST_P(olIsValidBinaryTest, Success) {
+
+  std::unique_ptr<llvm::MemoryBuffer> DeviceBin;
+  ASSERT_TRUE(TestEnvironment::loadDeviceBinary("foo", Device, DeviceBin));
+  ASSERT_GE(DeviceBin->getBufferSize(), 0lu);
+
+  bool IsValid = false;
+  ASSERT_SUCCESS(olIsValidBinary(Device, DeviceBin->getBufferStart(),
+                                 DeviceBin->getBufferSize(), &IsValid));
+  ASSERT_TRUE(IsValid);
+
+  ASSERT_SUCCESS(
+      olIsValidBinary(Device, DeviceBin->getBufferStart(), 0, &IsValid));
+  ASSERT_FALSE(IsValid);
+}
+
+TEST_P(olIsValidBinaryTest, Invalid) {
+
+  std::unique_ptr<llvm::MemoryBuffer> DeviceBin;
+  ASSERT_TRUE(TestEnvironment::loadDeviceBinary("foo", Device, DeviceBin));
+  ASSERT_GE(DeviceBin->getBufferSize(), 0lu);
+
+  bool IsValid = false;
+  ASSERT_SUCCESS(
+      olIsValidBinary(Device, DeviceBin->getBufferStart(), 0, &IsValid));
+  ASSERT_FALSE(IsValid);
+}
+
+TEST_P(olIsValidBinaryTest, NullPointer) {
+  bool IsValid = false;
+  ASSERT_ERROR(OL_ERRC_INVALID_NULL_POINTER,
+               olIsValidBinary(Device, nullptr, 42, &IsValid));
+  ASSERT_FALSE(IsValid);
+}