Reorganise the way the lifted reference implementation works
[packages/dph.git] / dph-lifted-boxed / Data / Array / Parallel / Lifted / Combinators.hs
1 {-# OPTIONS -fno-spec-constr #-}
2 #include "fusion-phases.h"
3
4 -- | Closure converted lifted array combinators.
5 -- The vectoriser produces code that uses these combinators directly.
6 --
7 -- All of the combinators in this module are polymorphic, work on `PArray`, and
8 -- take `PA` dictionaries. Combinators that are specific to a certain element type,
9 -- like `Int`, are defined in the corresponding prelude module,
10 -- eg "Data.Array.Parallel.Prelude.Int".
11 --
12 module Data.Array.Parallel.Lifted.Combinators
13 ( -- * Conversions
14 fromPArrayPP
15 , toPArrayPP
16 , fromNestedPArrayPP
17
18 -- * Constructors
19 , emptyPP
20 , singletonPP
21 , replicatePP
22 , appendPP
23
24 -- * Projections
25 , lengthPP
26 , indexPP
27 , slicePP
28
29 -- * Traversals
30 , mapPP
31 , zipWithPP
32 , crossMapPP
33
34 -- * Filtering
35 , filterPP
36
37 -- * Concatenation
38 , concatPP
39
40 -- * Tuple functions
41 , zipPP
42 , unzipPP)
43 where
44 import Data.Array.Parallel.Lifted.Closure
45 import Data.Array.Parallel.PArray.PData as PA
46 import Data.Array.Parallel.PArray.PRepr as PA
47 import Data.Array.Parallel.PArray as PA
48
49
50 -- Conversions ================================================================
51 -- The following identity functions are used as the vectorised versions of the
52 -- functions that convert between the source level array type [:a:] and the
53 -- PArray type which is used in the library.
54
55 -- | Identity function, used as the vectorised version of fromPArrayP.
56 fromPArrayPP :: PA a => PArray a :-> PArray a
57 fromPArrayPP = closure1 (\x -> x) (\_ xs -> xs)
58 {-# INLINE fromPArrayPP #-}
59
60
61 -- | Identity function, used as the vectorised version of toPArrayP.
62 toPArrayPP :: PA a => PArray a :-> PArray a
63 toPArrayPP = closure1 (\x -> x) (\_ xs -> xs)
64 {-# INLINE toPArrayPP #-}
65
66
67 -- | Identity function, used as the vectorised version of fromNestedPArrayP
68 fromNestedPArrayPP :: PA a => (PArray (PArray a) :-> PArray (PArray a))
69 fromNestedPArrayPP = closure1 (\xs -> xs) (\_ xss -> xss)
70 {-# INLINE fromNestedPArrayPP #-}
71
72
73 -- Combinators ================================================================
74 -- For each combinator:
75 -- The *PP_v version is the "vectorised" version that has had its parameters
76 -- closure converted. For first-order functions, the *PP_v version is
77 -- identical to the standard *PA version from D.A.P.PArray, so we can
78 -- just use that directly.
79 --
80 -- The *PP_l version is the "lifted" version that works on arrays of arrays.
81 -- Each of these functions also takes an integer as its first argument.
82 -- This is the "lifting context" that says now many element to expect in
83 -- each of the argument arrays.
84 --
85 -- The *PP version contains both the vectorised and lifted versions wrapped
86 -- up in a closure. The code produced by the vectoriser uses the *PP
87 -- versions directly.
88
89
90 -- Constructors ---------------------------------------------------------------
91 -- | O(1). Construct an empty array.
92 emptyPP :: PA a => PArray a
93 emptyPP = PA.empty
94 {-# INLINE_PA emptyPP #-}
95
96
97 -- | O(1). Construct an array containing a single element.
98 singletonPP :: PA a => a :-> PArray a
99 singletonPP = closure1' PA.singleton PA.singletonl
100 {-# INLINE_PA singletonPP #-}
101
102
103 -- | O(n). Construct an array of the given size, that maps all elements to the same value.
104 replicatePP :: PA a => Int :-> a :-> PArray a
105 replicatePP = closure2' PA.replicate PA.replicatel
106 {-# INLINE_PA replicatePP #-}
107
108
109 -- | O(len result). Append two arrays.
110 appendPP :: PA a => PArray a :-> PArray a :-> PArray a
111 appendPP = closure2' PA.append PA.appendl
112 {-# INLINE_PA appendPP #-}
113
114
115 -- | O(len result). Concatenate a nested array.
116 concatPP :: PA a => PArray (PArray a) :-> PArray a
117 concatPP = closure1' PA.concat PA.concatl
118 {-# INLINE_PA concatPP #-}
119
120
121 -- Projections ----------------------------------------------------------------
122 -- | O(1). Take the number of elements in an array.
123 lengthPP :: PA a => PArray a :-> Int
124 lengthPP = closure1' PA.length PA.lengthl
125 {-# INLINE_PA lengthPP #-}
126
127
128 -- | O(1). Lookup a single element from the source array.
129 indexPP :: PA a => PArray a :-> Int :-> a
130 indexPP = closure2' PA.index PA.indexl
131 {-# INLINE_PA indexPP #-}
132
133
134 -- | O(len slice). Extract a range of elements from an array.
135 slicePP :: PA a => Int :-> Int :-> PArray a :-> PArray a
136 slicePP = closure3' PA.slice PA.slicel
137 {-# INLINE_PA slicePP #-}
138
139
140 -- Traversals -----------------------------------------------------------------
141 -- | Apply a worker function to every element of an array.
142 mapPP :: (PA a, PA b)
143 => (a :-> b) :-> PArray a :-> PArray b
144
145 mapPP = closure2' mapPP_v mapPP_l
146 {-# INLINE_PA mapPP #-}
147
148
149 mapPP_v :: (PA a, PA b)
150 => (a :-> b) -> PArray a -> PArray b
151 mapPP_v f as
152 = PA.replicate (PA.length as) f $:^ as
153 {-# INLINE mapPP_v #-}
154
155
156 mapPP_l :: (PA a, PA b)
157 => (PArray (a :-> b)) -> PArray (PArray a) -> PArray (PArray b)
158 mapPP_l fs ass
159 = PA.unconcat ass
160 $ PA.replicates (PA.takeUSegd ass) fs
161 $:^ PA.concat ass
162 {-# INLINE mapPP_l #-}
163
164
165 -- | Apply a worker function to every pair of two arrays.
166 zipWithPP
167 :: (PA a, PA b, PA c)
168 => (a :-> b :-> c) :-> PArray a :-> PArray b :-> PArray c
169
170 zipWithPP = closure3' zipWithPP_v zipWithPP_l
171 where
172 {-# INLINE zipWithPP_v #-}
173 zipWithPP_v f as bs
174 = PA.replicate (PA.length as) f $:^ as $:^ bs
175
176 {-# INLINE zipWithPP_l #-}
177 zipWithPP_l fs ass bss
178 = PA.unconcat ass
179 $ PA.replicates (PA.takeUSegd ass) fs
180 $:^ PA.concat ass
181 $:^ PA.concat bss
182 {-# INLINE_PA zipWithPP #-}
183
184
185 -- |
186 crossMapPP
187 :: (PA a, PA b)
188 => PArray a :-> (a :-> PArray b) :-> PArray (a, b)
189
190 crossMapPP = closure2' crossMapPP_v crossMapPP_l
191 where
192 {-# INLINE crossMapPP_v #-}
193 crossMapPP_v _ _
194 = error "crossMapP: not implemented"
195
196 {-# INLINE crossMapPP_l #-}
197 crossMapPP_l _ _
198 = error "crossMapP: not implemented"
199
200 {-# INLINE_PA crossMapPP #-}
201
202 -- Filtering ------------------------------------------------------------------
203 -- | Extract the elements from an array that match the given predicate.
204 filterPP :: PA a => (a :-> Bool) :-> PArray a :-> PArray a
205 {-# INLINE filterPP #-}
206 filterPP = closure2' filterPP_v filterPP_l
207 where
208 {-# INLINE filterPP_v #-}
209 filterPP_v p xs = PA.pack xs (mapPP_v p xs)
210
211 {-# INLINE filterPP_l #-}
212 filterPP_l ps xss = PA.packl xss (mapPP_l ps xss)
213
214
215 -- Tuple Functions ------------------------------------------------------------
216 -- | Zip a pair of arrays into an array of pairs.
217 zipPP :: (PA a, PA b) => PArray a :-> PArray b :-> PArray (a, b)
218 zipPP = closure2' PA.zip PA.zipl
219 {-# INLINE_PA zipPP #-}
220
221
222 -- | Unzip an array of pairs into a pair of arrays.
223 unzipPP :: (PA a, PA b) => PArray (a, b) :-> (PArray a, PArray b)
224 unzipPP = closure1' PA.unzip PA.unzipl
225 {-# INLINE_PA unzipPP #-}
226