SWDEV-546287 - Implement hipLibrary load/unload (#975)

projects/clr/hipamd/src/CMakeLists.txt

@@ -127,7 +127,8 @@ target_sources(amdhip64 PRIVATE
   hip_api_trace.cpp
   hip_table_interface.cpp
   hip_table_interface_c.cpp
-  hip_comgr_helper.cpp)
+  hip_comgr_helper.cpp
+  hip_library.cpp)
 
 if(WIN32)
   target_sources(amdhip64 PRIVATE hip_runtime.cpp)

projects/clr/hipamd/src/amdhip.def

@@ -511,3 +511,8 @@ hipGetDriverEntryPoint_spt
 hipMemPrefetchAsync_v2
 hipMemAdvise_v2
 hipStreamGetId
+hipLibraryLoadData
+hipLibraryLoadFromFile
+hipLibraryUnload
+hipLibraryGetKernel
+hipLibraryGetKernelCount

projects/clr/hipamd/src/hip_api_trace.cpp

@@ -863,6 +863,17 @@ hipError_t hipMemcpy3DBatchAsync(size_t numOps, struct hipMemcpy3DBatchOp* opLis
                                  unsigned long long flags, hipStream_t stream);
 hipError_t hipMemcpy3DPeer(hipMemcpy3DPeerParms* p);
 hipError_t hipMemcpy3DPeerAsync(hipMemcpy3DPeerParms* p, hipStream_t stream);
+hipError_t hipLibraryLoadData(hipLibrary_t* library, const void* code, hipJitOption** jitOptions,
+                              void** jitOptionsValues, unsigned int numJitOptions,
+                              hipLibraryOption** libraryOptions, void** libraryOptionValues,
+                              unsigned int numLibraryOptions);
+hipError_t hipLibraryLoadFromFile(hipLibrary_t* library, const char* fileName,
+                                  hipJitOption** jitOptions, void** jitOptionsValues,
+                                  unsigned int numJitOptions, hipLibraryOption** libraryOptions,
+                                  void** libraryOptionValues, unsigned int numLibraryOptions);
+hipError_t hipLibraryUnload(hipLibrary_t library);
+hipError_t hipLibraryGetKernel(hipKernel_t* pKernel, hipLibrary_t library, const char* name);
+hipError_t hipLibraryGetKernelCount(unsigned int* count, hipLibrary_t library);
 }  // namespace hip
 
 namespace hip {
@@ -1398,6 +1409,11 @@ void UpdateDispatchTable(HipDispatchTable* ptrDispatchTable) {
   ptrDispatchTable->hipMemcpy3DBatchAsync_fn = hip::hipMemcpy3DBatchAsync;
   ptrDispatchTable->hipMemcpy3DPeer_fn = hip::hipMemcpy3DPeer;
   ptrDispatchTable->hipMemcpy3DPeerAsync_fn = hip::hipMemcpy3DPeerAsync;
+  ptrDispatchTable->hipLibraryLoadData_fn = hip::hipLibraryLoadData;
+  ptrDispatchTable->hipLibraryLoadFromFile_fn = hip::hipLibraryLoadFromFile;
+  ptrDispatchTable->hipLibraryUnload_fn = hip::hipLibraryUnload;
+  ptrDispatchTable->hipLibraryGetKernel_fn = hip::hipLibraryGetKernel;
+  ptrDispatchTable->hipLibraryGetKernelCount_fn = hip::hipLibraryGetKernelCount;
 }
 
 #if HIP_ROCPROFILER_REGISTER > 0
@@ -2062,15 +2078,21 @@ HIP_ENFORCE_ABI(HipDispatchTable, hipGetDriverEntryPoint_spt_fn, 492);
 HIP_ENFORCE_ABI(HipDispatchTable, hipMemPrefetchAsync_v2_fn, 493);
 HIP_ENFORCE_ABI(HipDispatchTable, hipMemAdvise_v2_fn, 494);
 HIP_ENFORCE_ABI(HipDispatchTable, hipStreamGetId_fn, 495);
+// HIP_RUNTIME_API_TABLE_STEP_VERSION == 15
+HIP_ENFORCE_ABI(HipDispatchTable, hipLibraryLoadData_fn, 496);
+HIP_ENFORCE_ABI(HipDispatchTable, hipLibraryLoadFromFile_fn, 497);
+HIP_ENFORCE_ABI(HipDispatchTable, hipLibraryUnload_fn, 498);
+HIP_ENFORCE_ABI(HipDispatchTable, hipLibraryGetKernel_fn, 499);
+HIP_ENFORCE_ABI(HipDispatchTable, hipLibraryGetKernelCount_fn, 500);
 // if HIP_ENFORCE_ABI entries are added for each new function pointer in the table, the number below
 // will be +1 of the number in the last HIP_ENFORCE_ABI line. E.g.:
 //
 //  HIP_ENFORCE_ABI(<table>, <functor>, 8)
 //
 //  HIP_ENFORCE_ABI_VERSIONING(<table>, 9) <- 8 + 1 = 9
-HIP_ENFORCE_ABI_VERSIONING(HipDispatchTable, 496)
+HIP_ENFORCE_ABI_VERSIONING(HipDispatchTable, 501)
 
-static_assert(HIP_RUNTIME_API_TABLE_MAJOR_VERSION == 0 && HIP_RUNTIME_API_TABLE_STEP_VERSION == 14,
+static_assert(HIP_RUNTIME_API_TABLE_MAJOR_VERSION == 0 && HIP_RUNTIME_API_TABLE_STEP_VERSION == 15,
               "If you get this error, add new HIP_ENFORCE_ABI(...) code for the new function "
               "pointers and then update this check so it is true");
 #endif

projects/clr/hipamd/src/hip_hcc.map.in

@@ -629,4 +629,15 @@ global:
     hipStreamGetId;
 local:
     *;
-} hip_6.5;
+} hip_6.5;
+
+hip_7.2 {
+global:
+    hipLibraryLoadData;
+    hipLibraryLoadFromFile;
+    hipLibraryUnload;
+    hipLibraryGetKernel;
+    hipLibraryGetKernelCount;
+local:
+    *;
+} hip_7.1;

projects/clr/hipamd/src/hip_library.cpp

@@ -0,0 +1,180 @@
+/*
+Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANNTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER INN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR INN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+*/
+
+#include <filesystem>
+#include <mutex>
+#include <string>
+#include <vector>
+
+#include "hip/hip_runtime.h"
+#include "hip_library.hpp"
+#include "hip_platform.hpp"
+#include "utils/debug.hpp"
+
+namespace hip {
+void LibraryContainer::Register(std::string name, int device, hipKernel_t k) {
+  std::scoped_lock<std::mutex> lock(lib_mutex_);
+  auto key = std::make_pair(name, device);
+  if (kernels_.find(key) == kernels_.end()) {
+    kernels_.insert(std::make_pair(std::make_pair(name, device), k));
+    if (!hip::PlatformState::instance().RegisterLibraryFunction(k)) {
+      LogPrintfInfo("Already registered: %p", k);
+    }
+  }
+}
+
+hipError_t LibraryContainer::Kernel(hipKernel_t* k, std::string name) {
+  auto device_id = hip::ihipGetDevice();
+  if (auto ki = kernels_.find(std::make_pair(name, device_id)); ki != kernels_.end()) {
+    *k = ki->second;
+    return hipSuccess;
+  }
+  auto m = fatbin_->Module(device_id);
+  auto f = functions_.find(name);
+  if (f == functions_.end()) {
+    return hipErrorNotFound;
+  }
+  auto ret = f->second.get()->getDynFunc(reinterpret_cast<hipFunction_t*>(k), m);
+
+  // Register it, basically make it available for query though the hip context.
+  Register(name, device_id, *k);
+  return hipSuccess;
+}
+
+LibraryContainer::LibraryContainer(const char* code_object) {
+  fatbin_ = std::make_shared<hip::FatBinaryInfo>(nullptr, code_object);
+}
+
+LibraryContainer::LibraryContainer(const std::string file_name) {
+  fatbin_ = std::make_shared<hip::FatBinaryInfo>(file_name.c_str(), nullptr);
+}
+
+LibraryContainer::~LibraryContainer() {
+  for (const auto& k : kernels_) {
+    (void)hip::PlatformState::instance().UnregisterLibraryFunction(k.second);
+  }
+  kernels_.clear();
+}
+
+// BuildIt builds and loads the Library, default behavior is lazy load.
+// This function needs to be called before any query on library.
+hipError_t LibraryContainer::BuildIt() {
+  std::scoped_lock<std::mutex> lock(lib_mutex_);
+  if (built_) {
+    return hipSuccess;
+  }
+
+  if (!fatbin_) {
+    return hipErrorInvalidValue;
+  }
+
+  int device_id = ihipGetDevice();
+  std::vector<hip::Device*> devices = {g_devices[device_id]};
+  IHIP_RETURN_ONFAIL(fatbin_->ExtractFatBinaryUsingCOMGR(devices));
+  IHIP_RETURN_ONFAIL(fatbin_->BuildProgram(device_id));
+
+  auto program =
+      fatbin_->GetProgram(device_id)->getDeviceProgram(*hip::getCurrentDevice()->devices()[0]);
+
+  // Process Functions
+  std::vector<std::string> function_names;
+  program->getGlobalFuncFromCodeObj(&function_names);
+  for (auto& name : function_names) {
+    functions_.emplace(std::make_pair(name, std::make_shared<hip::Function>(name)));
+  }
+
+  built_ = true;
+  return hipSuccess;
+}
+
+hipError_t hipLibraryLoadData(hipLibrary_t* library, const void* image, hipJitOption** jitOptions,
+                              void** jitOptionsValues, unsigned int numJitOptions,
+                              hipLibraryOption** libraryOptions, void** libraryOptionValues,
+                              unsigned int numLibraryOptions) {
+  HIP_INIT_API(hipLibraryLoadData, library, image, jitOptions, jitOptionsValues, numJitOptions,
+               libraryOptions, libraryOptionValues, numLibraryOptions);
+  if (library == nullptr || image == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+
+  // We do not support JIT options
+  if (numJitOptions > 0) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+
+  auto* l = new hip::LibraryContainer((const char*)image);
+  *library = reinterpret_cast<hipLibrary_t>(l);
+  HIP_RETURN(hipSuccess);
+}
+
+hipError_t hipLibraryLoadFromFile(hipLibrary_t* library, const char* fname,
+                                  hipJitOption** jitOptions, void** jitOptionsValues,
+                                  unsigned int numJitOptions, hipLibraryOption** libraryOptions,
+                                  void** libraryOptionValues, unsigned int numLibraryOptions) {
+  HIP_INIT_API(hipLibraryLoadFromFile, library, fname, jitOptions, jitOptionsValues, numJitOptions,
+               libraryOptions, libraryOptionValues, numLibraryOptions);
+  if (library == nullptr || !std::filesystem::exists(fname) || numJitOptions > 0) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto* l = new hip::LibraryContainer(std::string(fname));
+  *library = reinterpret_cast<hipLibrary_t>(l);
+  HIP_RETURN(hipSuccess);
+}
+
+hipError_t hipLibraryUnload(hipLibrary_t library) {
+  HIP_INIT_API(hipLibraryUnload, library);
+  if (library == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto l = reinterpret_cast<hip::LibraryContainer*>(library);
+  delete l;
+  HIP_RETURN(hipSuccess);
+}
+
+hipError_t hipLibraryGetKernelCount(unsigned int* count, hipLibrary_t library) {
+  HIP_INIT_API(hipLibraryGetKernelCount, count, library);
+  if (library == nullptr || count == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto l = reinterpret_cast<hip::LibraryContainer*>(library);
+  auto ret = l->BuildIt();
+  if (ret != hipSuccess) {
+    HIP_RETURN(ret);
+  }
+  *count = static_cast<int>(l->KernelCount());
+  HIP_RETURN(hipSuccess);
+}
+
+hipError_t hipLibraryGetKernel(hipKernel_t* kernel, hipLibrary_t library, const char* kname) {
+  HIP_INIT_API(hipLibraryGetKernel, kernel, library, kname);
+  if (library == nullptr || kname == nullptr || kernel == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto l = reinterpret_cast<hip::LibraryContainer*>(library);
+  auto ret = l->BuildIt();
+  if (ret != hipSuccess) {
+    HIP_RETURN(ret);
+  }
+  ret = l->Kernel(kernel, kname);
+  HIP_RETURN(ret);
+}
+}  // namespace hip

projects/clr/hipamd/src/hip_library.hpp

@@ -0,0 +1,74 @@
+/*
+Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANNTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER INN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR INN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+*/
+
+#pragma once
+
+#include <cstdlib>
+#include <memory>
+#include <mutex>
+#include <string>
+
+#include <hip/hip_runtime.h>
+
+#include "hip_code_object.hpp"
+#include "hip_fatbin.hpp"
+
+namespace hip {
+// An abstract Library container
+class LibraryContainer {
+ public:
+  // Create from pointer
+  explicit LibraryContainer(const char* code_object);  // from pointer
+  // Create from file
+  explicit LibraryContainer(const std::string file_name);  // deep copy from file
+  ~LibraryContainer();
+
+  // Load and build the library
+  hipError_t BuildIt();
+
+  // Get the total Kernel count in Library
+  size_t KernelCount() const { return functions_.size(); }
+
+  // Get the Kernel from name
+  hipError_t Kernel(hipKernel_t* k, std::string name);
+
+  // Get Fatbin pointer
+  inline FatBinaryInfo* FatBin() { return fatbin_.get(); }
+
+  // Register the kernel function, make an entry in global state
+  void Register(std::string name, int device, hipKernel_t k);
+
+ private:
+  LibraryContainer() = delete;
+  LibraryContainer(const LibraryContainer&) = delete;
+  LibraryContainer(const LibraryContainer&&) = delete;
+  LibraryContainer& operator=(const LibraryContainer&) = delete;
+  LibraryContainer& operator=(const LibraryContainer&&) = delete;
+
+  std::mutex lib_mutex_;
+  std::atomic_bool built_ = false;
+  std::shared_ptr<FatBinaryInfo> fatbin_;
+  std::map<std::string, std::shared_ptr<hip::Function>> functions_;
+  // Store already looked up kernels for certain devices
+  std::map<std::pair<std::string /* name */, int /* device */>, hipKernel_t> kernels_;
+};
+}  // namespace hip

projects/clr/hipamd/src/hip_platform.cpp

@@ -647,15 +647,19 @@ hipError_t ihipLaunchKernel(const void* hostFunction, dim3 gridDim, dim3 blockDi
   if (!hip::isValid(stream)) {
     return hipErrorInvalidValue;
   }
+  if (hostFunction == nullptr) {
+    return hipErrorInvalidDeviceFunction;
+  }
+
   hipFunction_t func = nullptr;
   int deviceId = hip::Stream::DeviceId(stream);
-  hipError_t hip_error = PlatformState::instance().getStatFunc(&func, hostFunction, deviceId);
+
+  hipError_t hip_error =
+      PlatformState::instance().getStatFunc(&func, hostFunction, deviceId);
   if ((hip_error != hipSuccess) || (func == nullptr)) {
-    if (hip_error == hipErrorNoBinaryForGpu) {
-      return hip_error;
-    } else {
-      return hipErrorInvalidDeviceFunction;
-    }
+    // assume its hip function type if we did not get a valid output from static
+    // func lookup
+    func = reinterpret_cast<hipFunction_t>(const_cast<void *>(hostFunction));
   }
 
   constexpr auto gridDimYZmax = static_cast<uint64_t>(std::numeric_limits<uint16_t>::max()) + 1;

projects/clr/hipamd/src/hip_platform.hpp

@@ -113,6 +113,23 @@ class PlatformState {
 
   size_t UfdMapSize() const { return ufd_map_.size(); }
 
+  inline bool RegisterLibraryFunction(const hipKernel_t f) {
+    amd::ScopedLock lock(lock_);
+    if (library_functions_.find(f) == library_functions_.end()) {
+      library_functions_.insert(f);
+      return true;
+    }
+    return false;
+  }
+  inline bool UnregisterLibraryFunction(const hipKernel_t f) {
+    amd::ScopedLock lock(lock_);
+    if (library_functions_.find(f) != library_functions_.end()) {
+      library_functions_.erase(f);
+      return true;
+    }
+    return false;
+  }
+
  private:
   // Dynamic Code Object map, keyin module to get the corresponding object
   std::unordered_map<hipModule_t, hip::DynCO*> dynCO_map_;
@@ -123,5 +140,6 @@ class PlatformState {
   std::unordered_map<std::string, std::shared_ptr<UniqueFD>> ufd_map_;  //!< Unique File Desc Map
 
   void* dynamicLibraryHandle_{nullptr};
+  std::unordered_set<hipKernel_t> library_functions_;
 };
 }  // namespace hip

projects/clr/hipamd/src/hip_table_interface.cpp

@@ -2010,4 +2010,31 @@ hipError_t hipGraphExecExternalSemaphoresWaitNodeSetParams(
     const hipExternalSemaphoreWaitNodeParams* nodeParams) {
   return hip::GetHipDispatchTable()->hipGraphExecExternalSemaphoresWaitNodeSetParams_fn(
       hGraphExec, hNode, nodeParams);
+}
+hipError_t hipLibraryLoadData(hipLibrary_t* library, const void* code, hipJitOption** jitOptions,
+                              void** jitOptionsValues, unsigned int numJitOptions,
+                              hipLibraryOption** libraryOptions, void** libraryOptionValues,
+                              unsigned int numLibraryOptions) {
+  return hip::GetHipDispatchTable()->hipLibraryLoadData_fn(
+      library, code, jitOptions, jitOptionsValues, numJitOptions, libraryOptions,
+      libraryOptionValues, numLibraryOptions);
+}
+hipError_t hipLibraryLoadFromFile(hipLibrary_t* library, const char* fileName,
+                                  hipJitOption** jitOptions, void** jitOptionsValues,
+                                  unsigned int numJitOptions, hipLibraryOption** libraryOptions,
+                                  void** libraryOptionValues, unsigned int numLibraryOptions) {
+  return hip::GetHipDispatchTable()->hipLibraryLoadFromFile_fn(
+      library, fileName, jitOptions, jitOptionsValues, numJitOptions, libraryOptions,
+      libraryOptionValues, numLibraryOptions);
+}
+hipError_t hipLibraryUnload(hipLibrary_t library) {
+  return hip::GetHipDispatchTable()->hipLibraryUnload_fn(library);
+}
+hipError_t hipLibraryGetKernel(hipKernel_t* pKernel, hipLibrary_t library, const char* name)  {
+  return hip::GetHipDispatchTable()->hipLibraryGetKernel_fn(pKernel, library,
+                                                            name);
+}
+hipError_t hipLibraryGetKernelCount(unsigned int *count, hipLibrary_t library) {
+  return hip::GetHipDispatchTable()->hipLibraryGetKernelCount_fn(count,
+                                                                 library);
 }