a91acab69d5ba145887db07e101f34c6a02404d5
[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 -- |Vectorise a type environment.
95 --
96 vectTypeEnv :: [TyCon] -- TyCons defined in this module
97 -> [CoreVect] -- All 'VECTORISE [SCALAR] type' declarations in this module
98 -> VM ( [TyCon] -- old TyCons ++ new TyCons
99 , [FamInst] -- New type family instances.
100 , [(Var, CoreExpr)]) -- New top level bindings.
101 vectTypeEnv tycons vectTypeDecls
102 = do { traceVt "** vectTypeEnv" $ ppr tycons
103
104 -- Build a map containing all vectorised type constructor. If they are scalar, they are
105 -- mapped to 'False' (vectorised type constructor == original type constructor).
106 ; allScalarTyConNames <- globalScalarTyCons -- covers both current and imported modules
107 ; vectTyCons <- globalVectTyCons
108 ; let vectTyConBase = mapNameEnv (const True) vectTyCons -- by default fully vectorised
109 vectTyConFlavour = foldNameSet (\n env -> extendNameEnv env n False) vectTyConBase
110 allScalarTyConNames
111
112 ; let -- {-# VECTORISE SCALAR type T -#} (imported and local tycons)
113 localScalarTyCons = [tycon | VectType True tycon Nothing <- vectTypeDecls]
114
115 -- {-# VECTORISE type T -#} (ONLY the imported tycons)
116 impVectTyCons = [tycon | VectType False tycon Nothing <- vectTypeDecls]
117 \\ tycons
118
119 -- {-# VECTORISE type T = ty -#} (imported and local tycons)
120 vectTyConsWithRHS = [ (tycon, rhs)
121 | VectType False tycon (Just rhs) <- vectTypeDecls]
122
123 -- filter VECTORISE SCALAR tycons and VECTORISE tycons with explicit rhses
124 vectSpecialTyConNames = mkNameSet . map tyConName $
125 localScalarTyCons ++ map fst vectTyConsWithRHS
126 notLocalScalarTyCon tc = not $ (tyConName tc) `elemNameSet` vectSpecialTyConNames
127
128 -- Split the list of 'TyCons' into the ones (1) that we must vectorise and those (2)
129 -- that we could, but don't need to vectorise. Type constructors that are not data
130 -- type constructors or use non-Haskell98 features are being dropped. They may not
131 -- appear in vectorised code. (We also drop the local type constructors appearing in a
132 -- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as
133 -- these are being handled separately.)
134 ; let maybeVectoriseTyCons = filter notLocalScalarTyCon tycons ++ impVectTyCons
135 (conv_tcs, keep_tcs) = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
136 orig_tcs = keep_tcs ++ conv_tcs
137
138 ; traceVt " VECT SCALAR : " $ ppr localScalarTyCons
139 ; traceVt " VECT with rhs : " $ ppr (map fst vectTyConsWithRHS)
140 ; traceVt " reuse : " $ ppr keep_tcs
141 ; traceVt " convert : " $ ppr conv_tcs
142
143 ; let defTyConDataCons origTyCon vectTyCon
144 = do { defTyCon origTyCon vectTyCon
145 ; MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
146 ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
147 }
148
149 -- For the type constructors that we don't need to vectorise, we use the original
150 -- representation in both unvectorised and vectorised code.
151 ; zipWithM_ defTyConDataCons keep_tcs keep_tcs
152
153 -- We do the same for type constructors declared VECTORISE SCALAR, while ignoring their
154 -- representation (data constructors) — see "Note [Pragmas to vectorise tycons]".
155 ; zipWithM_ defTyCon localScalarTyCons localScalarTyCons
156
157 -- For type constructors declared VECTORISE with an explicit vectorised type, we use the
158 -- explicitly given type in vectorised code and map data constructors one for one — see
159 -- "Note [Pragmas to vectorise tycons]".
160 ; mapM_ (uncurry defTyConDataCons) vectTyConsWithRHS
161
162 -- Vectorise all the data type declarations that we can and must vectorise.
163 ; new_tcs <- vectTyConDecls conv_tcs
164
165 -- We don't need new representation types for dictionary constructors. The constructors
166 -- are always fully applied, and we don't need to lift them to arrays as a dictionary
167 -- of a particular type always has the same value.
168 ; let vect_tcs = filter (not . isClassTyCon)
169 $ keep_tcs ++ new_tcs
170
171 -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
172 -- type constructors with vectorised representations.
173 ; reprs <- mapM tyConRepr vect_tcs
174 ; repr_tcs <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
175 ; pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
176 ; let inst_tcs = repr_tcs ++ pdata_tcs
177 fam_insts = map mkLocalFamInst inst_tcs
178 ; updGEnv $ extendFamEnv fam_insts
179
180 -- Generate dfuns for the 'PA' instances of the vectorised type constructors and
181 -- associate the type constructors with their dfuns in the global environment. We get
182 -- back the dfun bindings (which we will subsequently inject into the modules toplevel).
183 ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
184 do { defTyConPAs (zipLazy vect_tcs dfuns)
185 ; dfuns <- sequence
186 $ zipWith4 buildTyConBindings
187 orig_tcs
188 vect_tcs
189 repr_tcs
190 pdata_tcs
191
192 ; binds <- takeHoisted
193 ; return (dfuns, binds)
194 }
195
196 -- We return: (1) the vectorised type constructors, (2)
197 -- their 'PRepr' & 'PData' instance constructors two.
198 ; let new_tycons = tycons ++ new_tcs ++ inst_tcs
199
200 ; return (new_tycons, fam_insts, binds)
201 }
202
203
204 -- Helpers -------------------
205
206 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var
207 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc
208 = do { vectDataConWorkers orig_tc vect_tc pdata_tc
209 ; repr <- tyConRepr vect_tc
210 ; buildPADict vect_tc prepr_tc pdata_tc repr
211 }
212
213 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
214 vectDataConWorkers orig_tc vect_tc arr_tc
215 = do bs <- sequence
216 . zipWith3 def_worker (tyConDataCons orig_tc) rep_tys
217 $ zipWith4 mk_data_con (tyConDataCons vect_tc)
218 rep_tys
219 (inits rep_tys)
220 (tail $ tails rep_tys)
221 mapM_ (uncurry hoistBinding) bs
222 where
223 tyvars = tyConTyVars vect_tc
224 var_tys = mkTyVarTys tyvars
225 ty_args = map Type var_tys
226 res_ty = mkTyConApp vect_tc var_tys
227
228 cons = tyConDataCons vect_tc
229 arity = length cons
230 [arr_dc] = tyConDataCons arr_tc
231
232 rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
233
234
235 mk_data_con con tys pre post
236 = liftM2 (,) (vect_data_con con)
237 (lift_data_con tys pre post (mkDataConTag con))
238
239 sel_replicate len tag
240 | arity > 1 = do
241 rep <- builtin (selReplicate arity)
242 return [rep `mkApps` [len, tag]]
243
244 | otherwise = return []
245
246 vect_data_con con = return $ mkConApp con ty_args
247 lift_data_con tys pre_tys post_tys tag
248 = do
249 len <- builtin liftingContext
250 args <- mapM (newLocalVar (fsLit "xs"))
251 =<< mapM mkPDataType tys
252
253 sel <- sel_replicate (Var len) tag
254
255 pre <- mapM emptyPD (concat pre_tys)
256 post <- mapM emptyPD (concat post_tys)
257
258 return . mkLams (len : args)
259 . wrapFamInstBody arr_tc var_tys
260 . mkConApp arr_dc
261 $ ty_args ++ sel ++ pre ++ map Var args ++ post
262
263 def_worker data_con arg_tys mk_body
264 = do
265 arity <- polyArity tyvars
266 body <- closedV
267 . inBind orig_worker
268 . polyAbstract tyvars $ \args ->
269 liftM (mkLams (tyvars ++ args) . vectorised)
270 $ buildClosures tyvars [] arg_tys res_ty mk_body
271
272 raw_worker <- mkVectId orig_worker (exprType body)
273 let vect_worker = raw_worker `setIdUnfolding`
274 mkInlineUnfolding (Just arity) body
275 defGlobalVar orig_worker vect_worker
276 return (vect_worker, body)
277 where
278 orig_worker = dataConWorkId data_con