{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Massiv.Array.Delayed.Windowed (
DW (..),
Array (..),
Window (..),
insertWindow,
getWindow,
dropWindow,
makeWindowedArray,
) where
import Control.Monad (when)
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Manifest.Boxed
import Data.Massiv.Array.Manifest.Internal
import Data.Massiv.Core
import Data.Massiv.Core.Common
import Data.Massiv.Core.List (showArrayList, showsArrayPrec)
import Data.Maybe (fromMaybe)
import GHC.TypeLits
data DW = DW
data Window ix e = Window
{ forall ix e. Window ix e -> ix
windowStart :: !ix
, forall ix e. Window ix e -> Sz ix
windowSize :: !(Sz ix)
, forall ix e. Window ix e -> ix -> e
windowIndex :: ix -> e
, forall ix e. Window ix e -> Maybe Ix1
windowUnrollIx2 :: !(Maybe Int)
}
instance Functor (Window ix) where
fmap :: forall a b. (a -> b) -> Window ix a -> Window ix b
fmap a -> b
f arr :: Window ix a
arr@Window{ix -> a
windowIndex :: ix -> a
windowIndex :: forall ix e. Window ix e -> ix -> e
windowIndex} = Window ix a
arr{windowIndex :: ix -> b
windowIndex = a -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> a
windowIndex}
data instance Array DW ix e = DWArray
{ forall ix e. Array DW ix e -> Array D ix e
dwArray :: !(Array D ix e)
, forall ix e. Array DW ix e -> Maybe (Window ix e)
dwWindow :: !(Maybe (Window ix e))
}
instance (Ragged L ix e, Load DW ix e, Show e) => Show (Array DW ix e) where
showsPrec :: Ix1 -> Array DW ix e -> ShowS
showsPrec = forall r r' ix e.
(Ragged L ix e, Load r ix e, Load r' ix e, Source r' e, Show e) =>
(Array r ix e -> Array r' ix e) -> Ix1 -> Array r ix e -> ShowS
showsArrayPrec (forall r e r' ix.
(Manifest r e, Load r' ix e) =>
r -> Array r' ix e -> Array r ix e
computeAs B
B)
showList :: [Array DW ix e] -> ShowS
showList = forall arr. Show arr => [arr] -> ShowS
showArrayList
instance Strategy DW where
setComp :: forall ix e. Comp -> Array DW ix e -> Array DW ix e
setComp Comp
c Array DW ix e
arr = Array DW ix e
arr{dwArray :: Array D ix e
dwArray = (forall ix e. Array DW ix e -> Array D ix e
dwArray Array DW ix e
arr){dComp :: Comp
dComp = Comp
c}}
{-# INLINE setComp #-}
getComp :: forall ix e. Array DW ix e -> Comp
getComp = forall ix e. Array D ix e -> Comp
dComp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array DW ix e -> Array D ix e
dwArray
{-# INLINE getComp #-}
repr :: DW
repr = DW
DW
instance Functor (Array DW ix) where
fmap :: forall a b. (a -> b) -> Array DW ix a -> Array DW ix b
fmap a -> b
f arr :: Array DW ix a
arr@DWArray{Array D ix a
dwArray :: Array D ix a
dwArray :: forall ix e. Array DW ix e -> Array D ix e
dwArray, Maybe (Window ix a)
dwWindow :: Maybe (Window ix a)
dwWindow :: forall ix e. Array DW ix e -> Maybe (Window ix e)
dwWindow} =
Array DW ix a
arr
{ dwArray :: Array D ix b
dwArray = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Array D ix a
dwArray
, dwWindow :: Maybe (Window ix b)
dwWindow = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Window ix a)
dwWindow
}
{-# INLINE fmap #-}
makeWindowedArray
:: (Index ix, Source r e)
=> Array r ix e
-> ix
-> Sz ix
-> (ix -> e)
-> Array DW ix e
makeWindowedArray :: forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> ix -> Sz ix -> (ix -> e) -> Array DW ix e
makeWindowedArray !Array r ix e
arr ix
wStart Sz ix
wSize ix -> e
wIndex =
forall ix e.
Index ix =>
Array D ix e -> Window ix e -> Array DW ix e
insertWindow (forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Array D ix e
delay Array r ix e
arr) forall a b. (a -> b) -> a -> b
$
Window{windowStart :: ix
windowStart = ix
wStart, windowSize :: Sz ix
windowSize = Sz ix
wSize, windowIndex :: ix -> e
windowIndex = ix -> e
wIndex, windowUnrollIx2 :: Maybe Ix1
windowUnrollIx2 = forall a. Maybe a
Nothing}
{-# INLINE makeWindowedArray #-}
insertWindow
:: Index ix
=> Array D ix e
-> Window ix e
-> Array DW ix e
insertWindow :: forall ix e.
Index ix =>
Array D ix e -> Window ix e -> Array DW ix e
insertWindow !Array D ix e
arr !Window ix e
window =
DWArray
{ dwArray :: Array D ix e
dwArray = forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Array D ix e
delay Array D ix e
arr
, dwWindow :: Maybe (Window ix e)
dwWindow =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$!
Window
{ windowStart :: ix
windowStart = ix
wStart'
, windowSize :: Sz ix
windowSize = forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 forall a. Ord a => a -> a -> a
min ix
wSize (forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 (-) ix
sz ix
wStart'))
, windowIndex :: ix -> e
windowIndex = ix -> e
wIndex
, windowUnrollIx2 :: Maybe Ix1
windowUnrollIx2 = Maybe Ix1
wUnrollIx2
}
}
where
wStart' :: ix
wStart' = forall ix. Sz ix -> ix
unSz (forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 forall a. Ord a => a -> a -> a
min ix
wStart (forall ix. Index ix => (Ix1 -> Ix1) -> ix -> ix
liftIndex (forall a. Num a => a -> a -> a
subtract Ix1
1) ix
sz)))
Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array D ix e
arr
Window
{ windowStart :: forall ix e. Window ix e -> ix
windowStart = ix
wStart
, windowSize :: forall ix e. Window ix e -> Sz ix
windowSize = Sz ix
wSize
, windowIndex :: forall ix e. Window ix e -> ix -> e
windowIndex = ix -> e
wIndex
, windowUnrollIx2 :: forall ix e. Window ix e -> Maybe Ix1
windowUnrollIx2 = Maybe Ix1
wUnrollIx2
} = Window ix e
window
{-# INLINE insertWindow #-}
getWindow :: Array DW ix e -> Maybe (Window ix e)
getWindow :: forall ix e. Array DW ix e -> Maybe (Window ix e)
getWindow = forall ix e. Array DW ix e -> Maybe (Window ix e)
dwWindow
{-# INLINE getWindow #-}
dropWindow :: Array DW ix e -> Array D ix e
dropWindow :: forall ix e. Array DW ix e -> Array D ix e
dropWindow = forall ix e. Array DW ix e -> Array D ix e
dwArray
{-# INLINE dropWindow #-}
zeroWindow :: Index ix => Window ix e
zeroWindow :: forall ix e. Index ix => Window ix e
zeroWindow = forall ix e. ix -> Sz ix -> (ix -> e) -> Maybe Ix1 -> Window ix e
Window forall ix. Index ix => ix
zeroIndex forall ix. Index ix => Sz ix
zeroSz forall a. a
windowError forall a. Maybe a
Nothing
{-# INLINE zeroWindow #-}
data EmptyWindowException = EmptyWindowException deriving (EmptyWindowException -> EmptyWindowException -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: EmptyWindowException -> EmptyWindowException -> Bool
$c/= :: EmptyWindowException -> EmptyWindowException -> Bool
== :: EmptyWindowException -> EmptyWindowException -> Bool
$c== :: EmptyWindowException -> EmptyWindowException -> Bool
Eq, Ix1 -> EmptyWindowException -> ShowS
[EmptyWindowException] -> ShowS
EmptyWindowException -> String
forall a.
(Ix1 -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [EmptyWindowException] -> ShowS
$cshowList :: [EmptyWindowException] -> ShowS
show :: EmptyWindowException -> String
$cshow :: EmptyWindowException -> String
showsPrec :: Ix1 -> EmptyWindowException -> ShowS
$cshowsPrec :: Ix1 -> EmptyWindowException -> ShowS
Show)
instance Exception EmptyWindowException where
displayException :: EmptyWindowException -> String
displayException EmptyWindowException
_ = String
"Index of zero size Window"
windowError :: a
windowError :: forall a. a
windowError = forall e a. (HasCallStack, Exception e) => e -> a
throwImpossible EmptyWindowException
EmptyWindowException
{-# NOINLINE windowError #-}
loadWithIx1
:: (Monad m)
=> (m () -> m ())
-> Array DW Ix1 e
-> (Ix1 -> e -> m a)
-> m (Ix1 -> Ix1 -> m (), Ix1, Ix1)
loadWithIx1 :: forall (m :: * -> *) e a.
Monad m =>
(m () -> m ())
-> Array DW Ix1 e
-> (Ix1 -> e -> m a)
-> m (Ix1 -> Ix1 -> m (), Ix1, Ix1)
loadWithIx1 m () -> m ()
with (DWArray a :: Array D Ix1 e
a@(DArray Comp
_ Sz Ix1
sz PrefIndex Ix1 e
_) Maybe (Window Ix1 e)
mWindow) Ix1 -> e -> m a
uWrite = do
let Window Ix1
it Sz Ix1
wk Ix1 -> e
indexW Maybe Ix1
_ = forall a. a -> Maybe a -> a
fromMaybe forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window Ix1 e)
mWindow
wEnd :: Ix1
wEnd = Ix1
it forall a. Num a => a -> a -> a
+ forall ix. Sz ix -> ix
unSz Sz Ix1
wk
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ Ix1
0 Ix1
it Ix1
1 forall a. Ord a => a -> a -> Bool
(<) forall a b. (a -> b) -> a -> b
$ \ !Ix1
i -> Ix1 -> e -> m a
uWrite Ix1
i (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D Ix1 e
a Ix1
i)
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ Ix1
wEnd (forall ix. Sz ix -> ix
unSz Sz Ix1
sz) Ix1
1 forall a. Ord a => a -> a -> Bool
(<) forall a b. (a -> b) -> a -> b
$ \ !Ix1
i -> Ix1 -> e -> m a
uWrite Ix1
i (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D Ix1 e
a Ix1
i)
forall (m :: * -> *) a. Monad m => a -> m a
return (\Ix1
from Ix1
to -> m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ Ix1
from Ix1
to Ix1
1 forall a. Ord a => a -> a -> Bool
(<) forall a b. (a -> b) -> a -> b
$ \ !Ix1
i -> Ix1 -> e -> m a
uWrite Ix1
i (Ix1 -> e
indexW Ix1
i), Ix1
it, Ix1
wEnd)
{-# INLINE loadWithIx1 #-}
instance Index ix => Shape DW ix where
maxLinearSize :: forall e. Array DW ix e -> Maybe (Sz Ix1)
maxLinearSize = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix e. Shape r ix => Array r ix e -> Sz Ix1
linearSize
{-# INLINE maxLinearSize #-}
linearSize :: forall e. Array DW ix e -> Sz Ix1
linearSize = forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Sz ix -> Ix1
totalElem forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array D ix e -> Sz ix
dSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array DW ix e -> Array D ix e
dwArray
{-# INLINE linearSize #-}
outerSize :: forall e. Array DW ix e -> Sz ix
outerSize = forall ix e. Array D ix e -> Sz ix
dSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array DW ix e -> Array D ix e
dwArray
{-# INLINE outerSize #-}
instance Load DW Ix1 e where
makeArray :: Comp -> Sz Ix1 -> (Ix1 -> e) -> Array DW Ix1 e
makeArray Comp
c Sz Ix1
sz Ix1 -> e
f = forall ix e. Array D ix e -> Maybe (Window ix e) -> Array DW ix e
DWArray (forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
c Sz Ix1
sz Ix1 -> e
f) forall a. Maybe a
Nothing
{-# INLINE makeArray #-}
iterArrayLinearST_ :: forall s.
Scheduler s ()
-> Array DW Ix1 e -> (Ix1 -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler Array DW Ix1 e
arr Ix1 -> e -> ST s ()
uWrite = do
(Ix1 -> Ix1 -> ST s ()
loadWindow, Ix1
wStart, Ix1
wEnd) <- forall (m :: * -> *) e a.
Monad m =>
(m () -> m ())
-> Array DW Ix1 e
-> (Ix1 -> e -> m a)
-> m (Ix1 -> Ix1 -> m (), Ix1, Ix1)
loadWithIx1 (forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler) Array DW Ix1 e
arr Ix1 -> e -> ST s ()
uWrite
let (Ix1
chunkWidth, Ix1
slackWidth) = (Ix1
wEnd forall a. Num a => a -> a -> a
- Ix1
wStart) forall a. Integral a => a -> a -> (a, a)
`quotRem` forall s a. Scheduler s a -> Ix1
numWorkers Scheduler s ()
scheduler
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
0 (forall a. Ord a => a -> a -> Bool
< forall s a. Scheduler s a -> Ix1
numWorkers Scheduler s ()
scheduler) (forall a. Num a => a -> a -> a
+ Ix1
1) forall a b. (a -> b) -> a -> b
$ \ !Ix1
wid ->
let !it' :: Ix1
it' = Ix1
wid forall a. Num a => a -> a -> a
* Ix1
chunkWidth forall a. Num a => a -> a -> a
+ Ix1
wStart
in Ix1 -> Ix1 -> ST s ()
loadWindow Ix1
it' (Ix1
it' forall a. Num a => a -> a -> a
+ Ix1
chunkWidth)
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ix1
slackWidth forall a. Ord a => a -> a -> Bool
> Ix1
0) forall a b. (a -> b) -> a -> b
$
let !itSlack :: Ix1
itSlack = forall s a. Scheduler s a -> Ix1
numWorkers Scheduler s ()
scheduler forall a. Num a => a -> a -> a
* Ix1
chunkWidth forall a. Num a => a -> a -> a
+ Ix1
wStart
in Ix1 -> Ix1 -> ST s ()
loadWindow Ix1
itSlack (Ix1
itSlack forall a. Num a => a -> a -> a
+ Ix1
slackWidth)
{-# INLINE iterArrayLinearST_ #-}
instance StrideLoad DW Ix1 e where
iterArrayLinearWithStrideST_ :: forall s.
Scheduler s ()
-> Stride Ix1
-> Sz Ix1
-> Array DW Ix1 e
-> (Ix1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ Scheduler s ()
scheduler Stride Ix1
stride Sz Ix1
sz Array DW Ix1 e
arr Ix1 -> e -> ST s ()
uWrite = do
((Ix1, Ix1) -> ST s ()
loadWindow, (Ix1
wStart, Ix1
wEnd)) <- forall (m :: * -> *) e a.
Monad m =>
(m () -> m ())
-> Array DW Ix1 e
-> Stride Ix1
-> Sz Ix1
-> (Ix1 -> e -> m a)
-> m ((Ix1, Ix1) -> m (), (Ix1, Ix1))
loadArrayWithIx1 (forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler) Array DW Ix1 e
arr Stride Ix1
stride Sz Ix1
sz Ix1 -> e -> ST s ()
uWrite
let (Ix1
chunkWidth, Ix1
slackWidth) = (Ix1
wEnd forall a. Num a => a -> a -> a
- Ix1
wStart) forall a. Integral a => a -> a -> (a, a)
`quotRem` forall s a. Scheduler s a -> Ix1
numWorkers Scheduler s ()
scheduler
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
0 (forall a. Ord a => a -> a -> Bool
< forall s a. Scheduler s a -> Ix1
numWorkers Scheduler s ()
scheduler) (forall a. Num a => a -> a -> a
+ Ix1
1) forall a b. (a -> b) -> a -> b
$ \ !Ix1
wid ->
let !it' :: Ix1
it' = Ix1
wid forall a. Num a => a -> a -> a
* Ix1
chunkWidth forall a. Num a => a -> a -> a
+ Ix1
wStart
in (Ix1, Ix1) -> ST s ()
loadWindow (Ix1
it', Ix1
it' forall a. Num a => a -> a -> a
+ Ix1
chunkWidth)
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ix1
slackWidth forall a. Ord a => a -> a -> Bool
> Ix1
0) forall a b. (a -> b) -> a -> b
$
let !itSlack :: Ix1
itSlack = forall s a. Scheduler s a -> Ix1
numWorkers Scheduler s ()
scheduler forall a. Num a => a -> a -> a
* Ix1
chunkWidth forall a. Num a => a -> a -> a
+ Ix1
wStart
in (Ix1, Ix1) -> ST s ()
loadWindow (Ix1
itSlack, Ix1
itSlack forall a. Num a => a -> a -> a
+ Ix1
slackWidth)
{-# INLINE iterArrayLinearWithStrideST_ #-}
loadArrayWithIx1
:: (Monad m)
=> (m () -> m ())
-> Array DW Ix1 e
-> Stride Ix1
-> Sz1
-> (Ix1 -> e -> m a)
-> m ((Ix1, Ix1) -> m (), (Ix1, Ix1))
loadArrayWithIx1 :: forall (m :: * -> *) e a.
Monad m =>
(m () -> m ())
-> Array DW Ix1 e
-> Stride Ix1
-> Sz Ix1
-> (Ix1 -> e -> m a)
-> m ((Ix1, Ix1) -> m (), (Ix1, Ix1))
loadArrayWithIx1 m () -> m ()
with (DWArray darr :: Array D Ix1 e
darr@(DArray Comp
_ Sz Ix1
arrSz PrefIndex Ix1 e
_) Maybe (Window Ix1 e)
mWindow) Stride Ix1
stride Sz Ix1
_ Ix1 -> e -> m a
uWrite = do
let Window Ix1
it Sz Ix1
wk Ix1 -> e
indexW Maybe Ix1
_ = forall a. a -> Maybe a -> a
fromMaybe forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window Ix1 e)
mWindow
wEnd :: Ix1
wEnd = Ix1
it forall a. Num a => a -> a -> a
+ forall ix. Sz ix -> ix
unSz Sz Ix1
wk
strideIx :: Ix1
strideIx = forall ix. Stride ix -> ix
unStride Stride Ix1
stride
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ Ix1
0 Ix1
it Ix1
strideIx forall a. Ord a => a -> a -> Bool
(<) forall a b. (a -> b) -> a -> b
$ \ !Ix1
i -> Ix1 -> e -> m a
uWrite (Ix1
i forall a. Integral a => a -> a -> a
`div` Ix1
strideIx) (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D Ix1 e
darr Ix1
i)
m () -> m ()
with forall a b. (a -> b) -> a -> b
$
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix1
stride Ix1
wEnd) (forall ix. Sz ix -> ix
unSz Sz Ix1
arrSz) Ix1
strideIx forall a. Ord a => a -> a -> Bool
(<) forall a b. (a -> b) -> a -> b
$ \ !Ix1
i ->
Ix1 -> e -> m a
uWrite (Ix1
i forall a. Integral a => a -> a -> a
`div` Ix1
strideIx) (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D Ix1 e
darr Ix1
i)
forall (m :: * -> *) a. Monad m => a -> m a
return
( \(Ix1
from, Ix1
to) ->
m () -> m ()
with forall a b. (a -> b) -> a -> b
$
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix1
stride Ix1
from) Ix1
to Ix1
strideIx forall a. Ord a => a -> a -> Bool
(<) forall a b. (a -> b) -> a -> b
$ \ !Ix1
i ->
Ix1 -> e -> m a
uWrite (Ix1
i forall a. Integral a => a -> a -> a
`div` Ix1
strideIx) (Ix1 -> e
indexW Ix1
i)
, (Ix1
it, Ix1
wEnd)
)
{-# INLINE loadArrayWithIx1 #-}
loadWithIx2
:: Monad m
=> (m () -> m ())
-> Array DW Ix2 t1
-> (Int -> t1 -> m ())
-> m (Ix2 -> m (), Ix2)
loadWithIx2 :: forall (m :: * -> *) t1.
Monad m =>
(m () -> m ())
-> Array DW Ix2 t1 -> (Ix1 -> t1 -> m ()) -> m (Ix2 -> m (), Ix2)
loadWithIx2 m () -> m ()
with Array DW Ix2 t1
arr Ix1 -> t1 -> m ()
uWrite = do
let DWArray Array D Ix2 t1
darr Maybe (Window Ix2 t1)
window = Array DW Ix2 t1
arr
Sz (Ix1
m :. Ix1
n) = forall ix e. Array D ix e -> Sz ix
dSize Array D Ix2 t1
darr
Window (Ix1
it :. Ix1
jt) (Sz (Ix1
wm :. Ix1
wn)) Ix2 -> t1
indexW Maybe Ix1
mUnrollHeight = forall a. a -> Maybe a -> a
fromMaybe forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window Ix2 t1)
window
Ix1
ib :. Ix1
jb = (Ix1
wm forall a. Num a => a -> a -> a
+ Ix1
it) Ix1 -> Ix1 -> Ix2
:. (Ix1
wn forall a. Num a => a -> a -> a
+ Ix1
jt)
!blockHeight :: Ix1
blockHeight = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Ix1
1 (forall a. Ord a => a -> a -> a
min Ix1
7 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => a -> a -> a
max Ix1
1) Maybe Ix1
mUnrollHeight
stride :: Stride Ix2
stride = forall ix. Index ix => Stride ix
oneStride
!sz :: Sz Ix2
sz = forall ix. Index ix => Stride ix -> Sz ix -> Sz ix
strideSize Stride Ix2
stride forall a b. (a -> b) -> a -> b
$ forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array DW Ix2 t1
arr
writeB :: Ix2 -> m ()
writeB !Ix2
ix = Ix1 -> t1 -> m ()
uWrite (forall ix. Index ix => Sz ix -> ix -> Ix1
toLinearIndex Sz Ix2
sz Ix2
ix) (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D Ix2 t1
darr Ix2
ix)
{-# INLINE writeB #-}
writeW :: Ix2 -> m ()
writeW !Ix2
ix = Ix1 -> t1 -> m ()
uWrite (forall ix. Index ix => Sz ix -> ix -> Ix1
toLinearIndex Sz Ix2
sz Ix2
ix) (Ix2 -> t1
indexW Ix2
ix)
{-# INLINE writeW #-}
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (Ix1
0 Ix1 -> Ix1 -> Ix2
:. Ix1
0) (Ix1
it Ix1 -> Ix1 -> Ix2
:. Ix1
n) (Ix1
1 Ix1 -> Ix1 -> Ix2
:. Ix1
1) forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (Ix1
ib Ix1 -> Ix1 -> Ix2
:. Ix1
0) (Ix1
m Ix1 -> Ix1 -> Ix2
:. Ix1
n) (Ix1
1 Ix1 -> Ix1 -> Ix2
:. Ix1
1) forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (Ix1
it Ix1 -> Ix1 -> Ix2
:. Ix1
0) (Ix1
ib Ix1 -> Ix1 -> Ix2
:. Ix1
jt) (Ix1
1 Ix1 -> Ix1 -> Ix2
:. Ix1
1) forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (Ix1
it Ix1 -> Ix1 -> Ix2
:. Ix1
jb) (Ix1
ib Ix1 -> Ix1 -> Ix2
:. Ix1
n) (Ix1
1 Ix1 -> Ix1 -> Ix2
:. Ix1
1) forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
let f :: Ix2 -> m ()
f (Ix1
it' :. Ix1
ib') = m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
Monad m =>
Ix1 -> Ix2 -> Ix2 -> Ix1 -> (Ix2 -> m ()) -> m ()
unrollAndJam Ix1
blockHeight (Ix1
it' Ix1 -> Ix1 -> Ix2
:. Ix1
jt) (Ix1
ib' Ix1 -> Ix1 -> Ix2
:. Ix1
jb) Ix1
1 Ix2 -> m ()
writeW
{-# INLINE f #-}
forall (m :: * -> *) a. Monad m => a -> m a
return (Ix2 -> m ()
f, Ix1
it Ix1 -> Ix1 -> Ix2
:. Ix1
ib)
{-# INLINE loadWithIx2 #-}
loadArrayWithIx2
:: Monad m
=> (m () -> m ())
-> Array DW Ix2 e
-> Stride Ix2
-> Sz2
-> (Int -> e -> m ())
-> m (Ix2 -> m (), Ix2)
loadArrayWithIx2 :: forall (m :: * -> *) e.
Monad m =>
(m () -> m ())
-> Array DW Ix2 e
-> Stride Ix2
-> Sz Ix2
-> (Ix1 -> e -> m ())
-> m (Ix2 -> m (), Ix2)
loadArrayWithIx2 m () -> m ()
with Array DW Ix2 e
arr Stride Ix2
stride Sz Ix2
sz Ix1 -> e -> m ()
uWrite = do
let DWArray Array D Ix2 e
darr Maybe (Window Ix2 e)
window = Array DW Ix2 e
arr
Sz (Ix1
m :. Ix1
n) = forall ix e. Array D ix e -> Sz ix
dSize Array D Ix2 e
darr
Window (Ix1
it :. Ix1
jt) (Sz (Ix1
wm :. Ix1
wn)) Ix2 -> e
indexW Maybe Ix1
mUnrollHeight = forall a. a -> Maybe a -> a
fromMaybe forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window Ix2 e)
window
Ix1
ib :. Ix1
jb = (Ix1
wm forall a. Num a => a -> a -> a
+ Ix1
it) Ix1 -> Ix1 -> Ix2
:. (Ix1
wn forall a. Num a => a -> a -> a
+ Ix1
jt)
!blockHeight :: Ix1
blockHeight = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Ix1
1 (forall a. Ord a => a -> a -> a
min Ix1
7 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => a -> a -> a
max Ix1
1) Maybe Ix1
mUnrollHeight
strideIx :: Ix2
strideIx@(Ix1
is :. Ix1
js) = forall ix. Stride ix -> ix
unStride Stride Ix2
stride
writeB :: Ix2 -> m ()
writeB !Ix2
ix = Ix1 -> e -> m ()
uWrite (forall ix. Index ix => Stride ix -> Sz ix -> ix -> Ix1
toLinearIndexStride Stride Ix2
stride Sz Ix2
sz Ix2
ix) (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D Ix2 e
darr Ix2
ix)
{-# INLINE writeB #-}
writeW :: Ix2 -> m ()
writeW !Ix2
ix = Ix1 -> e -> m ()
uWrite (forall ix. Index ix => Stride ix -> Sz ix -> ix -> Ix1
toLinearIndexStride Stride Ix2
stride Sz Ix2
sz Ix2
ix) (Ix2 -> e
indexW Ix2
ix)
{-# INLINE writeW #-}
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (Ix1
0 Ix1 -> Ix1 -> Ix2
:. Ix1
0) (Ix1
it Ix1 -> Ix1 -> Ix2
:. Ix1
n) Ix2
strideIx forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix2
stride (Ix1
ib Ix1 -> Ix1 -> Ix2
:. Ix1
0)) (Ix1
m Ix1 -> Ix1 -> Ix2
:. Ix1
n) Ix2
strideIx forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix2
stride (Ix1
it Ix1 -> Ix1 -> Ix2
:. Ix1
0)) (Ix1
ib Ix1 -> Ix1 -> Ix2
:. Ix1
jt) Ix2
strideIx forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
m () -> m ()
with forall a b. (a -> b) -> a -> b
$ forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix2
stride (Ix1
it Ix1 -> Ix1 -> Ix2
:. Ix1
jb)) (Ix1
ib Ix1 -> Ix1 -> Ix2
:. Ix1
n) Ix2
strideIx forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
let f :: Ix2 -> m ()
f (Ix1
it' :. Ix1
ib')
| Ix1
is forall a. Ord a => a -> a -> Bool
> Ix1
1 Bool -> Bool -> Bool
|| Ix1
blockHeight forall a. Ord a => a -> a -> Bool
<= Ix1
1 =
forall ix (f :: * -> *) a.
(Index ix, Applicative f) =>
ix -> ix -> ix -> (Ix1 -> Ix1 -> Bool) -> (ix -> f a) -> f ()
iterA_ (forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix2
stride (Ix1
it' Ix1 -> Ix1 -> Ix2
:. Ix1
jt)) (Ix1
ib' Ix1 -> Ix1 -> Ix2
:. Ix1
jb) Ix2
strideIx forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeW
| Bool
otherwise =
forall (m :: * -> *).
Monad m =>
Ix1 -> Ix2 -> Ix2 -> Ix1 -> (Ix2 -> m ()) -> m ()
unrollAndJam Ix1
blockHeight (forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix2
stride (Ix1
it' Ix1 -> Ix1 -> Ix2
:. Ix1
jt)) (Ix1
ib' Ix1 -> Ix1 -> Ix2
:. Ix1
jb) Ix1
js Ix2 -> m ()
writeW
{-# INLINE f #-}
forall (m :: * -> *) a. Monad m => a -> m a
return (m () -> m ()
with forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix2 -> m ()
f, Ix1
it Ix1 -> Ix1 -> Ix2
:. Ix1
ib)
{-# INLINE loadArrayWithIx2 #-}
loadWindowIx2 :: Monad m => Int -> (Ix2 -> m ()) -> Ix2 -> m ()
loadWindowIx2 :: forall (m :: * -> *).
Monad m =>
Ix1 -> (Ix2 -> m ()) -> Ix2 -> m ()
loadWindowIx2 Ix1
nWorkers Ix2 -> m ()
loadWindow (Ix1
it :. Ix1
ib) = do
let !(Ix1
chunkHeight, Ix1
slackHeight) = (Ix1
ib forall a. Num a => a -> a -> a
- Ix1
it) forall a. Integral a => a -> a -> (a, a)
`quotRem` Ix1
nWorkers
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
0 (forall a. Ord a => a -> a -> Bool
< Ix1
nWorkers) (forall a. Num a => a -> a -> a
+ Ix1
1) forall a b. (a -> b) -> a -> b
$ \ !Ix1
wid ->
let !it' :: Ix1
it' = Ix1
wid forall a. Num a => a -> a -> a
* Ix1
chunkHeight forall a. Num a => a -> a -> a
+ Ix1
it
in Ix2 -> m ()
loadWindow (Ix1
it' Ix1 -> Ix1 -> Ix2
:. (Ix1
it' forall a. Num a => a -> a -> a
+ Ix1
chunkHeight))
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ix1
slackHeight forall a. Ord a => a -> a -> Bool
> Ix1
0) forall a b. (a -> b) -> a -> b
$
let !itSlack :: Ix1
itSlack = Ix1
nWorkers forall a. Num a => a -> a -> a
* Ix1
chunkHeight forall a. Num a => a -> a -> a
+ Ix1
it
in Ix2 -> m ()
loadWindow (Ix1
itSlack Ix1 -> Ix1 -> Ix2
:. (Ix1
itSlack forall a. Num a => a -> a -> a
+ Ix1
slackHeight))
{-# INLINE loadWindowIx2 #-}
instance Load DW Ix2 e where
makeArray :: Comp -> Sz Ix2 -> (Ix2 -> e) -> Array DW Ix2 e
makeArray Comp
c Sz Ix2
sz Ix2 -> e
f = forall ix e. Array D ix e -> Maybe (Window ix e) -> Array DW ix e
DWArray (forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
c Sz Ix2
sz Ix2 -> e
f) forall a. Maybe a
Nothing
{-# INLINE makeArray #-}
iterArrayLinearST_ :: forall s.
Scheduler s ()
-> Array DW Ix2 e -> (Ix1 -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler Array DW Ix2 e
arr Ix1 -> e -> ST s ()
uWrite =
forall (m :: * -> *) t1.
Monad m =>
(m () -> m ())
-> Array DW Ix2 t1 -> (Ix1 -> t1 -> m ()) -> m (Ix2 -> m (), Ix2)
loadWithIx2 (forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler) Array DW Ix2 e
arr Ix1 -> e -> ST s ()
uWrite
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall (m :: * -> *).
Monad m =>
Ix1 -> (Ix2 -> m ()) -> Ix2 -> m ()
loadWindowIx2 (forall s a. Scheduler s a -> Ix1
numWorkers Scheduler s ()
scheduler))
{-# INLINE iterArrayLinearST_ #-}
instance StrideLoad DW Ix2 e where
iterArrayLinearWithStrideST_ :: forall s.
Scheduler s ()
-> Stride Ix2
-> Sz Ix2
-> Array DW Ix2 e
-> (Ix1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ Scheduler s ()
scheduler Stride Ix2
stride Sz Ix2
sz Array DW Ix2 e
arr Ix1 -> e -> ST s ()
uWrite =
forall (m :: * -> *) e.
Monad m =>
(m () -> m ())
-> Array DW Ix2 e
-> Stride Ix2
-> Sz Ix2
-> (Ix1 -> e -> m ())
-> m (Ix2 -> m (), Ix2)
loadArrayWithIx2 (forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler) Array DW Ix2 e
arr Stride Ix2
stride Sz Ix2
sz Ix1 -> e -> ST s ()
uWrite
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall (m :: * -> *).
Monad m =>
Ix1 -> (Ix2 -> m ()) -> Ix2 -> m ()
loadWindowIx2 (forall s a. Scheduler s a -> Ix1
numWorkers Scheduler s ()
scheduler))
{-# INLINE iterArrayLinearWithStrideST_ #-}
instance (Index (IxN n), Load DW (Ix (n - 1)) e) => Load DW (IxN n) e where
makeArray :: Comp -> Sz (IxN n) -> (IxN n -> e) -> Array DW (IxN n) e
makeArray Comp
c Sz (IxN n)
sz IxN n -> e
f = forall ix e. Array D ix e -> Maybe (Window ix e) -> Array DW ix e
DWArray (forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
c Sz (IxN n)
sz IxN n -> e
f) forall a. Maybe a
Nothing
{-# INLINE makeArray #-}
iterArrayLinearST_ :: forall s.
Scheduler s ()
-> Array DW (IxN n) e -> (Ix1 -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ = forall ix e s.
(Index ix, Load DW (Lower ix) e) =>
Scheduler s () -> Array DW ix e -> (Ix1 -> e -> ST s ()) -> ST s ()
loadWithIxN
{-# INLINE iterArrayLinearST_ #-}
instance (Index (IxN n), StrideLoad DW (Ix (n - 1)) e) => StrideLoad DW (IxN n) e where
iterArrayLinearWithStrideST_ :: forall s.
Scheduler s ()
-> Stride (IxN n)
-> Sz (IxN n)
-> Array DW (IxN n) e
-> (Ix1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ = forall ix e s.
(Index ix, StrideLoad DW (Lower ix) e) =>
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array DW ix e
-> (Ix1 -> e -> ST s ())
-> ST s ()
loadArrayWithIxN
{-# INLINE iterArrayLinearWithStrideST_ #-}
loadArrayWithIxN
:: (Index ix, StrideLoad DW (Lower ix) e)
=> Scheduler s ()
-> Stride ix
-> Sz ix
-> Array DW ix e
-> (Int -> e -> ST s ())
-> ST s ()
loadArrayWithIxN :: forall ix e s.
(Index ix, StrideLoad DW (Lower ix) e) =>
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array DW ix e
-> (Ix1 -> e -> ST s ())
-> ST s ()
loadArrayWithIxN Scheduler s ()
scheduler Stride ix
stride Sz ix
szResult Array DW ix e
arr Ix1 -> e -> ST s ()
uWrite = do
let DWArray Array D ix e
darr Maybe (Window ix e)
window = Array DW ix e
arr
Window{ix
windowStart :: ix
windowStart :: forall ix e. Window ix e -> ix
windowStart, Sz ix
windowSize :: Sz ix
windowSize :: forall ix e. Window ix e -> Sz ix
windowSize, ix -> e
windowIndex :: ix -> e
windowIndex :: forall ix e. Window ix e -> ix -> e
windowIndex, Maybe Ix1
windowUnrollIx2 :: Maybe Ix1
windowUnrollIx2 :: forall ix e. Window ix e -> Maybe Ix1
windowUnrollIx2} = forall a. a -> Maybe a -> a
fromMaybe forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window ix e)
window
!(Sz Ix1
headSourceSize, Sz (Lower ix)
lowerSourceSize) = forall ix. Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz (forall ix e. Array D ix e -> Sz ix
dSize Array D ix e
darr)
!lowerSize :: Sz (Lower ix)
lowerSize = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz Sz ix
szResult
!(Ix1
s, Lower ix
lowerStrideIx) = forall ix. Index ix => ix -> (Ix1, Lower ix)
unconsDim forall a b. (a -> b) -> a -> b
$ forall ix. Stride ix -> ix
unStride Stride ix
stride
!(Ix1
curWindowStart, Lower ix
lowerWindowStart) = forall ix. Index ix => ix -> (Ix1, Lower ix)
unconsDim ix
windowStart
!(Sz Ix1
headWindowSz, Sz (Lower ix)
tailWindowSz) = forall ix. Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz Sz ix
windowSize
!curWindowEnd :: Ix1
curWindowEnd = Ix1
curWindowStart forall a. Num a => a -> a -> a
+ forall ix. Sz ix -> ix
unSz Sz Ix1
headWindowSz
!pageElements :: Ix1
pageElements = forall ix. Index ix => Sz ix -> Ix1
totalElem Sz (Lower ix)
lowerSize
mkLowerWindow :: Ix1 -> Window (Lower ix) e
mkLowerWindow Ix1
i =
Window
{ windowStart :: Lower ix
windowStart = Lower ix
lowerWindowStart
, windowSize :: Sz (Lower ix)
windowSize = Sz (Lower ix)
tailWindowSz
, windowIndex :: Lower ix -> e
windowIndex = ix -> e
windowIndex forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Ix1 -> Lower ix -> ix
consDim Ix1
i
, windowUnrollIx2 :: Maybe Ix1
windowUnrollIx2 = Maybe Ix1
windowUnrollIx2
}
mkLowerArray :: Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> Array DW (Lower ix) e
mkLowerArray Maybe (Ix1 -> Window (Lower ix) e)
mw Ix1
i =
DWArray
{ dwArray :: Array D (Lower ix) e
dwArray =
Array D ix e
darr
{ dComp :: Comp
dComp = Comp
Seq
, dSize :: Sz (Lower ix)
dSize = Sz (Lower ix)
lowerSourceSize
, dPrefIndex :: PrefIndex (Lower ix) e
dPrefIndex = forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D ix e
darr forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Ix1 -> Lower ix -> ix
consDim Ix1
i)
}
, dwWindow :: Maybe (Window (Lower ix) e)
dwWindow = (forall a b. (a -> b) -> a -> b
$ Ix1
i) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Ix1 -> Window (Lower ix) e)
mw
}
loadLower :: Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> ST s ()
loadLower Maybe (Ix1 -> Window (Lower ix) e)
mw !Ix1
i =
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
(forall ix. Index ix => ix -> Stride ix
Stride Lower ix
lowerStrideIx)
Sz (Lower ix)
lowerSize
(Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> Array DW (Lower ix) e
mkLowerArray Maybe (Ix1 -> Window (Lower ix) e)
mw Ix1
i)
(\Ix1
k -> Ix1 -> e -> ST s ()
uWrite (Ix1
k forall a. Num a => a -> a -> a
+ Ix1
pageElements forall a. Num a => a -> a -> a
* (Ix1
i forall a. Integral a => a -> a -> a
`div` Ix1
s)))
{-# NOINLINE loadLower #-}
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
0 (forall a. Ord a => a -> a -> Bool
< forall ix. Index ix => ix -> Ix1
headDim ix
windowStart) (forall a. Num a => a -> a -> a
+ Ix1
s) (Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> ST s ()
loadLower forall a. Maybe a
Nothing)
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_
(forall ix. Index ix => Stride ix -> ix -> ix
strideStart (forall ix. Index ix => ix -> Stride ix
Stride Ix1
s) Ix1
curWindowStart)
(forall a. Ord a => a -> a -> Bool
< Ix1
curWindowEnd)
(forall a. Num a => a -> a -> a
+ Ix1
s)
(Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> ST s ()
loadLower (forall a. a -> Maybe a
Just Ix1 -> Window (Lower ix) e
mkLowerWindow))
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ (forall ix. Index ix => Stride ix -> ix -> ix
strideStart (forall ix. Index ix => ix -> Stride ix
Stride Ix1
s) Ix1
curWindowEnd) (forall a. Ord a => a -> a -> Bool
< forall ix. Sz ix -> ix
unSz Sz Ix1
headSourceSize) (forall a. Num a => a -> a -> a
+ Ix1
s) (Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> ST s ()
loadLower forall a. Maybe a
Nothing)
{-# INLINE loadArrayWithIxN #-}
loadWithIxN
:: (Index ix, Load DW (Lower ix) e)
=> Scheduler s ()
-> Array DW ix e
-> (Int -> e -> ST s ())
-> ST s ()
loadWithIxN :: forall ix e s.
(Index ix, Load DW (Lower ix) e) =>
Scheduler s () -> Array DW ix e -> (Ix1 -> e -> ST s ()) -> ST s ()
loadWithIxN Scheduler s ()
scheduler Array DW ix e
arr Ix1 -> e -> ST s ()
uWrite = do
let DWArray Array D ix e
darr Maybe (Window ix e)
window = Array DW ix e
arr
Window{ix
windowStart :: ix
windowStart :: forall ix e. Window ix e -> ix
windowStart, Sz ix
windowSize :: Sz ix
windowSize :: forall ix e. Window ix e -> Sz ix
windowSize, ix -> e
windowIndex :: ix -> e
windowIndex :: forall ix e. Window ix e -> ix -> e
windowIndex, Maybe Ix1
windowUnrollIx2 :: Maybe Ix1
windowUnrollIx2 :: forall ix e. Window ix e -> Maybe Ix1
windowUnrollIx2} = forall a. a -> Maybe a -> a
fromMaybe forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window ix e)
window
!(Sz Ix1
si, Sz (Lower ix)
szL) = forall ix. Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz (forall ix e. Array D ix e -> Sz ix
dSize Array D ix e
darr)
!windowEnd :: ix
windowEnd = forall ix. Index ix => (Ix1 -> Ix1 -> Ix1) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(+) ix
windowStart (forall ix. Sz ix -> ix
unSz Sz ix
windowSize)
!(Ix1
t, Lower ix
windowStartL) = forall ix. Index ix => ix -> (Ix1, Lower ix)
unconsDim ix
windowStart
!pageElements :: Ix1
pageElements = forall ix. Index ix => Sz ix -> Ix1
totalElem Sz (Lower ix)
szL
mkLowerWindow :: Ix1 -> Window (Lower ix) e
mkLowerWindow Ix1
i =
Window
{ windowStart :: Lower ix
windowStart = Lower ix
windowStartL
, windowSize :: Sz (Lower ix)
windowSize = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix))
unconsSz Sz ix
windowSize
, windowIndex :: Lower ix -> e
windowIndex = ix -> e
windowIndex forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Ix1 -> Lower ix -> ix
consDim Ix1
i
, windowUnrollIx2 :: Maybe Ix1
windowUnrollIx2 = Maybe Ix1
windowUnrollIx2
}
mkLowerArray :: Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> Array DW (Lower ix) e
mkLowerArray Maybe (Ix1 -> Window (Lower ix) e)
mw Ix1
i =
DWArray
{ dwArray :: Array D (Lower ix) e
dwArray =
Array D ix e
darr{dComp :: Comp
dComp = Comp
Seq, dSize :: Sz (Lower ix)
dSize = Sz (Lower ix)
szL, dPrefIndex :: PrefIndex (Lower ix) e
dPrefIndex = forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D ix e
darr forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index ix => Ix1 -> Lower ix -> ix
consDim Ix1
i)}
, dwWindow :: Maybe (Window (Lower ix) e)
dwWindow = (forall a b. (a -> b) -> a -> b
$ Ix1
i) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Ix1 -> Window (Lower ix) e)
mw
}
loadLower :: Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> ST s ()
loadLower Maybe (Ix1 -> Window (Lower ix) e)
mw !Ix1
i =
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
forall r ix e s.
Load r ix e =>
Scheduler s () -> Array r ix e -> (Ix1 -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler (Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> Array DW (Lower ix) e
mkLowerArray Maybe (Ix1 -> Window (Lower ix) e)
mw Ix1
i) (\Ix1
k -> Ix1 -> e -> ST s ()
uWrite (Ix1
k forall a. Num a => a -> a -> a
+ Ix1
pageElements forall a. Num a => a -> a -> a
* Ix1
i))
{-# NOINLINE loadLower #-}
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
0 (forall a. Ord a => a -> a -> Bool
< forall ix. Index ix => ix -> Ix1
headDim ix
windowStart) (forall a. Num a => a -> a -> a
+ Ix1
1) (Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> ST s ()
loadLower forall a. Maybe a
Nothing)
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
t (forall a. Ord a => a -> a -> Bool
< forall ix. Index ix => ix -> Ix1
headDim ix
windowEnd) (forall a. Num a => a -> a -> a
+ Ix1
1) (Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> ST s ()
loadLower (forall a. a -> Maybe a
Just Ix1 -> Window (Lower ix) e
mkLowerWindow))
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ (forall ix. Index ix => ix -> Ix1
headDim ix
windowEnd) (forall a. Ord a => a -> a -> Bool
< forall ix. Sz ix -> ix
unSz Sz Ix1
si) (forall a. Num a => a -> a -> a
+ Ix1
1) (Maybe (Ix1 -> Window (Lower ix) e) -> Ix1 -> ST s ()
loadLower forall a. Maybe a
Nothing)
{-# INLINE loadWithIxN #-}
unrollAndJam
:: Monad m
=> Int
-> Ix2
-> Ix2
-> Int
-> (Ix2 -> m ())
-> m ()
unrollAndJam :: forall (m :: * -> *).
Monad m =>
Ix1 -> Ix2 -> Ix2 -> Ix1 -> (Ix2 -> m ()) -> m ()
unrollAndJam !Ix1
bH (Ix1
it :. Ix1
jt) (Ix1
ib :. Ix1
jb) Ix1
js Ix2 -> m ()
f = do
let f2 :: Ix2 -> m ()
f2 (Ix1
i :. Ix1
j) = Ix2 -> m ()
f (Ix1
i Ix1 -> Ix1 -> Ix2
:. Ix1
j) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f ((Ix1
i forall a. Num a => a -> a -> a
+ Ix1
1) Ix1 -> Ix1 -> Ix2
:. Ix1
j)
let f3 :: Ix2 -> m ()
f3 (Ix1
i :. Ix1
j) = Ix2 -> m ()
f (Ix1
i Ix1 -> Ix1 -> Ix2
:. Ix1
j) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f2 ((Ix1
i forall a. Num a => a -> a -> a
+ Ix1
1) Ix1 -> Ix1 -> Ix2
:. Ix1
j)
let f4 :: Ix2 -> m ()
f4 (Ix1
i :. Ix1
j) = Ix2 -> m ()
f (Ix1
i Ix1 -> Ix1 -> Ix2
:. Ix1
j) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f3 ((Ix1
i forall a. Num a => a -> a -> a
+ Ix1
1) Ix1 -> Ix1 -> Ix2
:. Ix1
j)
let f5 :: Ix2 -> m ()
f5 (Ix1
i :. Ix1
j) = Ix2 -> m ()
f (Ix1
i Ix1 -> Ix1 -> Ix2
:. Ix1
j) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f4 ((Ix1
i forall a. Num a => a -> a -> a
+ Ix1
1) Ix1 -> Ix1 -> Ix2
:. Ix1
j)
let f6 :: Ix2 -> m ()
f6 (Ix1
i :. Ix1
j) = Ix2 -> m ()
f (Ix1
i Ix1 -> Ix1 -> Ix2
:. Ix1
j) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f5 ((Ix1
i forall a. Num a => a -> a -> a
+ Ix1
1) Ix1 -> Ix1 -> Ix2
:. Ix1
j)
let f7 :: Ix2 -> m ()
f7 (Ix1
i :. Ix1
j) = Ix2 -> m ()
f (Ix1
i Ix1 -> Ix1 -> Ix2
:. Ix1
j) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f6 ((Ix1
i forall a. Num a => a -> a -> a
+ Ix1
1) Ix1 -> Ix1 -> Ix2
:. Ix1
j)
let f' :: Ix2 -> m ()
f' = case Ix1
bH of
Ix1
1 -> Ix2 -> m ()
f
Ix1
2 -> Ix2 -> m ()
f2
Ix1
3 -> Ix2 -> m ()
f3
Ix1
4 -> Ix2 -> m ()
f4
Ix1
5 -> Ix2 -> m ()
f5
Ix1
6 -> Ix2 -> m ()
f6
Ix1
_ -> Ix2 -> m ()
f7
let !ibS :: Ix1
ibS = Ix1
ib forall a. Num a => a -> a -> a
- ((Ix1
ib forall a. Num a => a -> a -> a
- Ix1
it) forall a. Integral a => a -> a -> a
`mod` Ix1
bH)
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
it (forall a. Ord a => a -> a -> Bool
< Ix1
ibS) (forall a. Num a => a -> a -> a
+ Ix1
bH) forall a b. (a -> b) -> a -> b
$ \ !Ix1
i ->
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
jt (forall a. Ord a => a -> a -> Bool
< Ix1
jb) (forall a. Num a => a -> a -> a
+ Ix1
js) forall a b. (a -> b) -> a -> b
$ \ !Ix1
j ->
Ix2 -> m ()
f' (Ix1
i Ix1 -> Ix1 -> Ix2
:. Ix1
j)
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
ibS (forall a. Ord a => a -> a -> Bool
< Ix1
ib) (forall a. Num a => a -> a -> a
+ Ix1
1) forall a b. (a -> b) -> a -> b
$ \ !Ix1
i ->
forall (f :: * -> *) a.
Applicative f =>
Ix1 -> (Ix1 -> Bool) -> (Ix1 -> Ix1) -> (Ix1 -> f a) -> f ()
loopA_ Ix1
jt (forall a. Ord a => a -> a -> Bool
< Ix1
jb) (forall a. Num a => a -> a -> a
+ Ix1
js) forall a b. (a -> b) -> a -> b
$ \ !Ix1
j ->
Ix2 -> m ()
f (Ix1
i Ix1 -> Ix1 -> Ix2
:. Ix1
j)
{-# INLINE unrollAndJam #-}