module Data.Array.Repa.Eval.Chunked
( fillChunkedP
, fillChunkedS
, fillChunkedS'
, fillChunkedIOP)
where
import Data.Array.Repa.Eval.Gang
import GHC.Exts
import Prelude as P
fillChunkedS
:: Int
-> (Int -> a -> IO ())
-> (Int -> a)
-> IO ()
fillChunkedS !(I# len) !write !getElem
= fill 0#
where fill !ix
| ix >=# len = return ()
| otherwise
= do write (I# ix) (getElem (I# ix))
fill (ix +# 1#)
fillChunkedS'
:: Int
-> (Int -> IO ())
-> IO ()
fillChunkedS' !(I# len) eat
= fill 0#
where fill !ix
| ix >=# len = return ()
| otherwise
= do eat (I# ix)
fill (ix +# 1#)
fillChunkedP
:: Int
-> (Int -> a -> IO ())
-> (Int -> a)
-> IO ()
fillChunkedP !(I# len) !write !getElem
= gangIO theGang
$ \(I# thread) ->
let !start = splitIx thread
!end = splitIx (thread +# 1#)
in fill start end
where
!(I# threads) = gangSize theGang
!chunkLen = len `quotInt#` threads
!chunkLeftover = len `remInt#` threads
splitIx thread
| thread <# chunkLeftover = thread *# (chunkLen +# 1#)
| otherwise = thread *# chunkLen +# chunkLeftover
fill !ix !end
| ix >=# end = return ()
| otherwise
= do write (I# ix) (getElem (I# ix))
fill (ix +# 1#) end
fillChunkedIOP
:: Int
-> (Int -> a -> IO ())
-> (Int -> IO (Int -> IO a))
-> IO ()
fillChunkedIOP !(I# len) !write !mkGetElem
= gangIO theGang
$ \(I# thread) ->
let !start = splitIx thread
!end = splitIx (thread +# 1#)
in fillChunk thread start end
where
!(I# threads) = gangSize theGang
!chunkLen = len `quotInt#` threads
!chunkLeftover = len `remInt#` threads
splitIx thread
| thread <# chunkLeftover = thread *# (chunkLen +# 1#)
| otherwise = thread *# chunkLen +# chunkLeftover
fillChunk !thread !ixStart !ixEnd
= do getElem <- mkGetElem (I# thread)
fill getElem ixStart ixEnd
fill !getElem !ix0 !end
= go ix0
where go !ix
| ix >=# end = return ()
| otherwise
= do x <- getElem (I# ix)
write (I# ix) x
go (ix +# 1#)