Add VECTORISE [SCALAR] type pragma
[ghc.git] / compiler / vectorise / Vectorise / Monad / Global.hs
1 -- Operations on the global state of the vectorisation monad.
2
3 module Vectorise.Monad.Global (
4 readGEnv,
5 setGEnv,
6 updGEnv,
7
8 -- * Vars
9 defGlobalVar,
10
11 -- * Vectorisation declarations
12 lookupVectDecl, noVectDecl,
13
14 -- * Scalars
15 globalScalarVars, isGlobalScalar, globalScalarTyCons,
16
17 -- * TyCons
18 lookupTyCon, lookupBoxedTyCon,
19 defTyCon, globalVectTyCons,
20
21 -- * Datacons
22 lookupDataCon,
23 defDataCon,
24
25 -- * PA Dictionaries
26 lookupTyConPA,
27 defTyConPAs,
28
29 -- * PR Dictionaries
30 lookupTyConPR
31 ) where
32
33 import Vectorise.Monad.Base
34 import Vectorise.Env
35
36 import CoreSyn
37 import Type
38 import TyCon
39 import DataCon
40 import NameEnv
41 import NameSet
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 globalScalarVars :: VM VarSet
97 globalScalarVars = 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 -> var `elemVarSet` global_scalar_vars env
103
104 -- |Get the set of global scalar type constructors including both those scalar type constructors
105 -- declared in an imported module and those declared in the current module.
106 --
107 globalScalarTyCons :: VM NameSet
108 globalScalarTyCons = readGEnv global_scalar_tycons
109
110
111 -- TyCons ---------------------------------------------------------------------
112
113 -- |Lookup the vectorised version of a `TyCon` from the global environment.
114 --
115 lookupTyCon :: TyCon -> VM (Maybe TyCon)
116 lookupTyCon tc
117 | isUnLiftedTyCon tc || isTupleTyCon tc
118 = return (Just tc)
119 | otherwise
120 = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
121
122 -- |Lookup the vectorised version of a boxed `TyCon` from the global environment.
123 --
124 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
125 lookupBoxedTyCon tc
126 = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
127 (tyConName tc)
128
129 -- |Add a mapping between plain and vectorised `TyCon`s to the global environment.
130 --
131 defTyCon :: TyCon -> TyCon -> VM ()
132 defTyCon tc tc' = updGEnv $ \env ->
133 env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
134
135 -- |Get the set of all vectorised type constructors.
136 --
137 globalVectTyCons :: VM (NameEnv TyCon)
138 globalVectTyCons = readGEnv global_tycons
139
140
141 -- DataCons -------------------------------------------------------------------
142
143 -- |Lookup the vectorised version of a `DataCon` from the global environment.
144 --
145 lookupDataCon :: DataCon -> VM (Maybe DataCon)
146 lookupDataCon dc
147 | isTupleTyCon (dataConTyCon dc)
148 = return (Just dc)
149
150 | otherwise
151 = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
152
153 -- |Add the mapping between plain and vectorised `DataCon`s to the global environment.
154 --
155 defDataCon :: DataCon -> DataCon -> VM ()
156 defDataCon dc dc' = updGEnv $ \env ->
157 env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
158
159
160 -- 'PA' dictionaries ------------------------------------------------------------
161
162 -- |Lookup the 'PA' dfun of a vectorised type constructor in the global environment.
163 --
164 lookupTyConPA :: TyCon -> VM (Maybe Var)
165 lookupTyConPA tc
166 = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
167
168 -- |Associate vectorised type constructors with the dfun of their 'PA' instances in the global
169 -- environment.
170 --
171 defTyConPAs :: [(TyCon, Var)] -> VM ()
172 defTyConPAs ps = updGEnv $ \env ->
173 env { global_pa_funs = extendNameEnvList (global_pa_funs env)
174 [(tyConName tc, pa) | (tc, pa) <- ps] }
175
176
177 -- PR Dictionaries ------------------------------------------------------------
178
179 lookupTyConPR :: TyCon -> VM (Maybe Var)
180 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
181
182