+{-# LANGUAGE NamedFieldPuns #-}
+
+-- | The Vectorisation monad.
module VectMonad (
- Scope(..),
VM,
noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
initV, cantVectorise, maybeCantVectorise, maybeCantVectoriseM,
liftDs,
cloneName, cloneId, cloneVar,
- newExportedVar, newLocalVar, newDummyVar, newTyVar,
+ newExportedVar, newLocalVar, newLocalVars, newDummyVar, newTyVar,
- Builtins(..), sumTyCon, prodTyCon,
- combinePAVar, scalarZip, closureCtrFun,
+ Builtins(..), sumTyCon, prodTyCon, prodDataCon,
+ selTy, selReplicate, selPick, selTags, selElements,
+ combinePDVar, scalarZip, closureCtrFun,
builtin, builtins,
- GlobalEnv(..),
setFamInstEnv,
readGEnv, setGEnv, updGEnv,
- LocalEnv(..),
readLEnv, setLEnv, updLEnv,
getBindName, inBind,
#include "HsVersions.h"
import VectBuiltIn
+import Vectorise.Env
import HscTypes hiding ( MonadThings(..) )
import Module ( PackageId )
import Control.Monad
-data Scope a b = Global a | Local b
-
--- ----------------------------------------------------------------------------
--- Vectorisation monad
-
-data GlobalEnv = GlobalEnv {
- -- Mapping from global variables to their vectorised versions.
- --
- global_vars :: VarEnv Var
-
- -- Purely scalar variables. Code which mentions only these
- -- variables doesn't have to be lifted.
- , global_scalars :: VarSet
-
- -- 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 DataCons to their vectorised versions
- --
- , global_datacons :: NameEnv DataCon
-
- -- Mapping from TyCons to their PA dfuns
- --
- , global_pa_funs :: NameEnv Var
-
- -- Mapping from TyCons to their PR dfuns
- , global_pr_funs :: NameEnv Var
-
- -- Mapping from unboxed TyCons to their boxed versions
- , global_boxed_tycons :: NameEnv TyCon
-
- -- External package inst-env & home-package inst-env for class
- -- instances
- --
- , global_inst_env :: (InstEnv, InstEnv)
-
- -- External package inst-env & home-package inst-env for family
- -- instances
- --
- , global_fam_inst_env :: FamInstEnvs
-
- -- Hoisted bindings
- , global_bindings :: [(Var, CoreExpr)]
- }
-
-data LocalEnv = LocalEnv {
- -- Mapping from local variables to their vectorised and
- -- lifted versions
- --
- local_vars :: VarEnv (Var, Var)
-
- -- In-scope type variables
- --
- , local_tyvars :: [TyVar]
-
- -- Mapping from tyvars to their PA dictionaries
- , local_tyvar_pa :: VarEnv CoreExpr
-
- -- Local binding name
- , local_bind_name :: FastString
- }
-
-initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
-initGlobalEnv info instEnvs famInstEnvs
- = GlobalEnv {
- global_vars = mapVarEnv snd $ vectInfoVar info
- , global_scalars = emptyVarSet
- , global_exported_vars = emptyVarEnv
- , global_tycons = mapNameEnv snd $ vectInfoTyCon info
- , global_datacons = mapNameEnv snd $ vectInfoDataCon info
- , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
- , global_pr_funs = emptyNameEnv
- , global_boxed_tycons = emptyNameEnv
- , global_inst_env = instEnvs
- , global_fam_inst_env = famInstEnvs
- , global_bindings = []
- }
-
-extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
-extendImportedVarsEnv ps genv
- = genv { global_vars = extendVarEnvList (global_vars genv) ps }
-
-extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
-extendScalars vs genv
- = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
-
-setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
-setFamInstEnv l_fam_inst genv
- = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
- where
- (g_fam_inst, _) = global_fam_inst_env genv
-
-extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
-extendTyConsEnv ps genv
- = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
-
-extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
-extendDataConsEnv ps genv
- = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
-
-extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
-extendPAFunsEnv ps genv
- = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
-
-setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
-setPRFunsEnv ps genv
- = genv { global_pr_funs = mkNameEnv ps }
-
-setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
-setBoxedTyConsEnv ps genv
- = genv { global_boxed_tycons = mkNameEnv ps }
-
-emptyLocalEnv :: LocalEnv
-emptyLocalEnv = LocalEnv {
- local_vars = emptyVarEnv
- , local_tyvars = []
- , local_tyvar_pa = emptyVarEnv
- , local_bind_name = fsLit "fn"
- }
-
--- FIXME
-updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
-updVectInfo env tyenv info
- = info {
- vectInfoVar = global_exported_vars env
- , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
- , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
- , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
- }
- where
- mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
- | from <- from_tyenv tyenv
- , let name = getName from
- , Just to <- [lookupNameEnv (from_env env) name]]
+-- The Vectorisation Monad ----------------------------------------------------
+
+-- Vectorisation can either succeed with new envionment and a value,
+-- or return with failure.
+--
data VResult a = Yes GlobalEnv LocalEnv a | No
newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
No -> return No
+-- | Throw an error saying we can't vectorise something
cantVectorise :: String -> SDoc -> a
cantVectorise s d = pgmError
. showSDocDump
Just x -> return x
Nothing -> cantVectorise s d
+
+-- Control --------------------------------------------------------------------
+-- | Return some result saying we've failed.
noV :: VM a
noV = VM $ \_ _ _ -> return No
traceNoV :: String -> SDoc -> VM a
traceNoV s d = pprTrace s d noV
+
+-- | If True then carry on, otherwise fail.
ensureV :: Bool -> VM ()
ensureV False = noV
ensureV True = return ()
+
+-- | If True then return the first argument, otherwise fail.
onlyIfV :: Bool -> VM a -> VM a
onlyIfV b p = ensureV b >> p
traceEnsureV s d False = traceNoV s d
traceEnsureV _ _ True = return ()
+
+-- | Try some vectorisation computaton.
+-- If it succeeds then return Just the result,
+-- otherwise return Nothing.
tryV :: VM a -> VM (Maybe a)
tryV (VM p) = VM $ \bi genv lenv ->
do
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
unYes (Yes _ _ x) = x
unYes No = panic "VectMonad.fixV: no result"
+
+-- Local Environments ---------------------------------------------------------
+-- | Perform a computation in its own local environment.
+-- This does not alter the environment of the current state.
localV :: VM a -> VM a
localV p = do
env <- readLEnv id
setLEnv env
return x
+-- | Perform a computation in an empty local environment.
closedV :: VM a -> VM a
closedV p = do
env <- readLEnv id
setLEnv env
return x
+-- Lifting --------------------------------------------------------------------
+-- | Lift a desugaring computation into the vectorisation monad.
liftDs :: DsM a -> VM a
liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
+
+
+-- Builtins -------------------------------------------------------------------
+-- Operations on Builtins
liftBuiltinDs :: (Builtins -> DsM a) -> VM a
liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
+
+-- | Project something from the set of builtins.
builtin :: (Builtins -> a) -> VM a
builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
builtins :: (a -> Builtins -> b) -> VM (a -> b)
builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
+
+-- Environments ---------------------------------------------------------------
+-- | Project something from the global environment.
readGEnv :: (GlobalEnv -> a) -> VM a
readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
+
+-- | Project something from the local environment.
readLEnv :: (LocalEnv -> a) -> VM a
readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
+-- | Set the local environment.
setLEnv :: LocalEnv -> VM ()
setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
+-- | Update the enviroment using a provided function.
updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
+
+-- InstEnv --------------------------------------------------------------------
getInstEnv :: VM (InstEnv, InstEnv)
getInstEnv = readGEnv global_inst_env
getFamInstEnv :: VM FamInstEnvs
getFamInstEnv = readGEnv global_fam_inst_env
+
+-- Names ----------------------------------------------------------------------
+-- | Get the name of the local binding currently being vectorised.
getBindName :: VM FastString
getBindName = readLEnv local_bind_name
| otherwise = Id.mkLocalId name ty
return id'
+-- Make a fresh instance of this var, with a new unique.
cloneVar :: Var -> VM Var
cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
u <- liftDs newUnique
return $ mkSysLocal fs u ty
+newLocalVars :: FastString -> [Type] -> VM [Var]
+newLocalVars fs = mapM (newLocalVar fs)
+
newDummyVar :: Type -> VM Var
newDummyVar = newLocalVar (fsLit "vv")
u <- liftDs newUnique
return $ mkTyVar (mkSysTvName u fs) k
+
+-- | Add a mapping between a global var and its vectorised version to the state.
defGlobalVar :: Var -> Var -> VM ()
defGlobalVar v v' = updGEnv $ \env ->
env { global_vars = extendVarEnv (global_vars env) v v'
upd env | isExportedId v = extendVarEnv env v (v, v')
| otherwise = env
+-- Var ------------------------------------------------------------------------
+-- | Lookup the vectorised and\/or lifted versions of this variable.
+-- If it's in the global environment we get the vectorised version.
+-- If it's in the local environment we get both the vectorised and lifted version.
+--
lookupVar :: Var -> VM (Scope Var (Var, Var))
lookupVar v
- = do
- r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
+ = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
case r of
Just e -> return (Local e)
Nothing -> liftM Global
- . maybeCantVectoriseM "Variable not vectorised:" (ppr v)
+ . maybeCantVectoriseVarM v
. readGEnv $ \env -> lookupVarEnv (global_vars env) v
+maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
+maybeCantVectoriseVarM v p
+ = do r <- p
+ case r of
+ Just x -> return x
+ Nothing -> dumpVar v
+
+dumpVar :: Var -> a
+dumpVar var
+ | Just _ <- isClassOpId_maybe var
+ = cantVectorise "ClassOpId not vectorised:" (ppr var)
+
+ | otherwise
+ = cantVectorise "Variable not vectorised:" (ppr var)
+
+-------------------------------------------------------------------------------
globalScalars :: VM VarSet
globalScalars = readGEnv global_scalars
(ppr $ mkTyConApp tycon tys)
}
+
+-- | Run a vectorisation computation.
initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
initV pkg hsc_env guts info p
= do
builtin_vars <- initBuiltinVars builtins
builtin_tycons <- initBuiltinTyCons builtins
let builtin_datacons = initBuiltinDataCons builtins
- builtin_pas <- initBuiltinPAs builtins
- builtin_prs <- initBuiltinPRs builtins
builtin_boxed <- initBuiltinBoxedTyCons builtins
builtin_scalars <- initBuiltinScalars builtins
let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
instEnvs = (eps_inst_env eps, mg_inst_env guts)
+ builtin_prs <- initBuiltinPRs builtins instEnvs
+ builtin_pas <- initBuiltinPAs builtins instEnvs
+
let genv = extendImportedVarsEnv builtin_vars
. extendScalars builtin_scalars
. extendTyConsEnv builtin_tycons
projects / ghc-hetmet.git / blobdiff