Env.hs 16.3 KB
Newer Older
1
-- Vectorise a modules type and class declarations.
2
--
3
4
-- This produces new type constructors and family instances top be included in the module toplevel
-- as well as bindings for worker functions, dfuns, and the like.
5

6
module Vectorise.Type.Env ( 
7
  vectTypeEnv,
8
9
) where
  
10
11
#include "HsVersions.h"

12
import Vectorise.Env
13
import Vectorise.Vect
14
15
import Vectorise.Monad
import Vectorise.Builtins
16
import Vectorise.Type.TyConDecl
17
import Vectorise.Type.Classify
18
19
import Vectorise.Generic.PADict
import Vectorise.Generic.PAMethods
20
import Vectorise.Generic.PData
21
import Vectorise.Generic.Description
22
import Vectorise.Utils
23

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
24
import CoreSyn
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
25
import CoreUtils
26
import CoreUnfold
27
import DataCon
28
29
import TyCon
import Type
30
import FamInstEnv
31
import Id
32
import MkId
33
import NameEnv
34
import NameSet
35

36
import Util
37
import Outputable
38
import FastString
39
40
import MonadUtils
import Control.Monad
41
import Data.List
42

43
44
45
46
47
48
49
50
51
-- Note [Pragmas to vectorise tycons]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
--
-- VECTORISE pragmas for type constructors cover three different flavours of vectorising data type
-- constructors:
--
-- (1) Data type constructor 'T' that may be used in vectorised code, where 'T' represents itself,
--     but the representation of 'T' is opaque in vectorised code.  
--
52
--     An example is the treatment of 'Int'.  'Int's can be used in vectorised code and remain
53
54
55
--     unchanged by vectorisation.  However, the representation of 'Int' by the 'I#' data
--     constructor wrapping an 'Int#' is not exposed in vectorised code.  Instead, computations
--     involving the representation need to be confined to scalar code.
56
--
57
58
59
60
61
62
--     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
--     by the vectoriser).
--
--     Type constructors declared with {-# VECTORISE SCALAR type T #-} are treated in this manner.
--     (The vectoriser never treats a type constructor automatically in this manner.)
--
63
64
65
66
67
68
-- (2) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
--     code, where 'T' and the 'Cn' are automatically vectorised in the same manner as data types
--     declared in a vectorised module.  This includes the case where the vectoriser determines that
--     the original representation of 'T' may be used in vectorised code (as it does not embed any
--     parallel arrays.)  This case is for type constructors that are *imported* from a non-
--     vectorised module, but that we want to use with full vectorisation support.
69
--
70
71
--     An example is the treatment of 'Ordering' and '[]'.  The former remains unchanged by
--     vectorisation, whereas the latter is fully vectorised.
72
73
74
75
76

--     'PData' and 'PRepr' instances are automatically generated by the vectoriser.
--
--     Type constructors declared with {-# VECTORISE type T #-} are treated in this manner.
--
77
78
79
-- (3) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
--     code, where 'T' is represented by an explicitly given 'Tv' whose constructors 'Cvn' represent
--     the original constructors in vectorised code.  As a special case, we can have 'Tv = T'
80
--
81
82
83
--     An example is the treatment of 'Bool', which is represented by itself in vectorised code
--     (as it cannot embed any parallel arrays).  However, we do not want any automatic generation
--     of class and family instances, which is why Case (2) does not apply.
84
--
85
86
--     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
--     by the vectoriser).
87
--
88
--     Type constructors declared with {-# VECTORISE type T = T' #-} are treated in this manner.
89
--
90
91
-- In addition, we have also got a single pragma form for type classes: {-# VECTORISE class C #-}.
-- It implies that the class type constructor may be used in vectorised code together with its data
92
-- constructor.  We generally produce a vectorised version of the data type and data constructor.
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
-- We do not generate 'PData' and 'PRepr' instances for class type constructors.  This pragma is the
-- default for all type classes declared in this module, but the pragma can also be used explitly on
-- imported classes.

-- Note [Vectorising classes]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~
--
-- We vectorise classes essentially by just vectorising their desugared Core representation, but we
-- do generate a 'Class' structure along the way (see 'Vectorise.Type.TyConDecl.vectTyConDecl').
--
-- Here is an example illustrating the mapping — assume
--
--   class Num a where
--     (+) :: a -> a -> a
--
-- It desugars to
--
110
--   data Num a = D:Num { (+) :: a -> a -> a }
111
112
113
--
-- which we vectorise to
--
114
--  data V:Num a = D:V:Num { ($v+) :: PArray a :-> PArray a :-> PArray a }
115
116
117
--
-- while adding the following entries to the vectorisation map:
--
118
119
--   tycon  : Num   --> V:Num
--   datacon: D:Num --> D:V:Num
120
--   var    : (+) --> ($v+)
121

122
-- |Vectorise type constructor including class type constructors.
123
--
124
vectTypeEnv :: [TyCon]                  -- Type constructors defined in this module
125
            -> [CoreVect]               -- All 'VECTORISE [SCALAR] type' declarations in this module
126
            -> [CoreVect]               -- All 'VECTORISE class' declarations in this module
127
            -> VM ( [TyCon]             -- old TyCons ++ new TyCons
128
129
                  , [FamInst]           -- New type family instances.
                  , [(Var, CoreExpr)])  -- New top level bindings.
130
vectTypeEnv tycons vectTypeDecls vectClassDecls
131
  = do { traceVt "** vectTypeEnv" $ ppr tycons
132
133
134
135
136
137
138
139
140

         -- Build a map containing all vectorised type constructor.  If they are scalar, they are
         -- mapped to 'False' (vectorised type constructor == original type constructor).
       ; allScalarTyConNames <- globalScalarTyCons  -- covers both current and imported modules
       ; vectTyCons          <- globalVectTyCons
       ; let vectTyConBase    = mapNameEnv (const True) vectTyCons   -- by default fully vectorised
             vectTyConFlavour = foldNameSet (\n env -> extendNameEnv env n False) vectTyConBase
                                            allScalarTyConNames

141
142
143
144
       ; let   -- {-# VECTORISE SCALAR type T -#} (imported and local tycons)
             localScalarTyCons      = [tycon | VectType True  tycon Nothing <- vectTypeDecls]

               -- {-# VECTORISE type T -#} (ONLY the imported tycons)
145
146
             impVectTyCons          = (   [tycon | VectType False tycon Nothing <- vectTypeDecls]
                                       ++ [tycon | VectClass tycon              <- vectClassDecls])
147
148
149
150
151
152
153
154
155
156
157
                                      \\ tycons

               -- {-# VECTORISE type T = ty -#} (imported and local tycons)
             vectTyConsWithRHS      = [ (tycon, rhs) 
                                      | VectType False tycon (Just rhs) <- vectTypeDecls]

               -- filter VECTORISE SCALAR tycons and VECTORISE tycons with explicit rhses
             vectSpecialTyConNames  = mkNameSet . map tyConName $
                                        localScalarTyCons ++ map fst vectTyConsWithRHS
             notLocalScalarTyCon tc = not $ (tyConName tc) `elemNameSet` vectSpecialTyConNames

158
159
160
161
           -- Split the list of 'TyCons' into the ones (1) that we must vectorise and those (2)
           -- that we could, but don't need to vectorise.  Type constructors that are not data
           -- type constructors or use non-Haskell98 features are being dropped.  They may not
           -- appear in vectorised code.  (We also drop the local type constructors appearing in a
162
163
           -- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as
           -- these are being handled separately.)
164
165
166
           -- Furthermore, 'drop_tcs' are those type constructors that we cannot vectorise.
       ; let maybeVectoriseTyCons           = filter notLocalScalarTyCon tycons ++ impVectTyCons
             (conv_tcs, keep_tcs, drop_tcs) = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
167
             
168
       ; traceVt " VECT SCALAR    : " $ ppr localScalarTyCons
169
       ; traceVt " VECT [class]   : " $ ppr impVectTyCons
170
       ; traceVt " VECT with rhs  : " $ ppr (map fst vectTyConsWithRHS)
171
       ; traceVt " -- after classification (local and VECT [class] tycons) --" empty
172
173
       ; traceVt " reuse          : " $ ppr keep_tcs
       ; traceVt " convert        : " $ ppr conv_tcs
174
175
       
           -- warn the user about unvectorised type constructors
176
177
178
179
       ; let explanation    = ptext (sLit "(They use unsupported language extensions") $$
                              ptext (sLit "or depend on type constructors that are not vectorised)")
             drop_tcs_nosyn = filter (not . isSynTyCon) drop_tcs
       ; unless (null drop_tcs_nosyn) $
180
           emitVt "Warning: cannot vectorise these type constructors:" $ 
181
             pprQuotedList drop_tcs_nosyn $$ explanation
182

183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
       ; let defTyConDataCons origTyCon vectTyCon
               = do { defTyCon origTyCon vectTyCon
                    ; MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
                    ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
                    }

           -- For the type constructors that we don't need to vectorise, we use the original
           -- representation in both unvectorised and vectorised code.
       ; zipWithM_ defTyConDataCons keep_tcs keep_tcs

           -- We do the same for type constructors declared VECTORISE SCALAR, while ignoring their
           -- representation (data constructors) — see "Note [Pragmas to vectorise tycons]".
       ; zipWithM_ defTyCon localScalarTyCons localScalarTyCons

           -- For type constructors declared VECTORISE with an explicit vectorised type, we use the
           -- explicitly given type in vectorised code and map data constructors one for one — see
           -- "Note [Pragmas to vectorise tycons]".
       ; mapM_ (uncurry defTyConDataCons) vectTyConsWithRHS
201

202
203
           -- Vectorise all the data type declarations that we can and must vectorise (enter the
           -- type and data constructors into the vectorisation map on-the-fly.)
204
205
206
207
208
       ; new_tcs <- vectTyConDecls conv_tcs

           -- We don't need new representation types for dictionary constructors. The constructors
           -- are always fully applied, and we don't need to lift them to arrays as a dictionary
           -- of a particular type always has the same value.
209
210
       ; let orig_tcs = filter (not . isClassTyCon) $ keep_tcs ++ conv_tcs
             vect_tcs = filter (not . isClassTyCon) $ keep_tcs ++ new_tcs
211
212
213

           -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
           -- type constructors with vectorised representations.
214
215
216
217
       ; reprs      <- mapM tyConRepr vect_tcs
       ; repr_tcs   <- zipWith3M buildPReprTyCon  orig_tcs vect_tcs reprs
       ; pdata_tcs  <- zipWith3M buildPDataTyCon  orig_tcs vect_tcs reprs
       ; pdatas_tcs <- zipWith3M buildPDatasTyCon orig_tcs vect_tcs reprs
218

219
       ; let inst_tcs  = repr_tcs ++ pdata_tcs ++ pdatas_tcs
220
221
222
             fam_insts = map mkLocalFamInst inst_tcs
       ; updGEnv $ extendFamEnv fam_insts

223
224
225
226
           -- Generate workers for the vectorised data constructors, dfuns for the 'PA' instances of
           -- the vectorised type constructors, and associate the type constructors with their dfuns
           -- in the global environment.  We get back the dfun bindings (which we will subsequently
           -- inject into the modules toplevel).
227
228
       ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
           do { defTyConPAs (zipLazy vect_tcs dfuns)
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252

                  -- query the 'PData' instance type constructors for type constructors that have a
                  -- VECTORISE pragma with an explicit right-hand side (this is Item (3) of
                  -- "Note [Pragmas to vectorise tycons]" above)
              ; pdata_withRHS_tcs <- mapM pdataReprTyConExact
                                          [ mkTyConApp tycon tys
                                          | (tycon, _) <- vectTyConsWithRHS
                                          , let tys = mkTyVarTys (tyConTyVars tycon)
                                          ]

                  -- build workers for all vectorised data constructors (except scalar ones)
              ; sequence_ $
                  zipWith3 vectDataConWorkers (orig_tcs  ++ map fst vectTyConsWithRHS)
                                              (vect_tcs  ++ map snd vectTyConsWithRHS)
                                              (pdata_tcs ++ pdata_withRHS_tcs)

                  -- build a 'PA' dictionary for all type constructors (except scalar ones and those
                  -- defined with an explicit right-hand side where the dictionary is user-supplied)
              ; dfuns <- sequence $
                           zipWith4 buildTyConPADict
                                    vect_tcs
                                    repr_tcs
                                    pdata_tcs
                                    pdatas_tcs
253
254
255
256
257

              ; binds <- takeHoisted
              ; return (dfuns, binds)
              }

258
259
           -- Return the vectorised variants of type constructors as well as the generated instance
           -- type constructors, family instances, and dfun bindings.
260
       ; return (new_tcs ++ inst_tcs, fam_insts, binds)
261
262
263
       }


264
-- Helpers --------------------------------------------------------------------
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
265

266
267
268
269
270
271
272
273
274
275
276
277
buildTyConPADict :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var
buildTyConPADict vect_tc prepr_tc pdata_tc pdatas_tc
 = tyConRepr vect_tc >>= buildPADict vect_tc prepr_tc pdata_tc pdatas_tc

-- Produce a custom-made worker for the data constructors of a vectorised data type.  This includes
-- all data constructors that may be used in vetcorised code — i.e., all data constructors of data
-- types other than scalar ones.  Also adds a mapping from the original to vectorised worker into
-- the vectorisation map.
--
-- FIXME: It's not nice that we need create a special worker after the data constructors has
--   already been constructed.  Also, I don't think the worker is properly added to the data
--   constructor.  Seems messy.
278
279
vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
vectDataConWorkers orig_tc vect_tc arr_tc
280
281
282
283
284
285
286
287
288
289
290
  = do { traceVt "Building vectorised worker for datatype" (ppr orig_tc)
  
       ; bs <- sequence
             . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
             $ zipWith4 mk_data_con (tyConDataCons vect_tc)
                                    rep_tys
                                    (inits rep_tys)
                                    (tail $ tails rep_tys)
        ; mapM_ (uncurry hoistBinding) bs
        }
  where
291
292
293
294
295
    tyvars   = tyConTyVars vect_tc
    var_tys  = mkTyVarTys tyvars
    ty_args  = map Type var_tys
    res_ty   = mkTyConApp vect_tc var_tys

296
297
298
299
    cons     = tyConDataCons vect_tc
    arity    = length cons
    [arr_dc] = tyConDataCons arr_tc

300
301
302
303
    rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc

    mk_data_con con tys pre post
      = liftM2 (,) (vect_data_con con)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
304
                   (lift_data_con tys pre post (mkDataConTag con))
305

306
307
308
309
310
311
312
    sel_replicate len tag
      | arity > 1 = do
                      rep <- builtin (selReplicate arity)
                      return [rep `mkApps` [len, tag]]

      | otherwise = return []

313
    vect_data_con con = return $ mkConApp con ty_args
314
    lift_data_con tys pre_tys post_tys tag
315
316
      = do
          len  <- builtin liftingContext
Ian Lynagh's avatar
Ian Lynagh committed
317
          args <- mapM (newLocalVar (fsLit "xs"))
318
                  =<< mapM mkPDataType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
319

320
          sel  <- sel_replicate (Var len) tag
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
321

322
323
          pre   <- mapM emptyPD (concat pre_tys)
          post  <- mapM emptyPD (concat post_tys)
324
325
326
327

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
328
                 $ ty_args ++ sel ++ pre ++ map Var args ++ post
329
330
331

    def_worker data_con arg_tys mk_body
      = do
332
          arity <- polyArity tyvars
333
334
          body <- closedV
                . inBind orig_worker
335
336
                . polyAbstract tyvars $ \args ->
                  liftM (mkLams (tyvars ++ args) . vectorised)
337
                $ buildClosures tyvars [] [] arg_tys res_ty mk_body
338

339
          raw_worker <- mkVectId orig_worker (exprType body)
340
          let vect_worker = raw_worker `setIdUnfolding`
341
                              mkInlineUnfolding (Just arity) body
342
343
344
345
          defGlobalVar orig_worker vect_worker
          return (vect_worker, body)
      where
        orig_worker = dataConWorkId data_con