{-# 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
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
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

-- | Delayed Windowed Array representation.
data DW = DW

data Window ix e = Window
  { forall ix e. Window ix e -> ix
windowStart :: !ix
  -- ^ Index of where window will start at.
  , forall ix e. Window ix e -> Sz ix
windowSize :: !(Sz ix)
  -- ^ Size of the window
  , forall ix e. Window ix e -> ix -> e
windowIndex :: ix -> e
  -- ^ Indexing function for the window
  , forall ix e. Window ix e -> Maybe Ix1
windowUnrollIx2 :: !(Maybe Int)
  -- ^ Setting this value during stencil application improves cache
  -- utilization by unrolling the loop for Ix2 and higher dimensions.
  -- Has no affect on arrays with one dimension.
  }

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 #-}

--
--
-- @since 0.3.0
-- _makeWindowedArrayM
--   :: Source r ix e
--   => Array r ix e -- ^ Source array that will have a window inserted into it
--   -> ix -- ^ Start index for the window
--   -> Sz ix -- ^ Size of the window
--   -> (ix -> e) -- ^ Inside window indexing function
--   -> Array DW ix e
-- _makeWindowedArrayM !arr !windowStart !windowSize windowIndex
--   | not (isSafeIndex sz windowStart) =
--     error $
--     "makeWindowedArray: Incorrect window starting index: (" ++
--     show windowStart ++ ") for array size: (" ++ show (size arr) ++ ")"
--   | totalElem windowSize == 0 =
--     error $
--     "makeWindowedArray: Window can't hold any elements with this size: (" ++ show windowSize ++ ")"
--   | not
--      (isSafeIndex
--         (Sz (liftIndex (+ 1) (unSz sz)))
--         (liftIndex2 (+) windowStart (unSz windowSize))) =
--     error $
--     "makeWindowedArray: Incorrect window size: (" ++
--     show windowSize ++
--     ") and/or starting index: (" ++
--     show windowStart ++ ") for array size: (" ++ show (size arr) ++ ")"
--   | otherwise =
--     DWArray {dwArray = delay arr, dwWindow = Just $! Window {..}}
--   where
--     windowUnrollIx2 = Nothing
--     sz = size arr
-- {-# INLINE _makeWindowedArrayM #-}

-- | Construct a delayed windowed array by supply a separate element producing function for the
-- interior of an array. This is very usful for stencil mapping, where interior function does not
-- perform boundary checks, thus significantly speeding up computation process.
--
-- @since 0.1.3
makeWindowedArray
  :: (Index ix, Source r e)
  => Array r ix e
  -- ^ Source array that will have a window inserted into it
  -> ix
  -- ^ Start index for the window
  -> Sz ix
  -- ^ Size of the window
  -> (ix -> e)
  -- ^ Indexing function foto use inside window
  -> 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 #-}

-- | Inserts a `Window` into a delayed array while scaling the window down if it doesn't fit inside
-- that array.
--
-- @since 0.3.0
insertWindow
  :: Index ix
  => Array D ix e
  -- ^ Source array that will have a window inserted into it
  -> Window ix e
  -- ^ Window to place inside the delayed array
  -> 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 #-}

-- | Get the `Window` from a windowed array.
--
-- @since 0.2.1
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 #-}

-- | Drop the `Window` from a windowed array.
--
-- @since 0.3.0
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 =
            -- Turn off unrolling for vertical strides
            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
  -- ^ Block height
  -> Ix2
  -- ^ Top corner
  -> Ix2
  -- ^ Bottom corner
  -> Int
  -- ^ Column Stride
  -> (Ix2 -> m ())
  -- ^ Writing function
  -> 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 #-}

-- TODO: Implement Hilbert curve