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
  11   pdataReprTyCon, pdataReprDataCon, mkVScrut,
  12   prDictOfType, prDFunOfTyCon,
  13   paDictArgType, paDictOfType, paDFunType,
  14   paMethod, wrapPR, replicatePD, emptyPD, packByTagPD,
  15   combinePD,
  16   liftPD,
  17   zipScalars, scalarClosure,
  18   polyAbstract, polyApply, polyVApply, polyArity,
  19   Inline(..), addInlineArity, inlineMe,
  20   hoistBinding, hoistExpr, hoistPolyVExpr, takeHoisted,
  21   buildClosure, buildClosures,
  22   mkClosureApp
  23 ) where
  24 import Vectorise.Monad
  25 import Vectorise.Env
  26 import Vectorise.Vect
  27 import Vectorise.Builtins
  28
  29 import MkCore ( mkCoreTup, mkWildCase )
  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 TysWiredIn
  42 import BasicTypes         ( Boxity(..), Arity )
  43 import Literal            ( Literal, mkMachInt )
  44
  45
  46 import Outputable
  47 import FastString
  48
  49 import Control.Monad
  50
  51 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
  52 collectAnnTypeArgs expr = go expr []
  53   where
  54     go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
  55     go e                             tys = (e, tys)
  56
  57 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  58 collectAnnTypeBinders expr = go [] expr
  59   where
  60     go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
  61     go bs e                           = (reverse bs, e)
  62
  63 collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  64 collectAnnValBinders expr = go [] expr
  65   where
  66     go bs (_, AnnLam b e) | isId b = go (b:bs) e
  67     go bs e                        = (reverse bs, e)
  68
  69 isAnnTypeArg :: AnnExpr b ann -> Bool
  70 isAnnTypeArg (_, AnnType _) = True
  71 isAnnTypeArg _              = False
  72
  73 dataConTagZ :: DataCon -> Int
  74 dataConTagZ con = dataConTag con - fIRST_TAG
  75
  76 mkDataConTagLit :: DataCon -> Literal
  77 mkDataConTagLit = mkMachInt . toInteger . dataConTagZ
  78
  79 mkDataConTag :: DataCon -> CoreExpr
  80 mkDataConTag = mkIntLitInt . dataConTagZ
  81
  82 splitPrimTyCon :: Type -> Maybe TyCon
  83 splitPrimTyCon ty
  84   | Just (tycon, []) <- splitTyConApp_maybe ty
  85   , isPrimTyCon tycon
  86   = Just tycon
  87
  88   | otherwise = Nothing
  89
  90 mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type
  91 mkBuiltinTyConApp get_tc tys
  92   = do
  93       tc <- builtin get_tc
  94       return $ mkTyConApp tc tys
  95
  96 mkBuiltinTyConApps :: (Builtins -> TyCon) -> [Type] -> Type -> VM Type
  97 mkBuiltinTyConApps get_tc tys ty
  98   = do
  99       tc <- builtin get_tc
 100       return $ foldr (mk tc) ty tys
 101   where
 102     mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
 103
 104 voidType :: VM Type
 105 voidType = mkBuiltinTyConApp voidTyCon []
 106
 107 mkWrapType :: Type -> VM Type
 108 mkWrapType ty = mkBuiltinTyConApp wrapTyCon [ty]
 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 -- Closures -------------------------------------------------------------------
 438 mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
 439 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
 440  = do Just dict <- paDictOfType env_ty
 441       mkv       <- builtin closureVar
 442       mkl       <- builtin liftedClosureVar
 443       return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
 444               Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
 445
 446
 447 mkClosureApp :: Type -> Type -> VExpr -> VExpr -> VM VExpr
 448 mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
 449  = do vapply <- builtin applyVar
 450       lapply <- builtin liftedApplyVar
 451       lc     <- builtin liftingContext
 452       return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
 453               Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [Var lc, lclo, larg])
 454
 455
 456 buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
 457 buildClosures _   _    [] _ mk_body
 458   = mk_body
 459 buildClosures tvs vars [arg_ty] res_ty mk_body
 460   = -- liftM vInlineMe $
 461       buildClosure tvs vars arg_ty res_ty mk_body
 462 buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
 463   = do
 464       res_ty' <- mkClosureTypes arg_tys res_ty
 465       arg <- newLocalVVar (fsLit "x") arg_ty
 466       -- liftM vInlineMe
 467       buildClosure tvs vars arg_ty res_ty'
 468         . hoistPolyVExpr tvs (Inline (length vars + 1))
 469         $ do
 470             lc <- builtin liftingContext
 471             clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
 472             return $ vLams lc (vars ++ [arg]) clo
 473
 474 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
 475 --   where
 476 --     f  = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
 477 --     f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
 478 --
 479 buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
 480 buildClosure tvs vars arg_ty res_ty mk_body
 481   = do
 482       (env_ty, env, bind) <- buildEnv vars
 483       env_bndr <- newLocalVVar (fsLit "env") env_ty
 484       arg_bndr <- newLocalVVar (fsLit "arg") arg_ty
 485
 486       fn <- hoistPolyVExpr tvs (Inline 2)
 487           $ do
 488               lc    <- builtin liftingContext
 489               body  <- mk_body
 490               return -- . vInlineMe
 491                      . vLams lc [env_bndr, arg_bndr]
 492                      $ bind (vVar env_bndr)
 493                             (vVarApps lc body (vars ++ [arg_bndr]))
 494
 495       mkClosure arg_ty res_ty env_ty fn env
 496
 497
 498 -- Environments ---------------------------------------------------------------
 499 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
 500 buildEnv [] = do
 501              ty    <- voidType
 502              void  <- builtin voidVar
 503              pvoid <- builtin pvoidVar
 504              return (ty, vVar (void, pvoid), \_ body -> body)
 505
 506 buildEnv [v] = return (vVarType v, vVar v,
 507                     \env body -> vLet (vNonRec v env) body)
 508
 509 buildEnv vs
 510   = do
 511
 512       (lenv_tc, lenv_tyargs) <- pdataReprTyCon ty
 513
 514       let venv_con   = tupleCon Boxed (length vs)
 515           [lenv_con] = tyConDataCons lenv_tc
 516
 517           venv       = mkCoreTup (map Var vvs)
 518           lenv       = Var (dataConWrapId lenv_con)
 519                        `mkTyApps` lenv_tyargs
 520                        `mkApps`   map Var lvs
 521
 522           vbind env body = mkWildCase env ty (exprType body)
 523                            [(DataAlt venv_con, vvs, body)]
 524
 525           lbind env body =
 526             let scrut = unwrapFamInstScrut lenv_tc lenv_tyargs env
 527             in
 528             mkWildCase scrut (exprType scrut) (exprType body)
 529               [(DataAlt lenv_con, lvs, body)]
 530
 531           bind (venv, lenv) (vbody, lbody) = (vbind venv vbody,
 532                                               lbind lenv lbody)
 533
 534       return (ty, (venv, lenv), bind)
 535   where
 536     (vvs, lvs) = unzip vs
 537     tys        = map vVarType vs
 538     ty         = mkBoxedTupleTy tys
 539