5f9a7311288ef563457f80a2712bf901e1d2f0fb
[packages/dph.git] / dph-lifted-vseg / 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
33 -- * Filtering
34 , filterPP
35
36 -- * Concatenation
37 , concatPP
38
39 -- * Tuple functions
40 , zipPP
41 , unzipPP)
42 where
43 import Data.Array.Parallel.Lifted.Closure
44 import Data.Array.Parallel.PArray.PData as PA
45 import Data.Array.Parallel.PArray.PRepr as PA
46 import Data.Array.Parallel.PArray as PA
47
48
49 -- Conversions ================================================================
50 -- The following identity functions are used as the vectorised versions of the
51 -- functions that convert between the source level array type [:a:] and the
52 -- PArray type which is used in the library.
53
54 -- | Identity function, used as the vectorised version of fromPArrayP.
55 fromPArrayPP :: PA a => PArray a :-> PArray a
56 fromPArrayPP = closure1 (\x -> x) (\_ xs -> xs)
57 {-# INLINE fromPArrayPP #-}
58
59
60 -- | Identity function, used as the vectorised version of toPArrayP.
61 toPArrayPP :: PA a => PArray a :-> PArray a
62 toPArrayPP = closure1 (\x -> x) (\_ xs -> xs)
63 {-# INLINE toPArrayPP #-}
64
65
66 -- | Identity function, used as the vectorised version of fromNestedPArrayP
67 fromNestedPArrayPP :: PA a => (PArray (PArray a) :-> PArray (PArray a))
68 fromNestedPArrayPP = closure1 (\xs -> xs) (\_ xss -> xss)
69 {-# INLINE fromNestedPArrayPP #-}
70
71
72 -- Combinators ================================================================
73 -- For each combinator:
74 -- The *PP_v version is the "vectorised" version that has had its parameters
75 -- closure converted. For first-order functions, the *PP_v version is
76 -- identical to the standard *PA version from D.A.P.PArray, so we can
77 -- just use that directly.
78 --
79 -- The *PP_l version is the "lifted" version that works on arrays of arrays.
80 -- Each of these functions also takes an integer as its first argument.
81 -- This is the "lifting context" that says now many element to expect in
82 -- each of the argument arrays.
83 --
84 -- The *PP version contains both the vectorised and lifted versions wrapped
85 -- up in a closure. The code produced by the vectoriser uses the *PP
86 -- versions directly.
87
88
89 -- Constructors ---------------------------------------------------------------
90 -- | O(1). Construct an empty array.
91 emptyPP :: PA a => PArray a
92 emptyPP = PA.empty
93 {-# INLINE_PA emptyPP #-}
94
95
96 -- | O(1). Construct an array containing a single element.
97 singletonPP :: PA a => a :-> PArray a
98 singletonPP = closure1' PA.singleton PA.singletonl
99 {-# INLINE_PA singletonPP #-}
100
101
102 -- | O(n). Construct an array of the given size, that maps all elements to the same value.
103 replicatePP :: PA a => Int :-> a :-> PArray a
104 replicatePP = closure2' PA.replicate PA.replicatel
105 {-# INLINE_PA replicatePP #-}
106
107
108 -- | O(len result). Append two arrays.
109 appendPP :: PA a => PArray a :-> PArray a :-> PArray a
110 appendPP = closure2' PA.append PA.appendl
111 {-# INLINE_PA appendPP #-}
112
113
114 -- | O(len result). Concatenate a nested array.
115 concatPP :: PA a => PArray (PArray a) :-> PArray a
116 concatPP = closure1' PA.concat PA.concatl
117 {-# INLINE_PA concatPP #-}
118
119
120 -- Projections ----------------------------------------------------------------
121 -- | O(1). Take the number of elements in an array.
122 lengthPP :: PA a => PArray a :-> Int
123 lengthPP = closure1' PA.length PA.lengthl
124 {-# INLINE_PA lengthPP #-}
125
126
127 -- | O(1). Lookup a single element from the source array.
128 indexPP :: PA a => PArray a :-> Int :-> a
129 indexPP = closure2' PA.index PA.indexl
130 {-# INLINE_PA indexPP #-}
131
132
133 -- | O(len slice). Extract a range of elements from an array.
134 slicePP :: PA a => Int :-> Int :-> PArray a :-> PArray a
135 slicePP = closure3' PA.slice PA.slicel
136 {-# INLINE_PA slicePP #-}
137
138
139 -- Traversals -----------------------------------------------------------------
140 -- | Apply a worker function to every element of an array.
141 mapPP :: (PA a, PA b)
142 => (a :-> b) :-> PArray a :-> PArray b
143
144 mapPP = closure2' mapPP_v mapPP_l
145 {-# INLINE_PA mapPP #-}
146
147
148 mapPP_v :: (PA a, PA b)
149 => (a :-> b) -> PArray a -> PArray b
150 mapPP_v f as
151 = PA.replicate (PA.length as) f $:^ as
152 {-# INLINE mapPP_v #-}
153
154
155 mapPP_l :: (PA a, PA b)
156 => (PArray (a :-> b)) -> PArray (PArray a) -> PArray (PArray b)
157 mapPP_l fs ass
158 = PA.unconcat ass
159 $ PA.replicates (PA.unsafeTakeSegd ass) fs
160 $:^ PA.concat ass
161 {-# INLINE mapPP_l #-}
162
163
164 -- | Apply a worker function to every pair of two arrays.
165 zipWithPP
166 :: (PA a, PA b, PA c)
167 => (a :-> b :-> c) :-> PArray a :-> PArray b :-> PArray c
168
169 {-# INLINE_PA zipWithPP #-}
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.unsafeTakeSegd ass) fs
180 $:^ PA.concat ass
181 $:^ PA.concat bss
182
183
184 -- Filtering ------------------------------------------------------------------
185 -- | Extract the elements from an array that match the given predicate.
186 filterPP :: PA a => (a :-> Bool) :-> PArray a :-> PArray a
187 {-# INLINE filterPP #-}
188 filterPP = closure2' filterPP_v filterPP_l
189 where
190 {-# INLINE filterPP_v #-}
191 filterPP_v p xs = PA.pack xs (mapPP_v p xs)
192
193 {-# INLINE filterPP_l #-}
194 filterPP_l ps xss = PA.packl xss (mapPP_l ps xss)
195
196
197 -- Tuple Functions ------------------------------------------------------------
198 -- | Zip a pair of arrays into an array of pairs.
199 zipPP :: (PA a, PA b) => PArray a :-> PArray b :-> PArray (a, b)
200 zipPP = closure2' PA.zip PA.zipl
201 {-# INLINE_PA zipPP #-}
202
203
204 -- | Unzip an array of pairs into a pair of arrays.
205 unzipPP :: (PA a, PA b) => PArray (a, b) :-> (PArray a, PArray b)
206 unzipPP = closure1' PA.unzip PA.unzipl
207 {-# INLINE_PA unzipPP #-}
208