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