#include "cuco_wrapper.h"
#include <cuco/dynamic_map.cuh>
#include <cuco/static_map.cuh>
#include <cuda_runtime.h>
#include <dlfcn.h>
#include <thrust/copy.h>
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/fill.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/sort.h>
#include <thrust/tuple.h>
#include <thrust/unique.h>
#include <cuda/std/functional>
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <limits>
#include <memory>
#include <type_traits>
#include <variant>
namespace {
// c10::cuda::CUDAStream (16-byte trivially copyable value wrapping c10::Stream)
struct RawTorchCUDAStream {
std::uint64_t lo;
std::uint64_t hi;
};
// Dynamically resolve LibTorch's (libc10_cuda.so) stream accessor functions at runtime.
// This avoids "DSO missing from command line" errors in environments where libtorch-ffi
// does not pass -lc10_cuda directly on the linker command line, as well as undefined
// symbol errors when linking against CPU-only LibTorch builds.
inline cudaStream_t get_current_torch_cuda_stream(int device_index = -1) {
using GetCurrentCUDAStreamFn = RawTorchCUDAStream (*)(std::int8_t);
using CUDAStreamStreamFn = cudaStream_t (*)(const RawTorchCUDAStream*);
static const auto get_current_fn = reinterpret_cast<GetCurrentCUDAStreamFn>(
dlsym(RTLD_DEFAULT, "_ZN3c104cuda20getCurrentCUDAStreamEa"));
static const auto stream_fn = reinterpret_cast<CUDAStreamStreamFn>(
dlsym(RTLD_DEFAULT, "_ZNK3c104cuda10CUDAStream6streamEv"));
if (get_current_fn != nullptr && stream_fn != nullptr) {
RawTorchCUDAStream raw = get_current_fn(static_cast<std::int8_t>(device_index));
return stream_fn(&raw);
}
return nullptr;
}
// RAII guard that switches to the target GPU device within the current scope (equivalent to c10::cuda::CUDAGuard)
struct ScopedCudaDevice {
int prev_device = -1;
bool switched = false;
explicit ScopedCudaDevice(int target_device) {
if (target_device >= 0) {
if (cudaGetDevice(&prev_device) == cudaSuccess && prev_device != target_device) {
cudaSetDevice(target_device);
switched = true;
}
}
}
~ScopedCudaDevice() {
if (switched && prev_device >= 0) {
cudaSetDevice(prev_device);
}
}
ScopedCudaDevice(const ScopedCudaDevice&) = delete;
ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete;
};
template <typename To, typename From>
inline To bit_cast_val(From src) noexcept {
static_assert(sizeof(To) == sizeof(From), "bit_cast_val requires matching sizes");
To dst;
std::memcpy(&dst, &src, sizeof(To));
return dst;
}
template <typename K>
K make_empty_key_sentinel(int key_dtype) noexcept {
if constexpr (sizeof(K) == 4) {
if (key_dtype == CUCO_DTYPE_FLOAT) {
return bit_cast_val<K>(std::numeric_limits<float>::quiet_NaN());
}
return std::numeric_limits<K>::min();
} else {
if (key_dtype == CUCO_DTYPE_DOUBLE) {
return bit_cast_val<K>(std::numeric_limits<double>::quiet_NaN());
}
return std::numeric_limits<K>::min();
}
}
template <typename V>
V make_empty_val_sentinel(int val_dtype) noexcept {
if constexpr (sizeof(V) == 4) {
if (val_dtype == CUCO_DTYPE_FLOAT) {
return bit_cast_val<V>(std::numeric_limits<float>::quiet_NaN());
}
return static_cast<V>(-1);
} else {
if (val_dtype == CUCO_DTYPE_DOUBLE) {
return bit_cast_val<V>(std::numeric_limits<double>::quiet_NaN());
}
return static_cast<V>(-1);
}
}
template <typename K, typename V>
struct TupleToCucoPair {
template <typename TupleLike>
__host__ __device__ cuco::pair<K, V> operator()(TupleLike const& t) const noexcept {
return cuco::pair<K, V>{
static_cast<K>(thrust::get<0>(t)),
static_cast<V>(thrust::get<1>(t))
};
}
};
template <typename K, typename V, typename KeyIt, typename ValIt>
auto make_cuco_pair_iterator(KeyIt k_it, ValIt v_it) {
return thrust::make_transform_iterator(
thrust::make_zip_iterator(thrust::make_tuple(k_it, v_it)),
TupleToCucoPair<K, V>{}
);
}
struct MapBase {
virtual ~MapBase() = default;
virtual int key_dtype() const noexcept = 0;
virtual int val_dtype() const noexcept = 0;
virtual size_t size(cudaStream_t stream) const = 0;
virtual void insert_batch(int num_elements, const void* d_keys, const void* d_values, cudaStream_t stream) = 0;
virtual void lookup_batch(int num_elements, const void* d_keys, void* d_values, cudaStream_t stream) const = 0;
virtual void insert_keys(int num_elements, const void* d_keys, int64_t* d_indices, cudaStream_t stream) = 0;
virtual void lookup_indices(int num_elements, const void* d_keys, int64_t* d_indices, cudaStream_t stream) const = 0;
virtual void finalize(cudaStream_t stream) = 0;
};
template <typename K, typename V>
class TypedCudaHashMap final : public MapBase {
using dynamic_map_t = cuco::dynamic_map<K, V>;
using static_map_t = cuco::static_map<K, V>;
int key_dtype_;
int val_dtype_;
K empty_key_;
V empty_val_;
std::variant<dynamic_map_t, static_map_t> map_storage_;
public:
TypedCudaHashMap(int k_dtype, int v_dtype, size_t initial_capacity, cudaStream_t stream)
: key_dtype_(k_dtype),
val_dtype_(v_dtype),
empty_key_(make_empty_key_sentinel<K>(k_dtype)),
empty_val_(make_empty_val_sentinel<V>(v_dtype)),
map_storage_(
std::in_place_type<dynamic_map_t>,
// cuCollections' dynamic_map requires an initial capacity of at least 100,000 due to its internal min_insert_size (1e4)
std::max<size_t>(initial_capacity, 100000),
cuco::empty_key<K>{empty_key_},
cuco::empty_value<V>{empty_val_},
cuda::std::equal_to<K>{},
cuco::linear_probing<4, cuco::default_hash_function<K>>{},
cuco::cuda_thread_scope<cuda::thread_scope_device>{},
cuco::storage<1>{},
cuco::cuda_allocator<cuco::pair<K, V>>{},
cuda::stream_ref{stream})
{}
int key_dtype() const noexcept override { return key_dtype_; }
int val_dtype() const noexcept override { return val_dtype_; }
size_t size(cudaStream_t stream) const override {
if (std::holds_alternative<dynamic_map_t>(map_storage_)) {
return std::get<dynamic_map_t>(map_storage_).size();
} else {
return std::get<static_map_t>(map_storage_).size(cuda::stream_ref{stream});
}
}
void insert_batch(int num_elements, const void* d_keys, const void* d_values, cudaStream_t stream) override {
if (num_elements <= 0) return;
const K* keys = static_cast<const K*>(d_keys);
const V* values = static_cast<const V*>(d_values);
auto pairs_begin = make_cuco_pair_iterator<K, V>(keys, values);
auto pairs_end = pairs_begin + num_elements;
// Overwrite with the new value if the key already exists (insert_or_assign)
std::visit([&](auto& m) {
m.insert_or_assign(pairs_begin, pairs_end, cuda::stream_ref{stream});
}, map_storage_);
}
void lookup_batch(int num_elements, const void* d_keys, void* d_values, cudaStream_t stream) const override {
if (num_elements <= 0) return;
const K* keys = static_cast<const K*>(d_keys);
V* values = static_cast<V*>(d_values);
std::visit([&](auto const& m) {
m.find_async(keys, keys + num_elements, values, cuda::stream_ref{stream});
}, map_storage_);
}
void insert_keys(int num_elements, const void* d_keys, int64_t* d_indices, cudaStream_t stream) override {
if (num_elements <= 0) return;
if constexpr (std::is_same_v<V, int64_t>) {
const K* keys = static_cast<const K*>(d_keys);
auto exec = thrust::cuda::par.on(stream);
// 1. Deduplicate keys within the input batch using sort + unique
thrust::device_vector<K> unique_keys(num_elements);
thrust::copy(exec, keys, keys + num_elements, unique_keys.begin());
thrust::sort(exec, unique_keys.begin(), unique_keys.end());
auto unique_end = thrust::unique(exec, unique_keys.begin(), unique_keys.end());
size_t num_unique = static_cast<size_t>(unique_end - unique_keys.begin());
// 2. Check in batch which unique keys are already present in the map
thrust::device_vector<bool> is_contained(num_unique);
std::visit([&](auto const& m) {
m.contains_async(unique_keys.begin(), unique_end, is_contained.begin(), cuda::stream_ref{stream});
}, map_storage_);
// 3. Extract only the unregistered keys
thrust::device_vector<K> new_keys(num_unique);
auto new_keys_end = thrust::copy_if(
exec,
unique_keys.begin(),
unique_end,
is_contained.begin(),
new_keys.begin(),
cuda::std::logical_not<bool>{}
);
size_t num_new = static_cast<size_t>(new_keys_end - new_keys.begin());
// 4. Assign sequential Int64 indices in [base_idx, base_idx + num_new) to the new keys and insert them
if (num_new > 0) {
int64_t base_idx = static_cast<int64_t>(size(stream));
auto idx_begin = thrust::make_counting_iterator<int64_t>(base_idx);
auto pairs_begin = make_cuco_pair_iterator<K, int64_t>(new_keys.begin(), idx_begin);
auto pairs_end = pairs_begin + num_new;
std::visit([&](auto& m) {
m.insert(pairs_begin, pairs_end, cuda::stream_ref{stream});
}, map_storage_);
}
// 5. Write the 0-based indices corresponding to all original input keys into d_indices
std::visit([&](auto const& m) {
m.find_async(keys, keys + num_elements, d_indices, cuda::stream_ref{stream});
}, map_storage_);
}
}
void lookup_indices(int num_elements, const void* d_keys, int64_t* d_indices, cudaStream_t stream) const override {
if (num_elements <= 0) return;
if constexpr (std::is_same_v<V, int64_t>) {
const K* keys = static_cast<const K*>(d_keys);
std::visit([&](auto const& m) {
m.find_async(keys, keys + num_elements, d_indices, cuda::stream_ref{stream});
}, map_storage_);
}
}
void finalize(cudaStream_t stream) override {
if (!std::holds_alternative<dynamic_map_t>(map_storage_)) {
return;
}
auto& d_map = std::get<dynamic_map_t>(map_storage_);
size_t current_size = d_map.size();
thrust::device_vector<K> all_keys(current_size);
thrust::device_vector<V> all_vals(current_size);
size_t actual_size = 0;
if (current_size > 0) {
auto [keys_end, vals_end] = d_map.retrieve_all(
all_keys.begin(), all_vals.begin(), cuda::stream_ref{stream});
actual_size = static_cast<size_t>(keys_end - all_keys.begin());
}
size_t optimal_capacity = std::max<size_t>(actual_size * 2, 1024);
// Construct static_map in-place while destroying the old dynamic_map
auto& s_map = map_storage_.template emplace<static_map_t>(
optimal_capacity,
cuco::empty_key<K>{empty_key_},
cuco::empty_value<V>{empty_val_},
cuda::std::equal_to<K>{},
cuco::linear_probing<4, cuco::default_hash_function<K>>{},
cuco::cuda_thread_scope<cuda::thread_scope_device>{},
cuco::storage<1>{},
cuco::cuda_allocator<cuco::pair<K, V>>{},
cuda::stream_ref{stream}
);
if (actual_size > 0) {
auto pairs_begin = make_cuco_pair_iterator<K, V>(all_keys.begin(), all_vals.begin());
auto pairs_end = pairs_begin + actual_size;
s_map.insert(pairs_begin, pairs_end, cuda::stream_ref{stream});
}
}
};
std::unique_ptr<MapBase> create_typed_map(int key_dtype, int val_dtype, size_t initial_capacity, cudaStream_t stream) {
bool key_is_64 = (key_dtype == CUCO_DTYPE_INT64 || key_dtype == CUCO_DTYPE_DOUBLE);
bool val_is_64 = (val_dtype == CUCO_DTYPE_INT64 || val_dtype == CUCO_DTYPE_DOUBLE);
if (!key_is_64 && !val_is_64) {
return std::make_unique<TypedCudaHashMap<int32_t, int32_t>>(key_dtype, val_dtype, initial_capacity, stream);
} else if (!key_is_64 && val_is_64) {
return std::make_unique<TypedCudaHashMap<int32_t, int64_t>>(key_dtype, val_dtype, initial_capacity, stream);
} else if (key_is_64 && !val_is_64) {
return std::make_unique<TypedCudaHashMap<int64_t, int32_t>>(key_dtype, val_dtype, initial_capacity, stream);
} else {
return std::make_unique<TypedCudaHashMap<int64_t, int64_t>>(key_dtype, val_dtype, initial_capacity, stream);
}
}
template <typename V>
void fill_uninitialized_lookup(int val_dtype, int num_elements, void* d_values, cudaStream_t stream) {
V sentinel = make_empty_val_sentinel<V>(val_dtype);
V* out = static_cast<V*>(d_values);
thrust::fill(thrust::cuda::par.on(stream), out, out + num_elements, sentinel);
}
} // namespace
struct CudaHashMapOpaque {
int device_index;
size_t initial_capacity;
std::unique_ptr<MapBase> impl;
explicit CudaHashMapOpaque(size_t cap = 100000, int dev = -1)
: device_index(dev), initial_capacity(cap), impl(nullptr) {}
MapBase* ensure_initialized(int key_dtype, int val_dtype, cudaStream_t stream) {
if (!impl) {
if (device_index < 0) {
int cur_dev = 0;
if (cudaGetDevice(&cur_dev) == cudaSuccess) {
device_index = cur_dev;
} else {
device_index = 0;
}
}
impl = create_typed_map(key_dtype, val_dtype, initial_capacity, stream);
}
return impl.get();
}
};
extern "C" {
CudaHashMapOpaque* custom_cuco_create(void) {
return new CudaHashMapOpaque(100000, -1);
}
CudaHashMapOpaque* custom_cuco_create_with_capacity(size_t initial_capacity) {
return new CudaHashMapOpaque(initial_capacity, -1);
}
CudaHashMapOpaque* custom_cuco_create_on_device(int device_index, size_t initial_capacity) {
return new CudaHashMapOpaque(initial_capacity, device_index);
}
void custom_cuco_destroy(CudaHashMapOpaque* map_ptr) {
if (!map_ptr) return;
if (map_ptr->impl && map_ptr->device_index >= 0) {
ScopedCudaDevice guard(map_ptr->device_index);
map_ptr->impl.reset();
}
delete map_ptr;
}
int custom_cuco_device_index(CudaHashMapOpaque* map_ptr) {
if (!map_ptr) return -1;
return map_ptr->device_index;
}
void custom_cuco_set_device_index(CudaHashMapOpaque* map_ptr, int device_index) {
if (!map_ptr) return;
if (!map_ptr->impl) {
map_ptr->device_index = device_index;
}
}
size_t custom_cuco_size(CudaHashMapOpaque* map_ptr) {
if (!map_ptr || !map_ptr->impl) return 0;
ScopedCudaDevice guard(map_ptr->device_index);
cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
return map_ptr->impl->size(stream);
}
int custom_cuco_key_dtype(CudaHashMapOpaque* map_ptr) {
if (!map_ptr || !map_ptr->impl) return -1;
return map_ptr->impl->key_dtype();
}
int custom_cuco_val_dtype(CudaHashMapOpaque* map_ptr) {
if (!map_ptr || !map_ptr->impl) return -1;
return map_ptr->impl->val_dtype();
}
void custom_cuco_insert_keys(
CudaHashMapOpaque* map_ptr,
int key_dtype,
int num_elements,
const void* d_keys,
int64_t* d_indices)
{
if (!map_ptr || num_elements <= 0) return;
ScopedCudaDevice guard(map_ptr->device_index);
cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
MapBase* m = map_ptr->ensure_initialized(key_dtype, CUCO_DTYPE_INT64, stream);
m->insert_keys(num_elements, d_keys, d_indices, stream);
}
void custom_cuco_lookup_indices(
CudaHashMapOpaque* map_ptr,
int key_dtype,
int num_elements,
const void* d_keys,
int64_t* d_indices)
{
if (!map_ptr || num_elements <= 0) return;
ScopedCudaDevice guard(map_ptr->device_index);
cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
if (!map_ptr->impl) {
fill_uninitialized_lookup<int64_t>(CUCO_DTYPE_INT64, num_elements, d_indices, stream);
return;
}
map_ptr->impl->lookup_indices(num_elements, d_keys, d_indices, stream);
}
void custom_cuco_insert_batch_typed(
CudaHashMapOpaque* map_ptr,
int key_dtype,
int val_dtype,
int num_elements,
const void* d_keys,
const void* d_values)
{
if (!map_ptr || num_elements <= 0) return;
ScopedCudaDevice guard(map_ptr->device_index);
cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
MapBase* m = map_ptr->ensure_initialized(key_dtype, val_dtype, stream);
m->insert_batch(num_elements, d_keys, d_values, stream);
}
void custom_cuco_lookup_batch_typed(
CudaHashMapOpaque* map_ptr,
int key_dtype,
int val_dtype,
int num_elements,
const void* d_keys,
void* d_values)
{
if (!map_ptr || num_elements <= 0) return;
ScopedCudaDevice guard(map_ptr->device_index);
cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
if (!map_ptr->impl) {
if (val_dtype == CUCO_DTYPE_INT64 || val_dtype == CUCO_DTYPE_DOUBLE) {
fill_uninitialized_lookup<int64_t>(val_dtype, num_elements, d_values, stream);
} else {
fill_uninitialized_lookup<int32_t>(val_dtype, num_elements, d_values, stream);
}
return;
}
(void)key_dtype;
map_ptr->impl->lookup_batch(num_elements, d_keys, d_values, stream);
}
void custom_cuco_insert_batch(
CudaHashMapOpaque* map_ptr,
int num_elements,
void* d_keys,
void* d_values)
{
if (!map_ptr || num_elements <= 0) return;
int k_dtype = map_ptr->impl ? map_ptr->impl->key_dtype() : CUCO_DTYPE_INT32;
int v_dtype = map_ptr->impl ? map_ptr->impl->val_dtype() : CUCO_DTYPE_INT32;
custom_cuco_insert_batch_typed(map_ptr, k_dtype, v_dtype, num_elements, d_keys, d_values);
}
void custom_cuco_lookup_batch(
CudaHashMapOpaque* map_ptr,
int num_elements,
void* d_keys,
void* d_values)
{
if (!map_ptr || num_elements <= 0) return;
int k_dtype = map_ptr->impl ? map_ptr->impl->key_dtype() : CUCO_DTYPE_INT32;
int v_dtype = map_ptr->impl ? map_ptr->impl->val_dtype() : CUCO_DTYPE_INT32;
custom_cuco_lookup_batch_typed(map_ptr, k_dtype, v_dtype, num_elements, d_keys, d_values);
}
void custom_cuco_finalize(CudaHashMapOpaque* map_ptr) {
if (!map_ptr || !map_ptr->impl) return;
ScopedCudaDevice guard(map_ptr->device_index);
cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
map_ptr->impl->finalize(stream);
}
} // extern "C"