Env.hs 22.2 KB
Newer Older
1 2
{-# LANGUAGE CPP #-}

3
-- Vectorise a modules type and class declarations.
4
--
5 6
-- 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.
7

8
module Vectorise.Type.Env (
9
  vectTypeEnv,
10
) where
11

12 13
#include "HsVersions.h"

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

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
26
import CoreSyn
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
27
import CoreUtils
28
import CoreUnfold
29
import DataCon
30
import TyCon
31
import CoAxiom
32
import Type
33
import FamInstEnv
34
import Id
35
import MkId
36
import NameEnv
37
import NameSet
38
import UniqFM
39
import OccName
40
import Unique
41

42
import Util
43
import Outputable
44
import DynFlags
45
import FastString
46
import MonadUtils
47

48
import Control.Monad
49
import Data.Maybe
50
import Data.List
51

52

53 54 55
-- Note [Pragmas to vectorise tycons]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
--
56 57 58 59 60 61 62 63 64 65 66 67 68
-- All imported type constructors that are not mapped to a vectorised type in the vectorisation map
-- (possibly because the defining module was not compiled with vectorisation) may be used in scalar
-- code encapsulated in vectorised code. If a such a type constructor 'T' is a member of the
-- 'Scalar' class (and hence also of 'PData' and 'PRepr'), it may also be used in vectorised code,
-- where 'T' represents itself, but the representation of 'T' still remains opaque in vectorised
-- code (i.e., it can only be used in scalar code).
--
-- An example is the treatment of 'Int'. 'Int's can be used in vectorised code and remain 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.
--
-- VECTORISE pragmas for type constructors cover four different flavours of vectorising data type
69 70
-- constructors:
--
71 72 73 74 75 76
-- (1) 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.
77
--
78 79
--     An example is the treatment of 'Ordering' and '[]'.  The former remains unchanged by
--     vectorisation, whereas the latter is fully vectorised.
80
--
81
--     'PData' and 'PRepr' instances are automatically generated by the vectoriser.
82
--
83
--     Type constructors declared with {-# VECTORISE type T #-} are treated in this manner.
84
--
85
-- (2) Data type constructor 'T' that may be used in vectorised code, where 'T' is represented by an
86
--     explicitly given 'Tv', but the representation of 'T' is opaque in vectorised code (i.e., the
87
--     constructors of 'T' may not occur in vectorised code).
88
--
89 90 91
--     An example is the treatment of '[::]'. The type '[::]' can be used in vectorised code and is
--     vectorised to 'PArray'. However, the representation of '[::]' is not exposed in vectorised
--     code. Instead, computations involving the representation need to be confined to scalar code.
92 93 94 95
--
--     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
--     by the vectoriser).
--
96
--     Type constructors declared with {-# VECTORISE type T = Tv #-} are treated in this manner
97 98
--     manner. (The vectoriser never treats a type constructor automatically in this manner.)
--
99 100 101 102
-- (3) Data type constructor 'T' that does not contain any parallel arrays and has explicitly
--     provided 'PData' and 'PRepr' instances (and maybe also a 'Scalar' instance), which together
--     with the type's constructors 'Cn' may be used in vectorised code. The type 'T' and its
--     constructors 'Cn' are represented by themselves in vectorised code.
103
--
104 105 106
--     An example is '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 (1) does not apply.
107
--
108 109 110 111
--     '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.
112
--
113 114 115 116
-- (4) Data type constructor 'T' that does not contain any parallel arrays and that, in vectorised
--     code, is represented by an explicitly given 'Tv', but the representation of 'T' is opaque in
--     vectorised code and 'T' is regarded to be scalar — i.e., it may be used in encapsulated
--     scalar subcomputations.
117
--
118 119 120 121
--     An example is the treatment of '(->)'. Types '(->)' can be used in vectorised code and are
--     vectorised to '(:->)'.  However, the representation of '(->)' is not exposed in vectorised
--     code. Instead, computations involving the representation need to be confined to scalar code
--     and may be part of encapsulated scalar computations.
122
--
123 124
--     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
--     by the vectoriser).
125
--
126
--     Type constructors declared with {-# VECTORISE SCALAR type T = Tv #-} are treated in this
127
--     manner. (The vectoriser never treats a type constructor automatically in this manner.)
128
--
129 130
-- 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
131
-- constructor.  We generally produce a vectorised version of the data type and data constructor.
132
-- We do not generate 'PData' and 'PRepr' instances for class type constructors.  This pragma is the
133 134
-- default for all type classes declared in a vectorised module, but the pragma can also be used
-- explitly on imported classes.
135 136 137 138 139 140 141 142 143 144 145 146 147 148

-- 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
--
149
--   data Num a = D:Num { (+) :: a -> a -> a }
150 151 152
--
-- which we vectorise to
--
153
--  data V:Num a = D:V:Num { ($v+) :: PArray a :-> PArray a :-> PArray a }
154 155 156
--
-- while adding the following entries to the vectorisation map:
--
157 158
--   tycon  : Num   --> V:Num
--   datacon: D:Num --> D:V:Num
159
--   var    : (+) --> ($v+)
160

161
-- |Vectorise type constructor including class type constructors.
162
--
163 164 165 166
vectTypeEnv :: [TyCon]                   -- Type constructors defined in this module
            -> [CoreVect]                -- All 'VECTORISE [SCALAR] type' declarations in this module
            -> [CoreVect]                -- All 'VECTORISE class' declarations in this module
            -> VM ( [TyCon]              -- old TyCons ++ new TyCons
167
                  , [FamInst]            -- New type family instances.
168
                  , [(Var, CoreExpr)])   -- New top level bindings.
169
vectTypeEnv tycons vectTypeDecls vectClassDecls
170
  = do { traceVt "** vectTypeEnv" $ ppr tycons
171

172
       ; let   -- {-# VECTORISE type T -#} (ONLY the imported tycons)
173 174
             impVectTyCons          = (   [tycon | VectType False tycon Nothing <- vectTypeDecls]
                                       ++ [tycon | VectClass tycon              <- vectClassDecls])
175
                                      \\ tycons
176

177 178 179 180 181
               -- {-# VECTORISE type T = Tv -#} (imported & local tycons with an /RHS/)
             vectTyConsWithRHS      = [ (tycon, rhs)
                                      | VectType False tycon (Just rhs) <- vectTypeDecls]

               -- {-# VECTORISE SCALAR type T = Tv -#} (imported & local tycons with an /RHS/)
182
             scalarTyConsWithRHS    = [ (tycon, rhs)
183
                                      | VectType True  tycon (Just rhs) <- vectTypeDecls]
184

185 186
               -- {-# VECTORISE SCALAR type T -#} (imported & local /scalar/ tycons without an RHS)
             scalarTyConsNoRHS      = [tycon | VectType True tycon Nothing <- vectTypeDecls]
187

188
               -- Check that is not a VECTORISE SCALAR tycon nor VECTORISE tycons with explicit rhs?
189 190
             vectSpecialTyConNames  = mkNameSet . map tyConName $
                                        scalarTyConsNoRHS ++
191
                                        map fst (vectTyConsWithRHS ++ scalarTyConsWithRHS)
192
             notVectSpecialTyCon tc = not $ (tyConName tc) `elemNameSet` vectSpecialTyConNames
193

194 195 196
         -- Build a map containing all vectorised type constructor. If the vectorised type
         -- constructor differs from the original one, then it is mapped to 'True'; if they are
         -- both the same, then it maps to 'False'.
197
       ; vectTyCons          <- globalVectTyCons
198 199
       ; let vectTyConBase    = mapUFM_Directly isDistinct vectTyCons    -- 'True' iff tc /= V[[tc]]
             isDistinct u tc  = u /= getUnique tc
200 201 202
             vectTyConFlavour = vectTyConBase
                                `plusNameEnv`
                                mkNameEnv [ (tyConName tycon, True)
203
                                          | (tycon, _) <- vectTyConsWithRHS ++ scalarTyConsWithRHS]
204 205
                                `plusNameEnv`
                                mkNameEnv [ (tyConName tycon, False)  -- original representation
206
                                          | tycon <- scalarTyConsNoRHS]
207

208

209 210 211 212
           -- 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
213
           -- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as
214 215 216
           -- these are being handled separately.  NB: Some type constructors may be marked SCALAR
           -- /and/ have an explicit right-hand side.)
           --
217 218 219
           -- Furthermore, 'par_tcs' are those type constructors (converted or not) whose
           -- definition, directly or indirectly, depends on parallel arrays. Finally, 'drop_tcs'
           -- are all type constructors that cannot be vectorised.
220
       ; parallelTyCons <- (`extendNameSetList` map (tyConName . fst) vectTyConsWithRHS) <$>
221
                             globalParallelTyCons
222
       ; let maybeVectoriseTyCons = filter notVectSpecialTyCon tycons ++ impVectTyCons
223
             (conv_tcs, keep_tcs, par_tcs, drop_tcs)
224
               = classifyTyCons vectTyConFlavour parallelTyCons maybeVectoriseTyCons
225 226 227

       ; traceVt " known parallel : " $ ppr parallelTyCons
       ; traceVt " VECT SCALAR    : " $ ppr (scalarTyConsNoRHS ++ map fst scalarTyConsWithRHS)
228
       ; traceVt " VECT [class]   : " $ ppr impVectTyCons
229
       ; traceVt " VECT with rhs  : " $ ppr (map fst (vectTyConsWithRHS ++ scalarTyConsWithRHS))
230
       ; traceVt " -- after classification (local and VECT [class] tycons) --" Outputable.empty
231 232
       ; traceVt " reuse          : " $ ppr keep_tcs
       ; traceVt " convert        : " $ ppr conv_tcs
233

234
           -- warn the user about unvectorised type constructors
235 236 237
       ; let explanation    = text "(They use unsupported language extensions"
                          $$  text "or depend on type constructors that are" <+>
                              text "not vectorised)"
238 239
             drop_tcs_nosyn = filter (not . isTypeFamilyTyCon) .
                              filter (not . isTypeSynonymTyCon) $ drop_tcs
240
       ; unless (null drop_tcs_nosyn) $
241
           emitVt "Warning: cannot vectorise these type constructors:" $
242
             pprQuotedList drop_tcs_nosyn $$ explanation
243

244
       ; mapM_ addParallelTyConAndCons $ par_tcs ++ map fst vectTyConsWithRHS
245

246
       ; let mapping =
247 248
                    -- Type constructors that we found we don't need to vectorise and those
                    -- declared VECTORISE SCALAR /without/ an explicit right-hand side, use the same
249 250
                    -- representation in both unvectorised and vectorised code; they are not
                    -- abstract.
251
                  [(tycon, tycon, False) | tycon <- keep_tcs ++ scalarTyConsNoRHS]
252
                    -- We do the same for type constructors declared VECTORISE SCALAR /without/
253
                    -- an explicit right-hand side
254
               ++ [(tycon, vTycon, True) | (tycon, vTycon) <- vectTyConsWithRHS ++ scalarTyConsWithRHS]
255
       ; syn_tcs <- catMaybes <$> mapM defTyConDataCons mapping
256

257 258
           -- 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.)
259
       ; new_tcs <- vectTyConDecls conv_tcs
260

261 262 263 264 265
       ; let dumpTc tc vTc = traceVt "---" (ppr tc <+> text "::" <+> ppr (dataConSig tc) $$
                                            ppr vTc <+> text "::" <+> ppr (dataConSig vTc))
             dataConSig tc | Just dc <- tyConSingleDataCon_maybe tc = dataConRepType dc
                           | otherwise                              = panic "dataConSig"
       ; zipWithM_ dumpTc (filter isClassTyCon conv_tcs) (filter isClassTyCon new_tcs)
266 267 268 269

           -- 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.
270 271
       ; let orig_tcs = filter (not . isClassTyCon) $ keep_tcs ++ conv_tcs
             vect_tcs = filter (not . isClassTyCon) $ keep_tcs ++ new_tcs
272 273 274

           -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
           -- type constructors with vectorised representations.
275
       ; reprs      <- mapM tyConRepr vect_tcs
276 277 278 279 280 281 282 283
       ; repr_fis   <- zipWith3M buildPReprTyCon  orig_tcs vect_tcs reprs
       ; pdata_fis  <- zipWith3M buildPDataTyCon  orig_tcs vect_tcs reprs
       ; pdatas_fis <- zipWith3M buildPDatasTyCon orig_tcs vect_tcs reprs

       ; let fam_insts  = repr_fis ++ pdata_fis ++ pdatas_fis
             repr_axs   = map famInstAxiom repr_fis
             pdata_tcs  = famInstsRepTyCons pdata_fis
             pdatas_tcs = famInstsRepTyCons pdatas_fis
284

285 286
       ; updGEnv $ extendFamEnv fam_insts

287 288 289 290
           -- 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).
291 292
       ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
           do { defTyConPAs (zipLazy vect_tcs dfuns)
293

294
                  -- Query the 'PData' instance type constructors for type constructors that have a
295 296 297
                  -- VECTORISE SCALAR type pragma without an explicit right-hand side (this is Item
                  --  (3) of "Note [Pragmas to vectorise tycons]" above).
              ; pdata_scalar_tcs <- mapM pdataReprTyConExact scalarTyConsNoRHS
298

299
                  -- Build workers for all vectorised data constructors (except abstract ones)
300
              ; sequence_ $
301 302 303
                  zipWith3 vectDataConWorkers (orig_tcs  ++ scalarTyConsNoRHS)
                                              (vect_tcs  ++ scalarTyConsNoRHS)
                                              (pdata_tcs ++ pdata_scalar_tcs)
304

305
                  -- Build a 'PA' dictionary for all type constructors (except abstract ones & those
306 307 308 309
                  -- defined with an explicit right-hand side where the dictionary is user-supplied)
              ; dfuns <- sequence $
                           zipWith4 buildTyConPADict
                                    vect_tcs
310
                                    repr_axs
311 312
                                    pdata_tcs
                                    pdatas_tcs
313 314 315 316 317

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

318 319
           -- Return the vectorised variants of type constructors as well as the generated instance
           -- type constructors, family instances, and dfun bindings.
320 321
       ; return ( new_tcs ++ pdata_tcs ++ pdatas_tcs ++ syn_tcs
                , fam_insts, binds)
322
       }
323
  where
324 325 326
    addParallelTyConAndCons tycon
      = do
        { addGlobalParallelTyCon tycon
327 328 329
        ; mapM_ addGlobalParallelVar [ id | dc <- tyConDataCons tycon
                                          , AnId id <- dataConImplicitTyThings dc ]
                                          -- Ignoring the promoted tycon; hope that's ok
330
        }
331

332
    -- Add a mapping from the original to vectorised type constructor to the vectorisation map.
Gabor Greif's avatar
Gabor Greif committed
333
    -- Unless the type constructor is abstract, also mappings from the original's data constructors
334 335 336 337 338 339 340 341 342
    -- to the vectorised type's data constructors.
    --
    -- We have three cases: (1) original and vectorised type constructor are the same, (2) the
    -- name of the vectorised type constructor is canonical (as prescribed by 'mkVectTyConOcc'), or
    -- (3) the name is not canonical.  In the third case, we additionally introduce a type synonym
    -- with the canonical name that is set equal to the non-canonical name (so that we find the
    -- right type constructor when reading vectorisation information from interface files).
    --
    defTyConDataCons (origTyCon, vectTyCon, isAbstract)
343 344 345 346
      = do
        { canonName <- mkLocalisedName mkVectTyConOcc origName
        ; if    origName == vectName                             -- Case (1)
             || vectName == canonName                            -- Case (2)
347
          then do
348 349 350 351 352 353 354 355 356 357 358
          { defTyCon origTyCon vectTyCon                         -- T  --> vT
          ; defDataCons                                          -- Ci --> vCi
          ; return Nothing
          }
          else do                                                 -- Case (3)
          { let synTyCon = mkSyn canonName (mkTyConTy vectTyCon)  -- type S = vT
          ; defTyCon origTyCon synTyCon                           -- T  --> S
          ; defDataCons                                           -- Ci --> vCi
          ; return $ Just synTyCon
          }
        }
359 360 361 362
      where
        origName  = tyConName origTyCon
        vectName  = tyConName vectTyCon

363
        mkSyn canonName ty = buildSynTyCon canonName [] (typeKind ty) [] ty
364

365 366
        defDataCons
          | isAbstract = return ()
367
          | otherwise
368 369 370
          = do { MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
               ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
               }
371 372


373
-- Helpers --------------------------------------------------------------------
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
374

375
buildTyConPADict :: TyCon -> CoAxiom Unbranched -> TyCon -> TyCon -> VM Var
376 377
buildTyConPADict vect_tc prepr_ax pdata_tc pdatas_tc
 = tyConRepr vect_tc >>= buildPADict vect_tc prepr_ax pdata_tc pdatas_tc
378 379

-- Produce a custom-made worker for the data constructors of a vectorised data type.  This includes
380 381 382
-- all data constructors that may be used in vectorised code — i.e., all data constructors of data
-- types with 'VECTORISE [SCALAR] type' pragmas with an explicit right-hand side.  Also adds a mapping
-- from the original to vectorised worker into the vectorisation map.
383 384 385 386
--
-- 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.
387 388
vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
vectDataConWorkers orig_tc vect_tc arr_tc
389
  = do { traceVt "Building vectorised worker for datatype" (ppr orig_tc)
390

391 392 393 394 395 396 397 398 399
       ; 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
400 401 402 403 404
    tyvars   = tyConTyVars vect_tc
    var_tys  = mkTyVarTys tyvars
    ty_args  = map Type var_tys
    res_ty   = mkTyConApp vect_tc var_tys

405 406 407 408
    cons     = tyConDataCons vect_tc
    arity    = length cons
    [arr_dc] = tyConDataCons arr_tc

409 410 411
    rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc

    mk_data_con con tys pre post
412 413 414
      = do dflags <- getDynFlags
           liftM2 (,) (vect_data_con con)
                      (lift_data_con tys pre post (mkDataConTag dflags con))
415

416 417 418 419 420 421 422
    sel_replicate len tag
      | arity > 1 = do
                      rep <- builtin (selReplicate arity)
                      return [rep `mkApps` [len, tag]]

      | otherwise = return []

423
    vect_data_con con = return $ mkConApp con ty_args
424
    lift_data_con tys pre_tys post_tys tag
425 426
      = do
          len  <- builtin liftingContext
Ian Lynagh's avatar
Ian Lynagh committed
427
          args <- mapM (newLocalVar (fsLit "xs"))
428
                  =<< mapM mkPDataType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
429

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

432 433
          pre   <- mapM emptyPD (concat pre_tys)
          post  <- mapM emptyPD (concat post_tys)
434 435 436 437

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
438
                 $ ty_args ++ sel ++ pre ++ map Var args ++ post
439 440 441

    def_worker data_con arg_tys mk_body
      = do
442
          arity <- polyArity tyvars
443 444
          body <- closedV
                . inBind orig_worker
445 446
                . polyAbstract tyvars $ \args ->
                  liftM (mkLams (tyvars ++ args) . vectorised)
447
                $ buildClosures tyvars [] [] arg_tys res_ty mk_body
448

449
          raw_worker <- mkVectId orig_worker (exprType body)
450
          let vect_worker = raw_worker `setIdUnfolding`
451
                              mkInlineUnfoldingWithArity arity body
452 453 454 455
          defGlobalVar orig_worker vect_worker
          return (vect_worker, body)
      where
        orig_worker = dataConWorkId data_con