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