#include "fusion-phases.h"
module Data.Array.Parallel.Unlifted.Stream.Segments
( streamSegsFromNestedUSSegd
, streamSegsFromVectorsUSSegd
, streamSegsFromVectorsUVSegd
, streamSegsFromVectorsUSSegdSegmap
, streamSegsFromVectorsUSSegd_split)
where
import Data.Vector.Fusion.Stream.Size
import Data.Vector.Fusion.Stream.Monadic
import Data.Array.Parallel.Unlifted.Sequential.Vector (Unbox, Vector, index)
import Data.Array.Parallel.Unlifted.Vectors (Unboxes, Vectors)
import Data.Array.Parallel.Unlifted.Sequential.USegd (USegd(..))
import Data.Array.Parallel.Unlifted.Sequential.USSegd (USSegd(..))
import Data.Array.Parallel.Unlifted.Sequential.UVSegd (UVSegd(..))
import qualified Data.Array.Parallel.Unlifted.Vectors as US
import qualified Data.Array.Parallel.Unlifted.Sequential.USegd as USegd
import qualified Data.Array.Parallel.Unlifted.Sequential.USSegd as USSegd
import qualified Data.Array.Parallel.Unlifted.Sequential.Vector as U
import qualified Data.Vector as V
import qualified Data.Primitive.ByteArray as P
import System.IO.Unsafe
streamSegsFromNestedUSSegd
:: (Unbox a, Monad m)
=> V.Vector (Vector a)
-> USSegd
-> Stream m a
streamSegsFromNestedUSSegd
pdatas
ussegd@(USSegd _ starts sources usegd)
= let
here = "streamSegsFromNestedUSSegd"
pseglens = USegd.takeLengths usegd
fn (pseg, ix)
| pseg >= USSegd.length ussegd
= return $ Done
| ix >= U.index here pseglens pseg
= return $ Skip (pseg + 1, 0)
| otherwise
= let !srcid = index here sources pseg
!pdata = pdatas `V.unsafeIndex` srcid
!start = index here starts pseg
!result = index here pdata (start + ix)
in return $ Yield result (pseg, ix + 1)
in Stream fn (0, 0) Unknown
streamSegsFromVectorsUSSegd
:: (Unboxes a, Monad m)
=> Vectors a
-> USSegd
-> Stream m a
streamSegsFromVectorsUSSegd
vectors
ussegd@(USSegd _ segStarts segSources usegd)
= segStarts `seq` segSources `seq` usegd `seq` vectors `seq`
let here = "stremSegsFromVectorsUSSegd"
!segLens = USegd.takeLengths usegd
!segsTotal = USSegd.length ussegd
!elements = USegd.takeElements usegd
fnSeg (ixSeg, baSeg, ixEnd, ixElem)
= ixSeg `seq` baSeg `seq`
if ixElem >= ixEnd
then if ixSeg + 1 >= segsTotal
then return $ Done
else let ixSeg' = ixSeg + 1
sourceSeg = index here segSources ixSeg'
startSeg = index here segStarts ixSeg'
lenSeg = index here segLens ixSeg'
(arr, startArr, _)
= US.unsafeIndexUnpack vectors sourceSeg
in return $ Skip
( ixSeg'
, arr
, startArr + startSeg + lenSeg
, startArr + startSeg)
else let !result = P.indexByteArray baSeg ixElem
in return $ Yield result (ixSeg, baSeg, ixEnd, ixElem + 1)
!dummy = unsafePerformIO
$ P.newByteArray 0 >>= P.unsafeFreezeByteArray
!initState
= ( 1
, dummy
, 0
, 0)
in Stream fnSeg initState (Exact elements)
streamSegsFromVectorsUVSegd
:: (Unboxes a, Monad m)
=> Vectors a
-> UVSegd
-> Stream m a
streamSegsFromVectorsUVSegd
vectors
(UVSegd _ _ segmap _ ussegd)
= streamSegsFromVectorsUSSegdSegmap vectors ussegd segmap
streamSegsFromVectorsUSSegdSegmap
:: (Unboxes a, Monad m)
=> Vectors a
-> USSegd
-> Vector Int
-> Stream m a
streamSegsFromVectorsUSSegdSegmap
vectors ussegd@(USSegd _ segStarts segSources usegd) segmap
= segStarts `seq` segSources `seq` usegd `seq` segmap `seq`
let here = "stremSegsFromVectorsUVSegd"
!lengths = USSegd.takeLengths ussegd
!elemsTotal = U.sum $ U.map (U.index here lengths) segmap
!segsTotal = U.length segmap
!segLens = USegd.takeLengths usegd
fnSeg (ixSeg, baSeg, ixEnd, ixElem)
= ixSeg `seq` baSeg `seq`
if ixElem >= ixEnd
then if ixSeg + 1 >= segsTotal
then return $ Done
else let ixSeg' = ixSeg + 1
ixPSeg = index here segmap ixSeg'
sourceSeg = index here segSources ixPSeg
startSeg = index here segStarts ixPSeg
lenSeg = index here segLens ixPSeg
(arr, startArr, _)
= US.unsafeIndexUnpack vectors sourceSeg
in return $ Skip
( ixSeg'
, arr
, startArr + startSeg + lenSeg
, startArr + startSeg)
else let !result = P.indexByteArray baSeg ixElem
in return $ Yield result (ixSeg, baSeg, ixEnd, ixElem + 1)
!dummy = unsafePerformIO
$ P.newByteArray 0 >>= P.unsafeFreezeByteArray
!initState
= ( 1
, dummy
, 0
, 0)
in Stream fnSeg initState (Exact elemsTotal)
streamSegsFromVectorsUSSegd_split
:: (Unboxes a, Monad m)
=> Vectors a
-> USSegd
-> Vector Int
-> ((USegd,Int),Int)
-> Stream m a
streamSegsFromVectorsUSSegd_split
!vectors !ussegd
!vsegids ((!segd,!seg_off),!el_off)
= let here = "streamSegsFromVectorsUSSegd_split"
!lengths = USegd.takeLengths segd
!elemsTotal = U.sum lengths
!segsTotal = U.length lengths
!segStarts = USSegd.takeStarts ussegd
!segSources = USSegd.takeSources ussegd
vsegid seg = index here vsegids (seg + seg_off)
source pseg = index here segSources pseg
start pseg = index here segStarts pseg
len seg = index here lengths seg
fnSeg (!ixSeg, !baSeg, !ixEnd, !ixElem)
= if ixElem >= ixEnd
then if ixSeg + 1 >= segsTotal
then return $ Done
else let ixSeg' = ixSeg + 1
ixPSeg = vsegid ixSeg'
sourceSeg = source ixPSeg
startSeg = start ixPSeg
lenSeg = len ixSeg'
el_off' = if ixSeg' == 0 then el_off else 0
(arr, startArr, _)
= US.unsafeIndexUnpack vectors sourceSeg
in return $ Skip
( ixSeg'
, arr
, startArr + startSeg + el_off' + lenSeg
, startArr + startSeg + el_off')
else let !result = P.indexByteArray baSeg ixElem
in return $ Yield result (ixSeg, baSeg, ixEnd, ixElem + 1)
!dummy = unsafePerformIO
$ P.newByteArray 0 >>= P.unsafeFreezeByteArray
!initState
= ( 1
, dummy
, 0
, 0)
in Stream fnSeg initState (Exact elemsTotal)