compiler/vectorise/Vectorise/Type/Env.hs

   1 {-# OPTIONS_GHC -XNoMonoLocalBinds #-}
   2
   3 -- Vectorise a modules type environment, the structure containing all type things defined in a
   4 -- module.
   5 --
   6 -- This extends the type environment with vectorised variants of data types and produces value
   7 -- bindings for worker functions and the like.
   8
   9 module Vectorise.Type.Env (
  10   vectTypeEnv,
  11 ) where
  12
  13 #include "HsVersions.h"
  14
  15 import Vectorise.Env
  16 import Vectorise.Vect
  17 import Vectorise.Monad
  18 import Vectorise.Builtins
  19 import Vectorise.Type.TyConDecl
  20 import Vectorise.Type.Classify
  21 import Vectorise.Type.PADict
  22 import Vectorise.Type.PData
  23 import Vectorise.Type.PRepr
  24 import Vectorise.Type.Repr
  25 import Vectorise.Utils
  26
  27 import CoreSyn
  28 import CoreUtils
  29 import CoreUnfold
  30 import DataCon
  31 import TyCon
  32 import Type
  33 import FamInstEnv
  34 import Id
  35 import MkId
  36 import NameEnv
  37 import NameSet
  38
  39 import Util
  40 import Outputable
  41 import FastString
  42 import MonadUtils
  43 import Control.Monad
  44 import Data.List
  45
  46
  47 -- Note [Pragmas to vectorise tycons]
  48 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  49 --
  50 -- VECTORISE pragmas for type constructors cover three different flavours of vectorising data type
  51 -- constructors:
  52 --
  53 -- (1) Data type constructor 'T' that may be used in vectorised code, where 'T' represents itself,
  54 --     but the representation of 'T' is opaque in vectorised code.
  55 --
  56 --     An example is the treatment of 'Int'.  'Int's can be used in vectorised code and remain
  57 --     unchanged by vectorisation.  However, the representation of 'Int' by the 'I#' data
  58 --     constructor wrapping an 'Int#' is not exposed in vectorised code.  Instead, computations
  59 --     involving the representation need to be confined to scalar code.
  60 --
  61 --     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
  62 --     by the vectoriser).
  63 --
  64 --     Type constructors declared with {-# VECTORISE SCALAR type T #-} are treated in this manner.
  65 --     (The vectoriser never treats a type constructor automatically in this manner.)
  66 --
  67 -- (2) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
  68 --     code, where 'T' and the 'Cn' are automatically vectorised in the same manner as data types
  69 --     declared in a vectorised module.  This includes the case where the vectoriser determines that
  70 --     the original representation of 'T' may be used in vectorised code (as it does not embed any
  71 --     parallel arrays.)  This case is for type constructors that are *imported* from a non-
  72 --     vectorised module, but that we want to use with full vectorisation support.
  73 --
  74 --     An example is the treatment of 'Ordering' and '[]'.  The former remains unchanged by
  75 --     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 -- (3) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
  82 --     code, where 'T' is represented by an explicitly given 'Tv' whose constructors 'Cvn' represent
  83 --     the original constructors in vectorised code.  As a special case, we can have 'Tv = T'
  84 --
  85 --     An example is the treatment of 'Bool', which is represented by itself in vectorised code
  86 --     (as it cannot embed any parallel arrays).  However, we do not want any automatic generation
  87 --     of class and family instances, which is why Case (2) does not apply.
  88 --
  89 --     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
  90 --     by the vectoriser).
  91 --
  92 --     Type constructors declared with {-# VECTORISE type T = T' #-} are treated in this manner.
  93 --
  94 -- In addition, we have also got a single pragma form for type classes: {-# VECTORISE class C #-}.  It
  95 -- implies that the class type constructor may be used in vectorised code together with its data
  96 -- constructor.  We generally produce a vectorised version of the data type and data constructor.
  97 -- We do not generate 'PData' and 'PRepr' instances for class type constructors.
  98
  99 -- |Vectorise a type environment.
 100 --
 101 vectTypeEnv :: [TyCon]                  -- TyCons defined in this module
 102             -> [CoreVect]               -- All 'VECTORISE [SCALAR] type' declarations in this module
 103             -> VM ( [TyCon]             -- old TyCons ++ new TyCons
 104                   , [FamInst]           -- New type family instances.
 105                   , [(Var, CoreExpr)])  -- New top level bindings.
 106 vectTypeEnv tycons vectTypeDecls
 107   = do { traceVt "** vectTypeEnv" $ ppr tycons
 108
 109          -- Build a map containing all vectorised type constructor.  If they are scalar, they are
 110          -- mapped to 'False' (vectorised type constructor == original type constructor).
 111        ; allScalarTyConNames <- globalScalarTyCons  -- covers both current and imported modules
 112        ; vectTyCons          <- globalVectTyCons
 113        ; let vectTyConBase    = mapNameEnv (const True) vectTyCons   -- by default fully vectorised
 114              vectTyConFlavour = foldNameSet (\n env -> extendNameEnv env n False) vectTyConBase
 115                                             allScalarTyConNames
 116
 117        ; let   -- {-# VECTORISE SCALAR type T -#} (imported and local tycons)
 118              localScalarTyCons      = [tycon | VectType True  tycon Nothing <- vectTypeDecls]
 119
 120                -- {-# VECTORISE type T -#} (ONLY the imported tycons)
 121              impVectTyCons          = [tycon | VectType False tycon Nothing <- vectTypeDecls]
 122                                       \\ tycons
 123
 124                -- {-# VECTORISE type T = ty -#} (imported and local tycons)
 125              vectTyConsWithRHS      = [ (tycon, rhs)
 126                                       | VectType False tycon (Just rhs) <- vectTypeDecls]
 127
 128                -- filter VECTORISE SCALAR tycons and VECTORISE tycons with explicit rhses
 129              vectSpecialTyConNames  = mkNameSet . map tyConName $
 130                                         localScalarTyCons ++ map fst vectTyConsWithRHS
 131              notLocalScalarTyCon tc = not $ (tyConName tc) `elemNameSet` vectSpecialTyConNames
 132
 133            -- Split the list of 'TyCons' into the ones (1) that we must vectorise and those (2)
 134            -- that we could, but don't need to vectorise.  Type constructors that are not data
 135            -- type constructors or use non-Haskell98 features are being dropped.  They may not
 136            -- appear in vectorised code.  (We also drop the local type constructors appearing in a
 137            -- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as
 138            -- these are being handled separately.)
 139        ; let maybeVectoriseTyCons   = filter notLocalScalarTyCon tycons ++ impVectTyCons
 140              (conv_tcs, keep_tcs)   = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
 141              orig_tcs               = keep_tcs ++ conv_tcs
 142
 143        ; traceVt " VECT SCALAR    : " $ ppr localScalarTyCons
 144        ; traceVt " VECT with rhs  : " $ ppr (map fst vectTyConsWithRHS)
 145        ; traceVt " reuse          : " $ ppr keep_tcs
 146        ; traceVt " convert        : " $ ppr conv_tcs
 147
 148        ; let defTyConDataCons origTyCon vectTyCon
 149                = do { defTyCon origTyCon vectTyCon
 150                     ; MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
 151                     ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
 152                     }
 153
 154            -- For the type constructors that we don't need to vectorise, we use the original
 155            -- representation in both unvectorised and vectorised code.
 156        ; zipWithM_ defTyConDataCons keep_tcs keep_tcs
 157
 158            -- We do the same for type constructors declared VECTORISE SCALAR, while ignoring their
 159            -- representation (data constructors) — see "Note [Pragmas to vectorise tycons]".
 160        ; zipWithM_ defTyCon localScalarTyCons localScalarTyCons
 161
 162            -- For type constructors declared VECTORISE with an explicit vectorised type, we use the
 163            -- explicitly given type in vectorised code and map data constructors one for one — see
 164            -- "Note [Pragmas to vectorise tycons]".
 165        ; mapM_ (uncurry defTyConDataCons) vectTyConsWithRHS
 166
 167            -- Vectorise all the data type declarations that we can and must vectorise.
 168        ; new_tcs <- vectTyConDecls conv_tcs
 169
 170            -- We don't need new representation types for dictionary constructors. The constructors
 171            -- are always fully applied, and we don't need to lift them to arrays as a dictionary
 172            -- of a particular type always has the same value.
 173        ; let vect_tcs = filter (not . isClassTyCon)
 174                       $ keep_tcs ++ new_tcs
 175
 176            -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
 177            -- type constructors with vectorised representations.
 178        ; reprs     <- mapM tyConRepr vect_tcs
 179        ; repr_tcs  <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
 180        ; pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
 181        ; let inst_tcs  = repr_tcs ++ pdata_tcs
 182              fam_insts = map mkLocalFamInst inst_tcs
 183        ; updGEnv $ extendFamEnv fam_insts
 184
 185            -- Generate dfuns for the 'PA' instances of the vectorised type constructors and
 186            -- associate the type constructors with their dfuns in the global environment.  We get
 187            -- back the dfun bindings (which we will subsequently inject into the modules toplevel).
 188        ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
 189            do { defTyConPAs (zipLazy vect_tcs dfuns)
 190               ; dfuns <- sequence
 191                       $  zipWith4 buildTyConBindings
 192                                   orig_tcs
 193                                   vect_tcs
 194                                   repr_tcs
 195                                   pdata_tcs
 196
 197               ; binds <- takeHoisted
 198               ; return (dfuns, binds)
 199               }
 200
 201            -- Return the vectorised variants of type constructors as well as the generated instance type
 202            -- constructors, family instances, and dfun bindings.
 203        ; return (new_tcs ++ inst_tcs, fam_insts, binds)
 204        }
 205
 206
 207 -- Helpers -------------------
 208
 209 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var
 210 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc
 211  = do { vectDataConWorkers orig_tc vect_tc pdata_tc
 212       ; repr <- tyConRepr vect_tc
 213       ; buildPADict vect_tc prepr_tc pdata_tc repr
 214       }
 215
 216 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
 217 vectDataConWorkers orig_tc vect_tc arr_tc
 218  = do bs <- sequence
 219           . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
 220           $ zipWith4 mk_data_con (tyConDataCons vect_tc)
 221                                  rep_tys
 222                                  (inits rep_tys)
 223                                  (tail $ tails rep_tys)
 224       mapM_ (uncurry hoistBinding) bs
 225  where
 226     tyvars   = tyConTyVars vect_tc
 227     var_tys  = mkTyVarTys tyvars
 228     ty_args  = map Type var_tys
 229     res_ty   = mkTyConApp vect_tc var_tys
 230
 231     cons     = tyConDataCons vect_tc
 232     arity    = length cons
 233     [arr_dc] = tyConDataCons arr_tc
 234
 235     rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc
 236
 237
 238     mk_data_con con tys pre post
 239       = liftM2 (,) (vect_data_con con)
 240                    (lift_data_con tys pre post (mkDataConTag con))
 241
 242     sel_replicate len tag
 243       | arity > 1 = do
 244                       rep <- builtin (selReplicate arity)
 245                       return [rep `mkApps` [len, tag]]
 246
 247       | otherwise = return []
 248
 249     vect_data_con con = return $ mkConApp con ty_args
 250     lift_data_con tys pre_tys post_tys tag
 251       = do
 252           len  <- builtin liftingContext
 253           args <- mapM (newLocalVar (fsLit "xs"))
 254                   =<< mapM mkPDataType tys
 255
 256           sel  <- sel_replicate (Var len) tag
 257
 258           pre   <- mapM emptyPD (concat pre_tys)
 259           post  <- mapM emptyPD (concat post_tys)
 260
 261           return . mkLams (len : args)
 262                  . wrapFamInstBody arr_tc var_tys
 263                  . mkConApp arr_dc
 264                  $ ty_args ++ sel ++ pre ++ map Var args ++ post
 265
 266     def_worker data_con arg_tys mk_body
 267       = do
 268           arity <- polyArity tyvars
 269           body <- closedV
 270                 . inBind orig_worker
 271                 . polyAbstract tyvars $ \args ->
 272                   liftM (mkLams (tyvars ++ args) . vectorised)
 273                 $ buildClosures tyvars [] arg_tys res_ty mk_body
 274
 275           raw_worker <- mkVectId orig_worker (exprType body)
 276           let vect_worker = raw_worker `setIdUnfolding`
 277                               mkInlineUnfolding (Just arity) body
 278           defGlobalVar orig_worker vect_worker
 279           return (vect_worker, body)
 280       where
 281         orig_worker = dataConWorkId data_con