compiler/vectorise/Vectorise/Monad/Global.hs

   1
   2 module Vectorise.Monad.Global (
   3         readGEnv,
   4         setGEnv,
   5         updGEnv,
   6
   7   -- * Vars
   8   defGlobalVar,
   9
  10   -- * Vectorisation declarations
  11   lookupVectDecl,
  12
  13   -- * Scalars
  14   globalScalars, isGlobalScalar,
  15
  16         -- * TyCons
  17         lookupTyCon,
  18         lookupBoxedTyCon,
  19         defTyCon,
  20
  21         -- * Datacons
  22         lookupDataCon,
  23         defDataCon,
  24
  25         -- * PA Dictionaries
  26         lookupTyConPA,
  27         defTyConPA,
  28         defTyConPAs,
  29
  30         -- * PR Dictionaries
  31         lookupTyConPR
  32 ) where
  33
  34 import Vectorise.Monad.Base
  35 import Vectorise.Env
  36
  37 import CoreSyn
  38 import Type
  39 import TyCon
  40 import DataCon
  41 import NameEnv
  42 import Var
  43 import VarEnv
  44 import VarSet
  45
  46
  47 -- Global Environment ---------------------------------------------------------
  48 -- | Project something from the global environment.
  49 readGEnv :: (GlobalEnv -> a) -> VM a
  50 readGEnv f      = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
  51
  52
  53 -- | Set the value of the global environment.
  54 setGEnv :: GlobalEnv -> VM ()
  55 setGEnv genv    = VM $ \_ _ lenv -> return (Yes genv lenv ())
  56
  57
  58 -- | Update the global environment using the provided function.
  59 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
  60 updGEnv f       = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
  61
  62
  63 -- Vars -----------------------------------------------------------------------
  64 -- | Add a mapping between a global var and its vectorised version to the state.
  65 defGlobalVar :: Var -> Var -> VM ()
  66 defGlobalVar v v' = updGEnv $ \env ->
  67   env { global_vars = extendVarEnv (global_vars env) v v'
  68       , global_exported_vars = upd (global_exported_vars env)
  69       }
  70   where
  71     upd env | isExportedId v = extendVarEnv env v (v, v')
  72             | otherwise      = env
  73
  74
  75 -- Vectorisation declarations -------------------------------------------------
  76
  77 -- |Check whether a variable has a (non-scalar) vectorisation declaration.
  78 --
  79 lookupVectDecl :: Var -> VM (Maybe (Type, CoreExpr))
  80 lookupVectDecl var = readGEnv $ \env -> lookupVarEnv (global_vect_decls env) var
  81
  82
  83 -- Scalars --------------------------------------------------------------------
  84
  85 -- |Get the set of global scalar variables.
  86 --
  87 globalScalars :: VM VarSet
  88 globalScalars = readGEnv global_scalar_vars
  89
  90 -- |Check whether a given variable is in the set of global scalar variables.
  91 --
  92 isGlobalScalar :: Var -> VM Bool
  93 isGlobalScalar var = readGEnv $ \env -> elemVarSet var (global_scalar_vars env)
  94
  95
  96 -- TyCons ---------------------------------------------------------------------
  97 -- | Lookup the vectorised version of a `TyCon` from the global environment.
  98 lookupTyCon :: TyCon -> VM (Maybe TyCon)
  99 lookupTyCon tc
 100   | isUnLiftedTyCon tc || isTupleTyCon tc
 101   = return (Just tc)
 102
 103   | otherwise
 104   = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 105
 106
 107 -- | Lookup the vectorised version of a boxed `TyCon` from the global environment.
 108 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
 109 lookupBoxedTyCon tc
 110         = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
 111                                            (tyConName tc)
 112
 113
 114 -- | Add a mapping between plain and vectorised `TyCon`s to the global environment.
 115 defTyCon :: TyCon -> TyCon -> VM ()
 116 defTyCon tc tc' = updGEnv $ \env ->
 117   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 118
 119
 120 -- DataCons -------------------------------------------------------------------
 121 -- | Lookup the vectorised version of a `DataCon` from the global environment.
 122 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 123 lookupDataCon dc
 124   | isTupleTyCon (dataConTyCon dc)
 125   = return (Just dc)
 126
 127   | otherwise
 128   = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 129
 130
 131 -- | Add the mapping between plain and vectorised `DataCon`s to the global environment.
 132 defDataCon :: DataCon -> DataCon -> VM ()
 133 defDataCon dc dc' = updGEnv $ \env ->
 134   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 135
 136
 137 -- PA dictionaries ------------------------------------------------------------
 138 -- | Lookup a PA `TyCon` from the global environment.
 139 lookupTyConPA :: TyCon -> VM (Maybe Var)
 140 lookupTyConPA tc
 141         = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 142
 143
 144 -- | Add a mapping between a PA TyCon and is vectorised version to the global environment.
 145 defTyConPA :: TyCon -> Var -> VM ()
 146 defTyConPA tc pa = updGEnv $ \env ->
 147   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 148
 149
 150 -- | Add several mapping between PA TyCons and their vectorised versions to the global environment.
 151 defTyConPAs :: [(TyCon, Var)] -> VM ()
 152 defTyConPAs ps = updGEnv $ \env ->
 153   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 154                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 155
 156
 157 -- PR Dictionaries ------------------------------------------------------------
 158 lookupTyConPR :: TyCon -> VM (Maybe Var)
 159 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 160
 161