{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE PatternSynonyms #-}
module Data.Massiv.Core.Index.Iterator (
Iterator (..),
iterTargetAccST,
iterTargetAccST_,
iterTargetFullWithStrideAccST,
iterTargetFullWithStrideAccST_,
iterTargetST_,
iterTargetFullWithStrideST_,
RowMajor (RowMajor),
defRowMajor,
RowMajorLinear (RowMajorLinear),
defRowMajorLinear,
RowMajorUnbalanced (RowMajorUnbalanced),
defRowMajorUnbalanced,
) where
import Control.Monad
import Control.Monad.ST
import Control.Scheduler
import Data.Massiv.Core.Index.Internal
import Data.Massiv.Core.Index.Stride
import Data.Massiv.Core.Loop
class Iterator it where
{-# MINIMAL (iterTargetM, iterTargetA_, iterTargetWithStrideAccST, iterTargetWithStrideAccST_) #-}
iterTargetA_
:: (Index ix, Applicative f)
=> it
-> Int
-> Sz ix
-> ix
-> Stride ix
-> (Ix1 -> ix -> f a)
-> f ()
iterTargetM
:: (Index ix, Monad m)
=> it
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
iterTargetWithStrideAccST
:: Index ix
=> it
-> Scheduler s a
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
iterTargetWithStrideAccST_
:: Index ix
=> it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterFullM
:: (Index ix, Monad m)
=> it
-> ix
-> Sz ix
-> a
-> (ix -> a -> m a)
-> m a
iterFullM it
it ix
start Sz ix
sz a
acc ix -> a -> m a
f =
forall it ix (m :: * -> *) a.
(Iterator it, Index ix, Monad m) =>
it
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterTargetM it
it Int
0 Sz ix
sz ix
start forall ix. Index ix => Stride ix
oneStride a
acc (forall a b. a -> b -> a
const ix -> a -> m a
f)
{-# INLINE iterFullM #-}
iterFullA_
:: (Index ix, Applicative f)
=> it
-> ix
-> Sz ix
-> (ix -> f a)
-> f ()
iterFullA_ it
it ix
start Sz ix
sz ix -> f a
f =
forall it ix (f :: * -> *) a.
(Iterator it, Index ix, Applicative f) =>
it -> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> f a) -> f ()
iterTargetA_ it
it Int
0 Sz ix
sz ix
start forall ix. Index ix => Stride ix
oneStride (forall a b. a -> b -> a
const ix -> f a
f)
{-# INLINE iterFullA_ #-}
iterFullAccST
:: Index ix
=> it
-> Scheduler s a
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> ST s a
iterFullAccST it
it Scheduler s a
scheduler ix
start Sz ix
sz a
acc a -> ST s (a, a)
splitAcc ix -> a -> ST s a
f =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s a
-> Int
-> Sz ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetAccST it
it Scheduler s a
scheduler Int
0 Sz ix
sz ix
start a
acc a -> ST s (a, a)
splitAcc (forall a b. a -> b -> a
const ix -> a -> ST s a
f)
{-# INLINE iterFullAccST #-}
iterTargetFullAccST
:: Index ix
=> it
-> Scheduler s a
-> Ix1
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
iterTargetFullAccST it
it Scheduler s a
scheduler Int
iStart Sz ix
sz =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s a
-> Int
-> Sz ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetFullWithStrideAccST it
it Scheduler s a
scheduler Int
iStart Sz ix
sz forall ix. Index ix => Stride ix
oneStride
{-# INLINE iterTargetFullAccST #-}
iterTargetFullAccST_
:: Index ix
=> it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterTargetFullAccST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetFullWithStrideAccST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz forall ix. Index ix => Stride ix
oneStride
{-# INLINE iterTargetFullAccST_ #-}
iterTargetFullST_
:: Index ix
=> it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> (Ix1 -> ix -> ST s ())
-> ST s ()
iterTargetFullST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz =
forall it ix s.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> (Int -> ix -> ST s ())
-> ST s ()
iterTargetST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz (forall ix. Index ix => Int -> ix
pureIndex Int
0)
{-# INLINE iterTargetFullST_ #-}
iterTargetWithStrideST_
:: Index ix
=> it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> (Ix1 -> ix -> ST s a)
-> ST s ()
iterTargetWithStrideST_ it
it Scheduler s ()
scheduler Int
i Sz ix
sz ix
ix Stride ix
stride Int -> ix -> ST s a
action =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetWithStrideAccST_ it
it Scheduler s ()
scheduler Int
i Sz ix
sz ix
ix Stride ix
stride () forall (m :: * -> *). Applicative m => () -> m ((), ())
noSplit forall a b. (a -> b) -> a -> b
$ \Int
j ix
jx ()
_ ->
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ Int -> ix -> ST s a
action Int
j ix
jx
{-# INLINE iterTargetWithStrideST_ #-}
newtype RowMajor = RowMajorInternal Int
defRowMajor :: RowMajor
defRowMajor :: RowMajor
defRowMajor = Int -> RowMajor
RowMajorInternal Int
8
pattern RowMajor
:: Int
-> RowMajor
pattern $bRowMajor :: Int -> RowMajor
$mRowMajor :: forall {r}. RowMajor -> (Int -> r) -> ((# #) -> r) -> r
RowMajor f <- RowMajorInternal f
where
RowMajor = Int -> RowMajor
RowMajorInternal forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => a -> a -> a
max Int
1
{-# COMPLETE RowMajor #-}
instance Iterator RowMajor where
iterFullM :: forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
RowMajor -> ix -> Sz ix -> a -> (ix -> a -> m a) -> m a
iterFullM RowMajor
_ ix
start (Sz ix
sz) = forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM ix
start ix
sz (forall ix. Index ix => Int -> ix
pureIndex Int
1) forall a. Ord a => a -> a -> Bool
(<)
{-# INLINE iterFullM #-}
iterFullA_ :: forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
RowMajor -> ix -> Sz ix -> (ix -> f a) -> f ()
iterFullA_ RowMajor
_ ix
start (Sz ix
sz) = forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> f a) -> f ()
iterA_ ix
start ix
sz (forall ix. Index ix => Int -> ix
pureIndex Int
1) forall a. Ord a => a -> a -> Bool
(<)
{-# INLINE iterFullA_ #-}
iterFullAccST :: forall ix s a.
Index ix =>
RowMajor
-> Scheduler s a
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> ST s a
iterFullAccST (RowMajorInternal Int
fact) Scheduler s a
scheduler ix
startIx =
forall ix s a.
Index ix =>
Int
-> Scheduler s a
-> ix
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> ST s a
iterRowMajorST Int
fact Scheduler s a
scheduler ix
startIx (forall ix. Index ix => Int -> ix
pureIndex Int
1)
{-# INLINE iterFullAccST #-}
iterTargetA_ :: forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
RowMajor
-> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> f a) -> f ()
iterTargetA_ RowMajor
_ Int
i Sz ix
sz ix
start (Stride ix
stride) =
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
Int -> Int -> Sz ix -> ix -> ix -> (Int -> ix -> f a) -> f ()
iterTargetRowMajorA_ Int
0 Int
i Sz ix
sz ix
start ix
stride
{-# INLINE iterTargetA_ #-}
iterTargetM :: forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
RowMajor
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterTargetM RowMajor
_ Int
i Sz ix
sz ix
start (Stride ix
stride) =
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Int
-> Int -> Sz ix -> ix -> ix -> a -> (Int -> ix -> a -> m a) -> m a
iterTargetRowMajorAccM Int
0 Int
i Sz ix
sz ix
start ix
stride
{-# INLINE iterTargetM #-}
iterTargetWithStrideAccST :: forall ix s a.
Index ix =>
RowMajor
-> Scheduler s a
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetWithStrideAccST (RowMajor Int
fact) Scheduler s a
scheduler Int
i Sz ix
sz ix
ix (Stride ix
stride) =
forall ix s a.
Index ix =>
Int
-> Int
-> Scheduler s a
-> Int
-> Sz ix
-> ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetRowMajorAccST Int
0 Int
fact Scheduler s a
scheduler Int
i Sz ix
sz ix
ix ix
stride
{-# INLINE iterTargetWithStrideAccST #-}
iterTargetWithStrideAccST_ :: forall ix s a.
Index ix =>
RowMajor
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetWithStrideAccST_ (RowMajor Int
fact) Scheduler s ()
scheduler Int
i Sz ix
sz ix
ix (Stride ix
stride) =
forall ix s a.
Index ix =>
Int
-> Int
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetRowMajorAccST_ Int
0 Int
fact Scheduler s ()
scheduler Int
i Sz ix
sz ix
ix ix
stride
{-# INLINE iterTargetWithStrideAccST_ #-}
newtype RowMajorLinear = RowMajorLinear Int
defRowMajorLinear :: RowMajorLinear
defRowMajorLinear :: RowMajorLinear
defRowMajorLinear = Int -> RowMajorLinear
RowMajorLinear Int
8
instance Iterator RowMajorLinear where
iterTargetM :: forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
RowMajorLinear
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterTargetM RowMajorLinear
_ Int
iStart Sz ix
sz ix
start (Stride ix
stride) a
acc Int -> ix -> a -> m a
action =
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
< forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) (forall a. Num a => a -> a -> a
+ Int
1) a
acc forall a b. (a -> b) -> a -> b
$ \Int
i ->
Int -> ix -> a -> m a
action (Int
iStart forall a. Num a => a -> a -> a
+ Int
i) (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
start (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(*) ix
stride (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i)))
{-# INLINE iterTargetM #-}
iterTargetA_ :: forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
RowMajorLinear
-> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> f a) -> f ()
iterTargetA_ RowMajorLinear
_ Int
iStart Sz ix
sz ix
start (Stride ix
stride) Int -> ix -> f a
action =
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (forall a. Ord a => a -> a -> Bool
< forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ \Int
i ->
Int -> ix -> f a
action (Int
iStart forall a. Num a => a -> a -> a
+ Int
i) (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
start (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(*) ix
stride (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i)))
{-# INLINE iterTargetA_ #-}
iterTargetFullAccST :: forall ix s a.
Index ix =>
RowMajorLinear
-> Scheduler s a
-> Int
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetFullAccST RowMajorLinear
it Scheduler s a
scheduler Int
iStart Sz ix
sz a
acc a -> ST s (a, a)
splitAcc Int -> ix -> a -> ST s a
action =
let !(RowMajorLinear Int
fact) = RowMajorLinear
it
in 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
iStart Int
1 (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) a
acc a -> ST s (a, a)
splitAcc forall a b. (a -> b) -> a -> b
$ \ !Int
i ->
Int -> ix -> a -> ST s a
action Int
i (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i)
{-# INLINE iterTargetFullAccST #-}
iterTargetFullAccST_ :: forall ix s a.
Index ix =>
RowMajorLinear
-> Scheduler s ()
-> Int
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetFullAccST_ RowMajorLinear
it Scheduler s ()
scheduler Int
iStart Sz ix
sz a
acc a -> ST s (a, a)
splitAcc Int -> ix -> a -> ST s a
action =
let !(RowMajorLinear Int
fact) = RowMajorLinear
it
in 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
iStart Int
1 (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) a
acc a -> ST s (a, a)
splitAcc forall a b. (a -> b) -> a -> b
$ \ !Int
i ->
Int -> ix -> a -> ST s a
action Int
i (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i)
{-# INLINE iterTargetFullAccST_ #-}
iterTargetFullST_ :: forall ix s.
Index ix =>
RowMajorLinear
-> Scheduler s ()
-> Int
-> Sz ix
-> (Int -> ix -> ST s ())
-> ST s ()
iterTargetFullST_ RowMajorLinear
it Scheduler s ()
scheduler Int
iStart Sz ix
sz Int -> ix -> ST s ()
action =
let !(RowMajorLinear Int
fact) = RowMajorLinear
it
in forall s a.
Int
-> Scheduler s ()
-> Int
-> Int
-> Int
-> (Int -> ST s a)
-> ST s ()
iterLinearST_ Int
fact Scheduler s ()
scheduler Int
iStart Int
1 (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) forall a b. (a -> b) -> a -> b
$ \ !Int
i ->
Int -> ix -> ST s ()
action Int
i (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i)
{-# INLINE iterTargetFullST_ #-}
iterTargetWithStrideAccST :: forall ix s a.
Index ix =>
RowMajorLinear
-> Scheduler s a
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetWithStrideAccST RowMajorLinear
it Scheduler s a
scheduler Int
iStart Sz ix
sz ix
start (Stride ix
stride) a
acc a -> ST s (a, a)
spliAcc Int -> ix -> a -> ST s a
action =
let RowMajorLinear Int
fact = RowMajorLinear
it
in 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
0 Int
1 (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) a
acc a -> ST s (a, a)
spliAcc forall a b. (a -> b) -> a -> b
$ \Int
i ->
Int -> ix -> a -> ST s a
action (Int
iStart forall a. Num a => a -> a -> a
+ Int
i) forall a b. (a -> b) -> a -> b
$
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
start (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(*) ix
stride (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i))
{-# INLINE iterTargetWithStrideAccST #-}
iterTargetWithStrideAccST_ :: forall ix s a.
Index ix =>
RowMajorLinear
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetWithStrideAccST_ RowMajorLinear
it Scheduler s ()
scheduler Int
iStart Sz ix
sz ix
start (Stride ix
stride) a
acc a -> ST s (a, a)
spliAcc Int -> ix -> a -> ST s a
action =
let RowMajorLinear Int
fact = RowMajorLinear
it
in 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
0 Int
1 (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) a
acc a -> ST s (a, a)
spliAcc forall a b. (a -> b) -> a -> b
$ \Int
i ->
Int -> ix -> a -> ST s a
action (Int
iStart forall a. Num a => a -> a -> a
+ Int
i) forall a b. (a -> b) -> a -> b
$
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
start (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(*) ix
stride (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i))
{-# INLINE iterTargetWithStrideAccST_ #-}
newtype RowMajorUnbalanced = RowMajorUnbalancedInternal Int
defRowMajorUnbalanced :: RowMajorUnbalanced
defRowMajorUnbalanced :: RowMajorUnbalanced
defRowMajorUnbalanced = Int -> RowMajorUnbalanced
RowMajorUnbalancedInternal Int
8
pattern RowMajorUnbalanced
:: Int
-> RowMajorUnbalanced
pattern $bRowMajorUnbalanced :: Int -> RowMajorUnbalanced
$mRowMajorUnbalanced :: forall {r}. RowMajorUnbalanced -> (Int -> r) -> ((# #) -> r) -> r
RowMajorUnbalanced f <- RowMajorUnbalancedInternal f
where
RowMajorUnbalanced = Int -> RowMajorUnbalanced
RowMajorUnbalancedInternal forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => a -> a -> a
max Int
1
{-# COMPLETE RowMajorUnbalanced #-}
instance Iterator RowMajorUnbalanced where
iterFullM :: forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
RowMajorUnbalanced -> ix -> Sz ix -> a -> (ix -> a -> m a) -> m a
iterFullM (RowMajorUnbalanced Int
fact) = forall it ix (m :: * -> *) a.
(Iterator it, Index ix, Monad m) =>
it -> ix -> Sz ix -> a -> (ix -> a -> m a) -> m a
iterFullM (Int -> RowMajor
RowMajor Int
fact)
{-# INLINE iterFullM #-}
iterFullA_ :: forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
RowMajorUnbalanced -> ix -> Sz ix -> (ix -> f a) -> f ()
iterFullA_ (RowMajorUnbalanced Int
fact) = forall it ix (f :: * -> *) a.
(Iterator it, Index ix, Applicative f) =>
it -> ix -> Sz ix -> (ix -> f a) -> f ()
iterFullA_ (Int -> RowMajor
RowMajor Int
fact)
{-# INLINE iterFullA_ #-}
iterTargetM :: forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
RowMajorUnbalanced
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterTargetM (RowMajorUnbalanced Int
fact) = forall it ix (m :: * -> *) a.
(Iterator it, Index ix, Monad m) =>
it
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterTargetM (Int -> RowMajor
RowMajor Int
fact)
{-# INLINE iterTargetM #-}
iterTargetA_ :: forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
RowMajorUnbalanced
-> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> f a) -> f ()
iterTargetA_ (RowMajorUnbalanced Int
fact) = forall it ix (f :: * -> *) a.
(Iterator it, Index ix, Applicative f) =>
it -> Int -> Sz ix -> ix -> Stride ix -> (Int -> ix -> f a) -> f ()
iterTargetA_ (Int -> RowMajor
RowMajor Int
fact)
{-# INLINE iterTargetA_ #-}
iterTargetWithStrideAccST :: forall ix s a.
Index ix =>
RowMajorUnbalanced
-> Scheduler s a
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetWithStrideAccST = forall ix a t s b.
Index ix =>
(Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> t) -> ST s b)
-> RowMajorUnbalanced
-> Scheduler s b
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> t)
-> ST s a
iterUnbalancedTargetWithStride forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM
{-# INLINE iterTargetWithStrideAccST #-}
iterTargetWithStrideAccST_ :: forall ix s a.
Index ix =>
RowMajorUnbalanced
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetWithStrideAccST_ RowMajorUnbalanced
it Scheduler s ()
scheduler Int
iStart Sz ix
sz ix
start Stride ix
stride a
acc a -> ST s (a, a)
splitAcc' Int -> ix -> a -> ST s a
action =
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$
forall ix a t s b.
Index ix =>
(Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> t) -> ST s b)
-> RowMajorUnbalanced
-> Scheduler s b
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> t)
-> ST s a
iterUnbalancedTargetWithStride forall {f :: * -> *} {a}.
Monad f =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> f a) -> f ()
innerLoop RowMajorUnbalanced
it Scheduler s ()
scheduler Int
iStart Sz ix
sz ix
start Stride ix
stride a
acc a -> ST s (a, a)
splitAcc' Int -> ix -> a -> ST s a
action
where
innerLoop :: Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> f a) -> f ()
innerLoop Int
initial Int -> Bool
condition Int -> Int
increment a
initAcc Int -> a -> f a
f =
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
initial Int -> Bool
condition Int -> Int
increment a
initAcc Int -> a -> f a
f
{-# INLINE innerLoop #-}
{-# INLINE iterTargetWithStrideAccST_ #-}
iterUnbalancedTargetWithStride
:: Index ix
=> (Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> t) -> ST s b)
-> RowMajorUnbalanced
-> Scheduler s b
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> t)
-> ST s a
iterUnbalancedTargetWithStride :: forall ix a t s b.
Index ix =>
(Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> t) -> ST s b)
-> RowMajorUnbalanced
-> Scheduler s b
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> t)
-> ST s a
iterUnbalancedTargetWithStride Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> t) -> ST s b
innerLoop RowMajorUnbalanced
it Scheduler s b
scheduler Int
iStart Sz ix
sz ix
start Stride ix
stride a
acc a -> ST s (a, a)
splitAcc Int -> ix -> t
action =
let RowMajorUnbalanced Int
fact = RowMajorUnbalanced
it
!n :: Int
n = forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
!step :: Int
step = forall a. Ord a => a -> a -> a
min (Int
fact forall a. Num a => a -> a -> a
* forall s a. Scheduler s a -> Int
numWorkers Scheduler s b
scheduler) Int
n
in 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
step) (forall a. Num a => a -> a -> a
+ Int
1) a
acc forall a b. (a -> b) -> a -> b
$ \ !Int
istep !a
a -> do
(a
curAcc, a
nextAcc) <- a -> ST s (a, a)
splitAcc a
a
forall (m :: * -> *) a.
PrimBase m =>
Scheduler (PrimState m) a -> m a -> m ()
scheduleMassivWork Scheduler s b
scheduler forall a b. (a -> b) -> a -> b
$
Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> t) -> ST s b
innerLoop Int
istep (forall a. Ord a => a -> a -> Bool
< Int
n) (forall a. Num a => a -> a -> a
+ Int
step) a
curAcc forall a b. (a -> b) -> a -> b
$ \Int
i ->
Int -> ix -> t
action (Int
iStart forall a. Num a => a -> a -> a
+ Int
i) forall a b. (a -> b) -> a -> b
$
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
start (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(*) (forall ix. Stride ix -> ix
unStride Stride ix
stride) (forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i))
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
nextAcc
{-# INLINE iterUnbalancedTargetWithStride #-}
noSplit :: Applicative m => () -> m ((), ())
noSplit :: forall (m :: * -> *). Applicative m => () -> m ((), ())
noSplit ()
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure ((), ())
iterTargetAccST
:: (Iterator it, Index ix)
=> it
-> Scheduler s a
-> Ix1
-> Sz ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
iterTargetAccST :: forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s a
-> Int
-> Sz ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetAccST it
it Scheduler s a
scheduler Int
iStart Sz ix
sz ix
ix =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s a
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetWithStrideAccST it
it Scheduler s a
scheduler Int
iStart Sz ix
sz ix
ix forall ix. Index ix => Stride ix
oneStride
{-# INLINE iterTargetAccST #-}
iterTargetAccST_
:: (Iterator it, Index ix)
=> it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterTargetAccST_ :: forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetAccST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz ix
ix =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetWithStrideAccST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz ix
ix forall ix. Index ix => Stride ix
oneStride
{-# INLINE iterTargetAccST_ #-}
iterTargetFullWithStrideST_
:: (Iterator it, Index ix)
=> it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> Stride ix
-> (Ix1 -> ix -> ST s ())
-> ST s ()
iterTargetFullWithStrideST_ :: forall it ix s.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> Stride ix
-> (Int -> ix -> ST s ())
-> ST s ()
iterTargetFullWithStrideST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> Stride ix
-> (Int -> ix -> ST s a)
-> ST s ()
iterTargetWithStrideST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz (forall ix. Index ix => Int -> ix
pureIndex Int
0)
{-# INLINE iterTargetFullWithStrideST_ #-}
iterTargetST_
:: (Iterator it, Index ix)
=> it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> (Ix1 -> ix -> ST s ())
-> ST s ()
iterTargetST_ :: forall it ix s.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> (Int -> ix -> ST s ())
-> ST s ()
iterTargetST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz ix
ix =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> Stride ix
-> (Int -> ix -> ST s a)
-> ST s ()
iterTargetWithStrideST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz ix
ix forall ix. Index ix => Stride ix
oneStride
{-# INLINE iterTargetST_ #-}
iterTargetFullWithStrideAccST
:: (Iterator it, Index ix)
=> it
-> Scheduler s a
-> Ix1
-> Sz ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
iterTargetFullWithStrideAccST :: forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s a
-> Int
-> Sz ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetFullWithStrideAccST it
it Scheduler s a
scheduler Int
iStart Sz ix
sz =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s a
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetWithStrideAccST it
it Scheduler s a
scheduler Int
iStart Sz ix
sz (forall ix. Index ix => Int -> ix
pureIndex Int
0)
{-# INLINE iterTargetFullWithStrideAccST #-}
iterTargetFullWithStrideAccST_
:: (Iterator it, Index ix)
=> it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterTargetFullWithStrideAccST_ :: forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetFullWithStrideAccST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz =
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s ()
iterTargetWithStrideAccST_ it
it Scheduler s ()
scheduler Int
iStart Sz ix
sz (forall ix. Index ix => Int -> ix
pureIndex Int
0)
{-# INLINE iterTargetFullWithStrideAccST_ #-}