X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectUtils.hs;h=3bf97fa7ffbe2f2cad7e29c38b90926c44c09cd7;hb=3f6a74eafcabc1f8d496937a33ec92e7b416f989;hp=7e331f316cd8225edeb83cab327738c0c518a1c4;hpb=675aada9c8cbe8bac3f48c40b4f95caf0fdd2871;p=ghc-hetmet.git diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs index 7e331f3..3bf97fa 100644 --- a/compiler/vectorise/VectUtils.hs +++ b/compiler/vectorise/VectUtils.hs @@ -1,29 +1,27 @@ module VectUtils ( collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg, collectAnnValBinders, - mkDataConTag, - splitClosureTy, + dataConTagZ, mkDataConTag, mkDataConTagLit, - TyConRepr(..), mkTyConRepr, - mkToPRepr, mkToArrPRepr, mkFromPRepr, mkFromArrPRepr, + newLocalVVar, + + mkBuiltinCo, mkPADictType, mkPArrayType, mkPReprType, - parrayCoerce, parrayReprTyCon, parrayReprDataCon, mkVScrut, - prDictOfType, prCoerce, + parrayReprTyCon, parrayReprDataCon, mkVScrut, + prDFunOfTyCon, paDictArgType, paDictOfType, paDFunType, - paMethod, lengthPA, replicatePA, emptyPA, liftPA, + paMethod, mkPR, lengthPA, replicatePA, emptyPA, packPA, combinePA, liftPA, polyAbstract, polyApply, polyVApply, hoistBinding, hoistExpr, hoistPolyVExpr, takeHoisted, buildClosure, buildClosures, mkClosureApp ) where -#include "HsVersions.h" - import VectCore import VectMonad -import DsUtils +import MkCore import CoreSyn import CoreUtils import Coercion @@ -34,16 +32,15 @@ import DataCon import Var import Id ( mkWildId ) import MkId ( unwrapFamInstScrut ) -import Name ( Name ) -import PrelNames import TysWiredIn import BasicTypes ( Boxity(..) ) +import Literal ( Literal, mkMachInt ) import Outputable import FastString -import Data.List ( zipWith4 ) -import Control.Monad ( liftM, liftM2, zipWithM_ ) +import Control.Monad + collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type]) collectAnnTypeArgs expr = go expr [] @@ -60,45 +57,29 @@ collectAnnTypeBinders expr = go [] expr 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) + go bs (_, AnnLam b e) | isIdVar b = go (b:bs) e + go bs e = (reverse bs, e) isAnnTypeArg :: AnnExpr b ann -> Bool -isAnnTypeArg (_, AnnType t) = True +isAnnTypeArg (_, AnnType _) = True isAnnTypeArg _ = False -mkDataConTag :: DataCon -> CoreExpr -mkDataConTag dc = mkConApp intDataCon [mkIntLitInt $ dataConTag dc] - -splitUnTy :: String -> Name -> Type -> Type -splitUnTy s name ty - | Just (tc, [ty']) <- splitTyConApp_maybe ty - , tyConName tc == name - = ty' - - | otherwise = pprPanic s (ppr ty) - -splitBinTy :: String -> Name -> Type -> (Type, Type) -splitBinTy s name ty - | Just (tc, [ty1, ty2]) <- splitTyConApp_maybe ty - , tyConName tc == name - = (ty1, ty2) - - | otherwise = pprPanic s (ppr ty) +dataConTagZ :: DataCon -> Int +dataConTagZ con = dataConTag con - fIRST_TAG -splitFixedTyConApp :: TyCon -> Type -> [Type] -splitFixedTyConApp tc ty - | Just (tc', tys) <- splitTyConApp_maybe ty - , tc == tc' - = tys +mkDataConTagLit :: DataCon -> Literal +mkDataConTagLit = mkMachInt . toInteger . dataConTagZ - | otherwise = pprPanic "splitFixedTyConApp" (ppr tc <+> ppr ty) +mkDataConTag :: DataCon -> CoreExpr +mkDataConTag = mkIntLitInt . dataConTagZ -splitClosureTy :: Type -> (Type, Type) -splitClosureTy = splitBinTy "splitClosureTy" closureTyConName +splitPrimTyCon :: Type -> Maybe TyCon +splitPrimTyCon ty + | Just (tycon, []) <- splitTyConApp_maybe ty + , isPrimTyCon tycon + = Just tycon -splitPArrayTy :: Type -> Type -splitPArrayTy = splitUnTy "splitPArrayTy" parrayTyConName + | otherwise = Nothing mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type mkBuiltinTyConApp get_tc tys @@ -114,9 +95,10 @@ mkBuiltinTyConApps get_tc tys ty where mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2] +{- mkBuiltinTyConApps1 :: (Builtins -> TyCon) -> Type -> [Type] -> VM Type -mkBuiltinTyConApps1 get_tc dft [] = return dft -mkBuiltinTyConApps1 get_tc dft tys +mkBuiltinTyConApps1 _ dft [] = return dft +mkBuiltinTyConApps1 get_tc _ tys = do tc <- builtin get_tc case tys of @@ -125,163 +107,9 @@ mkBuiltinTyConApps1 get_tc dft tys where mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2] -data TyConRepr = TyConRepr { - repr_tyvars :: [TyVar] - , repr_tys :: [[Type]] - - , repr_prod_tycons :: [Maybe TyCon] - , repr_prod_tys :: [Type] - , repr_sum_tycon :: Maybe TyCon - , repr_type :: Type - } - -mkTyConRepr :: TyCon -> VM TyConRepr -mkTyConRepr vect_tc - = do - prod_tycons <- mapM (mk_tycon prodTyCon) rep_tys - let prod_tys = zipWith mk_tc_app_maybe prod_tycons rep_tys - sum_tycon <- mk_tycon sumTyCon prod_tys - - return $ TyConRepr { - repr_tyvars = tyvars - , repr_tys = rep_tys - - , repr_prod_tycons = prod_tycons - , repr_prod_tys = prod_tys - , repr_sum_tycon = sum_tycon - , repr_type = mk_tc_app_maybe sum_tycon prod_tys - } - where - tyvars = tyConTyVars vect_tc - data_cons = tyConDataCons vect_tc - rep_tys = map dataConRepArgTys data_cons - - mk_tycon get_tc tys - | n > 1 = builtin (Just . get_tc n) - | otherwise = return Nothing - where n = length tys - - mk_tc_app_maybe Nothing [] = unitTy - mk_tc_app_maybe Nothing [ty] = ty - mk_tc_app_maybe (Just tc) tys = mkTyConApp tc tys - -{- -mkPRepr :: [[Type]] -> VM Type -mkPRepr tys - = do - embed_tc <- builtin embedTyCon - sum_tcs <- builtins sumTyCon - prod_tcs <- builtins prodTyCon - - let mk_sum [] = unitTy - mk_sum [ty] = ty - mk_sum tys = mkTyConApp (sum_tcs $ length tys) tys - - mk_prod [] = unitTy - mk_prod [ty] = ty - mk_prod tys = mkTyConApp (prod_tcs $ length tys) tys - - mk_embed ty = mkTyConApp embed_tc [ty] - - return . mk_sum - . map (mk_prod . map mk_embed) - $ tys --} - -mkToPRepr :: TyConRepr -> [[CoreExpr]] -> [CoreExpr] -mkToPRepr (TyConRepr { - repr_tys = repr_tys - , repr_prod_tycons = prod_tycons - , repr_prod_tys = prod_tys - , repr_sum_tycon = repr_sum_tycon - }) - = mk_sum . zipWith3 mk_prod prod_tycons repr_tys - where - Just sum_tycon = repr_sum_tycon - - mk_sum [] = [Var unitDataConId] - mk_sum [expr] = [expr] - mk_sum exprs = zipWith (mk_alt prod_tys) (tyConDataCons sum_tycon) exprs - - mk_alt tys dc expr = mk_con_app dc tys [expr] - - mk_prod _ _ [] = Var unitDataConId - mk_prod _ _ [expr] = expr - mk_prod (Just tc) tys exprs = mk_con_app dc tys exprs - where - [dc] = tyConDataCons tc - - mk_con_app dc tys exprs = mkConApp dc (map Type tys ++ exprs) - -mkToArrPRepr :: CoreExpr -> CoreExpr -> [[CoreExpr]] -> VM CoreExpr -mkToArrPRepr len sel ess - = do - let mk_sum [(expr, ty)] = return (expr, ty) - mk_sum es - = do - sum_tc <- builtin . sumTyCon $ length es - (sum_rtc, _) <- parrayReprTyCon (mkTyConApp sum_tc tys) - let [sum_rdc] = tyConDataCons sum_rtc - - return (mkConApp sum_rdc (map Type tys ++ (len : sel : exprs)), - mkTyConApp sum_tc tys) - where - (exprs, tys) = unzip es - - mk_prod [expr] = return (expr, splitPArrayTy (exprType expr)) - mk_prod exprs - = do - prod_tc <- builtin . prodTyCon $ length exprs - (prod_rtc, _) <- parrayReprTyCon (mkTyConApp prod_tc tys) - let [prod_rdc] = tyConDataCons prod_rtc - - return (mkConApp prod_rdc (map Type tys ++ (len : exprs)), - mkTyConApp prod_tc tys) - where - tys = map (splitPArrayTy . exprType) exprs - - liftM fst (mk_sum =<< mapM mk_prod ess) - -mkFromPRepr :: CoreExpr -> Type -> [([Var], CoreExpr)] -> VM CoreExpr -mkFromPRepr scrut res_ty alts - = do - sum_tcs <- builtins sumTyCon - prod_tcs <- builtins prodTyCon - - let un_sum expr ty [(vars, res)] = un_prod expr ty vars res - un_sum expr ty bs - = do - ps <- mapM (newLocalVar FSLIT("p")) tys - bodies <- sequence - $ zipWith4 un_prod (map Var ps) tys vars rs - return . Case expr (mkWildId ty) res_ty - $ zipWith3 mk_alt sum_dcs ps bodies - where - (vars, rs) = unzip bs - tys = splitFixedTyConApp sum_tc ty - sum_tc = sum_tcs $ length bs - sum_dcs = tyConDataCons sum_tc - - mk_alt dc p body = (DataAlt dc, [p], body) - - un_prod expr ty [] r = return r - un_prod expr ty [var] r = return $ Let (NonRec var expr) r - un_prod expr ty vars r - = return $ Case expr (mkWildId ty) res_ty - [(DataAlt prod_dc, vars, r)] - where - prod_tc = prod_tcs $ length vars - [prod_dc] = tyConDataCons prod_tc - - un_sum scrut (exprType scrut) alts - -mkFromArrPRepr :: CoreExpr -> Type -> Var -> Var -> [[Var]] -> CoreExpr - -> VM CoreExpr -mkFromArrPRepr scrut res_ty len sel vars res - = return (Var unitDataConId) - mkClosureType :: Type -> Type -> VM Type mkClosureType arg_ty res_ty = mkBuiltinTyConApp closureTyCon [arg_ty, res_ty] +-} mkClosureTypes :: [Type] -> Type -> VM Type mkClosureTypes = mkBuiltinTyConApps closureTyCon @@ -293,18 +121,20 @@ 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] -parrayCoerce :: TyCon -> [Type] -> CoreExpr -> VM CoreExpr -parrayCoerce repr_tc args expr - | Just arg_co <- tyConFamilyCoercion_maybe repr_tc +mkBuiltinCo :: (Builtins -> TyCon) -> VM Coercion +mkBuiltinCo get_tc = do - parray <- builtin parrayTyCon - - let co = mkAppCoercion (mkTyConApp parray []) - (mkSymCoercion (mkTyConApp arg_co args)) - - return $ mkCoerce co expr + tc <- builtin get_tc + return $ mkTyConApp tc [] parrayReprTyCon :: Type -> VM (TyCon, [Type]) parrayReprTyCon ty = builtin parrayTyCon >>= (`lookupFamInst` [ty]) @@ -322,46 +152,11 @@ mkVScrut (ve, le) (tc, arg_tys) <- parrayReprTyCon (exprType ve) return ((ve, unwrapFamInstScrut tc arg_tys le), tc, arg_tys) -prDictOfType :: Type -> VM CoreExpr -prDictOfType orig_ty - | Just (tycon, ty_args) <- splitTyConApp_maybe orig_ty - = do - dfun <- traceMaybeV "prDictOfType" (ppr tycon) (lookupTyConPR tycon) - prDFunApply (Var dfun) ty_args - -prDFunApply :: CoreExpr -> [Type] -> VM CoreExpr -prDFunApply dfun tys - = do - args <- mapM mkDFunArg arg_tys - return $ mkApps mono_dfun args - where - mono_dfun = mkTyApps dfun tys - (arg_tys, _) = splitFunTys (exprType mono_dfun) - -mkDFunArg :: Type -> VM CoreExpr -mkDFunArg ty - | Just (tycon, [arg]) <- splitTyConApp_maybe ty - - = let name = tyConName tycon - - get_dict | name == paTyConName = paDictOfType - | name == prTyConName = prDictOfType - | otherwise = pprPanic "mkDFunArg" (ppr ty) - - in get_dict arg - -mkDFunArg ty = pprPanic "mkDFunArg" (ppr ty) - -prCoerce :: TyCon -> [Type] -> CoreExpr -> VM CoreExpr -prCoerce repr_tc args expr - | Just arg_co <- tyConFamilyCoercion_maybe repr_tc - = do - pr_tc <- builtin prTyCon - - let co = mkAppCoercion (mkTyConApp pr_tc []) - (mkSymCoercion (mkTyConApp arg_co args)) - - return $ mkCoerce co expr +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) @@ -369,7 +164,7 @@ paDictArgType tv = go (TyVarTy tv) (tyVarKind tv) go ty k | Just k' <- kindView k = go ty k' go ty (FunTy k1 k2) = do - tv <- newTyVar FSLIT("a") k1 + tv <- newTyVar (fsLit "a") k1 mty1 <- go (TyVarTy tv) k1 case mty1 of Just ty1 -> do @@ -381,7 +176,7 @@ paDictArgType tv = go (TyVarTy tv) (tyVarKind tv) | isLiftedTypeKind k = liftM Just (mkPADictType ty) - go ty k = return Nothing + go _ _ = return Nothing paDictOfType :: Type -> VM CoreExpr paDictOfType ty = paDictOfTyApp ty_fn ty_args @@ -397,9 +192,11 @@ paDictOfTyApp (TyVarTy tv) ty_args paDFunApply dfun ty_args paDictOfTyApp (TyConApp tc _) ty_args = do - dfun <- traceMaybeV "paDictOfTyApp" (ppr tc) (lookupTyConPA tc) + dfun <- maybeCantVectoriseM "No PA dictionary for tycon" (ppr tc) + $ lookupTyConPA tc paDFunApply (Var dfun) ty_args -paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty) +paDictOfTyApp ty _ + = cantVectorise "Can't construct PA dictionary for type" (ppr ty) paDFunType :: TyCon -> VM Type paDFunType tc @@ -418,27 +215,55 @@ paDFunApply dfun tys dicts <- mapM paDictOfType tys return $ mkApps (mkTyApps dfun tys) dicts -paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr -paMethod method ty +type PAMethod = (Builtins -> Var, String) + +pa_length, pa_replicate, pa_empty, pa_pack :: (Builtins -> Var, String) +pa_length = (lengthPAVar, "lengthPA") +pa_replicate = (replicatePAVar, "replicatePA") +pa_empty = (emptyPAVar, "emptyPA") +pa_pack = (packPAVar, "packPA") + +paMethod :: PAMethod -> Type -> VM CoreExpr +paMethod (_method, name) ty + | Just tycon <- splitPrimTyCon ty + = liftM Var + . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon) + $ lookupPrimMethod tycon name + +paMethod (method, _name) ty = do fn <- builtin method dict <- paDictOfType ty return $ mkApps (Var fn) [Type ty, dict] mkPR :: Type -> VM CoreExpr -mkPR = paMethod mkPRVar +mkPR ty + = do + fn <- builtin mkPRVar + dict <- paDictOfType ty + return $ mkApps (Var fn) [Type ty, dict] -lengthPA :: CoreExpr -> VM CoreExpr -lengthPA x = liftM (`App` x) (paMethod lengthPAVar ty) - where - ty = splitPArrayTy (exprType x) +lengthPA :: Type -> CoreExpr -> VM CoreExpr +lengthPA ty x = liftM (`App` x) (paMethod pa_length ty) replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr replicatePA len x = liftM (`mkApps` [len,x]) - (paMethod replicatePAVar (exprType x)) + (paMethod pa_replicate (exprType x)) emptyPA :: Type -> VM CoreExpr -emptyPA = paMethod emptyPAVar +emptyPA = paMethod pa_empty + +packPA :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr +packPA ty xs len sel = liftM (`mkApps` [xs, len, sel]) + (paMethod pa_pack ty) + +combinePA :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> [CoreExpr] + -> VM CoreExpr +combinePA ty len sel is xs + = liftM (`mkApps` (len : sel : is : xs)) + (paMethod (combinePAVar n, "combine" ++ show n ++ "PA") ty) + where + n = length xs liftPA :: CoreExpr -> VM CoreExpr liftPA x @@ -465,7 +290,7 @@ polyAbstract tvs p mk_dict_var tv = do r <- paDictArgType tv case r of - Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty) + Just ty -> liftM Just (newLocalVar (fsLit "dPA") ty) Nothing -> return Nothing mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts]) @@ -516,6 +341,20 @@ takeHoisted setGEnv $ env { global_bindings = [] } return $ global_bindings env +{- +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) +-} + mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv) = do @@ -525,25 +364,23 @@ mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv) return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv], Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv]) -mkClosureApp :: VExpr -> VExpr -> VM VExpr -mkClosureApp (vclo, lclo) (varg, larg) +mkClosureApp :: Type -> Type -> VExpr -> VExpr -> VM VExpr +mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg) = do vapply <- builtin applyClosureVar lapply <- builtin applyClosurePVar return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg], Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [lclo, larg]) - where - (arg_ty, res_ty) = splitClosureTy (exprType vclo) buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr -buildClosures tvs vars [] res_ty mk_body +buildClosures _ _ [] _ mk_body = mk_body buildClosures tvs vars [arg_ty] res_ty mk_body = buildClosure tvs vars arg_ty res_ty mk_body buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body = do res_ty' <- mkClosureTypes arg_tys res_ty - arg <- newLocalVVar FSLIT("x") arg_ty + arg <- newLocalVVar (fsLit "x") arg_ty buildClosure tvs vars arg_ty res_ty' . hoistPolyVExpr tvs $ do @@ -560,8 +397,8 @@ buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr buildClosure tvs vars arg_ty res_ty mk_body = do (env_ty, env, bind) <- buildEnv vars - env_bndr <- newLocalVVar FSLIT("env") env_ty - arg_bndr <- newLocalVVar FSLIT("arg") arg_ty + env_bndr <- newLocalVVar (fsLit "env") env_ty + arg_bndr <- newLocalVVar (fsLit "arg") arg_ty fn <- hoistPolyVExpr tvs $ do @@ -587,10 +424,10 @@ buildEnv vvs return (vbody', lbody')) where (vs,ls) = unzip vvs - tys = map idType vs + tys = map varType vs mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr) -mkVectEnv [] [] = (unitTy, Var unitDataConId, \env body -> body) +mkVectEnv [] [] = (unitTy, Var unitDataConId, \_ body -> body) mkVectEnv [ty] [v] = (ty, Var v, \env body -> Let (NonRec v env) body) mkVectEnv tys vs = (ty, mkCoreTup (map Var vs), \env body -> Case env (mkWildId ty) (exprType body) @@ -602,7 +439,7 @@ mkLiftEnv :: Var -> [Type] -> [Var] -> VM (CoreExpr, CoreExpr -> CoreExpr -> VM mkLiftEnv lc [ty] [v] = return (Var v, \env body -> do - len <- lengthPA (Var v) + len <- lengthPA ty (Var v) return . Let (NonRec v env) $ Case len lc (exprType body) [(DEFAULT, [], body)]) @@ -610,11 +447,16 @@ mkLiftEnv lc [ty] [v] mkLiftEnv lc tys vs = do (env_tc, env_tyargs) <- parrayReprTyCon vty + + bndrs <- if null vs then do + v <- newDummyVar unitTy + return [v] + else return vs let [env_con] = tyConDataCons env_tc env = Var (dataConWrapId env_con) `mkTyApps` env_tyargs - `mkVarApps` (lc : vs) + `mkApps` (Var lc : args) bind env body = let scrut = unwrapFamInstScrut env_tc env_tyargs env in @@ -625,6 +467,6 @@ mkLiftEnv lc tys vs where vty = mkCoreTupTy tys - bndrs | null vs = [mkWildId unitTy] - | otherwise = vs + args | null vs = [Var unitDataConId] + | otherwise = map Var vs