2 module Vectorise.Monad.Global (
10 -- * Vectorisation declarations
14 globalScalars, isGlobalScalar,
34 import Vectorise.Monad.Base
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))
53 -- | Set the value of the global environment.
54 setGEnv :: GlobalEnv -> VM ()
55 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
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 ())
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)
71 upd env | isExportedId v = extendVarEnv env v (v, v')
75 -- Vectorisation declarations -------------------------------------------------
77 -- |Check whether a variable has a (non-scalar) vectorisation declaration.
79 lookupVectDecl :: Var -> VM (Maybe (Type, CoreExpr))
80 lookupVectDecl var = readGEnv $ \env -> lookupVarEnv (global_vect_decls env) var
83 -- Scalars --------------------------------------------------------------------
85 -- |Get the set of global scalar variables.
87 globalScalars :: VM VarSet
88 globalScalars = readGEnv global_scalar_vars
90 -- |Check whether a given variable is in the set of global scalar variables.
92 isGlobalScalar :: Var -> VM Bool
93 isGlobalScalar var = readGEnv $ \env -> elemVarSet var (global_scalar_vars env)
96 -- TyCons ---------------------------------------------------------------------
97 -- | Lookup the vectorised version of a `TyCon` from the global environment.
98 lookupTyCon :: TyCon -> VM (Maybe TyCon)
100 | isUnLiftedTyCon tc || isTupleTyCon tc
104 = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
107 -- | Lookup the vectorised version of a boxed `TyCon` from the global environment.
108 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
110 = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
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' }
120 -- DataCons -------------------------------------------------------------------
121 -- | Lookup the vectorised version of a `DataCon` from the global environment.
122 lookupDataCon :: DataCon -> VM (Maybe DataCon)
124 | isTupleTyCon (dataConTyCon dc)
128 = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
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' }
137 -- PA dictionaries ------------------------------------------------------------
138 -- | Lookup a PA `TyCon` from the global environment.
139 lookupTyConPA :: TyCon -> VM (Maybe Var)
141 = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
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 }
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] }
157 -- PR Dictionaries ------------------------------------------------------------
158 lookupTyConPR :: TyCon -> VM (Maybe Var)
159 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)