VECTORISE pragmas for type classes and instances
[ghc.git] / compiler / vectorise / Vectorise / Type / Env.hs
1 {-# OPTIONS_GHC -XNoMonoLocalBinds #-}
2
3 -- Vectorise a modules type environment, the structure containing all type things defined in a
4 -- module.
5 --
6 -- This extends the type environment with vectorised variants of data types and produces value
7 -- bindings for worker functions and the like.
8
9 module Vectorise.Type.Env (
10 vectTypeEnv,
11 ) where
12
13 #include "HsVersions.h"
14
15 import Vectorise.Env
16 import Vectorise.Vect
17 import Vectorise.Monad
18 import Vectorise.Builtins
19 import Vectorise.Type.TyConDecl
20 import Vectorise.Type.Classify
21 import Vectorise.Type.PADict
22 import Vectorise.Type.PData
23 import Vectorise.Type.PRepr
24 import Vectorise.Type.Repr
25 import Vectorise.Utils
26
27 import CoreSyn
28 import CoreUtils
29 import CoreUnfold
30 import DataCon
31 import TyCon
32 import Type
33 import FamInstEnv
34 import Id
35 import MkId
36 import NameEnv
37 import NameSet
38
39 import Util
40 import Outputable
41 import FastString
42 import MonadUtils
43 import Control.Monad
44 import Data.List
45
46
47 -- Note [Pragmas to vectorise tycons]
48 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
49 --
50 -- VECTORISE pragmas for type constructors cover three different flavours of vectorising data type
51 -- constructors:
52 --
53 -- (1) Data type constructor 'T' that may be used in vectorised code, where 'T' represents itself,
54 -- but the representation of 'T' is opaque in vectorised code.
55 --
56 -- An example is the treatment of 'Int'. 'Int's can be used in vectorised code and remain
57 -- unchanged by vectorisation. However, the representation of 'Int' by the 'I#' data
58 -- constructor wrapping an 'Int#' is not exposed in vectorised code. Instead, computations
59 -- involving the representation need to be confined to scalar code.
60 --
61 -- 'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
62 -- by the vectoriser).
63 --
64 -- Type constructors declared with {-# VECTORISE SCALAR type T #-} are treated in this manner.
65 -- (The vectoriser never treats a type constructor automatically in this manner.)
66 --
67 -- (2) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
68 -- code, where 'T' and the 'Cn' are automatically vectorised in the same manner as data types
69 -- declared in a vectorised module. This includes the case where the vectoriser determines that
70 -- the original representation of 'T' may be used in vectorised code (as it does not embed any
71 -- parallel arrays.) This case is for type constructors that are *imported* from a non-
72 -- vectorised module, but that we want to use with full vectorisation support.
73 --
74 -- An example is the treatment of 'Ordering' and '[]'. The former remains unchanged by
75 -- vectorisation, whereas the latter is fully vectorised.
76
77 -- 'PData' and 'PRepr' instances are automatically generated by the vectoriser.
78 --
79 -- Type constructors declared with {-# VECTORISE type T #-} are treated in this manner.
80 --
81 -- (3) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
82 -- code, where 'T' is represented by an explicitly given 'Tv' whose constructors 'Cvn' represent
83 -- the original constructors in vectorised code. As a special case, we can have 'Tv = T'
84 --
85 -- An example is the treatment of 'Bool', which is represented by itself in vectorised code
86 -- (as it cannot embed any parallel arrays). However, we do not want any automatic generation
87 -- of class and family instances, which is why Case (2) does not apply.
88 --
89 -- 'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
90 -- by the vectoriser).
91 --
92 -- Type constructors declared with {-# VECTORISE type T = T' #-} are treated in this manner.
93 --
94 -- In addition, we have also got a single pragma form for type classes: {-# VECTORISE class C #-}. It
95 -- implies that the class type constructor may be used in vectorised code together with its data
96 -- constructor. We generally produce a vectorised version of the data type and data constructor.
97 -- We do not generate 'PData' and 'PRepr' instances for class type constructors.
98
99 -- |Vectorise a type environment.
100 --
101 vectTypeEnv :: [TyCon] -- TyCons defined in this module
102 -> [CoreVect] -- All 'VECTORISE [SCALAR] type' declarations in this module
103 -> VM ( [TyCon] -- old TyCons ++ new TyCons
104 , [FamInst] -- New type family instances.
105 , [(Var, CoreExpr)]) -- New top level bindings.
106 vectTypeEnv tycons vectTypeDecls
107 = do { traceVt "** vectTypeEnv" $ ppr tycons
108
109 -- Build a map containing all vectorised type constructor. If they are scalar, they are
110 -- mapped to 'False' (vectorised type constructor == original type constructor).
111 ; allScalarTyConNames <- globalScalarTyCons -- covers both current and imported modules
112 ; vectTyCons <- globalVectTyCons
113 ; let vectTyConBase = mapNameEnv (const True) vectTyCons -- by default fully vectorised
114 vectTyConFlavour = foldNameSet (\n env -> extendNameEnv env n False) vectTyConBase
115 allScalarTyConNames
116
117 ; let -- {-# VECTORISE SCALAR type T -#} (imported and local tycons)
118 localScalarTyCons = [tycon | VectType True tycon Nothing <- vectTypeDecls]
119
120 -- {-# VECTORISE type T -#} (ONLY the imported tycons)
121 impVectTyCons = [tycon | VectType False tycon Nothing <- vectTypeDecls]
122 \\ tycons
123
124 -- {-# VECTORISE type T = ty -#} (imported and local tycons)
125 vectTyConsWithRHS = [ (tycon, rhs)
126 | VectType False tycon (Just rhs) <- vectTypeDecls]
127
128 -- filter VECTORISE SCALAR tycons and VECTORISE tycons with explicit rhses
129 vectSpecialTyConNames = mkNameSet . map tyConName $
130 localScalarTyCons ++ map fst vectTyConsWithRHS
131 notLocalScalarTyCon tc = not $ (tyConName tc) `elemNameSet` vectSpecialTyConNames
132
133 -- Split the list of 'TyCons' into the ones (1) that we must vectorise and those (2)
134 -- that we could, but don't need to vectorise. Type constructors that are not data
135 -- type constructors or use non-Haskell98 features are being dropped. They may not
136 -- appear in vectorised code. (We also drop the local type constructors appearing in a
137 -- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as
138 -- these are being handled separately.)
139 ; let maybeVectoriseTyCons = filter notLocalScalarTyCon tycons ++ impVectTyCons
140 (conv_tcs, keep_tcs) = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
141 orig_tcs = keep_tcs ++ conv_tcs
142
143 ; traceVt " VECT SCALAR : " $ ppr localScalarTyCons
144 ; traceVt " VECT with rhs : " $ ppr (map fst vectTyConsWithRHS)
145 ; traceVt " reuse : " $ ppr keep_tcs
146 ; traceVt " convert : " $ ppr conv_tcs
147
148 ; let defTyConDataCons origTyCon vectTyCon
149 = do { defTyCon origTyCon vectTyCon
150 ; MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
151 ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
152 }
153
154 -- For the type constructors that we don't need to vectorise, we use the original
155 -- representation in both unvectorised and vectorised code.
156 ; zipWithM_ defTyConDataCons keep_tcs keep_tcs
157
158 -- We do the same for type constructors declared VECTORISE SCALAR, while ignoring their
159 -- representation (data constructors) — see "Note [Pragmas to vectorise tycons]".
160 ; zipWithM_ defTyCon localScalarTyCons localScalarTyCons
161
162 -- For type constructors declared VECTORISE with an explicit vectorised type, we use the
163 -- explicitly given type in vectorised code and map data constructors one for one — see
164 -- "Note [Pragmas to vectorise tycons]".
165 ; mapM_ (uncurry defTyConDataCons) vectTyConsWithRHS
166
167 -- Vectorise all the data type declarations that we can and must vectorise.
168 ; new_tcs <- vectTyConDecls conv_tcs
169
170 -- We don't need new representation types for dictionary constructors. The constructors
171 -- are always fully applied, and we don't need to lift them to arrays as a dictionary
172 -- of a particular type always has the same value.
173 ; let vect_tcs = filter (not . isClassTyCon)
174 $ keep_tcs ++ new_tcs
175
176 -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
177 -- type constructors with vectorised representations.
178 ; reprs <- mapM tyConRepr vect_tcs
179 ; repr_tcs <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
180 ; pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
181 ; let inst_tcs = repr_tcs ++ pdata_tcs
182 fam_insts = map mkLocalFamInst inst_tcs
183 ; updGEnv $ extendFamEnv fam_insts
184
185 -- Generate dfuns for the 'PA' instances of the vectorised type constructors and
186 -- associate the type constructors with their dfuns in the global environment. We get
187 -- back the dfun bindings (which we will subsequently inject into the modules toplevel).
188 ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
189 do { defTyConPAs (zipLazy vect_tcs dfuns)
190 ; dfuns <- sequence
191 $ zipWith4 buildTyConBindings
192 orig_tcs
193 vect_tcs
194 repr_tcs
195 pdata_tcs
196
197 ; binds <- takeHoisted
198 ; return (dfuns, binds)
199 }
200
201 -- Return the vectorised variants of type constructors as well as the generated instance type
202 -- constructors, family instances, and dfun bindings.
203 ; return (new_tcs ++ inst_tcs, fam_insts, binds)
204 }
205
206
207 -- Helpers -------------------
208
209 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var
210 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc
211 = do { vectDataConWorkers orig_tc vect_tc pdata_tc
212 ; repr <- tyConRepr vect_tc
213 ; buildPADict vect_tc prepr_tc pdata_tc repr
214 }
215
216 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
217 vectDataConWorkers orig_tc vect_tc arr_tc
218 = do bs <- sequence
219 . zipWith3 def_worker (tyConDataCons orig_tc) rep_tys
220 $ zipWith4 mk_data_con (tyConDataCons vect_tc)
221 rep_tys
222 (inits rep_tys)
223 (tail $ tails rep_tys)
224 mapM_ (uncurry hoistBinding) bs
225 where
226 tyvars = tyConTyVars vect_tc
227 var_tys = mkTyVarTys tyvars
228 ty_args = map Type var_tys
229 res_ty = mkTyConApp vect_tc var_tys
230
231 cons = tyConDataCons vect_tc
232 arity = length cons
233 [arr_dc] = tyConDataCons arr_tc
234
235 rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
236
237
238 mk_data_con con tys pre post
239 = liftM2 (,) (vect_data_con con)
240 (lift_data_con tys pre post (mkDataConTag con))
241
242 sel_replicate len tag
243 | arity > 1 = do
244 rep <- builtin (selReplicate arity)
245 return [rep `mkApps` [len, tag]]
246
247 | otherwise = return []
248
249 vect_data_con con = return $ mkConApp con ty_args
250 lift_data_con tys pre_tys post_tys tag
251 = do
252 len <- builtin liftingContext
253 args <- mapM (newLocalVar (fsLit "xs"))
254 =<< mapM mkPDataType tys
255
256 sel <- sel_replicate (Var len) tag
257
258 pre <- mapM emptyPD (concat pre_tys)
259 post <- mapM emptyPD (concat post_tys)
260
261 return . mkLams (len : args)
262 . wrapFamInstBody arr_tc var_tys
263 . mkConApp arr_dc
264 $ ty_args ++ sel ++ pre ++ map Var args ++ post
265
266 def_worker data_con arg_tys mk_body
267 = do
268 arity <- polyArity tyvars
269 body <- closedV
270 . inBind orig_worker
271 . polyAbstract tyvars $ \args ->
272 liftM (mkLams (tyvars ++ args) . vectorised)
273 $ buildClosures tyvars [] arg_tys res_ty mk_body
274
275 raw_worker <- mkVectId orig_worker (exprType body)
276 let vect_worker = raw_worker `setIdUnfolding`
277 mkInlineUnfolding (Just arity) body
278 defGlobalVar orig_worker vect_worker
279 return (vect_worker, body)
280 where
281 orig_worker = dataConWorkId data_con