Add VECTORISE [SCALAR] type pragma
[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 import Vectorise.Env
14 import Vectorise.Vect
15 import Vectorise.Monad
16 import Vectorise.Builtins
17 import Vectorise.Type.TyConDecl
18 import Vectorise.Type.Classify
19 import Vectorise.Type.PADict
20 import Vectorise.Type.PData
21 import Vectorise.Type.PRepr
22 import Vectorise.Type.Repr
23 import Vectorise.Utils
24
25 import HscTypes
26 import CoreSyn
27 import CoreUtils
28 import CoreUnfold
29 import DataCon
30 import TyCon
31 import Type
32 import FamInstEnv
33 import OccName
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' represent themselves in vectorised code.
69 --
70 -- An example is the treatment of 'Bool'. 'Bool' together with 'False' and 'True' may appear in
71 -- vectorised code and they remain unchanged by vectorisation. (There is no need for a special
72 -- representation as the values cannot embed any arrays.)
73
74 -- 'PData' and 'PRepr' instances are automatically generated by the vectoriser.
75 --
76 -- Type constructors declared with {-# VECTORISE type T #-} are treated in this manner.
77 -- (This is the same treatment that type constructors receive that the vectoriser deems fit for
78 -- use in vectorised code, but for which no special vectorised variant needs to be generated.)
79 --
80 -- (3) [NOT IMPLEMENTED YET]
81 -- Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
82 -- code, where 'T' is represented by 'Tv' and the workers of the 'Cn' are represented 'vCn' in
83 -- vectorised code.
84 --
85 -- ??Example??
86 --
87 -- 'PData' and 'PRepr' instances are automatically generated by the vectoriser.
88 --
89 -- ??How declared??
90
91 -- |Vectorise a type environment.
92 --
93 vectTypeEnv :: TypeEnv -- Original type environment
94 -> [CoreVect] -- All 'VECTORISE [SCALAR] type' declarations in this module
95 -> VM ( TypeEnv -- Vectorised type environment.
96 , [FamInst] -- New type family instances.
97 , [(Var, CoreExpr)]) -- New top level bindings.
98 vectTypeEnv env vectTypeDecls
99 = do { traceVt "** vectTypeEnv" $ ppr env
100
101 -- Build a map containing all vectorised type constructor. If they are scalar, they are
102 -- mapped to 'False' (vectorised type constructor == original type constructor).
103 ; allScalarTyConNames <- globalScalarTyCons -- covers both current and imported modules
104 ; vectTyCons <- globalVectTyCons
105 ; let vectTyConBase = mapNameEnv (const True) vectTyCons -- by default fully vectorised
106 vectTyConFlavour = foldNameSet (\n env -> extendNameEnv env n False) vectTyConBase
107 allScalarTyConNames
108
109 -- Split the list of 'TyCons' into the ones (1) that we must vectorise and those (2)
110 -- that we could, but don't need to vectorise. Type constructors that are not data
111 -- type constructors or use non-Haskell98 features are being dropped. They may not
112 -- appear in vectorised code. (We also drop the local type constructors appearing in a
113 -- VECTORISE SCALAR pragma, as they are being handled separately.)
114 ; let localScalarTyCons = [tycon | VectType tycon Nothing <- vectTypeDecls]
115 localScalarTyConNames = mkNameSet (map tyConName localScalarTyCons)
116 notLocalScalarTyCon tc = not $ (tyConName tc) `elemNameSet` localScalarTyConNames
117
118 maybeVectoriseTyCons = filter notLocalScalarTyCon (typeEnvTyCons env)
119 (conv_tcs, keep_tcs) = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
120 orig_tcs = keep_tcs ++ conv_tcs
121 keep_dcs = concatMap tyConDataCons keep_tcs
122
123 keep_and_scalar_tcs = keep_tcs ++ localScalarTyCons
124
125 -- Of those type constructors that we don't need to vectorise, we use the original
126 -- representation in both unvectorised and vectorised code. For those declared VECTORISE
127 -- SCALAR, we ignore their represention — see "Note [Pragmas to vectorise tycons]".
128 ; zipWithM_ defTyCon keep_and_scalar_tcs keep_and_scalar_tcs
129 ; zipWithM_ defDataCon keep_dcs keep_dcs
130
131 -- Vectorise all the data type declarations that we can and must vectorise.
132 ; new_tcs <- vectTyConDecls conv_tcs
133
134 -- We don't need new representation types for dictionary constructors. The constructors
135 -- are always fully applied, and we don't need to lift them to arrays as a dictionary
136 -- of a particular type always has the same value.
137 ; let vect_tcs = filter (not . isClassTyCon)
138 $ keep_tcs ++ new_tcs
139
140 -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
141 -- type constructors with vectorised representations.
142 ; reprs <- mapM tyConRepr vect_tcs
143 ; repr_tcs <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
144 ; pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
145 ; let inst_tcs = repr_tcs ++ pdata_tcs
146 fam_insts = map mkLocalFamInst inst_tcs
147 ; updGEnv $ extendFamEnv fam_insts
148
149 -- Generate dfuns for the 'PA' instances of the vectorised type constructors and
150 -- associate the type constructors with their dfuns in the global environment. We get
151 -- back the dfun bindings (which we will subsequently inject into the modules toplevel).
152 ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
153 do { defTyConPAs (zipLazy vect_tcs dfuns)
154 ; dfuns <- sequence
155 $ zipWith4 buildTyConBindings
156 orig_tcs
157 vect_tcs
158 repr_tcs
159 pdata_tcs
160
161 ; binds <- takeHoisted
162 ; return (dfuns, binds)
163 }
164
165 -- We add to the type environment: (1) the vectorised type constructors, (2) their
166 -- 'PRepr' & 'PData' instance constructors, and (3) the data constructors of the fomer
167 -- two.
168 ; let all_new_tcs = new_tcs ++ inst_tcs
169 new_env = extendTypeEnvList env
170 $ map ATyCon all_new_tcs ++
171 [ADataCon dc | tc <- all_new_tcs
172 , dc <- tyConDataCons tc]
173
174 ; return (new_env, fam_insts, binds)
175 }
176
177
178 -- Helpers -------------------
179
180 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var
181 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc
182 = do { vectDataConWorkers orig_tc vect_tc pdata_tc
183 ; repr <- tyConRepr vect_tc
184 ; buildPADict vect_tc prepr_tc pdata_tc repr
185 }
186
187 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
188 vectDataConWorkers orig_tc vect_tc arr_tc
189 = do bs <- sequence
190 . zipWith3 def_worker (tyConDataCons orig_tc) rep_tys
191 $ zipWith4 mk_data_con (tyConDataCons vect_tc)
192 rep_tys
193 (inits rep_tys)
194 (tail $ tails rep_tys)
195 mapM_ (uncurry hoistBinding) bs
196 where
197 tyvars = tyConTyVars vect_tc
198 var_tys = mkTyVarTys tyvars
199 ty_args = map Type var_tys
200 res_ty = mkTyConApp vect_tc var_tys
201
202 cons = tyConDataCons vect_tc
203 arity = length cons
204 [arr_dc] = tyConDataCons arr_tc
205
206 rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
207
208
209 mk_data_con con tys pre post
210 = liftM2 (,) (vect_data_con con)
211 (lift_data_con tys pre post (mkDataConTag con))
212
213 sel_replicate len tag
214 | arity > 1 = do
215 rep <- builtin (selReplicate arity)
216 return [rep `mkApps` [len, tag]]
217
218 | otherwise = return []
219
220 vect_data_con con = return $ mkConApp con ty_args
221 lift_data_con tys pre_tys post_tys tag
222 = do
223 len <- builtin liftingContext
224 args <- mapM (newLocalVar (fsLit "xs"))
225 =<< mapM mkPDataType tys
226
227 sel <- sel_replicate (Var len) tag
228
229 pre <- mapM emptyPD (concat pre_tys)
230 post <- mapM emptyPD (concat post_tys)
231
232 return . mkLams (len : args)
233 . wrapFamInstBody arr_tc var_tys
234 . mkConApp arr_dc
235 $ ty_args ++ sel ++ pre ++ map Var args ++ post
236
237 def_worker data_con arg_tys mk_body
238 = do
239 arity <- polyArity tyvars
240 body <- closedV
241 . inBind orig_worker
242 . polyAbstract tyvars $ \args ->
243 liftM (mkLams (tyvars ++ args) . vectorised)
244 $ buildClosures tyvars [] arg_tys res_ty mk_body
245
246 raw_worker <- cloneId mkVectOcc orig_worker (exprType body)
247 let vect_worker = raw_worker `setIdUnfolding`
248 mkInlineUnfolding (Just arity) body
249 defGlobalVar orig_worker vect_worker
250 return (vect_worker, body)
251 where
252 orig_worker = dataConWorkId data_con