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