Separate length from data in array representation
[packages/dph.git] / dph-common / Data / Array / Parallel / Lifted / PArray.hs
1 {-# LANGUAGE CPP #-}
2
3 #include "fusion-phases.h"
4
5 module Data.Array.Parallel.Lifted.PArray (
6 PArray(..), PData,
7
8 PA(..),
9 lengthPA#, dataPA#, replicatePA#, replicatelPA#, repeatPA#, repeatcPA#,
10 emptyPA, indexPA#, extractPA#, bpermutePA#, appPA#, applPA#,
11 packPA#, combine2PA#, fromListPA#, fromListPA, nfPA,
12
13 replicatePD, replicatelPD, repeatPD, repeatcPD, emptyPD,
14 indexPD, extractPD, bpermutePD, appPD, applPD,
15 packPD, combine2PD, fromListPD, fromListPD, nfPD,
16
17 PRepr, PR(..), mkPR, mkReprPA,
18
19 T_replicatePR, T_replicatelPR, T_repeatPR, T_repeatcPR, T_emptyPR,
20 T_indexPR, T_extractPR, T_bpermutePR, T_appPR, T_applPR,
21 T_packPR, T_combine2PR, T_fromListPR, T_fromListPR, T_nfPR
22 ) where
23
24 import qualified Data.Array.Parallel.Unlifted as U
25 import Data.Array.Parallel.Lifted.Selector
26 import Data.Array.Parallel.Lifted.Unboxed ( elementsSegd# )
27 import GHC.Exts (Int#, Int(..), (+#), (*#))
28
29 -- |Lifted parallel arrays
30 --
31 data PArray a = PArray Int# (PData a)
32 data family PData a
33
34 -- |Representation types
35 --
36 type family PRepr a
37
38 -- |Dictionaries
39 --
40
41 data PA a = PA {
42 toPRepr :: a -> PRepr a
43 , fromPRepr :: PRepr a -> a
44 , toArrPRepr :: PData a -> PData (PRepr a)
45 , fromArrPRepr :: PData (PRepr a) -> PData a
46 , dictPRepr :: PR (PRepr a)
47 }
48
49
50 type T_emptyPR a = PData a
51
52 type T_replicatePR a = Int# -> a -> PData a
53
54 type T_replicatelPR a = U.Segd -- segd of result array
55 -> PData a -> PData a
56
57 type T_repeatPR a = Int# -- number of times to repeat
58 -> Int# -- length of src array
59 -> PData a -> PData a
60
61 type T_repeatcPR a = Int# -- length of result array
62 -> U.Array Int -- number of times each segment
63 -- is repeated
64 -> U.Segd -- src segd
65 -> PData a -> PData a
66
67 type T_indexPR a = PData a -> Int# -> a
68
69 type T_extractPR a = PData a
70 -> Int# -- starting index
71 -> Int# -- length of result array
72 -> PData a
73
74 type T_bpermutePR a = PData a
75 -> Int# -- result length
76 -> U.Array Int -- indices
77 -> PData a
78
79 type T_appPR a = PData a -> PData a -> PData a
80
81 type T_applPR a = U.Segd -> PData a -- src segd/data 1
82 -> U.Segd -> PData a -- src segd/data 2
83 -> PData a
84
85 type T_packPR a = PData a
86 -> Int# -- result length
87 -> U.Array Bool -- flags
88 -> PData a
89
90 type T_combine2PR a = Int# -- result length
91 -> Sel2 -- selector
92 -> PData a -> PData a -> PData a
93
94 type T_fromListPR a = Int# -> [a] -> PData a
95
96 type T_nfPR a = PData a -> ()
97
98 data PR a = PR {
99 emptyPR :: T_emptyPR a
100 , replicatePR :: T_replicatePR a
101 , replicatelPR :: T_replicatelPR a
102 , repeatPR :: T_repeatPR a
103 , repeatcPR :: T_repeatcPR a
104 , indexPR :: T_indexPR a
105 , extractPR :: T_extractPR a
106 , bpermutePR :: T_bpermutePR a
107 , appPR :: T_appPR a
108 , applPR :: T_applPR a
109 , packPR :: T_packPR a
110 , combine2PR :: T_combine2PR a
111 , fromListPR :: T_fromListPR a
112 , nfPR :: T_nfPR a
113 }
114
115 emptyPD :: PA a -> T_emptyPR a
116 {-# INLINE_PA emptyPD #-}
117 emptyPD pa = fromArrPRepr pa
118 $ emptyPR (dictPRepr pa)
119
120 replicatePD :: PA a -> T_replicatePR a
121 {-# INLINE_PA replicatePD #-}
122 replicatePD pa n# x = fromArrPRepr pa
123 . replicatePR (dictPRepr pa) n#
124 $ toPRepr pa x
125
126 replicatelPD :: PA a -> T_replicatelPR a
127 {-# INLINE_PA replicatelPD #-}
128 replicatelPD pa segd xs = fromArrPRepr pa
129 . replicatelPR (dictPRepr pa) segd
130 $ toArrPRepr pa xs
131
132 repeatPD :: PA a -> T_repeatPR a
133 {-# INLINE_PA repeatPD #-}
134 repeatPD pa n# len# xs = fromArrPRepr pa
135 . repeatPR (dictPRepr pa) n# len#
136 $ toArrPRepr pa xs
137
138 repeatcPD :: PA a -> T_repeatcPR a
139 {-# INLINE_PA repeatcPD #-}
140 repeatcPD pa n# ns segd xs = fromArrPRepr pa
141 . repeatcPR (dictPRepr pa) n# ns segd
142 $ toArrPRepr pa xs
143
144 indexPD :: PA a -> T_indexPR a
145 {-# INLINE_PA indexPD #-}
146 indexPD pa xs i# = fromPRepr pa
147 $ indexPR (dictPRepr pa) (toArrPRepr pa xs) i#
148
149 extractPD :: PA a -> T_extractPR a
150 {-# INLINE_PA extractPD #-}
151 extractPD pa xs i# m# = fromArrPRepr pa
152 $ extractPR (dictPRepr pa) (toArrPRepr pa xs) i# m#
153
154 bpermutePD :: PA a -> T_bpermutePR a
155 {-# INLINE bpermutePD #-}
156 bpermutePD pa xs n# is = fromArrPRepr pa
157 $ bpermutePR (dictPRepr pa) (toArrPRepr pa xs) n# is
158
159 appPD :: PA a -> T_appPR a
160 {-# INLINE_PA appPD #-}
161 appPD pa xs ys = fromArrPRepr pa
162 $ appPR (dictPRepr pa) (toArrPRepr pa xs) (toArrPRepr pa ys)
163
164 applPD :: PA a -> T_applPR a
165 {-# INLINE_PA applPD #-}
166 applPD pa is xs js ys = fromArrPRepr pa
167 $ applPR (dictPRepr pa) is (toArrPRepr pa xs)
168 js (toArrPRepr pa ys)
169
170 packPD :: PA a -> T_packPR a
171 {-# INLINE_PA packPD #-}
172 packPD pa xs n# bs = fromArrPRepr pa
173 $ packPR (dictPRepr pa) (toArrPRepr pa xs) n# bs
174
175 combine2PD :: PA a -> T_combine2PR a
176 {-# INLINE_PA combine2PD #-}
177 combine2PD pa n# sel as bs
178 = fromArrPRepr pa
179 $ combine2PR (dictPRepr pa) n# sel (toArrPRepr pa as)
180 (toArrPRepr pa bs)
181
182 fromListPD :: PA a -> T_fromListPR a
183 {-# INLINE_PA fromListPD #-}
184 fromListPD pa n# xs = fromArrPRepr pa
185 $ fromListPR (dictPRepr pa) n# (map (toPRepr pa) xs)
186
187 nfPD :: PA a -> T_nfPR a
188 {-# INLINE nfPD #-}
189 nfPD pa xs = nfPR (dictPRepr pa) (toArrPRepr pa xs)
190
191
192
193
194 lengthPA# :: PArray a -> Int#
195 {-# INLINE_PA lengthPA# #-}
196 lengthPA# (PArray n# _) = n#
197
198 dataPA# :: PArray a -> PData a
199 {-# INLINE_PA dataPA# #-}
200 dataPA# (PArray _ d) = d
201
202 emptyPA :: PA a -> PArray a
203 {-# INLINE_PA emptyPA #-}
204 emptyPA pa = PArray 0# (emptyPD pa)
205
206 replicatePA# :: PA a -> Int# -> a -> PArray a
207 {-# INLINE_PA replicatePA# #-}
208 replicatePA# pa n# x = PArray n# (replicatePD pa n# x)
209
210 replicatelPA# :: PA a -> U.Segd -> PArray a -> PArray a
211 {-# INLINE_PA replicatelPA# #-}
212 replicatelPA# pa segd (PArray n# xs)
213 = PArray (elementsSegd# segd) (replicatelPD pa segd xs)
214
215 repeatPA# :: PA a -> Int# -> PArray a -> PArray a
216 {-# INLINE_PA repeatPA# #-}
217 repeatPA# pa m# (PArray n# xs) = PArray (m# *# n#) (repeatPD pa m# n# xs)
218
219 repeatcPA# :: PA a -> U.Array Int -> U.Segd -> PArray a -> PArray a
220 {-# INLINE_PA repeatcPA# #-}
221 repeatcPA# pa ns segd (PArray n# xs)
222 = case U.sum (U.zipWith (*) ns (U.lengthsSegd segd)) of
223 I# m# -> PArray m# (repeatcPD pa m# ns segd xs)
224
225 indexPA# :: PA a -> PArray a -> Int# -> a
226 {-# INLINE_PA indexPA# #-}
227 indexPA# pa (PArray _ xs) i# = indexPD pa xs i#
228
229 extractPA# :: PA a -> PArray a -> Int# -> Int# -> PArray a
230 {-# INLINE_PA extractPA# #-}
231 extractPA# pa (PArray _ xs) i# n# = PArray n# (extractPD pa xs i# n#)
232
233 bpermutePA# :: PA a -> PArray a -> Int# -> U.Array Int -> PArray a
234 {-# INLINE bpermutePA# #-}
235 bpermutePA# pa (PArray _ xs) n# is = PArray n# (bpermutePD pa xs n# is)
236
237 appPA# :: PA a -> PArray a -> PArray a -> PArray a
238 {-# INLINE_PA appPA# #-}
239 appPA# pa (PArray m# xs) (PArray n# ys) = PArray (m# +# n#) (appPD pa xs ys)
240
241 applPA# :: PA a -> U.Segd -> PArray a -> U.Segd -> PArray a -> PArray a
242 {-# INLINE_PA applPA# #-}
243 applPA# pa is (PArray m# xs) js (PArray n# ys)
244 = PArray (m# +# n#) (applPD pa is xs js ys)
245
246 packPA# :: PA a -> PArray a -> Int# -> U.Array Bool -> PArray a
247 {-# INLINE_PA packPA# #-}
248 packPA# pa (PArray _ xs) n# bs = PArray n# (packPD pa xs n# bs)
249
250 combine2PA# :: PA a -> Int# -> Sel2 -> PArray a -> PArray a -> PArray a
251 {-# INLINE_PA combine2PA# #-}
252 combine2PA# pa n# sel (PArray _ as) (PArray _ bs)
253 = PArray n# (combine2PD pa n# sel as bs)
254
255 fromListPA# :: PA a -> Int# -> [a] -> PArray a
256 {-# INLINE_PA fromListPA# #-}
257 fromListPA# pa n# xs = PArray n# (fromListPD pa n# xs)
258
259 fromListPA :: PA a -> [a] -> PArray a
260 {-# INLINE fromListPA #-}
261 fromListPA pa xs = case length xs of
262 I# n# -> fromListPA# pa n# xs
263
264 nfPA :: PA a -> PArray a -> ()
265 {-# INLINE nfPA #-}
266 nfPA pa (PArray _ xs) = nfPD pa xs
267
268 mkPR :: PA a -> PR a
269 {-# INLINE mkPR #-}
270 mkPR pa = PR {
271 emptyPR = emptyPD pa
272 , replicatePR = replicatePD pa
273 , replicatelPR = replicatelPD pa
274 , repeatPR = repeatPD pa
275 , repeatcPR = repeatcPD pa
276 , indexPR = indexPD pa
277 , bpermutePR = bpermutePD pa
278 , appPR = appPD pa
279 , applPR = applPD pa
280 , packPR = packPD pa
281 , combine2PR = combine2PD pa
282 , fromListPR = fromListPD pa
283 , nfPR = nfPD pa
284 }
285
286 mkReprPA :: (a ~ PRepr a) => PR a -> PA a
287 {-# INLINE mkReprPA #-}
288 mkReprPA pr = PA {
289 toPRepr = id
290 , fromPRepr = id
291 , toArrPRepr = id
292 , fromArrPRepr = id
293 , dictPRepr = pr
294 }
295