6fbfc4e8865e5e78d1bdd4e424361a1435dd5213
[packages/dph.git] / dph-prim-par / Data / Array / Parallel / Unlifted / Parallel / UPSel.hs
1 {-# LANGUAGE CPP #-}
2 #include "fusion-phases.h"
3
4 -- | Parallel selectors.
5 module Data.Array.Parallel.Unlifted.Parallel.UPSel
6 ( -- * Types
7 UPSel2
8 , UPSelRep2
9
10 -- * Operations
11 , tagsUPSel2
12 , indicesUPSel2
13 , elementsUPSel2_0
14 , elementsUPSel2_1
15 , selUPSel2
16 , repUPSel2
17 , mkUPSel2
18 , mkUPSelRep2
19 , indicesUPSelRep2
20 , elementsUPSelRep2_0
21 , elementsUPSelRep2_1)
22 where
23 import Data.Array.Parallel.Unlifted.Sequential.Vector as US
24 import Data.Array.Parallel.Unlifted.Sequential.USel
25 import Data.Array.Parallel.Unlifted.Distributed
26 import Data.Array.Parallel.Base (Tag, tagToInt)
27
28
29 -- | Contains a selector `USel2`, as well as an `USelRep2` which says how
30 -- to distribute this selector across the PEs.
31 --
32 -- See @dph-prim-seq:Data.Array.Parallel.Unlifted.Sequential.Segmented.USel@
33 -- for more discussion of what selectors are for.
34 --
35 data UPSel2
36 = UPSel2
37 { upsel2_usel :: USel2
38 , upsel2_rep :: UPSelRep2 }
39
40
41 --
42 type UPSelRep2
43 = Dist ((Int,Int), (Int,Int))
44
45 -- Projections ----------------------------------------------------------------
46 -- INLINE trivial projections as they'll expand to a single record selector.
47
48 -- | O(1). Get the tags of a selector.
49 tagsUPSel2 :: UPSel2 -> Vector Tag
50 tagsUPSel2 = tagsUSel2 . upsel2_usel
51 {-# INLINE tagsUPSel2 #-}
52
53
54 -- | O(1). Get the indices of a selector.
55 indicesUPSel2 :: UPSel2 -> Vector Int
56 indicesUPSel2 = indicesUSel2 . upsel2_usel
57 {-# INLINE indicesUPSel2 #-}
58
59
60 -- | O(1). Get the number of elements that will be taken from the first array.
61 elementsUPSel2_0 :: UPSel2 -> Int
62 elementsUPSel2_0 = elementsUSel2_0 . upsel2_usel
63 {-# INLINE elementsUPSel2_0 #-}
64
65
66 -- | O(1). Get the number of elements that will be taken from the second array.
67 elementsUPSel2_1 :: UPSel2 -> Int
68 elementsUPSel2_1 = elementsUSel2_1 . upsel2_usel
69 {-# INLINE elementsUPSel2_1 #-}
70
71
72 -- | O(1). Take the sequential `USel2` from a `UPSel2`.
73 selUPSel2 :: UPSel2 -> USel2
74 selUPSel2 = upsel2_usel
75 {-# INLINE selUPSel2 #-}
76
77
78 -- | O(1). Take the `UPSelRep2` from a `UPSel2`.
79 repUPSel2 :: UPSel2 -> UPSelRep2
80 repUPSel2 = upsel2_rep
81 {-# INLINE repUPSel2 #-}
82
83
84 -- Representation selectors ---------------------------------------------------
85 -- | Computes a `UPSelRep2` from an array of tags.
86 -- This is used when parallelising a `combine` operation.
87 -- See the docs for `UPSelRep2` for details.
88 mkUPSelRep2 :: Vector Tag -> UPSelRep2
89 mkUPSelRep2 tags = zipD idxs lens
90 where
91 lens = mapD theGang count
92 $ splitD theGang balanced tags
93
94 idxs = fst
95 $ scanD theGang add (0,0) lens
96
97 count bs = let ones = US.sum (US.map tagToInt bs)
98 in (US.length bs - ones,ones)
99
100 add (x1,y1) (x2,y2) = (x1+x2, y1+y2)
101 {-# INLINE_UP mkUPSelRep2 #-}
102
103
104 indicesUPSelRep2 :: Vector Tag -> UPSelRep2 -> Vector Int
105 indicesUPSelRep2 tags rep
106 = joinD theGang balanced
107 $ zipWithD theGang indices
108 (splitD theGang balanced tags)
109 rep
110 where
111 indices tags' ((i,j), (m,n))
112 = US.combine2ByTag tags' (US.enumFromStepLen i 1 m)
113 (US.enumFromStepLen j 1 n)
114 {-# INLINE_UP indicesUPSelRep2 #-}
115
116
117 -- | O(n). Count the number of elements to take from the first array.
118 elementsUPSelRep2_0 :: Vector Tag -> UPSelRep2 -> Int
119 elementsUPSelRep2_0 _
120 = sumD theGang . fstD . sndD
121 {-# INLINE_UP elementsUPSelRep2_0 #-}
122
123
124 -- | O(n). Count the number of elements to take from the second array.
125 elementsUPSelRep2_1 :: Vector Tag -> UPSelRep2 -> Int
126 elementsUPSelRep2_1 _
127 = sumD theGang . sndD . sndD
128 {-# INLINE_UP elementsUPSelRep2_1 #-}
129
130
131 -- | O(1). Construct a selector. Wrapper for `UPSel2`.
132 mkUPSel2 :: Vector Tag -> Vector Int -> Int -> Int -> UPSelRep2 -> UPSel2
133 mkUPSel2 tags is n0 n1 rep
134 = UPSel2 (mkUSel2 tags is n0 n1) rep
135 {-# INLINE_UP mkUPSel2 #-}