{-# 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