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

Vectorise Case on products

parent 5eec4625
......@@ -7,13 +7,17 @@ module VectCore (
vNonRec, vRec,
vVar, vType, vNote, vLet,
vLams, vLamsWithoutLC, vVarApps
vLams, vLamsWithoutLC, vVarApps,
vCaseDEFAULT, vCaseProd
) where
#include "HsVersions.h"
import CoreSyn
import CoreUtils ( exprType )
import DataCon ( DataCon )
import Type ( Type )
import Id ( mkWildId )
import Var
type Vect a = (a,a)
......@@ -69,4 +73,20 @@ vVarApps lc (ve, le) vvs = (ve `mkVarApps` vs, le `mkVarApps` (lc : ls))
where
(vs,ls) = unzip vvs
vCaseDEFAULT :: VExpr -> VVar -> Type -> Type -> VExpr -> VExpr
vCaseDEFAULT (vscrut, lscrut) (vbndr, lbndr) vty lty (vbody, lbody)
= (Case vscrut vbndr vty (mkDEFAULT vbody),
Case lscrut lbndr lty (mkDEFAULT lbody))
where
mkDEFAULT e = [(DEFAULT, [], e)]
vCaseProd :: VExpr -> Type -> Type
-> DataCon -> DataCon -> [Var] -> [VVar] -> VExpr -> VExpr
vCaseProd (vscrut, lscrut) vty lty vdc ldc sh_bndrs bndrs
(vbody,lbody)
= (Case vscrut (mkWildId $ exprType vscrut) vty
[(DataAlt vdc, vbndrs, vbody)],
Case lscrut (mkWildId $ exprType lscrut) lty
[(DataAlt ldc, sh_bndrs ++ lbndrs, lbody)])
where
(vbndrs, lbndrs) = unzip bndrs
......@@ -4,7 +4,7 @@ module VectUtils (
mkDataConTag,
splitClosureTy,
mkPADictType, mkPArrayType,
parrayReprTyCon, parrayReprDataCon,
parrayReprTyCon, parrayReprDataCon, mkVScrut,
paDictArgType, paDictOfType, paDFunType,
paMethod, lengthPA, replicatePA, emptyPA, liftPA,
polyAbstract, polyApply, polyVApply,
......@@ -120,6 +120,12 @@ parrayReprDataCon ty
let [dc] = tyConDataCons tc
return (dc, arg_tys)
mkVScrut :: VExpr -> VM (VExpr, TyCon, [Type])
mkVScrut (ve, le)
= do
(tc, arg_tys) <- parrayReprTyCon (exprType ve)
return ((ve, unwrapFamInstScrut tc arg_tys le), tc, arg_tys)
paDictArgType :: TyVar -> VM (Maybe Type)
paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
where
......
......@@ -154,6 +154,14 @@ vectBndrIn v p
x <- p
return (vv, x)
vectBndrIn' :: Var -> (VVar -> VM a) -> VM (VVar, a)
vectBndrIn' v p
= localV
$ do
vv <- vectBndr v
x <- p vv
return (vv, x)
vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
vectBndrsIn vs p
= localV
......@@ -227,6 +235,12 @@ vectExpr (_, AnnApp fn arg)
arg' <- vectExpr arg
mkClosureApp fn' arg'
vectExpr (_, AnnCase scrut bndr ty alts)
| isAlgType scrut_ty
= vectAlgCase scrut bndr ty alts
where
scrut_ty = exprType (deAnnotate scrut)
vectExpr (_, AnnCase expr bndr ty alts)
= panic "vectExpr: case"
......@@ -279,3 +293,44 @@ vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
vectTyAppExpr e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)
type CoreAltWithFVs = AnnAlt Id VarSet
-- We convert
--
-- case e :: t of v { ... }
--
-- to
--
-- V: let v = e in case v of _ { ... }
-- L: let v = e in case v `cast` ... of _ { ... }
--
-- When lifting, we have to do it this way because v must have the type
-- [:V(T):] but the scrutinee must be cast to the representation type.
--
-- FIXME: this is too lazy
vectAlgCase scrut bndr ty [(DEFAULT, [], body)]
= do
vscrut <- vectExpr scrut
vty <- vectType ty
lty <- mkPArrayType vty
(vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
return $ vCaseDEFAULT vscrut vbndr vty lty vbody
vectAlgCase scrut bndr ty [(DataAlt dc, bndrs, body)]
= do
vty <- vectType ty
lty <- mkPArrayType vty
vexpr <- vectExpr scrut
(vbndr, (vbndrs, vbody)) <- vectBndrIn bndr
. vectBndrsIn bndrs
$ vectExpr body
(vscrut, arr_tc, arg_tys) <- mkVScrut (vVar vbndr)
vect_dc <- maybeV (lookupDataCon dc)
let [arr_dc] = tyConDataCons arr_tc
let shape_tys = take (dataConRepArity arr_dc - length bndrs)
(dataConRepArgTys arr_dc)
shape_bndrs <- mapM (newLocalVar FSLIT("s")) shape_tys
return . vLet (vNonRec vbndr vexpr)
$ vCaseProd vscrut vty lty vect_dc arr_dc shape_bndrs vbndrs vbody
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