Makefiles need real tab chars, ffs.
[packages/dph.git] / dph-lifted-copy / Data / Array / Parallel / Lifted / Closure.hs
1 {-# OPTIONS -fno-warn-missing-methods #-}
2 module Data.Array.Parallel.Lifted.Closure (
3 (:->)(..), PArray(..),
4 mkClosure, mkClosureP, ($:), ($:^),
5 closure, liftedClosure, liftedApply,
6
7 closure1, closure2, closure3, closure4
8 ) where
9 import Data.Array.Parallel.PArray.PReprInstances ()
10 import Data.Array.Parallel.PArray.PDataInstances
11 import Data.Array.Parallel.Lifted.PArray
12
13 import GHC.Exts (Int#)
14
15 infixr 0 :->
16 infixl 0 $:, $:^
17
18 -- | The type of closures.
19 -- This bundles up:
20 -- 1) the vectorised version of the function that takes an explicit environment
21 -- 2) the lifted version, that works on arrays.
22 -- the first parameter of this function is the 'lifting context'
23 -- that gives the length of the array.
24 -- 3) the environment of the closure.
25 --
26 -- The vectoriser closure-converts the source program so that all functions
27 -- types are expressed in this form.
28 --
29 data a :-> b
30 = forall e. PA e
31 => Clo !(e -> a -> b) -- vectorised function
32 !(Int# -> PData e -> PData a -> PData b) -- lifted function
33 e -- environment
34
35
36 -- | Apply a lifted function by wrapping up the provided array data
37 -- into some real `PArray`s, and passing it those.
38 lifted :: (PArray e -> PArray a -> PArray b) -- ^ lifted function to call.
39 -> Int# -- ^ lifting context
40 -> PData e -- ^ environments
41 -> PData a -- ^ arguments
42 -> PData b -- ^ returned elements
43 {-# INLINE lifted #-}
44 lifted f n# es as
45 = case f (PArray n# es) (PArray n# as) of
46 PArray _ bs -> bs
47
48
49 -- | Construct a closure.
50 mkClosure
51 :: forall a b e
52 . PA e
53 => (e -> a -> b) -- ^ vectorised function, with explicit environment.
54 -> (PArray e -> PArray a -> PArray b) -- ^ lifted function, taking an array of environments.
55 -> e -- ^ environment
56 -> (a :-> b)
57 {-# INLINE CONLIKE mkClosure #-}
58 mkClosure fv fl e
59 = Clo fv (lifted fl) e
60
61
62 -- | Construct a closure.
63 -- This is like the `mkClosure` function above, except that the provided
64 -- lifted version of the function can take raw array data, instead of
65 -- data wrapped up into a `PArray`.
66 closure :: forall a b e
67 . PA e
68 => (e -> a -> b) -- ^ vectorised function, with explicit environment.
69 -> (Int# -> PData e -> PData a -> PData b) -- ^ lifted function, taking an array of environments.
70 -> e -- ^ environment
71 -> (a :-> b)
72 {-# INLINE closure #-}
73 closure fv fl e = Clo fv fl e
74
75
76 -- | Apply a closure to its argument.
77 --
78 ($:) :: forall a b. (a :-> b) -> a -> b
79 {-# INLINE ($:) #-}
80 Clo f _ e $: a = f e a
81
82 {-# RULES
83
84 "mkClosure/($:)" forall fv fl e x.
85 mkClosure fv fl e $: x = fv e x
86
87 #-}
88
89
90 -- | Arrays of closures (aka array closures)
91 -- We need to represent arrays of closures when vectorising partial applications.
92 --
93 -- For example, consider:
94 -- @mapP (+) xs :: [: Int -> Int :]@
95 --
96 -- Representing this an array of thunks doesn't work because we can't evaluate
97 -- in a data parallel manner. Instead, we want *one* function applied to many
98 -- array elements.
99 --
100 -- Instead, such an array of closures is represented as the vectorised
101 -- and lifted versions of (+), along with an environment array xs that
102 -- contains the partially applied arguments.
103 --
104 -- @mapP (+) xs ==> AClo plus_v plus_l xs@
105 --
106 -- When we find out what the final argument is, we can then use the lifted
107 -- closure application function to compute the result:
108 --
109 -- @PArray n (AClo plus_v plus_l xs) $:^ (PArray n' ys)
110 -- => PArray n (plus_l n xs ys)@
111 --
112 data instance PData (a :-> b)
113 = forall e. PA e
114 => AClo !(e -> a -> b) -- vectorised function, with explicit environment.
115 !(Int# -> PData e -> PData a -> PData b) -- lifted function, taking an array of environments.
116 (PData e) -- array of environments.
117
118
119 -- |Lifted closure construction
120 --
121 mkClosureP :: forall a b e.
122 PA e => (e -> a -> b)
123 -> (PArray e -> PArray a -> PArray b)
124 -> PArray e -> PArray (a :-> b)
125 {-# INLINE mkClosureP #-}
126 mkClosureP fv fl (PArray n# es)
127 = PArray n# (AClo fv (lifted fl) es)
128
129
130 liftedClosure :: forall a b e.
131 PA e => (e -> a -> b)
132 -> (Int# -> PData e -> PData a -> PData b)
133 -> PData e
134 -> PData (a :-> b)
135 {-# INLINE liftedClosure #-}
136 liftedClosure fv fl es = AClo fv fl es
137
138
139 -- |Lifted closure application
140 --
141 ($:^) :: forall a b. PArray (a :-> b) -> PArray a -> PArray b
142 {-# INLINE ($:^) #-}
143 PArray n# (AClo _ f es) $:^ PArray _ as
144 = PArray n# (f n# es as)
145
146
147 liftedApply :: forall a b. Int# -> PData (a :-> b) -> PData a -> PData b
148 {-# INLINE liftedApply #-}
149 liftedApply n# (AClo _ f es) as
150 = f n# es as
151
152
153 -- PRepr instance for closures ------------------------------------------------
154 type instance PRepr (a :-> b) = a :-> b
155
156 instance (PA a, PA b) => PA (a :-> b) where
157 toPRepr = id
158 fromPRepr = id
159 toArrPRepr = id
160 fromArrPRepr = id
161
162 instance PR (a :-> b) where
163 {-# INLINE emptyPR #-}
164 emptyPR = AClo (\_ _ -> error "empty array closure")
165 (\_ _ -> error "empty array closure")
166 (emptyPD :: PData ())
167
168 {-# INLINE replicatePR #-}
169 replicatePR n# (Clo f f' e)
170 = AClo f f' (replicatePD n# e)
171
172 {-# INLINE replicatelPR #-}
173 replicatelPR segd (AClo f f' es)
174 = AClo f f' (replicatelPD segd es)
175
176 {-# INLINE indexPR #-}
177 indexPR (AClo f f' es) i#
178 = Clo f f' (indexPD es i#)
179
180 {-# INLINE bpermutePR #-}
181 bpermutePR (AClo f f' es) n# is
182 = AClo f f' (bpermutePD es n# is)
183
184 {-# INLINE packByTagPR #-}
185 packByTagPR (AClo f f' es) n# tags t#
186 = AClo f f' (packByTagPD es n# tags t#)
187
188
189 -- Closure construction -------------------------------------------------------
190 -- | Arity-1 closures.
191 closure1 :: (a -> b) -> (PArray a -> PArray b) -> (a :-> b)
192 {-# INLINE closure1 #-}
193 closure1 fv fl = mkClosure (\_ -> fv) (\_ -> fl) ()
194
195 -- | Arity-2 closures.
196 closure2 :: PA a
197 => (a -> b -> c)
198 -> (PArray a -> PArray b -> PArray c)
199 -> (a :-> b :-> c)
200
201 {-# INLINE closure2 #-}
202 closure2 fv fl = mkClosure fv_1 fl_1 ()
203 where
204 fv_1 _ x = mkClosure fv fl x
205 fl_1 _ xs = mkClosureP fv fl xs
206
207 -- | Arity-3 closures.
208 closure3 :: (PA a, PA b)
209 => (a -> b -> c -> d)
210 -> (PArray a -> PArray b -> PArray c -> PArray d)
211 -> (a :-> b :-> c :-> d)
212
213 {-# INLINE closure3 #-}
214 closure3 fv fl = mkClosure fv_1 fl_1 ()
215 where
216 fv_1 _ x = mkClosure fv_2 fl_2 x
217 fl_1 _ xs = mkClosureP fv_2 fl_2 xs
218
219 fv_2 x y = mkClosure fv_3 fl_3 (x,y)
220 fl_2 xs ys = mkClosureP fv_3 fl_3 (zipPA# xs ys)
221
222 fv_3 (x,y) z = fv x y z
223 fl_3 ps zs = case unzipPA# ps of (xs,ys) -> fl xs ys zs
224
225 -- | Arity-4 closures.
226 closure4 :: (PA a, PA b, PA c)
227 => (a -> b -> c -> d -> e)
228 -> (PArray a -> PArray b -> PArray c -> PArray d -> PArray e)
229 -> (a :-> b :-> c :-> d :-> e)
230
231 {-# INLINE closure4 #-}
232 closure4 fv fl = mkClosure fv_1 fl_1 ()
233 where
234 fv_1 _ x = mkClosure fv_2 fl_2 x
235 fl_1 _ xs = mkClosureP fv_2 fl_2 xs
236
237 fv_2 x y = mkClosure fv_3 fl_3 (x, y)
238 fl_2 xs ys = mkClosureP fv_3 fl_3 (zipPA# xs ys)
239
240 fv_3 (x, y) z = mkClosure fv_4 fl_4 (x, y, z)
241 fl_3 xys zs = case unzipPA# xys of (xs, ys) -> mkClosureP fv_4 fl_4 (zip3PA# xs ys zs)
242
243 fv_4 (x, y, z) v = fv x y z v
244 fl_4 ps vs = case unzip3PA# ps of (xs, ys, zs) -> fl xs ys zs vs