hasktorch-CudaHashMap

hasktorch-CudaHashMap provides a HaskTorch interface to the GPU-parallel hash map implementations (cuco::dynamic_map and cuco::static_map) in NVIDIA cuCollections (cuco).

It supports high-throughput batch insertions and lookups directly on GPU-resident Tensors without round-tripping through host (CPU) memory.

Features

• Two-Stage Lifecycle (dynamic_map \(\rightarrow\) static_map):
  • Starts as a dynamically resizing GPU hash map (cuco::dynamic_map) for workloads where the number of unique keys is not known in advance (e.g., training or incremental vocabulary/graph construction).
  • Supports one-time compaction into a single flat cuco::static_map via finalizeCudaHashMap for maximum lookup throughput during inference or read-heavy phases.
• Typed Key-Value Batch Operations (insertBatch / lookupBatch):
  • Supports Int32, Int64, Float, and Double Tensors for both keys and values.
  • Uses overwrite (insert_or_assign) semantics when duplicate keys are inserted across or within batches.
• Key-to-Index Mapping for indexPut / indexSelect (insertKeys / lookupIndices):
  • Automatically assigns 0-based sequential Int64 indices to unique keys on the GPU.
  • Produces index Tensors that can be passed directly to HaskTorch's indexSelect and indexPut (or via the convenience wrappers hashMapIndexPut and hashMapIndexSelect).
• CUDA Device & Stream Awareness:
  • Automatically moves input tensors to the target CUDA device and executes kernels on LibTorch's current CUDA stream.

System & Toolchain Requirements

Building the C++/CUDA wrapper (cbits/cuco_wrapper.cu) and linking it with GHC and LibTorch is sensitive to the versions of cmake, nvcc, gcc, and g++.

Important OS & Compiler Note: Ubuntu 26.04 LTS (or newer) is recommended. Ubuntu 24.04 LTS is not sufficient out of the box because its standard package repositories only provide gcc/g++ up to version 14, which is not enough for building this project's toolchain and dependencies. Ubuntu 26.04 provides GCC/G++ 15 alongside newer CMake and CUDA packages.

Verified Environment (Ubuntu 26.04.1 LTS)

Environment Setup & Building

1. LibTorch & LD_LIBRARY_PATH Configuration

cabal.project enables the +cuda flag on libtorch-ffi:

packages: .

package libtorch-ffi
  flags: +cuda

Both cbits/CMakeLists.txt and setenv resolve the LibTorch path using the following environment variables (shown with their defaults):

• LIBTORCH_HOME: defaults to $XDG_CACHE_HOME/libtorch or $HOME/.cache/libtorch
• LIBTORCH_VERSION: defaults to 2.9.1
• LIBTORCH_CUDA_VERSION: defaults to cu128

Before running executables, tests, or benchmarks, source setenv to add LibTorch's lib directory to LD_LIBRARY_PATH:

source ./setenv

2. Building with Cabal

The custom Setup.hs hook automatically runs CMake during cabal configure / cabal build to fetch NVIDIA/cuCollections (via CPM) and compile libcuco_wrapper.a into dist-newstyle/cbits-build:

cabal build

(Optional) If you prefer to configure and build the C++/CUDA wrapper manually (e.g., to pass custom CMake variables), create cbits/build before running cabal build. When cbits/build/CMakeCache.txt exists, Setup.hs uses cbits/build instead of dist-newstyle/cbits-build:

cmake -S cbits -B cbits/build
cmake --build cbits/build -j4
cabal build

Note that cabal clean removes both dist-newstyle and cbits/build.

3. Running the Demo, Tests, and Benchmarks

# Run the example executable (app/Main.hs)
cabal run CudaHashMap

# Run the test suite (test/Main.hs)
cabal test

# Run the benchmark comparing Torch.CudaHashMap vs Data.HashMap.Strict (bench/Main.hs)
cabal bench

Usage Examples

Pattern A: Key-Value Batch Insert & Lookup

import Data.Int (Int64)
import Torch.Tensor (asTensor)
import Torch.CudaHashMap

main :: IO ()
main = do
  -- 1. Create a dynamic GPU hash map
  cudaMap <- newCudaHashMap

  let keysTensor      = asTensor ([10, 20, 30] :: [Int64])
      valuesTensor    = asTensor ([1.5, 2.5, 3.5] :: [Float])
      queryKeysTensor = asTensor ([20, 10, 99] :: [Int64])

  -- 2. Insert key-value batches (overwrites existing keys with new values)
  insertBatch cudaMap keysTensor valuesTensor

  -- 3. (Optional) Compact into a static_map once insertions are finished
  finalizeCudaHashMap cudaMap

  -- 4. Batch lookup (unregistered keys yield -1 for Int32/Int64, or NaN for Float/Double)
  resultTensor <- lookupBatch cudaMap queryKeysTensor
  print resultTensor
  -- Output: Tensor Float [3] [ 2.5000,  1.5000,  NaN]

Pattern B: Key-to-Index Mapping (indexPut / indexSelect)

When storing multi-dimensional feature rows or embedding vectors in a separate Tensor table, CudaHashMap can maintain a mapping from arbitrary keys to 0-based sequential Int64 row indices:

import Data.Int (Int64)
import Torch.Tensor (asTensor)
import Torch.TensorFactories (zeros')
import Torch.CudaHashMap

main :: IO ()
main = do
  idxMap <- newCudaHashMap
  let table0     = zeros' [1000]
      keysTensor = asTensor ([500, 600, 500, 700] :: [Int64])
      valsTensor = asTensor ([5.0, 6.0, 5.5, 7.0] :: [Float])

  -- Assigns 0-based Int64 indices to unique keys and writes valsTensor into table0
  table1 <- hashMapIndexPut idxMap False keysTensor valsTensor table0

  finalizeCudaHashMap idxMap

  -- Look up indices for registered keys and select from table1 along dimension 0
  selected <- hashMapIndexSelect idxMap 0 (asTensor ([700, 500, 600] :: [Int64])) table1
  print selected

You can also obtain the raw Int64 index tensors directly via insertKeys and lookupIndices (which returns -1 for missing keys).

API Summary (Torch.CudaHashMap)

Sentinel Value Conventions

cuCollections reserves sentinel values in open-addressing hash tables:

• Empty key sentinel: std::numeric_limits<K>::min() for Int32/Int64, and quiet NaN for Float/Double. Avoid inserting these reserved values as keys.
• Missing value sentinel (on lookup): -1 for Int32/Int64 values/indices, and quiet NaN for Float/Double values.

License

BSD-3-Clause (see LICENSE).