<commit_before>{-# LANGUAGE FlexibleContexts, ParallelListComp #-}
--------------------------------------------------------------------------------
--
-- Module    : VectorAdd
-- Copyright : (c) 2009 <NAME>
-- License   : BSD
--
-- Element-wise addition of two vectors
--
--------------------------------------------------------------------------------

module Main where

import Foreign
import Foreign.CUDA.Driver (forceEither)
import qualified Foreign.CUDA.Driver as CUDA

import Data.Array.Storable
import System.Random
import Control.Monad
import Control.Exception


type Vector e = StorableArray Int e

withVector :: Vector e -> (Ptr e -> IO a) -> IO a
withVector = withStorableArray


-- To ensure the array is fully evaluated, force one element
--
evaluateVec :: Storable e => Vector e -> IO (Vector e)
evaluateVec vec = (join $ evaluate (vec `readArray` 0)) >> return vec

-- Generate a new random vector
--
randomVec :: (Num e, Storable e, Random e, MArray StorableArray e IO)
          => Int -> IO (Vector e)
randomVec n = do
  rg <- newStdGen
  let -- The standard random number generator is too slow to generate really
      -- large vectors. Instead, we generate a short vector and repeat that.
      k     = 1000
      rands = take k (randomRs (-100,100) rg)

  newListArray (0,n-1) [rands !! (i`mod`k) | i <- [0..n-1]] >>= evaluateVec


--------------------------------------------------------------------------------
-- Reference implementation
--------------------------------------------------------------------------------

testRef :: (Num e, Storable e) => Vector e -> Vector e -> IO (Vector e)
testRef xs ys = do
  (i,j) <- getBounds xs
  res   <- newArray_ (i,j)
  forM_ [i..j] (add res)
  evaluateVec res
  where
    add res idx = do
      a <- readArray xs idx
      b <- readArray ys idx
      writeArray res idx (a+b)

--------------------------------------------------------------------------------
-- CUDA
--------------------------------------------------------------------------------

initCUDA :: IO (CUDA.Fun)
initCUDA = do
  res <- CUDA.initialise []
  case res of
    Just e  -> error e
    Nothing -> do
      mdl <- forceEither `fmap` CUDA.loadFile   "data/VectorAdd.ptx"
      fun <- forceEither `fmap` CUDA.getFun mdl "VecAdd"
      return fun

initData :: (Num e, Storable e)
         => Vector e -> Vector e -> IO (CUDA.DevicePtr e, CUDA.DevicePtr e, CUDA.DevicePtr e)
initData xs ys = do
  (m,n) <- getBounds xs
  let len = (n-m+1)
  dxs   <- forceEither `fmap` CUDA.malloc len
  dys   <- forceEither `fmap` CUDA.malloc len
  res   <- forceEither `fmap` CUDA.malloc len
  withVector xs $ \p -> CUDA.pokeArray len p dxs
  withVector ys $ \p -> CUDA.pokeArray len p dys
  return (dxs, dys, res)


testCUDA :: (Num e, Storable e) => Vector e -> Vector e -> IO (Vector e)
testCUDA xs ys = do
  -- Initialise environment and copy over test data
  --
  addVec  <- initCUDA
  (x,y,z) <- initData xs ys<commit_msg>need to explicitly create the context<commit_after>{-# LANGUAGE FlexibleContexts, ParallelListComp #-}
--------------------------------------------------------------------------------
--
-- Module    : VectorAdd
-- Copyright : (c) 2009 <NAME>
-- License   : BSD
--
-- Element-wise addition of two vectors
--
--------------------------------------------------------------------------------

module Main where

import Foreign
import Foreign.CUDA.Driver (forceEither)
import qualified Foreign.CUDA.Driver as CUDA

import Data.Array.Storable
import System.Random
import Control.Monad
import Control.Exception


type Vector e = StorableArray Int e

withVector :: Vector e -> (Ptr e -> IO a) -> IO a
withVector = withStorableArray


-- To ensure the array is fully evaluated, force one element
--
evaluateVec :: Storable e => Vector e -> IO (Vector e)
evaluateVec vec = (join $ evaluate (vec `readArray` 0)) >> return vec

-- Generate a new random vector
--
randomVec :: (Num e, Storable e, Random e, MArray StorableArray e IO)
          => Int -> IO (Vector e)
randomVec n = do
  rg <- newStdGen
  let -- The standard random number generator is too slow to generate really
      -- large vectors. Instead, we generate a short vector and repeat that.
      k     = 1000
      rands = take k (randomRs (-100,100) rg)

  newListArray (0,n-1) [rands !! (i`mod`k) | i <- [0..n-1]] >>= evaluateVec


--------------------------------------------------------------------------------
-- Reference implementation
--------------------------------------------------------------------------------

testRef :: (Num e, Storable e) => Vector e -> Vector e -> IO (Vector e)
testRef xs ys = do
  (i,j) <- getBounds xs
  res   <- newArray_ (i,j)
  forM_ [i..j] (add res)
  evaluateVec res
  where
    add res idx = do
      a <- readArray xs idx
      b <- readArray ys idx
      writeArray res idx (a+b)

--------------------------------------------------------------------------------
-- CUDA
--------------------------------------------------------------------------------

initCUDA :: IO (CUDA.Fun)
initCUDA = do
  res <- CUDA.initialise []
  case res of
    Just e  -> error e
    Nothing -> do
      dev <- forceEither `fmap` CUDA.get 0
      ctx <- forceEither `fmap` CUDA.create dev []
      mdl <- forceEither `fmap` CUDA.loadFile   "data/VectorAdd.ptx"
      fun <- forceEither `fmap` CUDA.getFun mdl "VecAdd"
      return fun

initData :: (Num e, Storable e)
         => Vector e -> Vector e -> IO (CUDA.DevicePtr e, CUDA.DevicePtr e, CUDA.DevicePtr e)
initData xs ys = do
  (m,n) <- getBounds xs
  let len = (n-m+1)
  dxs   <- forceEither `fmap` CUDA.malloc len
  dys   <- forceEither `fmap` CUDA.malloc len
  res   <- forceEither `fmap` CUDA.malloc len
  withVector xs $ \p -> CUDA.pokeArray len p dxs
  withVector ys $ \p -> CUDA.pokeArray len p dys
  return (dxs, dys, res)


testCUDA :: (Num e, Storable e) => Vector e -> Vector e -> IO (Vector e)
testCUDA xs ys = do
  -- Initialise environment and copy over test data
  --
  addVec  <- initCUDA
  (x,y,z) <- initData xs ys
{'fraction_non_alphanumeric': 0.25906218822746924, 'fraction_numerical': 0.0068174260059860325, 'mean_word_length': 1.5804375804375805, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 53, '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': '7376265', 'n_tokens_mistral': 1856, 'n_tokens_neox': 1720, 'n_words': 634}