Add VECTORISE [SCALAR] type pragma
[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 and that appear in a
96 -- 'VECTORISE SCALAR type' pragma in the current or an imported module. Scalar code may
97 -- only operate on such data.
98
99 , global_novect_vars :: VarSet
100 -- ^Variables that are not vectorised. (They may be referenced in the right-hand sides
101 -- of vectorisation declarations, though.)
102
103 , global_exported_vars :: VarEnv (Var, Var)
104 -- ^Exported variables which have a vectorised version.
105
106 , global_tycons :: NameEnv TyCon
107 -- ^Mapping from TyCons to their vectorised versions.
108 -- TyCons which do not have to be vectorised are mapped to themselves.
109
110 , global_datacons :: NameEnv DataCon
111 -- ^Mapping from DataCons to their vectorised versions.
112
113 , global_pa_funs :: NameEnv Var
114 -- ^Mapping from TyCons to their PA dfuns.
115
116 , global_pr_funs :: NameEnv Var
117 -- ^Mapping from TyCons to their PR dfuns.
118
119 , global_boxed_tycons :: NameEnv TyCon
120 -- ^Mapping from unboxed TyCons to their boxed versions.
121
122 , global_inst_env :: (InstEnv, InstEnv)
123 -- ^External package inst-env & home-package inst-env for class instances.
124
125 , global_fam_inst_env :: FamInstEnvs
126 -- ^External package inst-env & home-package inst-env for family instances.
127
128 , global_bindings :: [(Var, CoreExpr)]
129 -- ^Hoisted bindings.
130 }
131
132 -- |Create an initial global environment.
133 --
134 initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
135 initGlobalEnv info vectDecls instEnvs famInstEnvs
136 = GlobalEnv
137 { global_vars = mapVarEnv snd $ vectInfoVar info
138 , global_vect_decls = mkVarEnv vects
139 , global_scalar_vars = vectInfoScalarVars info `extendVarSetList` scalar_vars
140 , global_scalar_tycons = vectInfoScalarTyCons info `addListToNameSet` scalar_tycons
141 , global_novect_vars = mkVarSet novects
142 , global_exported_vars = emptyVarEnv
143 , global_tycons = mapNameEnv snd $ vectInfoTyCon info
144 , global_datacons = mapNameEnv snd $ vectInfoDataCon info
145 , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
146 , global_pr_funs = emptyNameEnv
147 , global_boxed_tycons = emptyNameEnv
148 , global_inst_env = instEnvs
149 , global_fam_inst_env = famInstEnvs
150 , global_bindings = []
151 }
152 where
153 vects = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
154 scalar_vars = [var | Vect var Nothing <- vectDecls]
155 novects = [var | NoVect var <- vectDecls]
156 scalar_tycons = [tyConName tycon | VectType tycon Nothing <- vectDecls]
157
158
159 -- Operators on Global Environments -------------------------------------------
160
161 -- |Extend the list of global variables in an environment.
162 --
163 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
164 extendImportedVarsEnv ps genv
165 = genv { global_vars = extendVarEnvList (global_vars genv) ps }
166
167 -- |Set the list of type family instances in an environment.
168 --
169 setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
170 setFamEnv l_fam_inst genv
171 = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
172 where (g_fam_inst, _) = global_fam_inst_env genv
173
174 -- |Extend the list of type family instances.
175 --
176 extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
177 extendFamEnv new genv
178 = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
179 where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
180
181 -- |Extend the list of type constructors in an environment.
182 --
183 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
184 extendTyConsEnv ps genv
185 = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
186
187 -- |Extend the list of data constructors in an environment.
188 --
189 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
190 extendDataConsEnv ps genv
191 = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
192
193 -- |Extend the list of PA functions in an environment.
194 --
195 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
196 extendPAFunsEnv ps genv
197 = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
198
199 -- |Set the list of PR functions in an environment.
200 --
201 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
202 setPRFunsEnv ps genv
203 = genv { global_pr_funs = mkNameEnv ps }
204
205 -- |Set the list of boxed type constructor in an environment.
206 --
207 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
208 setBoxedTyConsEnv ps genv
209 = genv { global_boxed_tycons = mkNameEnv ps }
210
211 -- |Compute vectorisation information that goes into 'ModGuts' (and is stored in interface files).
212 -- The incoming 'vectInfo' is that from the 'HscEnv' and 'EPS'. The outgoing one contains only the
213 -- definitions for the currently compiled module.
214 --
215 modVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
216 modVectInfo env tyenv info
217 = info
218 { vectInfoVar = global_exported_vars env
219 , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
220 , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
221 , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
222 , vectInfoScalarVars = global_scalar_vars env `minusVarSet` vectInfoScalarVars info
223 , vectInfoScalarTyCons = global_scalar_tycons env `minusNameSet` vectInfoScalarTyCons info
224 }
225 where
226 mk_env from_tyenv from_env
227 = mkNameEnv [(name, (from,to))
228 | from <- from_tyenv tyenv
229 , let name = getName from
230 , Just to <- [lookupNameEnv (from_env env) name]]