dph-prim-seq: use ArrayArray for Vectors
[packages/dph.git] / dph-prim-seq / Data / Array / Parallel / Unlifted / Vectors.hs
1 {-# LANGUAGE BangPatterns, FlexibleInstances, UndecidableInstances, CPP #-}
2 #include "fusion-phases.h"
3
4 -- | Irregular two dimensional arrays.
5 --
6 -- * TODO: The inner arrays should be unboxed so we don't get an unboxing overhead
7 -- for every call to unsafeIndex2. This might need an extension to the GHC
8 -- runtime if we alwo want to convert a U.Vector directly to this form.
9 --
10 -- * TODO: We currently only allow primitive types to be in a Vectors, but
11 -- in future we'll want `Vectors` of tuples etc.
12 --
13 module Data.Array.Parallel.Unlifted.Vectors
14 ( Vectors(..)
15 , Unboxes
16 , empty
17 , singleton
18 , length
19 , unsafeIndex
20 , unsafeIndex2
21 , unsafeIndexUnpack
22 , append
23 , fromVector
24 , toVector)
25 where
26 import qualified Data.Array.Parallel.Unlifted.ArrayArray as AA
27 import qualified Data.Primitive.ByteArray as P
28 import qualified Data.Primitive.Types as P
29 import qualified Data.Primitive as P
30 import qualified Data.Vector.Generic as G
31 import qualified Data.Vector.Primitive as R
32 import qualified Data.Vector.Unboxed as U
33 import qualified Data.Vector as V
34 import Data.Vector.Unboxed (Unbox)
35 import Prelude hiding (length)
36 import Data.Word
37 import Control.Monad.ST
38
39 -- | Class of element types that can be used in a `Vectors`
40 class R.Prim a => Unboxes a
41 instance Unboxes Int
42 instance Unboxes Word8
43 instance Unboxes Float
44 instance Unboxes Double
45
46
47 -- | A 2-dimensional array,
48 -- where the inner arrays can all have different lengths.
49 data Vectors a
50 = Vectors
51 {-# UNPACK #-} !Int -- number of inner vectors
52 {-# UNPACK #-} !P.ByteArray -- starting index of each vector in its chunk
53 {-# UNPACK #-} !P.ByteArray -- lengths of each inner vector
54 {-# UNPACK #-} !(AA.ArrayArray P.ByteArray) -- chunks
55
56
57 instance (Unboxes a, Unbox a, Show a) => Show (Vectors a) where
58 show = show . toVector
59 {-# NOINLINE show #-}
60
61
62 -- | Construct an empty `Vectors` with no arrays of no elements.
63 empty :: Vectors a
64 empty
65 = runST
66 $ do mba <- P.newByteArray 0
67 ba <- P.unsafeFreezeByteArray mba
68
69 maa <- AA.newArrayArray 0
70 AA.writeArrayArray maa 0 ba
71 aa <- AA.unsafeFreezeArrayArray maa
72
73 return $ Vectors 0 ba ba aa
74 {-# INLINE_U empty #-}
75
76
77 -- | Construct a `Vectors` containing data from a single unboxed array.
78 singleton :: (Unboxes a, Unbox a) => U.Vector a -> Vectors a
79 singleton vec
80 = runST
81 $ do R.MVector start len mbaData <- R.unsafeThaw $ G.convert vec
82 baData <- P.unsafeFreezeByteArray mbaData
83
84 mbaStarts <- P.newByteArray (P.sizeOf (undefined :: Int))
85 P.writeByteArray mbaStarts 0 start
86 baStarts <- P.unsafeFreezeByteArray mbaStarts
87
88 mbaLengths <- P.newByteArray (P.sizeOf (undefined :: Int))
89 P.writeByteArray mbaLengths 0 len
90 baLengths <- P.unsafeFreezeByteArray mbaLengths
91
92 maaChunks <- AA.newArrayArray 1
93 AA.writeArrayArray maaChunks 0 baData
94 aaChunks <- AA.unsafeFreezeArrayArray maaChunks
95
96 return $ Vectors 1 baStarts baLengths aaChunks
97 {-# INLINE_U singleton #-}
98
99
100 -- | Yield the number of vectors in a `Vectors`.
101 length :: Unboxes a => Vectors a -> Int
102 length (Vectors len _ _ _) = len
103 {-# INLINE_U length #-}
104
105
106 -- | Take one of the outer vectors from a `Vectors`.
107 unsafeIndex :: (Unboxes a, Unbox a) => Vectors a -> Int -> U.Vector a
108 unsafeIndex (Vectors _ starts lens arrs) ix
109 = G.convert
110 $ runST
111 $ do let start = P.indexByteArray starts ix
112 let len = P.indexByteArray lens ix
113 let arr = AA.indexArrayArray arrs ix
114 marr <- P.unsafeThawByteArray arr
115 let mvec = R.MVector start len marr
116 R.unsafeFreeze mvec
117 {-# INLINE_U unsafeIndex #-}
118
119
120 -- | Retrieve a single element from a `Vectors`,
121 -- given the outer and inner indices.
122 unsafeIndex2 :: Unboxes a => Vectors a -> Int -> Int -> a
123 unsafeIndex2 (Vectors _ starts _ arrs) ix1 ix2
124 = (arrs `AA.indexArrayArray` ix1) `P.indexByteArray` ((starts `P.indexByteArray` ix1) + ix2)
125 {-# INLINE_U unsafeIndex2 #-}
126
127
128 -- | Retrieve an inner array from a `Vectors`, returning the array data,
129 -- starting index in the data, and vector length.
130 unsafeIndexUnpack :: Unboxes a => Vectors a -> Int -> (P.ByteArray, Int, Int)
131 unsafeIndexUnpack (Vectors _ starts lens arrs) ix
132 = ( arrs `AA.indexArrayArray` ix
133 , starts `P.indexByteArray` ix
134 , lens `P.indexByteArray` ix)
135 {-# INLINE_U unsafeIndexUnpack #-}
136
137
138 -- | Append two `Vectors`.
139 --
140 -- * Important: appending two `Vectors` involes work proportional to
141 -- the length of the outer arrays, not the size of the inner ones.
142 append :: (Unboxes a, Unbox a, Show a) => Vectors a -> Vectors a -> Vectors a
143 append (Vectors len1 starts1 lens1 chunks1)
144 (Vectors len2 starts2 lens2 chunks2)
145 = runST
146 $ do let len' = len1 + len2
147
148 -- append starts into result
149 let lenStarts1 = P.sizeofByteArray starts1
150 let lenStarts2 = P.sizeofByteArray starts2
151 maStarts <- P.newByteArray (lenStarts1 + lenStarts2)
152 P.copyByteArray maStarts 0 starts1 0 lenStarts1
153 P.copyByteArray maStarts lenStarts1 starts2 0 lenStarts2
154 starts' <- P.unsafeFreezeByteArray maStarts
155
156 -- append lens into result
157 let lenLens1 = P.sizeofByteArray lens1
158 let lenLens2 = P.sizeofByteArray lens2
159 maLens <- P.newByteArray (lenLens1 + lenLens2)
160 P.copyByteArray maLens 0 lens1 0 lenLens1
161 P.copyByteArray maLens lenStarts1 lens2 0 lenLens2
162 lens' <- P.unsafeFreezeByteArray maLens
163
164 -- append arrs into result
165 maChunks <- AA.newArrayArray len'
166 AA.copyArrayArray maChunks 0 chunks1 0 len1
167 AA.copyArrayArray maChunks len1 chunks2 0 len2
168 chunks' <- AA.unsafeFreezeArrayArray maChunks
169
170 let result = Vectors len' starts' lens' chunks'
171 return $ result
172 {-# INLINE_U append #-}
173
174
175 -- | Convert a boxed vector of unboxed vectors to a `Vectors`.
176 fromVector :: (Unboxes a, Unbox a) => V.Vector (U.Vector a) -> Vectors a
177 fromVector vecs
178 = runST
179 $ do let len = V.length vecs
180 let (_, vstarts, vlens) = V.unzip3 $ V.map unpackUVector vecs
181 let (baStarts, _, _) = unpackUVector $ V.convert vstarts
182 let (baLens, _, _) = unpackUVector $ V.convert vlens
183 mchunks <- AA.newArrayArray len
184 V.zipWithM_
185 (\i vec
186 -> let (ba, _, _) = unpackUVector vec
187 in AA.writeArrayArray mchunks i ba)
188 (V.enumFromN 0 len)
189 vecs
190
191 chunks <- AA.unsafeFreezeArrayArray mchunks
192 return $ Vectors len baStarts baLens chunks
193 {-# INLINE_U fromVector #-}
194
195
196 -- | Convert a `Vectors` to a boxed vector of unboxed vectors.
197 toVector :: (Unboxes a, Unbox a) => Vectors a -> V.Vector (U.Vector a)
198 toVector vectors
199 = V.map (unsafeIndex vectors)
200 $ V.enumFromN 0 (length vectors)
201 {-# INLINE_U toVector #-}
202
203
204 -- | Unpack an unboxed vector into array data, starting index, and vector length.
205 unpackUVector :: (Unbox a, P.Prim a) => U.Vector a -> (P.ByteArray, Int, Int)
206 unpackUVector vec
207 = runST
208 $ do let pvec = V.convert vec
209 R.MVector start len mba <- R.unsafeThaw pvec
210 ba <- P.unsafeFreezeByteArray mba
211 return (ba, start, len)
212 {-# INLINE_U unpackUVector #-}
213