[FEA] RTCX: Replace Software SHA256 with SIMD-accelerated xxHash128 for JIT cache

cpp/librtcx/README.md

@@ -12,7 +12,8 @@ RTCX (runtime-compiler extended) is a wrapper around NVRTC and nvJitLink designe
 
 ## Build-time Requirements
 - CMake >= 4.0
-- LibZSTD - for binary compression
+- LibZSTD
+- xxHash
 - CUDA >= 12.2
 
 # Dependencies

cpp/librtcx/embed.hpp

@@ -4,11 +4,14 @@
  */
 
 #pragma once
-#include "sha256.hpp"
+#include "hash.hpp"
 
 #include <fcntl.h>
 #include <sys/mman.h>
 #include <unistd.h>
+
+#define XXH_INLINE_ALL
+#include <xxhash.h>
 #include <zstd.h>
 
 #include <cstdint>
@@ -99,17 +102,20 @@ std::vector<uint8_t> compress_bytes(std::span<uint8_t const> bytes, std::string_
   return compressed;
 }
 
-rtcx::sha256 compute_embed_hash(std::span<uint8_t const> uncompressed_files_bytes,
-                                std::span<uint8_t const> merged_dests_bytes,
-                                std::span<uint8_t const> merged_include_dirs_bytes,
-                                std::string_view compression)
+rtcx::hash128 compute_embed_hash(std::span<uint8_t const> uncompressed_files_bytes,
+                                 std::span<uint8_t const> merged_dests_bytes,
+                                 std::span<uint8_t const> merged_include_dirs_bytes,
+                                 std::string_view compression)
 {
-  rtcx::sha256_context ctx;
-  ctx.update(uncompressed_files_bytes);
-  ctx.update(merged_dests_bytes);
-  ctx.update(merged_include_dirs_bytes);
-  ctx.update(std::span{reinterpret_cast<const uint8_t*>(compression.data()), compression.size()});
-  return ctx.finalize();
+  XXH3_state_t* state = XXH3_createState();
+  XXH3_128bits_reset(state);
+  XXH3_128bits_update(state, uncompressed_files_bytes.data(), uncompressed_files_bytes.size());
+  XXH3_128bits_update(state, merged_dests_bytes.data(), merged_dests_bytes.size());
+  XXH3_128bits_update(state, merged_include_dirs_bytes.data(), merged_include_dirs_bytes.size());
+  XXH3_128bits_update(state, compression.data(), compression.size());
+  XXH128_hash_t hash = XXH3_128bits_digest(state);
+  XXH3_freeState(state);
+  return rtcx::hash128{hash.high64, hash.low64};
 }
 
 template <typename Container, typename Formatter>

cpp/librtcx/hash.hpp

