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