removed duplicate import
[packages/dph.git] / dph-lifted-copy / Data / Array / Parallel / PArray / PDataInstances.hs
1 {-# LANGUAGE CPP, TemplateHaskell, EmptyDataDecls #-}
2 {-# OPTIONS -fno-warn-orphans -fno-warn-missing-methods #-}
3
4 #include "fusion-phases.h"
5
6 -- | Instances for the PData class
7 module Data.Array.Parallel.PArray.PDataInstances(
8 PData(..), PDatas(..), Sels2,
9 pvoid,
10 punit,
11
12 -- * Operators on arrays of tuples
13 zipPA#, unzipPA#, zip3PA#, unzip3PA#,
14 zip4PA#, unzip4PA#, zip5PA#, unzip5PA#, zip6PA#, unzip6PA#,
15 zip7PA#, unzip7PA#, zip8PA#, unzip8PA#,
16
17 -- * Operators on nested arrays
18 segdPA#, concatPA#, segmentPA#, copySegdPA#
19 )
20 where
21 import Data.Array.Parallel.PArray.Base
22 import Data.Array.Parallel.PArray.PData
23 import Data.Array.Parallel.PArray.PRepr
24 import Data.Array.Parallel.PArray.Types
25 import Data.Array.Parallel.Lifted.TH.Repr
26 import Data.Array.Parallel.Lifted.Unboxed (elementsSegd#, elementsSel2_0#, elementsSel2_1#)
27 import Data.Array.Parallel.Base.DTrace (traceFn)
28 import Data.Array.Parallel.Base (intToTag)
29 import qualified Data.Array.Parallel.Unlifted as U
30 import Data.List (unzip4, unzip5, unzip6, unzip7)
31 import GHC.Exts (Int(..), Int#)
32
33 -- Extra unzips ------------
34
35 -- We need unzips at large tuples (for tuple instances) as the closure environments generated by the
36 -- vectoriser are tuples whose arity is determined by the number of free variables.
37
38 unzip8 :: [(a,b,c,d,e,f,g,h)] -> ([a],[b],[c],[d],[e],[f],[g],[h])
39 unzip8 = foldr (\(a,b,c,d,e,f,g,h) ~(as,bs,cs,ds,es,fs,gs,hs) ->
40 (a:as,b:bs,c:cs,d:ds,e:es,f:fs,g:gs,h:hs))
41 ([],[],[],[],[],[],[],[])
42
43 unzip9 :: [(a,b,c,d,e,f,g,h,i)] -> ([a],[b],[c],[d],[e],[f],[g],[h],[i])
44 unzip9 = foldr (\(a,b,c,d,e,f,g,h,i) ~(as,bs,cs,ds,es,fs,gs,hs,is) ->
45 (a:as,b:bs,c:cs,d:ds,e:es,f:fs,g:gs,h:hs,i:is))
46 ([],[],[],[],[],[],[],[],[])
47
48 unzip10 :: [(a,b,c,d,e,f,g,h,i,j)] -> ([a],[b],[c],[d],[e],[f],[g],[h],[i],[j])
49 unzip10 = foldr (\(a,b,c,d,e,f,g,h,i,j) ~(as,bs,cs,ds,es,fs,gs,hs,is,js) ->
50 (a:as,b:bs,c:cs,d:ds,e:es,f:fs,g:gs,h:hs,i:is,j:js))
51 ([],[],[],[],[],[],[],[],[],[])
52
53 unzip11 :: [(a,b,c,d,e,f,g,h,i,j,k)] -> ([a],[b],[c],[d],[e],[f],[g],[h],[i],[j],[k])
54 unzip11 = foldr (\(a,b,c,d,e,f,g,h,i,j,k) ~(as,bs,cs,ds,es,fs,gs,hs,is,js,ks) ->
55 (a:as,b:bs,c:cs,d:ds,e:es,f:fs,g:gs,h:hs,i:is,j:js,k:ks))
56 ([],[],[],[],[],[],[],[],[],[],[])
57
58 unzip12 :: [(a,b,c,d,e,f,g,h,i,j,k,l)] -> ([a],[b],[c],[d],[e],[f],[g],[h],[i],[j],[k],[l])
59 unzip12 = foldr (\(a,b,c,d,e,f,g,h,i,j,k,l) ~(as,bs,cs,ds,es,fs,gs,hs,is,js,ks,ls) ->
60 (a:as,b:bs,c:cs,d:ds,e:es,f:fs,g:gs,h:hs,i:is,j:js,k:ks,l:ls))
61 ([],[],[],[],[],[],[],[],[],[],[],[])
62
63 unzip13 :: [(a,b,c,d,e,f,g,h,i,j,k,l,m)] -> ([a],[b],[c],[d],[e],[f],[g],[h],[i],[j],[k],[l],[m])
64 unzip13 = foldr (\(a,b,c,d,e,f,g,h,i,j,k,l,m) ~(as,bs,cs,ds,es,fs,gs,hs,is,js,ks,ls,ms) ->
65 (a:as,b:bs,c:cs,d:ds,e:es,f:fs,g:gs,h:hs,i:is,j:js,k:ks,l:ls,m:ms))
66 ([],[],[],[],[],[],[],[],[],[],[],[],[])
67
68 unzip14 :: [(a,b,c,d,e,f,g,h,i,j,k,l,m,n)]
69 -> ([a],[b],[c],[d],[e],[f],[g],[h],[i],[j],[k],[l],[m],[n])
70 unzip14 = foldr (\(a,b,c,d,e,f,g,h,i,j,k,l,m,n) ~(as,bs,cs,ds,es,fs,gs,hs,is,js,ks,ls,ms,ns) ->
71 (a:as,b:bs,c:cs,d:ds,e:es,f:fs,g:gs,h:hs,i:is,j:js,k:ks,l:ls,m:ms,n:ns))
72 ([],[],[],[],[],[],[],[],[],[],[],[],[],[])
73
74 unzip15 :: [(a,b,c,d,e,f,g,h,i,j,k,l,m,n,o)]
75 -> ([a],[b],[c],[d],[e],[f],[g],[h],[i],[j],[k],[l],[m],[n],[o])
76 unzip15 = foldr (\(a,b,c,d,e,f,g,h,i,j,k,l,m,n,o) ~(as,bs,cs,ds,es,fs,gs,hs,is,js,ks,ls,ms,ns,os) ->
77 (a:as,b:bs,c:cs,d:ds,e:es,f:fs,g:gs,h:hs,i:is,j:js,k:ks,l:ls,m:ms,n:ns,o:os))
78 ([],[],[],[],[],[],[],[],[],[],[],[],[],[],[])
79
80
81 -- Void -----------------------------------------------------------------------
82
83 -- | The Void type is used when representing enumerations.
84 -- A type like Bool is represented as @Sum2 Void Void@, meaning that we only
85 -- only care about the tag of the data constructor and not its argumnent.
86 --
87 data instance PData Void
88
89 pvoid :: PData Void
90 pvoid = error "Data.Array.Parallel.PData Void"
91
92 $(voidPRInstance ''Void 'void 'pvoid)
93
94
95 -- Unit -----------------------------------------------------------------------
96 -- | An array of unit values is represented by a single constructor.
97 -- There is only one possible value, so we only need to record it once.
98 --
99 -- We often uses arrays of unit values as the environmnent portion of a
100 -- lifted closure. For example, suppose we vectorise the unary function
101 -- @neg@. This function has no environment, so we construct the closure,
102 -- we fill in the environment field with @()@, which gives @Clo neg_v neg_l ()@.
103 --
104 -- Suppose we then compute @replicate n neg@. This results in an array of
105 -- closures. We only need one copy of the implementation functions neg_v and
106 -- neg_l, but the unit environment () is lifted to an array of units,
107 -- which we represent as PUnit.
108 --
109 -- Note that we need to store at least one real value, PUnit in this case,
110 -- because this value also represents the divergence behaviour of the whole
111 -- array. When evaluating a bulk-strict array, if any of the elements diverge
112 -- then the whole array does. We represent a diverging array of () by using
113 -- a diverging computation of type PUnit as its representation.
114 --
115 data instance PData ()
116 = PUnit
117
118 punit :: PData ()
119 punit = PUnit
120
121 $(unitPRInstance 'PUnit)
122
123
124 -- Wrap -----------------------------------------------------------------------
125 newtype instance PData (Wrap a)
126 = PWrap (PData a)
127
128 $(wrapPRInstance ''Wrap 'Wrap 'unWrap 'PWrap)
129
130 {- Generated code:
131 instance PA a => PR (Wrap a) where
132 ... INLINE pragmas ...
133 emptyPR = traceFn "emptyPR" "Wrap a" (PWrap emptyPD)
134
135 replicatePR n# (Wrap x)
136 = traceFn "replicatePR" "Wrap a" (PWrap (replicatePD n# x))
137
138 replicatelPR segd (PWrap xs)
139 = traceFn "replicatelPR" "Wrap a" (PWrap (replicatelPD segd xs))
140
141 repeatPR n# len# (PWrap xs)
142 = traceFn "repeatPR" "Wrap a" (PWrap (repeatPD n# len# xs))
143
144 indexPR (PWrap xs) i#
145 = traceFn "indexPR" "Wrap a" (Wrap (indexPD xs i#))
146
147 extractPR (PWrap xs) i# n#
148 = traceFn "extractPR" "Wrap a" (PWrap (extractPD xs i# n#))
149
150 bpermutePR (PWrap xs) n# is
151 = traceFn "bpermutePR" "Wrap a" (PWrap (bpermutePD xs n# is))
152
153 appPR (PWrap xs1) (PWrap xs2)
154 = traceFn "appPR" "Wrap a" (PWrap (appPD xs1 xs2))
155
156 applPR segd is (PWrap xs1) js (PWrap xs2)
157 = traceFn "applPR" "Wrap a" (PWrap (applPD segd is xs1 js xs2))
158
159 packByTagPR (PWrap xs) n# tags t#
160 = traceFn
161 "packByTagPR" "Wrap a" (PWrap (packByTagPD xs n# tags t#))
162
163 combine2PR n# sel (PWrap xs1) (PWrap xs2)
164 = traceFn "combine2PR" "Wrap a" (PWrap (combine2PD n# sel xs1 xs2))
165
166 updatePR (PWrap xs1) is (PWrap xs2)
167 = traceFn "updatePR" "Wrap a" (PWrap (updatePD xs1 is xs2))
168
169 fromListPR n# xs
170 = traceFn "fromListPR" "Wrap a" (PWrap (fromListPD n# (map unWrap xs)))
171
172 nfPR (PWrap xs)
173 = traceFn "nfPR" "Wrap a" (nfPD xs) }
174 -}
175
176
177 -- Tuples ---------------------------------------------------------------------
178
179 $(tupleInstances [2..15])
180
181 {- Generated code:
182
183 data instance PData (a,b)
184 = P_2 (PData a)
185 (PData b)
186
187 instance (PR a, PR b) => PR (a,b) where
188 {-# INLINE emptyPR #-}
189 emptyPR = P_2 emptyPR emptyPR
190
191 {-# INLINE replicatePR #-}
192 replicatePR n# (a,b) =
193 P_2 (replicatePR n# a)
194 (replicatePR n# b)
195
196 {-# INLINE replicatelPR #-}
197 replicatelPR segd (P_2 as bs) =
198 P_2 (replicatelPR segd as)
199 (replicatelPR segd bs)
200
201 {-# INLINE repeatPR #-}
202 repeatPR n# len# (P_2 as bs) =
203 P_2 (repeatPR n# len# as)
204 (repeatPR n# len# bs)
205
206 {-# INLINE indexPR #-}
207 indexPR (P_2 as bs) i# = (indexPR as i#, indexPR bs i#)
208
209 {-# INLINE extractPR #-}
210 extractPR (P_2 as bs) i# n# =
211 P_2 (extractPR as i# n#)
212 (extractPR bs i# n#)
213
214 {-# INLINE bpermutePR #-}
215 bpermutePR (P_2 as bs) n# is =
216 P_2 (bpermutePR as n# is)
217 (bpermutePR bs n# is)
218
219 {-# INLINE appPR #-}
220 appPR (P_2 as1 bs1) (P_2 as2 bs2) =
221 P_2 (appPR as1 as2) (appPR bs1 bs2)
222
223 {-# INLINE applPR #-}
224 applPR is (P_2 as1 bs1) js (P_2 as2 bs2) =
225 P_2 (applPR is as1 js as2)
226 (applPR is bs1 js bs2)
227
228 {-# INLINE packByTagPR #-}
229 packByTagPR (P_2 as bs) n# tags t# =
230 P_2 (packByTagPR as n# tags t#)
231 (packByTagPR bs n# tags t#)
232
233 {-# INLINE combine2PR #-}
234 combine2PR n# sel (P_2 as1 bs1) (P_2 as2 bs2) =
235 P_2 (combine2PR n# sel as1 as2)
236 (combine2PR n# sel bs1 bs2)
237
238 {-# INLINE updatePR #-}
239 updatePR (P_2 as1 bs1) is (P_2 as2 bs2) =
240 P_2 (updatePR as1 is as2)
241 (updatePR bs1 is bs2)
242
243 {-# INLINE fromListPR #-}
244 fromListPR n# xs = let (as,bs) = unzip xs in
245 P_2 (fromListPR n# as)
246 (fromListPR n# bs)
247
248 {-# INLINE nfPR #-}
249 nfPR (P_2 as bs) = nfPR as `seq` nfPR bs
250 -}
251
252 data instance PDatas (a, b)
253 = Ps_2 (PDatas a) (PDatas b)
254
255 -- Operators on arrays of tuples.
256 -- These are here instead of in "Data.Array.Parallel.PArray.Base" because
257 -- they need to know about the P_2 P_3 constructors. These are the representations
258 -- of tuple constructors that are generated by $(tupleInstances) above.
259 zipPA# :: PArray a -> PArray b -> PArray (a ,b)
260 {-# INLINE_PA zipPA# #-}
261 zipPA# (PArray n# xs) (PArray _ ys)
262 = PArray n# (P_2 xs ys)
263
264 unzipPA# :: PArray (a, b) -> (PArray a, PArray b)
265 {-# INLINE_PA unzipPA# #-}
266 unzipPA# (PArray n# (P_2 xs ys))
267 = (PArray n# xs, PArray n# ys)
268
269 zip3PA# :: PArray a -> PArray b -> PArray c -> PArray (a, b, c)
270 {-# INLINE_PA zip3PA# #-}
271 zip3PA# (PArray n# xs) (PArray _ ys) (PArray _ zs)
272 = PArray n# (P_3 xs ys zs)
273
274 unzip3PA# :: PArray (a, b, c) -> (PArray a, PArray b, PArray c)
275 {-# INLINE_PA unzip3PA# #-}
276 unzip3PA# (PArray n# (P_3 xs ys zs))
277 = (PArray n# xs, PArray n# ys, PArray n# zs)
278
279
280 zip4PA# :: PArray a -> PArray b -> PArray c -> PArray d -> PArray (a, b, c, d)
281 {-# INLINE_PA zip4PA# #-}
282 zip4PA# (PArray n# xs) (PArray _ ys) (PArray _ zs) (PArray _ as)
283 = PArray n# (P_4 xs ys zs as)
284
285 unzip4PA# :: PArray (a, b, c, d) -> (PArray a, PArray b, PArray c, PArray d)
286 {-# INLINE_PA unzip4PA# #-}
287 unzip4PA# (PArray n# (P_4 ws xs ys zs))
288 = (PArray n# ws, PArray n# xs, PArray n# ys, PArray n# zs)
289
290 zip5PA# :: PArray a -> PArray b -> PArray c -> PArray d -> PArray e -> PArray (a, b, c, d, e)
291 {-# INLINE_PA zip5PA# #-}
292 zip5PA# (PArray n# xs) (PArray _ ys) (PArray _ zs) (PArray _ as) (PArray _ bs)
293 = PArray n# (P_5 xs ys zs as bs)
294
295 unzip5PA# :: PArray (a, b, c, d, e) -> (PArray a, PArray b, PArray c, PArray d, PArray e)
296 {-# INLINE_PA unzip5PA# #-}
297 unzip5PA# (PArray n# (P_5 vs ws xs ys zs))
298 = (PArray n# vs, PArray n# ws, PArray n# xs, PArray n# ys, PArray n# zs)
299
300 zip6PA# :: PArray a -> PArray b -> PArray c -> PArray d -> PArray e -> PArray f -> PArray (a, b, c, d, e, f)
301 {-# INLINE_PA zip6PA# #-}
302 zip6PA# (PArray n# xs) (PArray _ ys) (PArray _ zs) (PArray _ as) (PArray _ bs) (PArray _ cs)
303 = PArray n# (P_6 xs ys zs as bs cs)
304
305 unzip6PA# :: PArray (a, b, c, d, e, f) -> (PArray a, PArray b, PArray c, PArray d, PArray e, PArray f)
306 {-# INLINE_PA unzip6PA# #-}
307 unzip6PA# (PArray n# (P_6 us vs ws xs ys zs))
308 = (PArray n# us, PArray n# vs, PArray n# ws, PArray n# xs, PArray n# ys, PArray n# zs)
309
310
311 zip7PA# :: PArray a -> PArray b -> PArray c -> PArray d -> PArray e -> PArray f -> PArray g ->
312 PArray (a, b, c, d, e, f, g)
313 {-# INLINE_PA zip7PA# #-}
314 zip7PA# (PArray n# xs) (PArray _ ys) (PArray _ zs) (PArray _ as) (PArray _ bs) (PArray _ cs) (PArray _ ds)
315 = PArray n# (P_7 xs ys zs as bs cs ds)
316
317 unzip7PA# :: PArray (a, b, c, d, e, f, g) -> (PArray a, PArray b, PArray c, PArray d, PArray e, PArray f, PArray g)
318 {-# INLINE_PA unzip7PA# #-}
319 unzip7PA# (PArray n# (P_7 ts us vs ws xs ys zs))
320 = (PArray n# ts, PArray n# us, PArray n# vs, PArray n# ws, PArray n# xs, PArray n# ys, PArray n# zs)
321
322 zip8PA# :: PArray a -> PArray b -> PArray c -> PArray d -> PArray e -> PArray f -> PArray g -> PArray h ->
323 PArray (a, b, c, d, e, f, g, h)
324 {-# INLINE_PA zip8PA# #-}
325 zip8PA# (PArray n# xs) (PArray _ ys) (PArray _ zs) (PArray _ as) (PArray _ bs) (PArray _ cs) (PArray _ ds) (PArray _ es)
326 = PArray n# (P_8 xs ys zs as bs cs ds es)
327
328 unzip8PA# :: PArray (a, b, c, d, e, f, g, h) -> (PArray a, PArray b, PArray c, PArray d, PArray e, PArray f, PArray g, PArray h)
329 {-# INLINE_PA unzip8PA# #-}
330 unzip8PA# (PArray n# (P_8 ss ts us vs ws xs ys zs))
331 = (PArray n# ss, PArray n# ts, PArray n# us, PArray n# vs, PArray n# ws, PArray n# xs, PArray n# ys, PArray n# zs)
332
333
334 -- Sums -----------------------------------------------------------------------
335 data instance PData (Sum2 a b)
336 = PSum2 U.Sel2 (PData a) (PData b)
337
338 data instance PDatas (Sum2 a b)
339 = PSums2 Sels2 (PDatas a) (PDatas b)
340
341 type Sels2 = ()
342
343 instance (PR a, PR b) => PR (Sum2 a b) where
344 {-# INLINE emptyPR #-}
345 emptyPR
346 = traceFn "emptyPR" "(Sum2 a b)" $
347 PSum2 (U.mkSel2 U.empty U.empty 0 0 (U.mkSelRep2 U.empty)) emptyPR emptyPR
348
349 {-# INLINE replicatePR #-}
350 replicatePR n# (Alt2_1 x)
351 = traceFn "replicatePR" "(Sum2 a b)" $
352 PSum2 (U.mkSel2 (U.replicate (I# n#) 0)
353 (U.enumFromStepLen 0 1 (I# n#))
354 (I# n#) 0
355 (U.mkSelRep2 (U.replicate (I# n#) 0)))
356 (replicatePR n# x)
357 emptyPR
358 replicatePR n# (Alt2_2 x)
359 = traceFn "replicatePR" "(Sum2 a b)" $
360 PSum2 (U.mkSel2 (U.replicate (I# n#) 1)
361 (U.enumFromStepLen 0 1 (I# n#))
362 0 (I# n#)
363 (U.mkSelRep2 (U.replicate (I# n#) 1)))
364 emptyPR
365 (replicatePR n# x)
366
367 {-# INLINE replicatelPR #-}
368 replicatelPR segd (PSum2 sel as bs)
369 = traceFn "replicatelPR" "(Sum2 a b)" $
370 PSum2 sel' as' bs'
371 where
372 tags = U.tagsSel2 sel
373 tags' = U.replicate_s segd tags
374 sel' = U.tagsToSel2 tags'
375
376 lens = U.lengthsSegd segd
377
378 asegd = U.lengthsToSegd (U.packByTag lens tags 0)
379 bsegd = U.lengthsToSegd (U.packByTag lens tags 1)
380
381 as' = replicatelPR asegd as
382 bs' = replicatelPR bsegd bs
383
384 {-# INLINE repeatPR #-}
385 repeatPR m# n# (PSum2 sel as bs)
386 = traceFn "repeatPR" "(Sum2 a b)" $
387 PSum2 sel' as' bs'
388 where
389 sel' = U.tagsToSel2
390 . U.repeat (I# m#) (I# n#)
391 $ U.tagsSel2 sel
392
393 as' = repeatPR m# (elementsSel2_0# sel) as
394 bs' = repeatPR n# (elementsSel2_1# sel) bs
395
396 {-# INLINE indexPR #-}
397 indexPR (PSum2 sel as bs) i#
398 = traceFn "indexPR" "(Sum2 a b)" $
399 case U.index "indexPR[Sum2]" (U.indicesSel2 sel) (I# i#) of
400 I# k# -> case U.index "indexPR[Sum2]" (U.tagsSel2 sel) (I# i#) of
401 0 -> Alt2_1 (indexPR as k#)
402 _ -> Alt2_2 (indexPR bs k#)
403
404 {-# INLINE appPR #-}
405 appPR (PSum2 sel1 as1 bs1)
406 (PSum2 sel2 as2 bs2)
407 = traceFn "appPR" "(Sum2 a b)" $
408 PSum2 sel (appPR as1 as2)
409 (appPR bs1 bs2)
410 where
411 sel = U.tagsToSel2
412 $ U.tagsSel2 sel1 U.+:+ U.tagsSel2 sel2
413
414 {-# INLINE packByTagPR #-}
415 packByTagPR (PSum2 sel as bs) _ tags t#
416 = PSum2 sel' as' bs'
417 where
418 my_tags = U.tagsSel2 sel
419 my_tags' = U.packByTag my_tags tags (intToTag (I# t#))
420 sel' = U.tagsToSel2 my_tags'
421
422 atags = U.packByTag tags my_tags 0
423 btags = U.packByTag tags my_tags 1
424
425 as' = packByTagPR as (elementsSel2_0# sel') atags t#
426 bs' = packByTagPR bs (elementsSel2_1# sel') btags t#
427
428 {-# INLINE combine2PR #-}
429 combine2PR _ sel (PSum2 sel1 as1 bs1)
430 (PSum2 sel2 as2 bs2)
431 = traceFn "combine2PR" "(Sum2 a b)" $
432 PSum2 sel' as bs
433 where
434 tags = U.tagsSel2 sel
435 tags' = U.combine2 (U.tagsSel2 sel) (U.repSel2 sel)
436 (U.tagsSel2 sel1)
437 (U.tagsSel2 sel2)
438 sel' = U.tagsToSel2 tags'
439
440 asel = U.tagsToSel2 (U.packByTag tags tags' 0)
441 bsel = U.tagsToSel2 (U.packByTag tags tags' 1)
442
443 as = combine2PR (elementsSel2_0# sel') asel as1 as2
444 bs = combine2PR (elementsSel2_1# sel') bsel bs1 bs2
445
446
447 -- Nested Arrays --------------------------------------------------------------
448 data instance PData (PArray a)
449 = PNested U.Segd (PData a)
450
451 instance PR a => PR (PArray a) where
452 {-# INLINE emptyPR #-}
453 emptyPR = traceFn "emptyPR" "(PArray a)" $
454 PNested (U.mkSegd U.empty U.empty 0) emptyPR
455
456 {-# INLINE replicatePR #-}
457 replicatePR n# (PArray m# xs)
458 = traceFn "replicatePR" "(PArray a)" $
459 PNested (U.mkSegd (U.replicate (I# n#) (I# m#))
460 (U.enumFromStepLen 0 (I# m#) (I# n#))
461 (I# n# * I# m#))
462 (repeatPR n# m# xs)
463
464 {-# INLINE indexPR #-}
465 indexPR (PNested segd xs) i#
466 = traceFn "indexPR" "(PArray a)" $
467 case U.index "indexPR[Nested]" (U.lengthsSegd segd) (I# i#) of { I# n# ->
468 case U.index "indexPR[Nested]" (U.indicesSegd segd) (I# i#) of { I# k# ->
469 PArray n# (extractPR xs k# n#) }}
470
471 {-# INLINE extractPR #-}
472 extractPR (PNested segd xs) i# n#
473 = traceFn "extractPR" "(PArray a)" $
474 PNested segd' (extractPR xs k# (elementsSegd# segd'))
475 where
476 segd' = U.lengthsToSegd
477 $ U.extract (U.lengthsSegd segd) (I# i#) (I# n#)
478
479 -- NB: not indicesSegd segd !: i because i might be one past the end
480 !(I# k#) | I# i# == 0 = 0
481 | otherwise = U.index "extractPR[Nested]" (U.indicesSegd segd) (I# i# - 1)
482 + U.index "extractPR[Nested]" (U.lengthsSegd segd) (I# i# - 1)
483
484 {-# INLINE bpermutePR #-}
485 bpermutePR (PNested segd xs) _ is
486 = traceFn "bpermutePR" "(PArray a)" $
487 PNested segd' (bpermutePR xs (elementsSegd# segd') js)
488 where
489 lens' = U.bpermute (U.lengthsSegd segd) is
490 starts = U.bpermute (U.indicesSegd segd) is
491
492 segd' = U.lengthsToSegd lens'
493
494 js = U.zipWith (+) (U.indices_s segd')
495 (U.replicate_s segd' starts)
496
497 {-# INLINE appPR #-}
498 appPR (PNested xsegd xs) (PNested ysegd ys)
499 = traceFn "appPR" "(PArray a)" $
500 PNested (U.lengthsToSegd (U.lengthsSegd xsegd U.+:+ U.lengthsSegd ysegd))
501 (appPR xs ys)
502
503 {-# INLINE applPR #-}
504 applPR rsegd segd1 (PNested xsegd xs) segd2 (PNested ysegd ys)
505 = traceFn "applPR" "(PArray a)"$
506 PNested segd (applPR (U.plusSegd xsegd' ysegd') xsegd' xs ysegd' ys)
507 where
508 segd = U.lengthsToSegd
509 $ U.append_s rsegd segd1 (U.lengthsSegd xsegd)
510 segd2 (U.lengthsSegd ysegd)
511
512 xsegd' = U.lengthsToSegd
513 $ U.sum_s segd1 (U.lengthsSegd xsegd)
514 ysegd' = U.lengthsToSegd
515 $ U.sum_s segd2 (U.lengthsSegd ysegd)
516
517 {-# INLINE repeatPR #-}
518 repeatPR n# len# (PNested segd xs)
519 = traceFn "repeatPR" "(PArray a)" $
520 PNested segd' (repeatPR n# (elementsSegd# segd) xs)
521 where
522 segd' = U.lengthsToSegd (U.repeat (I# n#) (I# len#) (U.lengthsSegd segd))
523
524 {-# INLINE replicatelPR #-}
525 replicatelPR segd (PNested xsegd xs)
526 = traceFn "replicatelPR" "(PArray a)" $
527 PNested xsegd' $ bpermutePR xs (elementsSegd# xsegd')
528 $ U.enumFromStepLenEach (U.elementsSegd xsegd')
529 is (U.replicate (U.elementsSegd segd) 1) ns
530 where
531 is = U.replicate_s segd (U.indicesSegd xsegd)
532 ns = U.replicate_s segd (U.lengthsSegd xsegd)
533 xsegd' = U.lengthsToSegd ns
534
535 {-# INLINE packByTagPR #-}
536 packByTagPR (PNested segd xs) _ tags t#
537 = traceFn "packByTagPR" "(PArray a)" $
538 PNested segd' xs'
539 where
540 segd' = U.lengthsToSegd
541 $ U.packByTag (U.lengthsSegd segd) tags (intToTag (I# t#))
542
543 xs' = packByTagPR xs (elementsSegd# segd') (U.replicate_s segd tags) t#
544
545 {-# INLINE combine2PR #-}
546 combine2PR _ sel (PNested xsegd xs) (PNested ysegd ys)
547 = traceFn "combine2PR" "(PArray a)" $
548 PNested segd xys
549 where
550 tags = U.tagsSel2 sel
551
552 segd = U.lengthsToSegd
553 $ U.combine2 (U.tagsSel2 sel) (U.repSel2 sel)
554 (U.lengthsSegd xsegd)
555 (U.lengthsSegd ysegd)
556
557 sel' = U.tagsToSel2
558 $ U.replicate_s segd tags
559
560 xys = combine2PR (elementsSegd# segd) sel' xs ys
561
562
563 -- Operators on Nested Arrays
564 -- These are here instead of in "Data.Array.Parallel.PArray.Base" because
565 -- they need to know about the PNested constructor which is defined above.
566
567 -- | O(1). Extract the segment descriptor from a nested array.
568 segdPA# :: PArray (PArray a) -> U.Segd
569 {-# INLINE_PA segdPA# #-}
570 segdPA# (PArray _ (PNested segd _))
571 = segd
572
573
574 -- | O(1). Concatenate a nested array. This is a constant time operation as
575 -- we can just discard the segment descriptor.
576 concatPA# :: PArray (PArray a) -> PArray a
577 {-# INLINE_PA concatPA# #-}
578 concatPA# (PArray _ (PNested segd xs))
579 = PArray (elementsSegd# segd) xs
580
581
582 -- | O(1). Create a nested array from an element count, segment descriptor,
583 -- and data elements.
584 segmentPA# :: Int# -> U.Segd -> PArray a -> PArray (PArray a)
585 {-# INLINE_PA segmentPA# #-}
586 segmentPA# n# segd (PArray _ xs)
587 = PArray n# (PNested segd xs)
588
589
590 -- | O(1). Create a nested array by using the element count and segment
591 -- descriptor from another, but use new data elements.
592 copySegdPA# :: PArray (PArray a) -> PArray b -> PArray (PArray b)
593 {-# INLINE copySegdPA# #-}
594 copySegdPA# (PArray n# (PNested segd _)) (PArray _ xs)
595 = PArray n# (PNested segd xs)