Upstream cuVS

AlgorithmLauncher.cpp

@@ -0,0 +1,65 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <AlgorithmLauncher.hpp>
+#include <raft/util/cuda_rt_essentials.hpp>
+
+AlgorithmLauncher::AlgorithmLauncher(cudaKernel_t k, cudaLibrary_t lib) : kernel{k}, library{lib} {}
+
+AlgorithmLauncher::~AlgorithmLauncher()
+{
+  if (library != nullptr) { (void)cudaLibraryUnload(library); }
+}
+
+AlgorithmLauncher::AlgorithmLauncher(AlgorithmLauncher&& other) noexcept
+  : kernel{other.kernel}, library{other.library}
+{
+  other.kernel  = nullptr;
+  other.library = nullptr;
+}
+
+AlgorithmLauncher& AlgorithmLauncher::operator=(AlgorithmLauncher&& other) noexcept
+{
+  if (this != &other) {
+    if (library != nullptr) { cudaLibraryUnload(library); }
+    kernel        = other.kernel;
+    library       = other.library;
+    other.kernel  = nullptr;
+    other.library = nullptr;
+  }
+  return *this;
+}
+
+void AlgorithmLauncher::call(
+  cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** kernel_args)
+{
+  cudaLaunchConfig_t config{};
+  config.gridDim          = grid;
+  config.blockDim         = block;
+  config.stream           = stream;
+  config.dynamicSmemBytes = shared_mem;
+  config.numAttrs         = 0;
+  config.attrs            = NULL;
+
+  RAFT_CUDA_TRY(cudaLaunchKernelExC(&config, kernel, kernel_args));
+}
+
+void AlgorithmLauncher::call_cooperative(
+  cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** kernel_args)
+{
+  cudaLaunchAttribute attribute[1];
+  attribute[0].id              = cudaLaunchAttributeCooperative;
+  attribute[0].val.cooperative = 1;
+
+  cudaLaunchConfig_t config{};
+  config.gridDim          = grid;
+  config.blockDim         = block;
+  config.stream           = stream;
+  config.dynamicSmemBytes = shared_mem;
+  config.numAttrs         = 1;
+  config.attrs            = attribute;
+
+  RAFT_CUDA_TRY(cudaLaunchKernelExC(&config, kernel, kernel_args));
+}

AlgorithmLauncher.hpp

@@ -0,0 +1,61 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+#include <driver_types.h>
+#include <vector_types.h>
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <unordered_map>
+
+struct AlgorithmLauncher {
+  AlgorithmLauncher() : kernel{nullptr}, library{nullptr} {}
+
+  AlgorithmLauncher(cudaKernel_t k, cudaLibrary_t lib);
+
+  ~AlgorithmLauncher();
+
+  AlgorithmLauncher(AlgorithmLauncher const&)            = delete;
+  AlgorithmLauncher& operator=(AlgorithmLauncher const&) = delete;
+
+  AlgorithmLauncher(AlgorithmLauncher&& other) noexcept;
+  AlgorithmLauncher& operator=(AlgorithmLauncher&& other) noexcept;
+
+  template <typename FuncT, typename... Args>
+  void dispatch(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, Args&&... args)
+  {
+    static_assert(std::is_same_v<FuncT, void(std::remove_reference_t<Args>...)>,
+                  "dispatch() argument types do not match the kernel function signature FuncT");
+
+    void* kernel_args[] = {const_cast<void*>(static_cast<void const*>(&args))...};
+    this->call(stream, grid, block, shared_mem, kernel_args);
+  }
+
+  template <typename FuncT, typename... Args>
+  void dispatch_cooperative(
+    cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, Args&&... args)
+  {
+    static_assert(
+      std::is_same_v<FuncT, void(std::remove_reference_t<Args>...)>,
+      "dispatch_cooperative() argument types do not match the kernel function signature FuncT");
+
+    void* kernel_args[] = {const_cast<void*>(static_cast<void const*>(&args))...};
+    this->call_cooperative(stream, grid, block, shared_mem, kernel_args);
+  }
+
+  cudaKernel_t get_kernel() { return this->kernel; }
+
+ private:
+  void call(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args);
+  void call_cooperative(
+    cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args);
+  cudaKernel_t kernel;
+  cudaLibrary_t library;
+};

AlgorithmPlanner.cpp

