hasktorch-CudaHashMap-0.1.0.0: cbits/cuco_wrapper.h

#pragma once

#include <stddef.h>
#include <stdint.h>

#ifdef __cplusplus
extern "C" {
#endif

// Opaque struct representing a C++ hash map instance
typedef struct CudaHashMapOpaque CudaHashMapOpaque;

// Supported tensor data types (DType)
typedef enum {
    CUCO_DTYPE_INT32  = 0,
    CUCO_DTYPE_INT64  = 1,
    CUCO_DTYPE_FLOAT  = 2,
    CUCO_DTYPE_DOUBLE = 3
} CucoDType;

// Hash map creation and destruction
// When device_index is -1, uses the device of the first tensor passed in or the current GPU device
CudaHashMapOpaque* custom_cuco_create(void);
CudaHashMapOpaque* custom_cuco_create_with_capacity(size_t initial_capacity);
CudaHashMapOpaque* custom_cuco_create_on_device(int device_index, size_t initial_capacity);
void custom_cuco_destroy(CudaHashMapOpaque* map_ptr);

// Get or set the CUDA device index associated with the map (-1 when unset)
int custom_cuco_device_index(CudaHashMapOpaque* map_ptr);
void custom_cuco_set_device_index(CudaHashMapOpaque* map_ptr, int device_index);

// Get the number of currently registered unique keys
size_t custom_cuco_size(CudaHashMapOpaque* map_ptr);

// Get the key and value DTypes configured for the map (-1 when uninitialized)
int custom_cuco_key_dtype(CudaHashMapOpaque* map_ptr);
int custom_cuco_val_dtype(CudaHashMapOpaque* map_ptr);

// Automatically assign 0-based sequential indices (Int64) to keys and write an Int64
// index tensor that can be used directly with indexSelect or indexPut
void custom_cuco_insert_keys(
    CudaHashMapOpaque* map_ptr,
    int key_dtype,
    int num_elements,
    const void* d_keys,
    int64_t* d_indices);

// Look up Int64 indices for keys without inserting new keys (unregistered keys get -1)
void custom_cuco_lookup_indices(
    CudaHashMapOpaque* map_ptr,
    int key_dtype,
    int num_elements,
    const void* d_keys,
    int64_t* d_indices);

// Typed batch insertion (existing keys are overwritten with new values: insert_or_assign)
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);

// Typed batch lookup (unregistered keys get -1 for Int32/Int64, or NaN for Float/Double)
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);

// Compatibility interface (operates with default Int32 or the already configured types)
void custom_cuco_insert_batch(
    CudaHashMapOpaque* map_ptr,
    int num_elements,
    void* d_keys,
    void* d_values);

void custom_cuco_lookup_batch(
    CudaHashMapOpaque* map_ptr,
    int num_elements,
    void* d_keys,
    void* d_values);

// Optimization function to call when entering a steady state (converts dynamic_map -> static_map)
void custom_cuco_finalize(CudaHashMapOpaque* map_ptr);

#ifdef __cplusplus
}
#endif