X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=489d6cc4bf34b72230718df686d451b51fd829b2;hb=02cff9dfe0e5f6b9a92949ee988989e16d764f8b;hp=c73564cd4d99ed35551aecd4e41ae8e4a14f3493;hpb=86193bcfc847f1a1f844508224489456f08d6b83;p=ghc-hetmet.git diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index c73564c..489d6cc 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 @@ -33,6 +34,7 @@ import OccName import DsMonad hiding (mapAndUnzipM) import DsUtils ( mkCoreTup, mkCoreTupTy ) +import Literal ( Literal ) import PrelNames import TysWiredIn import TysPrim ( intPrimTy ) @@ -112,7 +114,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) @@ -122,28 +124,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 -capply :: (CoreExpr, CoreExpr) -> (CoreExpr, CoreExpr) -> VM (CoreExpr, CoreExpr) +capply :: VExpr -> VExpr -> VM VExpr capply (vfn, lfn) (varg, larg) = do apply <- builtin applyClosureVar @@ -154,57 +156,56 @@ 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 <- replicatePA 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 <- replicatePA vexpr lc - return (vexpr, lexpr) - where - mk_app e = polyApply e =<< mapM vectType tys + 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) + +vectLiteral :: Var -> Literal -> VM VExpr +vectLiteral lc lit + = do + lexpr <- replicatePA (Lit lit) (Var lc) + return (Lit lit, lexpr) -vectPolyExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr) +vectPolyExpr :: Var -> CoreExprWithFVs -> VM VExpr vectPolyExpr lc expr - = polyAbstract 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) + = liftM vType (vectType ty) -vectExpr lc (_, AnnVar v) = vectVar lc v +vectExpr lc (_, AnnVar v) = vectVar lc v -vectExpr lc (_, AnnLit lit) - = do - let vexpr = Lit lit - lexpr <- replicatePA vexpr lc - return (vexpr, lexpr) +vectExpr lc (_, AnnLit lit) = vectLiteral lc lit vectExpr lc (_, AnnNote note expr) - = do - (vexpr, lexpr) <- vectExpr lc expr - return (Note note vexpr, Note note lexpr) + = liftM (vNote note) (vectExpr lc expr) vectExpr lc e@(_, AnnApp _ arg) | isAnnTypeArg arg @@ -223,23 +224,19 @@ 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) - return (Let (NonRec vbndr vrhs) vbody, - Let (NonRec lbndr lrhs) lbody) + vrhs <- vectPolyExpr lc rhs + (vbndr, vbody) <- vectBndrIn bndr (vectExpr lc body) + return $ vLet (vNonRec vbndr vrhs) vbody -vectExpr lc (_, AnnLet (AnnRec prs) body) +vectExpr lc (_, AnnLet (AnnRec bs) body) = do - (vbndrs, lbndrs, (vrhss, vbody, lrhss, lbody)) <- vectBndrsIn bndrs vect - return (Let (Rec (zip vbndrs vrhss)) vbody, - Let (Rec (zip lbndrs lrhss)) lbody) + (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs + $ liftM2 (,) + (mapM (vectExpr lc) rhss) + (vectPolyExpr lc body) + return $ vLet (vRec vbndrs vrhss) vbody where - (bndrs, rhss) = unzip prs - - vect = do - (vrhss, lrhss) <- mapAndUnzipM (vectExpr lc) rhss - (vbody, lbody) <- vectPolyExpr lc body - return (vrhss, vbody, lrhss, lbody) + (bndrs, rhss) = unzip bs vectExpr lc e@(_, AnnLam bndr body) | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e) @@ -253,7 +250,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 @@ -276,7 +273,6 @@ vectExpr lc (fvs, AnnLam bndr body) `mkApps` [pa_dict, mono_vfn, mono_lfn, lenv] return (vclo, lclo) - data CEnvInfo = CEnvInfo { cenv_vars :: [Var] @@ -294,8 +290,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 @@ -357,9 +353,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 @@ -401,7 +398,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)