X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectUtils.hs;h=73c986b4200d6419bd51ab4940099a60f085fd04;hb=0a21de62e274acc8e8e260298da4f6c1ee18ecc2;hp=630c425f44f49bd7e07c3bbf83bfee522815ad59;hpb=9c9ff443e39618f01e5a449b6f6049782b534921;p=ghc-hetmet.git diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs index 630c425..73c986b 100644 --- a/compiler/vectorise/VectUtils.hs +++ b/compiler/vectorise/VectUtils.hs @@ -3,23 +3,36 @@ module VectUtils ( splitClosureTy, mkPADictType, mkPArrayType, paDictArgType, paDictOfType, - lookupPArrayFamInst + paMethod, lengthPA, replicatePA, emptyPA, + polyAbstract, polyApply, polyVApply, + lookupPArrayFamInst, + 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 [] @@ -38,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 @@ -48,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 @@ -105,6 +129,176 @@ paDFunApply dfun tys dicts <- mapM paDictOfType tys return $ mkApps (mkTyApps dfun tys) dicts +paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr +paMethod method ty + = do + fn <- builtin method + 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]) +hoistExpr :: FastString -> CoreExpr -> VM Var +hoistExpr fs expr + = do + var <- newLocalVar fs (exprType expr) + updGEnv $ \env -> + 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 + env <- readGEnv id + 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) + (mkVVarApps lv body (vars ++ [arg_bndr])) + return (mkVLams [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 +