X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectUtils.hs;h=3abbe44b49b656ae5fbc6b0766837fcd0f277bcf;hb=76cec9c6231e5e73c5dd17e5c7111a79ffde0b03;hp=4982fccb7f30962c648700d3920b6711b7d159e0;hpb=a0a97c2bc86a334185b8c45ac377e2c1ff65f608;p=ghc-hetmet.git diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs index 4982fcc..3abbe44 100644 --- a/compiler/vectorise/VectUtils.hs +++ b/compiler/vectorise/VectUtils.hs @@ -2,27 +2,37 @@ module VectUtils ( collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg, splitClosureTy, mkPADictType, mkPArrayType, - paDictArgType, paDictOfType, paMethod, + paDictArgType, paDictOfType, + paMethod, lengthPA, replicatePA, emptyPA, + polyAbstract, polyApply, polyVApply, lookupPArrayFamInst, - hoistExpr, takeHoisted + hoistExpr, hoistPolyVExpr, takeHoisted, + buildClosure ) where #include "HsVersions.h" +import VectCore import VectMonad +import DsUtils import CoreSyn import CoreUtils import Type import TypeRep import TyCon +import DataCon ( dataConWrapId ) import Var +import Id ( mkWildId ) +import MkId ( unwrapFamInstScrut ) import PrelNames +import TysWiredIn +import BasicTypes ( Boxity(..) ) import Outputable import FastString -import Control.Monad ( liftM ) +import Control.Monad ( liftM, zipWithM_ ) collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type]) collectAnnTypeArgs expr = go expr [] @@ -41,7 +51,7 @@ isAnnTypeArg (_, AnnType t) = True isAnnTypeArg _ = False isClosureTyCon :: TyCon -> Bool -isClosureTyCon tc = tyConUnique tc == closureTyConKey +isClosureTyCon tc = tyConName tc == closureTyConName splitClosureTy :: Type -> (Type, Type) splitClosureTy ty @@ -51,6 +61,17 @@ splitClosureTy ty | otherwise = pprPanic "splitClosureTy" (ppr ty) +isPArrayTyCon :: TyCon -> Bool +isPArrayTyCon tc = tyConName tc == parrayTyConName + +splitPArrayTy :: Type -> Type +splitPArrayTy ty + | Just (tc, [arg_ty]) <- splitTyConApp_maybe ty + , isPArrayTyCon tc + = arg_ty + + | otherwise = pprPanic "splitPArrayTy" (ppr ty) + mkPADictType :: Type -> VM Type mkPADictType ty = do @@ -115,6 +136,54 @@ paMethod method ty 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) + +replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr +replicatePA len x = liftM (`mkApps` [len,x]) + (paMethod replicatePAVar (exprType x)) + +emptyPA :: Type -> VM CoreExpr +emptyPA = paMethod emptyPAVar + +newLocalVVar :: FastString -> Type -> VM VVar +newLocalVVar fs vty + = do + lty <- mkPArrayType vty + vv <- newLocalVar fs vty + lv <- newLocalVar fs lty + return (vv,lv) + +polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> 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_lams 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_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts]) + +polyApply :: CoreExpr -> [Type] -> VM CoreExpr +polyApply expr tys + = do + dicts <- mapM paDictOfType tys + return $ expr `mkTyApps` tys `mkApps` dicts + +polyVApply :: VExpr -> [Type] -> VM VExpr +polyVApply expr tys + = do + dicts <- mapM paDictOfType tys + return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr + lookupPArrayFamInst :: Type -> VM (TyCon, [Type]) lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty]) @@ -126,6 +195,21 @@ hoistExpr fs expr env { global_bindings = (var, expr) : global_bindings env } return var +hoistVExpr :: FastString -> VExpr -> VM VVar +hoistVExpr fs (ve, le) + = do + vv <- hoistExpr ('v' `consFS` fs) ve + lv <- hoistExpr ('l' `consFS` fs) le + return (vv, lv) + +hoistPolyVExpr :: FastString -> [TyVar] -> VM VExpr -> VM VExpr +hoistPolyVExpr fs tvs p + = do + expr <- closedV . polyAbstract tvs $ \abstract -> + liftM (mapVect abstract) p + fn <- hoistVExpr fs expr + polyVApply (vVar fn) (mkTyVarTys tvs) + takeHoisted :: VM [(Var, CoreExpr)] takeHoisted = do @@ -133,3 +217,88 @@ takeHoisted setGEnv $ env { global_bindings = [] } return $ global_bindings env + +mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr +mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv) + = do + dict <- paDictOfType env_ty + mkv <- builtin mkClosureVar + mkl <- builtin mkClosurePVar + 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]) + +-- (clo , aclo (Arr lc xs1 ... xsn) ) +-- where +-- f = \env v -> case env of -> e x1 ... xn v +-- f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v + +buildClosure :: [TyVar] -> Var -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr +buildClosure tvs lv vars arg_ty res_ty mk_body + = do + (env_ty, env, bind) <- buildEnv lv vars + env_bndr <- newLocalVVar FSLIT("env") env_ty + arg_bndr <- newLocalVVar FSLIT("arg") arg_ty + + fn <- hoistPolyVExpr FSLIT("fn") tvs + $ do + body <- mk_body + body' <- bind (vVar env_bndr) + (vVarApps lv body (vars ++ [arg_bndr])) + return (vLamsWithoutLC [env_bndr, arg_bndr] body') + + mkClosure arg_ty res_ty env_ty fn env + +buildEnv :: Var -> [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VM VExpr) +buildEnv lv vvs + = do + let (ty, venv, vbind) = mkVectEnv tys vs + (lenv, lbind) <- mkLiftEnv lv tys ls + return (ty, (venv, lenv), + \(venv,lenv) (vbody,lbody) -> + do + let vbody' = vbind venv vbody + lbody' <- lbind lenv lbody + return (vbody', lbody')) + where + (vs,ls) = unzip vvs + tys = map idType vs + +mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr) +mkVectEnv [] [] = (unitTy, Var unitDataConId, \env 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) + [(DataAlt (tupleCon Boxed (length vs)), vs, body)]) + where + ty = mkCoreTupTy tys + +mkLiftEnv :: Var -> [Type] -> [Var] -> VM (CoreExpr, CoreExpr -> CoreExpr -> VM CoreExpr) +mkLiftEnv lv [ty] [v] + = return (Var v, \env body -> + do + len <- lengthPA (Var v) + return . Let (NonRec v env) + $ Case len lv (exprType body) [(DEFAULT, [], body)]) + +-- NOTE: this transparently deals with empty environments +mkLiftEnv lv tys vs + = do + (env_tc, env_tyargs) <- lookupPArrayFamInst vty + let [env_con] = tyConDataCons env_tc + + env = Var (dataConWrapId env_con) + `mkTyApps` env_tyargs + `mkVarApps` (lv : vs) + + bind env body = let scrut = unwrapFamInstScrut env_tc env_tyargs env + in + return $ Case scrut (mkWildId (exprType scrut)) + (exprType body) + [(DataAlt env_con, lv : bndrs, body)] + return (env, bind) + where + vty = mkCoreTupTy tys + + bndrs | null vs = [mkWildId unitTy] + | otherwise = vs +