#include "fusion-phases.h"
module Data.Array.Parallel.Unlifted.Stream.Elems
( streamElemsFromVector
, streamElemsFromVectors
, streamElemsFromVectorsVSegd)
where
import Data.Array.Parallel.Unlifted.Stream.Ixs
import Data.Vector.Fusion.Stream.Monadic
import Data.Array.Parallel.Unlifted.Sequential.Vector (Unbox, Vector)
import Data.Array.Parallel.Unlifted.Vectors (Unboxes, Vectors)
import Data.Array.Parallel.Unlifted.Sequential.UVSegd (UVSegd(..))
import qualified Data.Array.Parallel.Unlifted.Sequential.Vector as U
import qualified Data.Array.Parallel.Unlifted.Vectors as US
import qualified Data.Array.Parallel.Unlifted.Sequential.UVSegd as UVSegd
streamElemsFromVector
:: (Monad m, Unbox a)
=> Vector a -> Stream m Int -> Stream m a
streamElemsFromVector vector (Stream mkStep s0 size0)
= vector `seq` Stream mkStep' s0 size0
where
mkStep' s
= do step <- mkStep s
case step of
Yield ix s'
-> let !result = U.index "streamElemsFromVector" vector ix
in return $ Yield result s'
Skip s' -> return $ Skip s'
Done -> return Done
streamElemsFromVectors
:: (Monad m, Unboxes a)
=> Vectors a -> Stream m (Int, Int) -> Stream m a
streamElemsFromVectors vectors (Stream mkStep s0 size0)
= vectors `seq` Stream mkStep' s0 size0
where
mkStep' s
= do step <- mkStep s
case step of
Yield (ix1, ix2) s'
-> let !result = US.unsafeIndex2 vectors ix1 ix2
in return $ Yield result s'
Skip s' -> return $ Skip s'
Done -> return Done
streamElemsFromVectorsVSegd
:: (Monad m, Unboxes a)
=> Vectors a -> UVSegd -> Stream m (Int, Int) -> Stream m a
streamElemsFromVectorsVSegd vectors uvsegd vsrcixs
= let
vsegids = UVSegd.takeVSegidsRedundant uvsegd
ussegd = UVSegd.takeUSSegdRedundant uvsegd
in streamElemsFromVectors vectors
$ streamSrcIxsThroughUSSegd ussegd
$ streamSrcIxsThroughVSegids vsegids
$ vsrcixs