d70f09affdbe22e266fd9479e8ac702100cbf923
[ghc.git] / compiler / vectorise / Vectorise / Env.hs
1
2 module Vectorise.Env (
3 Scope(..),
4
5 -- * Local Environments
6 LocalEnv(..),
7 emptyLocalEnv,
8
9 -- * Global Environments
10 GlobalEnv(..),
11 initGlobalEnv,
12 extendImportedVarsEnv,
13 setFamEnv,
14 extendFamEnv,
15 extendTyConsEnv,
16 extendDataConsEnv,
17 extendPAFunsEnv,
18 setPRFunsEnv,
19 setBoxedTyConsEnv,
20 modVectInfo
21 ) where
22
23 import HscTypes
24 import InstEnv
25 import FamInstEnv
26 import CoreSyn
27 import Type
28 import TyCon
29 import DataCon
30 import VarEnv
31 import VarSet
32 import Var
33 import NameSet
34 import Name
35 import NameEnv
36 import FastString
37
38
39 -- | Indicates what scope something (a variable) is in.
40 data Scope a b
41 = Global a
42 | Local b
43
44
45 -- LocalEnv -------------------------------------------------------------------
46 -- | The local environment.
47 data LocalEnv
48 = LocalEnv {
49 -- Mapping from local variables to their vectorised and lifted versions.
50 local_vars :: VarEnv (Var, Var)
51
52 -- In-scope type variables.
53 , local_tyvars :: [TyVar]
54
55 -- Mapping from tyvars to their PA dictionaries.
56 , local_tyvar_pa :: VarEnv CoreExpr
57
58 -- Local binding name.
59 , local_bind_name :: FastString
60 }
61
62
63 -- | Create an empty local environment.
64 emptyLocalEnv :: LocalEnv
65 emptyLocalEnv = LocalEnv {
66 local_vars = emptyVarEnv
67 , local_tyvars = []
68 , local_tyvar_pa = emptyVarEnv
69 , local_bind_name = fsLit "fn"
70 }
71
72
73 -- GlobalEnv ------------------------------------------------------------------
74
75 -- |The global environment: entities that exist at top-level.
76 --
77 data GlobalEnv
78 = GlobalEnv
79 { global_vars :: VarEnv Var
80 -- ^Mapping from global variables to their vectorised versions — aka the /vectorisation
81 -- map/.
82
83 , global_vect_decls :: VarEnv (Type, CoreExpr)
84 -- ^Mapping from global variables that have a vectorisation declaration to the right-hand
85 -- side of that declaration and its type. This mapping only applies to non-scalar
86 -- vectorisation declarations. All variables with a scalar vectorisation declaration are
87 -- mentioned in 'global_scalars_vars'.
88
89 , global_scalar_vars :: VarSet
90 -- ^Purely scalar variables. Code which mentions only these variables doesn't have to be
91 -- lifted. This includes variables from the current module that have a scalar
92 -- vectorisation declaration and those that the vectoriser determines to be scalar.
93
94 , global_scalar_tycons :: NameSet
95 -- ^Type constructors whose values can only contain scalar data. Scalar code may only
96 -- operate on such data.
97
98 , global_novect_vars :: VarSet
99 -- ^Variables that are not vectorised. (They may be referenced in the right-hand sides
100 -- of vectorisation declarations, though.)
101
102 , global_exported_vars :: VarEnv (Var, Var)
103 -- ^Exported variables which have a vectorised version.
104
105 , global_tycons :: NameEnv TyCon
106 -- ^Mapping from TyCons to their vectorised versions.
107 -- TyCons which do not have to be vectorised are mapped to themselves.
108
109 , global_datacons :: NameEnv DataCon
110 -- ^Mapping from DataCons to their vectorised versions.
111
112 , global_pa_funs :: NameEnv Var
113 -- ^Mapping from TyCons to their PA dfuns.
114
115 , global_pr_funs :: NameEnv Var
116 -- ^Mapping from TyCons to their PR dfuns.
117
118 , global_boxed_tycons :: NameEnv TyCon
119 -- ^Mapping from unboxed TyCons to their boxed versions.
120
121 , global_inst_env :: (InstEnv, InstEnv)
122 -- ^External package inst-env & home-package inst-env for class instances.
123
124 , global_fam_inst_env :: FamInstEnvs
125 -- ^External package inst-env & home-package inst-env for family instances.
126
127 , global_bindings :: [(Var, CoreExpr)]
128 -- ^Hoisted bindings.
129 }
130
131 -- |Create an initial global environment.
132 --
133 initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
134 initGlobalEnv info vectDecls instEnvs famInstEnvs
135 = GlobalEnv
136 { global_vars = mapVarEnv snd $ vectInfoVar info
137 , global_vect_decls = mkVarEnv vects
138 , global_scalar_vars = vectInfoScalarVars info `extendVarSetList` scalars
139 , global_scalar_tycons = vectInfoScalarTyCons info
140 , global_novect_vars = mkVarSet novects
141 , global_exported_vars = emptyVarEnv
142 , global_tycons = mapNameEnv snd $ vectInfoTyCon info
143 , global_datacons = mapNameEnv snd $ vectInfoDataCon info
144 , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
145 , global_pr_funs = emptyNameEnv
146 , global_boxed_tycons = emptyNameEnv
147 , global_inst_env = instEnvs
148 , global_fam_inst_env = famInstEnvs
149 , global_bindings = []
150 }
151 where
152 vects = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
153 scalars = [var | Vect var Nothing <- vectDecls]
154 novects = [var | NoVect var <- vectDecls]
155
156
157 -- Operators on Global Environments -------------------------------------------
158
159 -- |Extend the list of global variables in an environment.
160 --
161 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
162 extendImportedVarsEnv ps genv
163 = genv { global_vars = extendVarEnvList (global_vars genv) ps }
164
165 -- |Set the list of type family instances in an environment.
166 --
167 setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
168 setFamEnv l_fam_inst genv
169 = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
170 where (g_fam_inst, _) = global_fam_inst_env genv
171
172 -- |Extend the list of type family instances.
173 --
174 extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
175 extendFamEnv new genv
176 = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
177 where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
178
179 -- |Extend the list of type constructors in an environment.
180 --
181 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
182 extendTyConsEnv ps genv
183 = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
184
185 -- |Extend the list of data constructors in an environment.
186 --
187 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
188 extendDataConsEnv ps genv
189 = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
190
191 -- |Extend the list of PA functions in an environment.
192 --
193 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
194 extendPAFunsEnv ps genv
195 = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
196
197 -- |Set the list of PR functions in an environment.
198 --
199 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
200 setPRFunsEnv ps genv
201 = genv { global_pr_funs = mkNameEnv ps }
202
203 -- |Set the list of boxed type constructor in an environment.
204 --
205 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
206 setBoxedTyConsEnv ps genv
207 = genv { global_boxed_tycons = mkNameEnv ps }
208
209 -- |Compute vectorisation information that goes into 'ModGuts' (and is stored in interface files).
210 -- The incoming 'vectInfo' is that from the 'HscEnv' and 'EPS'. The outgoing one contains only the
211 -- definitions for the currently compiled module.
212 --
213 modVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
214 modVectInfo env tyenv info
215 = info
216 { vectInfoVar = global_exported_vars env
217 , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
218 , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
219 , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
220 , vectInfoScalarVars = global_scalar_vars env `minusVarSet` vectInfoScalarVars info
221 , vectInfoScalarTyCons = global_scalar_tycons env `minusNameSet` vectInfoScalarTyCons info
222 }
223 where
224 mk_env from_tyenv from_env
225 = mkNameEnv [(name, (from,to))
226 | from <- from_tyenv tyenv
227 , let name = getName from
228 , Just to <- [lookupNameEnv (from_env env) name]]