Env.hs 20.4 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 OccName
36

37
import Util
38
import Outputable
39
import FastString
40
import MonadUtils
41

42
import Control.Monad
43
import Data.Maybe
44
import Data.List
45

46

47 48 49 50 51 52 53 54 55
-- Note [Pragmas to vectorise tycons]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
--
-- VECTORISE pragmas for type constructors cover three different flavours of vectorising data type
-- constructors:
--
-- (1) Data type constructor 'T' that may be used in vectorised code, where 'T' represents itself,
--     but the representation of 'T' is opaque in vectorised code.  
--
56
--     An example is the treatment of 'Int'.  'Int's can be used in vectorised code and remain
57 58 59
--     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.
60
--
61 62 63 64 65 66
--     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
--     by the vectoriser).
--
--     Type constructors declared with {-# VECTORISE SCALAR type T #-} are treated in this manner.
--     (The vectoriser never treats a type constructor automatically in this manner.)
--
67 68 69 70 71 72 73 74 75 76 77 78 79 80
-- (2) Data type constructor 'T' that may be used in vectorised code, where 'T' is represented by an
--     explicitly given 'Tv', but the representation of 'T' is opaque in vectorised code.  
--
--     An example is the treatment of '[::]'.  '[::]'s 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.
--
--     '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 = T' #-} are treated in this 
--     manner. (The vectoriser never treats a type constructor automatically in this manner.)
--
-- (3) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
81 82 83 84 85
--     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.
86
--
87 88
--     An example is the treatment of 'Ordering' and '[]'.  The former remains unchanged by
--     vectorisation, whereas the latter is fully vectorised.
89 90 91 92 93

--     'PData' and 'PRepr' instances are automatically generated by the vectoriser.
--
--     Type constructors declared with {-# VECTORISE type T #-} are treated in this manner.
--
94
-- (4) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
95 96
--     code, where 'T' is represented by an explicitly given 'Tv' whose constructors 'Cvn' represent
--     the original constructors in vectorised code.  As a special case, we can have 'Tv = T'
97
--
98 99
--     An example is the treatment of 'Bool', which is represented by itself in vectorised code
--     (as it cannot embed any parallel arrays).  However, we do not want any automatic generation
100
--     of class and family instances, which is why Case (3) does not apply.
101
--
102 103
--     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
--     by the vectoriser).
104
--
105
--     Type constructors declared with {-# VECTORISE type T = T' #-} are treated in this manner.
106
--
107 108
-- 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
109
-- constructor.  We generally produce a vectorised version of the data type and data constructor.
110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126
-- We do not generate 'PData' and 'PRepr' instances for class type constructors.  This pragma is the
-- default for all type classes declared in this module, but the pragma can also be used explitly on
-- imported classes.

-- Note [Vectorising classes]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~
--
-- We vectorise classes essentially by just vectorising their desugared Core representation, but we
-- do generate a 'Class' structure along the way (see 'Vectorise.Type.TyConDecl.vectTyConDecl').
--
-- Here is an example illustrating the mapping — assume
--
--   class Num a where
--     (+) :: a -> a -> a
--
-- It desugars to
--
127
--   data Num a = D:Num { (+) :: a -> a -> a }
128 129 130
--
-- which we vectorise to
--
131
--  data V:Num a = D:V:Num { ($v+) :: PArray a :-> PArray a :-> PArray a }
132 133 134
--
-- while adding the following entries to the vectorisation map:
--
135 136
--   tycon  : Num   --> V:Num
--   datacon: D:Num --> D:V:Num
137
--   var    : (+) --> ($v+)
138

139
-- |Vectorise type constructor including class type constructors.
140
--
141
vectTypeEnv :: [TyCon]                  -- Type constructors defined in this module
142
            -> [CoreVect]               -- All 'VECTORISE [SCALAR] type' declarations in this module
143
            -> [CoreVect]               -- All 'VECTORISE class' declarations in this module
144
            -> VM ( [TyCon]             -- old TyCons ++ new TyCons
145 146
                  , [FamInst]           -- New type family instances.
                  , [(Var, CoreExpr)])  -- New top level bindings.
147
vectTypeEnv tycons vectTypeDecls vectClassDecls
148
  = do { traceVt "** vectTypeEnv" $ ppr tycons
149

150
       ; let   -- {-# VECTORISE SCALAR type T -#} (imported and local tycons)
151
             localAbstractTyCons    = [tycon | VectType True tycon Nothing <- vectTypeDecls]
152 153

               -- {-# VECTORISE type T -#} (ONLY the imported tycons)
154 155
             impVectTyCons          = (   [tycon | VectType False tycon Nothing <- vectTypeDecls]
                                       ++ [tycon | VectClass tycon              <- vectClassDecls])
156 157
                                      \\ tycons

158 159 160
               -- {-# VECTORISE [SCALAR] type T = T' -#} (imported and local tycons)
             vectTyConsWithRHS      = [ (tycon, rhs, isAbstract) 
                                      | VectType isAbstract tycon (Just rhs) <- vectTypeDecls]
161 162 163

               -- filter VECTORISE SCALAR tycons and VECTORISE tycons with explicit rhses
             vectSpecialTyConNames  = mkNameSet . map tyConName $
164 165
                                        localAbstractTyCons ++ map fst3 vectTyConsWithRHS
             notVectSpecialTyCon tc = not $ (tyConName tc) `elemNameSet` vectSpecialTyConNames
166

167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183
         -- Build a map containing all vectorised type constructor.  If they are scalar, they are
         -- mapped to 'False' (vectorised type constructor == original type constructor).
       ; allScalarTyConNames <- globalScalarTyCons  -- covers both current and imported modules
       ; vectTyCons          <- globalVectTyCons
       ; let vectTyConBase    = mapNameEnv (const True) vectTyCons    -- by default fully vectorised
             vectTyConFlavour = vectTyConBase 
                                `plusNameEnv` 
                                mkNameEnv [ (tyConName tycon, True) 
                                          | (tycon, _, _) <- vectTyConsWithRHS]
                                `plusNameEnv`
                                mkNameEnv [ (tcName, False)           -- original representation
                                          | tcName <- nameSetToList allScalarTyConNames]
                                `plusNameEnv`
                                mkNameEnv [ (tyConName tycon, False)  -- original representation
                                          | tycon <- localAbstractTyCons]
                                            

184 185 186 187
           -- 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
188
           -- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as
189 190 191
           -- these are being handled separately.  NB: Some type constructors may be marked SCALAR
           -- /and/ have an explicit right-hand side.)
           --
192
           -- Furthermore, 'drop_tcs' are those type constructors that we cannot vectorise.
193
       ; let maybeVectoriseTyCons           = filter notVectSpecialTyCon tycons ++ impVectTyCons
194
             (conv_tcs, keep_tcs, drop_tcs) = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
195
             
196
       ; traceVt " VECT SCALAR    : " $ ppr localAbstractTyCons
197
       ; traceVt " VECT [class]   : " $ ppr impVectTyCons
198
       ; traceVt " VECT with rhs  : " $ ppr (map fst3 vectTyConsWithRHS)
199
       ; traceVt " -- after classification (local and VECT [class] tycons) --" empty
200 201
       ; traceVt " reuse          : " $ ppr keep_tcs
       ; traceVt " convert        : " $ ppr conv_tcs
202 203
       
           -- warn the user about unvectorised type constructors
204 205 206 207
       ; 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) $
208
           emitVt "Warning: cannot vectorise these type constructors:" $ 
209
             pprQuotedList drop_tcs_nosyn $$ explanation
210

211
       ; mapM_ addGlobalScalarTyCon keep_tcs
212

213 214 215 216 217 218 219 220 221 222 223 224 225 226
       ; let mapping =      
                    -- Type constructors that we don't need to vectorise, use the same
                    -- representation in both unvectorised and vectorised code; they are not
                    -- abstract.
                  [(tycon, tycon, False) | tycon <- keep_tcs]
                    -- We do the same for type constructors declared VECTORISE SCALAR /without/
                    -- an explicit right-hand side, but ignore their representation (data
                    -- constructors) as they are abstract.
               ++ [(tycon, tycon, True) | tycon <- localAbstractTyCons]
                    -- Type constructors declared VECTORISE /with/ an explicit vectorised type,
                    -- we map from the original to the given type; whether they are abstract depends
                    -- on whether the vectorisation declaration was SCALAR.
               ++ vectTyConsWithRHS
       ; syn_tcs <- catMaybes <$> mapM defTyConDataCons mapping
227

228 229
           -- 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.)
230
       ; new_tcs <- vectTyConDecls conv_tcs
231 232 233 234 235 236
       
       ; 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)
237 238 239 240

           -- 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.
241 242
       ; let orig_tcs = filter (not . isClassTyCon) $ keep_tcs ++ conv_tcs
             vect_tcs = filter (not . isClassTyCon) $ keep_tcs ++ new_tcs
243 244 245

           -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
           -- type constructors with vectorised representations.
246
       ; reprs      <- mapM tyConRepr vect_tcs
247 248 249 250 251 252 253 254 255
       ; 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
             
256 257
       ; updGEnv $ extendFamEnv fam_insts

258 259 260 261
           -- 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).
262 263
       ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
           do { defTyConPAs (zipLazy vect_tcs dfuns)
264

265
                  -- Query the 'PData' instance type constructors for type constructors that have a
266
                  -- VECTORISE pragma with an explicit right-hand side (this is Item (4) of
267
                  -- "Note [Pragmas to vectorise tycons]" above).
268 269 270
              ; let (withRHS_non_abstract, vwithRHS_non_abstract) 
                      = unzip [(tycon, vtycon) | (tycon, vtycon, False) <- vectTyConsWithRHS]
              ; pdata_withRHS_tcs <- mapM pdataReprTyConExact withRHS_non_abstract
271

272
                  -- Build workers for all vectorised data constructors (except abstract ones)
273
              ; sequence_ $
274 275
                  zipWith3 vectDataConWorkers (orig_tcs  ++ withRHS_non_abstract)
                                              (vect_tcs  ++ vwithRHS_non_abstract)
276 277
                                              (pdata_tcs ++ pdata_withRHS_tcs)

278
                  -- Build a 'PA' dictionary for all type constructors (except abstract ones & those
279 280 281 282
                  -- defined with an explicit right-hand side where the dictionary is user-supplied)
              ; dfuns <- sequence $
                           zipWith4 buildTyConPADict
                                    vect_tcs
283
                                    repr_axs
284 285
                                    pdata_tcs
                                    pdatas_tcs
286 287 288 289 290

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

291 292
           -- Return the vectorised variants of type constructors as well as the generated instance
           -- type constructors, family instances, and dfun bindings.
293 294
       ; return ( new_tcs ++ pdata_tcs ++ pdatas_tcs ++ syn_tcs
                , fam_insts, binds)
295
       }
296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328
  where
    fst3 (a, _, _) = a

    -- 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)
      = 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
              }
           }
      where
        origName  = tyConName origTyCon
        vectName  = tyConName vectTyCon

