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

Put vectorisation monad into a separate file

parent f2e0a6b9
module VectMonad (
VM,
noV, tryV, maybeV, orElseV, localV, initV,
newLocalVar, newTyVar,
Builtins(..),
builtin,
GlobalEnv(..),
readGEnv, setGEnv, updGEnv,
LocalEnv(..),
readLEnv, setLEnv, updLEnv,
lookupTyCon, extendTyVarPA
) where
#include "HsVersions.h"
import HscTypes
import CoreSyn
import TyCon
import Type
import Var
import VarEnv
import Id
import Name
import NameEnv
import DsMonad
import PrelNames
import FastString
-- ----------------------------------------------------------------------------
-- Vectorisation monad
data Builtins = Builtins {
parrayTyCon :: TyCon
, paTyCon :: TyCon
, closureTyCon :: TyCon
, mkClosureVar :: Var
, applyClosureVar :: Var
, mkClosurePVar :: Var
, applyClosurePVar :: Var
, closurePAVar :: Var
, lengthPAVar :: Var
, replicatePAVar :: Var
}
initBuiltins :: DsM Builtins
initBuiltins
= do
parrayTyCon <- dsLookupTyCon parrayTyConName
paTyCon <- dsLookupTyCon paTyConName
closureTyCon <- dsLookupTyCon closureTyConName
mkClosureVar <- dsLookupGlobalId mkClosureName
applyClosureVar <- dsLookupGlobalId applyClosureName
mkClosurePVar <- dsLookupGlobalId mkClosurePName
applyClosurePVar <- dsLookupGlobalId applyClosurePName
closurePAVar <- dsLookupGlobalId closurePAName
lengthPAVar <- dsLookupGlobalId lengthPAName
replicatePAVar <- dsLookupGlobalId replicatePAName
return $ Builtins {
parrayTyCon = parrayTyCon
, paTyCon = paTyCon
, closureTyCon = closureTyCon
, mkClosureVar = mkClosureVar
, applyClosureVar = applyClosureVar
, mkClosurePVar = mkClosurePVar
, applyClosurePVar = applyClosurePVar
, closurePAVar = closurePAVar
, lengthPAVar = lengthPAVar
, replicatePAVar = replicatePAVar
}
data GlobalEnv = GlobalEnv {
-- Mapping from global variables to their vectorised versions.
--
global_vars :: VarEnv CoreExpr
-- Exported variables which have a vectorised version
--
, global_exported_vars :: VarEnv (Var, Var)
-- Mapping from TyCons to their vectorised versions.
-- TyCons which do not have to be vectorised are mapped to
-- themselves.
--
, global_tycons :: NameEnv TyCon
-- Mapping from TyCons to their PA dictionaries
--
, global_tycon_pa :: NameEnv CoreExpr
}
data LocalEnv = LocalEnv {
-- Mapping from local variables to their vectorised and
-- lifted versions
--
local_vars :: VarEnv (CoreExpr, CoreExpr)
-- Mapping from tyvars to their PA dictionaries
, local_tyvar_pa :: VarEnv CoreExpr
}
initGlobalEnv :: VectInfo -> GlobalEnv
initGlobalEnv info
= GlobalEnv {
global_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info
, global_exported_vars = emptyVarEnv
, global_tycons = mapNameEnv snd $ vectInfoCCTyCon info
, global_tycon_pa = emptyNameEnv
}
emptyLocalEnv = LocalEnv {
local_vars = emptyVarEnv
, local_tyvar_pa = emptyVarEnv
}
-- FIXME
updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
updVectInfo env tyenv info
= info {
vectInfoCCVar = global_exported_vars env
, vectInfoCCTyCon = tc_env
}
where
tc_env = mkNameEnv [(tc_name, (tc,tc'))
| tc <- typeEnvTyCons tyenv
, let tc_name = tyConName tc
, Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
data VResult a = Yes GlobalEnv LocalEnv a | No
newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
instance Monad VM where
return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
VM p >>= f = VM $ \bi genv lenv -> do
r <- p bi genv lenv
case r of
Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
No -> return No
noV :: VM a
noV = VM $ \_ _ _ -> return No
tryV :: VM a -> VM (Maybe a)
tryV (VM p) = VM $ \bi genv lenv ->
do
r <- p bi genv lenv
case r of
Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
No -> return (Yes genv lenv Nothing)
maybeV :: VM (Maybe a) -> VM a
maybeV p = maybe noV return =<< p
orElseV :: VM a -> VM a -> VM a
orElseV p q = maybe q return =<< tryV p
localV :: VM a -> VM a
localV p = do
env <- readLEnv id
x <- p
setLEnv env
return x
liftDs :: DsM a -> VM a
liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
builtin :: (Builtins -> a) -> VM a
builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
readGEnv :: (GlobalEnv -> a) -> VM a
readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
setGEnv :: GlobalEnv -> VM ()
setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
readLEnv :: (LocalEnv -> a) -> VM a
readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
setLEnv :: LocalEnv -> VM ()
setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
newLocalVar :: FastString -> Type -> VM Var
newLocalVar fs ty
= do
u <- liftDs newUnique
return $ mkSysLocal fs u ty
newTyVar :: FastString -> Kind -> VM Var
newTyVar fs k
= do
u <- liftDs newUnique
return $ mkTyVar (mkSysTvName u fs) k
lookupTyCon :: TyCon -> VM (Maybe TyCon)
lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
extendTyVarPA :: Var -> CoreExpr -> VM ()
extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
initV hsc_env guts info p
= do
Just r <- initDs hsc_env (mg_module guts)
(mg_rdr_env guts)
(mg_types guts)
go
return r
where
go = do
builtins <- initBuiltins
r <- runVM p builtins (initGlobalEnv info) emptyLocalEnv
case r of
Yes genv _ x -> return $ Just (new_info genv, x)
No -> return Nothing
new_info genv = updVectInfo genv (mg_types guts) info
......@@ -3,6 +3,8 @@ where
#include "HsVersions.h"
import VectMonad
import DynFlags
import HscTypes
......@@ -35,215 +37,17 @@ vectorise hsc_env guts
showPass dflags "Vectorisation"
eps <- hscEPS hsc_env
let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
Just guts' <- initDs hsc_env (mg_module guts)
(mg_rdr_env guts)
(mg_types guts)
(vectoriseModule info guts)
Just (info', guts') <- initV hsc_env guts info (vectModule guts)
endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
return guts'
return $ guts' { mg_vect_info = info' }
where
dflags = hsc_dflags hsc_env
-- ----------------------------------------------------------------------------
-- Vectorisation monad
data Builtins = Builtins {
parrayTyCon :: TyCon
, paTyCon :: TyCon
, closureTyCon :: TyCon
, mkClosureVar :: Var
, applyClosureVar :: Var
, mkClosurePVar :: Var
, applyClosurePVar :: Var
, closurePAVar :: Var
, lengthPAVar :: Var
, replicatePAVar :: Var
}
initBuiltins :: DsM Builtins
initBuiltins
= do
parrayTyCon <- dsLookupTyCon parrayTyConName
paTyCon <- dsLookupTyCon paTyConName
closureTyCon <- dsLookupTyCon closureTyConName
mkClosureVar <- dsLookupGlobalId mkClosureName
applyClosureVar <- dsLookupGlobalId applyClosureName
mkClosurePVar <- dsLookupGlobalId mkClosurePName
applyClosurePVar <- dsLookupGlobalId applyClosurePName
closurePAVar <- dsLookupGlobalId closurePAName
lengthPAVar <- dsLookupGlobalId lengthPAName
replicatePAVar <- dsLookupGlobalId replicatePAName
return $ Builtins {
parrayTyCon = parrayTyCon
, paTyCon = paTyCon
, closureTyCon = closureTyCon
, mkClosureVar = mkClosureVar
, applyClosureVar = applyClosureVar
, mkClosurePVar = mkClosurePVar
, applyClosurePVar = applyClosurePVar
, closurePAVar = closurePAVar
, lengthPAVar = lengthPAVar
, replicatePAVar = replicatePAVar
}
data GlobalEnv = GlobalEnv {
-- Mapping from global variables to their vectorised versions.
--
global_vars :: VarEnv CoreExpr
-- Exported variables which have a vectorised version
--
, global_exported_vars :: VarEnv (Var, Var)
-- Mapping from TyCons to their vectorised versions.
-- TyCons which do not have to be vectorised are mapped to
-- themselves.
--
, global_tycons :: NameEnv TyCon
-- Mapping from TyCons to their PA dictionaries
--
, global_tycon_pa :: NameEnv CoreExpr
}
data LocalEnv = LocalEnv {
-- Mapping from local variables to their vectorised and
-- lifted versions
--
local_vars :: VarEnv (CoreExpr, CoreExpr)
-- Mapping from tyvars to their PA dictionaries
, local_tyvar_pa :: VarEnv CoreExpr
}
initGlobalEnv :: VectInfo -> GlobalEnv
initGlobalEnv info
= GlobalEnv {
global_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info
, global_exported_vars = emptyVarEnv
, global_tycons = mapNameEnv snd $ vectInfoCCTyCon info
, global_tycon_pa = emptyNameEnv
}
emptyLocalEnv = LocalEnv {
local_vars = emptyVarEnv
, local_tyvar_pa = emptyVarEnv
}
-- FIXME
updVectInfo :: GlobalEnv -> ModGuts -> ModGuts
updVectInfo env guts = guts { mg_vect_info = info' }
where
info' = info {
vectInfoCCVar = global_exported_vars env
, vectInfoCCTyCon = tc_env
}
info = mg_vect_info guts
tyenv = mg_types guts
tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv
, let tc_name = tyConName tc
, Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
data VResult a = Yes GlobalEnv LocalEnv a | No
newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
instance Monad VM where
return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
VM p >>= f = VM $ \bi genv lenv -> do
r <- p bi genv lenv
case r of
Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
No -> return No
noV :: VM a
noV = VM $ \_ _ _ -> return No
tryV :: VM a -> VM (Maybe a)
tryV (VM p) = VM $ \bi genv lenv ->
do
r <- p bi genv lenv
case r of
Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
No -> return (Yes genv lenv Nothing)
maybeV :: VM (Maybe a) -> VM a
maybeV p = maybe noV return =<< p
orElseV :: VM a -> VM a -> VM a
orElseV p q = maybe q return =<< tryV p
localV :: VM a -> VM a
localV p = do
env <- readLEnv id
x <- p
setLEnv env
return x
liftDs :: DsM a -> VM a
liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
builtin :: (Builtins -> a) -> VM a
builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
readGEnv :: (GlobalEnv -> a) -> VM a
readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
setGEnv :: GlobalEnv -> VM ()
setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
readLEnv :: (LocalEnv -> a) -> VM a
readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
setLEnv :: LocalEnv -> VM ()
setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
newLocalVar :: FastString -> Type -> VM Var
newLocalVar fs ty
= do
u <- liftDs newUnique
return $ mkSysLocal fs u ty
newTyVar :: FastString -> Kind -> VM Var
newTyVar fs k
= do
u <- liftDs newUnique
return $ mkTyVar (mkSysTvName u fs) k
lookupTyCon :: TyCon -> VM (Maybe TyCon)
lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
extendTyVarPA :: Var -> CoreExpr -> VM ()
extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
-- ----------------------------------------------------------------------------
-- Bindings
vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
vectoriseModule info guts
= do
builtins <- initBuiltins
r <- runVM (vectModule guts) builtins (initGlobalEnv info) emptyLocalEnv
case r of
Yes genv _ guts' -> return $ updVectInfo genv guts'
No -> return guts
vectModule :: ModGuts -> VM ModGuts
vectModule guts = return guts
-- ----------------------------------------------------------------------------
-- Bindings
vectBndr :: Var -> VM (Var, Var)
vectBndr v
......
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