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 -------------------------------------------------
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
81 -- Scalars --------------------------------------------------------------------
82 -- | Get the set of global scalar variables.
83 globalScalars :: VM VarSet
84 globalScalars = readGEnv global_scalars
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)
91 -- TyCons ---------------------------------------------------------------------
92 -- | Lookup the vectorised version of a `TyCon` from the global environment.
93 lookupTyCon :: TyCon -> VM (Maybe TyCon)
95 | isUnLiftedTyCon tc || isTupleTyCon tc
99 = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
102 -- | Lookup the vectorised version of a boxed `TyCon` from the global environment.
103 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
105 = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
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' }
115 -- DataCons -------------------------------------------------------------------
116 -- | Lookup the vectorised version of a `DataCon` from the global environment.
117 lookupDataCon :: DataCon -> VM (Maybe DataCon)
119 | isTupleTyCon (dataConTyCon dc)
123 = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
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' }
132 -- PA dictionaries ------------------------------------------------------------
133 -- | Lookup a PA `TyCon` from the global environment.
134 lookupTyConPA :: TyCon -> VM (Maybe Var)
136 = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
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 }
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] }
152 -- PR Dictionaries ------------------------------------------------------------
153 lookupTyConPR :: TyCon -> VM (Maybe Var)
154 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)