@@ -0,0 +1,115 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "cuda_runtime.h"
+#include "nvJitLink.h"
+
+#include <AlgorithmPlanner.hpp>
+#include <nvjitlink_checker.hpp>
+#include <raft/core/logger.hpp>
+#include <raft/util/cuda_rt_essentials.hpp>
+
+#include <chrono>
+#include <iterator>
+#include <memory>
+#include <mutex>
+#include <new>
+#include <shared_mutex>
+#include <string>
+#include <vector>
+
+std::string AlgorithmPlanner::get_fragments_key() const
+{
+  std::string key = "";
+  for (auto const& fragment : this->fragments) {
+    key += fragment->get_key();
+  }
+  return key;
+}
+
+std::shared_ptr<AlgorithmLauncher> AlgorithmPlanner::read_cache(std::string const& launch_key) const
+{
+  auto& launchers = jit_cache_.launchers;
+  std::shared_lock<std::shared_mutex> read_lock(jit_cache_.mutex);
+  if (auto it = launchers.find(launch_key); it != launchers.end()) { return it->second; }
+  return nullptr;
+}
+
+std::shared_ptr<AlgorithmLauncher> AlgorithmPlanner::get_launcher()
+{
+  auto& launchers = jit_cache_.launchers;
+  auto launch_key = this->get_fragments_key();
+
+  if (auto hit = read_cache(launch_key)) { return hit; }
+
+  std::unique_lock<std::shared_mutex> write_lock(jit_cache_.mutex);
+  if (auto it = launchers.find(launch_key); it != launchers.end()) { return it->second; }
+
+  std::string log_message =
+    "JIT compiling launcher for kernel: " + this->entrypoint + " and device functions: ";
+  for (auto const& fragment : this->fragments) {
+    log_message += std::string{fragment->get_key()} + ",";
+  }
+  log_message.pop_back();
+  RAFT_LOG_DEBUG("%s", log_message.c_str());
+  auto launcher         = this->build();
+  launchers[launch_key] = launcher;
+  return launcher;
+}
+
+std::shared_ptr<AlgorithmLauncher> AlgorithmPlanner::build()
+{
+  int device = 0;
+  int major  = 0;
+  int minor  = 0;
+  RAFT_CUDA_TRY(cudaGetDevice(&device));
+  RAFT_CUDA_TRY(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device));
+  RAFT_CUDA_TRY(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device));
+
+  std::string archs = "-arch=sm_" + std::to_string((major * 10 + minor));
+
+  // Load the generated LTO IR and link them together
+  nvJitLinkHandle handle;
+  std::vector<char const*> lopts;
+  lopts.reserve(2 + linktime_extra_options.size());
+  lopts.push_back("-lto");
+  lopts.push_back(archs.c_str());
+  for (auto const& opt : linktime_extra_options) {
+    lopts.push_back(opt.c_str());
+  }
+  auto result = nvJitLinkCreate(&handle, static_cast<unsigned int>(lopts.size()), lopts.data());
+  check_nvjitlink_result(handle, result);
+
+  for (auto const& frag : this->fragments) {
+    frag->add_to(handle);
+  }
+
+  // Call to nvJitLinkComplete causes linker to link together all the LTO-IR
+  // modules perform any optimizations and generate cubin from it.
+  result = nvJitLinkComplete(handle);
+  check_nvjitlink_result(handle, result);
+
+  // get cubin from nvJitLink
+  size_t cubin_size;
+  result = nvJitLinkGetLinkedCubinSize(handle, &cubin_size);
+  check_nvjitlink_result(handle, result);
+
+  std::unique_ptr<char[]> cubin{new char[cubin_size]};
+  result = nvJitLinkGetLinkedCubin(handle, cubin.get());
+  check_nvjitlink_result(handle, result);
+
+  result = nvJitLinkDestroy(&handle);
+  RAFT_EXPECTS(result == NVJITLINK_SUCCESS, "nvJitLinkDestroy failed");
+
+  // cubin is linked, so now load it
+  cudaLibrary_t library;
+  RAFT_CUDA_TRY(
+    cudaLibraryLoadData(&library, cubin.get(), nullptr, nullptr, 0, nullptr, nullptr, 0));
+
+  cudaKernel_t kernel;
+  RAFT_CUDA_TRY(cudaLibraryGetKernel(&kernel, library, this->entrypoint.c_str()));
+
+  return std::make_shared<AlgorithmLauncher>(kernel, library);
+}

AlgorithmPlanner.hpp

@@ -0,0 +1,60 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include "AlgorithmLauncher.hpp"
+#include "FragmentEntry.hpp"
+
+#include <memory>
+#include <shared_mutex>
+#include <string>
+#include <type_traits>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+struct LauncherJitCache {
+  std::shared_mutex mutex;
+  std::unordered_map<std::string, std::shared_ptr<AlgorithmLauncher>> launchers;
+};
+
+struct AlgorithmPlanner {
+  AlgorithmPlanner(std::string entrypoint, LauncherJitCache& jit_cache)
+    : entrypoint(std::move(entrypoint)), jit_cache_(jit_cache)
+  {
+  }
+
+  std::shared_ptr<AlgorithmLauncher> get_launcher();
+
+  std::string entrypoint;
+  std::vector<std::unique_ptr<FragmentEntry>> fragments;
+
+  template <typename T, typename = std::enable_if_t<std::is_convertible_v<T*, FragmentEntry*>>>
+  void add_fragment(std::unique_ptr<T> fragment)
+  {
+    fragments.push_back(std::unique_ptr<FragmentEntry>(std::move(fragment)));
+  }
+
+  template <typename FragmentTag>
+  void add_static_fragment()
+  {
+    add_fragment(std::make_unique<StaticFatbinFragmentEntry<FragmentTag>>());
+  }
+
+ protected:
+  /** Extra link-time option strings passed to nvJitLink. Base build()
+   *  always passes "-lto" and "-arch=sm_XX" first; derived planners may append here in their
+   *  constructor body. */
+  std::vector<std::string> linktime_extra_options;
+
+ private:
+  std::string get_fragments_key() const;
+  std::shared_ptr<AlgorithmLauncher> build();
+
+  std::shared_ptr<AlgorithmLauncher> read_cache(std::string const& launch_key) const;
+
+  LauncherJitCache& jit_cache_;
+};