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