{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE PatternSynonyms #-}

-- |
-- Module      : Data.Massiv.Core.Index.Iterator
-- Copyright   : (c) Alexey Kuleshevich 2021-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Core.Index.Iterator (
  Iterator (..),

  -- * Extra iterator functions
  iterTargetAccST,
  iterTargetAccST_,
  iterTargetFullWithStrideAccST,
  iterTargetFullWithStrideAccST_,
  iterTargetST_,
  iterTargetFullWithStrideST_,

  -- * Iterator implementations
  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_) #-}

  -- | Iterate over a target region using linear index with access to the source
  -- index, which adjusted according to the stride. Use `iterTargetM` if you
  -- need an accumulator.
  --
  -- @since 1.0.2
  iterTargetA_
    :: (Index ix, Applicative f)
    => it
    -> Int
    -- ^ Target linear index start
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Source start index
    -> Stride ix
    -- ^ Source stride
    -> (Ix1 -> ix -> f a)
    -- ^ Action that accepts a linear index of the target and multi-dimensional
    -- index of the source.
    -> f ()

  -- | Iterate over a target region using linear index with access to the source
  -- index, which adjusted according to the stride.
  --
  -- @since 1.0.2
  iterTargetM
    :: (Index ix, Monad m)
    => it
    -> Ix1
    -- ^ Target linear index start
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Source start index
    -> Stride ix
    -- ^ Source stride
    -> a
    -- ^ Accumulator
    -> (Ix1 -> ix -> a -> m a)
    -- ^ Action that accepts a linear index of the target,
    -- multi-dimensional index of the source and accumulator
    -> m a

  iterTargetWithStrideAccST
    :: Index ix
    => it
    -> Scheduler s a
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Source start index
    -> Stride ix
    -- ^ Source stride
    -> a
    -- ^ Initial accumulator
    -> (a -> ST s (a, a))
    -- ^ Splitting action that produces new accumulators for separate worker threads.
    -> (Ix1 -> ix -> a -> ST s a)
    -- ^ Action
    -> ST s a

  iterTargetWithStrideAccST_
    :: Index ix
    => it
    -> Scheduler s ()
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Start
    -> Stride ix
    -- ^ Stride
    -> a
    -- ^ Initial accumulator
    -> (a -> ST s (a, a))
    -- ^ Splitting action that produces new accumulators for separate worker threads.
    -> (Ix1 -> ix -> a -> ST s a)
    -- ^ Action
    -> ST s ()

  -- | Iterate over a region with a monadic action and accumulator.
  --
  -- @since 1.0.2
  iterFullM
    :: (Index ix, Monad m)
    => it
    -> ix
    -- ^ Source start index
    -> Sz ix
    -- ^ Source size
    -> a
    -- ^ Accumulator
    -> (ix -> a -> m a)
    -- ^ Action that accepts a linear index of the target,
    -- multi-dimensional index of the source and accumulator
    -> 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 #-}

  -- | Iterate over a region with an applicative action ignoring the result.
  --
  -- @since 1.0.2
  iterFullA_
    :: (Index ix, Applicative f)
    => it
    -> ix
    -- ^ Source start index
    -> Sz ix
    -- ^ Source size
    -> (ix -> f a)
    -- ^ Action that accepts a linear index of the target,
    -- multi-dimensional index of the source and accumulator
    -> 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_ #-}

  -- | Iterate over a region in a ST monad with access to `Scheduler`.
  iterFullAccST
    :: Index ix
    => it
    -- ^ Scheduler multiplying factor. Must be positive
    -> Scheduler s a
    -- ^ Scheduler to use
    -> ix
    -- ^ Start index
    -> Sz ix
    -- ^ Size
    -> a
    -- ^ Initial accumulator
    -> (a -> ST s (a, a))
    -- ^ Function that splits accumulator for each scheduled job.
    -> (ix -> a -> ST s a)
    -- ^ Action
    -> 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
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> a
    -- ^ Initial accumulator
    -> (a -> ST s (a, a))
    -- ^ Function that splits accumulator for each scheduled job.
    -> (Ix1 -> ix -> a -> ST s a)
    -- ^ Action
    -> 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 ()
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> a
    -- ^ Initial accumulator
    -> (a -> ST s (a, a))
    -- ^ Function that splits accumulator for each scheduled job.
    -> (Ix1 -> ix -> a -> ST s a)
    -- ^ Action
    -> 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 ()
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> (Ix1 -> ix -> ST s ())
    -- ^ Action
    -> 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_ #-}

  -- NOTE: this function does not have to be part of the class, but for some
  -- reason it creates a severe regression when moved outside.

  -- | Iterate over a target array with a stride without an accumulator
  iterTargetWithStrideST_
    :: Index ix
    => it
    -> Scheduler s ()
    -- ^ Scheduler to use
    -> Ix1
    -- ^ Target linear start index
    -> Sz ix
    -- ^ Target size
    -> ix
    -- ^ Start
    -> Stride ix
    -- ^ Stride
    -> (Ix1 -> ix -> ST s a)
    -- ^ Action
    -> 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_ #-}

-- | Default iterator that parallelizes work in linear chunks. Supplied factor
-- will be used to schedule that many jobs per capability.
--
-- @since 1.0.2
newtype RowMajor = RowMajorInternal Int

-- | Default row major iterator with multiplying factor set to @8@.
defRowMajor :: RowMajor
defRowMajor :: RowMajor
defRowMajor = Int -> RowMajor
RowMajorInternal Int
8

pattern RowMajor
  :: Int
  -- ^ Multiplier that will be used to scale number of jobs.
  -> 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_ #-}

-- | Parallelizing unbalanced computation (i.e. computing some elements of the
-- array is much more expensive then the others) it can be benefitial to
-- interleave iteration. Perfect example of this would be a ray tracer or the
-- Mandelbrot set.
--
-- iteration without parallelization is equivalent to `RowMajor`
--
-- @since 1.0.2
newtype RowMajorUnbalanced = RowMajorUnbalancedInternal Int

defRowMajorUnbalanced :: RowMajorUnbalanced
defRowMajorUnbalanced :: RowMajorUnbalanced
defRowMajorUnbalanced = Int -> RowMajorUnbalanced
RowMajorUnbalancedInternal Int
8

pattern RowMajorUnbalanced
  :: Int
  -- ^ Multiplier that will be used to scale number of jobs.
  -> 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
  -- ^ Scheduler to use
  -> Ix1
  -- ^ Target linear start index
  -> Sz ix
  -- ^ Target size
  -> ix
  -- ^ Source start
  -> a
  -> (a -> ST s (a, a))
  -> (Ix1 -> ix -> a -> ST s a)
  -- ^ Action
  -> 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 ()
  -- ^ Scheduler to use
  -> Ix1
  -- ^ Target linear start index
  -> Sz ix
  -- ^ Target size
  -> ix
  -- ^ Source start
  -> a
  -> (a -> ST s (a, a))
  -> (Ix1 -> ix -> a -> ST s a)
  -- ^ Action
  -> 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 ()
  -- ^ Scheduler to use
  -> Ix1
  -- ^ Target linear start index
  -> Sz ix
  -- ^ Target size
  -> Stride ix
  -- ^ Stride
  -> (Ix1 -> ix -> ST s ())
  -- ^ Action
  -> 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 ()
  -- ^ Scheduler to use
  -> Ix1
  -- ^ Target linear start index
  -> Sz ix
  -- ^ Target size
  -> ix
  -- ^ Start
  -> (Ix1 -> ix -> ST s ())
  -- ^ Action
  -> 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
  -- ^ Scheduler to use
  -> Ix1
  -- ^ Target linear start index
  -> Sz ix
  -- ^ Target size
  -> Stride ix
  -- ^ Stride
  -> a
  -> (a -> ST s (a, a))
  -> (Ix1 -> ix -> a -> ST s a)
  -- ^ Action
  -> 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 ()
  -- ^ Scheduler to use
  -> Ix1
  -- ^ Target linear start index
  -> Sz ix
  -- ^ Target size
  -> Stride ix
  -- ^ Stride
  -> a
  -> (a -> ST s (a, a))
  -> (Ix1 -> ix -> a -> ST s a)
  -- ^ Action
  -> 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_ #-}