{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Massiv.Array.Delayed.Push
( DL(..)
, Array(..)
, Loader
, toLoadArray
, makeLoadArrayS
, makeLoadArray
, unsafeMakeLoadArray
, unsafeMakeLoadArrayAdjusted
, fromStrideLoad
, appendOuterM
, concatOuterM
) where
import Control.Monad
import Control.Scheduler as S (traverse_)
import Data.Foldable as F
import Data.Massiv.Core.Common
import Prelude hiding (map, zipWith)
#include "massiv.h"
data DL = DL deriving Ix1 -> DL -> ShowS
[DL] -> ShowS
DL -> String
forall a.
(Ix1 -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DL] -> ShowS
$cshowList :: [DL] -> ShowS
show :: DL -> String
$cshow :: DL -> String
showsPrec :: Ix1 -> DL -> ShowS
$cshowsPrec :: Ix1 -> DL -> ShowS
Show
type Loader e =
forall s. Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
data instance Array DL ix e = DLArray
{ forall ix e. Array DL ix e -> Comp
dlComp :: !Comp
, forall ix e. Array DL ix e -> Sz ix
dlSize :: !(Sz ix)
, forall ix e. Array DL ix e -> Loader e
dlLoad :: Loader e
}
instance Strategy DL where
getComp :: forall ix e. Array DL ix e -> Comp
getComp = forall ix e. Array DL ix e -> Comp
dlComp
{-# INLINE getComp #-}
setComp :: forall ix e. Comp -> Array DL ix e -> Array DL ix e
setComp Comp
c Array DL ix e
arr = Array DL ix e
arr {dlComp :: Comp
dlComp = Comp
c}
{-# INLINE setComp #-}
repr :: DL
repr = DL
DL
instance Index ix => Shape DL ix where
maxLinearSize :: forall e. Array DL ix e -> Maybe Sz1
maxLinearSize = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e. (Index ix, Size r) => Array r ix e -> Ix1
elemsCount
{-# INLINE maxLinearSize #-}
instance Size DL where
size :: forall ix e. Array DL ix e -> Sz ix
size = forall ix e. Array DL ix e -> Sz ix
dlSize
{-# INLINE size #-}
unsafeResize :: forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array DL ix e -> Array DL ix' e
unsafeResize !Sz ix'
sz !Array DL ix e
arr = Array DL ix e
arr { dlSize :: Sz ix'
dlSize = Sz ix'
sz }
{-# INLINE unsafeResize #-}
instance Semigroup (Array DL Ix1 e) where
<> :: Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e
(<>) = forall e. Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e
mappendDL
{-# INLINE (<>) #-}
instance Monoid (Array DL Ix1 e) where
mempty :: Array DL Ix1 e
mempty = DLArray {dlComp :: Comp
dlComp = forall a. Monoid a => a
mempty, dlSize :: Sz1
dlSize = forall ix. Index ix => Sz ix
zeroSz, dlLoad :: Loader e
dlLoad = \Scheduler s ()
_ Ix1
_ Ix1 -> e -> ST s ()
_ Ix1 -> Sz1 -> e -> ST s ()
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()}
{-# INLINE mempty #-}
#if !MIN_VERSION_base(4,11,0)
mappend = mappendDL
{-# INLINE mappend #-}
#endif
mconcat :: [Array DL Ix1 e] -> Array DL Ix1 e
mconcat [] = forall a. Monoid a => a
mempty
mconcat [Array DL Ix1 e
x] = Array DL Ix1 e
x
mconcat [Array DL Ix1 e
x, Array DL Ix1 e
y] = Array DL Ix1 e
x forall a. Semigroup a => a -> a -> a
<> Array DL Ix1 e
y
mconcat [Array DL Ix1 e]
xs = forall e. [Array DL Ix1 e] -> Array DL Ix1 e
mconcatDL [Array DL Ix1 e]
xs
{-# INLINE mconcat #-}
mconcatDL :: forall e . [Array DL Ix1 e] -> Array DL Ix1 e
mconcatDL :: forall e. [Array DL Ix1 e] -> Array DL Ix1 e
mconcatDL ![Array DL Ix1 e]
arrs =
DLArray {dlComp :: Comp
dlComp = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall r ix e. Strategy r => Array r ix e -> Comp
getComp [Array DL Ix1 e]
arrs, dlSize :: Sz1
dlSize = forall ix. ix -> Sz ix
SafeSz Ix1
k, dlLoad :: Loader e
dlLoad = Loader e
load}
where
!k :: Ix1
k = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' forall a. Num a => a -> a -> a
(+) Ix1
0 (forall ix. Sz ix -> ix
unSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix e. Size r => Array r ix e -> Sz ix
size forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Array DL Ix1 e]
arrs)
load :: forall s .
Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s ()
load :: Loader e
load Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet =
let loadArr :: Ix1 -> Array DL Ix1 e -> ST s Ix1
loadArr !Ix1
startAtCur DLArray {dlSize :: forall ix e. Array DL ix e -> Sz ix
dlSize = SafeSz Ix1
kCur, Loader e
dlLoad :: Loader e
dlLoad :: forall ix e. Array DL ix e -> Loader e
dlLoad} = do
let !endAtCur :: Ix1
endAtCur = Ix1
startAtCur forall a. Num a => a -> a -> a
+ Ix1
kCur
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Loader e
dlLoad Scheduler s ()
scheduler Ix1
startAtCur Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ix1
endAtCur
{-# INLINE loadArr #-}
in forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ Ix1 -> Array DL Ix1 e -> ST s Ix1
loadArr Ix1
startAt [Array DL Ix1 e]
arrs
{-# INLINE load #-}
{-# INLINE mconcatDL #-}
mappendDL :: forall e . Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e
mappendDL :: forall e. Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e
mappendDL (DLArray Comp
c1 Sz1
sz1 Loader e
load1) (DLArray Comp
c2 Sz1
sz2 Loader e
load2) =
DLArray {dlComp :: Comp
dlComp = Comp
c1 forall a. Semigroup a => a -> a -> a
<> Comp
c2, dlSize :: Sz1
dlSize = forall ix. ix -> Sz ix
SafeSz (Ix1
k1 forall a. Num a => a -> a -> a
+ Ix1
k2), dlLoad :: Loader e
dlLoad = Loader e
load}
where
!k1 :: Ix1
k1 = forall ix. Sz ix -> ix
unSz Sz1
sz1
!k2 :: Ix1
k2 = forall ix. Sz ix -> ix
unSz Sz1
sz2
load :: forall s.
Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s ()
load :: Loader e
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet = do
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Loader e
load1 Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Loader e
load2 Scheduler s ()
scheduler (Ix1
startAt forall a. Num a => a -> a -> a
+ Ix1
k1) Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet
{-# INLINE load #-}
{-# INLINE mappendDL #-}
appendOuterM ::
forall ix e m. (Index ix, MonadThrow m)
=> Array DL ix e
-> Array DL ix e
-> m (Array DL ix e)
appendOuterM :: forall ix e (m :: * -> *).
(Index ix, MonadThrow m) =>
Array DL ix e -> Array DL ix e -> m (Array DL ix e)
appendOuterM (DLArray Comp
c1 Sz ix
sz1 Loader e
load1) (DLArray Comp
c2 Sz ix
sz2 Loader e
load2) = do
let (!Sz1
i1, !Sz (Lower ix)
szl1) = forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz Sz ix
sz1
(!Sz1
i2, !Sz (Lower ix)
szl2) = forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz Sz ix
sz2
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Sz (Lower ix)
szl1 forall a. Eq a => a -> a -> Bool
== Sz (Lower ix)
szl2) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException Sz ix
sz1 Sz ix
sz2
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
DLArray {dlComp :: Comp
dlComp = Comp
c1 forall a. Semigroup a => a -> a -> a
<> Comp
c2, dlSize :: Sz ix
dlSize = forall ix. Index ix => Sz1 -> Sz (Lower ix) -> Sz ix
consSz (forall ix.
Index ix =>
(Ix1 -> Ix1 -> Ix1) -> Sz ix -> Sz ix -> Sz ix
liftSz2 forall a. Num a => a -> a -> a
(+) Sz1
i1 Sz1
i2) Sz (Lower ix)
szl1, dlLoad :: Loader e
dlLoad = Loader e
load}
where
load :: Loader e
load :: Loader e
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet = do
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Loader e
load1 Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Loader e
load2 Scheduler s ()
scheduler (Ix1
startAt forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz1) Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet
{-# INLINE load #-}
{-# INLINE appendOuterM #-}
concatOuterM ::
forall ix e m. (Index ix, MonadThrow m)
=> [Array DL ix e]
-> m (Array DL ix e)
concatOuterM :: forall ix e (m :: * -> *).
(Index ix, MonadThrow m) =>
[Array DL ix e] -> m (Array DL ix e)
concatOuterM =
\case
[] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall r ix e. Load r ix e => Array r ix e
empty
(Array DL ix e
x:[Array DL ix e]
xs) -> forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM forall ix e (m :: * -> *).
(Index ix, MonadThrow m) =>
Array DL ix e -> Array DL ix e -> m (Array DL ix e)
appendOuterM Array DL ix e
x [Array DL ix e]
xs
{-# INLINE concatOuterM #-}
makeLoadArrayS ::
forall ix e. Index ix
=> Sz ix
-> e
-> (forall m. Monad m => (ix -> e -> m Bool) -> m ())
-> Array DL ix e
makeLoadArrayS :: forall ix e.
Index ix =>
Sz ix
-> e
-> (forall (m :: * -> *). Monad m => (ix -> e -> m Bool) -> m ())
-> Array DL ix e
makeLoadArrayS Sz ix
sz e
defVal forall (m :: * -> *). Monad m => (ix -> e -> m Bool) -> m ()
writer = forall ix e.
Index ix =>
Comp
-> Sz ix
-> e
-> (forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ())
-> Array DL ix e
makeLoadArray Comp
Seq Sz ix
sz e
defVal (forall a b. a -> b -> a
const forall (m :: * -> *). Monad m => (ix -> e -> m Bool) -> m ()
writer)
{-# INLINE makeLoadArrayS #-}
makeLoadArray ::
forall ix e. Index ix
=> Comp
-> Sz ix
-> e
-> (forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ())
-> Array DL ix e
makeLoadArray :: forall ix e.
Index ix =>
Comp
-> Sz ix
-> e
-> (forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ())
-> Array DL ix e
makeLoadArray Comp
comp Sz ix
sz e
defVal forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ()
writer = forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray Comp
comp Sz ix
sz forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load
where
load :: forall s.
Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s ()
load :: forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
uWrite Ix1 -> Sz1 -> e -> ST s ()
uSet = do
Ix1 -> Sz1 -> e -> ST s ()
uSet Ix1
startAt (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
sz) e
defVal
let safeWrite :: ix -> e -> ST s Bool
safeWrite !ix
ix !e
e
| forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix = Bool
True forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Ix1 -> e -> ST s ()
uWrite (Ix1
startAt forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Ix1
toLinearIndex Sz ix
sz ix
ix) e
e
| Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
{-# INLINE safeWrite #-}
forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ()
writer Scheduler s ()
scheduler ix -> e -> ST s Bool
safeWrite
{-# INLINE load #-}
{-# INLINE makeLoadArray #-}
unsafeMakeLoadArray ::
forall ix e. Index ix
=> Comp
-> Sz ix
-> Maybe e
-> (forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ())
-> Array DL ix e
unsafeMakeLoadArray :: forall ix e.
Index ix =>
Comp
-> Sz ix
-> Maybe e
-> (forall s.
Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ())
-> Array DL ix e
unsafeMakeLoadArray Comp
comp Sz ix
sz Maybe e
mDefVal forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ()
writer = forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray Comp
comp Sz ix
sz Loader e
load
where
load :: Loader e
load :: Loader e
load Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
uWrite Ix1 -> Sz1 -> e -> ST s ()
uSet = do
forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
S.traverse_ (Ix1 -> Sz1 -> e -> ST s ()
uSet Ix1
startAt (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
sz)) Maybe e
mDefVal
forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ()
writer Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
uWrite
{-# INLINE load #-}
{-# INLINE unsafeMakeLoadArray #-}
unsafeMakeLoadArrayAdjusted ::
forall ix e. Index ix
=> Comp
-> Sz ix
-> Maybe e
-> (forall s. Scheduler s () -> (Ix1 -> e -> ST s ()) -> ST s ())
-> Array DL ix e
unsafeMakeLoadArrayAdjusted :: forall ix e.
Index ix =>
Comp
-> Sz ix
-> Maybe e
-> (forall s. Scheduler s () -> (Ix1 -> e -> ST s ()) -> ST s ())
-> Array DL ix e
unsafeMakeLoadArrayAdjusted Comp
comp Sz ix
sz Maybe e
mDefVal forall s. Scheduler s () -> (Ix1 -> e -> ST s ()) -> ST s ()
writer = forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray Comp
comp Sz ix
sz forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load
where
load :: forall s.
Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s ()
load :: forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
uWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet = do
forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
S.traverse_ (Ix1 -> Sz1 -> e -> ST s ()
dlSet Ix1
startAt (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
sz)) Maybe e
mDefVal
forall s. Scheduler s () -> (Ix1 -> e -> ST s ()) -> ST s ()
writer Scheduler s ()
scheduler (\Ix1
i -> Ix1 -> e -> ST s ()
uWrite (Ix1
startAt forall a. Num a => a -> a -> a
+ Ix1
i))
{-# INLINE load #-}
{-# INLINE unsafeMakeLoadArrayAdjusted #-}
toLoadArray ::
forall r ix e. (Size r, Load r ix e)
=> Array r ix e
-> Array DL ix e
toLoadArray :: forall r ix e.
(Size r, Load r ix e) =>
Array r ix e -> Array DL ix e
toLoadArray Array r ix e
arr = forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) Sz ix
sz forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load
where
!sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
load :: forall s.
Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s ()
load :: forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet =
forall r ix e s.
Load r ix e =>
Scheduler s ()
-> Array r ix e
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithSetST_ Scheduler s ()
scheduler Array r ix e
arr (Ix1 -> e -> ST s ()
dlWrite forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Num a => a -> a -> a
+ Ix1
startAt)) (\Ix1
offset -> Ix1 -> Sz1 -> e -> ST s ()
dlSet (Ix1
offset forall a. Num a => a -> a -> a
+ Ix1
startAt))
{-# INLINE load #-}
{-# INLINE[1] toLoadArray #-}
{-# RULES "toLoadArray/id" toLoadArray = id #-}
fromStrideLoad ::
forall r ix e. (StrideLoad r ix e)
=> Stride ix
-> Array r ix e
-> Array DL ix e
fromStrideLoad :: forall r ix e.
StrideLoad r ix e =>
Stride ix -> Array r ix e -> Array DL ix e
fromStrideLoad Stride ix
stride Array r ix e
arr =
forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) Sz ix
newsz Loader e
load
where
!newsz :: Sz ix
newsz = forall ix. Index ix => Stride ix -> Sz ix -> Sz ix
strideSize Stride ix
stride (forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r ix e
arr)
load :: Loader e
load :: Loader e
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
_ =
forall r ix e s.
StrideLoad r ix e =>
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array r ix e
-> (Ix1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ Scheduler s ()
scheduler Stride ix
stride Sz ix
newsz Array r ix e
arr (\ !Ix1
i -> Ix1 -> e -> ST s ()
dlWrite (Ix1
i forall a. Num a => a -> a -> a
+ Ix1
startAt))
{-# INLINE load #-}
{-# INLINE fromStrideLoad #-}
instance Index ix => Load DL ix e where
makeArrayLinear :: Comp -> Sz ix -> (Ix1 -> e) -> Array DL ix e
makeArrayLinear Comp
comp Sz ix
sz Ix1 -> e
f = forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray Comp
comp Sz ix
sz Loader e
load
where
load :: Loader e
load :: Loader e
load Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
_ =
forall s (m :: * -> *) b c.
MonadPrimBase s m =>
Scheduler s ()
-> Ix1 -> Ix1 -> (Ix1 -> m b) -> (Ix1 -> b -> m c) -> m ()
splitLinearlyWithStartAtM_ Scheduler s ()
scheduler Ix1
startAt (forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix1 -> e
f) Ix1 -> e -> ST s ()
dlWrite
{-# INLINE load #-}
{-# INLINE makeArrayLinear #-}
replicate :: Comp -> Sz ix -> e -> Array DL ix e
replicate Comp
comp !Sz ix
sz !e
e = forall ix e.
Index ix =>
Comp
-> Sz ix
-> e
-> (forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ())
-> Array DL ix e
makeLoadArray Comp
comp Sz ix
sz e
e forall a b. (a -> b) -> a -> b
$ \Scheduler s ()
_ ix -> e -> ST s Bool
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINE replicate #-}
iterArrayLinearWithSetST_ :: forall s.
Scheduler s ()
-> Array DL ix e
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithSetST_ Scheduler s ()
scheduler DLArray {Loader e
dlLoad :: Loader e
dlLoad :: forall ix e. Array DL ix e -> Loader e
dlLoad} = Loader e
dlLoad Scheduler s ()
scheduler Ix1
0
{-# INLINE iterArrayLinearWithSetST_ #-}
instance Index ix => Functor (Array DL ix) where
fmap :: forall a b. (a -> b) -> Array DL ix a -> Array DL ix b
fmap a -> b
f Array DL ix a
arr = Array DL ix a
arr {dlLoad :: Loader b
dlLoad = forall ix a b s.
Array DL ix a
-> (a -> b)
-> Scheduler s ()
-> Ix1
-> (Ix1 -> b -> ST s ())
-> (Ix1 -> Sz1 -> b -> ST s ())
-> ST s ()
loadFunctor Array DL ix a
arr a -> b
f}
{-# INLINE fmap #-}
<$ :: forall a b. a -> Array DL ix b -> Array DL ix a
(<$) = forall ix a b. Index ix => a -> Array DL ix b -> Array DL ix a
overwriteFunctor
{-# INLINE (<$) #-}
overwriteFunctor :: forall ix a b. Index ix => a -> Array DL ix b -> Array DL ix a
overwriteFunctor :: forall ix a b. Index ix => a -> Array DL ix b -> Array DL ix a
overwriteFunctor a
e Array DL ix b
arr = Array DL ix b
arr {dlLoad :: Loader a
dlLoad = Loader a
load}
where
load :: Loader a
load :: Loader a
load Scheduler s ()
_ !Ix1
startAt Ix1 -> a -> ST s ()
_ Ix1 -> Sz1 -> a -> ST s ()
dlSet = Ix1 -> Sz1 -> a -> ST s ()
dlSet Ix1
startAt (forall r ix e. Shape r ix => Array r ix e -> Sz1
linearSize Array DL ix b
arr) a
e
{-# INLINE load #-}
{-# INLINE overwriteFunctor #-}
loadFunctor ::
Array DL ix a
-> (a -> b)
-> Scheduler s ()
-> Ix1
-> (Ix1 -> b -> ST s ())
-> (Ix1 -> Sz1 -> b -> ST s ())
-> ST s ()
loadFunctor :: forall ix a b s.
Array DL ix a
-> (a -> b)
-> Scheduler s ()
-> Ix1
-> (Ix1 -> b -> ST s ())
-> (Ix1 -> Sz1 -> b -> ST s ())
-> ST s ()
loadFunctor Array DL ix a
arr a -> b
f Scheduler s ()
scheduler Ix1
startAt Ix1 -> b -> ST s ()
uWrite Ix1 -> Sz1 -> b -> ST s ()
uSet =
forall ix e. Array DL ix e -> Loader e
dlLoad Array DL ix a
arr Scheduler s ()
scheduler Ix1
startAt (\ !Ix1
i a
e -> Ix1 -> b -> ST s ()
uWrite Ix1
i (a -> b
f a
e)) (\Ix1
o Sz1
sz a
e -> Ix1 -> Sz1 -> b -> ST s ()
uSet Ix1
o Sz1
sz (a -> b
f a
e))
{-# INLINE loadFunctor #-}