{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : Data.Massiv.Array.Delayed.Interleaved
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Array.Delayed.Interleaved (
  DI (..),
  Array (..),
  toInterleaved,
  fromInterleaved,
) where

import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Core.Common
import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec)

-- | Delayed array that will be loaded in an interleaved fashion during parallel
-- computation.
--
-- /Warning/ - Will be deprecated in the next major version update.
data DI = DI

newtype instance Array DI ix e = DIArray
  { forall ix e. Array DI ix e -> Array D ix e
diArray :: Array D ix e
  }
  deriving (Array DI ix e -> Array DI ix e -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall ix e.
(Index ix, Eq e) =>
Array DI ix e -> Array DI ix e -> Bool
/= :: Array DI ix e -> Array DI ix e -> Bool
$c/= :: forall ix e.
(Index ix, Eq e) =>
Array DI ix e -> Array DI ix e -> Bool
== :: Array DI ix e -> Array DI ix e -> Bool
$c== :: forall ix e.
(Index ix, Eq e) =>
Array DI ix e -> Array DI ix e -> Bool
Eq, Array DI ix e -> Array DI ix e -> Bool
Array DI ix e -> Array DI ix e -> Ordering
Array DI ix e -> Array DI ix e -> Array DI ix e
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {ix} {e}. (Index ix, Ord e) => Eq (Array DI ix e)
forall ix e.
(Index ix, Ord e) =>
Array DI ix e -> Array DI ix e -> Bool
forall ix e.
(Index ix, Ord e) =>
Array DI ix e -> Array DI ix e -> Ordering
forall ix e.
(Index ix, Ord e) =>
Array DI ix e -> Array DI ix e -> Array DI ix e
min :: Array DI ix e -> Array DI ix e -> Array DI ix e
$cmin :: forall ix e.
(Index ix, Ord e) =>
Array DI ix e -> Array DI ix e -> Array DI ix e
max :: Array DI ix e -> Array DI ix e -> Array DI ix e
$cmax :: forall ix e.
(Index ix, Ord e) =>
Array DI ix e -> Array DI ix e -> Array DI ix e
>= :: Array DI ix e -> Array DI ix e -> Bool
$c>= :: forall ix e.
(Index ix, Ord e) =>
Array DI ix e -> Array DI ix e -> Bool
> :: Array DI ix e -> Array DI ix e -> Bool
$c> :: forall ix e.
(Index ix, Ord e) =>
Array DI ix e -> Array DI ix e -> Bool
<= :: Array DI ix e -> Array DI ix e -> Bool
$c<= :: forall ix e.
(Index ix, Ord e) =>
Array DI ix e -> Array DI ix e -> Bool
< :: Array DI ix e -> Array DI ix e -> Bool
$c< :: forall ix e.
(Index ix, Ord e) =>
Array DI ix e -> Array DI ix e -> Bool
compare :: Array DI ix e -> Array DI ix e -> Ordering
$ccompare :: forall ix e.
(Index ix, Ord e) =>
Array DI ix e -> Array DI ix e -> Ordering
Ord, forall a b. a -> Array DI ix b -> Array DI ix a
forall a b. (a -> b) -> Array DI ix a -> Array DI ix b
forall ix a b. a -> Array DI ix b -> Array DI ix a
forall ix a b. (a -> b) -> Array DI ix a -> Array DI ix b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Array DI ix b -> Array DI ix a
$c<$ :: forall ix a b. a -> Array DI ix b -> Array DI ix a
fmap :: forall a b. (a -> b) -> Array DI ix a -> Array DI ix b
$cfmap :: forall ix a b. (a -> b) -> Array DI ix a -> Array DI ix b
Functor, forall a. a -> Array DI ix a
forall {ix}. Index ix => Functor (Array DI ix)
forall ix a. Index ix => a -> Array DI ix a
forall ix a b.
Index ix =>
Array DI ix a -> Array DI ix b -> Array DI ix a
forall ix a b.
Index ix =>
Array DI ix a -> Array DI ix b -> Array DI ix b
forall ix a b.
Index ix =>
Array DI ix (a -> b) -> Array DI ix a -> Array DI ix b
forall ix a b c.
Index ix =>
(a -> b -> c) -> Array DI ix a -> Array DI ix b -> Array DI ix c
forall a b. Array DI ix a -> Array DI ix b -> Array DI ix a
forall a b. Array DI ix a -> Array DI ix b -> Array DI ix b
forall a b. Array DI ix (a -> b) -> Array DI ix a -> Array DI ix b
forall a b c.
(a -> b -> c) -> Array DI ix a -> Array DI ix b -> Array DI ix c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. Array DI ix a -> Array DI ix b -> Array DI ix a
$c<* :: forall ix a b.
Index ix =>
Array DI ix a -> Array DI ix b -> Array DI ix a
*> :: forall a b. Array DI ix a -> Array DI ix b -> Array DI ix b
$c*> :: forall ix a b.
Index ix =>
Array DI ix a -> Array DI ix b -> Array DI ix b
liftA2 :: forall a b c.
(a -> b -> c) -> Array DI ix a -> Array DI ix b -> Array DI ix c
$cliftA2 :: forall ix a b c.
Index ix =>
(a -> b -> c) -> Array DI ix a -> Array DI ix b -> Array DI ix c
<*> :: forall a b. Array DI ix (a -> b) -> Array DI ix a -> Array DI ix b
$c<*> :: forall ix a b.
Index ix =>
Array DI ix (a -> b) -> Array DI ix a -> Array DI ix b
pure :: forall a. a -> Array DI ix a
$cpure :: forall ix a. Index ix => a -> Array DI ix a
Applicative, forall a. Eq a => a -> Array DI ix a -> Bool
forall a. Num a => Array DI ix a -> a
forall a. Ord a => Array DI ix a -> a
forall m. Monoid m => Array DI ix m -> m
forall ix a. (Index ix, Eq a) => a -> Array DI ix a -> Bool
forall ix a. (Index ix, Num a) => Array DI ix a -> a
forall ix a. (Index ix, Ord a) => Array DI ix a -> a
forall ix m. (Index ix, Monoid m) => Array DI ix m -> m
forall ix a. Index ix => Array DI ix a -> Bool
forall ix a. Index ix => Array DI ix a -> Int
forall ix a. Index ix => Array DI ix a -> [a]
forall ix a. Index ix => (a -> a -> a) -> Array DI ix a -> a
forall ix m a.
(Index ix, Monoid m) =>
(a -> m) -> Array DI ix a -> m
forall ix b a. Index ix => (b -> a -> b) -> b -> Array DI ix a -> b
forall ix a b. Index ix => (a -> b -> b) -> b -> Array DI ix a -> b
forall a. Array DI ix a -> Bool
forall a. Array DI ix a -> Int
forall a. Array DI ix a -> [a]
forall a. (a -> a -> a) -> Array DI ix a -> a
forall m a. Monoid m => (a -> m) -> Array DI ix a -> m
forall b a. (b -> a -> b) -> b -> Array DI ix a -> b
forall a b. (a -> b -> b) -> b -> Array DI ix a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => Array DI ix a -> a
$cproduct :: forall ix a. (Index ix, Num a) => Array DI ix a -> a
sum :: forall a. Num a => Array DI ix a -> a
$csum :: forall ix a. (Index ix, Num a) => Array DI ix a -> a
minimum :: forall a. Ord a => Array DI ix a -> a
$cminimum :: forall ix a. (Index ix, Ord a) => Array DI ix a -> a
maximum :: forall a. Ord a => Array DI ix a -> a
$cmaximum :: forall ix a. (Index ix, Ord a) => Array DI ix a -> a
elem :: forall a. Eq a => a -> Array DI ix a -> Bool
$celem :: forall ix a. (Index ix, Eq a) => a -> Array DI ix a -> Bool
length :: forall a. Array DI ix a -> Int
$clength :: forall ix a. Index ix => Array DI ix a -> Int
null :: forall a. Array DI ix a -> Bool
$cnull :: forall ix a. Index ix => Array DI ix a -> Bool
toList :: forall a. Array DI ix a -> [a]
$ctoList :: forall ix a. Index ix => Array DI ix a -> [a]
foldl1 :: forall a. (a -> a -> a) -> Array DI ix a -> a
$cfoldl1 :: forall ix a. Index ix => (a -> a -> a) -> Array DI ix a -> a
foldr1 :: forall a. (a -> a -> a) -> Array DI ix a -> a
$cfoldr1 :: forall ix a. Index ix => (a -> a -> a) -> Array DI ix a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> Array DI ix a -> b
$cfoldl' :: forall ix b a. Index ix => (b -> a -> b) -> b -> Array DI ix a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Array DI ix a -> b
$cfoldl :: forall ix b a. Index ix => (b -> a -> b) -> b -> Array DI ix a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Array DI ix a -> b
$cfoldr' :: forall ix a b. Index ix => (a -> b -> b) -> b -> Array DI ix a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Array DI ix a -> b
$cfoldr :: forall ix a b. Index ix => (a -> b -> b) -> b -> Array DI ix a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> Array DI ix a -> m
$cfoldMap' :: forall ix m a.
(Index ix, Monoid m) =>
(a -> m) -> Array DI ix a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Array DI ix a -> m
$cfoldMap :: forall ix m a.
(Index ix, Monoid m) =>
(a -> m) -> Array DI ix a -> m
fold :: forall m. Monoid m => Array DI ix m -> m
$cfold :: forall ix m. (Index ix, Monoid m) => Array DI ix m -> m
Foldable)

instance (Ragged L ix e, Show e) => Show (Array DI ix e) where
  showsPrec :: Int -> Array DI 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) -> Int -> Array r ix e -> ShowS
showsArrayPrec forall ix e. Array DI ix e -> Array D ix e
diArray
  showList :: [Array DI ix e] -> ShowS
showList = forall arr. Show arr => [arr] -> ShowS
showArrayList

instance Strategy DI where
  setComp :: forall ix e. Comp -> Array DI ix e -> Array DI ix e
setComp Comp
c Array DI ix e
arr = Array DI ix e
arr{diArray :: Array D ix e
diArray = (forall ix e. Array DI ix e -> Array D ix e
diArray Array DI ix e
arr){dComp :: Comp
dComp = Comp
c}}
  {-# INLINE setComp #-}
  getComp :: forall ix e. Array DI 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 DI ix e -> Array D ix e
diArray
  {-# INLINE getComp #-}
  repr :: DI
repr = DI
DI

instance Index ix => Shape DI ix where
  maxLinearSize :: forall e. Array DI ix e -> Maybe Sz1
maxLinearSize = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount
  {-# INLINE maxLinearSize #-}

instance Size DI where
  size :: forall ix e. Array DI ix e -> Sz ix
size (DIArray Array D ix e
arr) = forall r ix e. Size r => Array r ix e -> Sz ix
size Array D ix e
arr
  {-# INLINE size #-}
  unsafeResize :: forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array DI ix e -> Array DI ix' e
unsafeResize Sz ix'
sz = forall ix e. Array D ix e -> Array DI ix e
DIArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz ix'
sz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Array DI ix e -> Array D ix e
diArray
  {-# INLINE unsafeResize #-}

instance Index ix => Load DI ix e where
  makeArray :: Comp -> Sz ix -> (ix -> e) -> Array DI ix e
makeArray Comp
c Sz ix
sz = forall ix e. Array D ix e -> Array DI ix e
DIArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
c Sz ix
sz
  {-# INLINE makeArray #-}
  iterArrayLinearST_ :: forall s.
Scheduler s () -> Array DI ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ Scheduler s ()
scheduler (DIArray darr :: Array D ix e
darr@(DArray Comp
_ Sz ix
sz PrefIndex ix e
_)) Int -> e -> ST s ()
uWrite =
    forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (forall a. Ord a => a -> a -> Bool
< forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
      forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
        forall ix (m :: * -> *).
(Index ix, Monad m) =>
Sz ix
-> Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> (Int -> ix -> m ())
-> m ()
iterLinearM_ Sz ix
sz Int
start (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) (forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) forall a. Ord a => a -> a -> Bool
(<) forall a b. (a -> b) -> a -> b
$ \ !Int
k ->
          Int -> e -> ST s ()
uWrite Int
k forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D ix e
darr
  {-# INLINE iterArrayLinearST_ #-}

instance Index ix => StrideLoad DI ix e where
  iterArrayLinearWithStrideST_ :: forall s.
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array DI ix e
-> (Int -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ Scheduler s ()
scheduler Stride ix
stride Sz ix
resultSize (DIArray Array D ix e
arr) Int -> e -> ST s ()
uWrite =
    forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (forall a. Ord a => a -> a -> Bool
< forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
      forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$
        forall ix (m :: * -> *).
(Index ix, Monad m) =>
Sz ix
-> Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> (Int -> ix -> m ())
-> m ()
iterLinearM_ Sz ix
resultSize Int
start (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
resultSize) (forall s a. Scheduler s a -> Int
numWorkers Scheduler s ()
scheduler) forall a. Ord a => a -> a -> Bool
(<) forall a b. (a -> b) -> a -> b
$
          \ !Int
i ix
ix -> Int -> e -> ST s ()
uWrite Int
i (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D ix e
arr (forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 forall a. Num a => a -> a -> a
(*) (forall ix. Stride ix -> ix
unStride Stride ix
stride) ix
ix))
  {-# INLINE iterArrayLinearWithStrideST_ #-}

-- | Convert a source array into an array that, when computed, will have its elemets evaluated out
-- of order (interleaved amongst cores), hence making unbalanced computation better parallelizable.
toInterleaved :: (Index ix, Source r e) => Array r ix e -> Array DI ix e
toInterleaved :: forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Array DI ix e
toInterleaved = forall ix e. Array D ix e -> Array DI ix e
DIArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Array D ix e
delay
{-# INLINE toInterleaved #-}

-- | /O(1)/ - Unwrap the interleved array.
--
-- @since 0.2.1
fromInterleaved :: Array DI ix e -> Array D ix e
fromInterleaved :: forall ix e. Array DI ix e -> Array D ix e
fromInterleaved = forall ix e. Array DI ix e -> Array D ix e
diArray
{-# INLINE fromInterleaved #-}