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

Generate conversion to PRepr during vectorisation

parent c2a7c5f5
......@@ -37,7 +37,7 @@ import Digraph ( SCC(..), stronglyConnComp )
import Outputable
import Control.Monad ( liftM, liftM2, zipWithM, zipWithM_ )
import Data.List ( inits, tails, zipWith4 )
import Data.List ( inits, tails, zipWith4, zipWith5 )
-- ----------------------------------------------------------------------------
-- Types
......@@ -101,8 +101,12 @@ vectTypeEnv env
parr_tcs <- zipWithM buildPArrayTyCon orig_tcs vect_tcs
dfuns <- mapM mkPADFun vect_tcs
defTyConPAs (zip vect_tcs dfuns)
binds <- sequence (zipWith4 buildTyConBindings orig_tcs vect_tcs parr_tcs dfuns)
binds <- sequence (zipWith5 buildTyConBindings orig_tcs
vect_tcs
repr_tcs
parr_tcs
dfuns)
let all_new_tcs = new_tcs ++ repr_tcs ++ parr_tcs
let new_env = extendTypeEnvList env
......@@ -195,7 +199,7 @@ buildPReprTyCon :: TyCon -> TyCon -> VM TyCon
buildPReprTyCon orig_tc vect_tc
= do
name <- cloneName mkPReprTyConOcc (tyConName orig_tc)
rhs_ty <- buildPReprRhsTy vect_tc
rhs_ty <- buildPReprType vect_tc
prepr_tc <- builtin preprTyCon
liftDs $ buildSynTyCon name
tyvars
......@@ -204,13 +208,27 @@ buildPReprTyCon orig_tc vect_tc
where
tyvars = tyConTyVars vect_tc
buildPReprRhsTy :: TyCon -> VM Type
buildPReprRhsTy = buildPReprTy . map dataConRepArgTys . tyConDataCons
buildPReprType :: TyCon -> VM Type
buildPReprType = mkPReprType . map dataConRepArgTys . tyConDataCons
buildToPRepr :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
buildToPRepr _ vect_tc prepr_tc _
= do
arg <- newLocalVar FSLIT("x") arg_ty
bndrss <- mapM (mapM (newLocalVar FSLIT("x"))) rep_tys
(alt_bodies, res_ty) <- mkPReprAlts $ map (map Var) bndrss
return . Lam arg
. wrapFamInstBody prepr_tc var_tys
. Case (Var arg) (mkWildId arg_ty) res_ty
$ zipWith3 mk_alt data_cons bndrss alt_bodies
where
var_tys = mkTyVarTys $ tyConTyVars vect_tc
arg_ty = mkTyConApp vect_tc var_tys
data_cons = tyConDataCons vect_tc
rep_tys = map dataConRepArgTys data_cons
buildPReprTy :: [[Type]] -> VM Type
buildPReprTy tys = mkPlusTypes unitTy
=<< mapM (mkCrossTypes unitTy)
=<< mapM (mapM mkEmbedType) tys
mk_alt data_con bndrs body = (DataAlt data_con, bndrs, body)
buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon
buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc ->
......@@ -293,8 +311,9 @@ tyConShape vect_tc
return [e]
}
buildTyConBindings :: TyCon -> TyCon -> TyCon -> Var -> VM [(Var, CoreExpr)]
buildTyConBindings orig_tc vect_tc arr_tc dfun
buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> Var
-> VM [(Var, CoreExpr)]
buildTyConBindings orig_tc vect_tc prepr_tc arr_tc dfun
= do
shape <- tyConShape vect_tc
sequence_ (zipWith4 (vectDataConWorker shape vect_tc arr_tc arr_dc)
......@@ -302,7 +321,7 @@ buildTyConBindings orig_tc vect_tc arr_tc dfun
vect_dcs
(inits repr_tys)
(tails repr_tys))
dict <- buildPADict shape vect_tc arr_tc dfun
dict <- buildPADict shape vect_tc prepr_tc arr_tc dfun
binds <- takeHoisted
return $ (dfun, dict) : binds
where
......@@ -354,8 +373,8 @@ vectDataConWorker shape vect_tc arr_tc arr_dc orig_dc vect_dc pre (dc_tys : post
++ map Var args
++ empty_post
buildPADict :: Shape -> TyCon -> TyCon -> Var -> VM CoreExpr
buildPADict shape vect_tc arr_tc dfun
buildPADict :: Shape -> TyCon -> TyCon -> TyCon -> Var -> VM CoreExpr
buildPADict shape vect_tc prepr_tc arr_tc dfun
= polyAbstract tvs $ \abstract ->
do
meth_binds <- mapM (mk_method shape) paMethods
......@@ -372,15 +391,16 @@ buildPADict shape vect_tc arr_tc dfun
mk_method shape (name, build)
= localV
$ do
body <- build shape vect_tc arr_tc
body <- build shape vect_tc prepr_tc arr_tc
var <- newLocalVar name (exprType body)
return (var, mkInlineMe body)
paMethods = [(FSLIT("lengthPA"), buildLengthPA),
(FSLIT("replicatePA"), buildReplicatePA)]
(FSLIT("replicatePA"), buildReplicatePA),
(FSLIT("toPRepr"), buildToPRepr)]
buildLengthPA :: Shape -> TyCon -> TyCon -> VM CoreExpr
buildLengthPA shape vect_tc arr_tc
buildLengthPA :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
buildLengthPA shape vect_tc _ arr_tc
= do
parr_ty <- mkPArrayType (mkTyConApp vect_tc arg_tys)
arg <- newLocalVar FSLIT("xs") parr_ty
......@@ -428,8 +448,8 @@ buildLengthPA shape vect_tc arr_tc
--
--
buildReplicatePA :: Shape -> TyCon -> TyCon -> VM CoreExpr
buildReplicatePA shape vect_tc arr_tc
buildReplicatePA :: Shape -> TyCon -> TyCon -> TyCon -> VM CoreExpr
buildReplicatePA shape vect_tc _ arr_tc
= do
len_var <- newLocalVar FSLIT("n") intPrimTy
val_var <- newLocalVar FSLIT("x") val_ty
......
......@@ -3,8 +3,7 @@ module VectUtils (
collectAnnValBinders,
mkDataConTag,
splitClosureTy,
mkPlusType, mkPlusTypes, mkCrossType, mkCrossTypes, mkEmbedType,
mkPlusAlts, mkCrosses, mkEmbed,
mkPReprType, mkPReprAlts,
mkPADictType, mkPArrayType,
parrayReprTyCon, parrayReprDataCon, mkVScrut,
paDictArgType, paDictOfType, paDFunType,
......@@ -111,64 +110,66 @@ mkBuiltinTyConApps1 get_tc dft tys
where
mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
mkBuiltinDataConApp :: (Builtins -> DataCon) -> [CoreExpr] -> VM CoreExpr
mkBuiltinDataConApp get_dc args
mkPReprType :: [[Type]] -> VM Type
mkPReprType [] = return unitTy
mkPReprType tys
= do
dc <- builtin get_dc
return $ mkConApp dc args
embed <- builtin embedTyCon
cross <- builtin crossTyCon
plus <- builtin plusTyCon
mkPlusType :: Type -> Type -> VM Type
mkPlusType ty1 ty2 = mkBuiltinTyConApp plusTyCon [ty1, ty2]
let mk_embed ty = mkTyConApp embed [ty]
mk_cross ty1 ty2 = mkTyConApp cross [ty1, ty2]
mk_plus ty1 ty2 = mkTyConApp plus [ty1, ty2]
mkPlusTypes :: Type -> [Type] -> VM Type
mkPlusTypes = mkBuiltinTyConApps1 plusTyCon
mk_tup [] = unitTy
mk_tup tys = foldr1 mk_cross tys
mkPlusAlts :: [CoreExpr] -> VM [CoreExpr]
mkPlusAlts [] = return []
mkPlusAlts exprs
= do
plus_tc <- builtin plusTyCon
left_dc <- builtin leftDataCon
right_dc <- builtin rightDataCon
let go [expr] = ([expr], exprType expr)
go (expr : exprs)
| (alts, right_ty) <- go exprs
= (mkConApp left_dc [Type left_ty, Type right_ty, expr]
: [mkConApp right_dc [Type left_ty, Type right_ty, alt]
| alt <- alts],
mkTyConApp plus_tc [left_ty, right_ty])
where
left_ty = exprType expr
return . fst $ go exprs
mkCrossType :: Type -> Type -> VM Type
mkCrossType ty1 ty2 = mkBuiltinTyConApp crossTyCon [ty1, ty2]
mk_sum [] = unitTy
mk_sum tys = foldr1 mk_plus tys
mkCrossTypes :: Type -> [Type] -> VM Type
mkCrossTypes = mkBuiltinTyConApps1 crossTyCon
return . mk_sum
. map (mk_tup . map mk_embed)
$ tys
mkCrosses :: [CoreExpr] -> VM CoreExpr
mkCrosses [] = return (Var unitDataConId)
mkCrosses exprs
mkPReprAlts :: [[CoreExpr]] -> VM ([CoreExpr], Type)
mkPReprAlts ess
= do
embed_tc <- builtin embedTyCon
embed_dc <- builtin embedDataCon
cross_tc <- builtin crossTyCon
cross_dc <- builtin crossDataCon
plus_tc <- builtin plusTyCon
left_dc <- builtin leftDataCon
right_dc <- builtin rightDataCon
let mk (left, left_ty) (right, right_ty)
= (mkConApp cross_dc [Type left_ty, Type right_ty, left, right],
mkTyConApp cross_tc [left_ty, right_ty])
return . fst
$ foldr1 mk [(expr, exprType expr) | expr <- exprs]
mkEmbedType :: Type -> VM Type
mkEmbedType ty = mkBuiltinTyConApp embedTyCon [ty]
mkEmbed :: CoreExpr -> VM CoreExpr
mkEmbed expr = mkBuiltinDataConApp embedDataCon
[Type $ exprType expr, expr]
let mk_embed (expr, ty, pa)
= (mkConApp embed_dc [Type ty, pa, expr],
mkTyConApp embed_tc [ty])
mk_cross (expr1, ty1) (expr2, ty2)
= (mkConApp cross_dc [Type ty1, Type ty2, expr1, expr2],
mkTyConApp cross_tc [ty1, ty2])
mk_tup [] = (Var unitDataConId, unitTy)
mk_tup es = foldr1 mk_cross es
mk_sum [] = ([Var unitDataConId], unitTy)
mk_sum [(expr, ty)] = ([expr], ty)
mk_sum ((expr, lty) : es)
= let (alts, rty) = mk_sum es
in
(mkConApp left_dc [Type lty, Type rty, expr]
: [mkConApp right_dc [Type lty, Type rty, alt] | alt <- alts],
mkTyConApp plus_tc [lty, rty])
liftM (mk_sum . map (mk_tup . map mk_embed))
(mapM (mapM init) ess)
where
init expr = let ty = exprType expr
in do
pa <- paDictOfType ty
return (expr, ty, pa)
mkClosureType :: Type -> Type -> VM Type
mkClosureType arg_ty res_ty = mkBuiltinTyConApp closureTyCon [arg_ty, res_ty]
......
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