{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Data.Massiv.Core.Loop (
loop,
loopF,
nextMaybeF,
loopA,
loopA_,
loopM,
loopM_,
iloopM,
iloopA_,
loopNextM,
loopNextA_,
loopDeepM,
splitLinearly,
splitLinearlyM,
splitLinearlyM_,
splitLinearlyWith_,
splitLinearlyWithM_,
splitLinearlyWithStartAtM_,
splitLinearlyWithStatefulM_,
iterLinearST_,
iterLinearAccST_,
iterLinearAccST,
splitNumChunks,
stepStartAdjust,
splitWorkWithFactorST,
scheduleMassivWork,
withMassivScheduler_,
) where
import Control.Monad (void, when)
import Control.Monad.IO.Unlift (MonadUnliftIO (..))
import Control.Monad.Primitive
import Control.Monad.ST (ST)
import Control.Scheduler (
Comp (..),
Scheduler,
SchedulerWS,
numWorkers,
scheduleWork,
scheduleWorkState_,
scheduleWork_,
trivialScheduler_,
unwrapSchedulerWS,
withScheduler_,
)
import Control.Scheduler.Global (globalScheduler, withGlobalScheduler_)
import Data.Coerce
import Data.Functor.Identity
loop :: Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a
loop :: forall a.
Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a
loop Int
initial Int -> Bool
condition Int -> Int
increment a
initAcc Int -> a -> a
f =
forall a. Identity a -> a
runIdentity (forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
initial Int -> Bool
condition Int -> Int
increment a
initAcc (coerce :: forall a b. Coercible a b => a -> b
coerce Int -> a -> a
f))
{-# INLINE loop #-}
loopM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM :: forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM !Int
initial Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> a -> m a
f =
Int -> a -> m a
go Int
initial a
initAcc
where
go :: Int -> a -> m a
go !Int
step !a
acc
| Int -> Bool
condition Int
step = Int -> a -> m a
f Int
step a
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> a -> m a
go (Int -> Int
increment Int
step)
| Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
{-# INLINE loopM #-}
iloopM
:: Monad m => Int -> Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> Int -> a -> m a) -> m a
iloopM :: forall (m :: * -> *) a.
Monad m =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
iloopM !Int
istart !Int
initIx Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> Int -> a -> m a
f = Int -> Int -> a -> m a
go Int
istart Int
initIx a
initAcc
where
go :: Int -> Int -> a -> m a
go !Int
i !Int
step !a
acc
| Int -> Bool
condition Int
step = Int -> Int -> a -> m a
f Int
i Int
step a
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Int -> a -> m a
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int
increment Int
step)
| Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
{-# INLINE iloopM #-}
loopM_ :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ :: forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ !Int
initial Int -> Bool
condition Int -> Int
increment Int -> m a
f = Int -> m ()
go Int
initial
where
go :: Int -> m ()
go !Int
step
| Int -> Bool
condition Int
step = Int -> m a
f Int
step forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
go (Int -> Int
increment Int
step)
| Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINE loopM_ #-}
{-# DEPRECATED loopM_ "In favor of `loopA_`" #-}
iloopA_
:: Applicative f => Int -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f ()
iloopA_ :: forall (f :: * -> *) a.
Applicative f =>
Int
-> Int
-> (Int -> Bool)
-> (Int -> Int)
-> (Int -> Int -> f a)
-> f ()
iloopA_ !Int
istart !Int
initIx Int -> Bool
condition Int -> Int
increment Int -> Int -> f a
f = Int -> Int -> f ()
go Int
istart Int
initIx
where
go :: Int -> Int -> f ()
go !Int
i !Int
step
| Int -> Bool
condition Int
step = Int -> Int -> f a
f Int
i Int
step forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> Int -> f ()
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int
increment Int
step)
| Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINE iloopA_ #-}
loopNextA_ :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f ()
loopNextA_ :: forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f ()
loopNextA_ !Int
initial Int -> Bool
condition Int -> Int
increment Int -> Int -> f a
f = Int -> f ()
go Int
initial
where
go :: Int -> f ()
go !Int
step
| Int -> Bool
condition Int
step =
let !next :: Int
next = Int -> Int
increment Int
step
in Int -> Int -> f a
f Int
step Int
next forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> f ()
go Int
next
| Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINE loopNextA_ #-}
loopNextM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> Int -> a -> m a) -> m a
loopNextM :: forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool)
-> (Int -> Int)
-> a
-> (Int -> Int -> a -> m a)
-> m a
loopNextM !Int
initial Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> Int -> a -> m a
f = Int -> a -> m a
go Int
initial a
initAcc
where
go :: Int -> a -> m a
go !Int
step !a
acc
| Int -> Bool
condition Int
step =
let !next :: Int
next = Int -> Int
increment Int
step
in Int -> Int -> a -> m a
f Int
step Int
next a
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> a -> m a
go Int
next
| Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
{-# INLINE loopNextM #-}
loopA_ :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ :: forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ !Int
initial Int -> Bool
condition Int -> Int
increment Int -> f a
f =
forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF Int
initial Int -> Bool
condition Int -> Int
increment (forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (\Int
i f ()
ma -> Int -> f a
f Int
i forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> f ()
ma)
{-# INLINE loopA_ #-}
loopA :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> f b -> (Int -> f (b -> b)) -> f b
loopA :: forall (f :: * -> *) b.
Applicative f =>
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f b
-> (Int -> f (b -> b))
-> f b
loopA !Int
initial Int -> Bool
condition Int -> Int
increment f b
lastAction Int -> f (b -> b)
f =
forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF Int
initial Int -> Bool
condition Int -> Int
increment f b
lastAction (\Int
i f b
ma -> Int -> f (b -> b)
f Int
i forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> f b
ma)
{-# INLINE loopA #-}
loopF :: Int -> (Int -> Bool) -> (Int -> Int) -> f a -> (Int -> f a -> f a) -> f a
loopF :: forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF !Int
initial Int -> Bool
condition Int -> Int
increment f a
lastAction Int -> f a -> f a
f = Int -> f a
go Int
initial
where
go :: Int -> f a
go !Int
step
| Int -> Bool
condition Int
step = Int -> f a -> f a
f Int
step (Int -> f a
go (Int -> Int
increment Int
step))
| Bool
otherwise = f a
lastAction
{-# INLINE loopF #-}
nextMaybeF :: Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
nextMaybeF :: forall (f :: * -> *) a.
Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a
nextMaybeF !Int
cur Int -> Bool
condition Int -> Int
increment Maybe Int -> f a
f =
let !i :: Int
i = Int -> Int
increment Int
cur
in Maybe Int -> f a
f forall a b. (a -> b) -> a -> b
$! if Int -> Bool
condition Int
i then forall a. a -> Maybe a
Just Int
i else forall a. Maybe a
Nothing
{-# INLINE nextMaybeF #-}
loopDeepM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM :: forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopDeepM !Int
initial Int -> Bool
condition Int -> Int
increment !a
initAcc Int -> a -> m a
f =
forall (f :: * -> *) a.
Int
-> (Int -> Bool)
-> (Int -> Int)
-> f a
-> (Int -> f a -> f a)
-> f a
loopF Int
initial Int -> Bool
condition Int -> Int
increment (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
initAcc) (\Int
i m a
ma -> m a
ma forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> a -> m a
f Int
i)
{-# INLINE loopDeepM #-}
splitLinearly
:: Int
-> Int
-> (Int -> Int -> a)
-> a
splitLinearly :: forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly Int
numChunks Int
totalLength Int -> Int -> a
action = Int -> Int -> a
action Int
chunkLength Int
slackStart
where
!chunkLength :: Int
chunkLength = Int
totalLength forall a. Integral a => a -> a -> a
`quot` Int
numChunks
!slackStart :: Int
slackStart = Int
chunkLength forall a. Num a => a -> a -> a
* Int
numChunks
{-# INLINE splitLinearly #-}
splitLinearlyM_
:: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ :: forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ Scheduler s ()
scheduler Int
totalLength Int -> Int -> m ()
action =
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f ()
loopNextA_ Int
0 (forall a. Ord a => a -> a -> Bool
< Int
slackStart) (forall a. Num a => a -> a -> a
+ Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \Int
start Int
next ->
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Int -> Int -> m ()
action Int
start Int
next
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart forall a. Ord a => a -> a -> Bool
< Int
totalLength) forall a b. (a -> b) -> a -> b
$
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
Int -> Int -> m ()
action Int
slackStart Int
totalLength
{-# INLINE splitLinearlyM_ #-}
splitLinearlyM
:: MonadPrimBase s m => Scheduler s a -> Int -> (Int -> Int -> m a) -> m ()
splitLinearlyM :: forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> Int -> (Int -> Int -> m a) -> m ()
splitLinearlyM Scheduler s a
scheduler Int
totalLength Int -> Int -> m a
action =
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler) Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f ()
loopNextA_ Int
0 (forall a. Ord a => a -> a -> Bool
< Int
slackStart) (forall a. Num a => a -> a -> a
+ Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \Int
start Int
next ->
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s a
scheduler (Int -> Int -> m a
action Int
start Int
next)
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart forall a. Ord a => a -> a -> Bool
< Int
totalLength) forall a b. (a -> b) -> a -> b
$
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s a
scheduler (Int -> Int -> m a
action Int
slackStart Int
totalLength)
{-# INLINE splitLinearlyM #-}
splitLinearlyWith_
:: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ :: forall s (m :: * -> *) b.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler s ()
scheduler Int
totalLength Int -> b
index =
forall s (m :: * -> *) b c.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithM_ Scheduler s ()
scheduler Int
totalLength (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> b
index)
{-# INLINE splitLinearlyWith_ #-}
splitLinearlyWithM_
:: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithM_ :: forall s (m :: * -> *) b c.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithM_ Scheduler s ()
scheduler Int
totalLength Int -> m b
make Int -> b -> m c
write =
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ Scheduler s ()
scheduler Int
totalLength Int -> Int -> m ()
go
where
go :: Int -> Int -> m ()
go Int
start Int
end = forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
start (forall a. Ord a => a -> a -> Bool
< Int
end) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ \Int
k -> Int -> m b
make Int
k forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
write Int
k
{-# INLINE go #-}
{-# INLINE splitLinearlyWithM_ #-}
splitLinearlyWithStartAtM_
:: MonadPrimBase s m => Scheduler s () -> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStartAtM_ :: forall s (m :: * -> *) b c.
MonadPrimBase s m =>
Scheduler s ()
-> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStartAtM_ Scheduler s ()
scheduler Int
startAt Int
totalLength Int -> m b
make Int -> b -> m c
write =
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
startAt (forall a. Ord a => a -> a -> Bool
< (Int
slackStart forall a. Num a => a -> a -> a
+ Int
startAt)) (forall a. Num a => a -> a -> a
+ Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
start (forall a. Ord a => a -> a -> Bool
< (Int
start forall a. Num a => a -> a -> a
+ Int
chunkLength)) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$
\ !Int
k -> Int -> m b
make Int
k forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
write Int
k
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart forall a. Ord a => a -> a -> Bool
< Int
totalLength) forall a b. (a -> b) -> a -> b
$
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ (Int
slackStart forall a. Num a => a -> a -> a
+ Int
startAt) (forall a. Ord a => a -> a -> Bool
< (Int
totalLength forall a. Num a => a -> a -> a
+ Int
startAt)) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$
\ !Int
k -> Int -> m b
make Int
k forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
write Int
k
{-# INLINE splitLinearlyWithStartAtM_ #-}
splitLinearlyWithStatefulM_
:: MonadUnliftIO m
=> SchedulerWS ws ()
-> Int
-> (Int -> ws -> m b)
-> (Int -> b -> m c)
-> m ()
splitLinearlyWithStatefulM_ :: forall (m :: * -> *) ws b c.
MonadUnliftIO m =>
SchedulerWS ws ()
-> Int -> (Int -> ws -> m b) -> (Int -> b -> m c) -> m ()
splitLinearlyWithStatefulM_ SchedulerWS ws ()
schedulerWS Int
totalLength Int -> ws -> m b
make Int -> b -> m c
store =
let nWorkers :: Int
nWorkers = forall s a. Scheduler s a -> Int
numWorkers (forall ws a. SchedulerWS ws a -> Scheduler RealWorld a
unwrapSchedulerWS SchedulerWS ws ()
schedulerWS)
in forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly Int
nWorkers Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (forall a. Ord a => a -> a -> Bool
< Int
slackStart) (forall a. Num a => a -> a -> a
+ Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
forall (m :: * -> *) ws.
MonadPrimBase RealWorld m =>
SchedulerWS ws () -> (ws -> m ()) -> m ()
scheduleWorkState_ SchedulerWS ws ()
schedulerWS forall a b. (a -> b) -> a -> b
$ \ws
s ->
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
start (forall a. Ord a => a -> a -> Bool
< (Int
start forall a. Num a => a -> a -> a
+ Int
chunkLength)) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ \ !Int
k ->
forall a. m a -> IO a
run (Int -> ws -> m b
make Int
k ws
s forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
store Int
k)
forall (m :: * -> *) ws.
MonadPrimBase RealWorld m =>
SchedulerWS ws () -> (ws -> m ()) -> m ()
scheduleWorkState_ SchedulerWS ws ()
schedulerWS forall a b. (a -> b) -> a -> b
$ \ws
s ->
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
slackStart (forall a. Ord a => a -> a -> Bool
< Int
totalLength) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ \ !Int
k ->
forall a. m a -> IO a
run (Int -> ws -> m b
make Int
k ws
s forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> b -> m c
store Int
k)
{-# INLINE splitLinearlyWithStatefulM_ #-}
splitWorkWithFactorST
:: Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST :: forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s a
scheduler Int
start Int
step Int
totalLength b
initAcc b -> ST s (b, b)
splitAcc Int -> Int -> Int -> Int -> b -> ST s a
action = do
let !(Int
chunkLength, Int
slackStart) = Int -> Int -> Int -> (Int, Int)
splitNumChunks Int
fact (forall s a. Scheduler s a -> Int
numWorkers Scheduler s a
scheduler) Int
totalLength
b
slackAcc <-
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
0 (forall a. Ord a => a -> a -> Bool
< Int
slackStart) (forall a. Num a => a -> a -> a
+ Int
chunkLength) b
initAcc forall a b. (a -> b) -> a -> b
$ \ !Int
chunkStart !b
acc -> do
(b
accCur, b
accNext) <- b -> ST s (b, b)
splitAcc b
acc
forall (m :: * -> *) a.
PrimBase m =>
Scheduler (PrimState m) a -> m a -> m ()
scheduleMassivWork Scheduler s a
scheduler forall a b. (a -> b) -> a -> b
$ do
let !chunkStartAdj :: Int
chunkStartAdj = Int
start forall a. Num a => a -> a -> a
+ Int
chunkStart forall a. Num a => a -> a -> a
* Int
step
!chunkStopAdj :: Int
chunkStopAdj = Int
chunkStartAdj forall a. Num a => a -> a -> a
+ Int
chunkLength forall a. Num a => a -> a -> a
* Int
step
Int -> Int -> Int -> Int -> b -> ST s a
action Int
chunkStart Int
chunkLength Int
chunkStartAdj Int
chunkStopAdj b
accCur
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
accNext
let !slackLength :: Int
slackLength = Int
totalLength forall a. Num a => a -> a -> a
- Int
slackStart
if Int
slackLength forall a. Ord a => a -> a -> Bool
> Int
0
then do
(b
curAcc, b
nextAcc) <- b -> ST s (b, b)
splitAcc b
slackAcc
forall (m :: * -> *) a.
PrimBase m =>
Scheduler (PrimState m) a -> m a -> m ()
scheduleMassivWork Scheduler s a
scheduler forall a b. (a -> b) -> a -> b
$ do
let !slackStartAdj :: Int
slackStartAdj = Int
start forall a. Num a => a -> a -> a
+ Int
slackStart forall a. Num a => a -> a -> a
* Int
step
!slackStopAdj :: Int
slackStopAdj = Int
slackStartAdj forall a. Num a => a -> a -> a
+ Int
slackLength forall a. Num a => a -> a -> a
* Int
step
Int -> Int -> Int -> Int -> b -> ST s a
action Int
slackStart Int
slackLength Int
slackStartAdj Int
slackStopAdj b
curAcc
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
nextAcc
else forall (f :: * -> *) a. Applicative f => a -> f a
pure b
slackAcc
{-# INLINE splitWorkWithFactorST #-}
iterLinearST_
:: Int
-> Scheduler s ()
-> Int
-> Int
-> Int
-> (Int -> ST s a)
-> ST s ()
iterLinearST_ :: forall s a.
Int
-> Scheduler s ()
-> Int
-> Int
-> Int
-> (Int -> ST s a)
-> ST s ()
iterLinearST_ Int
fact Scheduler s ()
scheduler Int
start Int
step Int
n Int -> ST s a
action = do
let totalLength :: Int
totalLength = (Int
n forall a. Num a => a -> a -> a
- Int
start) forall a. Integral a => a -> a -> a
`quot` Int
step
forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s ()
scheduler Int
start Int
step Int
totalLength () (\()
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ((), ())) forall a b. (a -> b) -> a -> b
$
\Int
_ Int
_ Int
chunkStartAdj Int
chunkStopAdj ()
_ ->
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (forall a. Num a => a -> a -> a
+ Int
step) Int -> ST s a
action
{-# INLINE iterLinearST_ #-}
iterLinearAccST_
:: Int
-> Scheduler s ()
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> a -> ST s a)
-> ST s ()
iterLinearAccST_ :: forall s a.
Int
-> Scheduler s ()
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> a -> ST s a)
-> ST s ()
iterLinearAccST_ Int
fact Scheduler s ()
scheduler Int
start Int
step Int
n a
initAcc a -> ST s (a, a)
splitAcc Int -> a -> ST s a
action = do
let totalLength :: Int
totalLength = (Int
n forall a. Num a => a -> a -> a
- Int
start) forall a. Integral a => a -> a -> a
`quot` Int
step
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$
forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s ()
scheduler Int
start Int
step Int
totalLength a
initAcc a -> ST s (a, a)
splitAcc forall a b. (a -> b) -> a -> b
$
\Int
_ Int
_ Int
chunkStartAdj Int
chunkStopAdj a
accCur ->
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (forall a. Num a => a -> a -> a
+ Int
step) a
accCur Int -> a -> ST s a
action
{-# INLINE iterLinearAccST_ #-}
iterLinearAccST
:: Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> a -> ST s a)
-> ST s a
iterLinearAccST :: forall s a.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> a
-> (a -> ST s (a, a))
-> (Int -> a -> ST s a)
-> ST s a
iterLinearAccST Int
fact Scheduler s a
scheduler Int
start Int
step Int
n a
initAcc a -> ST s (a, a)
splitAcc Int -> a -> ST s a
action = do
let totalLength :: Int
totalLength = (Int
n forall a. Num a => a -> a -> a
- Int
start) forall a. Integral a => a -> a -> a
`quot` Int
step
forall s a b.
Int
-> Scheduler s a
-> Int
-> Int
-> Int
-> b
-> (b -> ST s (b, b))
-> (Int -> Int -> Int -> Int -> b -> ST s a)
-> ST s b
splitWorkWithFactorST Int
fact Scheduler s a
scheduler Int
start Int
step Int
totalLength a
initAcc a -> ST s (a, a)
splitAcc forall a b. (a -> b) -> a -> b
$
\Int
_ Int
_ Int
chunkStartAdj Int
chunkStopAdj a
accCur ->
forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
chunkStartAdj (forall a. Ord a => a -> a -> Bool
< Int
chunkStopAdj) (forall a. Num a => a -> a -> a
+ Int
step) a
accCur Int -> a -> ST s a
action
{-# INLINE iterLinearAccST #-}
splitNumChunks :: Int -> Int -> Int -> (Int, Int)
splitNumChunks :: Int -> Int -> Int -> (Int, Int)
splitNumChunks Int
fact Int
nw Int
totalLength =
let maxNumChunks :: Int
maxNumChunks = Int
nw forall a. Num a => a -> a -> a
* forall a. Ord a => a -> a -> a
max Int
1 Int
fact
!numChunks :: Int
numChunks
| Int
nw forall a. Eq a => a -> a -> Bool
== Int
1 Bool -> Bool -> Bool
|| Int
totalLength forall a. Ord a => a -> a -> Bool
<= Int
0 = Int
1
| Int
totalLength forall a. Ord a => a -> a -> Bool
<= Int
nw = Int
totalLength
| Int
totalLength forall a. Ord a => a -> a -> Bool
>= Int
maxNumChunks = Int
maxNumChunks
| Bool
otherwise = Int
nw
!chunkLength :: Int
chunkLength = Int
totalLength forall a. Integral a => a -> a -> a
`quot` Int
numChunks
!slackStart :: Int
slackStart = Int
chunkLength forall a. Num a => a -> a -> a
* Int
numChunks
in (Int
chunkLength, Int
slackStart)
stepStartAdjust :: Int -> Int -> Int
stepStartAdjust :: Int -> Int -> Int
stepStartAdjust Int
step Int
ix = Int
ix forall a. Num a => a -> a -> a
+ ((Int
step forall a. Num a => a -> a -> a
- (Int
ix forall a. Integral a => a -> a -> a
`mod` Int
step)) forall a. Integral a => a -> a -> a
`mod` Int
step)
{-# INLINE stepStartAdjust #-}
scheduleMassivWork :: PrimBase m => Scheduler (PrimState m) a -> m a -> m ()
scheduleMassivWork :: forall (m :: * -> *) a.
PrimBase m =>
Scheduler (PrimState m) a -> m a -> m ()
scheduleMassivWork = forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork
{-# INLINE [0] scheduleMassivWork #-}
{-# RULES
"scheduleWork/scheduleWork_/ST" forall (scheduler :: Scheduler s ()) (action :: ST s ()). scheduleMassivWork scheduler action = scheduleWork_ scheduler action
"scheduleWork/scheduleWork_/IO" forall (scheduler :: Scheduler RealWorld ()) (action :: IO ()). scheduleMassivWork scheduler action = scheduleWork_ scheduler action
#-}
withMassivScheduler_ :: Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ :: Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ Comp
comp Scheduler RealWorld () -> IO ()
f =
case Comp
comp of
Comp
Par -> forall (m :: * -> *) a.
MonadUnliftIO m =>
GlobalScheduler -> (Scheduler RealWorld () -> m a) -> m ()
withGlobalScheduler_ GlobalScheduler
globalScheduler Scheduler RealWorld () -> IO ()
f
Comp
Seq -> Scheduler RealWorld () -> IO ()
f forall s. Scheduler s ()
trivialScheduler_
Comp
_ -> forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m ()
withScheduler_ Comp
comp Scheduler RealWorld () -> IO ()
f
{-# INLINE withMassivScheduler_ #-}