1 module Vectorise( vectorise )
4 #include "HsVersions.h"
9 import CoreLint ( showPass, endPass )
16 import Name ( mkSysTvName )
25 import Control.Monad ( liftM2 )
27 vectorise :: HscEnv -> ModGuts -> IO ModGuts
28 vectorise hsc_env guts
29 | not (Opt_Vectorise `dopt` dflags) = return guts
32 showPass dflags "Vectorisation"
34 let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
35 Just guts' <- initDs hsc_env (mg_module guts)
38 (vectoriseModule info guts)
39 endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
42 dflags = hsc_dflags hsc_env
44 -- ----------------------------------------------------------------------------
45 -- Vectorisation monad
47 data Builtins = Builtins {
50 , closureTyCon :: TyCon
52 , applyClosureVar :: Var
53 , mkClosurePVar :: Var
54 , applyClosurePVar :: Var
57 , replicatePAVar :: Var
60 initBuiltins :: DsM Builtins
63 parrayTyCon <- dsLookupTyCon parrayTyConName
64 paTyCon <- dsLookupTyCon paTyConName
65 closureTyCon <- dsLookupTyCon closureTyConName
67 mkClosureVar <- dsLookupGlobalId mkClosureName
68 applyClosureVar <- dsLookupGlobalId applyClosureName
69 mkClosurePVar <- dsLookupGlobalId mkClosurePName
70 applyClosurePVar <- dsLookupGlobalId applyClosurePName
71 closurePAVar <- dsLookupGlobalId closurePAName
72 lengthPAVar <- dsLookupGlobalId lengthPAName
73 replicatePAVar <- dsLookupGlobalId replicatePAName
76 parrayTyCon = parrayTyCon
78 , closureTyCon = closureTyCon
79 , mkClosureVar = mkClosureVar
80 , applyClosureVar = applyClosureVar
81 , mkClosurePVar = mkClosurePVar
82 , applyClosurePVar = applyClosurePVar
83 , closurePAVar = closurePAVar
84 , lengthPAVar = lengthPAVar
85 , replicatePAVar = replicatePAVar
89 -- Mapping from global variables to their vectorised versions.
91 vect_global_vars :: VarEnv CoreExpr
93 -- Mapping from local variables to their vectorised and lifted
96 , vect_local_vars :: VarEnv (CoreExpr, CoreExpr)
98 -- Exported variables which have a vectorised version
100 , vect_exported_vars :: VarEnv (Var, Var)
102 -- Mapping from TyCons to their vectorised versions.
103 -- TyCons which do not have to be vectorised are mapped to
106 , vect_tycons :: NameEnv TyCon
108 -- Mapping from TyCons to their PA dictionaries
110 , vect_tycon_pa :: NameEnv CoreExpr
112 -- Mapping from tyvars to their PA dictionaries
114 , vect_tyvar_pa :: VarEnv CoreExpr
117 initVEnv :: VectInfo -> DsM VEnv
120 vect_global_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info
121 , vect_local_vars = emptyVarEnv
122 , vect_exported_vars = emptyVarEnv
123 , vect_tycons = mapNameEnv snd $ vectInfoCCTyCon info
124 , vect_tycon_pa = emptyNameEnv
125 , vect_tyvar_pa = emptyVarEnv
129 updVectInfo :: VEnv -> ModGuts -> ModGuts
130 updVectInfo env guts = guts { mg_vect_info = info' }
133 vectInfoCCVar = vect_exported_vars env
134 , vectInfoCCTyCon = tc_env
137 info = mg_vect_info guts
138 tyenv = mg_types guts
140 tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv
141 , let tc_name = tyConName tc
142 , Just tc' <- [lookupNameEnv (vect_tycons env) tc_name]]
144 data VResult a = Yes VEnv a | No
146 newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VResult a) }
148 instance Monad VM where
149 return x = VM $ \bi env -> return (Yes env x)
150 VM p >>= f = VM $ \bi env -> do
153 Yes env' x -> runVM (f x) bi env'
157 noV = VM $ \bi env -> return No
159 tryV :: VM a -> VM (Maybe a)
160 tryV (VM p) = VM $ \bi env -> do
163 Yes env' x -> return (Yes env' (Just x))
164 No -> return (Yes env Nothing)
166 maybeV :: VM (Maybe a) -> VM a
167 maybeV p = maybe noV return =<< p
169 orElseV :: VM a -> VM a -> VM a
170 orElseV p q = maybe q return =<< tryV p
172 liftDs :: DsM a -> VM a
173 liftDs p = VM $ \bi env -> do { x <- p; return (Yes env x) }
175 builtin :: (Builtins -> a) -> VM a
176 builtin f = VM $ \bi env -> return (Yes env (f bi))
178 readEnv :: (VEnv -> a) -> VM a
179 readEnv f = VM $ \bi env -> return (Yes env (f env))
181 setEnv :: VEnv -> VM ()
182 setEnv env = VM $ \_ _ -> return (Yes env ())
184 updEnv :: (VEnv -> VEnv) -> VM ()
185 updEnv f = VM $ \_ env -> return (Yes (f env) ())
187 newTyVar :: FastString -> Kind -> VM Var
190 u <- liftDs newUnique
191 return $ mkTyVar (mkSysTvName u fs) k
193 lookupVar :: Var -> VM CoreExpr
194 lookupVar v = maybeV . readEnv $ \env -> lookupVarEnv (vect_vars env) v
196 lookupTyCon :: TyCon -> VM (Maybe TyCon)
197 lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc)
199 -- ----------------------------------------------------------------------------
202 vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
203 vectoriseModule info guts
205 builtins <- initBuiltins
207 r <- runVM (vectModule guts) builtins env
209 Yes env' guts' -> return $ updVectInfo env' guts'
212 vectModule :: ModGuts -> VM ModGuts
213 vectModule guts = return guts
215 -- ----------------------------------------------------------------------------
218 paArgType :: Type -> Kind -> VM (Maybe Type)
220 | Just k' <- kindView k = paArgType ty k'
222 -- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
223 -- be made up of * and (->), i.e., they can't be coercion kinds or #.
224 paArgType ty (FunTy k1 k2)
226 tv <- newTyVar FSLIT("a") k1
227 ty1 <- paArgType' (TyVarTy tv) k1
228 ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
229 return . Just $ ForAllTy tv (FunTy ty1 ty2)
234 tc <- builtin paTyCon
235 return . Just $ TyConApp tc [ty]
240 paArgType' :: Type -> Kind -> VM Type
245 Just ty' -> return ty'
246 Nothing -> pprPanic "paArgType'" (ppr ty)
248 vectTyCon :: TyCon -> VM TyCon
250 | isFunTyCon tc = builtin closureTyCon
251 | isBoxedTupleTyCon tc = return tc
252 | isUnLiftedTyCon tc = return tc
256 Just tc' -> return tc'
258 -- FIXME: just for now
259 Nothing -> pprTrace "ccTyCon:" (ppr tc) $ return tc
261 vectType :: Type -> VM Type
262 vectType ty | Just ty' <- coreView ty = vectType ty
263 vectType (TyVarTy tv) = return $ TyVarTy tv
264 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
265 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
266 vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon)
267 (mapM vectType [ty1,ty2])
268 vectType (ForAllTy tv ty)
270 r <- paArgType (TyVarTy tv) (tyVarKind tv)
272 return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
274 vectType ty = pprPanic "vectType:" (ppr ty)
276 isClosureTyCon :: TyCon -> Bool
277 isClosureTyCon tc = tyConUnique tc == closureTyConKey
279 splitClosureTy :: Type -> (Type, Type)
281 | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
285 | otherwise = pprPanic "splitClosureTy" (ppr ty)