hasktorch-CudaHashMap-0.1.0.0: app/Main.hs

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])