@@ -0,0 +1,106 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <array>
+#include <bit>
+#include <charconv>
+#include <cstdint>
+#include <cstring>
+#include <format>
+#include <span>
+#include <stdexcept>
+#include <string_view>
+
+namespace rtcx {
+
+struct [[nodiscard]] hash128_hex_string {
+  static constexpr std::size_t NUM_HEX_DIGITS = 32;
+  static constexpr std::size_t NUM_HEX_BYTES  = NUM_HEX_DIGITS / 2;
+
+  char data_[NUM_HEX_DIGITS + 1];  // NOLINT(modernize-avoid-c-arrays)
+
+  [[nodiscard]] constexpr std::string_view view() const
+  {
+    return std::string_view{data_, NUM_HEX_DIGITS};
+  }
+
+  [[nodiscard]] constexpr operator std::string_view() const { return view(); }
+
+  [[nodiscard]] char const* data() const { return data_; }
+
+  [[nodiscard]] char const* c_str() const { return data_; }
+
+  [[nodiscard]] static constexpr std::size_t size() { return NUM_HEX_DIGITS; }
+
+  static hash128_hex_string make(std::span<std::uint8_t const, NUM_HEX_BYTES> input)
+  {
+    constexpr char const HEX_CHARS[] = "0123456789abcdef";  // NOLINT(modernize-avoid-c-arrays)
+    hash128_hex_string hex;
+    for (std::size_t i = 0; i < NUM_HEX_BYTES; ++i) {
+      hex.data_[i * 2]     = HEX_CHARS[(input[i] >> 4) & 0x0F];
+      hex.data_[i * 2 + 1] = HEX_CHARS[input[i] & 0x0F];
+    }
+    hex.data_[NUM_HEX_DIGITS] = '\0';
+    return hex;
+  }
+
+  static hash128_hex_string make(__uint128_t hash)
+  {
+    auto array = std::bit_cast<std::array<std::uint8_t, NUM_HEX_BYTES>>(hash);
+    return make(array);
+  }
+};
+
+struct hash128 {
+  __uint128_t value;
+
+  constexpr hash128(__uint128_t v = 0) : value(v) {}
+
+  constexpr hash128(std::uint64_t high, std::uint64_t low)
+    : value((static_cast<__uint128_t>(high) << 64) | low)
+  {
+  }
+
+  [[nodiscard]] constexpr bool operator==(hash128 const&) const = default;
+
+  [[nodiscard]] std::uint8_t operator[](std::size_t index) const
+  {
+    return reinterpret_cast<std::uint8_t const*>(&value)[16 - index];
+  }
+
+  [[nodiscard]] std::size_t size() const { return 16; }
+
+  [[nodiscard]] std::uint8_t const* data() const
+  {
+    return reinterpret_cast<std::uint8_t const*>(&value);
+  }
+
+  hash128_hex_string to_hex_string() const { return hash128_hex_string::make(value); }
+
+  static hash128 parse(std::string_view hex)
+  {
+    if (hex.size() != hash128_hex_string::NUM_HEX_DIGITS) {
+      throw std::invalid_argument(
+        std::format("Invalid hash128 hex string length, expected {} got {} (hash: `{}`)",
+                    hash128_hex_string::NUM_HEX_DIGITS,
+                    hex.size(),
+                    hex));
+    }
+    std::array<std::uint8_t, hash128_hex_string::NUM_HEX_BYTES> data{};
+    for (std::size_t i = 0; i < hash128_hex_string::NUM_HEX_BYTES; ++i) {
+      auto hex_byte  = hex.substr(i * 2, 2);
+      auto [ptr, ec] = std::from_chars(hex_byte.begin(), hex_byte.end(), data[i], 16);
+      if (ec != std::errc()) {
+        throw std::invalid_argument(
+          std::format("Invalid hex character {} in HEX string: `{}`", hex_byte, hex));
+      }
+    }
+    return hash128{std::bit_cast<__uint128_t>(data)};
+  }
+};
+
+}  // namespace rtcx

cpp/librtcx/rtcx.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -15,6 +15,9 @@
 #include <sys/mman.h>
 #include <sys/stat.h>
 #include <sys/types.h>
+
+#define XXH_INLINE_ALL
+#include <xxhash.h>
 #include <zstd.h>
 
 #include <atomic>
