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, noVectDecl,
  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
  49 -- |Project something from the global environment.
  50 --
  51 readGEnv :: (GlobalEnv -> a) -> VM a
  52 readGEnv f      = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
  53
  54 -- |Set the value of the global environment.
  55 --
  56 setGEnv :: GlobalEnv -> VM ()
  57 setGEnv genv    = VM $ \_ _ lenv -> return (Yes genv lenv ())
  58
  59 -- |Update the global environment using the provided function.
  60 --
  61 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
  62 updGEnv f       = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
  63
  64
  65 -- Vars -----------------------------------------------------------------------
  66
  67 -- |Add a mapping between a global var and its vectorised version to the state.
  68 --
  69 defGlobalVar :: Var -> Var -> VM ()
  70 defGlobalVar v v' = updGEnv $ \env ->
  71   env { global_vars = extendVarEnv (global_vars env) v v'
  72       , global_exported_vars = upd (global_exported_vars env)
  73       }
  74   where
  75     upd env | isExportedId v = extendVarEnv env v (v, v')
  76             | otherwise      = env
  77
  78
  79 -- Vectorisation declarations -------------------------------------------------
  80
  81 -- |Check whether a variable has a (non-scalar) vectorisation declaration.
  82 --
  83 lookupVectDecl :: Var -> VM (Maybe (Type, CoreExpr))
  84 lookupVectDecl var = readGEnv $ \env -> lookupVarEnv (global_vect_decls env) var
  85
  86 -- |Check whether a variable has a 'NOVECTORISE' declaration.
  87 --
  88 noVectDecl :: Var -> VM Bool
  89 noVectDecl var = readGEnv $ \env -> elemVarSet var (global_novect_vars env)
  90
  91
  92 -- Scalars --------------------------------------------------------------------
  93
  94 -- |Get the set of global scalar variables.
  95 --
  96 globalScalars :: VM VarSet
  97 globalScalars = readGEnv global_scalar_vars
  98
  99 -- |Check whether a given variable is in the set of global scalar variables.
 100 --
 101 isGlobalScalar :: Var -> VM Bool
 102 isGlobalScalar var = readGEnv $ \env -> elemVarSet var (global_scalar_vars env)
 103
 104
 105 -- TyCons ---------------------------------------------------------------------
 106
 107 -- |Lookup the vectorised version of a `TyCon` from the global environment.
 108 --
 109 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 110 lookupTyCon tc
 111   | isUnLiftedTyCon tc || isTupleTyCon tc
 112   = return (Just tc)
 113
 114   | otherwise
 115   = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 116
 117 -- | Lookup the vectorised version of a boxed `TyCon` from the global environment.
 118 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
 119 lookupBoxedTyCon tc
 120         = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
 121                                            (tyConName tc)
 122
 123 -- | Add a mapping between plain and vectorised `TyCon`s to the global environment.
 124 defTyCon :: TyCon -> TyCon -> VM ()
 125 defTyCon tc tc' = updGEnv $ \env ->
 126   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 127
 128
 129 -- DataCons -------------------------------------------------------------------
 130
 131 -- | Lookup the vectorised version of a `DataCon` from the global environment.
 132 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 133 lookupDataCon dc
 134   | isTupleTyCon (dataConTyCon dc)
 135   = return (Just dc)
 136
 137   | otherwise
 138   = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 139
 140
 141 -- | Add the mapping between plain and vectorised `DataCon`s to the global environment.
 142 defDataCon :: DataCon -> DataCon -> VM ()
 143 defDataCon dc dc' = updGEnv $ \env ->
 144   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 145
 146
 147 -- PA dictionaries ------------------------------------------------------------
 148 -- | Lookup a PA `TyCon` from the global environment.
 149 lookupTyConPA :: TyCon -> VM (Maybe Var)
 150 lookupTyConPA tc
 151         = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 152
 153
 154 -- | Add a mapping between a PA TyCon and is vectorised version to the global environment.
 155 defTyConPA :: TyCon -> Var -> VM ()
 156 defTyConPA tc pa = updGEnv $ \env ->
 157   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 158
 159
 160 -- | Add several mapping between PA TyCons and their vectorised versions to the global environment.
 161 defTyConPAs :: [(TyCon, Var)] -> VM ()
 162 defTyConPAs ps = updGEnv $ \env ->
 163   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 164                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 165
 166
 167 -- PR Dictionaries ------------------------------------------------------------
 168 lookupTyConPR :: TyCon -> VM (Maybe Var)
 169 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 170
 171