a442b236ddf8b0be3363c31504adc33050538235
[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 , length
18 , singleton
19 , unsafeIndex
20 , unsafeIndex2
21 , unsafeIndexUnpack
22 , append
23 , fromVector
24 , toVector)
25 where
26 import qualified Data.Primitive.ByteArray as P
27 import qualified Data.Primitive.Array 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 System.IO.Unsafe
36 import Prelude hiding (length)
37 import Data.Word
38
39
40 -- | Class of element types that can be used in a `Vectors`
41 class R.Prim a => Unboxes a
42 instance Unboxes Int
43 instance Unboxes Word8
44 instance Unboxes Float
45 instance Unboxes Double
46
47
48 -- | A 2-dimensional array,
49 -- where the inner arrays can all have different lengths.
50 data Vectors a
51 = Vectors
52 {-# UNPACK #-} !Int -- number of inner vectors
53 {-# UNPACK #-} !P.ByteArray -- starting index of each vector in its chunk
54 {-# UNPACK #-} !P.ByteArray -- lengths of each inner vector
55 {-# UNPACK #-} !(P.Array P.ByteArray) -- chunks
56
57 instance (Unboxes a, Unbox a, Show a) => Show (Vectors a) where
58 show = show . toVector
59 {-# NOINLINE show #-}
60
61 -- | Construct an empty `Vectors` with no arrays of no elements.
62 empty :: Vectors a
63 empty
64 = unsafePerformIO
65 $ do mba <- P.newByteArray 0
66 ba <- P.unsafeFreezeByteArray mba
67 marr <- P.newArray 0 ba
68 arr <- P.unsafeFreezeArray marr
69 return $ Vectors 0 ba ba arr
70 {-# INLINE_U empty #-}
71
72
73 -- | Construct a `Vectors` containing data from a single unboxed array.
74 singleton :: (Unboxes a, Unbox a) => U.Vector a -> Vectors a
75 singleton vec
76 = unsafePerformIO
77 $ do R.MVector start len mbaData <- R.unsafeThaw $ G.convert vec
78 baData <- P.unsafeFreezeByteArray mbaData
79
80 mbaStarts <- P.newByteArray (P.sizeOf (undefined :: Int))
81 P.writeByteArray mbaStarts 0 start
82 baStarts <- P.unsafeFreezeByteArray mbaStarts
83
84 mbaLengths <- P.newByteArray (P.sizeOf (undefined :: Int))
85 P.writeByteArray mbaLengths 0 len
86 baLengths <- P.unsafeFreezeByteArray mbaLengths
87
88 maChunks <- P.newArray 1 baData
89 aChunks <- P.unsafeFreezeArray maChunks
90
91 return $ Vectors 1 baStarts baLengths aChunks
92 {-# INLINE_U singleton #-}
93
94
95 -- | Yield the number of vectors in a `Vectors`.
96 length :: Unboxes a => Vectors a -> Int
97 length (Vectors len _ _ _) = len
98 {-# INLINE_U length #-}
99
100
101 -- | Take one of the outer vectors from a `Vectors`.
102 unsafeIndex :: (Unboxes a, Unbox a) => Vectors a -> Int -> U.Vector a
103 unsafeIndex (Vectors _ starts lens arrs) ix
104 = G.convert
105 $ unsafePerformIO
106 $ do let start = P.indexByteArray starts ix
107 let len = P.indexByteArray lens ix
108 let arr = P.indexArray arrs ix
109 marr <- P.unsafeThawByteArray arr
110 let mvec = R.MVector start len marr
111 R.unsafeFreeze mvec
112 {-# INLINE_U unsafeIndex #-}
113
114
115 -- | Retrieve a single element from a `Vectors`,
116 -- given the outer and inner indices.
117 unsafeIndex2 :: Unboxes a => Vectors a -> Int -> Int -> a
118 unsafeIndex2 (Vectors _ starts _ arrs) ix1 ix2
119 = (arrs `P.indexArray` ix1) `P.indexByteArray` ((starts `P.indexByteArray` ix1) + ix2)
120 {-# INLINE_U unsafeIndex2 #-}
121
122
123 -- | Retrieve an inner array from a `Vectors`, returning the array data,
124 -- starting index in the data, and vector length.
125 unsafeIndexUnpack :: Unboxes a => Vectors a -> Int -> (P.ByteArray, Int, Int)
126 unsafeIndexUnpack (Vectors _ starts lens arrs) ix
127 = ( arrs `P.indexArray` ix
128 , starts `P.indexByteArray` ix
129 , lens `P.indexByteArray` ix)
130 {-# INLINE_U unsafeIndexUnpack #-}
131
132
133 -- | Append two `Vectors`.
134 --
135 -- * Important: appending two `Vectors` involes work proportional to
136 -- the length of the outer arrays, not the size of the inner ones.
137 append :: (Unboxes a, Unbox a, Show a) => Vectors a -> Vectors a -> Vectors a
138 append (Vectors len1 starts1 lens1 chunks1)
139 (Vectors len2 starts2 lens2 chunks2)
140 = unsafePerformIO
141 $ do let len' = len1 + len2
142
143 -- append starts into result
144 let lenStarts1 = P.sizeofByteArray starts1
145 let lenStarts2 = P.sizeofByteArray starts2
146 maStarts <- P.newByteArray (lenStarts1 + lenStarts2)
147 P.copyByteArray maStarts 0 starts1 0 lenStarts1
148 P.copyByteArray maStarts lenStarts1 starts2 0 lenStarts2
149 starts' <- P.unsafeFreezeByteArray maStarts
150
151 -- append lens into result
152 let lenLens1 = P.sizeofByteArray lens1
153 let lenLens2 = P.sizeofByteArray lens2
154 maLens <- P.newByteArray (lenLens1 + lenLens2)
155 P.copyByteArray maLens 0 lens1 0 lenLens1
156 P.copyByteArray maLens lenStarts1 lens2 0 lenLens2
157 lens' <- P.unsafeFreezeByteArray maLens
158
159 -- append arrs into result
160 maChunks <- P.newArray len' (error "Vectors: append argh!")
161 P.copyArray maChunks 0 chunks1 0 len1
162 P.copyArray maChunks len1 chunks2 0 len2
163 chunks' <- P.unsafeFreezeArray maChunks
164
165
166 let result = Vectors len' starts' lens' chunks'
167 return $ result
168 {-# INLINE_U append #-}
169
170
171 -- | Convert a boxed vector of unboxed vectors to a `Vectors`.
172 fromVector :: (Unboxes a, Unbox a) => V.Vector (U.Vector a) -> Vectors a
173 fromVector vecs
174 = unsafePerformIO
175 $ do let len = V.length vecs
176 let (_, vstarts, vlens) = V.unzip3 $ V.map unpackUVector vecs
177 let (baStarts, _, _) = unpackUVector $ V.convert vstarts
178 let (baLens, _, _) = unpackUVector $ V.convert vlens
179 mchunks <- P.newArray len (error "Vectors: fromVector argh!")
180 V.zipWithM_
181 (\i vec
182 -> let (ba, _, _) = unpackUVector vec
183 in P.writeArray mchunks i ba)
184 (V.enumFromN 0 len)
185 vecs
186
187 chunks <- P.unsafeFreezeArray mchunks
188 return $ Vectors len baStarts baLens chunks
189 {-# INLINE_U fromVector #-}
190
191
192 -- | Convert a `Vectors` to a boxed vector of unboxed vectors.
193 toVector :: (Unboxes a, Unbox a) => Vectors a -> V.Vector (U.Vector a)
194 toVector vectors
195 = V.map (unsafeIndex vectors)
196 $ V.enumFromN 0 (length vectors)
197 {-# INLINE_U toVector #-}
198
199
200 -- | Unpack an unboxed vector into array data, starting index, and vector length.
201 unpackUVector :: (Unbox a, P.Prim a) => U.Vector a -> (P.ByteArray, Int, Int)
202 unpackUVector vec
203 = unsafePerformIO
204 $ do let pvec = V.convert vec
205 R.MVector start len mba <- R.unsafeThaw pvec
206 ba <- P.unsafeFreezeByteArray mba
207 return (ba, start, len)
208 {-# INLINE_U unpackUVector #-}
209