007a67e407522a53adf87d0f4f3419c3ffaa0238
[packages/dph.git] / dph-prim-seq / Data / Array / Parallel / Unlifted / Sequential / Vectors.hs
1 {-# LANGUAGE BangPatterns, FlexibleInstances, UndecidableInstances, CPP #-}
2 #include "fusion-phases.h"
3 module Data.Array.Parallel.Unlifted.Sequential.Vectors
4 ( Unboxes
5 , Vectors(..)
6 , empty
7 , singleton
8 , unsafeIndexUnpack
9 , unsafeIndex2
10 , append
11 , fromVector
12 , toVector
13
14 , unsafeStreamVectors
15 , unsafeExtractsSU)
16 where
17 import qualified Data.Primitive.ByteArray as P
18 import qualified Data.Primitive.Array as P
19 import qualified Data.Primitive.Types as P
20
21 import qualified Data.Vector.Generic as G
22 import qualified Data.Vector.Fusion.Stream as S
23 import qualified Data.Vector.Fusion.Stream.Size as S
24 import qualified Data.Vector.Fusion.Stream.Monadic as M
25 import qualified Data.Vector.Primitive as R
26 import qualified Data.Vector.Unboxed as U
27 import qualified Data.Vector as V
28 import Data.Vector.Unboxed (Unbox)
29
30 import qualified Data.Array.Parallel.Unlifted.Sequential.USegd as USegd
31 import qualified Data.Array.Parallel.Unlifted.Sequential.USSegd as USSegd
32 import Data.Array.Parallel.Unlifted.Sequential.USSegd (USSegd(..))
33 import System.IO.Unsafe
34 import Prelude hiding (length)
35 import Debug.Trace
36
37
38 -- | Class of element types that can be used in a `Vectors`
39 class R.Prim a => Unboxes a
40 instance Unboxes Int
41 instance Unboxes Double
42
43
44 -- | A 2-dimensional array of vectors,
45 -- where the inner vectors can all have different lengths.
46 data Vectors a
47 = Vectors
48 {-# UNPACK #-} !Int -- number of vectors
49 {-# UNPACK #-} !P.ByteArray -- starting indices
50 {-# UNPACK #-} !P.ByteArray -- lengths
51 {-# UNPACK #-} !(P.Array P.ByteArray) -- chunks
52
53 instance (Unboxes a, Unbox a, Show a) => Show (Vectors a) where
54 show = show . toVector
55 {-# NOINLINE show #-}
56
57 -- | Conatruct an empty `Vectors` with no arrays of no elements.
58 empty :: Vectors a
59 empty
60 = unsafePerformIO
61 $ do mba <- P.newByteArray 0
62 ba <- P.unsafeFreezeByteArray mba
63 marr <- P.newArray 0 ba
64 arr <- P.unsafeFreezeArray marr
65 return $ Vectors 0 ba ba arr
66 {-# INLINE_U empty #-}
67
68
69 -- | Construct a `Vectors` containing data from a single unboxed array.
70 singleton :: P.Prim a => R.Vector a -> Vectors a
71 singleton vec
72 = unsafePerformIO
73 $ do R.MVector start len mbaData <- R.unsafeThaw vec
74 baData <- P.unsafeFreezeByteArray mbaData
75
76 mbaStarts <- P.newByteArray 1
77 P.writeByteArray mbaStarts 0 start
78 baStarts <- P.unsafeFreezeByteArray mbaStarts
79
80 mbaLengths <- P.newByteArray 1
81 P.writeByteArray mbaLengths 0 len
82 baLengths <- P.unsafeFreezeByteArray mbaLengths
83
84 maChunks <- P.newArray 1 baData
85 aChunks <- P.unsafeFreezeArray maChunks
86
87 return $ Vectors 1 baStarts baLengths aChunks
88 {-# INLINE_U singleton #-}
89
90
91 -- | Yield the number of vectors in a `Vectors`.
92 length :: P.Prim a => Vectors a -> Int
93 length (Vectors len _ _ _) = len
94 {-# INLINE_U length #-}
95
96
97 -- | Take one of the outer vectors from a `Vectors`.
98 unsafeIndex :: (Unboxes a, Unbox a) => Vectors a -> Int -> U.Vector a
99 unsafeIndex (Vectors _ starts lens arrs) ix
100 = G.convert
101 $ unsafePerformIO
102 $ do let start = P.indexByteArray starts ix
103 let len = P.indexByteArray lens ix
104 let arr = P.indexArray arrs ix
105 marr <- P.unsafeThawByteArray arr
106 let mvec = R.MVector start len marr
107 R.unsafeFreeze mvec
108 {-# INLINE_U unsafeIndex #-}
109
110
111 -- | Retrieve a single element from a `Vectors`,
112 -- given the outer and inner indices.
113 unsafeIndex2 :: Unboxes a => Vectors a -> Int -> Int -> a
114 unsafeIndex2 (Vectors _ starts lens arrs) ix1 ix2
115 = (arrs `P.indexArray` ix1) `P.indexByteArray` (starts `P.indexByteArray` ix1 + ix2)
116 {-# INLINE_U unsafeIndex2 #-}
117
118
119 -- | Retrieve the an inner vector from a `Vectors`.
120 --
121 -- Returns in unwrapped form:
122 -- array data, starting index in the array data, vector length.
123 unsafeIndexUnpack :: Unboxes a => Vectors a -> Int -> (P.ByteArray, Int, Int)
124 unsafeIndexUnpack (Vectors n starts lens arrs) ix
125 = ( arrs `P.indexArray` ix
126 , starts `P.indexByteArray` ix
127 , lens `P.indexByteArray` ix)
128 {-# INLINE_U unsafeIndexUnpack #-}
129
130
131 -- | Append two `Vectors`.
132 append :: Unboxes a => Vectors a -> Vectors a -> Vectors a
133 append (Vectors len1 starts1 lens1 chunks1)
134 (Vectors len2 starts2 lens2 chunks2)
135 = unsafePerformIO
136 $ do let len' = len1 + len2
137
138 -- append starts into result
139 let lenStarts1 = P.sizeofByteArray starts1
140 let lenStarts2 = P.sizeofByteArray starts2
141 maStarts <- P.newByteArray (lenStarts1 + lenStarts2)
142 P.copyByteArray maStarts 0 starts1 0 lenStarts1
143 P.copyByteArray maStarts lenStarts1 starts2 0 lenStarts2
144 starts' <- P.unsafeFreezeByteArray maStarts
145
146 -- append lens into result
147 let lenLens1 = P.sizeofByteArray lens1
148 let lenLens2 = P.sizeofByteArray lens2
149 maLens <- P.newByteArray (lenLens1 + lenLens2)
150 P.copyByteArray maLens 0 lens1 0 lenLens1
151 P.copyByteArray maLens lenStarts1 lens2 0 lenLens2
152 lens' <- P.unsafeFreezeByteArray maLens
153
154 -- append arrs into result
155 maChunks <- P.newArray len' (error "Vectors: append argh!")
156 P.copyArray maChunks 0 chunks1 0 len1
157 P.copyArray maChunks len1 chunks2 0 len2
158 chunks' <- P.unsafeFreezeArray maChunks
159
160 return $ Vectors len' starts' lens' chunks'
161 {-# INLINE_U append #-}
162
163
164 -- | Convert a boxed vector of unboxed vectors to a `Vectors`.
165 fromVector :: (Unboxes a, Unbox a) => V.Vector (U.Vector a) -> Vectors a
166 fromVector vecs
167 = unsafePerformIO
168 $ do let len = V.length vecs
169 let (barrs, vstarts, vlens) = V.unzip3 $ V.map unpackUVector vecs
170 let (baStarts, _, _) = unpackUVector $ V.convert vstarts
171 let (baLens, _, _) = unpackUVector $ V.convert vlens
172 mchunks <- P.newArray len (error "Vectors: fromVector argh!")
173 V.zipWithM_
174 (\i vec
175 -> let (ba, _, _) = unpackUVector vec
176 in P.writeArray mchunks i ba)
177 (V.enumFromN 0 len)
178 vecs
179
180 chunks <- P.unsafeFreezeArray mchunks
181 return $ Vectors len baStarts baLens chunks
182 {-# INLINE_U fromVector #-}
183
184
185 -- | Convert a `Vectors` to a boxed vector of unboxed vectors.
186 toVector :: (Unboxes a, Unbox a) => Vectors a -> V.Vector (U.Vector a)
187 toVector vectors
188 = V.map (unsafeIndex vectors)
189 $ V.enumFromN 0 (length vectors)
190 {-# INLINE_U toVector #-}
191
192
193 -- | Unpack an unboxed vector into array data, starting index, and vector length.
194 unpackUVector :: (Unbox a, P.Prim a) => U.Vector a -> (P.ByteArray, Int, Int)
195 unpackUVector vec
196 = unsafePerformIO
197 $ do let pvec = V.convert vec
198 R.MVector start len mba <- R.unsafeThaw pvec
199 ba <- P.unsafeFreezeByteArray mba
200 return (ba, start, len)
201 {-# INLINE_U unpackUVector #-}
202
203
204 -- | Pack some array data, starting index and vector length unto an unboxed vector.
205 packUVector :: (Unbox a, P.Prim a) => P.ByteArray -> Int -> Int -> U.Vector a
206 packUVector ba start len
207 = unsafePerformIO
208 $ do mba <- P.unsafeThawByteArray ba
209 pvec <- R.unsafeFreeze $ R.MVector start len mba
210 return $ G.convert pvec
211 {-# INLINE_U packUVector #-}
212
213
214 -- Stream -----------------------------------------------------------------------------------------
215 -- | Stream segments from a `Vectors`.
216 --
217 -- * There must be at least one segment in the `USSegd`, but this is not checked.
218 --
219 -- * No bounds checking is done for the `USSegd`.
220 unsafeStreamVectors :: (Unboxes a, Show a) => USSegd -> Vectors a -> S.Stream a
221 unsafeStreamVectors ussegd@(USSegd _ segStarts segSources usegd) vectors
222 = segStarts `seq` segSources `seq` usegd `seq` vectors `seq`
223 let -- Length of each segment
224 !segLens = USegd.takeLengths usegd
225
226 -- Total number of segments.
227 !segsTotal = USSegd.length ussegd
228
229 -- Total number of elements to stream.
230 !elements = USegd.takeElements usegd
231
232 -- seg, ix of that seg in usegd, length of seg, elem in seg
233 {-# INLINE_INNER fnSeg #-}
234 fnSeg (ixSeg, baSeg, ixEnd, ixElem)
235 = ixSeg `seq` baSeg `seq`
236 if ixElem >= ixEnd -- Was that the last elem in the current seg?
237 then if ixSeg + 1 >= segsTotal -- Was that last seg?
238
239 -- That was the last seg, we're done.
240 then return $ S.Done
241
242 -- Move to the next seg.
243 else let ixSeg' = ixSeg + 1
244 sourceSeg = U.unsafeIndex segSources ixSeg'
245 startSeg = U.unsafeIndex segStarts ixSeg'
246 lenSeg = U.unsafeIndex segLens ixSeg'
247 (arr, startArr, lenArr) = unsafeIndexUnpack vectors sourceSeg
248 in return $ S.Skip
249 ( ixSeg'
250 , arr
251 , startArr + startSeg + lenSeg
252 , startArr + startSeg)
253
254 -- Stream the next element from the segment.
255 else let !result = P.indexByteArray baSeg ixElem
256 in return $ S.Yield result (ixSeg, baSeg, ixEnd, ixElem + 1)
257
258 -- Starting state of the stream.
259 !initState
260 = let sourceSeg = U.unsafeIndex segSources 0
261 startSeg = U.unsafeIndex segStarts 0
262 lenSeg = U.unsafeIndex segLens 0
263 (arr, startArr, lenArr) = unsafeIndexUnpack vectors sourceSeg
264 in ( 0 -- starting segment id
265 , arr -- starting segment data
266 , startArr + startSeg + lenSeg -- segment end
267 , startArr + startSeg) -- segment start ix
268
269 -- It's important that we set the result stream size, so Data.Vector
270 -- doesn't need to add code to grow the result when it overflows.
271 in M.Stream fnSeg initState (S.Exact elements)
272 {-# INLINE_STREAM unsafeStreamVectors #-}
273
274
275 unsafeExtractsSU
276 :: (Unboxes a, U.Unbox a, Show a)
277 => USSegd -> Vectors a -> U.Vector a
278
279 unsafeExtractsSU ussegd vectors
280 = G.unstream $ unsafeStreamVectors ussegd vectors
281 {-# INLINE_U unsafeExtractsSU #-}
282