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         { global_vars           :: VarEnv Var
  81           -- ^Mapping from global variables to their vectorised versions — aka the /vectorisation
  82           -- map/.
  83
  84         , global_vect_decls     :: VarEnv (Type, CoreExpr)
  85           -- ^Mapping from global variables that have a vectorisation declaration to the right-hand
  86           -- side of that declaration and its type.  This mapping only applies to non-scalar
  87           -- vectorisation declarations.  All variables with a scalar vectorisation declaration are
  88           -- mentioned in 'global_scalars_vars'.
  89
  90         , global_scalar_vars    :: VarSet
  91           -- ^Purely scalar variables. Code which mentions only these variables doesn't have to be
  92           -- lifted.  This includes variables from the current module that have a scalar
  93           -- vectorisation declaration and those that the vectoriser determines to be scalar.
  94
  95         , global_scalar_tycons  :: NameSet
  96           -- ^Type constructors whose values can only contain scalar data.  Scalar code may only
  97           -- operate on such data.
  98
  99         , global_novect_vars    :: VarSet
 100           -- ^Variables that are not vectorised.  (They may be referenced in the right-hand sides
 101           -- of vectorisation declarations, though.)
 102
 103         , global_exported_vars  :: VarEnv (Var, Var)
 104           -- ^Exported variables which have a vectorised version.
 105
 106         , global_tycons         :: NameEnv TyCon
 107           -- ^Mapping from TyCons to their vectorised versions.
 108           -- TyCons which do not have to be vectorised are mapped to themselves.
 109
 110         , global_datacons       :: NameEnv DataCon
 111           -- ^Mapping from DataCons to their vectorised versions.
 112
 113         , global_pa_funs        :: NameEnv Var
 114           -- ^Mapping from TyCons to their PA dfuns.
 115
 116         , global_pr_funs        :: NameEnv Var
 117           -- ^Mapping from TyCons to their PR dfuns.
 118
 119         , global_boxed_tycons   :: NameEnv TyCon
 120           -- ^Mapping from unboxed TyCons to their boxed versions.
 121
 122         , global_inst_env       :: (InstEnv, InstEnv)
 123           -- ^External package inst-env & home-package inst-env for class instances.
 124
 125         , global_fam_inst_env   :: FamInstEnvs
 126           -- ^External package inst-env & home-package inst-env for family instances.
 127
 128         , global_bindings       :: [(Var, CoreExpr)]
 129           -- ^Hoisted bindings.
 130         }
 131
 132 -- |Create an initial global environment.
 133 --
 134 initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
 135 initGlobalEnv info vectDecls instEnvs famInstEnvs
 136   = GlobalEnv
 137   { global_vars          = mapVarEnv snd $ vectInfoVar info
 138   , global_vect_decls    = mkVarEnv vects
 139   , global_scalar_vars   = vectInfoScalarVars   info `extendVarSetList` scalars
 140   , global_scalar_tycons = vectInfoScalarTyCons info
 141   , global_novect_vars   = mkVarSet novects
 142   , global_exported_vars = emptyVarEnv
 143   , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
 144   , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
 145   , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
 146   , global_pr_funs       = emptyNameEnv
 147   , global_boxed_tycons  = emptyNameEnv
 148   , global_inst_env      = instEnvs
 149   , global_fam_inst_env  = famInstEnvs
 150   , global_bindings      = []
 151   }
 152   where
 153     vects   = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
 154     scalars = [var                       | Vect var Nothing    <- vectDecls]
 155     novects = [var                       | NoVect var          <- vectDecls]
 156
 157
 158 -- Operators on Global Environments -------------------------------------------
 159
 160 -- |Extend the list of global variables in an environment.
 161 --
 162 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
 163 extendImportedVarsEnv ps genv
 164   = genv { global_vars = extendVarEnvList (global_vars genv) ps }
 165
 166 -- |Extend the set of scalar variables in an environment.
 167 --
 168 extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
 169 extendScalars vs genv
 170   = genv { global_scalar_vars = extendVarSetList (global_scalar_vars genv) vs }
 171
 172 -- |Set the list of type family instances in an environment.
 173 --
 174 setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
 175 setFamEnv l_fam_inst genv
 176   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
 177   where (g_fam_inst, _) = global_fam_inst_env genv
 178
 179 -- |Extend the list of type family instances.
 180 --
 181 extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
 182 extendFamEnv new genv
 183   = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
 184   where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
 185
 186 -- |Extend the list of type constructors in an environment.
 187 --
 188 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 189 extendTyConsEnv ps genv
 190   = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
 191
 192 -- |Extend the list of data constructors in an environment.
 193 --
 194 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
 195 extendDataConsEnv ps genv
 196   = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
 197
 198 -- |Extend the list of PA functions in an environment.
 199 --
 200 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 201 extendPAFunsEnv ps genv
 202   = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
 203
 204 -- |Set the list of PR functions in an environment.
 205 --
 206 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 207 setPRFunsEnv ps genv
 208   = genv { global_pr_funs = mkNameEnv ps }
 209
 210 -- |Set the list of boxed type constructor in an environment.
 211 --
 212 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 213 setBoxedTyConsEnv ps genv
 214   = genv { global_boxed_tycons = mkNameEnv ps }
 215
 216 -- |Compute vectorisation information that goes into 'ModGuts' (and is stored in interface files).
 217 -- The incoming 'vectInfo' is that from the 'HscEnv' and 'EPS'.  The outgoing one contains only the
 218 -- definitions for the currently compiled module.
 219 --
 220 modVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 221 modVectInfo env tyenv info
 222   = info
 223     { vectInfoVar          = global_exported_vars env
 224     , vectInfoTyCon        = mk_env typeEnvTyCons global_tycons
 225     , vectInfoDataCon      = mk_env typeEnvDataCons global_datacons
 226     , vectInfoPADFun       = mk_env typeEnvTyCons global_pa_funs
 227     , vectInfoScalarVars   = global_scalar_vars   env `minusVarSet`  vectInfoScalarVars   info
 228     , vectInfoScalarTyCons = global_scalar_tycons env `minusNameSet` vectInfoScalarTyCons info
 229     }
 230   where
 231     mk_env from_tyenv from_env
 232       = mkNameEnv [(name, (from,to))
 233                   | from     <- from_tyenv tyenv
 234                   , let name =  getName from
 235                   , Just to  <- [lookupNameEnv (from_env env) name]]