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