<commit_before>--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Utils
-- Copyright : (c) 2010 <NAME>
-- License   : BSD
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Utils where

import Foreign
import Control.Exception

import Data.Vector.Generic       as G
import Data.Vector.Fusion.Stream as Stream
import Foreign.CUDA              as CUDA


--
-- Copy a vector to the device and perform a computation, returning the result.
-- This requires the data to be marshalled to a heap allocated array so that it
-- can be copied to the device.
--
withVector :: (G.Vector v a, Storable a) => v a -> (DevicePtr a -> IO b) -> IO b
withVector vec action = let l = G.length vec in
  CUDA.allocaArray l  $ \d_ptr -> do

    -- Transfer via a pinned heap array, which the device can read directly at
    -- much higher bandwidth than pageable memory.
    --
    bracket (CUDA.mallocHostArray [] l) CUDA.freeHost $
      flip withHostPtr $ \ptr ->
        Stream.foldM' (\p e -> poke p e >> return (p `advancePtr` 1)) ptr (G.stream vec) >>
        CUDA.pokeArray l ptr d_ptr

    -- Release host memory and perform computation on device array
    --
    action d_ptr
<commit_msg>specialised vector marshalling utility for storable-based vectors<commit_after>--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Utils
-- Copyright : (c) 2010 <NAME>
-- License   : BSD
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Utils where

import Foreign
import Control.Exception

import Data.Vector.Storable      as S
import Data.Vector.Generic       as G
import Data.Vector.Fusion.Stream as Stream
import Foreign.CUDA              as CUDA


--
-- Copy a vector to the device and perform a computation, returning the result.
-- This requires the data to be marshalled to a heap allocated array so that it
-- can be copied to the device.
--
withVector :: (G.Vector v a, Storable a) => v a -> (DevicePtr a -> IO b) -> IO b
withVector vec action = let l = G.length vec in
  CUDA.allocaArray l  $ \d_ptr -> do

    -- Transfer via a pinned heap array, which the device can read directly at
    -- much higher bandwidth than pageable memory.
    --
    bracket (CUDA.mallocHostArray [] l) CUDA.freeHost $
      flip withHostPtr $ \ptr ->
        Stream.foldM' (\p e -> poke p e >> return (p `advancePtr` 1)) ptr (G.stream vec) >>
        CUDA.pokeArray l ptr d_ptr

    -- Release host memory and perform computation on device array
    --
    action d_ptr


--
-- Copy the contents of a storable-based vector to the device and perform a
-- computation, returning the result. As storable vectors are stored as foreign
-- pointers which are visible to the C heap, it is not necessary to marshal the
-- data via a temporary array.
--
withVectorS :: Storable a => S.Vector a -> (DevicePtr a -> IO b) -> IO b
withVectorS vec action = let l = S.length vec in
  CUDA.allocaArray l $ \d_ptr -> do
    S.unsafeWith vec $ \p -> CUDA.pokeArray l p d_ptr
    action d_ptr

{'fraction_non_alphanumeric': 0.18557025834908633, 'fraction_numerical': 0.00315059861373661, 'mean_word_length': 1.8424350940017906, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7370917', 'n_tokens_mistral': 881, 'n_tokens_neox': 840, 'n_words': 370}