[UR][HIP] Select the correct binary in a multi-arch bundle (#19092)

This commit changes the HIP adapter to select the correct binary for the
device when a bundle contains binaries built for multiple AMDGPU
architectures.

Similarly to other adapaters, the HIP adapter would previously select
the first 'amdgcn' binary it came across. This works fine for the common
case where the program was compiled for one architecture but may fail
otherwise.

To aid in this, the SYCL runtime passes some extra information into
urDeviceSelectBinary via the pre-existing 'pNext' field of
ur_device_binary_t. It does this only for the HIP backend. The HIP
adapater then parses this binary information as a clang offload bundle,
which conveniently contains specific triple & architecture information
for each binary. For this we re-use the code that the offload adapter
was using, making it common and fixing a bug in the version matching
logic.

Author: frasercrmck

sycl/source/detail/program_manager/program_manager.cpp

@@ -1529,11 +1529,24 @@ RTDeviceBinaryImage *getBinImageFromMultiMap(
   if (DeviceFilteredImgs.empty())
     return nullptr;
 
-  std::vector<ur_device_binary_t> UrBinaries(DeviceFilteredImgs.size());
-  for (uint32_t BinaryCount = 0; BinaryCount < DeviceFilteredImgs.size();
-       BinaryCount++) {
-    UrBinaries[BinaryCount].pDeviceTargetSpec = getUrDeviceTarget(
-        getRawImg(DeviceFilteredImgs[BinaryCount])->DeviceTargetSpec);
+  const size_t NumImgs = DeviceFilteredImgs.size();
+  // Pass extra information to the HIP adapter to aid in binary selection. We
+  // pass it the raw binary as a {ptr, length} pair.
+  std::vector<std::pair<const unsigned char *, size_t>> UrBinariesStorage;
+  if (DeviceImpl.getBackend() == backend::ext_oneapi_hip)
+    UrBinariesStorage.reserve(NumImgs);
+
+  std::vector<ur_device_binary_t> UrBinaries(NumImgs);
+  for (uint32_t BinaryCount = 0; BinaryCount < NumImgs; BinaryCount++) {
+    sycl_device_binary RawImg = getRawImg(DeviceFilteredImgs[BinaryCount]);
+    UrBinaries[BinaryCount].pDeviceTargetSpec =
+        getUrDeviceTarget(RawImg->DeviceTargetSpec);
+    if (DeviceImpl.getBackend() == backend::ext_oneapi_hip) {
+      UrBinariesStorage.emplace_back(
+          RawImg->BinaryStart,
+          std::distance(RawImg->BinaryStart, RawImg->BinaryEnd));
+      UrBinaries[BinaryCount].pNext = &UrBinariesStorage[BinaryCount];
+    }
   }
 
   uint32_t ImgInd = 0;

unified-runtime/source/adapters/hip/CMakeLists.txt

@@ -157,6 +157,7 @@ if("${UR_HIP_PLATFORM}" STREQUAL "AMD")
     ${PROJECT_NAME}::headers
     ${PROJECT_NAME}::common
     ${PROJECT_NAME}::umf
+    ur_common
     rocmdrv
     )
 

unified-runtime/source/adapters/hip/device.cpp

@@ -12,6 +12,7 @@
 #include "adapter.hpp"
 #include "context.hpp"
 #include "event.hpp"
+#include "offload_bundle_parser.hpp"
 
 #include <array>
 #include <hip/hip_runtime.h>
@@ -1158,25 +1159,57 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceCreateWithNativeHandle(
 
 /// \return UR_RESULT_SUCCESS If available, the first binary that is PTX
 ///
-UR_APIEXPORT ur_result_t UR_APICALL
-urDeviceSelectBinary(ur_device_handle_t, const ur_device_binary_t *pBinaries,
-                     uint32_t NumBinaries, uint32_t *pSelectedBinary) {
-  // Look for an image for the HIP target, and return the first one that is
-  // found
+UR_APIEXPORT ur_result_t UR_APICALL urDeviceSelectBinary(
+    ur_device_handle_t hDevice, const ur_device_binary_t *pBinaries,
+    uint32_t NumBinaries, uint32_t *pSelectedBinary) {
+  // Look for an image for the HIP target. If we have a clang offload bundle,
+  // try to return an exact match on the architecture to ensure we select the
+  // correct binary in a multi-architecture bundle. If we can't find an exact
+  // match, we return the first one that is found.
 #if defined(__HIP_PLATFORM_AMD__)
   const char *BinaryType = UR_DEVICE_BINARY_TARGET_AMDGCN;
 #elif defined(__HIP_PLATFORM_NVIDIA__)
   const char *BinaryType = UR_DEVICE_BINARY_TARGET_NVPTX64;
 #else
 #error("Must define exactly one of __HIP_PLATFORM_AMD__ or __HIP_PLATFORM_NVIDIA__");
 #endif
+  std::optional<uint32_t> FirstBackupCandidate;
+
+  hipDeviceProp_t Props;
+  UR_CHECK_ERROR(hipGetDeviceProperties(&Props, hDevice->get()));
+
+  // The arch name is the 'gfxABC' architecture, and occasionally ends with
+  // architecture feature strings separated by colons. We're only interested in
+  // the first part, so take up until the (optional) first colon.
+  std::string_view ArchName = Props.gcnArchName;
+  ArchName = ArchName.substr(0, ArchName.find_first_of(":"));
+
   for (uint32_t i = 0; i < NumBinaries; i++) {
-    if (strcmp(pBinaries[i].pDeviceTargetSpec, BinaryType) == 0) {
-      *pSelectedBinary = i;
-      return UR_RESULT_SUCCESS;
+    if (strcmp(pBinaries[i].pDeviceTargetSpec, BinaryType) != 0) {
+      continue;
+    }
+    // If we've been given the actual binary by the SYCL runtime to inspect,
+    // attempt to parse it as a clang offload bundle.
+    using BinaryBlobTy = std::pair<const unsigned char *, size_t>;
+    if (auto *const BinaryBlob = (const BinaryBlobTy *)pBinaries[i].pNext) {
+      if (auto Parser = HipOffloadBundleParser::load(BinaryBlob->first,
+                                                     BinaryBlob->second)) {
+        if (Parser->containsBundle(ArchName)) {
+          *pSelectedBinary = i;
+          return UR_RESULT_SUCCESS;
+        }
+      }
+    }
+    if (!FirstBackupCandidate) {
+      FirstBackupCandidate = i;
     }
   }
 
+  if (FirstBackupCandidate) {
+    *pSelectedBinary = *FirstBackupCandidate;
+    return UR_RESULT_SUCCESS;
+  }
+
   // No image can be loaded for the given device
   return UR_RESULT_ERROR_INVALID_BINARY;
 }

unified-runtime/source/adapters/offload/CMakeLists.txt

@@ -62,6 +62,7 @@ target_link_libraries(${TARGET_NAME} PRIVATE
     ${PROJECT_NAME}::headers
     ${PROJECT_NAME}::common
     ${PROJECT_NAME}::umf
+    ur_common
     ${UR_OFFLOAD_INSTALL_DIR}/lib/libLLVMOffload.so
     ${ADDITIONAL_LINK_LIBS}
 )

unified-runtime/source/adapters/offload/program.cpp

@@ -14,6 +14,7 @@
 
 #include "context.hpp"
 #include "device.hpp"
+#include "offload_bundle_parser.hpp"
 #include "platform.hpp"
 #include "program.hpp"
 #include "ur2offload.hpp"
@@ -70,85 +71,6 @@ ur_result_t ProgramCreateCudaWorkaround(ur_context_handle_t, const uint8_t *,
 }
 #endif
 
-// https://clang.llvm.org/docs/ClangOffloadBundler.html#bundled-binary-file-layout
-class HipOffloadBundleParser {
-  static constexpr std::string_view Magic = "__CLANG_OFFLOAD_BUNDLE__";
-  const uint8_t *Buff;
-  size_t Length;
-
-  struct __attribute__((packed)) BundleEntry {
-    uint64_t ObjectOffset;
-    uint64_t ObjectSize;
-    uint64_t EntryIdSize;
-    char EntryIdStart;
-  };
-
-  struct __attribute__((packed)) BundleHeader {
-    const char HeaderMagic[Magic.size()];
-    uint64_t EntryCount;
-    BundleEntry FirstEntry;
-  };
-
-  HipOffloadBundleParser() = delete;
-  HipOffloadBundleParser(const uint8_t *Buff, size_t Length)
-      : Buff(Buff), Length(Length) {}
-
-public:
-  static std::optional<HipOffloadBundleParser> load(const uint8_t *Buff,
-                                                    size_t Length) {
-    if (std::string_view{reinterpret_cast<const char *>(Buff), Length}.find(
-            Magic) != 0) {
-      return std::nullopt;
-    }
-    return HipOffloadBundleParser(Buff, Length);
-  }
-
-  ur_result_t extract(std::string_view SearchTargetId,
-                      const uint8_t *&OutBinary, size_t &OutLength) {
-    const char *Limit = reinterpret_cast<const char *>(&Buff[Length]);
-
-    // The different check here means that a binary consisting of only the magic
-    // bytes (but nothing else) will result in INVALID_PROGRAM rather than being
-    // treated as a non-bundle
-    auto *Header = reinterpret_cast<const BundleHeader *>(Buff);
-    if (reinterpret_cast<const char *>(&Header->FirstEntry) > Limit) {
-      return UR_RESULT_ERROR_INVALID_PROGRAM;
-    }
-
-    const auto *CurrentEntry = &Header->FirstEntry;
-    for (uint64_t I = 0; I < Header->EntryCount; I++) {
-      if (&CurrentEntry->EntryIdStart > Limit) {
-        return UR_RESULT_ERROR_INVALID_PROGRAM;
-      }
-      auto EntryId = std::string_view(&CurrentEntry->EntryIdStart,
-                                      CurrentEntry->EntryIdSize);
-      if (EntryId.end() > Limit) {
-        return UR_RESULT_ERROR_INVALID_PROGRAM;
-      }
-
-      // Will match either "hip" or "hipv4"
-      bool isHip = EntryId.find("hip") == 0;
-      bool VersionMatches =
-          EntryId.find_last_of(SearchTargetId) == EntryId.size() - 1;
-
-      if (isHip && VersionMatches) {
-        OutBinary = reinterpret_cast<const uint8_t *>(
-            &Buff[CurrentEntry->ObjectOffset]);
-        OutLength = CurrentEntry->ObjectSize;
-
-        if (reinterpret_cast<const char *>(&OutBinary[OutLength]) > Limit) {
-          return UR_RESULT_ERROR_INVALID_PROGRAM;
-        }
-        return UR_RESULT_SUCCESS;
-      }
-
-      CurrentEntry = reinterpret_cast<const BundleEntry *>(EntryId.end());
-    }
-
-    return UR_RESULT_ERROR_INVALID_PROGRAM;
-  }
-};
-
 } // namespace
 
 UR_APIEXPORT ur_result_t UR_APICALL urProgramCreateWithBinary(

unified-runtime/source/common/CMakeLists.txt

@@ -23,6 +23,8 @@ add_ur_library(ur_common STATIC
     ur_util.cpp
     ur_util.hpp
     latency_tracker.hpp
+    offload_bundle_parser.cpp
+    offload_bundle_parser.hpp
     $<$<PLATFORM_ID:Windows>:windows/ur_lib_loader.cpp>
     $<$<PLATFORM_ID:Linux,Darwin>:linux/ur_lib_loader.cpp>
 )

unified-runtime/source/common/offload_bundle_parser.cpp

@@ -0,0 +1,80 @@
+/*
+ *
+ * Copyright (C) 2025 Intel Corporation
+ *
+ * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM
+ * Exceptions. See LICENSE.TXT
+ *
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ */
+
+#include "offload_bundle_parser.hpp"
+#include <string>
+
+std::optional<HipOffloadBundleParser>
+HipOffloadBundleParser::load(const uint8_t *Buff, size_t Length) {
+  if (std::string_view{reinterpret_cast<const char *>(Buff), Length}.find(
+          Magic) != 0) {
+    return std::nullopt;
+  }
+  return HipOffloadBundleParser(Buff, Length);
+}
+
+ur_result_t HipOffloadBundleParser::extract(std::string_view SearchTargetId,
+                                            const uint8_t *&OutBinary,
+                                            size_t &OutLength) {
+  if (auto Entry = containsBundle(SearchTargetId)) {
+    OutBinary = &Buff[Entry->ObjectOffset];
+    OutLength = Entry->ObjectSize;
+
+    if (const uint8_t *Limit = &Buff[Length]; &OutBinary[OutLength] <= Limit) {
+      return UR_RESULT_SUCCESS;
+    }
+  }
+  return UR_RESULT_ERROR_INVALID_PROGRAM;
+}
+
+std::optional<HipOffloadBundleParser::BundleEntry>
+HipOffloadBundleParser::containsBundle(std::string_view SearchTargetId) {
+  const uint8_t *Limit = &Buff[Length];
+
+  // The different check here means that a binary consisting of only the magic
+  // bytes (but nothing else) will result in INVALID_PROGRAM rather than being
+  // treated as a non-bundle
+  auto *Header = reinterpret_cast<const BundleHeader *>(Buff);
+  if (reinterpret_cast<const uint8_t *>(&Header->FirstEntry) > Limit) {
+    return std::nullopt;
+  }
+
+  // std::string_view::ends_with is C++20. Until then, roll our own. Note then
+  // this is the equivalent form listed on en.cppreference.com.
+  auto ends_with = [](std::string_view str, std::string_view sv) {
+    return str.size() >= sv.size() &&
+           str.compare(str.size() - sv.size(), std::string::npos, sv) == 0;
+  };
+
+  const auto *CurrentEntry = &Header->FirstEntry;
+  for (uint64_t I = 0; I < Header->EntryCount; I++) {
+    const uint8_t *EntryBytes = &CurrentEntry->EntryIdStart;
+    if (EntryBytes > Limit ||
+        (EntryBytes + CurrentEntry->EntryIdSize) > Limit) {
+      return std::nullopt;
+    }
+    auto EntryId = std::string_view(reinterpret_cast<const char *>(EntryBytes),
+                                    CurrentEntry->EntryIdSize);
+
+    // Will match either "hip" or "hipv4"
+    bool isHip = EntryId.find("hip") == 0;
+    bool VersionMatches = ends_with(EntryId, SearchTargetId);
+
+    if (isHip && VersionMatches) {
+      return *CurrentEntry;
+    }
+
+    CurrentEntry = reinterpret_cast<const BundleEntry *>(
+        EntryBytes + CurrentEntry->EntryIdSize);
+  }
+
+  return std::nullopt;
+}

unified-runtime/source/common/offload_bundle_parser.hpp

@@ -0,0 +1,55 @@
+/*
+ *
+ * Copyright (C) 2025 Intel Corporation
+ *
+ * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM
+ * Exceptions. See LICENSE.TXT
+ *
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ */
+
+#include <ur_api.h>
+
+#include <cstdint>
+#include <optional>
+#include <string_view>
+
+#ifdef _MSC_VER
+#define PACKED(d) __pragma(pack(push, 1)) d __pragma(pack(pop))
+#else
+#define PACKED(d) d __attribute__((packed))
+#endif
+
+// https://clang.llvm.org/docs/ClangOffloadBundler.html#bundled-binary-file-layout
+class HipOffloadBundleParser {
+  static constexpr std::string_view Magic = "__CLANG_OFFLOAD_BUNDLE__";
+  const uint8_t *Buff;
+  size_t Length;
+
+  PACKED(struct BundleEntry {
+    uint64_t ObjectOffset;
+    uint64_t ObjectSize;
+    uint64_t EntryIdSize;
+    uint8_t EntryIdStart;
+  });
+
+  PACKED(struct BundleHeader {
+    const char HeaderMagic[Magic.size()];
+    uint64_t EntryCount;
+    BundleEntry FirstEntry;
+  });
+
+  HipOffloadBundleParser() = delete;
+  HipOffloadBundleParser(const uint8_t *Buff, size_t Length)
+      : Buff(Buff), Length(Length) {}
+
+public:
+  static std::optional<HipOffloadBundleParser> load(const uint8_t *Buff,
+                                                    size_t Length);
+
+  ur_result_t extract(std::string_view SearchTargetId,
+                      const uint8_t *&OutBinary, size_t &OutLength);
+
+  std::optional<BundleEntry> containsBundle(std::string_view SearchTargetId);
+};