Commit b715bd16 authored by rl@cse.unsw.edu.au's avatar rl@cse.unsw.edu.au
Browse files

Allow variables to be mapped to arbitrary CoreExprs in vectorisation monad

parent 3ccad9ff
...@@ -7,6 +7,7 @@ import DynFlags ...@@ -7,6 +7,7 @@ import DynFlags
import HscTypes import HscTypes
import CoreLint ( showPass, endPass ) import CoreLint ( showPass, endPass )
import CoreSyn
import TyCon import TyCon
import Type import Type
import TypeRep import TypeRep
...@@ -85,9 +86,10 @@ initBuiltins ...@@ -85,9 +86,10 @@ initBuiltins
} }
data VEnv = VEnv { data VEnv = VEnv {
-- Mapping from variables to their vectorised versions -- Mapping from variables to their vectorised versions. Mapping
-- -- to expressions instead of just Vars gives us more freedom.
vect_vars :: VarEnv Var --
vect_vars :: VarEnv CoreExpr
-- Exported variables which have a vectorised version -- Exported variables which have a vectorised version
-- --
...@@ -102,7 +104,7 @@ data VEnv = VEnv { ...@@ -102,7 +104,7 @@ data VEnv = VEnv {
initVEnv :: VectInfo -> DsM VEnv initVEnv :: VectInfo -> DsM VEnv
initVEnv info initVEnv info
= return $ VEnv { = return $ VEnv {
vect_vars = mapVarEnv snd $ vectInfoCCVar info vect_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info
, vect_exported_vars = emptyVarEnv , vect_exported_vars = emptyVarEnv
, vect_tycons = mapNameEnv snd $ vectInfoCCTyCon info , vect_tycons = mapNameEnv snd $ vectInfoCCTyCon info
} }
...@@ -145,6 +147,9 @@ tryV (VM p) = VM $ \bi env -> do ...@@ -145,6 +147,9 @@ tryV (VM p) = VM $ \bi env -> do
Yes env' x -> return (Yes env' (Just x)) Yes env' x -> return (Yes env' (Just x))
No -> return (Yes env Nothing) No -> return (Yes env Nothing)
maybeV :: VM (Maybe a) -> VM a
maybeV p = maybe noV return =<< p
liftDs :: DsM a -> VM a liftDs :: DsM a -> VM a
liftDs p = VM $ \bi env -> do { x <- p; return (Yes env x) } liftDs p = VM $ \bi env -> do { x <- p; return (Yes env x) }
...@@ -166,6 +171,9 @@ newTyVar fs k ...@@ -166,6 +171,9 @@ newTyVar fs k
u <- liftDs newUnique u <- liftDs newUnique
return $ mkTyVar (mkSysTvName u fs) k return $ mkTyVar (mkSysTvName u fs) k
lookupVar :: Var -> VM CoreExpr
lookupVar v = maybeV . readEnv $ \env -> lookupVarEnv (vect_vars env) v
lookupTyCon :: TyCon -> VM (Maybe TyCon) lookupTyCon :: TyCon -> VM (Maybe TyCon)
lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc) lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment