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 -- | Check whether a variable has a (non-scalar) vectorisation declaration.
  77 lookupVectDecl :: Var -> VM (Maybe (Type, CoreExpr))
  78 lookupVectDecl var = readGEnv $ \env -> lookupVarEnv (global_vect_decls env) var
  79
  80
  81 -- Scalars --------------------------------------------------------------------
  82 -- | Get the set of global scalar variables.
  83 globalScalars :: VM VarSet
  84 globalScalars = readGEnv global_scalars
  85
  86 -- | Check whether a given variable is in the set of global scalar variables.
  87 isGlobalScalar :: Var -> VM Bool
  88 isGlobalScalar var = readGEnv $ \env -> elemVarSet var (global_scalars env)
  89
  90
  91 -- TyCons ---------------------------------------------------------------------
  92 -- | Lookup the vectorised version of a `TyCon` from the global environment.
  93 lookupTyCon :: TyCon -> VM (Maybe TyCon)
  94 lookupTyCon tc
  95   | isUnLiftedTyCon tc || isTupleTyCon tc
  96   = return (Just tc)
  97
  98   | otherwise
  99   = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 100
 101
 102 -- | Lookup the vectorised version of a boxed `TyCon` from the global environment.
 103 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
 104 lookupBoxedTyCon tc
 105         = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
 106                                            (tyConName tc)
 107
 108
 109 -- | Add a mapping between plain and vectorised `TyCon`s to the global environment.
 110 defTyCon :: TyCon -> TyCon -> VM ()
 111 defTyCon tc tc' = updGEnv $ \env ->
 112   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 113
 114
 115 -- DataCons -------------------------------------------------------------------
 116 -- | Lookup the vectorised version of a `DataCon` from the global environment.
 117 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 118 lookupDataCon dc
 119   | isTupleTyCon (dataConTyCon dc)
 120   = return (Just dc)
 121
 122   | otherwise
 123   = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 124
 125
 126 -- | Add the mapping between plain and vectorised `DataCon`s to the global environment.
 127 defDataCon :: DataCon -> DataCon -> VM ()
 128 defDataCon dc dc' = updGEnv $ \env ->
 129   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 130
 131
 132 -- PA dictionaries ------------------------------------------------------------
 133 -- | Lookup a PA `TyCon` from the global environment.
 134 lookupTyConPA :: TyCon -> VM (Maybe Var)
 135 lookupTyConPA tc
 136         = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 137
 138
 139 -- | Add a mapping between a PA TyCon and is vectorised version to the global environment.
 140 defTyConPA :: TyCon -> Var -> VM ()
 141 defTyConPA tc pa = updGEnv $ \env ->
 142   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 143
 144
 145 -- | Add several mapping between PA TyCons and their vectorised versions to the global environment.
 146 defTyConPAs :: [(TyCon, Var)] -> VM ()
 147 defTyConPAs ps = updGEnv $ \env ->
 148   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 149                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 150
 151
 152 -- PR Dictionaries ------------------------------------------------------------
 153 lookupTyConPR :: TyCon -> VM (Maybe Var)
 154 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 155
 156