Env.hs 22 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
import TyCon
29
import CoAxiom
30
import Type
31
import FamInstEnv
32
import Id
33
import MkId
34
import NameEnv
35
import NameSet
36
import UniqFM
37
import OccName
38
import Unique
39

40
import Util
41
import Outputable
42
import DynFlags
43
import FastString
44
import MonadUtils
45

46
import Control.Monad
47
import Data.Maybe
48
import Data.List
49

50

51 52 53
-- Note [Pragmas to vectorise tycons]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
--
54 55 56 57 58 59 60 61 62 63 64 65 66
-- 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
67 68
-- constructors:
--
69 70 71 72 73 74
-- (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.
75
--
76 77
--     An example is the treatment of 'Ordering' and '[]'.  The former remains unchanged by
--     vectorisation, whereas the latter is fully vectorised.
78
--
79
--     'PData' and 'PRepr' instances are automatically generated by the vectoriser.
80
--
81
--     Type constructors declared with {-# VECTORISE type T #-} are treated in this manner.
82
--
83
-- (2) Data type constructor 'T' that may be used in vectorised code, where 'T' is represented by an
84 85
--     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).  
86
--
87 88 89
--     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.
90 91 92 93
--
--     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
--     by the vectoriser).
--
94
--     Type constructors declared with {-# VECTORISE type T = Tv #-} are treated in this manner
95 96
--     manner. (The vectoriser never treats a type constructor automatically in this manner.)
--
97 98 99 100
-- (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.
101
--
102 103 104
--     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.
105
--
106 107 108 109
--     '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.
110
--
111 112 113 114
-- (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.
115
--
116 117 118 119
--     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.
120
--
121 122
--     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
--     by the vectoriser).
123
--
124 125
--     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.)
126
--
127 128
-- 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
129
-- constructor.  We generally produce a vectorised version of the data type and data constructor.
130
-- We do not generate 'PData' and 'PRepr' instances for class type constructors.  This pragma is the
131 132
-- default for all type classes declared in a vectorised module, but the pragma can also be used
-- explitly on imported classes.
133 134 135 136 137 138 139 140 141 142 143 144 145 146

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

159
-- |Vectorise type constructor including class type constructors.
160
--
161 162 163 164
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
165
                  , [FamInst]            -- New type family instances.
166
                  , [(Var, CoreExpr)])   -- New top level bindings.
167
vectTypeEnv tycons vectTypeDecls vectClassDecls
168
  = do { traceVt "** vectTypeEnv" $ ppr tycons
169

170
       ; let   -- {-# VECTORISE type T -#} (ONLY the imported tycons)
171 172
             impVectTyCons          = (   [tycon | VectType False tycon Nothing <- vectTypeDecls]
                                       ++ [tycon | VectClass tycon              <- vectClassDecls])
173
                                      \\ tycons
174
               
175 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/)
             scalarTyConsWithRHS    = [ (tycon, rhs) 
                                      | VectType True  tycon (Just rhs) <- vectTypeDecls]
182

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

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

192 193 194
         -- 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'.
195
       ; vectTyCons          <- globalVectTyCons
196 197
       ; let vectTyConBase    = mapUFM_Directly isDistinct vectTyCons    -- 'True' iff tc /= V[[tc]]
             isDistinct u tc  = u /= getUnique tc
198 199 200
             vectTyConFlavour = vectTyConBase 
                                `plusNameEnv` 
                                mkNameEnv [ (tyConName tycon, True) 
201
                                          | (tycon, _) <- vectTyConsWithRHS ++ scalarTyConsWithRHS]
202 203
                                `plusNameEnv`
                                mkNameEnv [ (tyConName tycon, False)  -- original representation
204
                                          | tycon <- scalarTyConsNoRHS]
205 206
                                            

207 208 209 210
           -- 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
211
           -- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as
212 213 214
           -- these are being handled separately.  NB: Some type constructors may be marked SCALAR
           -- /and/ have an explicit right-hand side.)
           --
215 216 217
           -- 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.
218
       ; parallelTyCons <- (`addListToNameSet` map (tyConName . fst) vectTyConsWithRHS) <$>
219
                             globalParallelTyCons
220
       ; let maybeVectoriseTyCons = filter notVectSpecialTyCon tycons ++ impVectTyCons
221
             (conv_tcs, keep_tcs, par_tcs, drop_tcs)
222
               = classifyTyCons vectTyConFlavour parallelTyCons maybeVectoriseTyCons
223 224 225

       ; traceVt " known parallel : " $ ppr parallelTyCons
       ; traceVt " VECT SCALAR    : " $ ppr (scalarTyConsNoRHS ++ map fst scalarTyConsWithRHS)
226
       ; traceVt " VECT [class]   : " $ ppr impVectTyCons
227
       ; traceVt " VECT with rhs  : " $ ppr (map fst (vectTyConsWithRHS ++ scalarTyConsWithRHS))
228
       ; traceVt " -- after classification (local and VECT [class] tycons) --" empty
229 230
       ; traceVt " reuse          : " $ ppr keep_tcs
       ; traceVt " convert        : " $ ppr conv_tcs
231 232
       
           -- warn the user about unvectorised type constructors
233 234 235 236
       ; 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) $
237
           emitVt "Warning: cannot vectorise these type constructors:" $ 
238
             pprQuotedList drop_tcs_nosyn $$ explanation
239

240
       ; mapM_ addParallelTyConAndCons $ par_tcs ++ map fst vectTyConsWithRHS
241

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

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

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

           -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
           -- type constructors with vectorised representations.
271
       ; reprs      <- mapM tyConRepr vect_tcs
272 273 274 275 276 277 278 279 280
       ; 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
             
281 282
       ; updGEnv $ extendFamEnv fam_insts

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

290
                  -- Query the 'PData' instance type constructors for type constructors that have a
291 292 293
                  -- 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
294

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

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

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

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

    -- 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)
337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352
      = 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
          }
        }
353 354 355 356
      where
        origName  = tyConName origTyCon
        vectName  = tyConName vectTyCon

357
        mkSyn canonName ty = mkSynTyCon canonName (typeKind ty) [] [] (SynonymTyCon ty) NoParentTyCon
358 359 360 361 362 363 364
        
        defDataCons
          | isAbstract = return ()
          | otherwise  
          = do { MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
               ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
               }
365 366


367
-- Helpers --------------------------------------------------------------------
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
368

369
buildTyConPADict :: TyCon -> CoAxiom Unbranched -> TyCon -> TyCon -> VM Var
370 371
buildTyConPADict vect_tc prepr_ax pdata_tc pdatas_tc
 = tyConRepr vect_tc >>= buildPADict vect_tc prepr_ax pdata_tc pdatas_tc
372 373

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

399 400 401 402
    cons     = tyConDataCons vect_tc
    arity    = length cons
    [arr_dc] = tyConDataCons arr_tc

403 404 405
    rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc

    mk_data_con con tys pre post
406 407 408
      = do dflags <- getDynFlags
           liftM2 (,) (vect_data_con con)
                      (lift_data_con tys pre post (mkDataConTag dflags con))
409

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

      | otherwise = return []

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

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

426 427
          pre   <- mapM emptyPD (concat pre_tys)
          post  <- mapM emptyPD (concat post_tys)
428 429 430 431

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
432
                 $ ty_args ++ sel ++ pre ++ map Var args ++ post
433 434 435

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

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