@@ -979,24 +982,26 @@ std::optional<blob_t> blob_t::from_file(char const* path)
 
 namespace {
 
-/// @brief retrieves a blob from disk based on the given sha256 hash and object type (e.g. "blob",
+/// @brief retrieves a blob from disk based on the given hash and object type (e.g. "blob",
 /// "cuLibrary"). Returns nullopt if the file doesn't exist on disk, and throws if any other error
 /// occurs.
-std::optional<blob> get_disk_blob(std::string const& cache_dir, object_type type, sha256 const& sha)
+std::optional<blob> get_disk_blob(std::string const& cache_dir,
+                                  object_type type,
+                                  hash128 const& hash)
 {
-  auto hex  = sha.to_hex_string();
+  auto hex  = hash.to_hex_string();
   auto path = std::format("{}/{}.{}.bin", cache_dir, hex.view(), object_tag(type));
   auto blob = blob_t::from_file(path.c_str());
 
   if (!blob.has_value()) { return std::nullopt; }
   return std::make_shared<blob_t>(std::move(*blob));
 }
 
-std::optional<library> get_disk_library(std::string const& cache_dir, sha256 const& sha)
+std::optional<library> get_disk_library(std::string const& cache_dir, hash128 const& hash)
 {
   RTCX_FUNC_RANGE();
 
-  auto hex  = sha.to_hex_string();
+  auto hex  = hash.to_hex_string();
   auto path = std::format("{}/{}.{}.bin", cache_dir, hex.view(), object_tag(object_type::LIBRARY));
 
   // WAR: avoid a driver API call when the cache file is not present
@@ -1033,7 +1038,7 @@ std::vector<std::string> get_disk_entries(std::string const& cache_dir)
 void cache_blob_to_disk(std::string const& cache_dir,
                         std::string const& tmp_dir,
                         object_type type,
-                        sha256 const& sha,
+                        hash128 const& hash,
                         std::span<std::uint8_t const> binary)
 {
   RTCX_FUNC_RANGE();
@@ -1061,7 +1066,7 @@ void cache_blob_to_disk(std::string const& cache_dir,
     }
   }
 
-  auto hex        = sha.to_hex_string();
+  auto hex        = hash.to_hex_string();
   auto final_path = std::format("{}/{}.{}.bin", cache_dir, hex.view(), object_tag(type));
 
   std::filesystem::create_directories(std::filesystem::path{final_path}.parent_path());
@@ -1084,7 +1089,7 @@ void cache_blob_to_disk(std::string const& cache_dir,
 
 }  // namespace
 
-std::shared_future<blob> cache_t::get_or_add_blob(sha256 const& sha, blob_compile_func compile)
+std::shared_future<blob> cache_t::get_or_add_blob(hash128 const& hash, blob_compile_func compile)
 {
   RTCX_FUNC_RANGE();
 
@@ -1094,7 +1099,7 @@ std::shared_future<blob> cache_t::get_or_add_blob(sha256 const& sha, blob_compil
   std::unique_lock lock{lock_};
 
   // check memory cache
-  if (auto it = enabled_ ? blobs_cache_.entries_.find(sha) : blobs_cache_.entries_.end();
+  if (auto it = enabled_ ? blobs_cache_.entries_.find(hash) : blobs_cache_.entries_.end();
       it != blobs_cache_.entries_.end()) {
     counter_.blob_mem_hits.incr();
 
@@ -1108,7 +1113,7 @@ std::shared_future<blob> cache_t::get_or_add_blob(sha256 const& sha, blob_compil
 
     // check disk cache
     std::optional<blob> disk_blob = std::nullopt;
-    if (enabled_) { disk_blob = get_disk_blob(cache_dir_, object_type::BLOB, sha); }
+    if (enabled_) { disk_blob = get_disk_blob(cache_dir_, object_type::BLOB, hash); }
 
     std::promise<blob> promise;
     auto fut       = promise.get_future().share();
@@ -1121,30 +1126,30 @@ std::shared_future<blob> cache_t::get_or_add_blob(sha256 const& sha, blob_compil
       promise.set_value(std::move(*disk_blob));
 
       // insert into cache
-      blobs_cache_.insert(sha, std::move(cache_fut), current_tick);
+      blobs_cache_.insert(hash, std::move(cache_fut), current_tick);
 
       return ret_fut;
 
     } else {
       counter_.blob_disk_misses.incr();
 
-      blobs_cache_.insert(sha, std::move(cache_fut), current_tick);
+      blobs_cache_.insert(hash, std::move(cache_fut), current_tick);
 
       // we can release the lock while calling the maker function since it may be expensive and we
-      // have already reserved a spot in the cache for this sha
+      // have already reserved a spot in the cache for this hash
       lock.unlock();
 
       auto result = compile();
       promise.set_value(result);
 
-      cache_blob_to_disk(cache_dir_, tmp_dir_, object_type::BLOB, sha, result->view());
+      cache_blob_to_disk(cache_dir_, tmp_dir_, object_type::BLOB, hash, result->view());
 
       return ret_fut;
     }
   }
 }
 
-std::shared_future<library> cache_t::get_or_add_library(sha256 const& sha,
+std::shared_future<library> cache_t::get_or_add_library(hash128 const& hash,
                                                         library_compile_func compile)
 {
   RTCX_FUNC_RANGE();
@@ -1155,7 +1160,7 @@ std::shared_future<library> cache_t::get_or_add_library(sha256 const& sha,
   std::unique_lock lock{lock_};
 
   // check memory cache
-  if (auto it = enabled_ ? libraries_cache_.entries_.find(sha) : libraries_cache_.entries_.end();
+  if (auto it = enabled_ ? libraries_cache_.entries_.find(hash) : libraries_cache_.entries_.end();
       it != libraries_cache_.entries_.end()) {
     counter_.library_mem_hits.incr();
 
@@ -1169,7 +1174,7 @@ std::shared_future<library> cache_t::get_or_add_library(sha256 const& sha,
 
     // check disk cache
     std::optional<library> disk_library = std::nullopt;
-    if (enabled_) { disk_library = get_disk_library(cache_dir_, sha); }
+    if (enabled_) { disk_library = get_disk_library(cache_dir_, hash); }
 
     std::promise<library> promise;
     auto fut       = promise.get_future().share();
@@ -1179,10 +1184,10 @@ std::shared_future<library> cache_t::get_or_add_library(sha256 const& sha,
     if (disk_library.has_value()) {
       counter_.library_disk_hits.incr();
 
-      libraries_cache_.insert(sha, std::move(cache_fut), current_tick);
+      libraries_cache_.insert(hash, std::move(cache_fut), current_tick);
 
       // we can release the lock while calling the maker function since it may be expensive and we
-      // have already reserved a spot in the cache for this sha
+      // have already reserved a spot in the cache for this hash
       lock.unlock();
 
       promise.set_value(std::move(*disk_library));
@@ -1192,17 +1197,17 @@ std::shared_future<library> cache_t::get_or_add_library(sha256 const& sha,
     } else {
       counter_.library_disk_misses.incr();
 
-      libraries_cache_.insert(sha, std::move(cache_fut), current_tick);
+      libraries_cache_.insert(hash, std::move(cache_fut), current_tick);
 
       // we can release the lock while calling the maker function since it may be expensive and we
-      // have already reserved a spot in the cache for this sha
+      // have already reserved a spot in the cache for this hash
       lock.unlock();
 
       auto [library, blob] = compile();
       promise.set_value(library);
 
       // store result to disk
-      cache_blob_to_disk(cache_dir_, tmp_dir_, object_type::LIBRARY, sha, blob->view());
+      cache_blob_to_disk(cache_dir_, tmp_dir_, object_type::LIBRARY, hash, blob->view());
 
       return ret_fut;
     }
@@ -1290,8 +1295,8 @@ void cache_t::preload_from_disk()
     for (auto const& path : entries) {
       try {
         auto file_name = std::filesystem::path{path}.filename().string();
-        auto sha_str   = file_name.substr(0, file_name.find('.'));
-        auto sha       = sha256::parse(sha_str);
+        auto hash_str  = file_name.substr(0, file_name.find('.'));
+        auto hash      = hash128::parse(hash_str);
 
         if (path.ends_with(".blob.bin")) {
           auto data = blob_t::from_file(path.c_str());
@@ -1300,14 +1305,14 @@ void cache_t::preload_from_disk()
           std::promise<rtcx::blob> promise;
           auto fut = promise.get_future().share();
           promise.set_value(std::move(blob));
-          blobs_cache_.insert(sha, std::move(fut), tick_);
+          blobs_cache_.insert(hash, std::move(fut), tick_);
         } else if (path.ends_with(".cuLibrary.bin")) {
-          auto lib = get_disk_library(cache_dir_, sha);
+          auto lib = get_disk_library(cache_dir_, hash);
           if (!lib.has_value()) { continue; }
           std::promise<library> promise;
           auto fut = promise.get_future().share();
           promise.set_value(std::move(*lib));
-          libraries_cache_.insert(sha, std::move(fut), tick_);
+          libraries_cache_.insert(hash, std::move(fut), tick_);
         }
       } catch (std::exception const& e) {
         // ignore any errors during preload