e471ebbc03f7ef702d1204b983b0933a83d67264
[ghc.git] / 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