From 8bae351221fbd5eabe562641499c14d379816875 Mon Sep 17 00:00:00 2001 From: Roman Leshchinskiy Date: Thu, 2 Aug 2007 02:25:03 +0000 Subject: [PATCH] Thread lifting context implicitly in the vectorisation monad --- compiler/vectorise/VectMonad.hs | 6 +++ compiler/vectorise/VectUtils.hs | 47 +++++++++++++--------- compiler/vectorise/Vectorise.hs | 83 +++++++++++++++++++-------------------- 3 files changed, 75 insertions(+), 61 deletions(-) diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs index c244f0a..a991b8c 100644 --- a/compiler/vectorise/VectMonad.hs +++ b/compiler/vectorise/VectMonad.hs @@ -39,6 +39,7 @@ import Id import OccName import Name import NameEnv +import TysPrim ( intPrimTy ) import DsMonad import PrelNames @@ -69,6 +70,7 @@ data Builtins = Builtins { , lengthPAVar :: Var , replicatePAVar :: Var , emptyPAVar :: Var + , liftingContext :: Var } paDictTyCon :: Builtins -> TyCon @@ -92,6 +94,9 @@ initBuiltins replicatePAVar <- dsLookupGlobalId replicatePAName emptyPAVar <- dsLookupGlobalId emptyPAName + liftingContext <- liftM (\u -> mkSysLocal FSLIT("lc") u intPrimTy) + newUnique + return $ Builtins { parrayTyCon = parrayTyCon , paClass = paClass @@ -103,6 +108,7 @@ initBuiltins , lengthPAVar = lengthPAVar , replicatePAVar = replicatePAVar , emptyPAVar = emptyPAVar + , liftingContext = liftingContext } data GlobalEnv = GlobalEnv { diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs index 718db85..0727c94 100644 --- a/compiler/vectorise/VectUtils.hs +++ b/compiler/vectorise/VectUtils.hs @@ -4,7 +4,7 @@ module VectUtils ( splitClosureTy, mkPADictType, mkPArrayType, paDictArgType, paDictOfType, - paMethod, lengthPA, replicatePA, emptyPA, + paMethod, lengthPA, replicatePA, emptyPA, liftPA, polyAbstract, polyApply, polyVApply, lookupPArrayFamInst, hoistExpr, hoistPolyVExpr, takeHoisted, @@ -170,6 +170,12 @@ replicatePA len x = liftM (`mkApps` [len,x]) emptyPA :: Type -> VM CoreExpr emptyPA = paMethod emptyPAVar +liftPA :: CoreExpr -> VM CoreExpr +liftPA x + = do + lc <- builtin liftingContext + replicatePA (Var lc) x + newLocalVVar :: FastString -> Type -> VM VVar newLocalVVar fs vty = do @@ -259,17 +265,18 @@ mkClosureApp (vclo, lclo) (varg, larg) where (arg_ty, res_ty) = splitClosureTy (exprType vclo) -buildClosures :: [TyVar] -> Var -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr -buildClosures tvs lc vars [arg_ty] res_ty mk_body - = buildClosure tvs lc vars arg_ty res_ty mk_body -buildClosures tvs lc vars (arg_ty : arg_tys) res_ty mk_body +buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr +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 - buildClosure tvs lc vars arg_ty res_ty' + buildClosure tvs vars arg_ty res_ty' . hoistPolyVExpr tvs $ do - clo <- buildClosures tvs lc (vars ++ [arg]) arg_tys res_ty mk_body + lc <- builtin liftingContext + clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body return $ vLams lc (vars ++ [arg]) clo -- (clo , aclo (Arr lc xs1 ... xsn) ) @@ -277,27 +284,29 @@ buildClosures tvs lc vars (arg_ty : arg_tys) res_ty mk_body -- 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 +buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr +buildClosure tvs vars arg_ty res_ty mk_body = do - (env_ty, env, bind) <- buildEnv lv vars + (env_ty, env, bind) <- buildEnv vars env_bndr <- newLocalVVar FSLIT("env") env_ty arg_bndr <- newLocalVVar FSLIT("arg") arg_ty fn <- hoistPolyVExpr tvs $ do + lc <- builtin liftingContext body <- mk_body body' <- bind (vVar env_bndr) - (vVarApps lv body (vars ++ [arg_bndr])) + (vVarApps lc 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 +buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VM VExpr) +buildEnv vvs = do + lc <- builtin liftingContext let (ty, venv, vbind) = mkVectEnv tys vs - (lenv, lbind) <- mkLiftEnv lv tys ls + (lenv, lbind) <- mkLiftEnv lc tys ls return (ty, (venv, lenv), \(venv,lenv) (vbody,lbody) -> do @@ -318,28 +327,28 @@ mkVectEnv tys vs = (ty, mkCoreTup (map Var vs), ty = mkCoreTupTy tys mkLiftEnv :: Var -> [Type] -> [Var] -> VM (CoreExpr, CoreExpr -> CoreExpr -> VM CoreExpr) -mkLiftEnv lv [ty] [v] +mkLiftEnv lc [ty] [v] = return (Var v, \env body -> do len <- lengthPA (Var v) return . Let (NonRec v env) - $ Case len lv (exprType body) [(DEFAULT, [], body)]) + $ Case len lc (exprType body) [(DEFAULT, [], body)]) -- NOTE: this transparently deals with empty environments -mkLiftEnv lv tys vs +mkLiftEnv lc 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) + `mkVarApps` (lc : 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)] + [(DataAlt env_con, lc : bndrs, body)] return (env, bind) where vty = mkCoreTupTy tys diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index 7a29a7b..e3f8008 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -111,10 +111,9 @@ vectTopBinder var vectTopRhs :: Var -> CoreExpr -> VM CoreExpr vectTopRhs var expr = do - lc <- newLocalVar FSLIT("lc") intPrimTy closedV . liftM vectorised . inBind var - $ vectPolyExpr lc (freeVars expr) + $ vectPolyExpr (freeVars expr) -- ---------------------------------------------------------------------------- -- Bindings @@ -150,19 +149,19 @@ vectBndrsIn vs p -- ---------------------------------------------------------------------------- -- Expressions -vectVar :: Var -> Var -> VM VExpr -vectVar lc v +vectVar :: Var -> VM VExpr +vectVar v = do r <- lookupVar v case r of Local (vv,lv) -> return (Var vv, Var lv) Global vv -> do let vexpr = Var vv - lexpr <- replicatePA (Var lc) vexpr + lexpr <- liftPA vexpr return (vexpr, lexpr) -vectPolyVar :: Var -> Var -> [Type] -> VM VExpr -vectPolyVar lc v tys +vectPolyVar :: Var -> [Type] -> VM VExpr +vectPolyVar v tys = do vtys <- mapM vectType tys r <- lookupVar v @@ -171,79 +170,78 @@ vectPolyVar lc v tys (polyApply (Var lv) vtys) Global poly -> do vexpr <- polyApply (Var poly) vtys - lexpr <- replicatePA (Var lc) vexpr + lexpr <- liftPA vexpr return (vexpr, lexpr) -vectLiteral :: Var -> Literal -> VM VExpr -vectLiteral lc lit +vectLiteral :: Literal -> VM VExpr +vectLiteral lit = do - lexpr <- replicatePA (Var lc) (Lit lit) + lexpr <- liftPA (Lit lit) return (Lit lit, lexpr) -vectPolyExpr :: Var -> CoreExprWithFVs -> VM VExpr -vectPolyExpr lc expr +vectPolyExpr :: CoreExprWithFVs -> VM VExpr +vectPolyExpr expr = polyAbstract tvs $ \abstract -> - -- FIXME: shadowing (tvs in lc) do - mono' <- vectExpr lc mono + mono' <- vectExpr mono return $ mapVect abstract mono' where (tvs, mono) = collectAnnTypeBinders expr -vectExpr :: Var -> CoreExprWithFVs -> VM VExpr -vectExpr lc (_, AnnType ty) +vectExpr :: CoreExprWithFVs -> VM VExpr +vectExpr (_, AnnType ty) = liftM vType (vectType ty) -vectExpr lc (_, AnnVar v) = vectVar lc v +vectExpr (_, AnnVar v) = vectVar v -vectExpr lc (_, AnnLit lit) = vectLiteral lc lit +vectExpr (_, AnnLit lit) = vectLiteral lit -vectExpr lc (_, AnnNote note expr) - = liftM (vNote note) (vectExpr lc expr) +vectExpr (_, AnnNote note expr) + = liftM (vNote note) (vectExpr expr) -vectExpr lc e@(_, AnnApp _ arg) +vectExpr e@(_, AnnApp _ arg) | isAnnTypeArg arg - = vectTyAppExpr lc fn tys + = vectTyAppExpr fn tys where (fn, tys) = collectAnnTypeArgs e -vectExpr lc (_, AnnApp fn arg) +vectExpr (_, AnnApp fn arg) = do - fn' <- vectExpr lc fn - arg' <- vectExpr lc arg + fn' <- vectExpr fn + arg' <- vectExpr arg mkClosureApp fn' arg' -vectExpr lc (_, AnnCase expr bndr ty alts) +vectExpr (_, AnnCase expr bndr ty alts) = panic "vectExpr: case" -vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body) +vectExpr (_, AnnLet (AnnNonRec bndr rhs) body) = do - vrhs <- localV . inBind bndr $ vectPolyExpr lc rhs - (vbndr, vbody) <- vectBndrIn bndr (vectExpr lc body) + vrhs <- localV . inBind bndr $ vectPolyExpr rhs + (vbndr, vbody) <- vectBndrIn bndr (vectExpr body) return $ vLet (vNonRec vbndr vrhs) vbody -vectExpr lc (_, AnnLet (AnnRec bs) body) +vectExpr (_, AnnLet (AnnRec bs) body) = do (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs $ liftM2 (,) (zipWithM vect_rhs bndrs rhss) - (vectPolyExpr lc body) + (vectPolyExpr body) return $ vLet (vRec vbndrs vrhss) vbody where (bndrs, rhss) = unzip bs vect_rhs bndr rhs = localV . inBind bndr - $ vectExpr lc rhs + $ vectExpr rhs -vectExpr lc e@(fvs, AnnLam bndr _) +vectExpr e@(fvs, AnnLam bndr _) | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e) - | otherwise = vectLam lc fvs bs body + | otherwise = vectLam fvs bs body where (bs,body) = collectAnnValBinders e -vectLam :: Var -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr -vectLam lc fvs bs body +vectLam :: VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr +vectLam fvs bs body = do tyvars <- localTyVars (vs, vvs) <- readLEnv $ \env -> @@ -253,14 +251,15 @@ vectLam lc fvs bs body arg_tys <- mapM (vectType . idType) bs res_ty <- vectType (exprType $ deAnnotate body) - buildClosures tyvars lc vvs arg_tys res_ty + buildClosures tyvars vvs arg_tys res_ty . hoistPolyVExpr tyvars $ do + lc <- builtin liftingContext (vbndrs, vbody) <- vectBndrsIn (vs ++ bs) - (vectExpr lc body) + (vectExpr body) return $ vLams lc vbndrs vbody -vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM VExpr -vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys -vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e) +vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr +vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys +vectTyAppExpr e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e) -- 1.7.10.4