2 module Vectorise.Monad.Global (
10 -- * Vectorisation declarations
11 lookupVectDecl, noVectDecl,
14 globalScalars, isGlobalScalar,
34 import Vectorise.Monad.Base
47 -- Global Environment ---------------------------------------------------------
49 -- |Project something from the global environment.
51 readGEnv :: (GlobalEnv -> a) -> VM a
52 readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
54 -- |Set the value of the global environment.
56 setGEnv :: GlobalEnv -> VM ()
57 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
59 -- |Update the global environment using the provided function.
61 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
62 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
65 -- Vars -----------------------------------------------------------------------
67 -- |Add a mapping between a global var and its vectorised version to the state.
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)
75 upd env | isExportedId v = extendVarEnv env v (v, v')
79 -- Vectorisation declarations -------------------------------------------------
81 -- |Check whether a variable has a (non-scalar) vectorisation declaration.
83 lookupVectDecl :: Var -> VM (Maybe (Type, CoreExpr))
84 lookupVectDecl var = readGEnv $ \env -> lookupVarEnv (global_vect_decls env) var
86 -- |Check whether a variable has a 'NOVECTORISE' declaration.
88 noVectDecl :: Var -> VM Bool
89 noVectDecl var = readGEnv $ \env -> elemVarSet var (global_novect_vars env)
92 -- Scalars --------------------------------------------------------------------
94 -- |Get the set of global scalar variables.
96 globalScalars :: VM VarSet
97 globalScalars = readGEnv global_scalar_vars
99 -- |Check whether a given variable is in the set of global scalar variables.
101 isGlobalScalar :: Var -> VM Bool
102 isGlobalScalar var = readGEnv $ \env -> elemVarSet var (global_scalar_vars env)
105 -- TyCons ---------------------------------------------------------------------
107 -- |Lookup the vectorised version of a `TyCon` from the global environment.
109 lookupTyCon :: TyCon -> VM (Maybe TyCon)
111 | isUnLiftedTyCon tc || isTupleTyCon tc
115 = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
117 -- | Lookup the vectorised version of a boxed `TyCon` from the global environment.
118 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
120 = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
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' }
129 -- DataCons -------------------------------------------------------------------
131 -- | Lookup the vectorised version of a `DataCon` from the global environment.
132 lookupDataCon :: DataCon -> VM (Maybe DataCon)
134 | isTupleTyCon (dataConTyCon dc)
138 = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
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' }
147 -- PA dictionaries ------------------------------------------------------------
148 -- | Lookup a PA `TyCon` from the global environment.
149 lookupTyConPA :: TyCon -> VM (Maybe Var)
151 = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
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 }
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] }
167 -- PR Dictionaries ------------------------------------------------------------
168 lookupTyConPR :: TyCon -> VM (Maybe Var)
169 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)