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 lookupVar :: Var -> VM CoreExpr
196 lookupVar v = maybeV . readEnv $ \env -> lookupVarEnv (vect_vars env) v
198 lookupTyCon :: TyCon -> VM (Maybe TyCon)
199 lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc)
201 -- ----------------------------------------------------------------------------
204 vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
205 vectoriseModule info guts
207 builtins <- initBuiltins
209 r <- runVM (vectModule guts) builtins env
211 Yes env' guts' -> return $ updVectInfo env' guts'
214 vectModule :: ModGuts -> VM ModGuts
215 vectModule guts = return guts
217 -- ----------------------------------------------------------------------------
220 paArgType :: Type -> Kind -> VM (Maybe Type)
222 | Just k' <- kindView k = paArgType ty k'
224 -- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
225 -- be made up of * and (->), i.e., they can't be coercion kinds or #.
226 paArgType ty (FunTy k1 k2)
228 tv <- newTyVar FSLIT("a") k1
229 ty1 <- paArgType' (TyVarTy tv) k1
230 ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
231 return . Just $ ForAllTy tv (FunTy ty1 ty2)
236 tc <- builtin paTyCon
237 return . Just $ TyConApp tc [ty]
242 paArgType' :: Type -> Kind -> VM Type
247 Just ty' -> return ty'
248 Nothing -> pprPanic "paArgType'" (ppr ty)
250 paOfTyCon :: TyCon -> VM CoreExpr
251 -- FIXME: just for now
252 paOfTyCon tc = maybeV (readEnv $ \env -> lookupNameEnv (vect_tycon_pa env) (tyConName tc))
254 paOfType :: Type -> VM CoreExpr
255 paOfType ty | Just ty' <- coreView ty = paOfType ty'
257 paOfType (TyVarTy tv) = maybeV (readEnv $ \env -> lookupVarEnv (vect_tyvar_pa env) tv)
258 paOfType (AppTy ty1 ty2)
262 return $ mkApps e1 [Type ty2, e2]
263 paOfType (TyConApp tc tys)
266 es <- mapM paOfType tys
267 return $ mkApps e [arg | (t,e) <- zip tys es, arg <- [Type t, e]]
268 paOfType (FunTy ty1 ty2) = paOfType (TyConApp funTyCon [ty1,ty2])
269 paOfType t@(ForAllTy tv ty) = pprPanic "paOfType:" (ppr t)
270 paOfType ty = pprPanic "paOfType:" (ppr ty)
274 -- ----------------------------------------------------------------------------
277 vectTyCon :: TyCon -> VM TyCon
279 | isFunTyCon tc = builtin closureTyCon
280 | isBoxedTupleTyCon tc = return tc
281 | isUnLiftedTyCon tc = return tc
285 Just tc' -> return tc'
287 -- FIXME: just for now
288 Nothing -> pprTrace "ccTyCon:" (ppr tc) $ return tc
290 vectType :: Type -> VM Type
291 vectType ty | Just ty' <- coreView ty = vectType ty
292 vectType (TyVarTy tv) = return $ TyVarTy tv
293 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
294 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
295 vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon)
296 (mapM vectType [ty1,ty2])
297 vectType (ForAllTy tv ty)
299 r <- paArgType (TyVarTy tv) (tyVarKind tv)
301 return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
303 vectType ty = pprPanic "vectType:" (ppr ty)
305 isClosureTyCon :: TyCon -> Bool
306 isClosureTyCon tc = tyConUnique tc == closureTyConKey
308 splitClosureTy :: Type -> (Type, Type)
310 | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
314 | otherwise = pprPanic "splitClosureTy" (ppr ty)