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         updVectInfo
  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 Name
  35 import NameEnv
  36 import FastString
  37
  38
  39 -- | Indicates what scope something (a variable) is in.
  40 data Scope a b
  41         = Global a
  42         | Local  b
  43
  44
  45 -- LocalEnv -------------------------------------------------------------------
  46 -- | The local environment.
  47 data LocalEnv
  48         = LocalEnv {
  49         -- Mapping from local variables to their vectorised and lifted versions.
  50             local_vars          :: VarEnv (Var, Var)
  51
  52         -- In-scope type variables.
  53         , local_tyvars          :: [TyVar]
  54
  55         -- Mapping from tyvars to their PA dictionaries.
  56         , local_tyvar_pa        :: VarEnv CoreExpr
  57
  58         -- Local binding name.
  59         , local_bind_name       :: FastString
  60         }
  61
  62
  63 -- | Create an empty local environment.
  64 emptyLocalEnv :: LocalEnv
  65 emptyLocalEnv = LocalEnv {
  66                    local_vars     = emptyVarEnv
  67                  , local_tyvars   = []
  68                  , local_tyvar_pa = emptyVarEnv
  69                  , local_bind_name  = fsLit "fn"
  70                  }
  71
  72
  73 -- GlobalEnv ------------------------------------------------------------------
  74 -- | The global environment.
  75 --      These are things the exist at top-level.
  76 data GlobalEnv
  77         = GlobalEnv {
  78         -- | Mapping from global variables to their vectorised versions — aka the /vectorisation
  79         --   map/.
  80           global_vars           :: VarEnv Var
  81
  82         -- | Mapping from global variables that have a vectorisation declaration to the right-hand
  83         --   side of that declaration and its type.  This mapping only applies to non-scalar
  84         --   vectorisation declarations.  All variables with a scalar vectorisation declaration are
  85         --   mentioned in 'global_scalars'.
  86         , global_vect_decls     :: VarEnv (Type, CoreExpr)
  87
  88         -- | Purely scalar variables. Code which mentions only these variables doesn't have to be
  89         --   lifted.  This includes variables from the current module that have a scalar
  90         --   vectorisation declaration and those that the vectoriser determines to be scalar.
  91         , global_scalars        :: VarSet
  92
  93         -- | Exported variables which have a vectorised version.
  94         , global_exported_vars  :: VarEnv (Var, Var)
  95
  96         -- | Mapping from TyCons to their vectorised versions.
  97         --   TyCons which do not have to be vectorised are mapped to themselves.
  98         , global_tycons         :: NameEnv TyCon
  99
 100         -- | Mapping from DataCons to their vectorised versions.
 101         , global_datacons       :: NameEnv DataCon
 102
 103         -- | Mapping from TyCons to their PA dfuns.
 104         , global_pa_funs        :: NameEnv Var
 105
 106         -- | Mapping from TyCons to their PR dfuns.
 107         , global_pr_funs        :: NameEnv Var
 108
 109         -- | Mapping from unboxed TyCons to their boxed versions.
 110         , global_boxed_tycons   :: NameEnv TyCon
 111
 112         -- | External package inst-env & home-package inst-env for class instances.
 113         , global_inst_env       :: (InstEnv, InstEnv)
 114
 115         -- | External package inst-env & home-package inst-env for family instances.
 116         , global_fam_inst_env   :: FamInstEnvs
 117
 118         -- | Hoisted bindings.
 119         , global_bindings       :: [(Var, CoreExpr)]
 120         }
 121
 122 -- | Create an initial global environment
 123 initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
 124 initGlobalEnv info vectDecls instEnvs famInstEnvs
 125   = GlobalEnv
 126   { global_vars          = mapVarEnv snd $ vectInfoVar info
 127   , global_vect_decls    = mkVarEnv vects
 128   , global_scalars       = mkVarSet scalars
 129   , global_exported_vars = emptyVarEnv
 130   , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
 131   , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
 132   , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
 133   , global_pr_funs       = emptyNameEnv
 134   , global_boxed_tycons  = emptyNameEnv
 135   , global_inst_env      = instEnvs
 136   , global_fam_inst_env  = famInstEnvs
 137   , global_bindings      = []
 138   }
 139   where
 140     vects   = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
 141     scalars = [var                       | Vect var Nothing    <- vectDecls]
 142
 143
 144 -- Operators on Global Environments -------------------------------------------
 145 -- | Extend the list of global variables in an environment.
 146 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
 147 extendImportedVarsEnv ps genv
 148   = genv { global_vars   = extendVarEnvList (global_vars genv) ps }
 149
 150 -- | Extend the set of scalar variables in an environment.
 151 extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
 152 extendScalars vs genv
 153   = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
 154
 155 -- | Set the list of type family instances in an environment.
 156 setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
 157 setFamEnv l_fam_inst genv
 158   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
 159   where (g_fam_inst, _) = global_fam_inst_env genv
 160
 161 extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
 162 extendFamEnv new genv
 163   = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
 164   where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
 165
 166
 167 -- | Extend the list of type constructors in an environment.
 168 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 169 extendTyConsEnv ps genv
 170   = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
 171
 172
 173 -- | Extend the list of data constructors in an environment.
 174 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
 175 extendDataConsEnv ps genv
 176   = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
 177
 178
 179 -- | Extend the list of PA functions in an environment.
 180 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 181 extendPAFunsEnv ps genv
 182   = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
 183
 184
 185 -- | Set the list of PR functions in an environment.
 186 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 187 setPRFunsEnv ps genv
 188   = genv { global_pr_funs = mkNameEnv ps }
 189
 190
 191 -- | Set the list of boxed type constructor in an environment.
 192 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 193 setBoxedTyConsEnv ps genv
 194   = genv { global_boxed_tycons = mkNameEnv ps }
 195
 196
 197 -- | TODO: What is this for?
 198 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 199 updVectInfo env tyenv info
 200   = info
 201     { vectInfoVar     = global_exported_vars env
 202     , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
 203     , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
 204     , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
 205     }
 206   where
 207     mk_env from_tyenv from_env
 208         = mkNameEnv [(name, (from,to))
 209                         | from     <- from_tyenv tyenv
 210                         , let name =  getName from
 211                         , Just to  <- [lookupNameEnv (from_env env) name]]
 212