Skip to content

GitLab

Commit 8bae3512 authored by rl@cse.unsw.edu.au's avatar rl@cse.unsw.edu.au
Browse files

Thread lifting context implicitly in the vectorisation monad

parent 6ed5e6a3
Hide whitespace changes
Inline Side-by-side
compiler/vectorise/VectMonad.hs
View file @ 8bae3512
......@@ -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 {
......
compiler/vectorise/VectUtils.hs
View file @ 8bae3512
......@@ -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 <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
......@@ -277,27 +284,29 @@ buildClosures tvs lc vars (arg_ty : arg_tys) res_ty mk_body
-- f = \env v -> case env of <x1,...,xn> -> 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
......
compiler/vectorise/Vectorise.hs
View file @ 8bae3512
......@@ -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)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment