module Data.Repa.Eval.Generic.Par.Interleaved
        (fillInterleaved)
where
import Data.Repa.Eval.Gang
import GHC.Exts


-- | Fill something in parallel, using a round-robin order.
-- 
--   * Threads handle elements in row major, round-robin order.
--
--   * Using this method helps even out unbalanced workloads.
--
fillInterleaved
        :: Gang                 -- ^ Gang to run the operation on.
        -> (Int# -> a -> IO ()) -- ^ Update function to write into result buffer.
        -> (Int# -> a)          -- ^ Function to get the value at a given index.
        -> Int#                 -- ^ Number of elements.
        -> IO ()

fillInterleaved :: forall a.
Gang -> (Int# -> a -> IO ()) -> (Int# -> a) -> Int# -> IO ()
fillInterleaved Gang
gang Int# -> a -> IO ()
write Int# -> a
getElem Int#
len 
 = Gang -> (Int# -> IO ()) -> IO ()
gangIO Gang
gang
 ((Int# -> IO ()) -> IO ()) -> (Int# -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$  \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.
        !threads :: Int#
threads        = Gang -> Int#
gangSize Gang
gang

        -- 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 a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
             | Bool
otherwise
             = do Int# -> a -> IO ()
write Int#
ix (Int# -> a
getElem Int#
ix)
                  Int# -> Int# -> IO ()
go (Int#
ix Int# -> Int# -> Int#
+# Int#
step) (Int#
count Int# -> Int# -> Int#
-# Int#
1#)
        {-# INLINE fill #-}
{-# INLINE [0] fillInterleaved #-}