Env.hs 22.1 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 11
) where
  
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 87
--     explicitly given 'Tv', but the representation of 'T' is opaque in vectorised code (i.e., the
--     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 127
--     Type constructors declared with {-# VECTORISE SCALAR type T = Tv #-} are treated in this 
--     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 182 183
               -- {-# 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/)
             scalarTyConsWithRHS    = [ (tycon, rhs) 
                                      | 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 189
               -- Check that is not a VECTORISE SCALAR tycon nor VECTORISE tycons with explicit rhs?
             vectSpecialTyConNames  = mkNameSet . map tyConName $ 
190 191
                                        scalarTyConsNoRHS ++ 
                                        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
       ; let explanation    = ptext (sLit "(They use unsupported language extensions") $$
                              ptext (sLit "or depend on type constructors that are not vectorised)")
237 238
             drop_tcs_nosyn = filter (not . isTypeFamilyTyCon) .
                              filter (not . isTypeSynonymTyCon) $ drop_tcs
239
       ; unless (null drop_tcs_nosyn) $
240
           emitVt "Warning: cannot vectorise these type constructors:" $ 
241
             pprQuotedList drop_tcs_nosyn $$ explanation
242

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

245
       ; let mapping =      
246 247
                    -- 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
248 249
                    -- representation in both unvectorised and vectorised code; they are not
                    -- abstract.
250
                  [(tycon, tycon, False) | tycon <- keep_tcs ++ scalarTyConsNoRHS]
251
                    -- We do the same for type constructors declared VECTORISE SCALAR /without/
252
                    -- an explicit right-hand side
253
               ++ [(tycon, vTycon, True) | (tycon, vTycon) <- vectTyConsWithRHS ++ scalarTyConsWithRHS]
254
       ; syn_tcs <- catMaybes <$> mapM defTyConDataCons mapping
255

256 257
           -- 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.)
258
       ; new_tcs <- vectTyConDecls conv_tcs
259 260 261 262 263 264
       
       ; 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)
265 266 267 268

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

           -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
           -- type constructors with vectorised representations.
274
       ; reprs      <- mapM tyConRepr vect_tcs
275 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
       ; updGEnv $ extendFamEnv fam_insts

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

293
                  -- Query the 'PData' instance type constructors for type constructors that have a
294 295 296
                  -- 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
297

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

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

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

317 318
           -- Return the vectorised variants of type constructors as well as the generated instance
           -- type constructors, family instances, and dfun bindings.
319 320
       ; return ( new_tcs ++ pdata_tcs ++ pdatas_tcs ++ syn_tcs
                , fam_insts, binds)
321
       }
322
  where
323 324 325
    addParallelTyConAndCons tycon
      = do
        { addGlobalParallelTyCon tycon
326
        ; mapM_ addGlobalParallelVar . concatMap dataConImplicitIds . tyConDataCons $ tycon
327
        }
328 329 330 331 332 333 334 335 336 337 338 339

    -- Add a mapping from the original to vectorised type constructor to the vectorisation map.  
    -- Unless the type constructor is abstract, also mappings from the orignal's data constructors
    -- 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)
340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355
      = do
        { canonName <- mkLocalisedName mkVectTyConOcc origName
        ; if    origName == vectName                             -- Case (1)
             || vectName == canonName                            -- Case (2)
          then do 
          { 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
          }
        }
356 357 358 359
      where
        origName  = tyConName origTyCon
        vectName  = tyConName vectTyCon

360
        mkSyn canonName ty = mkSynonymTyCon canonName (typeKind ty) [] [] ty
361 362 363 364 365 366 367
        
        defDataCons
          | isAbstract = return ()
          | otherwise  
          = do { MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
               ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
               }
368 369


370
-- Helpers --------------------------------------------------------------------
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
371

372
buildTyConPADict :: TyCon -> CoAxiom Unbranched -> TyCon -> TyCon -> VM Var
373 374
buildTyConPADict vect_tc prepr_ax pdata_tc pdatas_tc
 = tyConRepr vect_tc >>= buildPADict vect_tc prepr_ax pdata_tc pdatas_tc
375 376

-- Produce a custom-made worker for the data constructors of a vectorised data type.  This includes
377 378 379
-- 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.
380 381 382 383
--
-- 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.
384 385
vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
vectDataConWorkers orig_tc vect_tc arr_tc
386 387 388 389 390 391 392 393 394 395 396
  = 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
397 398 399 400 401
    tyvars   = tyConTyVars vect_tc
    var_tys  = mkTyVarTys tyvars
    ty_args  = map Type var_tys
    res_ty   = mkTyConApp vect_tc var_tys

402 403 404 405
    cons     = tyConDataCons vect_tc
    arity    = length cons
    [arr_dc] = tyConDataCons arr_tc

406 407 408
    rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc

    mk_data_con con tys pre post
409 410 411
      = do dflags <- getDynFlags
           liftM2 (,) (vect_data_con con)
                      (lift_data_con tys pre post (mkDataConTag dflags con))
412

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

      | otherwise = return []

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

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

429 430
          pre   <- mapM emptyPD (concat pre_tys)
          post  <- mapM emptyPD (concat post_tys)
431 432 433 434

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
435
                 $ ty_args ++ sel ++ pre ++ map Var args ++ post
436 437 438

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

446
          raw_worker <- mkVectId orig_worker (exprType body)
447
          let vect_worker = raw_worker `setIdUnfolding`
448
                              mkInlineUnfolding (Just arity) body
449 450 451 452
          defGlobalVar orig_worker vect_worker
          return (vect_worker, body)
      where
        orig_worker = dataConWorkId data_con