X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=ce4f175bf36993daeb03d3952844954cdc7a4bc2;hb=f3114b4ac0b4316bf7acb9469c400b3d1c483c6c;hp=286680fa27c57b3760cc5288ad6250f158c9a46e;hpb=4ee5e14f8b27ec88d31a8a61c987a72eebe52c70;p=ghc-hetmet.git diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index 286680f..ce4f175 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -6,6 +6,7 @@ where import VectMonad import VectUtils import VectType +import VectCore import DynFlags import HscTypes @@ -19,6 +20,8 @@ import Rules ( RuleBase ) import DataCon import TyCon import Type +import FamInstEnv ( extendFamInstEnvList ) +import InstEnv ( extendInstEnvList ) import Var import VarEnv import VarSet @@ -38,7 +41,7 @@ import BasicTypes ( Boxity(..) ) import Outputable import FastString -import Control.Monad ( liftM, liftM2, mapAndUnzipM, zipWithM_ ) +import Control.Monad ( liftM, liftM2, mapAndUnzipM ) vectorise :: HscEnv -> UniqSupply -> RuleBase -> ModGuts -> IO (SimplCount, ModGuts) @@ -56,10 +59,22 @@ vectorise hsc_env _ _ guts vectModule :: ModGuts -> VM ModGuts vectModule guts = do - types' <- vectTypeEnv (mg_types guts) + (types', fam_insts, pa_insts) <- vectTypeEnv (mg_types guts) + + let insts = map painstInstance pa_insts + fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts + inst_env' = extendInstEnvList (mg_inst_env guts) insts + updGEnv (setInstEnvs inst_env' fam_inst_env') + + dicts <- mapM buildPADict pa_insts binds' <- mapM vectTopBind (mg_binds guts) - return $ guts { mg_types = types' - , mg_binds = binds' } + return $ guts { mg_types = types' + , mg_binds = Rec (concat dicts) : binds' + , mg_inst_env = inst_env' + , mg_fam_inst_env = fam_inst_env' + , mg_insts = mg_insts guts ++ insts + , mg_fam_insts = mg_fam_insts guts ++ fam_insts + } vectTopBind :: CoreBind -> VM CoreBind vectTopBind b@(NonRec var expr) @@ -98,7 +113,7 @@ vectTopRhs = liftM fst . closedV . vectPolyExpr (panic "Empty lifting context") -- ---------------------------------------------------------------------------- -- Bindings -vectBndr :: Var -> VM (Var, Var) +vectBndr :: Var -> VM VVar vectBndr v = do vty <- vectType (idType v) @@ -108,37 +123,28 @@ vectBndr v updLEnv (mapTo vv lv) return (vv, lv) where - mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (Var vv, Var lv) } + mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) } -vectBndrIn :: Var -> VM a -> VM (Var, Var, a) +vectBndrIn :: Var -> VM a -> VM (VVar, a) vectBndrIn v p = localV $ do - (vv, lv) <- vectBndr v + vv <- vectBndr v x <- p - return (vv, lv, x) + return (vv, x) -vectBndrsIn :: [Var] -> VM a -> VM ([Var], [Var], a) +vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a) vectBndrsIn vs p = localV $ do - (vvs, lvs) <- mapAndUnzipM vectBndr vs + vvs <- mapM vectBndr vs x <- p - return (vvs, lvs, x) + return (vvs, x) -- ---------------------------------------------------------------------------- -- Expressions -replicateP :: CoreExpr -> CoreExpr -> VM CoreExpr -replicateP expr len - = do - dict <- paDictOfType ty - rep <- builtin replicatePAVar - return $ mkApps (Var rep) [Type ty, dict, expr, len] - where - ty = exprType expr - -capply :: (CoreExpr, CoreExpr) -> (CoreExpr, CoreExpr) -> VM (CoreExpr, CoreExpr) +capply :: VExpr -> VExpr -> VM VExpr capply (vfn, lfn) (varg, larg) = do apply <- builtin applyClosureVar @@ -149,73 +155,52 @@ capply (vfn, lfn) (varg, larg) fn_ty = exprType vfn (arg_ty, res_ty) = splitClosureTy fn_ty -vectVar :: CoreExpr -> Var -> VM (CoreExpr, CoreExpr) +vectVar :: Var -> Var -> VM VExpr vectVar lc v = do r <- lookupVar v case r of - Local es -> return es - Global vexpr -> do - lexpr <- replicateP vexpr lc - return (vexpr, lexpr) + Local (vv,lv) -> return (Var vv, Var lv) + Global vv -> do + let vexpr = Var vv + lexpr <- replicatePA vexpr (Var lc) + return (vexpr, lexpr) -vectPolyVar :: CoreExpr -> Var -> [Type] -> VM (CoreExpr, CoreExpr) +vectPolyVar :: Var -> Var -> [Type] -> VM VExpr vectPolyVar lc v tys = do + vtys <- mapM vectType tys r <- lookupVar v case r of - Local (vexpr, lexpr) -> liftM2 (,) (mk_app vexpr) (mk_app lexpr) - Global poly -> do - vexpr <- mk_app poly - lexpr <- replicateP vexpr lc - return (vexpr, lexpr) - where - mk_app e = applyToTypes e =<< mapM vectType tys - -abstractOverTyVars :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a -abstractOverTyVars tvs p - = 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]) - -applyToTypes :: CoreExpr -> [Type] -> VM CoreExpr -applyToTypes expr tys - = do - dicts <- mapM paDictOfType tys - return $ expr `mkTyApps` tys `mkApps` dicts - -vectPolyExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr) + Local (vv, lv) -> liftM2 (,) (polyApply (Var vv) vtys) + (polyApply (Var lv) vtys) + Global poly -> do + vexpr <- polyApply (Var poly) vtys + lexpr <- replicatePA vexpr (Var lc) + return (vexpr, lexpr) + +vectPolyExpr :: Var -> CoreExprWithFVs -> VM VExpr vectPolyExpr lc expr - = localV - . abstractOverTyVars tvs $ \mk_lams -> + = polyAbstract tvs $ \abstract -> -- FIXME: shadowing (tvs in lc) do - (vmono, lmono) <- vectExpr lc mono - return $ (mk_lams vmono, mk_lams lmono) + mono' <- vectExpr lc mono + return $ mapVect abstract mono' where (tvs, mono) = collectAnnTypeBinders expr -vectExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr) +vectExpr :: Var -> CoreExprWithFVs -> VM VExpr vectExpr lc (_, AnnType ty) = do vty <- vectType ty return (Type vty, Type vty) -vectExpr lc (_, AnnVar v) = vectVar lc v +vectExpr lc (_, AnnVar v) = vectVar lc v vectExpr lc (_, AnnLit lit) = do let vexpr = Lit lit - lexpr <- replicateP vexpr lc + lexpr <- replicatePA vexpr (Var lc) return (vexpr, lexpr) vectExpr lc (_, AnnNote note expr) @@ -241,13 +226,14 @@ vectExpr lc (_, AnnCase expr bndr ty alts) vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body) = do (vrhs, lrhs) <- vectPolyExpr lc rhs - (vbndr, lbndr, (vbody, lbody)) <- vectBndrIn bndr (vectExpr lc body) + ((vbndr, lbndr), (vbody, lbody)) <- vectBndrIn bndr (vectExpr lc body) return (Let (NonRec vbndr vrhs) vbody, Let (NonRec lbndr lrhs) lbody) vectExpr lc (_, AnnLet (AnnRec prs) body) = do - (vbndrs, lbndrs, (vrhss, vbody, lrhss, lbody)) <- vectBndrsIn bndrs vect + (bndrs, (vrhss, vbody, lrhss, lbody)) <- vectBndrsIn bndrs vect + let (vbndrs, lbndrs) = unzip bndrs return (Let (Rec (zip vbndrs vrhss)) vbody, Let (Rec (zip lbndrs lrhss)) lbody) where @@ -270,7 +256,7 @@ vectExpr lc (fvs, AnnLam bndr body) vfn_var <- hoistExpr FSLIT("vfn") poly_vfn lfn_var <- hoistExpr FSLIT("lfn") poly_lfn - let (venv, lenv) = mkClosureEnvs info lc + let (venv, lenv) = mkClosureEnvs info (Var lc) let env_ty = cenv_vty info @@ -280,8 +266,8 @@ vectExpr lc (fvs, AnnLam bndr body) res_ty <- vectType (exprType $ deAnnotate body) -- FIXME: move the functions to the top level - mono_vfn <- applyToTypes (Var vfn_var) (mkTyVarTys tyvars) - mono_lfn <- applyToTypes (Var lfn_var) (mkTyVarTys tyvars) + mono_vfn <- polyApply (Var vfn_var) (mkTyVarTys tyvars) + mono_lfn <- polyApply (Var lfn_var) (mkTyVarTys tyvars) mk_clo <- builtin mkClosureVar mk_cloP <- builtin mkClosurePVar @@ -293,7 +279,6 @@ vectExpr lc (fvs, AnnLam bndr body) `mkApps` [pa_dict, mono_vfn, mono_lfn, lenv] return (vclo, lclo) - data CEnvInfo = CEnvInfo { cenv_vars :: [Var] @@ -311,8 +296,8 @@ mkCEnvInfo fvs arg body locals <- readLEnv local_vars let (vars, vals) = unzip - [(var, val) | var <- varSetElems fvs - , Just val <- [lookupVarEnv locals var]] + [(var, (Var v, Var v')) | var <- varSetElems fvs + , Just (v,v') <- [lookupVarEnv locals var]] vtys <- mapM (vectType . varType) vars (vty, repr_tycon, repr_tyargs, repr_datacon) <- mk_env_ty vtys @@ -364,7 +349,7 @@ mkClosureFns :: CEnvInfo -> [TyVar] -> Var -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr) mkClosureFns info tyvars arg body = closedV - . abstractOverTyVars tyvars + . polyAbstract tyvars $ \mk_tlams -> do (vfn, lfn) <- mkClosureMonoFns info arg body @@ -374,9 +359,10 @@ mkClosureMonoFns :: CEnvInfo -> Var -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr mkClosureMonoFns info arg body = do lc_bndr <- newLocalVar FSLIT("lc") intPrimTy - (varg : vbndrs, larg : lbndrs, (vbody, lbody)) + (bndrs, (vbody, lbody)) <- vectBndrsIn (arg : cenv_vars info) - (vectExpr (Var lc_bndr) body) + (vectExpr lc_bndr body) + let (varg : vbndrs, larg : lbndrs) = unzip bndrs venv_bndr <- newLocalVar FSLIT("env") vty lenv_bndr <- newLocalVar FSLIT("env") lty @@ -398,10 +384,9 @@ mkClosureMonoFns info arg body bind_lenv lenv lbody lc_bndr [lbndr] = do - lengthPA <- builtin lengthPAVar - pa_dict <- paDictOfType vty + len <- lengthPA (Var lbndr) return . Let (NonRec lbndr lenv) - $ Case (mkApps (Var lengthPA) [Type vty, pa_dict, (Var lbndr)]) + $ Case len lc_bndr (exprType lbody) [(DEFAULT, [], lbody)] @@ -419,7 +404,7 @@ mkClosureMonoFns info arg body (exprType lbody) [(DataAlt (cenv_repr_datacon info), lc_bndr : lbndrs', lbody)] -vectTyAppExpr :: CoreExpr -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr) +vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr) vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)