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
91 -- Mapping from global variables to their vectorised versions.
93 vect_global_vars :: VarEnv CoreExpr
95 -- Mapping from local variables to their vectorised and lifted
98 , vect_local_vars :: VarEnv (CoreExpr, CoreExpr)
100 -- Exported variables which have a vectorised version
102 , vect_exported_vars :: VarEnv (Var, Var)
104 -- Mapping from TyCons to their vectorised versions.
105 -- TyCons which do not have to be vectorised are mapped to
108 , vect_tycons :: NameEnv TyCon
110 -- Mapping from TyCons to their PA dictionaries
112 , vect_tycon_pa :: NameEnv CoreExpr
114 -- Mapping from tyvars to their PA dictionaries
116 , vect_tyvar_pa :: VarEnv CoreExpr
119 initVEnv :: VectInfo -> DsM VEnv
122 vect_global_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info
123 , vect_local_vars = emptyVarEnv
124 , vect_exported_vars = emptyVarEnv
125 , vect_tycons = mapNameEnv snd $ vectInfoCCTyCon info
126 , vect_tycon_pa = emptyNameEnv
127 , vect_tyvar_pa = emptyVarEnv
131 updVectInfo :: VEnv -> ModGuts -> ModGuts
132 updVectInfo env guts = guts { mg_vect_info = info' }
135 vectInfoCCVar = vect_exported_vars env
136 , vectInfoCCTyCon = tc_env
139 info = mg_vect_info guts
140 tyenv = mg_types guts
142 tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv
143 , let tc_name = tyConName tc
144 , Just tc' <- [lookupNameEnv (vect_tycons env) tc_name]]
146 data VResult a = Yes VEnv a | No
148 newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VResult a) }
150 instance Monad VM where
151 return x = VM $ \bi env -> return (Yes env x)
152 VM p >>= f = VM $ \bi env -> do
155 Yes env' x -> runVM (f x) bi env'
159 noV = VM $ \bi env -> return No
161 tryV :: VM a -> VM (Maybe a)
162 tryV (VM p) = VM $ \bi env -> do
165 Yes env' x -> return (Yes env' (Just x))
166 No -> return (Yes env Nothing)
168 maybeV :: VM (Maybe a) -> VM a
169 maybeV p = maybe noV return =<< p
171 orElseV :: VM a -> VM a -> VM a
172 orElseV p q = maybe q return =<< tryV p
174 liftDs :: DsM a -> VM a
175 liftDs p = VM $ \bi env -> do { x <- p; return (Yes env x) }
177 builtin :: (Builtins -> a) -> VM a
178 builtin f = VM $ \bi env -> return (Yes env (f bi))
180 readEnv :: (VEnv -> a) -> VM a
181 readEnv f = VM $ \bi env -> return (Yes env (f env))
183 setEnv :: VEnv -> VM ()
184 setEnv env = VM $ \_ _ -> return (Yes env ())
186 updEnv :: (VEnv -> VEnv) -> VM ()
187 updEnv f = VM $ \_ env -> return (Yes (f env) ())
189 newTyVar :: FastString -> Kind -> VM Var
192 u <- liftDs newUnique
193 return $ mkTyVar (mkSysTvName u fs) k
195 lookupTyCon :: TyCon -> VM (Maybe TyCon)
196 lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc)
198 -- ----------------------------------------------------------------------------
201 vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
202 vectoriseModule info guts
204 builtins <- initBuiltins
206 r <- runVM (vectModule guts) builtins env
208 Yes env' guts' -> return $ updVectInfo env' guts'
211 vectModule :: ModGuts -> VM ModGuts
212 vectModule guts = return guts
214 -- ----------------------------------------------------------------------------
217 replicateP :: CoreExpr -> CoreExpr -> VM CoreExpr
221 rep <- builtin replicatePAVar
222 return $ mkApps (Var rep) [Type ty, pa, expr, len]
226 capply :: (CoreExpr, CoreExpr) -> (CoreExpr, CoreExpr) -> VM (CoreExpr, CoreExpr)
227 capply (vfn, lfn) (varg, larg)
229 apply <- builtin applyClosureVar
230 applyP <- builtin applyClosurePVar
231 return (mkApps (Var apply) [Type arg_ty, Type res_ty, vfn, varg],
232 mkApps (Var applyP) [Type arg_ty, Type res_ty, lfn, larg])
235 (arg_ty, res_ty) = splitClosureTy fn_ty
237 vectVar :: CoreExpr -> Var -> VM (CoreExpr, CoreExpr)
238 vectVar lc v = local v `orElseV` global v
240 local v = maybeV (readEnv $ \env -> lookupVarEnv (vect_local_vars env) v)
242 vexpr <- maybeV (readEnv $ \env -> lookupVarEnv (vect_global_vars env) v)
243 lexpr <- replicateP vexpr lc
244 return (vexpr, lexpr)
246 vectExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
247 vectExpr lc (_, AnnType ty)
250 return (Type vty, Type vty)
251 vectExpr lc (_, AnnVar v) = vectVar lc v
252 vectExpr lc (_, AnnLit lit)
255 lexpr <- replicateP vexpr lc
256 return (vexpr, lexpr)
257 vectExpr lc (_, AnnNote note expr)
259 (vexpr, lexpr) <- vectExpr lc expr
260 return (Note note vexpr, Note note lexpr)
261 vectExpr lc (_, AnnApp fn arg)
263 fn' <- vectExpr lc fn
264 arg' <- vectExpr lc arg
267 -- ----------------------------------------------------------------------------
270 paArgType :: Type -> Kind -> VM (Maybe Type)
272 | Just k' <- kindView k = paArgType ty k'
274 -- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
275 -- be made up of * and (->), i.e., they can't be coercion kinds or #.
276 paArgType ty (FunTy k1 k2)
278 tv <- newTyVar FSLIT("a") k1
279 ty1 <- paArgType' (TyVarTy tv) k1
280 ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
281 return . Just $ ForAllTy tv (FunTy ty1 ty2)
286 tc <- builtin paTyCon
287 return . Just $ TyConApp tc [ty]
292 paArgType' :: Type -> Kind -> VM Type
297 Just ty' -> return ty'
298 Nothing -> pprPanic "paArgType'" (ppr ty)
300 paOfTyCon :: TyCon -> VM CoreExpr
301 -- FIXME: just for now
302 paOfTyCon tc = maybeV (readEnv $ \env -> lookupNameEnv (vect_tycon_pa env) (tyConName tc))
304 paOfType :: Type -> VM CoreExpr
305 paOfType ty | Just ty' <- coreView ty = paOfType ty'
307 paOfType (TyVarTy tv) = maybeV (readEnv $ \env -> lookupVarEnv (vect_tyvar_pa env) tv)
308 paOfType (AppTy ty1 ty2)
312 return $ mkApps e1 [Type ty2, e2]
313 paOfType (TyConApp tc tys)
316 es <- mapM paOfType tys
317 return $ mkApps e [arg | (t,e) <- zip tys es, arg <- [Type t, e]]
318 paOfType (FunTy ty1 ty2) = paOfType (TyConApp funTyCon [ty1,ty2])
319 paOfType t@(ForAllTy tv ty) = pprPanic "paOfType:" (ppr t)
320 paOfType ty = pprPanic "paOfType:" (ppr ty)
324 -- ----------------------------------------------------------------------------
327 vectTyCon :: TyCon -> VM TyCon
329 | isFunTyCon tc = builtin closureTyCon
330 | isBoxedTupleTyCon tc = return tc
331 | isUnLiftedTyCon tc = return tc
335 Just tc' -> return tc'
337 -- FIXME: just for now
338 Nothing -> pprTrace "ccTyCon:" (ppr tc) $ return tc
340 vectType :: Type -> VM Type
341 vectType ty | Just ty' <- coreView ty = vectType ty
342 vectType (TyVarTy tv) = return $ TyVarTy tv
343 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
344 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
345 vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon)
346 (mapM vectType [ty1,ty2])
347 vectType (ForAllTy tv ty)
349 r <- paArgType (TyVarTy tv) (tyVarKind tv)
351 return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
353 vectType ty = pprPanic "vectType:" (ppr ty)
355 isClosureTyCon :: TyCon -> Bool
356 isClosureTyCon tc = tyConUnique tc == closureTyConKey
358 splitClosureTy :: Type -> (Type, Type)
360 | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
364 | otherwise = pprPanic "splitClosureTy" (ppr ty)