[llvm] [offload][AMDGPU] Add HSA AMD extension version checks (PR #228569)

[Robert Imschweiler via llvm-commits]

https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/228569

>From 8673f8d8827cb895b0b39cb2c3dcb992591b3f7b Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Sat, 3 Oct 2026 02:47:31 -0500
Subject: [PATCH] [offload][AMDGPU] Add dynamic HSA version and symbol checks

---
 .../amdgpu/dynamic_hsa/hsa.cpp                | 41 ++++++++++++++++++-
 .../plugins-nextgen/amdgpu/dynamic_hsa/hsa.h  |  2 +
 offload/plugins-nextgen/amdgpu/src/rtl.cpp    | 15 +++++++
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
index 13bd2fd064002b..724a5711a6be5e 100644
--- a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
+++ b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.cpp
@@ -11,6 +11,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/DynamicLibrary.h"
 
 #include "Shared/Debug.h"
@@ -19,6 +21,7 @@
 #include "hsa.h"
 #include "hsa_ext_amd.h"
 #include <memory>
+#include <utility>
 
 using namespace llvm::offload::debug;
 
@@ -96,14 +99,48 @@ DLWRAP_FINALIZE()
 #define DEBUG_PREFIX "Target " GETNAME(TARGET_NAME) " RTL"
 #endif
 
+/// Stand-in for an optional function that the HSA runtime does not provide.
+template <typename FuncTy> struct UnavailableTy;
+template <typename... ArgsTy>
+struct UnavailableTy<hsa_status_t (*)(ArgsTy...)> {
+  static hsa_status_t call(ArgsTy...) { return HSA_STATUS_ERROR; }
+};
+
+#define OPTIONAL_SYMBOL(SYMBOL)                                                \
+  {#SYMBOL, reinterpret_cast<void *>(&UnavailableTy<decltype(&SYMBOL)>::call)}
+
+/// The functions that are newer than version 1.0 of the AMD extension. If the
+/// runtime lacks one, calls to it fail instead of preventing the plugin from
+/// loading.
+static const std::pair<llvm::StringRef, void *> OptionalSymbols[] = {
+    OPTIONAL_SYMBOL(hsa_amd_vmem_address_reserve),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_address_free),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_handle_create),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_handle_release),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_map),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_unmap),
+    OPTIONAL_SYMBOL(hsa_amd_vmem_set_access),
+};
+
+#undef OPTIONAL_SYMBOL
+
 static bool resolveSymbols(llvm::sys::DynamicLibrary &Lib, const char *Name) {
   for (size_t I = 0; I < dlwrap::size(); I++) {
     const char *Sym = dlwrap::symbol(I);
 
     void *P = Lib.getAddressOfSymbol(Sym);
     if (P == nullptr) {
-      ODBG(OLDT_Init) << "Unable to find '" << Sym << "' in '" << Name << "'!";
-      return false;
+      const auto *It = llvm::find_if(
+          OptionalSymbols, [&](const auto &S) { return S.first == Sym; });
+      if (It == std::end(OptionalSymbols)) {
+        ODBG(OLDT_Init) << "Unable to find '" << Sym << "' in '" << Name
+                        << "'!";
+        return false;
+      }
+      ODBG(OLDT_Init) << "Optional '" << Sym << "' is unavailable in '" << Name
+                      << "'";
+      *dlwrap::pointer(I) = It->second;
+      continue;
     }
     ODBG(OLDT_Init) << "Implementing " << Sym << " with dlsym(" << Sym
                     << ") -> " << P;
diff --git a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h
index f66326a7f240e4..c303902ca3da7f 100644
--- a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h
+++ b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa.h
@@ -123,6 +123,8 @@ typedef enum {
   HSA_SYSTEM_INFO_VERSION_MINOR = 1,
   HSA_SYSTEM_INFO_TIMESTAMP = 2,
   HSA_SYSTEM_INFO_TIMESTAMP_FREQUENCY = 3,
+  HSA_AMD_SYSTEM_INFO_EXT_VERSION_MAJOR = 0x207,
+  HSA_AMD_SYSTEM_INFO_EXT_VERSION_MINOR = 0x208,
 } hsa_system_info_t;
 
 typedef enum {
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index de38c2a68e4ab4..5b52849af75da1 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -48,6 +48,7 @@
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Signals.h"
+#include "llvm/Support/VersionTuple.h"
 #include "llvm/Support/raw_ostream.h"
 
 #if !defined(__BYTE_ORDER__) || !defined(__ORDER_LITTLE_ENDIAN__) ||           \
@@ -232,6 +233,18 @@ static Error getTargetTripleAndFeatures(hsa_agent_t Agent,
   });
   return Err;
 }
+
+/// Get the version of the AMD extension that the HSA runtime implements, if the
+/// runtime reports it.
+static std::optional<llvm::VersionTuple> getAMDExtensionVersion() {
+  uint16_t Major, Minor;
+  if (hsa_system_get_info(HSA_AMD_SYSTEM_INFO_EXT_VERSION_MAJOR, &Major) !=
+          HSA_STATUS_SUCCESS ||
+      hsa_system_get_info(HSA_AMD_SYSTEM_INFO_EXT_VERSION_MINOR, &Minor) !=
+          HSA_STATUS_SUCCESS)
+    return std::nullopt;
+  return llvm::VersionTuple(Major, Minor);
+}
 } // namespace hsa_utils
 
 /// Utility class representing generic resource references to AMDGPU resources.
@@ -3241,6 +3254,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
       Info.add("HSA Runtime Version",
                std::to_string(Major) + "." + std::to_string(Minor), "",
                DeviceInfo::DRIVER_VERSION);
+    if (auto Version = hsa_utils::getAMDExtensionVersion())
+      Info.add("HSA AMD Extension Version", Version->getAsString());
 
     Info.add("HSA OpenMP Device Number", DeviceId);