compiler/vectorise/Vectorise/Env.hs

   1
   2 module Vectorise.Env (
   3   Scope(..),
   4
   5   -- * Local Environments
   6   LocalEnv(..),
   7   emptyLocalEnv,
   8
   9   -- * Global Environments
  10   GlobalEnv(..),
  11   initGlobalEnv,
  12   extendImportedVarsEnv,
  13   extendScalars,
  14   setFamEnv,
  15   extendFamEnv,
  16   extendTyConsEnv,
  17   extendDataConsEnv,
  18   extendPAFunsEnv,
  19   setPRFunsEnv,
  20   setBoxedTyConsEnv,
  21   modVectInfo
  22 ) where
  23
  24 import HscTypes
  25 import InstEnv
  26 import FamInstEnv
  27 import CoreSyn
  28 import Type
  29 import TyCon
  30 import DataCon
  31 import VarEnv
  32 import VarSet
  33 import Var
  34 import NameSet
  35 import Name
  36 import NameEnv
  37 import FastString
  38
  39
  40 -- | Indicates what scope something (a variable) is in.
  41 data Scope a b
  42         = Global a
  43         | Local  b
  44
  45
  46 -- LocalEnv -------------------------------------------------------------------
  47 -- | The local environment.
  48 data LocalEnv
  49         = LocalEnv {
  50         -- Mapping from local variables to their vectorised and lifted versions.
  51             local_vars          :: VarEnv (Var, Var)
  52
  53         -- In-scope type variables.
  54         , local_tyvars          :: [TyVar]
  55
  56         -- Mapping from tyvars to their PA dictionaries.
  57         , local_tyvar_pa        :: VarEnv CoreExpr
  58
  59         -- Local binding name.
  60         , local_bind_name       :: FastString
  61         }
  62
  63
  64 -- | Create an empty local environment.
  65 emptyLocalEnv :: LocalEnv
  66 emptyLocalEnv = LocalEnv {
  67                    local_vars     = emptyVarEnv
  68                  , local_tyvars   = []
  69                  , local_tyvar_pa = emptyVarEnv
  70                  , local_bind_name  = fsLit "fn"
  71                  }
  72
  73
  74 -- GlobalEnv ------------------------------------------------------------------
  75
  76 -- |The global environment: entities that exist at top-level.
  77 --
  78 data GlobalEnv
  79         = GlobalEnv
  80         -- |Mapping from global variables to their vectorised versions — aka the /vectorisation
  81         --  map/.
  82         { global_vars           :: VarEnv Var
  83
  84         -- |Mapping from global variables that have a vectorisation declaration to the right-hand
  85         --  side of that declaration and its type.  This mapping only applies to non-scalar
  86         --  vectorisation declarations.  All variables with a scalar vectorisation declaration are
  87         --  mentioned in 'global_scalars_vars'.
  88         , global_vect_decls     :: VarEnv (Type, CoreExpr)
  89
  90         -- |Purely scalar variables. Code which mentions only these variables doesn't have to be
  91         --  lifted.  This includes variables from the current module that have a scalar
  92         --  vectorisation declaration and those that the vectoriser determines to be scalar.
  93         , global_scalar_vars    :: VarSet
  94
  95         -- |Type constructors whose values can only contain scalar data.  Scalar code may only
  96         -- operate on such data.
  97         , global_scalar_tycons  :: NameSet
  98
  99         -- |Exported variables which have a vectorised version.
 100         , global_exported_vars  :: VarEnv (Var, Var)
 101
 102         -- |Mapping from TyCons to their vectorised versions.
 103         --  TyCons which do not have to be vectorised are mapped to themselves.
 104         , global_tycons         :: NameEnv TyCon
 105
 106         -- |Mapping from DataCons to their vectorised versions.
 107         , global_datacons       :: NameEnv DataCon
 108
 109         -- |Mapping from TyCons to their PA dfuns.
 110         , global_pa_funs        :: NameEnv Var
 111
 112         -- |Mapping from TyCons to their PR dfuns.
 113         , global_pr_funs        :: NameEnv Var
 114
 115         -- |Mapping from unboxed TyCons to their boxed versions.
 116         , global_boxed_tycons   :: NameEnv TyCon
 117
 118         -- |External package inst-env & home-package inst-env for class instances.
 119         , global_inst_env       :: (InstEnv, InstEnv)
 120
 121         -- |External package inst-env & home-package inst-env for family instances.
 122         , global_fam_inst_env   :: FamInstEnvs
 123
 124         -- |Hoisted bindings.
 125         , global_bindings       :: [(Var, CoreExpr)]
 126         }
 127
 128 -- |Create an initial global environment.
 129 --
 130 initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
 131 initGlobalEnv info vectDecls instEnvs famInstEnvs
 132   = GlobalEnv
 133   { global_vars          = mapVarEnv snd $ vectInfoVar info
 134   , global_vect_decls    = mkVarEnv vects
 135   , global_scalar_vars   = vectInfoScalarVars   info `extendVarSetList` scalars
 136   , global_scalar_tycons = vectInfoScalarTyCons info
 137   , global_exported_vars = emptyVarEnv
 138   , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
 139   , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
 140   , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
 141   , global_pr_funs       = emptyNameEnv
 142   , global_boxed_tycons  = emptyNameEnv
 143   , global_inst_env      = instEnvs
 144   , global_fam_inst_env  = famInstEnvs
 145   , global_bindings      = []
 146   }
 147   where
 148     vects   = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
 149     scalars = [var                       | Vect var Nothing    <- vectDecls]
 150
 151
 152 -- Operators on Global Environments -------------------------------------------
 153
 154 -- |Extend the list of global variables in an environment.
 155 --
 156 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
 157 extendImportedVarsEnv ps genv
 158   = genv { global_vars = extendVarEnvList (global_vars genv) ps }
 159
 160 -- |Extend the set of scalar variables in an environment.
 161 --
 162 extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
 163 extendScalars vs genv
 164   = genv { global_scalar_vars = extendVarSetList (global_scalar_vars genv) vs }
 165
 166 -- |Set the list of type family instances in an environment.
 167 --
 168 setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
 169 setFamEnv l_fam_inst genv
 170   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
 171   where (g_fam_inst, _) = global_fam_inst_env genv
 172
 173 -- |Extend the list of type family instances.
 174 --
 175 extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
 176 extendFamEnv new genv
 177   = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
 178   where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
 179
 180 -- |Extend the list of type constructors in an environment.
 181 --
 182 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 183 extendTyConsEnv ps genv
 184   = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
 185
 186 -- |Extend the list of data constructors in an environment.
 187 --
 188 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
 189 extendDataConsEnv ps genv
 190   = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
 191
 192 -- |Extend the list of PA functions in an environment.
 193 --
 194 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 195 extendPAFunsEnv ps genv
 196   = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
 197
 198 -- |Set the list of PR functions in an environment.
 199 --
 200 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 201 setPRFunsEnv ps genv
 202   = genv { global_pr_funs = mkNameEnv ps }
 203
 204 -- |Set the list of boxed type constructor in an environment.
 205 --
 206 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 207 setBoxedTyConsEnv ps genv
 208   = genv { global_boxed_tycons = mkNameEnv ps }
 209
 210 -- |Compute vectorisation information that goes into 'ModGuts' (and is stored in interface files).
 211 -- The incoming 'vectInfo' is that from the 'HscEnv' and 'EPS'.  The outgoing one contains only the
 212 -- definitions for the currently compiled module.
 213 --
 214 modVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 215 modVectInfo env tyenv info
 216   = info
 217     { vectInfoVar          = global_exported_vars env
 218     , vectInfoTyCon        = mk_env typeEnvTyCons global_tycons
 219     , vectInfoDataCon      = mk_env typeEnvDataCons global_datacons
 220     , vectInfoPADFun       = mk_env typeEnvTyCons global_pa_funs
 221     , vectInfoScalarVars   = global_scalar_vars   env `minusVarSet`  vectInfoScalarVars   info
 222     , vectInfoScalarTyCons = global_scalar_tycons env `minusNameSet` vectInfoScalarTyCons info
 223     }
 224   where
 225     mk_env from_tyenv from_env
 226       = mkNameEnv [(name, (from,to))
 227                   | from     <- from_tyenv tyenv
 228                   , let name =  getName from
 229                   , Just to  <- [lookupNameEnv (from_env env) name]]