packages feed

hasktorch-CudaHashMap (empty) → 0.1.0.0

raw patch · 13 files changed

+1777/−0 lines, 13 filesdep +basedep +clockdep +deepseqbuild-type:Customsetup-changed

Dependencies added: base, clock, deepseq, hasktorch, hasktorch-CudaHashMap, libtorch-ffi, unordered-containers

Files

+ CHANGELOG.md view
@@ -0,0 +1,5 @@+# Revision history for CudaHashMap++## 0.1.0.0 -- 2026-11-10++* First version. Released on an unsuspecting world.
+ LICENSE view
@@ -0,0 +1,29 @@+Copyright (c) 2026, Susumu Katayama+++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++    * Redistributions of source code must retain the above copyright+      notice, this list of conditions and the following disclaimer.++    * Redistributions in binary form must reproduce the above+      copyright notice, this list of conditions and the following+      disclaimer in the documentation and/or other materials provided+      with the distribution.++    * Neither the name of the copyright holder nor the names of its+      contributors may be used to endorse or promote products derived+      from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ README.md view
@@ -0,0 +1,190 @@+# hasktorch-CudaHashMap++`hasktorch-CudaHashMap` provides a [HaskTorch](https://github.com/hasktorch/hasktorch) interface to the GPU-parallel hash map implementations (`cuco::dynamic_map` and `cuco::static_map`) in NVIDIA [cuCollections (`cuco`)](https://github.com/NVIDIA/cuCollections).++It supports high-throughput batch insertions and lookups directly on GPU-resident `Tensor`s 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` `Tensor`s 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 `Tensor`s 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)++| Component | Verified Version | Notes |+| :--- | :--- | :--- |+| **OS** | Ubuntu 26.04.1 LTS | Ubuntu 24.04 (max GCC 14) is insufficient. |+| **`gcc` / `g++`** | **15.2.0** (`>= 15`) | Used both for host C++ compilation and as GHC's linker driver (`-pgml g++`). |+| **`cmake`** | **4.4.5** | Minimum `3.23.1` required by `CMakeLists.txt`; a recent CMake is needed for CUDA architecture handling and CPM (`cuCollections` / CCCL). |+| **`nvcc` (CUDA)** | **12.4** (or newer, e.g. `12.8`) | Must support C++17 (`-std=c++17`) and extended device lambdas (`--expt-extended-lambda`). CUDA libraries (`libcuda`, `libcudart`, `libnvrtc`) are required at link time. |+| **GPU Architecture** | **Volta (`sm_70`) or newer** | `cuCollections` requires `sm_70+` (`70;75;80;86;89;90`) for scoped atomic operations. |+| **GHC / Cabal** | GHC 9.10.3, Cabal `>= 3.0` | Configured via `Setup.hs` (`build-type: Custom`). |+| **LibTorch** | **2.9.1** (CUDA build) | Installed via `libtorch-ffi` with the `+cuda` flag enabled (configured in `cabal.project`). |++---++## Environment Setup & Building++### 1. LibTorch & `LD_LIBRARY_PATH` Configuration++`cabal.project` enables the `+cuda` flag on `libtorch-ffi`:++```cabal+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`:++```bash+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`:++```bash+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`:++```bash+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++```bash+# 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++```haskell+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:++```haskell+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`)++| Function | Type | Description |+| :--- | :--- | :--- |+| `newCudaHashMap` | `IO CudaHashMap` | Create a new GPU hash map with default initial capacity (`100,000`). Device is inferred from the first tensor (or `cuda:0` for CPU tensors). |+| `newCudaHashMapWithCapacity` | `Int -> IO CudaHashMap` | Create a new GPU hash map with a specified minimum capacity. |+| `newCudaHashMapOnDevice` | `Device -> Int -> IO CudaHashMap` | Create a new GPU hash map pinned to a specific CUDA `Device` (e.g. `Device CUDA 1`). |+| `hashMapDevice` | `CudaHashMap -> IO Device` | Query the CUDA `Device` associated with the hash map. |+| `sizeCudaHashMap` | `CudaHashMap -> IO Int` | Return the number of unique keys currently stored in the map. |+| `finalizeCudaHashMap` | `CudaHashMap -> IO ()` | Convert the internal `cuco::dynamic_map` into a single `cuco::static_map` for faster lookups. |+| `insertBatch` | `CudaHashMap -> Tensor -> Tensor -> IO ()` | Insert or overwrite (`insert_or_assign`) key-value pairs (`Int32`, `Int64`, `Float`, `Double`). |+| `lookupBatch` | `CudaHashMap -> Tensor -> IO Tensor` | Look up values for a tensor of keys. Missing keys yield `-1` (`Int32`/`Int64`) or `NaN` (`Float`/`Double`). |+| `insertKeys` | `CudaHashMap -> Tensor -> IO Tensor` | Assign 0-based sequential `Int64` indices to unseen keys and return the `Int64` index tensor for all input keys. |+| `lookupIndices` | `CudaHashMap -> Tensor -> IO Tensor` | Look up 0-based `Int64` indices for input keys without inserting (`-1` for missing keys). |+| `hashMapIndexPut` | `CudaHashMap -> Bool -> Tensor -> Tensor -> Tensor -> IO Tensor` | Map keys via `insertKeys` and update a table tensor using `indexPut`. |+| `hashMapIndexSelect` | `CudaHashMap -> Int -> Tensor -> Tensor -> IO Tensor` | Map keys via `lookupIndices` and gather from a table tensor using `indexSelect`. |++### 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](LICENSE)).
+ Setup.hs view
@@ -0,0 +1,52 @@+{-# LANGUAGE CPP #-}++import Control.Monad (filterM, unless, when)+import Distribution.Simple+import Distribution.Simple.Setup (ConfigFlags (..))+import System.Directory+    ( doesDirectoryExist+    , doesFileExist+    , makeAbsolute+    , removePathForcibly+    )+import System.Process (callProcess)++#if MIN_VERSION_Cabal(3,14,0)+import Distribution.Utils.Path (makeSymbolicPath)+#else+makeSymbolicPath :: a -> a+makeSymbolicPath = id+#endif++main :: IO ()+main = defaultMainWithHooks simpleUserHooks+    { confHook = \(gpd, hbi) flags -> do+        -- Use cbits/build if it was created manually;+        -- otherwise use dist-newstyle/cbits-build, which is removed automatically by 'cabal clean'.+        manualBuildExists <- doesFileExist "cbits/build/CMakeCache.txt"+        let buildDir = if manualBuildExists+                         then "cbits/build"+                         else "dist-newstyle/cbits-build"++        cmakeCacheExists <- doesFileExist (buildDir ++ "/CMakeCache.txt")+        unless cmakeCacheExists $+            callProcess "cmake" ["-S", "cbits", "-B", buildDir]+        callProcess "cmake" ["--build", buildDir, "-j4"]++        absBuildDir <- makeAbsolute buildDir+        cudaDirs <- filterM doesDirectoryExist+            [ "/usr/local/cuda/lib64"+            , "/usr/lib/x86_64-linux-gnu"+            ]+        let updatedFlags = flags+                { configExtraLibDirs =+                    map makeSymbolicPath (absBuildDir : cudaDirs)+                    ++ configExtraLibDirs flags+                }+        confHook simpleUserHooks (gpd, hbi) updatedFlags++    , cleanHook = \pd z uh cf -> do+        exists <- doesDirectoryExist "cbits/build"+        when exists $ removePathForcibly "cbits/build"+        cleanHook simpleUserHooks pd z uh cf+    }
+ app/Main.hs view
@@ -0,0 +1,53 @@+module Main (main) where++import Control.Monad (forM_)+import Data.Int (Int64)+import Torch.Functional (indexPut)+import Torch.Tensor (Tensor, asTensor, device, indexSelect, toDevice)+import Torch.TensorFactories (zeros')+import Torch.CudaHashMap++main :: IO ()+main = do+  -- === Pattern A: Key-Value Hash Map (insertBatch / lookupBatch) ===+  -- 1. Initialize the hash map (internally a dynamic_map)+  --    To use a specific GPU, you can also call: newCudaHashMapOnDevice (Device CUDA 1) 100000+  cudaMap <- newCudaHashMap++  -- 2. Training loop (frequent insertions when total size is unpredictable)+  forM_ ([1..100] :: [Int]) $ \_epoch -> do+    -- Overwrites with the new value if the key already exists (insert_or_assign)+    insertBatch cudaMap keysTensor valuesTensor++  -- 3. Once insertions have stabilized, run the one-time optimization+  finalizeCudaHashMap cudaMap  -- Converts the internal representation into an ultra-fast single static_map++  -- 4. Inference and subsequent fixed lookups (runs at very high speed)+  resultTensor <- lookupBatch cudaMap queryKeysTensor+  print resultTensor++  -- === Pattern B: Obtain sequential indices from keys for indexPut / indexSelect ===+  idxMap <- newCudaHashMap+  let table0 = zeros' [1000]+  -- Automatically assign and retrieve 0-based Int64 indices for the keys+  idx <- insertKeys idxMap keysTensor+  -- Move valuesTensor and table0 to the same CUDA device as idx and write via indexPut+  -- (alternatively, you can use: hashMapIndexPut idxMap False keysTensor valuesTensor table0)+  let dev    = device idx+      table1 = indexPut False [idx] (toDevice dev valuesTensor) (toDevice dev table0)+  finalizeCudaHashMap idxMap++  -- During inference, look up indices with lookupIndices (unregistered key 99 becomes -1)+  qIdx <- lookupIndices idxMap queryKeysTensor+  print qIdx++  -- indexSelect for registered keys (or: hashMapIndexSelect idxMap 0 validQueryKeysTensor table1)+  validQIdx <- lookupIndices idxMap validQueryKeysTensor+  let selected = indexSelect 0 validQIdx table1+  print selected++keysTensor, valuesTensor, queryKeysTensor, validQueryKeysTensor :: Tensor+keysTensor           = asTensor ([10, 20, 30] :: [Int64])+valuesTensor         = asTensor ([1.5, 2.5, 3.5] :: [Float])+queryKeysTensor      = asTensor ([20, 10, 99] :: [Int64])+validQueryKeysTensor = asTensor ([20, 10, 30] :: [Int64])
+ bench/Main.hs view
@@ -0,0 +1,199 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ForeignFunctionInterface #-}+{-# LANGUAGE ScopedTypeVariables #-}++module Main (main) where++import Control.DeepSeq (force)+import Control.Exception (evaluate)+import Control.Monad (forM_, replicateM_)+import Data.HashMap.Strict (HashMap)+import qualified Data.HashMap.Strict as HM+import Data.Int (Int64)+import Data.List (foldl')+import Foreign.C.Types (CInt(..))+import System.Clock (Clock(Monotonic), getTime, toNanoSecs)+import Text.Printf (printf)++import Torch.Device (Device(..), DeviceType(..))+import Torch.Tensor (Tensor, asTensor, toDevice)+import Torch.CudaHashMap+  ( finalizeCudaHashMap+  , insertBatch+  , insertKeys+  , lookupBatch+  , lookupIndices+  , newCudaHashMapWithCapacity+  )++foreign import ccall "cudaDeviceSynchronize"+  c_cudaDeviceSynchronize :: IO CInt++syncCuda :: IO ()+syncCuda = do+  _ <- c_cudaDeviceSynchronize+  return ()++-- | Measure wall-clock execution time in milliseconds (averaged over @iters@ runs).+timeMs :: forall a. Int -> IO a -> IO (Double, a)+timeMs iters action = do+  syncCuda+  t0 <- getTime Monotonic+  let loop :: Int -> IO a+      loop k+        | k <= 1    = action+        | otherwise = action >> loop (k - 1)+  !res <- loop (max 1 iters)+  syncCuda+  t1 <- getTime Monotonic+  let ns = fromIntegral (toNanoSecs t1 - toNanoSecs t0) :: Double+      ms = (ns / 1.0e6) / fromIntegral (max 1 iters)+  return (ms, res)++-- | Deterministic pseudo-random Int64 key generator.+genKeys :: Int -> Int64 -> [Int64]+genKeys n seed = take n $ drop 1 $ iterate step seed+  where+    step :: Int64 -> Int64+    step x = (x * 6364136223846793005 + 1442695040888963407) `mod` 1000000007++-- | CPU reference implementation for 'insertKeys' (assigns 0-based sequential Int64 indices).+cpuInsertKeys :: [Int64] -> (HashMap Int64 Int64, [Int64])+cpuInsertKeys ks =+  let step :: (HashMap Int64 Int64, Int64, [Int64]) -> Int64 -> (HashMap Int64 Int64, Int64, [Int64])+      step (!hm, !nextIdx, acc) k =+        case HM.lookup k hm of+          Just idx -> (hm, nextIdx, idx : acc)+          Nothing  ->+            let !hm'      = HM.insert k nextIdx hm+                !nextIdx' = nextIdx + 1+            in (hm', nextIdx', nextIdx : acc)+      (!finalHm, _, revIndices) = foldl' step (HM.empty, 0, []) ks+  in (finalHm, reverse revIndices)++printRow :: String -> Int -> Double -> Double -> IO ()+printRow label n baseMs ms = do+  let mops    = (fromIntegral n / 1.0e6) / (ms / 1.0e3)+      speedup = baseMs / ms+  putStrLn $ printf "  %-42s %10.3f ms  (%8.2f Mops/s, %6.2fx)" label ms mops speedup++runBenchmarkForSize :: Int -> IO ()+runBenchmarkForSize n = do+  putStrLn $ replicate 86 '='+  putStrLn $ printf "Benchmark size: N = %d elements (Key: Int64, Value: Float / Index: Int64)" n+  putStrLn $ replicate 86 '='++  let iters   = if n <= 100000 then 5 else 3+      cudaDev = Device CUDA 0+      cpuDev  = Device CPU 0++  -- Prepare input data on CPU and pre-materialize tensors+  !keysList  <- evaluate $ force $ genKeys n 42+  !valsList  <- evaluate $ force [fromIntegral (i `mod` 1000) * 0.5 :: Float | i <- [1 .. n]]+  -- Query keys: 50% hit, 50% miss+  let half = n `div` 2+  !queryList <- evaluate $ force (take half keysList ++ genKeys (n - half) 99999)+  !kvPairs   <- evaluate $ force (zip keysList valsList)++  !cpuKeysTensor  <- evaluate (asTensor keysList :: Tensor)+  !cpuValsTensor  <- evaluate (asTensor valsList :: Tensor)+  !cpuQueryTensor <- evaluate (asTensor queryList :: Tensor)++  !gpuKeysTensor  <- evaluate (toDevice cudaDev cpuKeysTensor)+  !gpuValsTensor  <- evaluate (toDevice cudaDev cpuValsTensor)+  !gpuQueryTensor <- evaluate (toDevice cudaDev cpuQueryTensor)+  syncCuda++  ----------------------------------------------------------------------------+  -- 1. Key-Value Batch Insert+  ----------------------------------------------------------------------------+  putStrLn "\n[1] Batch Insert (Int64 -> Float)"++  (cpuInsertMs, cpuMap) <- timeMs iters $+    evaluate $ force $ HM.fromList kvPairs+  printRow "Data.HashMap.Strict (fromList)" n cpuInsertMs cpuInsertMs++  -- Benchmark CudaHashMap insert (GPU-resident tensors)+  (gpuInsertMs, gpuMap) <- timeMs iters $ do+    m <- newCudaHashMapWithCapacity (n * 2)+    insertBatch m gpuKeysTensor gpuValsTensor+    return m+  printRow "CudaHashMap insertBatch (GPU tensor)" n cpuInsertMs gpuInsertMs++  -- Benchmark CudaHashMap insert (including CPU -> GPU transfer)+  (e2eInsertMs, _) <- timeMs iters $ do+    m <- newCudaHashMapWithCapacity (n * 2)+    insertBatch m cpuKeysTensor cpuValsTensor+    return m+  printRow "CudaHashMap insertBatch (CPU -> GPU)" n cpuInsertMs e2eInsertMs++  ----------------------------------------------------------------------------+  -- 2. Key-Value Batch Lookup (dynamic_map vs static_map vs Data.HashMap)+  ----------------------------------------------------------------------------+  putStrLn "\n[2] Batch Lookup (Int64 -> Float, 50% hit / 50% miss)"++  (cpuLookupMs, _) <- timeMs iters $+    evaluate $ force $ map (\k -> HM.lookupDefault (-1.0) k cpuMap) queryList+  printRow "Data.HashMap.Strict (lookupDefault)" n cpuLookupMs cpuLookupMs++  -- Warmup GPU lookup+  replicateM_ 2 $ lookupBatch gpuMap gpuQueryTensor+  syncCuda++  (gpuDynLookupMs, _) <- timeMs iters $+    lookupBatch gpuMap gpuQueryTensor+  printRow "CudaHashMap lookupBatch (dynamic_map, GPU)" n cpuLookupMs gpuDynLookupMs++  -- Finalize (dynamic_map -> static_map)+  (finalizeMs, _) <- timeMs 1 $+    finalizeCudaHashMap gpuMap+  putStrLn $ printf "  [finalizeCudaHashMap: %.3f ms]" finalizeMs++  (gpuStaticLookupMs, _) <- timeMs iters $+    lookupBatch gpuMap gpuQueryTensor+  printRow "CudaHashMap lookupBatch (static_map, GPU)" n cpuLookupMs gpuStaticLookupMs++  (e2eStaticLookupMs, _) <- timeMs iters $ do+    resGpu <- lookupBatch gpuMap cpuQueryTensor+    evaluate (toDevice cpuDev resGpu)+  printRow "CudaHashMap lookupBatch (static_map, CPU<->GPU)" n cpuLookupMs e2eStaticLookupMs++  ----------------------------------------------------------------------------+  -- 3. Key-to-Index Assignment & Lookup (insertKeys / lookupIndices)+  ----------------------------------------------------------------------------+  putStrLn "\n[3] Key-to-Index Assignment (insertKeys) & Lookup (lookupIndices)"++  (cpuIdxInsertMs, (cpuIdxMap, _)) <- timeMs iters $+    evaluate $ force $ cpuInsertKeys keysList+  printRow "Data.HashMap.Strict (key -> 0..K-1 insert)" n cpuIdxInsertMs cpuIdxInsertMs++  (gpuIdxInsertMs, gpuIdxMap) <- timeMs iters $ do+    m <- newCudaHashMapWithCapacity (n * 2)+    _ <- insertKeys m gpuKeysTensor+    return m+  printRow "CudaHashMap insertKeys (GPU tensor)" n cpuIdxInsertMs gpuIdxInsertMs++  finalizeCudaHashMap gpuIdxMap+  syncCuda++  (cpuIdxLookupMs, _) <- timeMs iters $+    evaluate $ force $ map (\k -> HM.lookupDefault (-1) k cpuIdxMap) queryList+  printRow "Data.HashMap.Strict (index lookup)" n cpuIdxLookupMs cpuIdxLookupMs++  (gpuIdxLookupMs, _) <- timeMs iters $+    lookupIndices gpuIdxMap gpuQueryTensor+  printRow "CudaHashMap lookupIndices (static_map, GPU)" n cpuIdxLookupMs gpuIdxLookupMs+  putStrLn ""++main :: IO ()+main = do+  putStrLn "Benchmarking Torch.CudaHashMap vs Data.HashMap.Strict (unordered-containers)"+  -- Warm up CUDA context and allocator+  warmupMap <- newCudaHashMapWithCapacity 100000+  let wKeys = asTensor ([1, 2, 3] :: [Int64])+      wVals = asTensor ([1.0, 2.0, 3.0] :: [Float])+  insertBatch warmupMap wKeys wVals+  _ <- lookupBatch warmupMap wKeys+  syncCuda++  forM_ [10000, 100000, 1000000] runBenchmarkForSize
+ cbits/CMakeLists.txt view
@@ -0,0 +1,91 @@+cmake_minimum_required(VERSION 3.23.1)++# cuCollections requires Volta (sm_70) or newer for scoped atomic operations+if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES OR CMAKE_CUDA_ARCHITECTURES STREQUAL "" OR CMAKE_CUDA_ARCHITECTURES STREQUAL "52")+    set(CMAKE_CUDA_ARCHITECTURES "70;75;80;86;89;90" CACHE STRING "CUDA architectures" FORCE)+endif()++project(cuco_wrapper LANGUAGES CXX CUDA)++# nvcc 12.x Extended device lambdas support+set(CMAKE_CUDA_STANDARD 17)+set(CMAKE_CXX_STANDARD 17)+set(CMAKE_CUDA_STANDARD_REQUIRED ON)+set(CMAKE_CXX_STANDARD_REQUIRED ON)+set(CMAKE_CUDA_EXTENSIONS ON)+set(CMAKE_CUDA_STANDARD_COMPUSE ON)++set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-extended-lambda")++include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/CPM.cmake")++# Fetch cuCollections via CPM+CPMAddPackage(+    NAME cucollections+    GITHUB_REPOSITORY NVIDIA/cuCollections+    GIT_TAG dev+)++# Configure LibTorch include paths (matching the LibTorch version used under the hood by HaskTorch).+#    Resolved in the order: CMake variables (-DLIBTORCH_HOME=..., etc.) -> environment variables -> default values.++if(NOT DEFINED LIBTORCH_HOME OR LIBTORCH_HOME STREQUAL "")+    if(DEFINED ENV{LIBTORCH_HOME} AND NOT "$ENV{LIBTORCH_HOME}" STREQUAL "")+        set(LIBTORCH_HOME "$ENV{LIBTORCH_HOME}")+    elseif(DEFINED ENV{XDG_CACHE_HOME} AND NOT "$ENV{XDG_CACHE_HOME}" STREQUAL "")+        set(LIBTORCH_HOME "$ENV{XDG_CACHE_HOME}/libtorch")+    elseif(DEFINED ENV{HOME} AND NOT "$ENV{HOME}" STREQUAL "")+        set(LIBTORCH_HOME "$ENV{HOME}/.cache/libtorch")+    endif()+endif()++if(NOT DEFINED LIBTORCH_VERSION OR LIBTORCH_VERSION STREQUAL "")+    if(DEFINED ENV{LIBTORCH_VERSION} AND NOT "$ENV{LIBTORCH_VERSION}" STREQUAL "")+        set(LIBTORCH_VERSION "$ENV{LIBTORCH_VERSION}")+    else()+        set(LIBTORCH_VERSION "2.9.1")+    endif()+endif()++if(NOT DEFINED LIBTORCH_CUDA_VERSION OR LIBTORCH_CUDA_VERSION STREQUAL "")+    if(DEFINED ENV{LIBTORCH_CUDA_VERSION} AND NOT "$ENV{LIBTORCH_CUDA_VERSION}" STREQUAL "")+        set(LIBTORCH_CUDA_VERSION "$ENV{LIBTORCH_CUDA_VERSION}")+    else()+        set(LIBTORCH_CUDA_VERSION "cu128")+    endif()+endif()++if(NOT DEFINED OS OR OS STREQUAL "")+    if(DEFINED ENV{OS} AND NOT "$ENV{OS}" STREQUAL "")+        set(OS "$ENV{OS}")+    else()+        execute_process(COMMAND uname -s OUTPUT_VARIABLE OS OUTPUT_STRIP_TRAILING_WHITESPACE)+    endif()+endif()+string(TOLOWER "${OS}" OS)++execute_process(COMMAND uname -m OUTPUT_VARIABLE LIBTORCH_ARCH OUTPUT_STRIP_TRAILING_WHITESPACE)++set(LIBTORCH_TARGET_DIR "${LIBTORCH_HOME}/${LIBTORCH_VERSION}/${OS}-${LIBTORCH_ARCH}/${LIBTORCH_CUDA_VERSION}")+if(NOT EXISTS "${LIBTORCH_TARGET_DIR}/include" AND EXISTS "${LIBTORCH_HOME}/${LIBTORCH_VERSION}/${OS}-${LIBTORCH_ARCH}/cpu/include")+    set(LIBTORCH_TARGET_DIR "${LIBTORCH_HOME}/${LIBTORCH_VERSION}/${OS}-${LIBTORCH_ARCH}/cpu")+endif()++include_directories("${LIBTORCH_HOME}/include")+include_directories("${LIBTORCH_HOME}/include/torch/csrc/api/include")+include_directories("${LIBTORCH_TARGET_DIR}/include")+include_directories("${LIBTORCH_TARGET_DIR}/include/torch/csrc/api/include")++# Define the wrapper library+add_library(cuco_wrapper STATIC cuco_wrapper.cu)++# Automatically configure header paths for cuCollections+target_link_libraries(cuco_wrapper PRIVATE cuco::cuco)++# Place output directly in the build directory and also remove _deps and CMakeCache.txt on 'make clean'+set_target_properties(cuco_wrapper PROPERTIES+    POSITION_INDEPENDENT_CODE ON+    CUDA_SEPARABLE_COMPILATION OFF+    ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}"+    ADDITIONAL_CLEAN_FILES "${CMAKE_BINARY_DIR}/_deps;${CMAKE_BINARY_DIR}/CMakeCache.txt"+)
+ cbits/cmake/CPM.cmake view
@@ -0,0 +1,28 @@+# SPDX-License-Identifier: MIT+#+# SPDX-FileCopyrightText: Copyright (c) 2019-2023 Lars Melchior and contributors++set(CPM_DOWNLOAD_VERSION 0.43.2)+set(CPM_HASH_SUM "49a3bef91ceb65bb66d57255e12d1ffd22abc2f6408fa9fe4534c544a2f232aa")++if(CPM_PATH)+  set(CPM_DOWNLOAD_LOCATION "${CPM_PATH}/CPM.cmake")+elseif(DEFINED ENV{CPM_PATH})+  file(TO_CMAKE_PATH "$ENV{CPM_PATH}/CPM.cmake" CPM_DOWNLOAD_LOCATION)+elseif(CPM_SOURCE_CACHE)+  set(CPM_DOWNLOAD_LOCATION "${CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")+elseif(DEFINED ENV{CPM_SOURCE_CACHE})+  set(CPM_DOWNLOAD_LOCATION "$ENV{CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")+else()+  set(CPM_DOWNLOAD_LOCATION "${CMAKE_BINARY_DIR}/cmake/CPM_${CPM_DOWNLOAD_VERSION}.cmake")+endif()++# Expand relative path. This is important if the provided path contains a tilde (~)+get_filename_component(CPM_DOWNLOAD_LOCATION ${CPM_DOWNLOAD_LOCATION} ABSOLUTE)++file(DOWNLOAD+     https://github.com/cpm-cmake/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake+     ${CPM_DOWNLOAD_LOCATION} EXPECTED_HASH SHA256=${CPM_HASH_SUM}+)++include(${CPM_DOWNLOAD_LOCATION})
+ cbits/cuco_wrapper.cu view
@@ -0,0 +1,520 @@+#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"
+ cbits/cuco_wrapper.h view
@@ -0,0 +1,92 @@+#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
+ hasktorch-CudaHashMap.cabal view
@@ -0,0 +1,158 @@+cabal-version:      3.0+-- The cabal-version field refers to the version of the .cabal specification,+-- and can be different from the cabal-install (the tool) version and the+-- Cabal (the library) version you are using. As such, the Cabal (the library)+-- version used must be equal or greater than the version stated in this field.+-- Starting from the specification version 2.2, the cabal-version field must be+-- the first thing in the cabal file.++-- Initial package description 'hasktorch-CudaHashMap' generated by+-- 'cabal init'. For further documentation, see:+--   http://haskell.org/cabal/users-guide/+--+-- The name of the package.+name:               hasktorch-CudaHashMap++-- The package version.+-- See the Haskell package versioning policy (PVP) for standards+-- guiding when and how versions should be incremented.+-- https://pvp.haskell.org+-- PVP summary:     +-+------- breaking API changes+--                  | | +----- non-breaking API additions+--                  | | | +--- code changes with no API change+version:            0.1.0.0++-- A short (one-line) description of the package.+synopsis:           HaskTorch interface to the GPU-parallel hashmap in cuCollections++-- A longer description of the package.+-- description:++license:            BSD-3-Clause+license-file:       LICENSE+author:             Susumu Katayama+maintainer:         skata@cs.miyazaki-u.ac.jp++-- A copyright notice.+-- copyright:+category:           Data+build-type:         Custom++-- Extra doc files to be distributed with the package, such as a CHANGELOG or a README.+extra-doc-files:    README.md, CHANGELOG.md++-- Extra source files to be distributed with the package, such as examples, or a tutorial module.+extra-source-files: cbits/cuco_wrapper.cu, cbits/cuco_wrapper.h, cbits/CMakeLists.txt, cbits/cmake/CPM.cmake++custom-setup+    setup-depends:+        base >= 4.7 && < 5,+        Cabal >= 3.0 && < 4,+        directory >= 1.3 && < 2,+        process >= 1.6 && < 2++common warnings+    ghc-options: -Wall++library+    -- Import common warning flags.+    import:           warnings++    exposed-modules:     Torch.CudaHashMap+    hs-source-dirs:      src+    default-language:    Haskell2010+--    ghc-options:         -Wall++    build-depends:       base >= 4.7 && < 5+                     , hasktorch+                     , libtorch-ffi++    include-dirs:        cbits++    -- Paths to the precompiled custom library (libcuco_wrapper.a) and CUDA libraries+    -- are automatically resolved as absolute paths in the confHook of Setup.hs.++    -- Libraries to link against+    --    cuco_wrapper: custom C++ wrapper (libcuco_wrapper.a)+    --    cuda, cudart: CUDA runtime+    --    nvrtc: CUDA runtime compilation (required internally by cuco)+    extra-libraries:     cuco_wrapper+                       , cuda+                       , cudart+                       , nvrtc+                       , stdc++++    -- GHC options for C++ and CUDA integration+    --    -pgml g++ : use g++ instead of gcc as the linker driver (required for C++ exception handling and standard library linking)+    --    --copy-dt-needed-entries : allow resolving indirect DSO symbols such as libc10_cuda.so required by libtorch_cuda.so+    ghc-options:         -pgml g+++                         -optl-Wl,--copy-dt-needed-entries+++    -- Modules included in this library but not exported.+    -- other-modules:++    -- LANGUAGE extensions used by modules in this package.+    -- other-extensions:+++executable CudaHashMap+    import:           warnings+    main-is:          Main.hs+    build-depends:+        base >= 4.7 && < 5,+        hasktorch,+        hasktorch-CudaHashMap++    ghc-options:      -pgml g+++                      -optl-no-pie+                      -optl-Lcbits/build+                      -optl-L/usr/local/cuda/lib64+                      -optl-Wl,--copy-dt-needed-entries++    hs-source-dirs:   app+    default-language: Haskell2010++test-suite CudaHashMap-test+    import:           warnings+    default-language: Haskell2010++    -- The interface type and version of the test suite.+    type:             exitcode-stdio-1.0++    hs-source-dirs:   test+    main-is:          Main.hs++    build-depends:+        base >= 4.7 && < 5,+        hasktorch,+        hasktorch-CudaHashMap++    ghc-options:      -pgml g+++                      -optl-no-pie+                      -optl-Lcbits/build+                      -optl-L/usr/local/cuda/lib64+                      -optl-Wl,--copy-dt-needed-entries++benchmark CudaHashMap-bench+    import:           warnings+    type:             exitcode-stdio-1.0+    hs-source-dirs:   bench+    main-is:          Main.hs+    default-language: Haskell2010++    build-depends:+        base >= 4.7 && < 5,+        clock >= 0.8,+        deepseq >= 1.4,+        hasktorch,+        hasktorch-CudaHashMap,+        unordered-containers >= 0.2++    ghc-options:      -O2+                      -pgml g+++                      -optl-no-pie+                      -optl-Lcbits/build+                      -optl-L/usr/local/cuda/lib64+                      -optl-Wl,--copy-dt-needed-entries+
+ src/Torch/CudaHashMap.hs view
@@ -0,0 +1,318 @@+{-# LANGUAGE ForeignFunctionInterface #-}++module Torch.CudaHashMap+  ( -- * Data Types+    CudaHashMap(..)+  , CudaHashMapOpaque+    -- * Creation, Optimization, and Metadata+  , newCudaHashMap+  , newCudaHashMapWithCapacity+  , newCudaHashMapOnDevice+  , hashMapDevice+  , finalizeCudaHashMap+  , sizeCudaHashMap+    -- * Key-Value Batch Insert & Lookup (supports Int32, Int64, Float, Double)+  , insertBatch+  , lookupBatch+    -- * Index Conversion API for indexSelect / indexPut+  , insertKeys+  , lookupIndices+  , hashMapIndexPut+  , hashMapIndexSelect+    -- * Low-Level Helper+  , withDevicePointer+  ) where++import Foreign.C.Types (CInt(..), CSize(..))+import Foreign.ForeignPtr (ForeignPtr, newForeignPtr, withForeignPtr)+import Foreign.Ptr (FunPtr, Ptr, castPtr)+import Data.Int (Int64)++import Torch.Device (Device(..), DeviceType(..))+import Torch.DType (DType(..))+import qualified Torch.DType as DType+import Torch.Functional (indexPut)+import Torch.Tensor+  ( Tensor(..)+  , contiguous+  , device+  , dtype+  , indexSelect+  , isContiguous+  , numel+  , reshape+  , shape+  , toDevice+  )+import Torch.TensorFactories (mkFactory)+import Torch.TensorOptions (defaultOpts, withDevice, withDType)+import qualified Torch.Internal.Managed.TensorFactories as LibTorch+import qualified Torch.Internal.Managed.Type.Tensor.Tensor0 as ATen++-- | Opaque type representing a pointer to a C++ hash map instance+data CudaHashMapOpaque++-- | GPU hash map type handled on the Haskell side+newtype CudaHashMap = CudaHashMap (ForeignPtr CudaHashMapOpaque)++-- FFI declarations for C++ creation, destruction, and metadata functions+foreign import ccall "custom_cuco_create"+  c_create :: IO (Ptr CudaHashMapOpaque)++foreign import ccall "custom_cuco_create_with_capacity"+  c_create_with_capacity :: CSize -> IO (Ptr CudaHashMapOpaque)++foreign import ccall "custom_cuco_create_on_device"+  c_create_on_device :: CInt -> CSize -> IO (Ptr CudaHashMapOpaque)++foreign import ccall "&custom_cuco_destroy"+  c_destroy :: FunPtr (Ptr CudaHashMapOpaque -> IO ())++foreign import ccall "custom_cuco_device_index"+  c_device_index :: Ptr CudaHashMapOpaque -> IO CInt++foreign import ccall "custom_cuco_set_device_index"+  c_set_device_index :: Ptr CudaHashMapOpaque -> CInt -> IO ()++foreign import ccall "custom_cuco_size"+  c_size :: Ptr CudaHashMapOpaque -> IO CSize++foreign import ccall "custom_cuco_val_dtype"+  c_val_dtype :: Ptr CudaHashMapOpaque -> IO CInt++foreign import ccall "custom_cuco_finalize"+  c_finalize :: Ptr CudaHashMapOpaque -> IO ()++-- FFI declarations for C++ Key-to-Index functions+foreign import ccall "custom_cuco_insert_keys"+  c_insert_keys+    :: Ptr CudaHashMapOpaque+    -> CInt+    -> CInt+    -> Ptr ()+    -> Ptr Int64+    -> IO ()++foreign import ccall "custom_cuco_lookup_indices"+  c_lookup_indices+    :: Ptr CudaHashMapOpaque+    -> CInt+    -> CInt+    -> Ptr ()+    -> Ptr Int64+    -> IO ()++-- FFI declarations for C++ typed batch Insert / Lookup functions+foreign import ccall "custom_cuco_insert_batch_typed"+  c_insert_batch_typed+    :: Ptr CudaHashMapOpaque+    -> CInt+    -> CInt+    -> CInt+    -> Ptr ()+    -> Ptr ()+    -> IO ()++foreign import ccall "custom_cuco_lookup_batch_typed"+  c_lookup_batch_typed+    :: Ptr CudaHashMapOpaque+    -> CInt+    -> CInt+    -> CInt+    -> Ptr ()+    -> Ptr ()+    -> IO ()++-- | Convert a HaskTorch DType to the C++ wrapper's CucoDType code+toCucoDType :: DType -> CInt+toCucoDType DType.Int32  = 0+toCucoDType DType.Int64  = 1+toCucoDType DType.Float  = 2+toCucoDType DType.Double = 3+toCucoDType dt =+  error $ "Torch.CudaHashMap: unsupported Tensor DType " ++ show dt+       ++ " (supported: Int32, Int64, Float, Double)"++fromCucoDType :: CInt -> Maybe DType+fromCucoDType 0 = Just DType.Int32+fromCucoDType 1 = Just DType.Int64+fromCucoDType 2 = Just DType.Float+fromCucoDType 3 = Just DType.Double+fromCucoDType _ = Nothing++-- | Determine or retrieve the CUDA device associated with the map, and ensure that the input tensor+-- resides in contiguous memory on that CUDA device.+-- If the map's device is not yet determined (-1), adopts the input tensor's device index if it is+-- a CUDA tensor, or defaults to cuda:0 (Device CUDA 0) if it is a CPU tensor, and records it in the map.+ensureCudaOnMap :: Ptr CudaHashMapOpaque -> Tensor -> IO Tensor+ensureCudaOnMap mapPtr t = do+  devIdx <- c_device_index mapPtr+  targetDev <-+    if devIdx >= 0+      then return $ Device CUDA (fromIntegral devIdx)+      else case device t of+             d@(Device CUDA idx) -> do+               c_set_device_index mapPtr (fromIntegral idx)+               return d+             _ -> do+               c_set_device_index mapPtr 0+               return $ Device CUDA 0+  let tCuda = if device t == targetDev then t else toDevice targetDev t+  return $ if isContiguous tCuda then tCuda else contiguous tCuda++-- | Allocate an output tensor with the specified shape, DType, and Device+allocTensor :: [Int] -> DType -> Device -> IO Tensor+allocTensor dims dt dev =+  let opts = withDType dt $ withDevice dev defaultOpts+  in mkFactory LibTorch.empty_lo dims opts++-- | Extract the raw device pointer (GPU pointer) from a Tensor and run an action with it+withDevicePointer :: Tensor -> (Ptr () -> IO a) -> IO a+withDevicePointer t action = do+  let Unsafe fptr = if isContiguous t then t else contiguous t+  withForeignPtr fptr $ \_ -> do+    ptr <- ATen.tensor_data_ptr fptr+    action ptr++-- | Initialize a hash map (starts as a dynamically resizable dynamic_map).+-- The target GPU device is automatically set to the device of the first CUDA tensor inserted or queried+-- (or cuda:0 if the first tensor is on the CPU).+newCudaHashMap :: IO CudaHashMap+newCudaHashMap = do+  ptr <- c_create+  fptr <- newForeignPtr c_destroy ptr+  return $ CudaHashMap fptr++-- | Initialize a hash map with a specified initial capacity+newCudaHashMapWithCapacity :: Int -> IO CudaHashMap+newCudaHashMapWithCapacity cap = do+  ptr <- c_create_with_capacity (fromIntegral (max 100000 cap))+  fptr <- newForeignPtr c_destroy ptr+  return $ CudaHashMap fptr++-- | Initialize a hash map with an explicit CUDA device (e.g., @Device CUDA 1@) and initial capacity+newCudaHashMapOnDevice :: Device -> Int -> IO CudaHashMap+newCudaHashMapOnDevice (Device CUDA devIdx) cap = do+  ptr <- c_create_on_device (fromIntegral devIdx) (fromIntegral (max 100000 cap))+  fptr <- newForeignPtr c_destroy ptr+  return $ CudaHashMap fptr+newCudaHashMapOnDevice dev _ =+  error $ "Torch.CudaHashMap.newCudaHashMapOnDevice: expected CUDA device, got " ++ show dev++-- | Get the CUDA device associated with the hash map (returns @Device CUDA 0@ if not yet determined)+hashMapDevice :: CudaHashMap -> IO Device+hashMapDevice (CudaHashMap fptr) =+  withForeignPtr fptr $ \mapPtr -> do+    devIdx <- c_device_index mapPtr+    return $ Device CUDA (if devIdx >= 0 then fromIntegral devIdx else 0)++-- | Get the number of unique keys currently registered in the hash map+sizeCudaHashMap :: CudaHashMap -> IO Int+sizeCudaHashMap (CudaHashMap fptr) =+  withForeignPtr fptr $ \mapPtr ->+    fromIntegral <$> c_size mapPtr++-- | Convert and freeze the internal dynamic_map into a single high-speed static_map once insertions have stabilized+finalizeCudaHashMap :: CudaHashMap -> IO ()+finalizeCudaHashMap (CudaHashMap fptr) =+  withForeignPtr fptr c_finalize++-- | Batch-insert key and value tensors (overwrites with the new value if a key already exists: insert_or_assign).+-- Both keys and values support Int32, Int64, Float, and Double.+insertBatch :: CudaHashMap -> Tensor -> Tensor -> IO ()+insertBatch (CudaHashMap fptr) keysTensor valuesTensor =+  withForeignPtr fptr $ \mapPtr -> do+    cudaKeys   <- ensureCudaOnMap mapPtr keysTensor+    cudaValues <- ensureCudaOnMap mapPtr valuesTensor+    let kDType  = toCucoDType (dtype cudaKeys)+        vDType  = toCucoDType (dtype cudaValues)+        numKeys = numel cudaKeys+        numVals = numel cudaValues+    if numKeys /= numVals+      then error $ "Torch.CudaHashMap.insertBatch: key and value element counts do not match ("+                ++ show numKeys ++ " vs " ++ show numVals ++ ")"+      else do+        let numElements = fromIntegral numKeys+        withDevicePointer cudaKeys $ \keyPtr ->+          withDevicePointer cudaValues $ \valPtr ->+            c_insert_batch_typed mapPtr kDType vDType numElements keyPtr valPtr++-- | Batch-lookup a tensor of keys and return the corresponding tensor of values.+-- Unregistered keys receive -1 for Int32 / Int64, or NaN for Float / Double.+lookupBatch :: CudaHashMap -> Tensor -> IO Tensor+lookupBatch (CudaHashMap fptr) keysTensor =+  withForeignPtr fptr $ \mapPtr -> do+    cudaKeys <- ensureCudaOnMap mapPtr keysTensor+    let kDType      = toCucoDType (dtype cudaKeys)+        numElements = fromIntegral (numel cudaKeys)+    storedValCode <- c_val_dtype mapPtr+    let outDType = case fromCucoDType storedValCode of+                     Just dt -> dt+                     Nothing -> dtype cudaKeys+        vDType   = toCucoDType outDType+    valuesTensor <- allocTensor (shape cudaKeys) outDType (device cudaKeys)+    withDevicePointer cudaKeys $ \keyPtr ->+      withDevicePointer valuesTensor $ \valPtr ->+        c_lookup_batch_typed mapPtr kDType vDType numElements keyPtr valPtr+    return valuesTensor++-- | Assign and return 0-based sequential indices (DType.Int64) for a tensor of keys.+-- Returns the existing index for already registered keys, and assigns new sequential indices to unregistered keys.+-- The returned tensor can be used directly as the index tensor for HaskTorch's 'indexSelect' or 'indexPut'.+insertKeys :: CudaHashMap -> Tensor -> IO Tensor+insertKeys (CudaHashMap fptr) keysTensor =+  withForeignPtr fptr $ \mapPtr -> do+    cudaKeys <- ensureCudaOnMap mapPtr keysTensor+    let kDType      = toCucoDType (dtype cudaKeys)+        numElements = fromIntegral (numel cudaKeys)+    indicesTensor <- allocTensor (shape cudaKeys) DType.Int64 (device cudaKeys)+    withDevicePointer cudaKeys $ \keyPtr ->+      withDevicePointer indicesTensor $ \idxPtr ->+        c_insert_keys mapPtr kDType numElements keyPtr (castPtr idxPtr)+    return indicesTensor++-- | Look up 0-based Int64 indices corresponding to a tensor of keys without registering new keys.+-- Unregistered keys receive -1.+lookupIndices :: CudaHashMap -> Tensor -> IO Tensor+lookupIndices (CudaHashMap fptr) keysTensor =+  withForeignPtr fptr $ \mapPtr -> do+    cudaKeys <- ensureCudaOnMap mapPtr keysTensor+    let kDType      = toCucoDType (dtype cudaKeys)+        numElements = fromIntegral (numel cudaKeys)+    indicesTensor <- allocTensor (shape cudaKeys) DType.Int64 (device cudaKeys)+    withDevicePointer cudaKeys $ \keyPtr ->+      withDevicePointer indicesTensor $ \idxPtr ->+        c_lookup_indices mapPtr kDType numElements keyPtr (castPtr idxPtr)+    return indicesTensor++-- | Obtain sequential indices from a tensor of keys (registering any unseen keys)+-- and write valuesTensor into the corresponding indices of tableTensor via indexPut.+hashMapIndexPut+  :: CudaHashMap+  -> Bool    -- ^ accumulate (True to add, False to overwrite)+  -> Tensor  -- ^ keysTensor+  -> Tensor  -- ^ valuesTensor+  -> Tensor  -- ^ tableTensor (target tensor to write into)+  -> IO Tensor+hashMapIndexPut hmap@(CudaHashMap fptr) accumulate keysTensor valuesTensor tableTensor = do+  indices <- insertKeys hmap keysTensor+  withForeignPtr fptr $ \mapPtr -> do+    cudaValues <- ensureCudaOnMap mapPtr valuesTensor+    cudaTable  <- ensureCudaOnMap mapPtr tableTensor+    let flatIndices = reshape [-1] indices+    return $ indexPut accumulate [flatIndices] cudaValues cudaTable++-- | Look up registered indices from a tensor of keys and extract elements+-- from the specified dimension of tableTensor via indexSelect.+hashMapIndexSelect+  :: CudaHashMap+  -> Int     -- ^ dimension (typically 0)+  -> Tensor  -- ^ keysTensor+  -> Tensor  -- ^ tableTensor+  -> IO Tensor+hashMapIndexSelect hmap@(CudaHashMap fptr) selectDim keysTensor tableTensor = do+  indices <- lookupIndices hmap keysTensor+  withForeignPtr fptr $ \mapPtr -> do+    cudaTable <- ensureCudaOnMap mapPtr tableTensor+    let flatIndices = reshape [-1] indices+    return $ indexSelect selectDim flatIndices cudaTable
+ test/Main.hs view
@@ -0,0 +1,42 @@+module Main (main) where++import Data.Int (Int32, Int64)+import Torch.Functional (indexPut)+import Torch.Tensor (asTensor, device, indexSelect, toDevice)+import Torch.TensorFactories (zeros')+import Torch.CudaHashMap++main :: IO ()+main = do+  putStrLn "=== Test 1: Int32 Key -> Int32 Value ==="+  m1 <- newCudaHashMap+  let k1 = asTensor ([1, 2, 3] :: [Int32])+      v1 = asTensor ([10, 20, 30] :: [Int32])+      v1_overwrite = asTensor ([100, 200, 300] :: [Int32])+  insertBatch m1 k1 v1+  insertBatch m1 k1 v1_overwrite+  finalizeCudaHashMap m1+  res1 <- lookupBatch m1 (asTensor ([3, 1, 99] :: [Int32]))+  print res1++  putStrLn "=== Test 2: Int64 Key -> Float Value ==="+  m2 <- newCudaHashMap+  let k2 = asTensor ([100, 200, 300] :: [Int64])+      v2 = asTensor ([1.25, 2.5, 3.75] :: [Float])+  insertBatch m2 k2 v2+  finalizeCudaHashMap m2+  res2 <- lookupBatch m2 (asTensor ([200, 100, 999] :: [Int64]))+  print res2++  putStrLn "=== Test 3: Key -> Index (Int64) for indexPut / indexSelect ==="+  m3 <- newCudaHashMap+  let table0 = zeros' [10]+      batchKeys1 = asTensor ([500, 600, 500, 700] :: [Int64])+      batchVals1 = asTensor ([5.0, 6.0, 5.5, 7.0] :: [Float])+  idx1 <- insertKeys m3 batchKeys1+  let dev    = device idx1+      table1 = indexPut False [idx1] (toDevice dev batchVals1) (toDevice dev table0)+  finalizeCudaHashMap m3+  qIdx <- lookupIndices m3 (asTensor ([700, 500, 600] :: [Int64]))+  let out = indexSelect 0 qIdx table1+  print out