Env.hs 21.5 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
import UniqFM
36
import OccName
37
import Unique
38

39
import Util
40
import Outputable
41
import FastString
42
import MonadUtils
43

44
import Control.Monad
45
import Data.Maybe
46
import Data.List
47

48

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

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

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

168
       ; let   -- {-# VECTORISE type T -#} (ONLY the imported tycons)
169 170
             impVectTyCons          = (   [tycon | VectType False tycon Nothing <- vectTypeDecls]
                                       ++ [tycon | VectClass tycon              <- vectClassDecls])
171
                                      \\ tycons
172 173 174 175
               
               -- {-# VECTORISE [SCALAR] type T = Tv -#} (imported & local tycons with an /RHS/)
             vectTyConsWithRHS      = [ (tycon, rhs, isScalar) 
                                      | VectType isScalar tycon (Just rhs) <- vectTypeDecls]
176

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

180 181 182
               -- Check that is not a VECTORISE SCALAR tycon nor VECTORISE tycons with explicit rhs?
             vectSpecialTyConNames  = mkNameSet . map tyConName $ 
                                        scalarTyConsNoRHS ++ map fst3 vectTyConsWithRHS
183
             notVectSpecialTyCon tc = not $ (tyConName tc) `elemNameSet` vectSpecialTyConNames
184

185 186 187
         -- 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'.
188
       ; vectTyCons          <- globalVectTyCons
189 190
       ; let vectTyConBase    = mapUFM_Directly isDistinct vectTyCons    -- 'True' iff tc /= V[[tc]]
             isDistinct u tc  = u /= getUnique tc
191 192 193 194 195 196
             vectTyConFlavour = vectTyConBase 
                                `plusNameEnv` 
                                mkNameEnv [ (tyConName tycon, True) 
                                          | (tycon, _, _) <- vectTyConsWithRHS]
                                `plusNameEnv`
                                mkNameEnv [ (tyConName tycon, False)  -- original representation
197
                                          | tycon <- scalarTyConsNoRHS]
198 199
                                            

200 201 202 203
           -- 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
204
           -- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as
205 206 207
           -- these are being handled separately.  NB: Some type constructors may be marked SCALAR
           -- /and/ have an explicit right-hand side.)
           --
208 209 210 211 212 213
           -- Furthermore, 'par_tcs' and 'drop_tcs' are those type constructors that we cannot
           -- vectorise, and of those, only the 'par_tcs' involve parallel arrays.
       ; parallelTyCons <- globalParallelTyCons
       ; let maybeVectoriseTyCons = filter notVectSpecialTyCon tycons ++ impVectTyCons
             (conv_tcs, keep_tcs, par_tcs, drop_tcs) 
               = classifyTyCons vectTyConFlavour parallelTyCons maybeVectoriseTyCons
214
             
215 216
       ; traceVt " VECT SCALAR    : " $ ppr (scalarTyConsNoRHS ++ 
                                             [tycon | (tycon, _, True) <- vectTyConsWithRHS])
217
       ; traceVt " VECT [class]   : " $ ppr impVectTyCons
218
       ; traceVt " VECT with rhs  : " $ ppr (map fst3 vectTyConsWithRHS)
219
       ; traceVt " -- after classification (local and VECT [class] tycons) --" empty
220 221
       ; traceVt " reuse          : " $ ppr keep_tcs
       ; traceVt " convert        : " $ ppr conv_tcs
222 223
       
           -- warn the user about unvectorised type constructors
224 225
       ; let explanation    = ptext (sLit "(They use unsupported language extensions") $$
                              ptext (sLit "or depend on type constructors that are not vectorised)")
226
             drop_tcs_nosyn = filter (not . isSynTyCon) (par_tcs ++ drop_tcs)
227
       ; unless (null drop_tcs_nosyn) $
228
           emitVt "Warning: cannot vectorise these type constructors:" $ 
229
             pprQuotedList drop_tcs_nosyn $$ explanation
230

231
       ; mapM_ addParallelTyConAndCons $ conv_tcs ++ par_tcs
232

233
       ; let mapping =      
234 235
                    -- 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
236 237
                    -- representation in both unvectorised and vectorised code; they are not
                    -- abstract.
238
                  [(tycon, tycon, False) | tycon <- keep_tcs ++ scalarTyConsNoRHS]
239
                    -- We do the same for type constructors declared VECTORISE SCALAR /without/
240 241
                    -- an explicit right-hand side
               ++ [(tycon, vTycon, True) | (tycon, vTycon, _) <- vectTyConsWithRHS]
242
       ; syn_tcs <- catMaybes <$> mapM defTyConDataCons mapping
243

244 245
           -- 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.)
246
       ; new_tcs <- vectTyConDecls conv_tcs
247 248 249 250 251 252
       
       ; 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)
