4 noV, tryV, maybeV, orElseV, localV, initV,
11 readGEnv, setGEnv, updGEnv,
14 readLEnv, setLEnv, updLEnv,
16 lookupTyCon, extendTyVarPA
19 #include "HsVersions.h"
36 -- ----------------------------------------------------------------------------
37 -- Vectorisation monad
39 data Builtins = Builtins {
42 , closureTyCon :: TyCon
44 , applyClosureVar :: Var
45 , mkClosurePVar :: Var
46 , applyClosurePVar :: Var
49 , replicatePAVar :: Var
52 initBuiltins :: DsM Builtins
55 parrayTyCon <- dsLookupTyCon parrayTyConName
56 paTyCon <- dsLookupTyCon paTyConName
57 closureTyCon <- dsLookupTyCon closureTyConName
59 mkClosureVar <- dsLookupGlobalId mkClosureName
60 applyClosureVar <- dsLookupGlobalId applyClosureName
61 mkClosurePVar <- dsLookupGlobalId mkClosurePName
62 applyClosurePVar <- dsLookupGlobalId applyClosurePName
63 closurePAVar <- dsLookupGlobalId closurePAName
64 lengthPAVar <- dsLookupGlobalId lengthPAName
65 replicatePAVar <- dsLookupGlobalId replicatePAName
68 parrayTyCon = parrayTyCon
70 , closureTyCon = closureTyCon
71 , mkClosureVar = mkClosureVar
72 , applyClosureVar = applyClosureVar
73 , mkClosurePVar = mkClosurePVar
74 , applyClosurePVar = applyClosurePVar
75 , closurePAVar = closurePAVar
76 , lengthPAVar = lengthPAVar
77 , replicatePAVar = replicatePAVar
80 data GlobalEnv = GlobalEnv {
81 -- Mapping from global variables to their vectorised versions.
83 global_vars :: VarEnv CoreExpr
85 -- Exported variables which have a vectorised version
87 , global_exported_vars :: VarEnv (Var, Var)
89 -- Mapping from TyCons to their vectorised versions.
90 -- TyCons which do not have to be vectorised are mapped to
93 , global_tycons :: NameEnv TyCon
95 -- Mapping from TyCons to their PA dictionaries
97 , global_tycon_pa :: NameEnv CoreExpr
100 data LocalEnv = LocalEnv {
101 -- Mapping from local variables to their vectorised and
104 local_vars :: VarEnv (CoreExpr, CoreExpr)
106 -- Mapping from tyvars to their PA dictionaries
107 , local_tyvar_pa :: VarEnv CoreExpr
111 initGlobalEnv :: VectInfo -> GlobalEnv
114 global_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info
115 , global_exported_vars = emptyVarEnv
116 , global_tycons = mapNameEnv snd $ vectInfoCCTyCon info
117 , global_tycon_pa = emptyNameEnv
120 emptyLocalEnv = LocalEnv {
121 local_vars = emptyVarEnv
122 , local_tyvar_pa = emptyVarEnv
126 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
127 updVectInfo env tyenv info
129 vectInfoCCVar = global_exported_vars env
130 , vectInfoCCTyCon = tc_env
133 tc_env = mkNameEnv [(tc_name, (tc,tc'))
134 | tc <- typeEnvTyCons tyenv
135 , let tc_name = tyConName tc
136 , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
138 data VResult a = Yes GlobalEnv LocalEnv a | No
140 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
142 instance Monad VM where
143 return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
144 VM p >>= f = VM $ \bi genv lenv -> do
147 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
151 noV = VM $ \_ _ _ -> return No
153 tryV :: VM a -> VM (Maybe a)
154 tryV (VM p) = VM $ \bi genv lenv ->
158 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
159 No -> return (Yes genv lenv Nothing)
161 maybeV :: VM (Maybe a) -> VM a
162 maybeV p = maybe noV return =<< p
164 orElseV :: VM a -> VM a -> VM a
165 orElseV p q = maybe q return =<< tryV p
167 localV :: VM a -> VM a
174 liftDs :: DsM a -> VM a
175 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
177 builtin :: (Builtins -> a) -> VM a
178 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
180 readGEnv :: (GlobalEnv -> a) -> VM a
181 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
183 setGEnv :: GlobalEnv -> VM ()
184 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
186 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
187 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
189 readLEnv :: (LocalEnv -> a) -> VM a
190 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
192 setLEnv :: LocalEnv -> VM ()
193 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
195 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
196 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
198 newLocalVar :: FastString -> Type -> VM Var
201 u <- liftDs newUnique
202 return $ mkSysLocal fs u ty
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 extendTyVarPA :: Var -> CoreExpr -> VM ()
214 extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
216 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
217 initV hsc_env guts info p
219 Just r <- initDs hsc_env (mg_module guts)
226 builtins <- initBuiltins
227 r <- runVM p builtins (initGlobalEnv info) emptyLocalEnv
229 Yes genv _ x -> return $ Just (new_info genv, x)
232 new_info genv = updVectInfo genv (mg_types guts) info