From: Roman Leshchinskiy Date: Tue, 10 Jul 2007 06:05:28 +0000 (+0000) Subject: Put vectorisation monad into a separate file X-Git-Url: http://git.megacz.com/?p=ghc-hetmet.git;a=commitdiff_plain;h=621bc50eeb52dad05750d77581c4829e37424741 Put vectorisation monad into a separate file --- diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs new file mode 100644 index 0000000..46204b0 --- /dev/null +++ b/compiler/vectorise/VectMonad.hs @@ -0,0 +1,233 @@ +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 + diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index 0358aca..df9fdb9 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -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