module Data.Yarr.Repr.Separate (
SE,
UArray(..),
fromSlices, unsafeMapSlices,
Data.Yarr.Repr.Separate.convert,
dmapElems, dmapElemsM,
dzipElems2, dzipElems2M, dzipElems3, dzipElems3M,
dzipElems, dzipElemsM,
fmapElems, fmapElemsM,
fzipElems2, fzipElems2M, fzipElems3, fzipElems3M,
fzipElems, fzipElemsM
) where
import Prelude as P
import Data.Functor ((<$>))
import Data.Yarr.Base as B
import Data.Yarr.Shape
import Data.Yarr.Repr.Delayed
import Data.Yarr.Utils.FixedVector as V
data SE r
instance (Regular r l sh e, Vector v e) => Regular (SE r) l sh (v e) where
data UArray (SE r) l sh (v e) =
Separate !sh (VecList (Dim v) (UArray r l sh e))
extent (Separate sh _) = sh
touchArray (Separate _ slices) = V.mapM_ touchArray slices
force (Separate sh slices) = do
sh `deepseq` return ()
V.mapM_ force slices
instance (NFData (UArray r l sh e), Shape sh, Vector v e) =>
NFData (UArray (SE r) l sh (v e)) where
rnf (Separate sh slices) = sh `deepseq` slices `deepseq` ()
instance (Regular r l sh e, Shape sh, Vector v e) =>
VecRegular (SE r) r l sh v e where
slices (Separate _ slices) = slices
instance (USource r l sh e, Vector v e) => USource (SE r) l sh (v e) where
index (Separate _ slices) sh =
V.convert <$> V.mapM (\el -> index el sh) slices
linearIndex (Separate _ slices) i =
V.convert <$> V.mapM (\el -> linearIndex el i) slices
instance (USource r l sh e, Vector v e) => UVecSource (SE r) r l sh v e
instance (DefaultFusion r D l, Fusion (SE r) D l) => DefaultFusion (SE r) D l
fmapElems
:: (VecRegular r slr l sh v a,
USource slr l sh a, USource fslr l sh b, Fusion slr fslr l,
Vector v2 b, Dim v ~ Dim v2)
=> VecList (Dim v) (a -> b)
-> UArray r l sh (v a)
-> UArray (SE fslr) l sh (v2 b)
fmapElems fs = fmapElemsM $ V.map (return .) fs
fmapElemsM
:: (VecRegular r slr l sh v a,
USource slr l sh a, USource fslr l sh b, Fusion slr fslr l,
Vector v2 b, Dim v ~ Dim v2)
=> VecList (Dim v) (a -> IO b)
-> UArray r l sh (v a)
-> UArray (SE fslr) l sh (v2 b)
fmapElemsM fs arr = Separate (extent arr) $ V.zipWith fmapM fs (slices arr)
fzipElems2
:: (VecRegular r slr l sh v a, USource slr l sh a,
VecRegular r slr l sh v b, USource slr l sh b,
USource fslr l sh c, Fusion slr fslr l, Vector v c)
=> VecList (Dim v) (a -> b -> c)
-> UArray r l sh (v a)
-> UArray r l sh (v b)
-> UArray (SE fslr) l sh (v c)
fzipElems2 fs arr1 arr2 =
let fMs = V.map (\f -> \x y -> return (f x y)) fs
in fzipElems2M fMs arr1 arr2
fzipElems2M
:: (VecRegular r slr l sh v a, USource slr l sh a,
VecRegular r slr l sh v b, USource slr l sh b,
USource fslr l sh c, Fusion slr fslr l, Vector v c)
=> VecList (Dim v) (a -> b -> IO c)
-> UArray r l sh (v a)
-> UArray r l sh (v b)
-> UArray (SE fslr) l sh (v c)
fzipElems2M fs arr1 arr2 =
let sh = intersect (vl_2 (extent arr1) (extent arr2))
slices1 = slices arr1
slices2 = slices arr2
makeSlice i f =
let sl1 = slices1 V.! i
sl2 = slices2 V.! i
in fzip2M f sl1 sl2
in Separate sh $ V.imap makeSlice fs
fzipElems3
:: (VecRegular r slr l sh v a, USource slr l sh a,
VecRegular r slr l sh v b, USource slr l sh b,
VecRegular r slr l sh v c, USource slr l sh c,
USource fslr l sh d, Fusion slr fslr l, Vector v d)
=> VecList (Dim v) (a -> b -> c -> d)
-> UArray r l sh (v a)
-> UArray r l sh (v b)
-> UArray r l sh (v c)
-> UArray (SE fslr) l sh (v d)
fzipElems3 fs arr1 arr2 arr3 =
let fMs = V.map (\f -> \x y z -> return (f x y z)) fs
in fzipElems3M fMs arr1 arr2 arr3
fzipElems3M
:: (VecRegular r slr l sh v a, USource slr l sh a,
VecRegular r slr l sh v b, USource slr l sh b,
VecRegular r slr l sh v c, USource slr l sh c,
USource fslr l sh d, Fusion slr fslr l, Vector v d)
=> VecList (Dim v) (a -> b -> c -> IO d)
-> UArray r l sh (v a)
-> UArray r l sh (v b)
-> UArray r l sh (v c)
-> UArray (SE fslr) l sh (v d)
fzipElems3M fs arr1 arr2 arr3 =
let sh = intersect (vl_3 (extent arr1) (extent arr2) (extent arr3))
slices1 = slices arr1
slices2 = slices arr2
slices3 = slices arr3
makeSlice i f =
let sl1 = slices1 V.! i
sl2 = slices2 V.! i
sl3 = slices3 V.! i
in fzip3M f sl1 sl2 sl3
in Separate sh $ V.imap makeSlice fs
fzipElems
:: (Vector v2 b, Arity m, m ~ S m0,
VecRegular r slr l sh v a,
USource slr l sh a, USource fslr l sh b, Fusion slr fslr l)
=> VecList (Dim v2) (Fun m a b)
-> VecList m (UArray r l sh (v a))
-> UArray (SE fslr) l sh (v2 b)
fzipElems funs arrs =
let funMs = V.map (P.fmap return) funs
in fzipElemsM funMs arrs
fzipElemsM
:: (Vector v2 b, Arity m, m ~ S m0,
VecRegular r slr l sh v a,
USource slr l sh a, USource fslr l sh b, Fusion slr fslr l)
=> VecList (Dim v2) (Fun m a (IO b))
-> VecList m (UArray r l sh (v a))
-> UArray (SE fslr) l sh (v2 b)
fzipElemsM funs arrs =
let sh = intersect $ V.map extent arrs
!allSlices = V.map slices arrs
makeSlice i fun =
let slices = V.map (V.! i) allSlices
in fzipM fun slices
in Separate sh $ V.imap makeSlice funs
dmapElems
:: (VecRegular r slr l sh v a,
USource slr l sh a, USource fslr l sh b, DefaultFusion slr fslr l,
Vector v2 b, Dim v ~ Dim v2)
=> VecList (Dim v) (a -> b)
-> UArray r l sh (v a)
-> UArray (SE fslr) l sh (v2 b)
dmapElems = fmapElems
dmapElemsM
:: (VecRegular r slr l sh v a,
USource slr l sh a, USource fslr l sh b, DefaultFusion slr fslr l,
Vector v2 b, Dim v ~ Dim v2)
=> VecList (Dim v) (a -> IO b)
-> UArray r l sh (v a)
-> UArray (SE fslr) l sh (v2 b)
dmapElemsM = fmapElemsM
dzipElems2
:: (VecRegular r slr l sh v a, USource slr l sh a,
VecRegular r slr l sh v b, USource slr l sh b,
USource fslr l sh c, DefaultFusion slr fslr l, Vector v c)
=> VecList (Dim v) (a -> b -> c)
-> UArray r l sh (v a)
-> UArray r l sh (v b)
-> UArray (SE fslr) l sh (v c)
dzipElems2 = fzipElems2
dzipElems2M
:: (VecRegular r slr l sh v a, USource slr l sh a,
VecRegular r slr l sh v b, USource slr l sh b,
USource fslr l sh c, DefaultFusion slr fslr l, Vector v c)
=> VecList (Dim v) (a -> b -> IO c)
-> UArray r l sh (v a)
-> UArray r l sh (v b)
-> UArray (SE fslr) l sh (v c)
dzipElems2M = fzipElems2M
dzipElems3
:: (VecRegular r slr l sh v a, USource slr l sh a,
VecRegular r slr l sh v b, USource slr l sh b,
VecRegular r slr l sh v c, USource slr l sh c,
USource fslr l sh d, DefaultFusion slr fslr l, Vector v d)
=> VecList (Dim v) (a -> b -> c -> d)
-> UArray r l sh (v a)
-> UArray r l sh (v b)
-> UArray r l sh (v c)
-> UArray (SE fslr) l sh (v d)
dzipElems3 = fzipElems3
dzipElems3M
:: (VecRegular r slr l sh v a, USource slr l sh a,
VecRegular r slr l sh v b, USource slr l sh b,
VecRegular r slr l sh v c, USource slr l sh c,
USource fslr l sh d, DefaultFusion slr fslr l, Vector v d)
=> VecList (Dim v) (a -> b -> c -> IO d)
-> UArray r l sh (v a)
-> UArray r l sh (v b)
-> UArray r l sh (v c)
-> UArray (SE fslr) l sh (v d)
dzipElems3M = fzipElems3M
dzipElems
:: (Vector v2 b, Arity m, m ~ S m0,
VecRegular r slr l sh v a,
USource slr l sh a, USource fslr l sh b, DefaultFusion slr fslr l)
=> VecList (Dim v2) (Fun m a b)
-> VecList m (UArray r l sh (v a))
-> UArray (SE fslr) l sh (v2 b)
dzipElems = fzipElems
dzipElemsM
:: (Vector v2 b, Arity m, m ~ S m0,
VecRegular r slr l sh v a,
USource slr l sh a, USource fslr l sh b, DefaultFusion slr fslr l)
=> VecList (Dim v2) (Fun m a (IO b))
-> VecList m (UArray r l sh (v a))
-> UArray (SE fslr) l sh (v2 b)
dzipElemsM = fzipElemsM
instance (UTarget tr tl sh e, Vector v e) => UTarget (SE tr) tl sh (v e) where
write (Separate _ slices) sh v =
V.zipWithM_ (\el x -> write el sh x) slices (V.convert v)
linearWrite (Separate _ slices) i v =
V.zipWithM_ (\el x -> linearWrite el i x) slices (V.convert v)
instance (Manifest r mr l sh e, Vector v e) =>
Manifest (SE r) (SE mr) l sh (v e) where
new sh = P.fmap (Separate sh) (V.replicateM (B.new sh))
freeze (Separate sh mslices) = P.fmap (Separate sh) (V.mapM freeze mslices)
thaw (Separate sh slices) = P.fmap (Separate sh) (V.mapM thaw slices)
instance (UTarget tr tl sh e, Vector v e) => UVecTarget (SE tr) tr tl sh v e
fromSlices
:: (Regular r l sh e, Vector v e, Dim v ~ S n0)
=> VecList (Dim v) (UArray r l sh e)
-> UArray (SE r) l sh (v e)
fromSlices slices =
let shapes = V.map extent slices
sh0 = V.head shapes
in if V.any (/= sh0) shapes
then error "Separate Repr: all slices must be of the same extent"
else Separate sh0 slices
unsafeMapSlices
:: (USource r l sh a, Vector v a,
USource r2 l2 sh2 b, Vector v b, Dim v ~ S n0)
=> (UArray r l sh a -> UArray r2 l2 sh2 b)
-> UArray (SE r) l sh (v a)
-> UArray (SE r2) l2 sh2 (v b)
unsafeMapSlices f (Separate sh slices) =
let slices' = V.map f slices
in Separate (extent (V.head slices')) slices'
convert
:: (Regular r l sh e, Vector v e, Vector v2 e, Dim v ~ Dim v2)
=> UArray (SE r) l sh (v e) -> UArray (SE r) l sh (v2 e)
convert (Separate sh slices) = Separate sh slices