1 module Vectorise( vectorise )
4 #include "HsVersions.h"
9 import CoreLint ( showPass, endPass )
18 import Name ( mkSysTvName )
27 import Control.Monad ( liftM2 )
29 vectorise :: HscEnv -> ModGuts -> IO ModGuts
30 vectorise hsc_env guts
31 | not (Opt_Vectorise `dopt` dflags) = return guts
34 showPass dflags "Vectorisation"
36 let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
37 Just guts' <- initDs hsc_env (mg_module guts)
40 (vectoriseModule info guts)
41 endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
44 dflags = hsc_dflags hsc_env
46 -- ----------------------------------------------------------------------------
47 -- Vectorisation monad
49 data Builtins = Builtins {
52 , closureTyCon :: TyCon
54 , applyClosureVar :: Var
55 , mkClosurePVar :: Var
56 , applyClosurePVar :: Var
59 , replicatePAVar :: Var
62 initBuiltins :: DsM Builtins
65 parrayTyCon <- dsLookupTyCon parrayTyConName
66 paTyCon <- dsLookupTyCon paTyConName
67 closureTyCon <- dsLookupTyCon closureTyConName
69 mkClosureVar <- dsLookupGlobalId mkClosureName
70 applyClosureVar <- dsLookupGlobalId applyClosureName
71 mkClosurePVar <- dsLookupGlobalId mkClosurePName
72 applyClosurePVar <- dsLookupGlobalId applyClosurePName
73 closurePAVar <- dsLookupGlobalId closurePAName
74 lengthPAVar <- dsLookupGlobalId lengthPAName
75 replicatePAVar <- dsLookupGlobalId replicatePAName
78 parrayTyCon = parrayTyCon
80 , closureTyCon = closureTyCon
81 , mkClosureVar = mkClosureVar
82 , applyClosureVar = applyClosureVar
83 , mkClosurePVar = mkClosurePVar
84 , applyClosurePVar = applyClosurePVar
85 , closurePAVar = closurePAVar
86 , lengthPAVar = lengthPAVar
87 , replicatePAVar = replicatePAVar
90 data GlobalEnv = GlobalEnv {
91 -- Mapping from global variables to their vectorised versions.
93 global_vars :: VarEnv CoreExpr
95 -- Exported variables which have a vectorised version
97 , global_exported_vars :: VarEnv (Var, Var)
99 -- Mapping from TyCons to their vectorised versions.
100 -- TyCons which do not have to be vectorised are mapped to
103 , global_tycons :: NameEnv TyCon
105 -- Mapping from TyCons to their PA dictionaries
107 , global_tycon_pa :: NameEnv CoreExpr
110 data LocalEnv = LocalEnv {
111 -- Mapping from local variables to their vectorised and
114 local_vars :: VarEnv (CoreExpr, CoreExpr)
116 -- Mapping from tyvars to their PA dictionaries
117 , local_tyvar_pa :: VarEnv CoreExpr
121 initGlobalEnv :: VectInfo -> GlobalEnv
124 global_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info
125 , global_exported_vars = emptyVarEnv
126 , global_tycons = mapNameEnv snd $ vectInfoCCTyCon info
127 , global_tycon_pa = emptyNameEnv
130 emptyLocalEnv = LocalEnv {
131 local_vars = emptyVarEnv
132 , local_tyvar_pa = emptyVarEnv
136 updVectInfo :: GlobalEnv -> ModGuts -> ModGuts
137 updVectInfo env guts = guts { mg_vect_info = info' }
140 vectInfoCCVar = global_exported_vars env
141 , vectInfoCCTyCon = tc_env
144 info = mg_vect_info guts
145 tyenv = mg_types guts
147 tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv
148 , let tc_name = tyConName tc
149 , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
151 data VResult a = Yes GlobalEnv LocalEnv a | No
153 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
155 instance Monad VM where
156 return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
157 VM p >>= f = VM $ \bi genv lenv -> do
160 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
164 noV = VM $ \_ _ _ -> return No
166 tryV :: VM a -> VM (Maybe a)
167 tryV (VM p) = VM $ \bi genv lenv ->
171 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
172 No -> return (Yes genv lenv Nothing)
174 maybeV :: VM (Maybe a) -> VM a
175 maybeV p = maybe noV return =<< p
177 orElseV :: VM a -> VM a -> VM a
178 orElseV p q = maybe q return =<< tryV p
180 liftDs :: DsM a -> VM a
181 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
183 builtin :: (Builtins -> a) -> VM a
184 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
186 readGEnv :: (GlobalEnv -> a) -> VM a
187 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
189 setGEnv :: GlobalEnv -> VM ()
190 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
192 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
193 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
195 readLEnv :: (LocalEnv -> a) -> VM a
196 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
198 setLEnv :: LocalEnv -> VM ()
199 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
201 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
202 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
204 newTyVar :: FastString -> Kind -> VM Var
207 u <- liftDs newUnique
208 return $ mkTyVar (mkSysTvName u fs) k
210 lookupTyCon :: TyCon -> VM (Maybe TyCon)
211 lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
213 -- ----------------------------------------------------------------------------
216 vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
217 vectoriseModule info guts
219 builtins <- initBuiltins
220 r <- runVM (vectModule guts) builtins (initGlobalEnv info) emptyLocalEnv
222 Yes genv _ guts' -> return $ updVectInfo genv guts'
225 vectModule :: ModGuts -> VM ModGuts
226 vectModule guts = return guts
228 -- ----------------------------------------------------------------------------
231 replicateP :: CoreExpr -> CoreExpr -> VM CoreExpr
235 rep <- builtin replicatePAVar
236 return $ mkApps (Var rep) [Type ty, pa, expr, len]
240 capply :: (CoreExpr, CoreExpr) -> (CoreExpr, CoreExpr) -> VM (CoreExpr, CoreExpr)
241 capply (vfn, lfn) (varg, larg)
243 apply <- builtin applyClosureVar
244 applyP <- builtin applyClosurePVar
245 return (mkApps (Var apply) [Type arg_ty, Type res_ty, vfn, varg],
246 mkApps (Var applyP) [Type arg_ty, Type res_ty, lfn, larg])
249 (arg_ty, res_ty) = splitClosureTy fn_ty
251 vectVar :: CoreExpr -> Var -> VM (CoreExpr, CoreExpr)
252 vectVar lc v = local v `orElseV` global v
254 local v = maybeV (readLEnv $ \env -> lookupVarEnv (local_vars env) v)
256 vexpr <- maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
257 lexpr <- replicateP vexpr lc
258 return (vexpr, lexpr)
260 vectExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
261 vectExpr lc (_, AnnType ty)
264 return (Type vty, Type vty)
265 vectExpr lc (_, AnnVar v) = vectVar lc v
266 vectExpr lc (_, AnnLit lit)
269 lexpr <- replicateP vexpr lc
270 return (vexpr, lexpr)
271 vectExpr lc (_, AnnNote note expr)
273 (vexpr, lexpr) <- vectExpr lc expr
274 return (Note note vexpr, Note note lexpr)
275 vectExpr lc (_, AnnApp fn arg)
277 fn' <- vectExpr lc fn
278 arg' <- vectExpr lc arg
281 -- ----------------------------------------------------------------------------
284 paArgType :: Type -> Kind -> VM (Maybe Type)
286 | Just k' <- kindView k = paArgType ty k'
288 -- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
289 -- be made up of * and (->), i.e., they can't be coercion kinds or #.
290 paArgType ty (FunTy k1 k2)
292 tv <- newTyVar FSLIT("a") k1
293 ty1 <- paArgType' (TyVarTy tv) k1
294 ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
295 return . Just $ ForAllTy tv (FunTy ty1 ty2)
300 tc <- builtin paTyCon
301 return . Just $ TyConApp tc [ty]
306 paArgType' :: Type -> Kind -> VM Type
311 Just ty' -> return ty'
312 Nothing -> pprPanic "paArgType'" (ppr ty)
314 paOfTyCon :: TyCon -> VM CoreExpr
315 -- FIXME: just for now
316 paOfTyCon tc = maybeV (readGEnv $ \env -> lookupNameEnv (global_tycon_pa env) (tyConName tc))
318 paOfType :: Type -> VM CoreExpr
319 paOfType ty | Just ty' <- coreView ty = paOfType ty'
321 paOfType (TyVarTy tv) = maybeV (readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv)
322 paOfType (AppTy ty1 ty2)
326 return $ mkApps e1 [Type ty2, e2]
327 paOfType (TyConApp tc tys)
330 es <- mapM paOfType tys
331 return $ mkApps e [arg | (t,e) <- zip tys es, arg <- [Type t, e]]
332 paOfType (FunTy ty1 ty2) = paOfType (TyConApp funTyCon [ty1,ty2])
333 paOfType t@(ForAllTy tv ty) = pprPanic "paOfType:" (ppr t)
334 paOfType ty = pprPanic "paOfType:" (ppr ty)
338 -- ----------------------------------------------------------------------------
341 vectTyCon :: TyCon -> VM TyCon
343 | isFunTyCon tc = builtin closureTyCon
344 | isBoxedTupleTyCon tc = return tc
345 | isUnLiftedTyCon tc = return tc
349 Just tc' -> return tc'
351 -- FIXME: just for now
352 Nothing -> pprTrace "ccTyCon:" (ppr tc) $ return tc
354 vectType :: Type -> VM Type
355 vectType ty | Just ty' <- coreView ty = vectType ty
356 vectType (TyVarTy tv) = return $ TyVarTy tv
357 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
358 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
359 vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon)
360 (mapM vectType [ty1,ty2])
361 vectType (ForAllTy tv ty)
363 r <- paArgType (TyVarTy tv) (tyVarKind tv)
365 return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
367 vectType ty = pprPanic "vectType:" (ppr ty)
369 isClosureTyCon :: TyCon -> Bool
370 isClosureTyCon tc = tyConUnique tc == closureTyConKey
372 splitClosureTy :: Type -> (Type, Type)
374 | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
378 | otherwise = pprPanic "splitClosureTy" (ppr ty)