From 907fa8af43e420e59ad1b78623f0ffe445c09e87 Mon Sep 17 00:00:00 2001 From: "benl@ouroborus.net" Date: Thu, 9 Sep 2010 06:13:11 +0000 Subject: [PATCH] Finish breaking up vectoriser utils --- compiler/ghc.cabal.in | 38 +-- compiler/vectorise/VectUtils.hs | 371 ------------------------ compiler/vectorise/Vectorise/Exp.hs | 4 +- compiler/vectorise/Vectorise/Type/Env.hs | 4 +- compiler/vectorise/Vectorise/Type/PADict.hs | 3 +- compiler/vectorise/Vectorise/Type/PData.hs | 2 +- compiler/vectorise/Vectorise/Type/PRDict.hs | 2 +- compiler/vectorise/Vectorise/Type/PRepr.hs | 2 +- compiler/vectorise/Vectorise/Type/Repr.hs | 2 +- compiler/vectorise/Vectorise/Type/Type.hs | 2 +- compiler/vectorise/Vectorise/Utils.hs | 131 +++++++++ compiler/vectorise/Vectorise/Utils/Base.hs | 154 ++++++++++ compiler/vectorise/Vectorise/Utils/Closure.hs | 5 +- compiler/vectorise/Vectorise/Utils/Hoisting.hs | 2 +- compiler/vectorise/Vectorise/Utils/PADict.hs | 113 ++++++++ compiler/vectorise/Vectorise/Utils/PRDict.hs | 55 ++++ compiler/vectorise/Vectorise/Utils/Poly.hs | 53 ++++ compiler/vectorise/Vectorise/Var.hs | 2 +- 18 files changed, 540 insertions(+), 405 deletions(-) delete mode 100644 compiler/vectorise/VectUtils.hs create mode 100644 compiler/vectorise/Vectorise/Utils.hs create mode 100644 compiler/vectorise/Vectorise/Utils/Base.hs create mode 100644 compiler/vectorise/Vectorise/Utils/PADict.hs create mode 100644 compiler/vectorise/Vectorise/Utils/PRDict.hs create mode 100644 compiler/vectorise/Vectorise/Utils/Poly.hs diff --git a/compiler/ghc.cabal.in b/compiler/ghc.cabal.in index bd8004a..279e672 100644 --- a/compiler/ghc.cabal.in +++ b/compiler/ghc.cabal.in @@ -454,23 +454,6 @@ Library UniqFM UniqSet Util - VectUtils - Vectorise.Var - Vectorise.Convert - Vectorise.Env - Vectorise.Vect - Vectorise.Exp - Vectorise.Type.Env - Vectorise.Type.Repr - Vectorise.Type.PData - Vectorise.Type.PRepr - Vectorise.Type.PADict - Vectorise.Type.PRDict - Vectorise.Type.Type - Vectorise.Type.TyConDecl - Vectorise.Type.Classify - Vectorise.Utils.Closure - Vectorise.Utils.Hoisting Vectorise.Builtins.Base Vectorise.Builtins.Initialise Vectorise.Builtins.Modules @@ -482,6 +465,27 @@ Library Vectorise.Monad.Global Vectorise.Monad.InstEnv Vectorise.Monad + Vectorise.Utils.Base + Vectorise.Utils.Closure + Vectorise.Utils.Hoisting + Vectorise.Utils.PADict + Vectorise.Utils.PRDict + Vectorise.Utils.Poly + Vectorise.Utils + Vectorise.Type.Env + Vectorise.Type.Repr + Vectorise.Type.PData + Vectorise.Type.PRepr + Vectorise.Type.PADict + Vectorise.Type.PRDict + Vectorise.Type.Type + Vectorise.Type.TyConDecl + Vectorise.Type.Classify + Vectorise.Convert + Vectorise.Vect + Vectorise.Var + Vectorise.Env + Vectorise.Exp Vectorise -- We only need to expose more modules as some of the ncg code is used diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs deleted file mode 100644 index 9c50d4a..0000000 --- a/compiler/vectorise/VectUtils.hs +++ /dev/null @@ -1,371 +0,0 @@ -module VectUtils ( - collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg, - collectAnnValBinders, - dataConTagZ, mkDataConTag, mkDataConTagLit, - - newLocalVVar, - - mkBuiltinCo, voidType, mkWrapType, - mkPADictType, mkPArrayType, mkPDataType, mkPReprType, mkPArray, - mkBuiltinTyConApps, mkClosureTypes, - - pdataReprTyCon, pdataReprDataCon, mkVScrut, - prDictOfType, prDFunOfTyCon, - paDictArgType, paDictOfType, paDFunType, - paMethod, wrapPR, replicatePD, emptyPD, packByTagPD, - combinePD, - liftPD, - zipScalars, scalarClosure, - polyAbstract, polyApply, polyVApply, polyArity -) where -import Vectorise.Monad -import Vectorise.Vect -import Vectorise.Builtins - -import CoreSyn -import CoreUtils -import Coercion -import Type -import TypeRep -import TyCon -import DataCon -import Var -import MkId -import Literal -import Outputable -import FastString -import Control.Monad - - -collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type]) -collectAnnTypeArgs expr = go expr [] - where - go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys) - go e tys = (e, tys) - -collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann) -collectAnnTypeBinders expr = go [] expr - where - go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e - go bs e = (reverse bs, e) - -collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann) -collectAnnValBinders expr = go [] expr - where - go bs (_, AnnLam b e) | isId b = go (b:bs) e - go bs e = (reverse bs, e) - -isAnnTypeArg :: AnnExpr b ann -> Bool -isAnnTypeArg (_, AnnType _) = True -isAnnTypeArg _ = False - -dataConTagZ :: DataCon -> Int -dataConTagZ con = dataConTag con - fIRST_TAG - -mkDataConTagLit :: DataCon -> Literal -mkDataConTagLit = mkMachInt . toInteger . dataConTagZ - -mkDataConTag :: DataCon -> CoreExpr -mkDataConTag = mkIntLitInt . dataConTagZ - -splitPrimTyCon :: Type -> Maybe TyCon -splitPrimTyCon ty - | Just (tycon, []) <- splitTyConApp_maybe ty - , isPrimTyCon tycon - = Just tycon - - | otherwise = Nothing - -mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type -mkBuiltinTyConApp get_tc tys - = do - tc <- builtin get_tc - return $ mkTyConApp tc tys - -mkBuiltinTyConApps :: (Builtins -> TyCon) -> [Type] -> Type -> VM Type -mkBuiltinTyConApps get_tc tys ty - = do - tc <- builtin get_tc - return $ foldr (mk tc) ty tys - where - mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2] - -voidType :: VM Type -voidType = mkBuiltinTyConApp voidTyCon [] - -mkWrapType :: Type -> VM Type -mkWrapType ty = mkBuiltinTyConApp wrapTyCon [ty] - - -mkClosureTypes :: [Type] -> Type -> VM Type -mkClosureTypes = mkBuiltinTyConApps closureTyCon - -mkPReprType :: Type -> VM Type -mkPReprType ty = mkBuiltinTyConApp preprTyCon [ty] - -mkPADictType :: Type -> VM Type -mkPADictType ty = mkBuiltinTyConApp paTyCon [ty] - -mkPArrayType :: Type -> VM Type -mkPArrayType ty - | Just tycon <- splitPrimTyCon ty - = do - r <- lookupPrimPArray tycon - case r of - Just arr -> return $ mkTyConApp arr [] - Nothing -> cantVectorise "Primitive tycon not vectorised" (ppr tycon) -mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty] - -mkPDataType :: Type -> VM Type -mkPDataType ty = mkBuiltinTyConApp pdataTyCon [ty] - -mkPArray :: Type -> CoreExpr -> CoreExpr -> VM CoreExpr -mkPArray ty len dat = do - tc <- builtin parrayTyCon - let [dc] = tyConDataCons tc - return $ mkConApp dc [Type ty, len, dat] - -mkBuiltinCo :: (Builtins -> TyCon) -> VM Coercion -mkBuiltinCo get_tc - = do - tc <- builtin get_tc - return $ mkTyConApp tc [] - -pdataReprTyCon :: Type -> VM (TyCon, [Type]) -pdataReprTyCon ty = builtin pdataTyCon >>= (`lookupFamInst` [ty]) - -pdataReprDataCon :: Type -> VM (DataCon, [Type]) -pdataReprDataCon ty - = do - (tc, arg_tys) <- pdataReprTyCon ty - let [dc] = tyConDataCons tc - return (dc, arg_tys) - -mkVScrut :: VExpr -> VM (CoreExpr, CoreExpr, TyCon, [Type]) -mkVScrut (ve, le) - = do - (tc, arg_tys) <- pdataReprTyCon ty - return (ve, unwrapFamInstScrut tc arg_tys le, tc, arg_tys) - where - ty = exprType ve - -prDFunOfTyCon :: TyCon -> VM CoreExpr -prDFunOfTyCon tycon - = liftM Var - . maybeCantVectoriseM "No PR dictionary for tycon" (ppr tycon) - $ lookupTyConPR tycon - - -paDictArgType :: TyVar -> VM (Maybe Type) -paDictArgType tv = go (TyVarTy tv) (tyVarKind tv) - where - go ty k | Just k' <- kindView k = go ty k' - go ty (FunTy k1 k2) - = do - tv <- newTyVar (fsLit "a") k1 - mty1 <- go (TyVarTy tv) k1 - case mty1 of - Just ty1 -> do - mty2 <- go (AppTy ty (TyVarTy tv)) k2 - return $ fmap (ForAllTy tv . FunTy ty1) mty2 - Nothing -> go ty k2 - - go ty k - | isLiftedTypeKind k - = liftM Just (mkPADictType ty) - - go _ _ = return Nothing - - --- | Get the PA dictionary for some type, or `Nothing` if there isn't one. -paDictOfType :: Type -> VM (Maybe CoreExpr) -paDictOfType ty - = paDictOfTyApp ty_fn ty_args - where - (ty_fn, ty_args) = splitAppTys ty - - paDictOfTyApp :: Type -> [Type] -> VM (Maybe CoreExpr) - paDictOfTyApp ty_fn ty_args - | Just ty_fn' <- coreView ty_fn - = paDictOfTyApp ty_fn' ty_args - - paDictOfTyApp (TyVarTy tv) ty_args - = do dfun <- maybeV (lookupTyVarPA tv) - liftM Just $ paDFunApply dfun ty_args - - paDictOfTyApp (TyConApp tc _) ty_args - = do mdfun <- lookupTyConPA tc - case mdfun of - Nothing - -> pprTrace "VectUtils.paDictOfType" - (vcat [ text "No PA dictionary" - , text "for tycon: " <> ppr tc - , text "in type: " <> ppr ty]) - $ return Nothing - - Just dfun -> liftM Just $ paDFunApply (Var dfun) ty_args - - paDictOfTyApp ty _ - = cantVectorise "Can't construct PA dictionary for type" (ppr ty) - - - -paDFunType :: TyCon -> VM Type -paDFunType tc - = do - margs <- mapM paDictArgType tvs - res <- mkPADictType (mkTyConApp tc arg_tys) - return . mkForAllTys tvs - $ mkFunTys [arg | Just arg <- margs] res - where - tvs = tyConTyVars tc - arg_tys = mkTyVarTys tvs - -paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr -paDFunApply dfun tys - = do Just dicts <- liftM sequence $ mapM paDictOfType tys - return $ mkApps (mkTyApps dfun tys) dicts - - -paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr -paMethod _ name ty - | Just tycon <- splitPrimTyCon ty - = liftM Var - . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon) - $ lookupPrimMethod tycon name - -paMethod method _ ty - = do - fn <- builtin method - Just dict <- paDictOfType ty - return $ mkApps (Var fn) [Type ty, dict] - -prDictOfType :: Type -> VM CoreExpr -prDictOfType ty = prDictOfTyApp ty_fn ty_args - where - (ty_fn, ty_args) = splitAppTys ty - -prDictOfTyApp :: Type -> [Type] -> VM CoreExpr -prDictOfTyApp ty_fn ty_args - | Just ty_fn' <- coreView ty_fn = prDictOfTyApp ty_fn' ty_args -prDictOfTyApp (TyConApp tc _) ty_args - = do - dfun <- liftM Var $ maybeV (lookupTyConPR tc) - prDFunApply dfun ty_args -prDictOfTyApp _ _ = noV - -prDFunApply :: CoreExpr -> [Type] -> VM CoreExpr -prDFunApply dfun tys - = do - dicts <- mapM prDictOfType tys - return $ mkApps (mkTyApps dfun tys) dicts - -wrapPR :: Type -> VM CoreExpr -wrapPR ty - = do - Just pa_dict <- paDictOfType ty - pr_dfun <- prDFunOfTyCon =<< builtin wrapTyCon - return $ mkApps pr_dfun [Type ty, pa_dict] - -replicatePD :: CoreExpr -> CoreExpr -> VM CoreExpr -replicatePD len x = liftM (`mkApps` [len,x]) - (paMethod replicatePDVar "replicatePD" (exprType x)) - -emptyPD :: Type -> VM CoreExpr -emptyPD = paMethod emptyPDVar "emptyPD" - -packByTagPD :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr - -> VM CoreExpr -packByTagPD ty xs len tags t - = liftM (`mkApps` [xs, len, tags, t]) - (paMethod packByTagPDVar "packByTagPD" ty) - -combinePD :: Type -> CoreExpr -> CoreExpr -> [CoreExpr] - -> VM CoreExpr -combinePD ty len sel xs - = liftM (`mkApps` (len : sel : xs)) - (paMethod (combinePDVar n) ("combine" ++ show n ++ "PD") ty) - where - n = length xs - --- | Like `replicatePD` but use the lifting context in the vectoriser state. -liftPD :: CoreExpr -> VM CoreExpr -liftPD x - = do - lc <- builtin liftingContext - replicatePD (Var lc) x - -zipScalars :: [Type] -> Type -> VM CoreExpr -zipScalars arg_tys res_ty - = do - scalar <- builtin scalarClass - (dfuns, _) <- mapAndUnzipM (\ty -> lookupInst scalar [ty]) ty_args - zipf <- builtin (scalarZip $ length arg_tys) - return $ Var zipf `mkTyApps` ty_args `mkApps` map Var dfuns - where - ty_args = arg_tys ++ [res_ty] - -scalarClosure :: [Type] -> Type -> CoreExpr -> CoreExpr -> VM CoreExpr -scalarClosure arg_tys res_ty scalar_fun array_fun - = do - ctr <- builtin (closureCtrFun $ length arg_tys) - Just pas <- liftM sequence $ mapM paDictOfType (init arg_tys) - return $ Var ctr `mkTyApps` (arg_tys ++ [res_ty]) - `mkApps` (pas ++ [scalar_fun, array_fun]) - -newLocalVVar :: FastString -> Type -> VM VVar -newLocalVVar fs vty - = do - lty <- mkPDataType vty - vv <- newLocalVar fs vty - lv <- newLocalVar fs lty - return (vv,lv) - -polyAbstract :: [TyVar] -> ([Var] -> VM a) -> VM a -polyAbstract tvs p - = localV - $ do - mdicts <- mapM mk_dict_var tvs - zipWithM_ (\tv -> maybe (defLocalTyVar tv) - (defLocalTyVarWithPA tv . Var)) tvs mdicts - p (mk_args mdicts) - where - mk_dict_var tv = do - r <- paDictArgType tv - case r of - Just ty -> liftM Just (newLocalVar (fsLit "dPA") ty) - Nothing -> return Nothing - - mk_args mdicts = [dict | Just dict <- mdicts] - -polyArity :: [TyVar] -> VM Int -polyArity tvs = do - tys <- mapM paDictArgType tvs - return $ length [() | Just _ <- tys] - -polyApply :: CoreExpr -> [Type] -> VM CoreExpr -polyApply expr tys - = do Just dicts <- liftM sequence $ mapM paDictOfType tys - return $ expr `mkTyApps` tys `mkApps` dicts - -polyVApply :: VExpr -> [Type] -> VM VExpr -polyVApply expr tys - = do Just dicts <- liftM sequence $ mapM paDictOfType tys - return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr - - -{- -boxExpr :: Type -> VExpr -> VM VExpr -boxExpr ty (vexpr, lexpr) - | Just (tycon, []) <- splitTyConApp_maybe ty - , isUnLiftedTyCon tycon - = do - r <- lookupBoxedTyCon tycon - case r of - Just tycon' -> let [dc] = tyConDataCons tycon' - in - return (mkConApp dc [vexpr], lexpr) - Nothing -> return (vexpr, lexpr) --} - - diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs index 95dbfb4..7831c93 100644 --- a/compiler/vectorise/Vectorise/Exp.hs +++ b/compiler/vectorise/Vectorise/Exp.hs @@ -3,10 +3,8 @@ module Vectorise.Exp (vectPolyExpr) where -import VectUtils +import Vectorise.Utils import Vectorise.Type.Type -import Vectorise.Utils.Closure -import Vectorise.Utils.Hoisting import Vectorise.Var import Vectorise.Vect import Vectorise.Env diff --git a/compiler/vectorise/Vectorise/Type/Env.hs b/compiler/vectorise/Vectorise/Type/Env.hs index 2bc7177..851fb79 100644 --- a/compiler/vectorise/Vectorise/Type/Env.hs +++ b/compiler/vectorise/Vectorise/Type/Env.hs @@ -4,7 +4,6 @@ module Vectorise.Type.Env ( vectTypeEnv, ) where -import VectUtils import Vectorise.Env import Vectorise.Vect import Vectorise.Monad @@ -15,8 +14,7 @@ import Vectorise.Type.PADict import Vectorise.Type.PData import Vectorise.Type.PRepr import Vectorise.Type.Repr -import Vectorise.Utils.Closure -import Vectorise.Utils.Hoisting +import Vectorise.Utils import HscTypes import CoreSyn diff --git a/compiler/vectorise/Vectorise/Type/PADict.hs b/compiler/vectorise/Vectorise/Type/PADict.hs index 7fdc31a..7799384 100644 --- a/compiler/vectorise/Vectorise/Type/PADict.hs +++ b/compiler/vectorise/Vectorise/Type/PADict.hs @@ -2,13 +2,12 @@ module Vectorise.Type.PADict (buildPADict) where -import VectUtils import Vectorise.Monad import Vectorise.Builtins import Vectorise.Type.Repr import Vectorise.Type.PRepr import Vectorise.Type.PRDict -import Vectorise.Utils.Hoisting +import Vectorise.Utils import BasicTypes import CoreSyn diff --git a/compiler/vectorise/Vectorise/Type/PData.hs b/compiler/vectorise/Vectorise/Type/PData.hs index 72a6163..2541bc1 100644 --- a/compiler/vectorise/Vectorise/Type/PData.hs +++ b/compiler/vectorise/Vectorise/Type/PData.hs @@ -2,10 +2,10 @@ module Vectorise.Type.PData (buildPDataTyCon) where -import VectUtils import Vectorise.Monad import Vectorise.Builtins import Vectorise.Type.Repr +import Vectorise.Utils import BasicTypes import BuildTyCl diff --git a/compiler/vectorise/Vectorise/Type/PRDict.hs b/compiler/vectorise/Vectorise/Type/PRDict.hs index 9343d2e..1a55116 100644 --- a/compiler/vectorise/Vectorise/Type/PRDict.hs +++ b/compiler/vectorise/Vectorise/Type/PRDict.hs @@ -2,7 +2,7 @@ module Vectorise.Type.PRDict (buildPRDict) where -import VectUtils +import Vectorise.Utils import Vectorise.Monad import Vectorise.Builtins import Vectorise.Type.Repr diff --git a/compiler/vectorise/Vectorise/Type/PRepr.hs b/compiler/vectorise/Vectorise/Type/PRepr.hs index 086969b..81edaab 100644 --- a/compiler/vectorise/Vectorise/Type/PRepr.hs +++ b/compiler/vectorise/Vectorise/Type/PRepr.hs @@ -6,7 +6,7 @@ module Vectorise.Type.PRepr , buildToArrPRepr , buildFromArrPRepr) where -import VectUtils +import Vectorise.Utils import Vectorise.Monad import Vectorise.Builtins import Vectorise.Type.Repr diff --git a/compiler/vectorise/Vectorise/Type/Repr.hs b/compiler/vectorise/Vectorise/Type/Repr.hs index ea06d1f..40242ae 100644 --- a/compiler/vectorise/Vectorise/Type/Repr.hs +++ b/compiler/vectorise/Vectorise/Type/Repr.hs @@ -12,7 +12,7 @@ module Vectorise.Type.Repr , compReprType , compOrigType) where -import VectUtils +import Vectorise.Utils import Vectorise.Monad import Vectorise.Builtins diff --git a/compiler/vectorise/Vectorise/Type/Type.hs b/compiler/vectorise/Vectorise/Type/Type.hs index 00df5d5..e62f45a 100644 --- a/compiler/vectorise/Vectorise/Type/Type.hs +++ b/compiler/vectorise/Vectorise/Type/Type.hs @@ -4,7 +4,7 @@ module Vectorise.Type.Type , vectAndLiftType , vectType) where -import VectUtils +import Vectorise.Utils import Vectorise.Monad import Vectorise.Builtins import TypeRep diff --git a/compiler/vectorise/Vectorise/Utils.hs b/compiler/vectorise/Vectorise/Utils.hs new file mode 100644 index 0000000..31bb508 --- /dev/null +++ b/compiler/vectorise/Vectorise/Utils.hs @@ -0,0 +1,131 @@ + +module Vectorise.Utils ( + module Vectorise.Utils.Base, + module Vectorise.Utils.Closure, + module Vectorise.Utils.Hoisting, + module Vectorise.Utils.PADict, + module Vectorise.Utils.PRDict, + module Vectorise.Utils.Poly, + + -- * Annotated Exprs + collectAnnTypeArgs, + collectAnnTypeBinders, + collectAnnValBinders, + isAnnTypeArg, + + -- * PD Functions + replicatePD, emptyPD, packByTagPD, + combinePD, liftPD, + + -- * Scalars + zipScalars, scalarClosure, + + -- * Naming + newLocalVar +) +where +import Vectorise.Utils.Base +import Vectorise.Utils.Closure +import Vectorise.Utils.Hoisting +import Vectorise.Utils.PADict +import Vectorise.Utils.PRDict +import Vectorise.Utils.Poly +import Vectorise.Monad +import Vectorise.Builtins +import CoreSyn +import CoreUtils +import Type +import Var +import Control.Monad + + +-- Annotated Exprs ------------------------------------------------------------ +collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type]) +collectAnnTypeArgs expr = go expr [] + where + go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys) + go e tys = (e, tys) + +collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann) +collectAnnTypeBinders expr = go [] expr + where + go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e + go bs e = (reverse bs, e) + +collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann) +collectAnnValBinders expr = go [] expr + where + go bs (_, AnnLam b e) | isId b = go (b:bs) e + go bs e = (reverse bs, e) + +isAnnTypeArg :: AnnExpr b ann -> Bool +isAnnTypeArg (_, AnnType _) = True +isAnnTypeArg _ = False + + +-- PD Functions --------------------------------------------------------------- +replicatePD :: CoreExpr -> CoreExpr -> VM CoreExpr +replicatePD len x = liftM (`mkApps` [len,x]) + (paMethod replicatePDVar "replicatePD" (exprType x)) + +emptyPD :: Type -> VM CoreExpr +emptyPD = paMethod emptyPDVar "emptyPD" + + +packByTagPD :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr +packByTagPD ty xs len tags t + = liftM (`mkApps` [xs, len, tags, t]) + (paMethod packByTagPDVar "packByTagPD" ty) + + +combinePD :: Type -> CoreExpr -> CoreExpr -> [CoreExpr] -> VM CoreExpr +combinePD ty len sel xs + = liftM (`mkApps` (len : sel : xs)) + (paMethod (combinePDVar n) ("combine" ++ show n ++ "PD") ty) + where + n = length xs + + +-- | Like `replicatePD` but use the lifting context in the vectoriser state. +liftPD :: CoreExpr -> VM CoreExpr +liftPD x + = do + lc <- builtin liftingContext + replicatePD (Var lc) x + + +-- Scalars -------------------------------------------------------------------- +zipScalars :: [Type] -> Type -> VM CoreExpr +zipScalars arg_tys res_ty + = do + scalar <- builtin scalarClass + (dfuns, _) <- mapAndUnzipM (\ty -> lookupInst scalar [ty]) ty_args + zipf <- builtin (scalarZip $ length arg_tys) + return $ Var zipf `mkTyApps` ty_args `mkApps` map Var dfuns + where + ty_args = arg_tys ++ [res_ty] + + +scalarClosure :: [Type] -> Type -> CoreExpr -> CoreExpr -> VM CoreExpr +scalarClosure arg_tys res_ty scalar_fun array_fun + = do + ctr <- builtin (closureCtrFun $ length arg_tys) + Just pas <- liftM sequence $ mapM paDictOfType (init arg_tys) + return $ Var ctr `mkTyApps` (arg_tys ++ [res_ty]) + `mkApps` (pas ++ [scalar_fun, array_fun]) + + + +{- +boxExpr :: Type -> VExpr -> VM VExpr +boxExpr ty (vexpr, lexpr) + | Just (tycon, []) <- splitTyConApp_maybe ty + , isUnLiftedTyCon tycon + = do + r <- lookupBoxedTyCon tycon + case r of + Just tycon' -> let [dc] = tyConDataCons tycon' + in + return (mkConApp dc [vexpr], lexpr) + Nothing -> return (vexpr, lexpr) +-} diff --git a/compiler/vectorise/Vectorise/Utils/Base.hs b/compiler/vectorise/Vectorise/Utils/Base.hs new file mode 100644 index 0000000..490eba6 --- /dev/null +++ b/compiler/vectorise/Vectorise/Utils/Base.hs @@ -0,0 +1,154 @@ + +module Vectorise.Utils.Base ( + voidType, + newLocalVVar, + + mkDataConTagLit, + mkDataConTag, dataConTagZ, + mkBuiltinTyConApp, + mkBuiltinTyConApps, + mkWrapType, + mkClosureTypes, + mkPReprType, + mkPArrayType, splitPrimTyCon, + mkPArray, + mkPDataType, + mkBuiltinCo, + mkVScrut, + + pdataReprTyCon, + pdataReprDataCon, +) +where +import Vectorise.Monad +import Vectorise.Vect +import Vectorise.Builtins + +import CoreSyn +import CoreUtils +import Coercion +import Type +import TyCon +import DataCon +import MkId +import Literal +import Outputable +import FastString + + +-- Simple Types --------------------------------------------------------------- +voidType :: VM Type +voidType = mkBuiltinTyConApp voidTyCon [] + + +-- Name Generation ------------------------------------------------------------ +newLocalVVar :: FastString -> Type -> VM VVar +newLocalVVar fs vty + = do + lty <- mkPDataType vty + vv <- newLocalVar fs vty + lv <- newLocalVar fs lty + return (vv,lv) + + +-- Constructors --------------------------------------------------------------- +mkDataConTagLit :: DataCon -> Literal +mkDataConTagLit = mkMachInt . toInteger . dataConTagZ + + +mkDataConTag :: DataCon -> CoreExpr +mkDataConTag = mkIntLitInt . dataConTagZ + + +dataConTagZ :: DataCon -> Int +dataConTagZ con = dataConTag con - fIRST_TAG + + +mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type +mkBuiltinTyConApp get_tc tys + = do + tc <- builtin get_tc + return $ mkTyConApp tc tys + + +mkBuiltinTyConApps :: (Builtins -> TyCon) -> [Type] -> Type -> VM Type +mkBuiltinTyConApps get_tc tys ty + = do + tc <- builtin get_tc + return $ foldr (mk tc) ty tys + where + mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2] + + +mkWrapType :: Type -> VM Type +mkWrapType ty = mkBuiltinTyConApp wrapTyCon [ty] + + +mkClosureTypes :: [Type] -> Type -> VM Type +mkClosureTypes = mkBuiltinTyConApps closureTyCon + + +mkPReprType :: Type -> VM Type +mkPReprType ty = mkBuiltinTyConApp preprTyCon [ty] + + +----- +mkPArrayType :: Type -> VM Type +mkPArrayType ty + | Just tycon <- splitPrimTyCon ty + = do + r <- lookupPrimPArray tycon + case r of + Just arr -> return $ mkTyConApp arr [] + Nothing -> cantVectorise "Primitive tycon not vectorised" (ppr tycon) + +mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty] + +splitPrimTyCon :: Type -> Maybe TyCon +splitPrimTyCon ty + | Just (tycon, []) <- splitTyConApp_maybe ty + , isPrimTyCon tycon + = Just tycon + + | otherwise = Nothing + + +------ +mkPArray :: Type -> CoreExpr -> CoreExpr -> VM CoreExpr +mkPArray ty len dat = do + tc <- builtin parrayTyCon + let [dc] = tyConDataCons tc + return $ mkConApp dc [Type ty, len, dat] + + +mkPDataType :: Type -> VM Type +mkPDataType ty = mkBuiltinTyConApp pdataTyCon [ty] + + +mkBuiltinCo :: (Builtins -> TyCon) -> VM Coercion +mkBuiltinCo get_tc + = do + tc <- builtin get_tc + return $ mkTyConApp tc [] + + +mkVScrut :: VExpr -> VM (CoreExpr, CoreExpr, TyCon, [Type]) +mkVScrut (ve, le) + = do + (tc, arg_tys) <- pdataReprTyCon ty + return (ve, unwrapFamInstScrut tc arg_tys le, tc, arg_tys) + where + ty = exprType ve + +pdataReprTyCon :: Type -> VM (TyCon, [Type]) +pdataReprTyCon ty = builtin pdataTyCon >>= (`lookupFamInst` [ty]) + + +pdataReprDataCon :: Type -> VM (DataCon, [Type]) +pdataReprDataCon ty + = do + (tc, arg_tys) <- pdataReprTyCon ty + let [dc] = tyConDataCons tc + return (dc, arg_tys) + + diff --git a/compiler/vectorise/Vectorise/Utils/Closure.hs b/compiler/vectorise/Vectorise/Utils/Closure.hs index d8be668..b70ecb4 100644 --- a/compiler/vectorise/Vectorise/Utils/Closure.hs +++ b/compiler/vectorise/Vectorise/Utils/Closure.hs @@ -8,11 +8,12 @@ module Vectorise.Utils.Closure ( buildEnv ) where -import VectUtils -import Vectorise.Utils.Hoisting import Vectorise.Builtins import Vectorise.Vect import Vectorise.Monad +import Vectorise.Utils.Base +import Vectorise.Utils.PADict +import Vectorise.Utils.Hoisting import CoreSyn import Type diff --git a/compiler/vectorise/Vectorise/Utils/Hoisting.hs b/compiler/vectorise/Vectorise/Utils/Hoisting.hs index a604927..9cce416 100644 --- a/compiler/vectorise/Vectorise/Utils/Hoisting.hs +++ b/compiler/vectorise/Vectorise/Utils/Hoisting.hs @@ -11,10 +11,10 @@ module Vectorise.Utils.Hoisting ( takeHoisted ) where -import VectUtils import Vectorise.Monad import Vectorise.Env import Vectorise.Vect +import Vectorise.Utils.Poly import CoreSyn import CoreUtils diff --git a/compiler/vectorise/Vectorise/Utils/PADict.hs b/compiler/vectorise/Vectorise/Utils/PADict.hs new file mode 100644 index 0000000..44faa2e --- /dev/null +++ b/compiler/vectorise/Vectorise/Utils/PADict.hs @@ -0,0 +1,113 @@ + +module Vectorise.Utils.PADict ( + mkPADictType, + paDictArgType, + paDictOfType, + paDFunType, + paDFunApply, + paMethod +) +where +import Vectorise.Monad +import Vectorise.Builtins +import Vectorise.Utils.Base + +import CoreSyn +import Coercion +import Type +import TypeRep +import TyCon +import Var +import Outputable +import FastString +import Control.Monad + + +mkPADictType :: Type -> VM Type +mkPADictType ty = mkBuiltinTyConApp paTyCon [ty] + + +paDictArgType :: TyVar -> VM (Maybe Type) +paDictArgType tv = go (TyVarTy tv) (tyVarKind tv) + where + go ty k | Just k' <- kindView k = go ty k' + go ty (FunTy k1 k2) + = do + tv <- newTyVar (fsLit "a") k1 + mty1 <- go (TyVarTy tv) k1 + case mty1 of + Just ty1 -> do + mty2 <- go (AppTy ty (TyVarTy tv)) k2 + return $ fmap (ForAllTy tv . FunTy ty1) mty2 + Nothing -> go ty k2 + + go ty k + | isLiftedTypeKind k + = liftM Just (mkPADictType ty) + + go _ _ = return Nothing + + +-- | Get the PA dictionary for some type, or `Nothing` if there isn't one. +paDictOfType :: Type -> VM (Maybe CoreExpr) +paDictOfType ty + = paDictOfTyApp ty_fn ty_args + where + (ty_fn, ty_args) = splitAppTys ty + + paDictOfTyApp :: Type -> [Type] -> VM (Maybe CoreExpr) + paDictOfTyApp ty_fn ty_args + | Just ty_fn' <- coreView ty_fn + = paDictOfTyApp ty_fn' ty_args + + paDictOfTyApp (TyVarTy tv) ty_args + = do dfun <- maybeV (lookupTyVarPA tv) + liftM Just $ paDFunApply dfun ty_args + + paDictOfTyApp (TyConApp tc _) ty_args + = do mdfun <- lookupTyConPA tc + case mdfun of + Nothing + -> pprTrace "VectUtils.paDictOfType" + (vcat [ text "No PA dictionary" + , text "for tycon: " <> ppr tc + , text "in type: " <> ppr ty]) + $ return Nothing + + Just dfun -> liftM Just $ paDFunApply (Var dfun) ty_args + + paDictOfTyApp ty _ + = cantVectorise "Can't construct PA dictionary for type" (ppr ty) + + + +paDFunType :: TyCon -> VM Type +paDFunType tc + = do + margs <- mapM paDictArgType tvs + res <- mkPADictType (mkTyConApp tc arg_tys) + return . mkForAllTys tvs + $ mkFunTys [arg | Just arg <- margs] res + where + tvs = tyConTyVars tc + arg_tys = mkTyVarTys tvs + +paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr +paDFunApply dfun tys + = do Just dicts <- liftM sequence $ mapM paDictOfType tys + return $ mkApps (mkTyApps dfun tys) dicts + + +paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr +paMethod _ name ty + | Just tycon <- splitPrimTyCon ty + = liftM Var + . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon) + $ lookupPrimMethod tycon name + +paMethod method _ ty + = do + fn <- builtin method + Just dict <- paDictOfType ty + return $ mkApps (Var fn) [Type ty, dict] + diff --git a/compiler/vectorise/Vectorise/Utils/PRDict.hs b/compiler/vectorise/Vectorise/Utils/PRDict.hs new file mode 100644 index 0000000..e573232 --- /dev/null +++ b/compiler/vectorise/Vectorise/Utils/PRDict.hs @@ -0,0 +1,55 @@ + +module Vectorise.Utils.PRDict ( + prDFunOfTyCon, + prDictOfType, + prDictOfTyApp, + prDFunApply, + wrapPR +) +where +import Vectorise.Monad +import Vectorise.Builtins +import Vectorise.Utils.PADict + +import CoreSyn +import Type +import TypeRep +import TyCon +import Outputable +import Control.Monad + + +prDFunOfTyCon :: TyCon -> VM CoreExpr +prDFunOfTyCon tycon + = liftM Var + . maybeCantVectoriseM "No PR dictionary for tycon" (ppr tycon) + $ lookupTyConPR tycon + + + +prDictOfType :: Type -> VM CoreExpr +prDictOfType ty = prDictOfTyApp ty_fn ty_args + where + (ty_fn, ty_args) = splitAppTys ty + +prDictOfTyApp :: Type -> [Type] -> VM CoreExpr +prDictOfTyApp ty_fn ty_args + | Just ty_fn' <- coreView ty_fn = prDictOfTyApp ty_fn' ty_args +prDictOfTyApp (TyConApp tc _) ty_args + = do + dfun <- liftM Var $ maybeV (lookupTyConPR tc) + prDFunApply dfun ty_args +prDictOfTyApp _ _ = noV + +prDFunApply :: CoreExpr -> [Type] -> VM CoreExpr +prDFunApply dfun tys + = do + dicts <- mapM prDictOfType tys + return $ mkApps (mkTyApps dfun tys) dicts + +wrapPR :: Type -> VM CoreExpr +wrapPR ty + = do + Just pa_dict <- paDictOfType ty + pr_dfun <- prDFunOfTyCon =<< builtin wrapTyCon + return $ mkApps pr_dfun [Type ty, pa_dict] diff --git a/compiler/vectorise/Vectorise/Utils/Poly.hs b/compiler/vectorise/Vectorise/Utils/Poly.hs new file mode 100644 index 0000000..04237f8 --- /dev/null +++ b/compiler/vectorise/Vectorise/Utils/Poly.hs @@ -0,0 +1,53 @@ + +module Vectorise.Utils.Poly ( + polyAbstract, + polyApply, + polyVApply, + polyArity +) +where +import Vectorise.Vect +import Vectorise.Monad +import Vectorise.Utils.PADict +import CoreSyn +import Type +import Var +import FastString +import Control.Monad + + +-- Poly Functions ------------------------------------------------------------- +polyAbstract :: [TyVar] -> ([Var] -> VM a) -> VM a +polyAbstract tvs p + = localV + $ do + mdicts <- mapM mk_dict_var tvs + zipWithM_ (\tv -> maybe (defLocalTyVar tv) + (defLocalTyVarWithPA tv . Var)) tvs mdicts + p (mk_args mdicts) + where + mk_dict_var tv = do + r <- paDictArgType tv + case r of + Just ty -> liftM Just (newLocalVar (fsLit "dPA") ty) + Nothing -> return Nothing + + mk_args mdicts = [dict | Just dict <- mdicts] + + +polyArity :: [TyVar] -> VM Int +polyArity tvs = do + tys <- mapM paDictArgType tvs + return $ length [() | Just _ <- tys] + + +polyApply :: CoreExpr -> [Type] -> VM CoreExpr +polyApply expr tys + = do Just dicts <- liftM sequence $ mapM paDictOfType tys + return $ expr `mkTyApps` tys `mkApps` dicts + + +polyVApply :: VExpr -> [Type] -> VM VExpr +polyVApply expr tys + = do Just dicts <- liftM sequence $ mapM paDictOfType tys + return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr diff --git a/compiler/vectorise/Vectorise/Var.hs b/compiler/vectorise/Vectorise/Var.hs index 5a1dcb3..f32cf78 100644 --- a/compiler/vectorise/Vectorise/Var.hs +++ b/compiler/vectorise/Vectorise/Var.hs @@ -10,7 +10,7 @@ module Vectorise.Var ( vectPolyVar, vectLiteral ) where -import VectUtils +import Vectorise.Utils import Vectorise.Monad import Vectorise.Env import Vectorise.Vect -- 1.7.10.4