module Vectorise.Env (
- Scope(..),
-
- -- * Local Environments
- LocalEnv(..),
- emptyLocalEnv,
-
- -- * Global Environments
- GlobalEnv(..),
- initGlobalEnv,
- extendImportedVarsEnv,
- extendScalars,
- setFamInstEnv,
- extendTyConsEnv,
- extendDataConsEnv,
- extendPAFunsEnv,
- setPRFunsEnv,
- setBoxedTyConsEnv,
- updVectInfo
+ Scope(..),
+
+ -- * Local Environments
+ LocalEnv(..),
+ emptyLocalEnv,
+
+ -- * Global Environments
+ GlobalEnv(..),
+ initGlobalEnv,
+ extendImportedVarsEnv,
+ extendScalars,
+ setFamEnv,
+ extendFamEnv,
+ extendTyConsEnv,
+ extendDataConsEnv,
+ extendPAFunsEnv,
+ setPRFunsEnv,
+ setBoxedTyConsEnv,
+ modVectInfo
) where
+
import HscTypes
import InstEnv
import FamInstEnv
import CoreSyn
+import Type
import TyCon
import DataCon
import VarEnv
import VarSet
import Var
+import NameSet
import Name
import NameEnv
import FastString
-- | Indicates what scope something (a variable) is in.
data Scope a b
- = Global a
- | Local b
+ = Global a
+ | Local b
-- LocalEnv -------------------------------------------------------------------
-- GlobalEnv ------------------------------------------------------------------
--- | The global environment.
--- These are things the exist at top-level.
+
+-- |The global environment: entities that exist at top-level.
+--
data GlobalEnv
- = GlobalEnv {
- -- | Mapping from global variables to their vectorised versions.
- global_vars :: VarEnv Var
+ = GlobalEnv
+ -- |Mapping from global variables to their vectorised versions — aka the /vectorisation
+ -- map/.
+ { global_vars :: VarEnv Var
- -- | Purely scalar variables. Code which mentions only these
- -- variables doesn't have to be lifted.
- , global_scalars :: VarSet
+ -- |Mapping from global variables that have a vectorisation declaration to the right-hand
+ -- side of that declaration and its type. This mapping only applies to non-scalar
+ -- vectorisation declarations. All variables with a scalar vectorisation declaration are
+ -- mentioned in 'global_scalars_vars'.
+ , global_vect_decls :: VarEnv (Type, CoreExpr)
- -- | Exported variables which have a vectorised version.
- , global_exported_vars :: VarEnv (Var, Var)
+ -- |Purely scalar variables. Code which mentions only these variables doesn't have to be
+ -- lifted. This includes variables from the current module that have a scalar
+ -- vectorisation declaration and those that the vectoriser determines to be scalar.
+ , global_scalar_vars :: VarSet
- -- | Mapping from TyCons to their vectorised versions.
- -- TyCons which do not have to be vectorised are mapped to themselves.
- , global_tycons :: NameEnv TyCon
+ -- |Type constructors whose values can only contain scalar data. Scalar code may only
+ -- operate on such data.
+ , global_scalar_tycons :: NameSet
- -- | Mapping from DataCons to their vectorised versions.
- , global_datacons :: NameEnv DataCon
+ -- |Exported variables which have a vectorised version.
+ , global_exported_vars :: VarEnv (Var, Var)
- -- | Mapping from TyCons to their PA dfuns.
- , global_pa_funs :: NameEnv 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 PR dfuns.
- , global_pr_funs :: NameEnv Var
+ -- |Mapping from DataCons to their vectorised versions.
+ , global_datacons :: NameEnv DataCon
- -- | Mapping from unboxed TyCons to their boxed versions.
- , global_boxed_tycons :: NameEnv TyCon
+ -- |Mapping from TyCons to their PA dfuns.
+ , global_pa_funs :: NameEnv Var
- -- | External package inst-env & home-package inst-env for class instances.
- , global_inst_env :: (InstEnv, InstEnv)
+ -- |Mapping from TyCons to their PR dfuns.
+ , global_pr_funs :: NameEnv Var
- -- | External package inst-env & home-package inst-env for family instances.
- , global_fam_inst_env :: FamInstEnvs
+ -- |Mapping from unboxed TyCons to their boxed versions.
+ , global_boxed_tycons :: NameEnv TyCon
- -- | Hoisted bindings.
- , global_bindings :: [(Var, CoreExpr)]
- }
+ -- |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
--- | Create an initial global environment
-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 = []
- }
+ -- |Hoisted bindings.
+ , global_bindings :: [(Var, CoreExpr)]
+ }
+-- |Create an initial global environment.
+--
+initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
+initGlobalEnv info vectDecls instEnvs famInstEnvs
+ = GlobalEnv
+ { global_vars = mapVarEnv snd $ vectInfoVar info
+ , global_vect_decls = mkVarEnv vects
+ , global_scalar_vars = vectInfoScalarVars info `extendVarSetList` scalars
+ , global_scalar_tycons = vectInfoScalarTyCons info
+ , 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 = []
+ }
+ where
+ vects = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
+ scalars = [var | Vect var Nothing <- vectDecls]
-- Operators on Global Environments -------------------------------------------
--- | Extend the list of global variables in an environment.
+
+-- |Extend the list of global variables in an environment.
+--
extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
extendImportedVarsEnv ps genv
- = genv { global_vars = extendVarEnvList (global_vars genv) ps }
-
+ = genv { global_vars = extendVarEnvList (global_vars genv) ps }
--- | Extend the set of scalar variables in an environment.
+-- |Extend the set of scalar variables in an environment.
+--
extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
extendScalars vs genv
- = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
+ = genv { global_scalar_vars = extendVarSetList (global_scalar_vars genv) vs }
-
--- | Set the list of type family instances in an environment.
-setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
-setFamInstEnv l_fam_inst genv
+-- |Set the list of type family instances in an environment.
+--
+setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
+setFamEnv l_fam_inst genv
= genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
where (g_fam_inst, _) = global_fam_inst_env genv
+-- |Extend the list of type family instances.
+--
+extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
+extendFamEnv new genv
+ = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
+ where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
--- | Extend the list of type constructors in an environment.
+-- |Extend the list of type constructors in an environment.
+--
extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
extendTyConsEnv ps genv
= genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
-
--- | Extend the list of data constructors in an environment.
+-- |Extend the list of data constructors in an environment.
+--
extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
extendDataConsEnv ps genv
= genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
-
--- | Extend the list of PA functions in an environment.
+-- |Extend the list of PA functions in an environment.
+--
extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
extendPAFunsEnv ps genv
= genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
-
--- | Set the list of PR functions in an environment.
+-- |Set the list of PR functions in an environment.
+--
setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
setPRFunsEnv ps genv
= genv { global_pr_funs = mkNameEnv ps }
-
--- | Set the list of boxed type constructor in an environment.
+-- |Set the list of boxed type constructor in an environment.
+--
setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
setBoxedTyConsEnv ps genv
= genv { global_boxed_tycons = mkNameEnv ps }
-
--- | TODO: What is this for?
-updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
-updVectInfo env tyenv info
+-- |Compute vectorisation information that goes into 'ModGuts' (and is stored in interface files).
+-- The incoming 'vectInfo' is that from the 'HscEnv' and 'EPS'. The outgoing one contains only the
+-- definitions for the currently compiled module.
+--
+modVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
+modVectInfo 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
+ { vectInfoVar = global_exported_vars env
+ , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
+ , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
+ , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
+ , vectInfoScalarVars = global_scalar_vars env `minusVarSet` vectInfoScalarVars info
+ , vectInfoScalarTyCons = global_scalar_tycons env `minusNameSet` vectInfoScalarTyCons info
}
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]]
-
+ = mkNameEnv [(name, (from,to))
+ | from <- from_tyenv tyenv
+ , let name = getName from
+ , Just to <- [lookupNameEnv (from_env env) name]]
projects / ghc-hetmet.git / blobdiff