{-# 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-2021
-- 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 { Window ix e -> ix
windowStart     :: !ix
                          -- ^ Index of where window will start at.
                          , Window ix e -> Sz ix
windowSize      :: !(Sz ix)
                          -- ^ Size of the window
                          , Window ix e -> ix -> e
windowIndex     :: ix -> e
                          -- ^ Indexing function for the window
                          , Window ix e -> Maybe Int
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 :: (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 (a -> b) -> (ix -> a) -> ix -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> a
windowIndex }

data instance Array DW ix e = DWArray { Array DW ix e -> Array D ix e
dwArray :: !(Array D 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 :: Int -> Array DW ix e -> ShowS
showsPrec = (Array DW ix e -> Array B ix e) -> Int -> Array DW ix e -> ShowS
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) -> Int -> Array r ix e -> ShowS
showsArrayPrec (B -> Array DW ix e -> Array B ix e
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 = [Array DW ix e] -> ShowS
forall arr. Show arr => [arr] -> ShowS
showArrayList

instance Strategy DW where
  setComp :: Comp -> Array DW ix e -> Array DW ix e
setComp Comp
c Array DW ix e
arr = Array DW ix e
R:ArrayDWixe ix e
arr { dwArray :: Array D ix e
dwArray = (Array DW ix e -> Array D ix e
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 :: Array DW ix e -> Comp
getComp = Array D ix e -> Comp
forall ix e. Array D ix e -> Comp
dComp (Array D ix e -> Comp)
-> (Array DW ix e -> Array D ix e) -> Array DW ix e -> Comp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DW ix e -> Array D ix e
forall ix e. Array DW ix e -> Array D ix e
dwArray
  {-# INLINE getComp #-}


instance Functor (Array DW ix) where
  fmap :: (a -> b) -> Array DW ix a -> Array DW ix b
fmap a -> b
f arr :: Array DW ix a
arr@DWArray{dwArray, dwWindow} =
    Array DW ix a
R:ArrayDWixe ix a
arr
    { dwArray :: Array D ix b
dwArray = (a -> b) -> Array D ix a -> Array D ix b
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 = (a -> b) -> Window ix a -> Window ix b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (Window ix a -> Window ix b)
-> Maybe (Window ix a) -> Maybe (Window ix b)
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 :: 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 =
  Array D ix e -> Window ix e -> Array DW ix e
forall ix e.
Index ix =>
Array D ix e -> Window ix e -> Array DW ix e
insertWindow (Array r ix e -> Array D ix e
forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Array D ix e
delay Array r ix e
arr) (Window ix e -> Array DW ix e) -> Window ix e -> Array DW ix e
forall a b. (a -> b) -> a -> b
$
  Window :: forall ix e. ix -> Sz ix -> (ix -> e) -> Maybe Int -> Window ix e
Window {windowStart :: ix
windowStart = ix
wStart, windowSize :: Sz ix
windowSize = Sz ix
wSize, windowIndex :: ix -> e
windowIndex = ix -> e
wIndex, windowUnrollIx2 :: Maybe Int
windowUnrollIx2 = Maybe Int
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 :: Array D ix e -> Window ix e -> Array DW ix e
insertWindow !Array D ix e
arr !Window ix e
window =
  DWArray :: forall ix e. Array D ix e -> Maybe (Window ix e) -> Array DW ix e
DWArray
    { dwArray :: Array D ix e
dwArray = Array D ix e -> Array D ix e
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 =
        Window ix e -> Maybe (Window ix e)
forall a. a -> Maybe a
Just (Window ix e -> Maybe (Window ix e))
-> Window ix e -> Maybe (Window ix e)
forall a b. (a -> b) -> a -> b
$!
        Window :: forall ix e. ix -> Sz ix -> (ix -> e) -> Maybe Int -> Window ix e
Window
          { windowStart :: ix
windowStart = ix
wStart'
          , windowSize :: Sz ix
windowSize = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min ix
wSize ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) ix
sz ix
wStart'))
          , windowIndex :: ix -> e
windowIndex = ix -> e
wIndex
          , windowUnrollIx2 :: Maybe Int
windowUnrollIx2 = Maybe Int
wUnrollIx2
          }
    }
  where
    wStart' :: ix
wStart' = Sz ix -> ix
forall ix. Sz ix -> ix
unSz (ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min ix
wStart ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1) ix
sz)))
    Sz ix
sz = Array D ix e -> Sz ix
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 Int
windowUnrollIx2 = Maybe Int
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 :: Array DW ix e -> Maybe (Window ix e)
getWindow = Array DW ix e -> Maybe (Window ix e)
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 :: Array DW ix e -> Array D ix e
dropWindow = Array DW ix e -> Array D ix e
forall ix e. Array DW ix e -> Array D ix e
dwArray
{-# INLINE dropWindow #-}


zeroWindow :: Index ix => Window ix e
zeroWindow :: Window ix e
zeroWindow = ix -> Sz ix -> (ix -> e) -> Maybe Int -> Window ix e
forall ix e. ix -> Sz ix -> (ix -> e) -> Maybe Int -> Window ix e
Window ix
forall ix. Index ix => ix
zeroIndex Sz ix
forall ix. Index ix => Sz ix
zeroSz ix -> e
forall a. a
windowError Maybe Int
forall a. Maybe a
Nothing
{-# INLINE zeroWindow #-}

data EmptyWindowException = EmptyWindowException deriving (EmptyWindowException -> EmptyWindowException -> Bool
(EmptyWindowException -> EmptyWindowException -> Bool)
-> (EmptyWindowException -> EmptyWindowException -> Bool)
-> Eq EmptyWindowException
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, Int -> EmptyWindowException -> ShowS
[EmptyWindowException] -> ShowS
EmptyWindowException -> String
(Int -> EmptyWindowException -> ShowS)
-> (EmptyWindowException -> String)
-> ([EmptyWindowException] -> ShowS)
-> Show EmptyWindowException
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [EmptyWindowException] -> ShowS
$cshowList :: [EmptyWindowException] -> ShowS
show :: EmptyWindowException -> String
$cshow :: EmptyWindowException -> String
showsPrec :: Int -> EmptyWindowException -> ShowS
$cshowsPrec :: Int -> EmptyWindowException -> ShowS
Show)

instance Exception EmptyWindowException where

  displayException :: EmptyWindowException -> String
displayException EmptyWindowException
_ = String
"Index of zero size Window"

windowError :: a
windowError :: a
windowError = EmptyWindowException -> a
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 :: (m () -> m ())
-> Array DW Int e
-> (Int -> e -> m a)
-> m (Int -> Int -> m (), Int, Int)
loadWithIx1 m () -> m ()
with (DWArray (DArray _ sz indexB) mWindow) Int -> e -> m a
uWrite = do
  let Window Int
it Sz Int
wk Int -> e
indexW Maybe Int
_ = Window Int e -> Maybe (Window Int e) -> Window Int e
forall a. a -> Maybe a -> a
fromMaybe Window Int e
forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window Int e)
mWindow
      wEnd :: Int
wEnd = Int
it Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
wk
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> (Int -> Int -> Bool) -> (Int -> m a) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ Int
0 Int
it Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((Int -> m a) -> m ()) -> (Int -> m a) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> Int -> e -> m a
uWrite Int
i (Int -> e
indexB Int
i)
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> (Int -> Int -> Bool) -> (Int -> m a) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ Int
wEnd (Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
sz) Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((Int -> m a) -> m ()) -> (Int -> m a) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> Int -> e -> m a
uWrite Int
i (Int -> e
indexB Int
i)
  (Int -> Int -> m (), Int, Int) -> m (Int -> Int -> m (), Int, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (\Int
from Int
to -> m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> (Int -> Int -> Bool) -> (Int -> m a) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ Int
from Int
to Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((Int -> m a) -> m ()) -> (Int -> m a) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> Int -> e -> m a
uWrite Int
i (Int -> e
indexW Int
i), Int
it, Int
wEnd)
{-# INLINE loadWithIx1 #-}

instance Index ix => Shape DW ix where
  maxLinearSize :: Array DW ix e -> Maybe (Sz Int)
maxLinearSize = Sz Int -> Maybe (Sz Int)
forall a. a -> Maybe a
Just (Sz Int -> Maybe (Sz Int))
-> (Array DW ix e -> Sz Int) -> Array DW ix e -> Maybe (Sz Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DW ix e -> Sz Int
forall r ix e. Shape r ix => Array r ix e -> Sz Int
linearSize
  {-# INLINE maxLinearSize #-}
  linearSize :: Array DW ix e -> Sz Int
linearSize = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Int -> Sz Int)
-> (Array DW ix e -> Int) -> Array DW ix e -> Sz Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Sz ix -> Int) -> (Array DW ix e -> Sz ix) -> Array DW ix e -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array D ix e -> Sz ix
forall ix e. Array D ix e -> Sz ix
dSize (Array D ix e -> Sz ix)
-> (Array DW ix e -> Array D ix e) -> Array DW ix e -> Sz ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DW ix e -> Array D ix e
forall ix e. Array DW ix e -> Array D ix e
dwArray
  {-# INLINE linearSize #-}
  outerSize :: Array DW ix e -> Sz ix
outerSize = Array D ix e -> Sz ix
forall ix e. Array D ix e -> Sz ix
dSize (Array D ix e -> Sz ix)
-> (Array DW ix e -> Array D ix e) -> Array DW ix e -> Sz ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array DW ix e -> Array D ix e
forall ix e. Array DW ix e -> Array D ix e
dwArray
  {-# INLINE outerSize #-}

instance Load DW Ix1 e where
  makeArray :: Comp -> Sz Int -> (Int -> e) -> Array DW Int e
makeArray Comp
c Sz Int
sz Int -> e
f = Array D Int e -> Maybe (Window Int e) -> Array DW Int e
forall ix e. Array D ix e -> Maybe (Window ix e) -> Array DW ix e
DWArray (Comp -> Sz Int -> (Int -> e) -> Array D Int e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
c Sz Int
sz Int -> e
f) Maybe (Window Int e)
forall a. Maybe a
Nothing
  {-# INLINE makeArray #-}
  iterArrayLinearST_ :: Scheduler s ()
-> Array DW Int e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler Array DW Int e
arr Int -> e -> ST s ()
uWrite = do
    (Int -> Int -> ST s ()
loadWindow, Int
wStart, Int
wEnd) <- (ST s () -> ST s ())
-> Array DW Int e
-> (Int -> e -> ST s ())
-> ST s (Int -> Int -> ST s (), Int, Int)
forall (m :: * -> *) e a.
Monad m =>
(m () -> m ())
-> Array DW Int e
-> (Int -> e -> m a)
-> m (Int -> Int -> m (), Int, Int)
loadWithIx1 (Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler) Array DW Int e
arr Int -> e -> ST s ()
uWrite
    let (Int
chunkWidth, Int
slackWidth) = (Int
wEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
wStart) Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler
    Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Int
wid ->
      let !it' :: Int
it' = Int
wid Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
chunkWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wStart
       in Int -> Int -> ST s ()
loadWindow Int
it' (Int
it' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkWidth)
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackWidth Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
      let !itSlack :: Int
itSlack = Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
chunkWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wStart
       in Int -> Int -> ST s ()
loadWindow Int
itSlack (Int
itSlack Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
slackWidth)
  {-# INLINE iterArrayLinearST_ #-}

instance StrideLoad DW Ix1 e where
  iterArrayLinearWithStrideST_ :: Scheduler s ()
-> Stride Int
-> Sz Int
-> Array DW Int e
-> (Int -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ Scheduler s ()
scheduler Stride Int
stride Sz Int
sz Array DW Int e
arr Int -> e -> ST s ()
uWrite = do
      ((Int, Int) -> ST s ()
loadWindow, (Int
wStart, Int
wEnd)) <- (ST s () -> ST s ())
-> Array DW Int e
-> Stride Int
-> Sz Int
-> (Int -> e -> ST s ())
-> ST s ((Int, Int) -> ST s (), (Int, Int))
forall (m :: * -> *) e a.
Monad m =>
(m () -> m ())
-> Array DW Int e
-> Stride Int
-> Sz Int
-> (Int -> e -> m a)
-> m ((Int, Int) -> m (), (Int, Int))
loadArrayWithIx1 (Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler) Array DW Int e
arr Stride Int
stride Sz Int
sz Int -> e -> ST s ()
uWrite
      let (Int
chunkWidth, Int
slackWidth) = (Int
wEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
wStart) Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler
      Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Int
wid ->
        let !it' :: Int
it' = Int
wid Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
chunkWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wStart
         in (Int, Int) -> ST s ()
loadWindow (Int
it', Int
it' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkWidth)
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackWidth Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
        let !itSlack :: Int
itSlack = Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
chunkWidth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wStart
         in (Int, Int) -> ST s ()
loadWindow (Int
itSlack, Int
itSlack Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
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 :: (m () -> m ())
-> Array DW Int e
-> Stride Int
-> Sz Int
-> (Int -> e -> m a)
-> m ((Int, Int) -> m (), (Int, Int))
loadArrayWithIx1 m () -> m ()
with (DWArray (DArray _ arrSz indexB) mWindow) Stride Int
stride Sz Int
_ Int -> e -> m a
uWrite = do
  let Window Int
it Sz Int
wk Int -> e
indexW Maybe Int
_ = Window Int e -> Maybe (Window Int e) -> Window Int e
forall a. a -> Maybe a -> a
fromMaybe Window Int e
forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window Int e)
mWindow
      wEnd :: Int
wEnd = Int
it Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
wk
      strideIx :: Int
strideIx = Stride Int -> Int
forall ix. Stride ix -> ix
unStride Stride Int
stride
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> (Int -> Int -> Bool) -> (Int -> m a) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ Int
0 Int
it Int
strideIx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((Int -> m a) -> m ()) -> (Int -> m a) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> Int -> e -> m a
uWrite (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
strideIx) (Int -> e
indexB Int
i)
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
    Int -> Int -> Int -> (Int -> Int -> Bool) -> (Int -> m a) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Stride Int -> Int -> Int
forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Int
stride Int
wEnd) (Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
arrSz) Int
strideIx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((Int -> m a) -> m ()) -> (Int -> m a) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i ->
      Int -> e -> m a
uWrite (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
strideIx) (Int -> e
indexB Int
i)
  ((Int, Int) -> m (), (Int, Int))
-> m ((Int, Int) -> m (), (Int, Int))
forall (m :: * -> *) a. Monad m => a -> m a
return
    ( \(Int
from, Int
to) ->
        m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
        Int -> Int -> Int -> (Int -> Int -> Bool) -> (Int -> m a) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Stride Int -> Int -> Int
forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Int
stride Int
from) Int
to Int
strideIx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) ((Int -> m a) -> m ()) -> (Int -> m a) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i ->
          Int -> e -> m a
uWrite (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
strideIx) (Int -> e
indexW Int
i)
    , (Int
it, Int
wEnd))
{-# INLINE loadArrayWithIx1 #-}



loadWithIx2 ::
     Monad m
  => (m () -> m ())
  -> Array DW Ix2 t1
  -> (Int -> t1 -> m ())
  -> m (Ix2 -> m (), Ix2)
loadWithIx2 :: (m () -> m ())
-> Array DW Ix2 t1 -> (Int -> t1 -> m ()) -> m (Ix2 -> m (), Ix2)
loadWithIx2 m () -> m ()
with Array DW Ix2 t1
arr Int -> t1 -> m ()
uWrite = do
  let DWArray (DArray _ (Sz (m :. n)) indexB) window = Array DW Ix2 t1
arr
  let Window (Int
it :. Int
jt) (Sz (Int
wm :. Int
wn)) Ix2 -> t1
indexW Maybe Int
mUnrollHeight = Window Ix2 t1 -> Maybe (Window Ix2 t1) -> Window Ix2 t1
forall a. a -> Maybe a -> a
fromMaybe Window Ix2 t1
forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window Ix2 t1)
window
  let Int
ib :. Int
jb = (Int
wm Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
it) Int -> Int -> Ix2
:. (Int
wn Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
jt)
      !blockHeight :: Int
blockHeight = Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
1 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
7 (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1) Maybe Int
mUnrollHeight
      stride :: Stride Ix2
stride = Stride Ix2
forall ix. Index ix => Stride ix
oneStride
      !sz :: Sz Ix2
sz = Stride Ix2 -> Sz Ix2 -> Sz Ix2
forall ix. Index ix => Stride ix -> Sz ix -> Sz ix
strideSize Stride Ix2
stride (Sz Ix2 -> Sz Ix2) -> Sz Ix2 -> Sz Ix2
forall a b. (a -> b) -> a -> b
$ Array DW Ix2 t1 -> Sz Ix2
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 = Int -> t1 -> m ()
uWrite (Sz Ix2 -> Ix2 -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz Ix2
sz Ix2
ix) (Ix2 -> t1
indexB Ix2
ix)
      {-# INLINE writeB #-}
      writeW :: Ix2 -> m ()
writeW !Ix2
ix = Int -> t1 -> m ()
uWrite (Sz Ix2 -> Ix2 -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz Ix2
sz Ix2
ix) (Ix2 -> t1
indexW Ix2
ix)
      {-# INLINE writeW #-}
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Ix2 -> Ix2 -> Ix2 -> (Int -> Int -> Bool) -> (Ix2 -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Int
0 Int -> Int -> Ix2
:. Int
0) (Int
it Int -> Int -> Ix2
:. Int
n) (Int
1 Int -> Int -> Ix2
:. Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Ix2 -> Ix2 -> Ix2 -> (Int -> Int -> Bool) -> (Ix2 -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Int
ib Int -> Int -> Ix2
:. Int
0) (Int
m Int -> Int -> Ix2
:. Int
n) (Int
1 Int -> Int -> Ix2
:. Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Ix2 -> Ix2 -> Ix2 -> (Int -> Int -> Bool) -> (Ix2 -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Int
it Int -> Int -> Ix2
:. Int
0) (Int
ib Int -> Int -> Ix2
:. Int
jt) (Int
1 Int -> Int -> Ix2
:. Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Ix2 -> Ix2 -> Ix2 -> (Int -> Int -> Bool) -> (Ix2 -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Int
it Int -> Int -> Ix2
:. Int
jb) (Int
ib Int -> Int -> Ix2
:. Int
n) (Int
1 Int -> Int -> Ix2
:. Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
  let f :: Ix2 -> m ()
f (Int
it' :. Int
ib') = m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Ix2 -> Ix2 -> Int -> (Ix2 -> m ()) -> m ()
forall (m :: * -> *).
Monad m =>
Int -> Ix2 -> Ix2 -> Int -> (Ix2 -> m ()) -> m ()
unrollAndJam Int
blockHeight (Int
it' Int -> Int -> Ix2
:. Int
jt) (Int
ib' Int -> Int -> Ix2
:. Int
jb) Int
1 Ix2 -> m ()
writeW
      {-# INLINE f #-}
  (Ix2 -> m (), Ix2) -> m (Ix2 -> m (), Ix2)
forall (m :: * -> *) a. Monad m => a -> m a
return (Ix2 -> m ()
f, Int
it Int -> Int -> Ix2
:. Int
ib)
{-# INLINE loadWithIx2 #-}

loadArrayWithIx2 ::
     Monad m
  => (m () -> m ())
  -> Array DW Ix2 e
  -> Stride Ix2
  -> Sz2
  -> (Int -> e -> m ())
  -> m (Ix2 -> m (), Ix2)
loadArrayWithIx2 :: (m () -> m ())
-> Array DW Ix2 e
-> Stride Ix2
-> Sz Ix2
-> (Int -> e -> m ())
-> m (Ix2 -> m (), Ix2)
loadArrayWithIx2 m () -> m ()
with Array DW Ix2 e
arr Stride Ix2
stride Sz Ix2
sz Int -> e -> m ()
uWrite = do
  let DWArray (DArray _ (Sz (m :. n)) indexB) window = Array DW Ix2 e
arr
  let Window (Int
it :. Int
jt) (Sz (Int
wm :. Int
wn)) Ix2 -> e
indexW Maybe Int
mUnrollHeight = Window Ix2 e -> Maybe (Window Ix2 e) -> Window Ix2 e
forall a. a -> Maybe a -> a
fromMaybe Window Ix2 e
forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window Ix2 e)
window
  let Int
ib :. Int
jb = (Int
wm Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
it) Int -> Int -> Ix2
:. (Int
wn Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
jt)
      !blockHeight :: Int
blockHeight = Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
1 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
7 (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1) Maybe Int
mUnrollHeight
      strideIx :: Ix2
strideIx@(Int
is :. Int
js) = Stride Ix2 -> Ix2
forall ix. Stride ix -> ix
unStride Stride Ix2
stride
      writeB :: Ix2 -> m ()
writeB !Ix2
ix = Int -> e -> m ()
uWrite (Stride Ix2 -> Sz Ix2 -> Ix2 -> Int
forall ix. Index ix => Stride ix -> Sz ix -> ix -> Int
toLinearIndexStride Stride Ix2
stride Sz Ix2
sz Ix2
ix) (Ix2 -> e
indexB Ix2
ix)
      {-# INLINE writeB #-}
      writeW :: Ix2 -> m ()
writeW !Ix2
ix = Int -> e -> m ()
uWrite (Stride Ix2 -> Sz Ix2 -> Ix2 -> Int
forall ix. Index ix => Stride ix -> Sz ix -> ix -> Int
toLinearIndexStride Stride Ix2
stride Sz Ix2
sz Ix2
ix) (Ix2 -> e
indexW Ix2
ix)
      {-# INLINE writeW #-}
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Ix2 -> Ix2 -> Ix2 -> (Int -> Int -> Bool) -> (Ix2 -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Int
0 Int -> Int -> Ix2
:. Int
0) (Int
it Int -> Int -> Ix2
:. Int
n) Ix2
strideIx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Ix2 -> Ix2 -> Ix2 -> (Int -> Int -> Bool) -> (Ix2 -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Stride Ix2 -> Ix2 -> Ix2
forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix2
stride (Int
ib Int -> Int -> Ix2
:. Int
0)) (Int
m Int -> Int -> Ix2
:. Int
n) Ix2
strideIx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Ix2 -> Ix2 -> Ix2 -> (Int -> Int -> Bool) -> (Ix2 -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Stride Ix2 -> Ix2 -> Ix2
forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix2
stride (Int
it Int -> Int -> Ix2
:. Int
0)) (Int
ib Int -> Int -> Ix2
:. Int
jt) Ix2
strideIx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
  m () -> m ()
with (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Ix2 -> Ix2 -> Ix2 -> (Int -> Int -> Bool) -> (Ix2 -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Stride Ix2 -> Ix2 -> Ix2
forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix2
stride (Int
it Int -> Int -> Ix2
:. Int
jb)) (Int
ib Int -> Int -> Ix2
:. Int
n) Ix2
strideIx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeB
  let f :: Ix2 -> m ()
f (Int
it' :. Int
ib')
        | Int
is Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
|| Int
blockHeight Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 -- Turn off unrolling for vertical strides
         = Ix2 -> Ix2 -> Ix2 -> (Int -> Int -> Bool) -> (Ix2 -> m ()) -> m ()
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m ()
iterM_ (Stride Ix2 -> Ix2 -> Ix2
forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix2
stride (Int
it' Int -> Int -> Ix2
:. Int
jt)) (Int
ib' Int -> Int -> Ix2
:. Int
jb) Ix2
strideIx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) Ix2 -> m ()
writeW
        | Bool
otherwise =
          Int -> Ix2 -> Ix2 -> Int -> (Ix2 -> m ()) -> m ()
forall (m :: * -> *).
Monad m =>
Int -> Ix2 -> Ix2 -> Int -> (Ix2 -> m ()) -> m ()
unrollAndJam Int
blockHeight (Stride Ix2 -> Ix2 -> Ix2
forall ix. Index ix => Stride ix -> ix -> ix
strideStart Stride Ix2
stride (Int
it' Int -> Int -> Ix2
:. Int
jt)) (Int
ib' Int -> Int -> Ix2
:. Int
jb) Int
js Ix2 -> m ()
writeW
      {-# INLINE f #-}
  (Ix2 -> m (), Ix2) -> m (Ix2 -> m (), Ix2)
forall (m :: * -> *) a. Monad m => a -> m a
return (m () -> m ()
with (m () -> m ()) -> (Ix2 -> m ()) -> Ix2 -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix2 -> m ()
f, Int
it Int -> Int -> Ix2
:. Int
ib)
{-# INLINE loadArrayWithIx2 #-}


loadWindowIx2 :: Monad m => Int -> (Ix2 -> m ()) -> Ix2 -> m ()
loadWindowIx2 :: Int -> (Ix2 -> m ()) -> Ix2 -> m ()
loadWindowIx2 Int
nWorkers Ix2 -> m ()
loadWindow (Int
it :. Int
ib) = do
  let !(Int
chunkHeight, Int
slackHeight) = (Int
ib Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
it) Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
nWorkers
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
nWorkers) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
wid ->
    let !it' :: Int
it' = Int
wid Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
chunkHeight Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
it
     in Ix2 -> m ()
loadWindow (Int
it' Int -> Int -> Ix2
:. (Int
it' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkHeight))
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackHeight Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
    let !itSlack :: Int
itSlack = Int
nWorkers Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
chunkHeight Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
it
     in Ix2 -> m ()
loadWindow (Int
itSlack Int -> Int -> Ix2
:. (Int
itSlack Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
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 = Array D Ix2 e -> Maybe (Window Ix2 e) -> Array DW Ix2 e
forall ix e. Array D ix e -> Maybe (Window ix e) -> Array DW ix e
DWArray (Comp -> Sz Ix2 -> (Ix2 -> e) -> Array D Ix2 e
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) Maybe (Window Ix2 e)
forall a. Maybe a
Nothing
  {-# INLINE makeArray #-}
  iterArrayLinearST_ :: Scheduler s ()
-> Array DW Ix2 e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler Array DW Ix2 e
arr Int -> e -> ST s ()
uWrite =
    (ST s () -> ST s ())
-> Array DW Ix2 e
-> (Int -> e -> ST s ())
-> ST s (Ix2 -> ST s (), Ix2)
forall (m :: * -> *) t1.
Monad m =>
(m () -> m ())
-> Array DW Ix2 t1 -> (Int -> t1 -> m ()) -> m (Ix2 -> m (), Ix2)
loadWithIx2 (Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler s ()
scheduler) Array DW Ix2 e
arr Int -> e -> ST s ()
uWrite ST s (Ix2 -> ST s (), Ix2)
-> ((Ix2 -> ST s (), Ix2) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
    ((Ix2 -> ST s ()) -> Ix2 -> ST s ())
-> (Ix2 -> ST s (), Ix2) -> ST s ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Int -> (Ix2 -> ST s ()) -> Ix2 -> ST s ()
forall (m :: * -> *).
Monad m =>
Int -> (Ix2 -> m ()) -> Ix2 -> m ()
loadWindowIx2 (Scheduler s () -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler))
  {-# INLINE iterArrayLinearST_ #-}

instance StrideLoad DW Ix2 e where
  iterArrayLinearWithStrideST_ :: Scheduler s ()
-> Stride Ix2
-> Sz Ix2
-> Array DW Ix2 e
-> (Int -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ Scheduler s ()
scheduler Stride Ix2
stride Sz Ix2
sz Array DW Ix2 e
arr Int -> e -> ST s ()
uWrite =
    (ST s () -> ST s ())
-> Array DW Ix2 e
-> Stride Ix2
-> Sz Ix2
-> (Int -> e -> ST s ())
-> ST s (Ix2 -> ST s (), Ix2)
forall (m :: * -> *) e.
Monad m =>
(m () -> m ())
-> Array DW Ix2 e
-> Stride Ix2
-> Sz Ix2
-> (Int -> e -> m ())
-> m (Ix2 -> m (), Ix2)
loadArrayWithIx2 (Scheduler s () -> ST s () -> ST s ()
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 Int -> e -> ST s ()
uWrite ST s (Ix2 -> ST s (), Ix2)
-> ((Ix2 -> ST s (), Ix2) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
    ((Ix2 -> ST s ()) -> Ix2 -> ST s ())
-> (Ix2 -> ST s (), Ix2) -> ST s ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Int -> (Ix2 -> ST s ()) -> Ix2 -> ST s ()
forall (m :: * -> *).
Monad m =>
Int -> (Ix2 -> m ()) -> Ix2 -> m ()
loadWindowIx2 (Scheduler s () -> Int
forall s a. Scheduler s a -> Int
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 = Array D (IxN n) e -> Maybe (Window (IxN n) e) -> Array DW (IxN n) e
forall ix e. Array D ix e -> Maybe (Window ix e) -> Array DW ix e
DWArray (Comp -> Sz (IxN n) -> (IxN n -> e) -> Array D (IxN n) e
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) Maybe (Window (IxN n) e)
forall a. Maybe a
Nothing
  {-# INLINE makeArray #-}
  iterArrayLinearST_ :: Scheduler s ()
-> Array DW (IxN n) e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ = Scheduler s ()
-> Array DW (IxN n) e -> (Int -> e -> ST s ()) -> ST s ()
forall ix e s.
(Index ix, Load DW (Lower ix) e) =>
Scheduler s () -> Array DW ix e -> (Int -> 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_ :: Scheduler s ()
-> Stride (IxN n)
-> Sz (IxN n)
-> Array DW (IxN n) e
-> (Int -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ = Scheduler s ()
-> Stride (IxN n)
-> Sz (IxN n)
-> Array DW (IxN n) e
-> (Int -> e -> ST s ())
-> ST s ()
forall ix e s.
(Index ix, StrideLoad DW (Lower ix) e) =>
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array DW ix e
-> (Int -> 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 :: Scheduler s ()
-> Stride ix
-> Sz ix
-> Array DW ix e
-> (Int -> e -> ST s ())
-> ST s ()
loadArrayWithIxN Scheduler s ()
scheduler Stride ix
stride Sz ix
szResult Array DW ix e
arr Int -> e -> ST s ()
uWrite = do
  let DWArray darr window = Array DW ix e
arr
      DArray {dSize = szSource, dIndex = indexBorder} = Array D ix e
darr
      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 Int
windowUnrollIx2 :: Maybe Int
windowUnrollIx2 :: forall ix e. Window ix e -> Maybe Int
windowUnrollIx2} = Window ix e -> Maybe (Window ix e) -> Window ix e
forall a. a -> Maybe a -> a
fromMaybe Window ix e
forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window ix e)
window
      !(Sz Int
headSourceSize, Sz (Lower ix)
lowerSourceSize) = Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
szSource
      !lowerSize :: Sz (Lower ix)
lowerSize = (Sz Int, Sz (Lower ix)) -> Sz (Lower ix)
forall a b. (a, b) -> b
snd ((Sz Int, Sz (Lower ix)) -> Sz (Lower ix))
-> (Sz Int, Sz (Lower ix)) -> Sz (Lower ix)
forall a b. (a -> b) -> a -> b
$ Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
szResult
      !(Int
s, Lower ix
lowerStrideIx) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim (ix -> (Int, Lower ix)) -> ix -> (Int, Lower ix)
forall a b. (a -> b) -> a -> b
$ Stride ix -> ix
forall ix. Stride ix -> ix
unStride Stride ix
stride
      !(Int
curWindowStart, Lower ix
lowerWindowStart) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
windowStart
      !(Sz Int
headWindowSz, Sz (Lower ix)
tailWindowSz) = Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
windowSize
      !curWindowEnd :: Int
curWindowEnd = Int
curWindowStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
headWindowSz
      !pageElements :: Int
pageElements = Sz (Lower ix) -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz (Lower ix)
lowerSize
      mkLowerWindow :: Int -> Window (Lower ix) e
mkLowerWindow Int
i =
        Window :: forall ix e. ix -> Sz ix -> (ix -> e) -> Maybe Int -> Window ix e
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 (ix -> e) -> (Lower ix -> ix) -> Lower ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i
          , windowUnrollIx2 :: Maybe Int
windowUnrollIx2 = Maybe Int
windowUnrollIx2
          }
      mkLowerArray :: Maybe (Int -> Window (Lower ix) e) -> Int -> Array DW (Lower ix) e
mkLowerArray Maybe (Int -> Window (Lower ix) e)
mw Int
i =
        DWArray :: forall ix e. Array D ix e -> Maybe (Window ix e) -> Array DW ix e
DWArray
          {dwArray :: Array D (Lower ix) e
dwArray = Comp -> Sz (Lower ix) -> (Lower ix -> e) -> Array D (Lower ix) e
forall ix e. Comp -> Sz ix -> (ix -> e) -> Array D ix e
DArray Comp
Seq Sz (Lower ix)
lowerSourceSize (ix -> e
indexBorder (ix -> e) -> (Lower ix -> ix) -> Lower ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i), dwWindow :: Maybe (Window (Lower ix) e)
dwWindow = ((Int -> Window (Lower ix) e) -> Int -> Window (Lower ix) e
forall a b. (a -> b) -> a -> b
$ Int
i) ((Int -> Window (Lower ix) e) -> Window (Lower ix) e)
-> Maybe (Int -> Window (Lower ix) e)
-> Maybe (Window (Lower ix) e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Int -> Window (Lower ix) e)
mw}
      loadLower :: Maybe (Int -> Window (Lower ix) e) -> Int -> ST s ()
loadLower Maybe (Int -> Window (Lower ix) e)
mw !Int
i =
        Scheduler s ()
-> Stride (Lower ix)
-> Sz (Lower ix)
-> Array DW (Lower ix) e
-> (Int -> e -> ST s ())
-> ST s ()
forall r ix e s.
StrideLoad r ix e =>
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array r ix e
-> (Int -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_
          Scheduler s ()
scheduler
          (Lower ix -> Stride (Lower ix)
forall ix. Index ix => ix -> Stride ix
Stride Lower ix
lowerStrideIx)
          Sz (Lower ix)
lowerSize
          (Maybe (Int -> Window (Lower ix) e) -> Int -> Array DW (Lower ix) e
mkLowerArray Maybe (Int -> Window (Lower ix) e)
mw Int
i)
          (\Int
k -> Int -> e -> ST s ()
uWrite (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
pageElements Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
s)))
      {-# NOINLINE loadLower #-}
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< ix -> Int
forall ix. Index ix => ix -> Int
headDim ix
windowStart) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s) (Maybe (Int -> Window (Lower ix) e) -> Int -> ST s ()
loadLower Maybe (Int -> Window (Lower ix) e)
forall a. Maybe a
Nothing)
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_
    (Stride Int -> Int -> Int
forall ix. Index ix => Stride ix -> ix -> ix
strideStart (Int -> Stride Int
forall ix. Index ix => ix -> Stride ix
Stride Int
s) Int
curWindowStart)
    (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
curWindowEnd)
    (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s)
    (Maybe (Int -> Window (Lower ix) e) -> Int -> ST s ()
loadLower ((Int -> Window (Lower ix) e) -> Maybe (Int -> Window (Lower ix) e)
forall a. a -> Maybe a
Just Int -> Window (Lower ix) e
mkLowerWindow))
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ (Stride Int -> Int -> Int
forall ix. Index ix => Stride ix -> ix -> ix
strideStart (Int -> Stride Int
forall ix. Index ix => ix -> Stride ix
Stride Int
s) Int
curWindowEnd) (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
headSourceSize) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s) (Maybe (Int -> Window (Lower ix) e) -> Int -> ST s ()
loadLower Maybe (Int -> Window (Lower ix) e)
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 :: Scheduler s () -> Array DW ix e -> (Int -> e -> ST s ()) -> ST s ()
loadWithIxN Scheduler s ()
scheduler Array DW ix e
arr Int -> e -> ST s ()
uWrite = do
  let DWArray darr window = Array DW ix e
arr
      DArray {dSize = sz, dIndex = indexBorder} = Array D ix e
darr
      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 Int
windowUnrollIx2 :: Maybe Int
windowUnrollIx2 :: forall ix e. Window ix e -> Maybe Int
windowUnrollIx2} = Window ix e -> Maybe (Window ix e) -> Window ix e
forall a. a -> Maybe a -> a
fromMaybe Window ix e
forall ix e. Index ix => Window ix e
zeroWindow Maybe (Window ix e)
window
      !(Sz Int
si, Sz (Lower ix)
szL) = Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
sz
      !windowEnd :: ix
windowEnd = (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ix
windowStart (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
windowSize)
      !(Int
t, Lower ix
windowStartL) = ix -> (Int, Lower ix)
forall ix. Index ix => ix -> (Int, Lower ix)
unconsDim ix
windowStart
      !pageElements :: Int
pageElements = Sz (Lower ix) -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz (Lower ix)
szL
      mkLowerWindow :: Int -> Window (Lower ix) e
mkLowerWindow Int
i =
        Window :: forall ix e. ix -> Sz ix -> (ix -> e) -> Maybe Int -> Window ix e
Window
          { windowStart :: Lower ix
windowStart = Lower ix
windowStartL
          , windowSize :: Sz (Lower ix)
windowSize = (Sz Int, Sz (Lower ix)) -> Sz (Lower ix)
forall a b. (a, b) -> b
snd ((Sz Int, Sz (Lower ix)) -> Sz (Lower ix))
-> (Sz Int, Sz (Lower ix)) -> Sz (Lower ix)
forall a b. (a -> b) -> a -> b
$ Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz Sz ix
windowSize
          , windowIndex :: Lower ix -> e
windowIndex = ix -> e
windowIndex (ix -> e) -> (Lower ix -> ix) -> Lower ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i
          , windowUnrollIx2 :: Maybe Int
windowUnrollIx2 = Maybe Int
windowUnrollIx2
          }
      mkLowerArray :: Maybe (Int -> Window (Lower ix) e) -> Int -> Array DW (Lower ix) e
mkLowerArray Maybe (Int -> Window (Lower ix) e)
mw Int
i =
        DWArray :: forall ix e. Array D ix e -> Maybe (Window ix e) -> Array DW ix e
DWArray {dwArray :: Array D (Lower ix) e
dwArray = Comp -> Sz (Lower ix) -> (Lower ix -> e) -> Array D (Lower ix) e
forall ix e. Comp -> Sz ix -> (ix -> e) -> Array D ix e
DArray Comp
Seq Sz (Lower ix)
szL (ix -> e
indexBorder (ix -> e) -> (Lower ix -> ix) -> Lower ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i), dwWindow :: Maybe (Window (Lower ix) e)
dwWindow = ((Int -> Window (Lower ix) e) -> Int -> Window (Lower ix) e
forall a b. (a -> b) -> a -> b
$ Int
i) ((Int -> Window (Lower ix) e) -> Window (Lower ix) e)
-> Maybe (Int -> Window (Lower ix) e)
-> Maybe (Window (Lower ix) e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Int -> Window (Lower ix) e)
mw}
      loadLower :: Maybe (Int -> Window (Lower ix) e) -> Int -> ST s ()
loadLower Maybe (Int -> Window (Lower ix) e)
mw !Int
i =
        Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
        Scheduler s ()
-> Array DW (Lower ix) e -> (Int -> e -> ST s ()) -> ST s ()
forall r ix e s.
Load r ix e =>
Scheduler s () -> Array r ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler (Maybe (Int -> Window (Lower ix) e) -> Int -> Array DW (Lower ix) e
mkLowerArray Maybe (Int -> Window (Lower ix) e)
mw Int
i) (\Int
k -> Int -> e -> ST s ()
uWrite (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
pageElements Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i))
      {-# NOINLINE loadLower #-}
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< ix -> Int
forall ix. Index ix => ix -> Int
headDim ix
windowStart) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Maybe (Int -> Window (Lower ix) e) -> Int -> ST s ()
loadLower Maybe (Int -> Window (Lower ix) e)
forall a. Maybe a
Nothing)
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
t (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< ix -> Int
forall ix. Index ix => ix -> Int
headDim ix
windowEnd) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Maybe (Int -> Window (Lower ix) e) -> Int -> ST s ()
loadLower ((Int -> Window (Lower ix) e) -> Maybe (Int -> Window (Lower ix) e)
forall a. a -> Maybe a
Just Int -> Window (Lower ix) e
mkLowerWindow))
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ (ix -> Int
forall ix. Index ix => ix -> Int
headDim ix
windowEnd) (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
si) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Maybe (Int -> Window (Lower ix) e) -> Int -> ST s ()
loadLower Maybe (Int -> Window (Lower ix) e)
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 :: Int -> Ix2 -> Ix2 -> Int -> (Ix2 -> m ()) -> m ()
unrollAndJam !Int
bH (Int
it :. Int
jt) (Int
ib :. Int
jb) Int
js Ix2 -> m ()
f = do
  let f2 :: Ix2 -> m ()
f2 (Int
i :. Int
j) = Ix2 -> m ()
f (Int
i Int -> Int -> Ix2
:. Int
j) m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f  ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Ix2
:. Int
j)
  let f3 :: Ix2 -> m ()
f3 (Int
i :. Int
j) = Ix2 -> m ()
f (Int
i Int -> Int -> Ix2
:. Int
j) m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f2 ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Ix2
:. Int
j)
  let f4 :: Ix2 -> m ()
f4 (Int
i :. Int
j) = Ix2 -> m ()
f (Int
i Int -> Int -> Ix2
:. Int
j) m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f3 ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Ix2
:. Int
j)
  let f5 :: Ix2 -> m ()
f5 (Int
i :. Int
j) = Ix2 -> m ()
f (Int
i Int -> Int -> Ix2
:. Int
j) m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f4 ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Ix2
:. Int
j)
  let f6 :: Ix2 -> m ()
f6 (Int
i :. Int
j) = Ix2 -> m ()
f (Int
i Int -> Int -> Ix2
:. Int
j) m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f5 ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Ix2
:. Int
j)
  let f7 :: Ix2 -> m ()
f7 (Int
i :. Int
j) = Ix2 -> m ()
f (Int
i Int -> Int -> Ix2
:. Int
j) m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ix2 -> m ()
f6 ((Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Ix2
:. Int
j)
  let f' :: Ix2 -> m ()
f' = case Int
bH of
             Int
1 -> Ix2 -> m ()
f
             Int
2 -> Ix2 -> m ()
f2
             Int
3 -> Ix2 -> m ()
f3
             Int
4 -> Ix2 -> m ()
f4
             Int
5 -> Ix2 -> m ()
f5
             Int
6 -> Ix2 -> m ()
f6
             Int
_ -> Ix2 -> m ()
f7
  let !ibS :: Int
ibS = Int
ib Int -> Int -> Int
forall a. Num a => a -> a -> a
- ((Int
ib Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
it) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
bH)
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
it (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ibS) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bH) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i ->
    Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
jt (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
jb) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
js) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
j ->
      Ix2 -> m ()
f' (Int
i Int -> Int -> Ix2
:. Int
j)
  Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
ibS (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ib) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i ->
    Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
jt (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
jb) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
js) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
j ->
      Ix2 -> m ()
f (Int
i Int -> Int -> Ix2
:. Int
j)
{-# INLINE unrollAndJam #-}



-- TODO: Implement Hilbert curve