253 254 255 256

           -- 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.
257 258
       ; let orig_tcs = filter (not . isClassTyCon) $ keep_tcs ++ conv_tcs
             vect_tcs = filter (not . isClassTyCon) $ keep_tcs ++ new_tcs
259 260 261

           -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
           -- type constructors with vectorised representations.
262
       ; reprs      <- mapM tyConRepr vect_tcs
263 264 265 266 267 268 269 270 271
       ; 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
             
272 273
       ; updGEnv $ extendFamEnv fam_insts

274 275 276 277
           -- 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).
278 279
       ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
           do { defTyConPAs (zipLazy vect_tcs dfuns)
280

281
                  -- Query the 'PData' instance type constructors for type constructors that have a
282 283 284
                  -- 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
285

286
                  -- Build workers for all vectorised data constructors (except abstract ones)
287
              ; sequence_ $
288 289 290
                  zipWith3 vectDataConWorkers (orig_tcs  ++ scalarTyConsNoRHS)
                                              (vect_tcs  ++ scalarTyConsNoRHS)
                                              (pdata_tcs ++ pdata_scalar_tcs)
291

292
                  -- Build a 'PA' dictionary for all type constructors (except abstract ones & those
293 294 295 296
                  -- defined with an explicit right-hand side where the dictionary is user-supplied)
              ; dfuns <- sequence $
                           zipWith4 buildTyConPADict
                                    vect_tcs
297
                                    repr_axs
298 299
                                    pdata_tcs
                                    pdatas_tcs
300 301 302 303 304

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

305 306
           -- Return the vectorised variants of type constructors as well as the generated instance
           -- type constructors, family instances, and dfun bindings.
307 308
       ; return ( new_tcs ++ pdata_tcs ++ pdatas_tcs ++ syn_tcs
                , fam_insts, binds)
309
       }
310 311
  where
    fst3 (a, _, _) = a
312 313 314 315 316 317
    
    addParallelTyConAndCons tycon
      = do
        { addGlobalParallelTyCon tycon
        ; mapM_ addGlobalParallelVar . concatMap dataConImplicitIds . tyConDataCons $ tycon
        }
318 319 320 321 322 323 324 325 326 327 328 329

    -- 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)
330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345
      = 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
          }
        }
346 347 348 349
      where
        origName  = tyConName origTyCon
        vectName  = tyConName vectTyCon

350
        mkSyn canonName ty = mkSynTyCon canonName (typeKind ty) [] (SynonymTyCon ty) NoParentTyCon
351 352 353 354 355 356 357
        
        defDataCons
          | isAbstract = return ()
          | otherwise  
          = do { MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
               ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
               }
358 359


360
-- Helpers --------------------------------------------------------------------
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
361

362 363 364
buildTyConPADict :: TyCon -> CoAxiom -> TyCon -> TyCon -> VM Var
buildTyConPADict vect_tc prepr_ax pdata_tc pdatas_tc
 = tyConRepr vect_tc >>= buildPADict vect_tc prepr_ax pdata_tc pdatas_tc
365 366

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

392 393 394 395
    cons     = tyConDataCons vect_tc
    arity    = length cons
    [arr_dc] = tyConDataCons arr_tc

396 397 398 399
    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
400
                   (lift_data_con tys pre post (mkDataConTag con))
401

402 403 404 405 406 407 408
    sel_replicate len tag
      | arity > 1 = do
                      rep <- builtin (selReplicate arity)
                      return [rep `mkApps` [len, tag]]

      | otherwise = return []

409
    vect_data_con con = return $ mkConApp con ty_args
410
    lift_data_con tys pre_tys post_tys tag
411 412
      = do
          len  <- builtin liftingContext
Ian Lynagh's avatar
Ian Lynagh committed
413
          args <- mapM (newLocalVar (fsLit "xs"))
414
                  =<< mapM mkPDataType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
415

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

418 419
          pre   <- mapM emptyPD (concat pre_tys)
          post  <- mapM emptyPD (concat post_tys)
420 421 422 423

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
424
                 $ ty_args ++ sel ++ pre ++ map Var args ++ post
425 426 427

    def_worker data_con arg_tys mk_body
      = do
428
          arity <- polyArity tyvars
429 430
          body <- closedV
                . inBind orig_worker
431 432
                . polyAbstract tyvars $ \args ->
                  liftM (mkLams (tyvars ++ args) . vectorised)
433
                $ buildClosures tyvars [] [] arg_tys res_ty mk_body
434

435
          raw_worker <- mkVectId orig_worker (exprType body)
436
          let vect_worker = raw_worker `setIdUnfolding`
437
                              mkInlineUnfolding (Just arity) body
438 439 440 441
          defGlobalVar orig_worker vect_worker
          return (vect_worker, body)
      where
        orig_worker = dataConWorkId data_con