{-# 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 =
    it
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
forall it ix (m :: * -> *) a.
(Iterator it, Index ix, Monad m) =>
it
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
iterTargetM it
it Ix1
0 Sz ix
sz ix
start Stride ix
forall ix. Index ix => Stride ix
oneStride a
acc ((ix -> a -> m a) -> Ix1 -> ix -> a -> m a
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 =
    it -> Ix1 -> Sz ix -> ix -> Stride ix -> (Ix1 -> ix -> f a) -> f ()
forall it ix (f :: * -> *) a.
(Iterator it, Index ix, Applicative f) =>
it -> Ix1 -> Sz ix -> ix -> Stride ix -> (Ix1 -> ix -> f a) -> f ()
iterTargetA_ it
it Ix1
0 Sz ix
sz ix
start Stride ix
forall ix. Index ix => Stride ix
oneStride ((ix -> f a) -> Ix1 -> ix -> f a
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 =
    it
-> Scheduler s a
-> Ix1
-> Sz ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
forall it ix s a.
(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 it
it Scheduler s a
scheduler Ix1
0 Sz ix
sz ix
start a
acc a -> ST s (a, a)
splitAcc ((ix -> a -> ST s a) -> Ix1 -> ix -> a -> ST s a
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 Ix1
iStart Sz ix
sz =
    it
-> Scheduler s a
-> Ix1
-> Sz ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
forall it ix s a.
(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 it
it Scheduler s a
scheduler Ix1
iStart Sz ix
sz Stride ix
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 Ix1
iStart Sz ix
sz =
    it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
forall it ix s a.
(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_ it
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz Stride ix
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 Ix1
iStart Sz ix
sz =
    it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> (Ix1 -> ix -> ST s ())
-> ST s ()
forall it ix s.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> (Ix1 -> ix -> ST s ())
-> ST s ()
iterTargetST_ it
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz (Ix1 -> ix
forall ix. Index ix => Ix1 -> ix
pureIndex Ix1
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 Ix1
i Sz ix
sz ix
ix Stride ix
stride Ix1 -> ix -> ST s a
action =
    it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> ()
-> (() -> ST s ((), ()))
-> (Ix1 -> ix -> () -> ST s ())
-> ST s ()
forall it ix s a.
(Iterator it, 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 ()
iterTargetWithStrideAccST_ it
it Scheduler s ()
scheduler Ix1
i Sz ix
sz ix
ix Stride ix
stride () () -> ST s ((), ())
forall (m :: * -> *). Applicative m => () -> m ((), ())
noSplit ((Ix1 -> ix -> () -> ST s ()) -> ST s ())
-> (Ix1 -> ix -> () -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Ix1
j ix
jx ()
_ ->
      ST s a -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s a -> ST s ()) -> ST s a -> ST s ()
forall a b. (a -> b) -> a -> b
$ Ix1 -> ix -> ST s a
action Ix1
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 = Ix1 -> RowMajor
RowMajorInternal Ix1
8

pattern RowMajor :: Int
                 -- ^ Multiplier that will be used to scale number of jobs.
                 -> RowMajor
pattern $bRowMajor :: Ix1 -> RowMajor
$mRowMajor :: forall r. RowMajor -> (Ix1 -> r) -> (Void# -> r) -> r
RowMajor f <- RowMajorInternal f
  where RowMajor = Ix1 -> RowMajor
RowMajorInternal (Ix1 -> RowMajor) -> (Ix1 -> Ix1) -> Ix1 -> RowMajor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix1 -> Ix1 -> Ix1
forall a. Ord a => a -> a -> a
max Ix1
1
{-# COMPLETE RowMajor #-}

instance Iterator RowMajor where
  iterFullM :: RowMajor -> ix -> Sz ix -> a -> (ix -> a -> m a) -> m a
iterFullM RowMajor
_ ix
start (Sz ix
sz) = ix
-> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> a -> (ix -> a -> m a) -> m a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM ix
start ix
sz (Ix1 -> ix
forall ix. Index ix => Ix1 -> ix
pureIndex Ix1
1) Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
(<)
  {-# INLINE iterFullM #-}
  iterFullA_ :: RowMajor -> ix -> Sz ix -> (ix -> f a) -> f ()
iterFullA_ RowMajor
_ ix
start (Sz ix
sz) = ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ ix
start ix
sz (Ix1 -> ix
forall ix. Index ix => Ix1 -> ix
pureIndex Ix1
1) Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
(<)
  {-# INLINE iterFullA_ #-}
  iterFullAccST :: RowMajor
-> Scheduler s a
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> ST s a
iterFullAccST (RowMajorInternal Ix1
fact) Scheduler s a
scheduler ix
startIx =
    Ix1
-> Scheduler s a
-> ix
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> ST s a
forall ix s a.
Index ix =>
Ix1
-> Scheduler s a
-> ix
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> ST s a
iterRowMajorST Ix1
fact Scheduler s a
scheduler ix
startIx (Ix1 -> ix
forall ix. Index ix => Ix1 -> ix
pureIndex Ix1
1)
  {-# INLINE iterFullAccST #-}
  iterTargetA_ :: RowMajor
-> Ix1 -> Sz ix -> ix -> Stride ix -> (Ix1 -> ix -> f a) -> f ()
iterTargetA_ RowMajor
_ Ix1
i Sz ix
sz ix
start (Stride ix
stride) =
    Ix1 -> Ix1 -> Sz ix -> ix -> ix -> (Ix1 -> ix -> f a) -> f ()
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
Ix1 -> Ix1 -> Sz ix -> ix -> ix -> (Ix1 -> ix -> f a) -> f ()
iterTargetRowMajorA_ Ix1
0 Ix1
i Sz ix
sz ix
start ix
stride
  {-# INLINE iterTargetA_ #-}
  iterTargetM :: RowMajor
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
iterTargetM RowMajor
_ Ix1
i Sz ix
sz ix
start (Stride ix
stride) =
    Ix1
-> Ix1 -> Sz ix -> ix -> ix -> a -> (Ix1 -> ix -> a -> m a) -> m a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Ix1
-> Ix1 -> Sz ix -> ix -> ix -> a -> (Ix1 -> ix -> a -> m a) -> m a
iterTargetRowMajorAccM Ix1
0 Ix1
i Sz ix
sz ix
start ix
stride
  {-# INLINE iterTargetM #-}
  iterTargetWithStrideAccST :: RowMajor
-> 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 (RowMajor Ix1
fact) Scheduler s a
scheduler Ix1
i Sz ix
sz ix
ix (Stride ix
stride) =
    Ix1
-> Ix1
-> Scheduler s a
-> Ix1
-> Sz ix
-> ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
forall ix s a.
Index ix =>
Ix1
-> Ix1
-> Scheduler s a
-> Ix1
-> Sz ix
-> ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
iterTargetRowMajorAccST Ix1
0 Ix1
fact Scheduler s a
scheduler Ix1
i Sz ix
sz ix
ix ix
stride
  {-# INLINE iterTargetWithStrideAccST #-}
  iterTargetWithStrideAccST_ :: RowMajor
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterTargetWithStrideAccST_ (RowMajor Ix1
fact) Scheduler s ()
scheduler Ix1
i Sz ix
sz ix
ix (Stride ix
stride) =
    Ix1
-> Ix1
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
forall ix s a.
Index ix =>
Ix1
-> Ix1
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterTargetRowMajorAccST_ Ix1
0 Ix1
fact Scheduler s ()
scheduler Ix1
i Sz ix
sz ix
ix ix
stride
  {-# INLINE iterTargetWithStrideAccST_ #-}


newtype RowMajorLinear = RowMajorLinear Int

defRowMajorLinear :: RowMajorLinear
defRowMajorLinear :: RowMajorLinear
defRowMajorLinear = Ix1 -> RowMajorLinear
RowMajorLinear Ix1
8

instance Iterator RowMajorLinear where
  iterTargetM :: RowMajorLinear
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
iterTargetM RowMajorLinear
_ Ix1
iStart Sz ix
sz ix
start (Stride ix
stride) a
acc Ix1 -> ix -> a -> m a
action =
    Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopM Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) a
acc ((Ix1 -> a -> m a) -> m a) -> (Ix1 -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Ix1
i ->
      Ix1 -> ix -> a -> m a
action (Ix1
iStart Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
i) ((Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
(+) ix
start ((Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
(*) ix
stride (Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz Ix1
i)))
  {-# INLINE iterTargetM #-}
  iterTargetA_ :: RowMajorLinear
-> Ix1 -> Sz ix -> ix -> Stride ix -> (Ix1 -> ix -> f a) -> f ()
iterTargetA_ RowMajorLinear
_ Ix1
iStart Sz ix
sz ix
start (Stride ix
stride) Ix1 -> ix -> f a
action =
    Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) ((Ix1 -> f a) -> f ()) -> (Ix1 -> f a) -> f ()
forall a b. (a -> b) -> a -> b
$ \Ix1
i ->
      Ix1 -> ix -> f a
action (Ix1
iStart Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
i) ((Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
(+) ix
start ((Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
(*) ix
stride (Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz Ix1
i)))
  {-# INLINE iterTargetA_ #-}
  iterTargetFullAccST :: RowMajorLinear
-> Scheduler s a
-> Ix1
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
iterTargetFullAccST RowMajorLinear
it Scheduler s a
scheduler Ix1
iStart Sz ix
sz a
acc a -> ST s (a, a)
splitAcc Ix1 -> ix -> a -> ST s a
action =
    let !(RowMajorLinear Ix1
fact) = RowMajorLinear
it
    in Ix1
-> Scheduler s a
-> Ix1
-> Ix1
-> Ix1
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> a -> ST s a)
-> ST s a
forall s a.
Ix1
-> Scheduler s a
-> Ix1
-> Ix1
-> Ix1
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> a -> ST s a)
-> ST s a
iterLinearAccST Ix1
fact Scheduler s a
scheduler Ix1
iStart Ix1
1 (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) a
acc a -> ST s (a, a)
splitAcc ((Ix1 -> a -> ST s a) -> ST s a) -> (Ix1 -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \ !Ix1
i ->
      Ix1 -> ix -> a -> ST s a
action Ix1
i (Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz Ix1
i)
  {-# INLINE iterTargetFullAccST #-}
  iterTargetFullAccST_ :: RowMajorLinear
-> Scheduler s ()
-> Ix1
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterTargetFullAccST_ RowMajorLinear
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz a
acc a -> ST s (a, a)
splitAcc Ix1 -> ix -> a -> ST s a
action =
    let !(RowMajorLinear Ix1
fact) = RowMajorLinear
it
    in Ix1
-> Scheduler s ()
-> Ix1
-> Ix1
-> Ix1
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> a -> ST s a)
-> ST s ()
forall s a.
Ix1
-> Scheduler s ()
-> Ix1
-> Ix1
-> Ix1
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> a -> ST s a)
-> ST s ()
iterLinearAccST_ Ix1
fact Scheduler s ()
scheduler Ix1
iStart Ix1
1 (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) a
acc a -> ST s (a, a)
splitAcc ((Ix1 -> a -> ST s a) -> ST s ())
-> (Ix1 -> a -> ST s a) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Ix1
i ->
      Ix1 -> ix -> a -> ST s a
action Ix1
i (Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz Ix1
i)
  {-# INLINE iterTargetFullAccST_ #-}
  iterTargetFullST_ :: RowMajorLinear
-> Scheduler s ()
-> Ix1
-> Sz ix
-> (Ix1 -> ix -> ST s ())
-> ST s ()
iterTargetFullST_ RowMajorLinear
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz Ix1 -> ix -> ST s ()
action =
    let !(RowMajorLinear Ix1
fact) = RowMajorLinear
it
    in Ix1
-> Scheduler s ()
-> Ix1
-> Ix1
-> Ix1
-> (Ix1 -> ST s ())
-> ST s ()
forall s a.
Ix1
-> Scheduler s ()
-> Ix1
-> Ix1
-> Ix1
-> (Ix1 -> ST s a)
-> ST s ()
iterLinearST_ Ix1
fact Scheduler s ()
scheduler Ix1
iStart Ix1
1 (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) ((Ix1 -> ST s ()) -> ST s ()) -> (Ix1 -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Ix1
i ->
      Ix1 -> ix -> ST s ()
action Ix1
i (Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz Ix1
i)
  {-# INLINE iterTargetFullST_ #-}
  iterTargetWithStrideAccST :: RowMajorLinear
-> 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 RowMajorLinear
it Scheduler s a
scheduler Ix1
iStart Sz ix
sz ix
start (Stride ix
stride) a
acc a -> ST s (a, a)
spliAcc Ix1 -> ix -> a -> ST s a
action =
    let RowMajorLinear Ix1
fact = RowMajorLinear
it
     in Ix1
-> Scheduler s a
-> Ix1
-> Ix1
-> Ix1
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> a -> ST s a)
-> ST s a
forall s a.
Ix1
-> Scheduler s a
-> Ix1
-> Ix1
-> Ix1
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> a -> ST s a)
-> ST s a
iterLinearAccST Ix1
fact Scheduler s a
scheduler Ix1
0 Ix1
1 (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) a
acc a -> ST s (a, a)
spliAcc ((Ix1 -> a -> ST s a) -> ST s a) -> (Ix1 -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \Ix1
i ->
          Ix1 -> ix -> a -> ST s a
action (Ix1
iStart Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
i) (ix -> a -> ST s a) -> ix -> a -> ST s a
forall a b. (a -> b) -> a -> b
$
            (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
(+) ix
start ((Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
(*) ix
stride (Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz Ix1
i))
  {-# INLINE iterTargetWithStrideAccST #-}
  iterTargetWithStrideAccST_ :: RowMajorLinear
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterTargetWithStrideAccST_ RowMajorLinear
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz ix
start (Stride ix
stride) a
acc a -> ST s (a, a)
spliAcc Ix1 -> ix -> a -> ST s a
action =
    let RowMajorLinear Ix1
fact = RowMajorLinear
it
     in Ix1
-> Scheduler s ()
-> Ix1
-> Ix1
-> Ix1
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> a -> ST s a)
-> ST s ()
forall s a.
Ix1
-> Scheduler s ()
-> Ix1
-> Ix1
-> Ix1
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> a -> ST s a)
-> ST s ()
iterLinearAccST_ Ix1
fact Scheduler s ()
scheduler Ix1
0 Ix1
1 (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) a
acc a -> ST s (a, a)
spliAcc ((Ix1 -> a -> ST s a) -> ST s ())
-> (Ix1 -> a -> ST s a) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Ix1
i ->
          Ix1 -> ix -> a -> ST s a
action (Ix1
iStart Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
i) (ix -> a -> ST s a) -> ix -> a -> ST s a
forall a b. (a -> b) -> a -> b
$
            (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
(+) ix
start ((Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
(*) ix
stride (Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz Ix1
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 = Ix1 -> RowMajorUnbalanced
RowMajorUnbalancedInternal Ix1
8

pattern RowMajorUnbalanced :: Int
                 -- ^ Multiplier that will be used to scale number of jobs.
                 -> RowMajorUnbalanced
pattern $bRowMajorUnbalanced :: Ix1 -> RowMajorUnbalanced
$mRowMajorUnbalanced :: forall r. RowMajorUnbalanced -> (Ix1 -> r) -> (Void# -> r) -> r
RowMajorUnbalanced f <- RowMajorUnbalancedInternal f
  where RowMajorUnbalanced = Ix1 -> RowMajorUnbalanced
RowMajorUnbalancedInternal (Ix1 -> RowMajorUnbalanced)
-> (Ix1 -> Ix1) -> Ix1 -> RowMajorUnbalanced
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix1 -> Ix1 -> Ix1
forall a. Ord a => a -> a -> a
max Ix1
1
{-# COMPLETE RowMajorUnbalanced #-}


instance Iterator RowMajorUnbalanced where
  iterFullM :: RowMajorUnbalanced -> ix -> Sz ix -> a -> (ix -> a -> m a) -> m a
iterFullM (RowMajorUnbalanced Ix1
fact) = RowMajor -> ix -> Sz ix -> a -> (ix -> a -> m a) -> m a
forall it ix (m :: * -> *) a.
(Iterator it, Index ix, Monad m) =>
it -> ix -> Sz ix -> a -> (ix -> a -> m a) -> m a
iterFullM (Ix1 -> RowMajor
RowMajor Ix1
fact)
  {-# INLINE iterFullM #-}
  iterFullA_ :: RowMajorUnbalanced -> ix -> Sz ix -> (ix -> f a) -> f ()
iterFullA_ (RowMajorUnbalanced Ix1
fact) = RowMajor -> ix -> Sz ix -> (ix -> f a) -> f ()
forall it ix (f :: * -> *) a.
(Iterator it, Index ix, Applicative f) =>
it -> ix -> Sz ix -> (ix -> f a) -> f ()
iterFullA_ (Ix1 -> RowMajor
RowMajor Ix1
fact)
  {-# INLINE iterFullA_ #-}
  iterTargetM :: RowMajorUnbalanced
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
iterTargetM (RowMajorUnbalanced Ix1
fact) = RowMajor
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
forall it ix (m :: * -> *) a.
(Iterator it, Index ix, Monad m) =>
it
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Ix1 -> ix -> a -> m a)
-> m a
iterTargetM (Ix1 -> RowMajor
RowMajor Ix1
fact)
  {-# INLINE iterTargetM #-}
  iterTargetA_ :: RowMajorUnbalanced
-> Ix1 -> Sz ix -> ix -> Stride ix -> (Ix1 -> ix -> f a) -> f ()
iterTargetA_ (RowMajorUnbalanced Ix1
fact) = RowMajor
-> Ix1 -> Sz ix -> ix -> Stride ix -> (Ix1 -> ix -> f a) -> f ()
forall it ix (f :: * -> *) a.
(Iterator it, Index ix, Applicative f) =>
it -> Ix1 -> Sz ix -> ix -> Stride ix -> (Ix1 -> ix -> f a) -> f ()
iterTargetA_ (Ix1 -> RowMajor
RowMajor Ix1
fact)
  {-# INLINE iterTargetA_ #-}
  iterTargetWithStrideAccST :: RowMajorUnbalanced
-> 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 = (Ix1
 -> (Ix1 -> Bool)
 -> (Ix1 -> Ix1)
 -> a
 -> (Ix1 -> a -> ST s a)
 -> ST s a)
-> RowMajorUnbalanced
-> Scheduler s a
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
forall ix a t s b.
Index ix =>
(Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> t) -> ST s b)
-> RowMajorUnbalanced
-> Scheduler s b
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> t)
-> ST s a
iterUnbalancedTargetWithStride Ix1
-> (Ix1 -> Bool)
-> (Ix1 -> Ix1)
-> a
-> (Ix1 -> a -> ST s a)
-> ST s a
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopM
  {-# INLINE iterTargetWithStrideAccST #-}
  iterTargetWithStrideAccST_ :: RowMajorUnbalanced
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterTargetWithStrideAccST_ RowMajorUnbalanced
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz ix
start Stride ix
stride a
acc a -> ST s (a, a)
splitAcc' Ix1 -> ix -> a -> ST s a
action =
    ST s a -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s a -> ST s ()) -> ST s a -> ST s ()
forall a b. (a -> b) -> a -> b
$
    (Ix1
 -> (Ix1 -> Bool)
 -> (Ix1 -> Ix1)
 -> a
 -> (Ix1 -> a -> ST s a)
 -> ST s ())
-> RowMajorUnbalanced
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
forall ix a t s b.
Index ix =>
(Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> t) -> ST s b)
-> RowMajorUnbalanced
-> Scheduler s b
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> t)
-> ST s a
iterUnbalancedTargetWithStride Ix1
-> (Ix1 -> Bool)
-> (Ix1 -> Ix1)
-> a
-> (Ix1 -> a -> ST s a)
-> ST s ()
forall (f :: * -> *) a.
Monad f =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> f a) -> f ()
innerLoop RowMajorUnbalanced
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz ix
start Stride ix
stride a
acc a -> ST s (a, a)
splitAcc' Ix1 -> ix -> a -> ST s a
action
    where
      innerLoop :: Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> f a) -> f ()
innerLoop Ix1
initial Ix1 -> Bool
condition Ix1 -> Ix1
increment a
initAcc Ix1 -> a -> f a
f =
        f a -> f ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (f a -> f ()) -> f a -> f ()
forall a b. (a -> b) -> a -> b
$ Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> f a) -> f a
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopM Ix1
initial Ix1 -> Bool
condition Ix1 -> Ix1
increment a
initAcc Ix1 -> 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 :: (Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> t) -> ST s b)
-> RowMajorUnbalanced
-> Scheduler s b
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> t)
-> ST s a
iterUnbalancedTargetWithStride Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> t) -> ST s b
innerLoop RowMajorUnbalanced
it Scheduler s b
scheduler Ix1
iStart Sz ix
sz ix
start Stride ix
stride a
acc a -> ST s (a, a)
splitAcc Ix1 -> ix -> t
action =
  let RowMajorUnbalanced Ix1
fact = RowMajorUnbalanced
it
      !n :: Ix1
n = Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz
      !step :: Ix1
step = Ix1 -> Ix1 -> Ix1
forall a. Ord a => a -> a -> a
min (Ix1
fact Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
* Scheduler s b -> Ix1
forall s a. Scheduler s a -> Ix1
numWorkers Scheduler s b
scheduler) Ix1
n
  in Ix1
-> (Ix1 -> Bool)
-> (Ix1 -> Ix1)
-> a
-> (Ix1 -> a -> ST s a)
-> ST s a
forall (m :: * -> *) a.
Monad m =>
Ix1
-> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> a -> m a) -> m a
loopM Ix1
0 (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Ix1
step) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1) a
acc ((Ix1 -> a -> ST s a) -> ST s a) -> (Ix1 -> a -> ST s a) -> ST s a
forall a b. (a -> b) -> a -> b
$ \ !Ix1
istep !a
a -> do
       (a
curAcc, a
nextAcc) <- a -> ST s (a, a)
splitAcc a
a
       Scheduler (PrimState (ST s)) b -> ST s b -> ST s ()
forall (m :: * -> *) a.
PrimBase m =>
Scheduler (PrimState m) a -> m a -> m ()
scheduleMassivWork Scheduler s b
Scheduler (PrimState (ST s)) b
scheduler (ST s b -> ST s ()) -> ST s b -> ST s ()
forall a b. (a -> b) -> a -> b
$
         Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> a -> (Ix1 -> t) -> ST s b
innerLoop Ix1
istep (Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Ix1
n) (Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
step) a
curAcc ((Ix1 -> t) -> ST s b) -> (Ix1 -> t) -> ST s b
forall a b. (a -> b) -> a -> b
$ \Ix1
i ->
           Ix1 -> ix -> t
action (Ix1
iStart Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
i) (ix -> t) -> ix -> t
forall a b. (a -> b) -> a -> b
$
             (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
(+) ix
start ((Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
(*) (Stride ix -> ix
forall ix. Stride ix -> ix
unStride Stride ix
stride) (Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz Ix1
i))
       a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
nextAcc
{-# INLINE iterUnbalancedTargetWithStride #-}


noSplit :: Applicative m => () -> m ((), ())
noSplit :: () -> m ((), ())
noSplit ()
_ = ((), ()) -> m ((), ())
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 :: it
-> Scheduler s a
-> Ix1
-> Sz ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s a
iterTargetAccST it
it Scheduler s a
scheduler Ix1
iStart Sz ix
sz ix
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
forall it ix s a.
(Iterator it, 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 it
it Scheduler s a
scheduler Ix1
iStart Sz ix
sz ix
ix Stride 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_ :: it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterTargetAccST_ it
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz ix
ix =
  it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
forall it ix s a.
(Iterator it, 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 ()
iterTargetWithStrideAccST_ it
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz ix
ix Stride 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_ :: it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> Stride ix
-> (Ix1 -> ix -> ST s ())
-> ST s ()
iterTargetFullWithStrideST_ it
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz =
  it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> (Ix1 -> ix -> ST s ())
-> ST s ()
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> (Ix1 -> ix -> ST s a)
-> ST s ()
iterTargetWithStrideST_ it
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz (Ix1 -> ix
forall ix. Index ix => Ix1 -> ix
pureIndex Ix1
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_ :: it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> (Ix1 -> ix -> ST s ())
-> ST s ()
iterTargetST_ it
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz ix
ix =
  it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> (Ix1 -> ix -> ST s ())
-> ST s ()
forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> (Ix1 -> ix -> ST s a)
-> ST s ()
iterTargetWithStrideST_ it
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz ix
ix Stride 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 :: 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 it
it Scheduler s a
scheduler Ix1
iStart Sz ix
sz =
  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
forall it ix s a.
(Iterator it, 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 it
it Scheduler s a
scheduler Ix1
iStart Sz ix
sz (Ix1 -> ix
forall ix. Index ix => Ix1 -> ix
pureIndex Ix1
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_ :: it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
iterTargetFullWithStrideAccST_ it
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz =
  it
-> Scheduler s ()
-> Ix1
-> Sz ix
-> ix
-> Stride ix
-> a
-> (a -> ST s (a, a))
-> (Ix1 -> ix -> a -> ST s a)
-> ST s ()
forall it ix s a.
(Iterator it, 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 ()
iterTargetWithStrideAccST_ it
it Scheduler s ()
scheduler Ix1
iStart Sz ix
sz (Ix1 -> ix
forall ix. Index ix => Ix1 -> ix
pureIndex Ix1
0)
{-# INLINE iterTargetFullWithStrideAccST_ #-}