329
        mkSyn canonName ty = mkSynTyCon canonName (typeKind ty) [] (SynonymTyCon ty) NoParentTyCon
330 331 332 333 334 335 336
        
        defDataCons
          | isAbstract = return ()
          | otherwise  
          = do { MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
               ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
               }
337 338


339
-- Helpers --------------------------------------------------------------------
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
340

341 342 343
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
344 345 346 347 348 349 350 351 352

-- Produce a custom-made worker for the data constructors of a vectorised data type.  This includes
-- all data constructors that may be used in vetcorised code — i.e., all data constructors of data
-- types other than scalar ones.  Also adds a mapping from the original to vectorised worker into
-- the vectorisation map.
--
-- FIXME: It's not nice that we need create a special worker after the data constructors has
--   already been constructed.  Also, I don't think the worker is properly added to the data
--   constructor.  Seems messy.
353 354
vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
vectDataConWorkers orig_tc vect_tc arr_tc
355 356 357 358 359 360 361 362 363 364 365
  = 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
366 367 368 369 370
    tyvars   = tyConTyVars vect_tc
    var_tys  = mkTyVarTys tyvars
    ty_args  = map Type var_tys
    res_ty   = mkTyConApp vect_tc var_tys

371 372 373 374
    cons     = tyConDataCons vect_tc
    arity    = length cons
    [arr_dc] = tyConDataCons arr_tc

