compiler/vectorise/VectUtils.hs

   1 module VectUtils (
   2   collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
   3   collectAnnValBinders,
   4   dataConTagZ, mkDataConTag, mkDataConTagLit,
   5
   6   newLocalVVar,
   7
   8   mkBuiltinCo, voidType, mkWrapType,
   9   mkPADictType, mkPArrayType, mkPDataType, mkPReprType, mkPArray,
  10   mkBuiltinTyConApps, mkClosureTypes,
  11
  12   pdataReprTyCon, pdataReprDataCon, mkVScrut,
  13   prDictOfType, prDFunOfTyCon,
  14   paDictArgType, paDictOfType, paDFunType,
  15   paMethod, wrapPR, replicatePD, emptyPD, packByTagPD,
  16   combinePD,
  17   liftPD,
  18   zipScalars, scalarClosure,
  19   polyAbstract, polyApply, polyVApply, polyArity,
  20   Inline(..), addInlineArity, inlineMe,
  21   hoistBinding, hoistExpr, hoistPolyVExpr, takeHoisted,
  22 ) where
  23
  24
  25 import Vectorise.Monad
  26 import Vectorise.Env
  27 import Vectorise.Vect
  28 import Vectorise.Builtins
  29
  30 import CoreSyn
  31 import CoreUtils
  32 import CoreUnfold         ( mkInlineRule )
  33 import Coercion
  34 import Type
  35 import TypeRep
  36 import TyCon
  37 import DataCon
  38 import Var
  39 import MkId               ( unwrapFamInstScrut )
  40 import Id                 ( setIdUnfolding )
  41 import BasicTypes
  42 import Literal            ( Literal, mkMachInt )
  43
  44
  45 import Outputable
  46 import FastString
  47
  48 import Control.Monad
  49
  50 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
  51 collectAnnTypeArgs expr = go expr []
  52   where
  53     go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
  54     go e                             tys = (e, tys)
  55
  56 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  57 collectAnnTypeBinders expr = go [] expr
  58   where
  59     go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
  60     go bs e                           = (reverse bs, e)
  61
  62 collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  63 collectAnnValBinders expr = go [] expr
  64   where
  65     go bs (_, AnnLam b e) | isId b = go (b:bs) e
  66     go bs e                        = (reverse bs, e)
  67
  68 isAnnTypeArg :: AnnExpr b ann -> Bool
  69 isAnnTypeArg (_, AnnType _) = True
  70 isAnnTypeArg _              = False
  71
  72 dataConTagZ :: DataCon -> Int
  73 dataConTagZ con = dataConTag con - fIRST_TAG
  74
  75 mkDataConTagLit :: DataCon -> Literal
  76 mkDataConTagLit = mkMachInt . toInteger . dataConTagZ
  77
  78 mkDataConTag :: DataCon -> CoreExpr
  79 mkDataConTag = mkIntLitInt . dataConTagZ
  80
  81 splitPrimTyCon :: Type -> Maybe TyCon
  82 splitPrimTyCon ty
  83   | Just (tycon, []) <- splitTyConApp_maybe ty
  84   , isPrimTyCon tycon
  85   = Just tycon
  86
  87   | otherwise = Nothing
  88
  89 mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type
  90 mkBuiltinTyConApp get_tc tys
  91   = do
  92       tc <- builtin get_tc
  93       return $ mkTyConApp tc tys
  94
  95 mkBuiltinTyConApps :: (Builtins -> TyCon) -> [Type] -> Type -> VM Type
  96 mkBuiltinTyConApps get_tc tys ty
  97   = do
  98       tc <- builtin get_tc
  99       return $ foldr (mk tc) ty tys
 100   where
 101     mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
 102
 103 voidType :: VM Type
 104 voidType = mkBuiltinTyConApp voidTyCon []
 105
 106 mkWrapType :: Type -> VM Type
 107 mkWrapType ty = mkBuiltinTyConApp wrapTyCon [ty]
 108
 109
 110 mkClosureTypes :: [Type] -> Type -> VM Type
 111 mkClosureTypes = mkBuiltinTyConApps closureTyCon
 112
 113 mkPReprType :: Type -> VM Type
 114 mkPReprType ty = mkBuiltinTyConApp preprTyCon [ty]
 115
 116 mkPADictType :: Type -> VM Type
 117 mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
 118
 119 mkPArrayType :: Type -> VM Type
 120 mkPArrayType ty
 121   | Just tycon <- splitPrimTyCon ty
 122   = do
 123       r <- lookupPrimPArray tycon
 124       case r of
 125         Just arr -> return $ mkTyConApp arr []
 126         Nothing  -> cantVectorise "Primitive tycon not vectorised" (ppr tycon)
 127 mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty]
 128
 129 mkPDataType :: Type -> VM Type
 130 mkPDataType ty = mkBuiltinTyConApp pdataTyCon [ty]
 131
 132 mkPArray :: Type -> CoreExpr -> CoreExpr -> VM CoreExpr
 133 mkPArray ty len dat = do
 134                         tc <- builtin parrayTyCon
 135                         let [dc] = tyConDataCons tc
 136                         return $ mkConApp dc [Type ty, len, dat]
 137
 138 mkBuiltinCo :: (Builtins -> TyCon) -> VM Coercion
 139 mkBuiltinCo get_tc
 140   = do
 141       tc <- builtin get_tc
 142       return $ mkTyConApp tc []
 143
 144 pdataReprTyCon :: Type -> VM (TyCon, [Type])
 145 pdataReprTyCon ty = builtin pdataTyCon >>= (`lookupFamInst` [ty])
 146
 147 pdataReprDataCon :: Type -> VM (DataCon, [Type])
 148 pdataReprDataCon ty
 149   = do
 150       (tc, arg_tys) <- pdataReprTyCon ty
 151       let [dc] = tyConDataCons tc
 152       return (dc, arg_tys)
 153
 154 mkVScrut :: VExpr -> VM (CoreExpr, CoreExpr, TyCon, [Type])
 155 mkVScrut (ve, le)
 156   = do
 157       (tc, arg_tys) <- pdataReprTyCon ty
 158       return (ve, unwrapFamInstScrut tc arg_tys le, tc, arg_tys)
 159   where
 160     ty = exprType ve
 161
 162 prDFunOfTyCon :: TyCon -> VM CoreExpr
 163 prDFunOfTyCon tycon
 164   = liftM Var
 165   . maybeCantVectoriseM "No PR dictionary for tycon" (ppr tycon)
 166   $ lookupTyConPR tycon
 167
 168
 169 paDictArgType :: TyVar -> VM (Maybe Type)
 170 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
 171   where
 172     go ty k | Just k' <- kindView k = go ty k'
 173     go ty (FunTy k1 k2)
 174       = do
 175           tv   <- newTyVar (fsLit "a") k1
 176           mty1 <- go (TyVarTy tv) k1
 177           case mty1 of
 178             Just ty1 -> do
 179                           mty2 <- go (AppTy ty (TyVarTy tv)) k2
 180                           return $ fmap (ForAllTy tv . FunTy ty1) mty2
 181             Nothing  -> go ty k2
 182
 183     go ty k
 184       | isLiftedTypeKind k
 185       = liftM Just (mkPADictType ty)
 186
 187     go _ _ = return Nothing
 188
 189
 190 -- | Get the PA dictionary for some type, or `Nothing` if there isn't one.
 191 paDictOfType :: Type -> VM (Maybe CoreExpr)
 192 paDictOfType ty
 193   = paDictOfTyApp ty_fn ty_args
 194   where
 195     (ty_fn, ty_args) = splitAppTys ty
 196
 197     paDictOfTyApp :: Type -> [Type] -> VM (Maybe CoreExpr)
 198     paDictOfTyApp ty_fn ty_args
 199         | Just ty_fn' <- coreView ty_fn
 200         = paDictOfTyApp ty_fn' ty_args
 201
 202     paDictOfTyApp (TyVarTy tv) ty_args
 203      = do dfun <- maybeV (lookupTyVarPA tv)
 204           liftM Just $ paDFunApply dfun ty_args
 205
 206     paDictOfTyApp (TyConApp tc _) ty_args
 207      = do mdfun <- lookupTyConPA tc
 208           case mdfun of
 209             Nothing
 210              -> pprTrace "VectUtils.paDictOfType"
 211                          (vcat [ text "No PA dictionary"
 212                                , text "for tycon: " <> ppr tc
 213                                , text "in type:   " <> ppr ty])
 214              $ return Nothing
 215
 216             Just dfun   -> liftM Just $ paDFunApply (Var dfun) ty_args
 217
 218     paDictOfTyApp ty _
 219      = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
 220
 221
 222
 223 paDFunType :: TyCon -> VM Type
 224 paDFunType tc
 225   = do
 226       margs <- mapM paDictArgType tvs
 227       res   <- mkPADictType (mkTyConApp tc arg_tys)
 228       return . mkForAllTys tvs
 229              $ mkFunTys [arg | Just arg <- margs] res
 230   where
 231     tvs = tyConTyVars tc
 232     arg_tys = mkTyVarTys tvs
 233
 234 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
 235 paDFunApply dfun tys
 236  = do Just dicts <- liftM sequence $ mapM paDictOfType tys
 237       return $ mkApps (mkTyApps dfun tys) dicts
 238
 239
 240 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
 241 paMethod _ name ty
 242   | Just tycon <- splitPrimTyCon ty
 243   = liftM Var
 244   . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
 245   $ lookupPrimMethod tycon name
 246
 247 paMethod method _ ty
 248   = do
 249       fn        <- builtin method
 250       Just dict <- paDictOfType ty
 251       return $ mkApps (Var fn) [Type ty, dict]
 252
 253 prDictOfType :: Type -> VM CoreExpr
 254 prDictOfType ty = prDictOfTyApp ty_fn ty_args
 255   where
 256     (ty_fn, ty_args) = splitAppTys ty
 257
 258 prDictOfTyApp :: Type -> [Type] -> VM CoreExpr
 259 prDictOfTyApp ty_fn ty_args
 260   | Just ty_fn' <- coreView ty_fn = prDictOfTyApp ty_fn' ty_args
 261 prDictOfTyApp (TyConApp tc _) ty_args
 262   = do
 263       dfun <- liftM Var $ maybeV (lookupTyConPR tc)
 264       prDFunApply dfun ty_args
 265 prDictOfTyApp _ _ = noV
 266
 267 prDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
 268 prDFunApply dfun tys
 269   = do
 270       dicts <- mapM prDictOfType tys
 271       return $ mkApps (mkTyApps dfun tys) dicts
 272
 273 wrapPR :: Type -> VM CoreExpr
 274 wrapPR ty
 275   = do
 276       Just  pa_dict <- paDictOfType ty
 277       pr_dfun       <- prDFunOfTyCon =<< builtin wrapTyCon
 278       return $ mkApps pr_dfun [Type ty, pa_dict]
 279
 280 replicatePD :: CoreExpr -> CoreExpr -> VM CoreExpr
 281 replicatePD len x = liftM (`mkApps` [len,x])
 282                           (paMethod replicatePDVar "replicatePD" (exprType x))
 283
 284 emptyPD :: Type -> VM CoreExpr
 285 emptyPD = paMethod emptyPDVar "emptyPD"
 286
 287 packByTagPD :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr
 288                  -> VM CoreExpr
 289 packByTagPD ty xs len tags t
 290   = liftM (`mkApps` [xs, len, tags, t])
 291           (paMethod packByTagPDVar "packByTagPD" ty)
 292
 293 combinePD :: Type -> CoreExpr -> CoreExpr -> [CoreExpr]
 294           -> VM CoreExpr
 295 combinePD ty len sel xs
 296   = liftM (`mkApps` (len : sel : xs))
 297           (paMethod (combinePDVar n) ("combine" ++ show n ++ "PD") ty)
 298   where
 299     n = length xs
 300
 301 -- | Like `replicatePD` but use the lifting context in the vectoriser state.
 302 liftPD :: CoreExpr -> VM CoreExpr
 303 liftPD x
 304   = do
 305       lc <- builtin liftingContext
 306       replicatePD (Var lc) x
 307
 308 zipScalars :: [Type] -> Type -> VM CoreExpr
 309 zipScalars arg_tys res_ty
 310   = do
 311       scalar <- builtin scalarClass
 312       (dfuns, _) <- mapAndUnzipM (\ty -> lookupInst scalar [ty]) ty_args
 313       zipf <- builtin (scalarZip $ length arg_tys)
 314       return $ Var zipf `mkTyApps` ty_args `mkApps` map Var dfuns
 315     where
 316       ty_args = arg_tys ++ [res_ty]
 317
 318 scalarClosure :: [Type] -> Type -> CoreExpr -> CoreExpr -> VM CoreExpr
 319 scalarClosure arg_tys res_ty scalar_fun array_fun
 320   = do
 321       ctr      <- builtin (closureCtrFun $ length arg_tys)
 322       Just pas <- liftM sequence $ mapM paDictOfType (init arg_tys)
 323       return $ Var ctr `mkTyApps` (arg_tys ++ [res_ty])
 324                        `mkApps`   (pas ++ [scalar_fun, array_fun])
 325
 326 newLocalVVar :: FastString -> Type -> VM VVar
 327 newLocalVVar fs vty
 328   = do
 329       lty <- mkPDataType vty
 330       vv  <- newLocalVar fs vty
 331       lv  <- newLocalVar fs lty
 332       return (vv,lv)
 333
 334 polyAbstract :: [TyVar] -> ([Var] -> VM a) -> VM a
 335 polyAbstract tvs p
 336   = localV
 337   $ do
 338       mdicts <- mapM mk_dict_var tvs
 339       zipWithM_ (\tv -> maybe (defLocalTyVar tv)
 340                               (defLocalTyVarWithPA tv . Var)) tvs mdicts
 341       p (mk_args mdicts)
 342   where
 343     mk_dict_var tv = do
 344                        r <- paDictArgType tv
 345                        case r of
 346                          Just ty -> liftM Just (newLocalVar (fsLit "dPA") ty)
 347                          Nothing -> return Nothing
 348
 349     mk_args mdicts = [dict | Just dict <- mdicts]
 350
 351 polyArity :: [TyVar] -> VM Int
 352 polyArity tvs = do
 353                   tys <- mapM paDictArgType tvs
 354                   return $ length [() | Just _ <- tys]
 355
 356 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
 357 polyApply expr tys
 358  = do Just dicts <- liftM sequence $ mapM paDictOfType tys
 359       return $ expr `mkTyApps` tys `mkApps` dicts
 360
 361 polyVApply :: VExpr -> [Type] -> VM VExpr
 362 polyVApply expr tys
 363  = do Just dicts <- liftM sequence $ mapM paDictOfType tys
 364       return     $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr
 365
 366 -- Inline ---------------------------------------------------------------------
 367 -- | Records whether we should inline a particular binding.
 368 data Inline
 369         = Inline Arity
 370         | DontInline
 371
 372 -- | Add to the arity contained within an `Inline`, if any.
 373 addInlineArity :: Inline -> Int -> Inline
 374 addInlineArity (Inline m) n = Inline (m+n)
 375 addInlineArity DontInline _ = DontInline
 376
 377 -- | Says to always inline a binding.
 378 inlineMe :: Inline
 379 inlineMe = Inline 0
 380
 381
 382 -- Hoising --------------------------------------------------------------------
 383 hoistBinding :: Var -> CoreExpr -> VM ()
 384 hoistBinding v e = updGEnv $ \env ->
 385   env { global_bindings = (v,e) : global_bindings env }
 386
 387 hoistExpr :: FastString -> CoreExpr -> Inline -> VM Var
 388 hoistExpr fs expr inl
 389   = do
 390       var <- mk_inline `liftM` newLocalVar fs (exprType expr)
 391       hoistBinding var expr
 392       return var
 393   where
 394     mk_inline var = case inl of
 395                       Inline arity -> var `setIdUnfolding`
 396                                       mkInlineRule expr (Just arity)
 397                       DontInline   -> var
 398
 399 hoistVExpr :: VExpr -> Inline -> VM VVar
 400 hoistVExpr (ve, le) inl
 401   = do
 402       fs <- getBindName
 403       vv <- hoistExpr ('v' `consFS` fs) ve inl
 404       lv <- hoistExpr ('l' `consFS` fs) le (addInlineArity inl 1)
 405       return (vv, lv)
 406
 407 hoistPolyVExpr :: [TyVar] -> Inline -> VM VExpr -> VM VExpr
 408 hoistPolyVExpr tvs inline p
 409   = do
 410       inline' <- liftM (addInlineArity inline) (polyArity tvs)
 411       expr <- closedV . polyAbstract tvs $ \args ->
 412               liftM (mapVect (mkLams $ tvs ++ args)) p
 413       fn   <- hoistVExpr expr inline'
 414       polyVApply (vVar fn) (mkTyVarTys tvs)
 415
 416 takeHoisted :: VM [(Var, CoreExpr)]
 417 takeHoisted
 418   = do
 419       env <- readGEnv id
 420       setGEnv $ env { global_bindings = [] }
 421       return $ global_bindings env
 422
 423 {-
 424 boxExpr :: Type -> VExpr -> VM VExpr
 425 boxExpr ty (vexpr, lexpr)
 426   | Just (tycon, []) <- splitTyConApp_maybe ty
 427   , isUnLiftedTyCon tycon
 428   = do
 429       r <- lookupBoxedTyCon tycon
 430       case r of
 431         Just tycon' -> let [dc] = tyConDataCons tycon'
 432                        in
 433                        return (mkConApp dc [vexpr], lexpr)
 434         Nothing     -> return (vexpr, lexpr)
 435 -}
 436
 437