5639c238e3a985be3784d945048629a84e8ab16e
[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,
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 VarEnv
43 import VarSet
44 import Outputable
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'
71 = do { traceVt "add global var mapping:" (ppr v <+> text "-->" <+> ppr v')
72
73 ; updGEnv $ \env -> env { global_vars = extendVarEnv (global_vars env) v v' }
74 }
75
76
77 -- Vectorisation declarations -------------------------------------------------
78
79 -- |Check whether a variable has a (non-scalar) vectorisation declaration.
80 --
81 lookupVectDecl :: Var -> VM (Maybe (Type, CoreExpr))
82 lookupVectDecl var = readGEnv $ \env -> lookupVarEnv (global_vect_decls env) var
83
84 -- |Check whether a variable has a 'NOVECTORISE' declaration.
85 --
86 noVectDecl :: Var -> VM Bool
87 noVectDecl var = readGEnv $ \env -> elemVarSet var (global_novect_vars env)
88
89
90 -- Scalars --------------------------------------------------------------------
91
92 -- |Get the set of global scalar variables.
93 --
94 globalScalarVars :: VM VarSet
95 globalScalarVars = readGEnv global_scalar_vars
96
97 -- |Check whether a given variable is in the set of global scalar variables.
98 --
99 isGlobalScalar :: Var -> VM Bool
100 isGlobalScalar var = readGEnv $ \env -> var `elemVarSet` global_scalar_vars env
101
102 -- |Get the set of global scalar type constructors including both those scalar type constructors
103 -- declared in an imported module and those declared in the current module.
104 --
105 globalScalarTyCons :: VM NameSet
106 globalScalarTyCons = readGEnv global_scalar_tycons
107
108
109 -- TyCons ---------------------------------------------------------------------
110
111 -- |Lookup the vectorised version of a `TyCon` from the global environment.
112 --
113 lookupTyCon :: TyCon -> VM (Maybe TyCon)
114 lookupTyCon tc
115 | isUnLiftedTyCon tc || isTupleTyCon tc
116 = return (Just tc)
117 | otherwise
118 = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
119
120 -- |Add a mapping between plain and vectorised `TyCon`s to the global environment.
121 --
122 defTyCon :: TyCon -> TyCon -> VM ()
123 defTyCon tc tc' = updGEnv $ \env ->
124 env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
125
126 -- |Get the set of all vectorised type constructors.
127 --
128 globalVectTyCons :: VM (NameEnv TyCon)
129 globalVectTyCons = readGEnv global_tycons
130
131
132 -- DataCons -------------------------------------------------------------------
133
134 -- |Lookup the vectorised version of a `DataCon` from the global environment.
135 --
136 lookupDataCon :: DataCon -> VM (Maybe DataCon)
137 lookupDataCon dc
138 | isTupleTyCon (dataConTyCon dc)
139 = return (Just dc)
140
141 | otherwise
142 = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
143
144 -- |Add the mapping between plain and vectorised `DataCon`s to the global environment.
145 --
146 defDataCon :: DataCon -> DataCon -> VM ()
147 defDataCon dc dc' = updGEnv $ \env ->
148 env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
149
150
151 -- 'PA' dictionaries ------------------------------------------------------------
152
153 -- |Lookup the 'PA' dfun of a vectorised type constructor in the global environment.
154 --
155 lookupTyConPA :: TyCon -> VM (Maybe Var)
156 lookupTyConPA tc
157 = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
158
159 -- |Associate vectorised type constructors with the dfun of their 'PA' instances in the global
160 -- environment.
161 --
162 defTyConPAs :: [(TyCon, Var)] -> VM ()
163 defTyConPAs ps = updGEnv $ \env ->
164 env { global_pa_funs = extendNameEnvList (global_pa_funs env)
165 [(tyConName tc, pa) | (tc, pa) <- ps] }
166
167
168 -- PR Dictionaries ------------------------------------------------------------
169
170 lookupTyConPR :: TyCon -> VM (Maybe Var)
171 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
172
173