375 376 377 378
    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
379
                   (lift_data_con tys pre post (mkDataConTag con))
380

381 382 383 384 385 386 387
    sel_replicate len tag
      | arity > 1 = do
                      rep <- builtin (selReplicate arity)
                      return [rep `mkApps` [len, tag]]

      | otherwise = return []

388
    vect_data_con con = return $ mkConApp con ty_args
389
    lift_data_con tys pre_tys post_tys tag
390 391
      = do
          len  <- builtin liftingContext
Ian Lynagh's avatar
Ian Lynagh committed
392
          args <- mapM (newLocalVar (fsLit "xs"))
393
                  =<< mapM mkPDataType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
394

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

397 398
          pre   <- mapM emptyPD (concat pre_tys)
          post  <- mapM emptyPD (concat post_tys)
399 400 401 402

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
403
                 $ ty_args ++ sel ++ pre ++ map Var args ++ post
404 405 406

    def_worker data_con arg_tys mk_body
      = do
407
          arity <- polyArity tyvars
408 409
          body <- closedV
                . inBind orig_worker
410 411
                . polyAbstract tyvars $ \args ->
                  liftM (mkLams (tyvars ++ args) . vectorised)
412
                $ buildClosures tyvars [] [] arg_tys res_ty mk_body
413

414
          raw_worker <- mkVectId orig_worker (exprType body)
415
          let vect_worker = raw_worker `setIdUnfolding`
416
                              mkInlineUnfolding (Just arity) body
417 418 419 420
          defGlobalVar orig_worker vect_worker
          return (vect_worker, body)
      where
        orig_worker = dataConWorkId data_con