{-# LANGUAGE MagicHash #-}
-- | Evaluate an array in parallel in an interleaved fashion,
--  with each by having each processor computing alternate elements.
module Data.Array.Repa.Eval.Interleaved
        ( fillInterleavedP)
where
import Data.Array.Repa.Eval.Gang
import GHC.Exts
import Prelude          as P


-- | Fill something in parallel.
-- 
--   * The array is split into linear chunks and each thread fills one chunk.
-- 
fillInterleavedP
        :: Int                  -- ^ Number of elements.
        -> (Int -> a -> IO ())  -- ^ Update function to write into result buffer.
        -> (Int -> a)           -- ^ Fn to get the value at a given index.
        -> IO ()

{-# INLINE [0] fillInterleavedP #-}
fillInterleavedP :: Int -> (Int -> a -> IO ()) -> (Int -> a) -> IO ()
fillInterleavedP !(I# Int#
len) Int -> a -> IO ()
write Int -> a
getElem
 =      Gang -> (Int -> IO ()) -> IO ()
gangIO Gang
theGang
         ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$  \(I# Int#
thread) -> 
              let !step :: Int#
step    = Int#
threads
                  !start :: Int#
start   = Int#
thread
                  !count :: Int#
count   = Int# -> Int#
elemsForThread Int#
thread
              in  Int# -> Int# -> Int# -> IO ()
fill Int#
step Int#
start Int#
count

 where
        -- Decide now to split the work across the threads.
        !(I# Int#
threads)   = Gang -> Int
gangSize Gang
theGang

        -- All threads get this many elements.
        !chunkLenBase :: Int#
chunkLenBase   = Int#
len Int# -> Int# -> Int#
`quotInt#` Int#
threads

        -- Leftover elements to divide between first few threads.
        !chunkLenSlack :: Int#
chunkLenSlack  = Int#
len Int# -> Int# -> Int#
`remInt#`  Int#
threads

        -- How many elements to compute with this thread.
        elemsForThread :: Int# -> Int#
elemsForThread Int#
thread
         | Int#
1# <- Int#
thread Int# -> Int# -> Int#
<# Int#
chunkLenSlack
         = Int#
chunkLenBase Int# -> Int# -> Int#
+# Int#
1#

         | Bool
otherwise
         = Int#
chunkLenBase
        {-# INLINE elemsForThread #-}

        -- Evaluate the elements of a single chunk.
        fill :: Int# -> Int# -> Int# -> IO ()
fill !Int#
step !Int#
ix0 !Int#
count0
         = Int# -> Int# -> IO ()
go Int#
ix0 Int#
count0
         where
          go :: Int# -> Int# -> IO ()
go !Int#
ix !Int#
count
             | Int#
1# <- Int#
count Int# -> Int# -> Int#
<=# Int#
0# 
             = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

             | Bool
otherwise
             = do Int -> a -> IO ()
write (Int# -> Int
I# Int#
ix) (Int -> a
getElem (Int# -> Int
I# Int#
ix))
                  Int# -> Int# -> IO ()
go (Int#
ix Int# -> Int# -> Int#
+# Int#
step) (Int#
count Int# -> Int# -> Int#
-# Int#
1#)
        {-# INLINE fill #-}