Collect hoisted vectorised functions
[ghc.git] / compiler / vectorise / VectMonad.hs
1 module VectMonad (
2 VM,
3
4 noV, tryV, maybeV, orElseV, localV, closedV, initV,
5 newLocalVar, newTyVar,
6
7 Builtins(..), paDictTyCon,
8 builtin,
9
10 GlobalEnv(..),
11 readGEnv, setGEnv, updGEnv,
12
13 LocalEnv(..),
14 readLEnv, setLEnv, updLEnv,
15
16 lookupTyCon,
17 lookupTyVarPA, extendTyVarPA, deleteTyVarPA,
18
19 lookupInst, lookupFamInst
20 ) where
21
22 #include "HsVersions.h"
23
24 import HscTypes
25 import CoreSyn
26 import Class
27 import TyCon
28 import Type
29 import Var
30 import VarEnv
31 import Id
32 import Name
33 import NameEnv
34
35 import DsMonad
36 import PrelNames
37
38 import InstEnv
39 import FamInstEnv
40
41 import Panic
42 import Outputable
43 import FastString
44
45 -- ----------------------------------------------------------------------------
46 -- Vectorisation monad
47
48 data Builtins = Builtins {
49 parrayTyCon :: TyCon
50 , paClass :: Class
51 , closureTyCon :: TyCon
52 , mkClosureVar :: Var
53 , applyClosureVar :: Var
54 , mkClosurePVar :: Var
55 , applyClosurePVar :: Var
56 , lengthPAVar :: Var
57 , replicatePAVar :: Var
58 }
59
60 paDictTyCon :: Builtins -> TyCon
61 paDictTyCon = classTyCon . paClass
62
63 initBuiltins :: DsM Builtins
64 initBuiltins
65 = do
66 parrayTyCon <- dsLookupTyCon parrayTyConName
67 paClass <- dsLookupClass paClassName
68 closureTyCon <- dsLookupTyCon closureTyConName
69
70 mkClosureVar <- dsLookupGlobalId mkClosureName
71 applyClosureVar <- dsLookupGlobalId applyClosureName
72 mkClosurePVar <- dsLookupGlobalId mkClosurePName
73 applyClosurePVar <- dsLookupGlobalId applyClosurePName
74 lengthPAVar <- dsLookupGlobalId lengthPAName
75 replicatePAVar <- dsLookupGlobalId replicatePAName
76
77 return $ Builtins {
78 parrayTyCon = parrayTyCon
79 , paClass = paClass
80 , closureTyCon = closureTyCon
81 , mkClosureVar = mkClosureVar
82 , applyClosureVar = applyClosureVar
83 , mkClosurePVar = mkClosurePVar
84 , applyClosurePVar = applyClosurePVar
85 , lengthPAVar = lengthPAVar
86 , replicatePAVar = replicatePAVar
87 }
88
89 data GlobalEnv = GlobalEnv {
90 -- Mapping from global variables to their vectorised versions.
91 --
92 global_vars :: VarEnv CoreExpr
93
94 -- Exported variables which have a vectorised version
95 --
96 , global_exported_vars :: VarEnv (Var, Var)
97
98 -- Mapping from TyCons to their vectorised versions.
99 -- TyCons which do not have to be vectorised are mapped to
100 -- themselves.
101 --
102 , global_tycons :: NameEnv TyCon
103
104 -- Mapping from TyCons to their PA dictionaries
105 --
106 , global_tycon_pa :: NameEnv CoreExpr
107
108 -- External package inst-env & home-package inst-env for class
109 -- instances
110 --
111 , global_inst_env :: (InstEnv, InstEnv)
112
113 -- External package inst-env & home-package inst-env for family
114 -- instances
115 --
116 , global_fam_inst_env :: FamInstEnvs
117 }
118
119 data LocalEnv = LocalEnv {
120 -- Mapping from local variables to their vectorised and
121 -- lifted versions
122 --
123 local_vars :: VarEnv (CoreExpr, CoreExpr)
124
125 -- Mapping from tyvars to their PA dictionaries
126 , local_tyvar_pa :: VarEnv CoreExpr
127
128 -- Hoisted bindings
129 , local_bindings :: [(Var, CoreExpr)]
130 }
131
132
133 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
134 initGlobalEnv info instEnvs famInstEnvs
135 = GlobalEnv {
136 global_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info
137 , global_exported_vars = emptyVarEnv
138 , global_tycons = mapNameEnv snd $ vectInfoCCTyCon info
139 , global_tycon_pa = emptyNameEnv
140 , global_inst_env = instEnvs
141 , global_fam_inst_env = famInstEnvs
142 }
143
144 emptyLocalEnv = LocalEnv {
145 local_vars = emptyVarEnv
146 , local_tyvar_pa = emptyVarEnv
147 , local_bindings = []
148 }
149
150 -- FIXME
151 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
152 updVectInfo env tyenv info
153 = info {
154 vectInfoCCVar = global_exported_vars env
155 , vectInfoCCTyCon = tc_env
156 }
157 where
158 tc_env = mkNameEnv [(tc_name, (tc,tc'))
159 | tc <- typeEnvTyCons tyenv
160 , let tc_name = tyConName tc
161 , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
162
163 data VResult a = Yes GlobalEnv LocalEnv a | No
164
165 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
166
167 instance Monad VM where
168 return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
169 VM p >>= f = VM $ \bi genv lenv -> do
170 r <- p bi genv lenv
171 case r of
172 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
173 No -> return No
174
175 noV :: VM a
176 noV = VM $ \_ _ _ -> return No
177
178 tryV :: VM a -> VM (Maybe a)
179 tryV (VM p) = VM $ \bi genv lenv ->
180 do
181 r <- p bi genv lenv
182 case r of
183 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
184 No -> return (Yes genv lenv Nothing)
185
186 maybeV :: VM (Maybe a) -> VM a
187 maybeV p = maybe noV return =<< p
188
189 orElseV :: VM a -> VM a -> VM a
190 orElseV p q = maybe q return =<< tryV p
191
192 localV :: VM a -> VM a
193 localV p = do
194 env <- readLEnv id
195 x <- p
196 setLEnv env
197 return x
198
199 closedV :: VM a -> VM a
200 closedV p = do
201 env <- readLEnv id
202 setLEnv emptyLocalEnv
203 x <- p
204 setLEnv env
205 return x
206
207 liftDs :: DsM a -> VM a
208 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
209
210 builtin :: (Builtins -> a) -> VM a
211 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
212
213 readGEnv :: (GlobalEnv -> a) -> VM a
214 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
215
216 setGEnv :: GlobalEnv -> VM ()
217 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
218
219 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
220 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
221
222 readLEnv :: (LocalEnv -> a) -> VM a
223 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
224
225 setLEnv :: LocalEnv -> VM ()
226 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
227
228 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
229 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
230
231 getInstEnv :: VM (InstEnv, InstEnv)
232 getInstEnv = readGEnv global_inst_env
233
234 getFamInstEnv :: VM FamInstEnvs
235 getFamInstEnv = readGEnv global_fam_inst_env
236
237 newLocalVar :: FastString -> Type -> VM Var
238 newLocalVar fs ty
239 = do
240 u <- liftDs newUnique
241 return $ mkSysLocal fs u ty
242
243 newTyVar :: FastString -> Kind -> VM Var
244 newTyVar fs k
245 = do
246 u <- liftDs newUnique
247 return $ mkTyVar (mkSysTvName u fs) k
248
249 lookupTyCon :: TyCon -> VM (Maybe TyCon)
250 lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
251
252 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
253 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
254
255 extendTyVarPA :: Var -> CoreExpr -> VM ()
256 extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
257
258 deleteTyVarPA :: Var -> VM ()
259 deleteTyVarPA tv = updLEnv $ \env -> env { local_tyvar_pa = delVarEnv (local_tyvar_pa env) tv }
260
261 -- Look up the dfun of a class instance.
262 --
263 -- The match must be unique - ie, match exactly one instance - but the
264 -- type arguments used for matching may be more specific than those of
265 -- the class instance declaration. The found class instances must not have
266 -- any type variables in the instance context that do not appear in the
267 -- instances head (i.e., no flexi vars); for details for what this means,
268 -- see the docs at InstEnv.lookupInstEnv.
269 --
270 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
271 lookupInst cls tys
272 = do { instEnv <- getInstEnv
273 ; case lookupInstEnv instEnv cls tys of
274 ([(inst, inst_tys)], _)
275 | noFlexiVar -> return (instanceDFunId inst, inst_tys')
276 | otherwise -> pprPanic "VectMonad.lookupInst: flexi var: "
277 (ppr $ mkTyConApp (classTyCon cls) tys)
278 where
279 inst_tys' = [ty | Right ty <- inst_tys]
280 noFlexiVar = all isRight inst_tys
281 _other -> noV
282 }
283 where
284 isRight (Left _) = False
285 isRight (Right _) = True
286
287 -- Look up the representation tycon of a family instance.
288 --
289 -- The match must be unique - ie, match exactly one instance - but the
290 -- type arguments used for matching may be more specific than those of
291 -- the family instance declaration.
292 --
293 -- Return the instance tycon and its type instance. For example, if we have
294 --
295 -- lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
296 --
297 -- then we have a coercion (ie, type instance of family instance coercion)
298 --
299 -- :Co:R42T Int :: T [Int] ~ :R42T Int
300 --
301 -- which implies that :R42T was declared as 'data instance T [a]'.
302 --
303 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
304 lookupFamInst tycon tys
305 = ASSERT( isOpenTyCon tycon )
306 do { instEnv <- getFamInstEnv
307 ; case lookupFamInstEnv instEnv tycon tys of
308 [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
309 _other ->
310 pprPanic "VectMonad.lookupFamInst: not found: "
311 (ppr $ mkTyConApp tycon tys)
312 }
313
314 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
315 initV hsc_env guts info p
316 = do
317 eps <- hscEPS hsc_env
318 let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
319 let instEnvs = (eps_inst_env eps, mg_inst_env guts)
320
321 Just r <- initDs hsc_env (mg_module guts)
322 (mg_rdr_env guts)
323 (mg_types guts)
324 (go instEnvs famInstEnvs)
325 return r
326 where
327
328 go instEnvs famInstEnvs =
329 do
330 builtins <- initBuiltins
331 r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs)
332 emptyLocalEnv
333 case r of
334 Yes genv _ x -> return $ Just (new_info genv, x)
335 No -> return Nothing
336
337 new_info genv = updVectInfo genv (mg_types guts) info
338