Commit e4b4159b authored by benl's avatar benl
Browse files

vectoriser: refactoring and cleanups for PRepr

parent b400e86d
......@@ -21,21 +21,11 @@ import Var
import Name
-- debug = False
-- dtrace s x = if debug then pprTrace "Vectoris.Type.PADict" s x else x
-- |Build the PA dictionary function for some type and hoist it to top level.
--
-- The PA dictionary holds fns that convert values to and from their vectorised representations.
-- The PA dictionary holds fns that convert values to and from their vectorised representations.
--
buildPADict
:: TyCon -- ^ tycon of the type being vectorised.
-> TyCon -- ^ tycon of the type used for the vectorised representation.
-> TyCon -- ^ PRepr instance tycon
-> SumRepr -- ^ representation used for the type being vectorised.
-> VM Var -- ^ name of the top-level dictionary function.
-- Recall the definition:
-- @Recall the definition:
-- class class PR (PRepr a) => PA a where
-- toPRepr :: a -> PRepr a
-- fromPRepr :: PRepr a -> a
......@@ -50,8 +40,17 @@ buildPADict
-- $toRepr :: forall a. PA a -> T a -> PRepr (T a)
-- $toPRepr = ...
-- The "..." stuff is filled in by buildPAScAndMethods
-- @
--
buildPADict
:: TyCon -- ^ tycon of the type being vectorised.
-> TyCon -- ^ tycon of the type used for the vectorised representation.
-> TyCon -- ^ PData instance tycon
-> TyCon -- ^ PDatas instance tycon
-> SumRepr -- ^ representation used for the type being vectorised.
-> VM Var -- ^ name of the top-level dictionary function.
buildPADict vect_tc prepr_tc arr_tc repr
buildPADict vect_tc prepr_tc pdata_tc pdatas_tc repr
= polyAbstract tvs $ \args -> -- The args are the dictionaries we lambda
-- abstract over; and they are put in the
-- envt, so when we need a (PA a) we can
......@@ -88,23 +87,21 @@ buildPADict vect_tc prepr_tc arr_tc repr
; return dfun
}
where
tvs = tyConTyVars vect_tc
arg_tys = mkTyVarTys tvs
inst_ty = mkTyConApp vect_tc arg_tys
tvs = tyConTyVars vect_tc
arg_tys = mkTyVarTys tvs
inst_ty = mkTyConApp vect_tc arg_tys
vect_tc_name = getName vect_tc
method args dfun_name (name, build)
= localV
$ do
expr <- build vect_tc prepr_tc arr_tc repr
let body = mkLams (tvs ++ args) expr
raw_var <- newExportedVar (method_name dfun_name name) (exprType body)
let var = raw_var
= localV
$ do expr <- build vect_tc prepr_tc pdata_tc pdatas_tc repr
let body = mkLams (tvs ++ args) expr
raw_var <- newExportedVar (method_name dfun_name name) (exprType body)
let var = raw_var
`setIdUnfolding` mkInlineUnfolding (Just (length args)) body
`setInlinePragma` alwaysInlinePragma
hoistBinding var body
return var
hoistBinding var body
return var
method_call args id = mkApps (Var id) (map Type arg_tys ++ map Var args)
method_name dfun_name name = mkVarOcc $ occNameString dfun_name ++ ('$' : name)
......@@ -64,7 +64,16 @@ mk_fam_inst fam_tc arg_tc
-- Not all lifted backends use the 'toArrPReprs' and 'fromArrPReprs' methods,
-- so we only generate these if the 'PDatas' type family is defined.
--
buildPAScAndMethods :: VM [( String, TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr)]
type PAInstanceBuilder
= TyCon -- ^ Vectorised TyCon
-> TyCon -- ^ Representation TyCon
-> TyCon -- ^ 'PData' TyCon
-> TyCon -- ^ 'PDatas' TyCon
-> SumRepr -- ^ Description of generic representation.
-> VM CoreExpr -- ^ Instance function.
buildPAScAndMethods :: VM [(String, PAInstanceBuilder)]
buildPAScAndMethods
= do hasPDatas <- liftM isJust $ builtin pdatasTyCon
return
......@@ -76,12 +85,11 @@ buildPAScAndMethods
++ (if hasPDatas then
[ ("toArrPReprs", buildToArrPReprs)
, ("fromArrPReprs", buildFromArrPReprs)]
else [])
else [])
buildPRDict :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
buildPRDict vect_tc prepr_tc _ _
buildPRDict :: PAInstanceBuilder
buildPRDict vect_tc prepr_tc _ _ _
= prDictOfPReprInstTyCon inst_ty prepr_tc arg_tys
where
arg_tys = mkTyVarTys (tyConTyVars vect_tc)
......@@ -90,8 +98,8 @@ buildPRDict vect_tc prepr_tc _ _
-- buildToPRepr ---------------------------------------------------------------
-- | Build the 'toRepr' method of the PA class.
buildToPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
buildToPRepr vect_tc repr_tc _ repr
buildToPRepr :: PAInstanceBuilder
buildToPRepr vect_tc repr_tc _ _ repr
= do let arg_ty = mkTyConApp vect_tc ty_args
-- Get the representation type of the argument.
......@@ -164,8 +172,8 @@ buildToPRepr vect_tc repr_tc _ repr
-- buildFromPRepr -------------------------------------------------------------
-- | Build the 'fromPRepr' method of the PA class.
buildFromPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
buildFromPRepr vect_tc repr_tc _ repr
buildFromPRepr :: PAInstanceBuilder
buildFromPRepr vect_tc repr_tc _ _ repr
= do
arg_ty <- mkPReprType res_ty
arg <- newLocalVar (fsLit "x") arg_ty
......@@ -218,8 +226,8 @@ buildFromPRepr vect_tc repr_tc _ repr
-- buildToArrRepr -------------------------------------------------------------
-- | Build the 'toArrRepr' method of the PA class.
buildToArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
buildToArrPRepr vect_tc prepr_tc pdata_tc r
buildToArrPRepr :: PAInstanceBuilder
buildToArrPRepr vect_tc prepr_tc pdata_tc _ r
= do arg_ty <- mkPDataType el_ty
res_ty <- mkPDataType =<< mkPReprType el_ty
arg <- newLocalVar (fsLit "xs") arg_ty
......@@ -301,8 +309,8 @@ buildToArrPRepr vect_tc prepr_tc pdata_tc r
-- buildFromArrPRepr ----------------------------------------------------------
-- | Build the 'fromArrPRepr' method for the PA class.
buildFromArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
buildFromArrPRepr vect_tc prepr_tc pdata_tc r
buildFromArrPRepr :: PAInstanceBuilder
buildFromArrPRepr vect_tc prepr_tc pdata_tc _ r
= do arg_ty <- mkPDataType =<< mkPReprType el_ty
res_ty <- mkPDataType el_ty
arg <- newLocalVar (fsLit "xs") arg_ty
......@@ -378,8 +386,54 @@ buildFromArrPRepr vect_tc prepr_tc pdata_tc r
= do (res', args') <- f res_ty res expr r
return (res', args' ++ args)
-- buildToArrPReprs -----------------------------------------------------------
buildToArrPReprs = error "buildToArrPReprs not done yet"
-- buildToArrPReprs -----------------------------------------------------------
-- | Build the 'toArrPReprs' instance for the PA class.
-- This converts a PData of elements into the generic representation.
buildToArrPReprs :: PAInstanceBuilder
buildToArrPReprs vect_tc prepr_tc _ pdatas_tc r
= do
-- The element type of the argument.
-- eg: 'Tree a b'.
let ty_args = mkTyVarTys $ tyConTyVars vect_tc
let el_ty = mkTyConApp vect_tc ty_args
-- The argument type of the instance.
-- eg: 'PDatas (Tree a b)'
Just arg_ty <- mkPDatasType el_ty
-- The result type.
-- eg: 'PDatas (PRepr (Tree a b))'
Just res_ty <- mkPDatasType =<< mkPReprType el_ty
-- Variable to bind the argument to the instance
-- eg: (xss :: PDatas (Tree a b))
varg <- newLocalVar (fsLit "xss") arg_ty
return $ Lam varg (Var varg)
-- buildFromArrPReprs ---------------------------------------------------------
buildFromArrPReprs = error "buildFromArrPReprs not done yet"
\ No newline at end of file
buildFromArrPReprs :: PAInstanceBuilder
buildFromArrPReprs vect_tc prepr_tc _ pdatas_tc r
= do
-- The element type of the argument.
-- eg: 'Tree a b'.
let ty_args = mkTyVarTys $ tyConTyVars vect_tc
let el_ty = mkTyConApp vect_tc ty_args
-- The argument type of the instance.
-- eg: 'PDatas (PRepr (Tree a b))'
Just arg_ty <- mkPDatasType =<< mkPReprType el_ty
-- The result type.
-- eg: 'PDatas (Tree a b)'
Just res_ty <- mkPDatasType el_ty
-- Variable to bind the argument to the instance
-- eg: (xss :: PDatas (PRepr (Tree a b)))
varg <- newLocalVar (fsLit "xss") arg_ty
return $ Lam varg (Var varg)
......@@ -204,9 +204,10 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
-- Build 'PRepr' and 'PData' instance type constructors and family instances for all
-- type constructors with vectorised representations.
; reprs <- mapM tyConRepr vect_tcs
; repr_tcs <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
; pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
; reprs <- mapM tyConRepr vect_tcs
; repr_tcs <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
; pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
; pdatas_tcs <- zipWith3M buildPDatasTyCon orig_tcs vect_tcs reprs
; let inst_tcs = repr_tcs ++ pdata_tcs
fam_insts = map mkLocalFamInst inst_tcs
; updGEnv $ extendFamEnv fam_insts
......@@ -217,11 +218,12 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
; (_, binds) <- fixV $ \ ~(dfuns, _) ->
do { defTyConPAs (zipLazy vect_tcs dfuns)
; dfuns <- sequence
$ zipWith4 buildTyConBindings
$ zipWith5 buildTyConBindings
orig_tcs
vect_tcs
repr_tcs
pdata_tcs
pdatas_tcs
; binds <- takeHoisted
; return (dfuns, binds)
......@@ -233,14 +235,13 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
}
-- Helpers -------------------
buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var
buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc
= do { vectDataConWorkers orig_tc vect_tc pdata_tc
; repr <- tyConRepr vect_tc
; buildPADict vect_tc prepr_tc pdata_tc repr
}
-- Helpers --------------------------------------------------------------------
buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> TyCon -> VM Var
buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc pdatas_tc
= do vectDataConWorkers orig_tc vect_tc pdata_tc
repr <- tyConRepr vect_tc
buildPADict vect_tc prepr_tc pdata_tc pdatas_tc repr
vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
vectDataConWorkers orig_tc vect_tc arr_tc
......
module Vectorise.Type.PData
( buildPDataTyCon
)
, buildPDatasTyCon )
where
import Vectorise.Monad
......@@ -20,10 +20,13 @@ import MonadUtils
import Control.Monad
buildPDatasTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
buildPDatasTyCon = buildPDataTyCon -- error "buildPDatasTyCon: not finished"
buildPDataTyCon :: TyCon -> TyCon -> SumRepr -> VM TyCon
buildPDataTyCon orig_tc vect_tc repr = fixV $ \repr_tc ->
do
name' <- mkLocalisedName mkPDataTyConOcc orig_name
buildPDataTyCon orig_tc vect_tc repr
= fixV $ \repr_tc ->
do name' <- mkLocalisedName mkPDataTyConOcc orig_name
rhs <- buildPDataTyConRhs orig_name vect_tc repr_tc repr
pdata <- builtin pdataTyCon
......@@ -35,22 +38,20 @@ buildPDataTyCon orig_tc vect_tc repr = fixV $ \repr_tc ->
False -- not GADT syntax
NoParentTyCon
(Just $ mk_fam_inst pdata vect_tc)
where
where
orig_name = tyConName orig_tc
tyvars = tyConTyVars vect_tc
rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
tyvars = tyConTyVars vect_tc
rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
buildPDataTyConRhs :: Name -> TyCon -> TyCon -> SumRepr -> VM AlgTyConRhs
buildPDataTyConRhs orig_name vect_tc repr_tc repr
= do
data_con <- buildPDataDataCon orig_name vect_tc repr_tc repr
= do data_con <- buildPDataDataCon orig_name vect_tc repr_tc repr
return $ DataTyCon { data_cons = [data_con], is_enum = False }
buildPDataDataCon :: Name -> TyCon -> TyCon -> SumRepr -> VM DataCon
buildPDataDataCon orig_name vect_tc repr_tc repr
= do
dc_name <- mkLocalisedName mkPDataDataConOcc orig_name
= do dc_name <- mkLocalisedName mkPDataDataConOcc orig_name
comp_tys <- sum_tys repr
liftDs $ buildDataCon dc_name
......@@ -64,10 +65,10 @@ buildPDataDataCon orig_name vect_tc repr_tc repr
comp_tys
(mkFamilyTyConApp repr_tc (mkTyVarTys tvs))
repr_tc
where
where
tvs = tyConTyVars vect_tc
sum_tys EmptySum = return []
sum_tys EmptySum = return []
sum_tys (UnarySum r) = con_tys r
sum_tys (Sum { repr_sel_ty = sel_ty
, repr_cons = cons })
......@@ -75,7 +76,7 @@ buildPDataDataCon orig_name vect_tc repr_tc repr
con_tys (ConRepr _ r) = prod_tys r
prod_tys EmptyProd = return []
prod_tys EmptyProd = return []
prod_tys (UnaryProd r) = liftM singleton (comp_ty r)
prod_tys (Prod { repr_comps = comps }) = mapM comp_ty comps
......
......@@ -11,7 +11,7 @@ module Vectorise.Utils.Base (
mkPReprType,
mkPArrayType, splitPrimTyCon,
mkPArray,
mkPDataType,
mkPDataType, mkPDatasType,
mkBuiltinCo,
mkVScrut,
......@@ -37,7 +37,7 @@ import Outputable
import FastString
import Control.Monad (liftM)
import Data.Maybe
-- Simple Types ---------------------------------------------------------------
voidType :: VM Type
......@@ -67,36 +67,38 @@ dataConTagZ :: DataCon -> Int
dataConTagZ con = dataConTag con - fIRST_TAG
-- Type Construction ----------------------------------------------------------
-- | Make an application of a builtin type constructor to some arguments.
mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type
mkBuiltinTyConApp get_tc tys
= do
tc <- builtin get_tc
= do tc <- builtin get_tc
return $ mkTyConApp tc tys
mkBuiltinTyConApps :: (Builtins -> TyCon) -> [Type] -> Type -> VM Type
mkBuiltinTyConApps get_tc tys ty
= do
tc <- builtin get_tc
= do tc <- builtin get_tc
return $ foldr (mk tc) ty tys
where
mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
-- | Make an application of the 'Wrap' type constructor.
mkWrapType :: Type -> VM Type
mkWrapType ty = mkBuiltinTyConApp wrapTyCon [ty]
mkWrapType ty = mkBuiltinTyConApp wrapTyCon [ty]
-- | Make an application of the closure type constructor.
mkClosureTypes :: [Type] -> Type -> VM Type
mkClosureTypes = mkBuiltinTyConApps closureTyCon
-- | Make an application of the 'PRepr' type constructor.
mkPReprType :: Type -> VM Type
mkPReprType ty = mkBuiltinTyConApp preprTyCon [ty]
-- |Wrap a type into 'PArray', treating unboxed types specially.
--
-- | Wrap a type into 'PArray', treating unboxed types specially.
mkPArrayType :: Type -> VM Type
mkPArrayType ty
| Just tycon <- splitPrimTyCon ty
......@@ -105,8 +107,23 @@ mkPArrayType ty
}
mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty]
-- | Make an appliction of the 'PData' tycon to some argument.
mkPDataType :: Type -> VM Type
mkPDataType ty
= mkBuiltinTyConApp pdataTyCon [ty]
-- | Make an application of the 'PDatas' tycon to some argument.
mkPDatasType :: Type -> VM (Maybe Type)
mkPDatasType ty
= do mtc <- builtin pdatasTyCon
case mtc of
Nothing -> return Nothing
Just tc' -> return $ Just $ mkTyConApp tc' [ty]
-- |Checks if a type constructor is defined in 'GHC.Prim' (e.g., 'Int#'); if so, returns it.
--
splitPrimTyCon :: Type -> Maybe TyCon
splitPrimTyCon ty
| Just (tycon, []) <- splitTyConApp_maybe ty
......@@ -115,22 +132,30 @@ splitPrimTyCon ty
| otherwise = Nothing
------
mkPArray :: Type -> CoreExpr -> CoreExpr -> VM CoreExpr
mkPArray ty len dat = do
tc <- builtin parrayTyCon
let [dc] = tyConDataCons tc
return $ mkConApp dc [Type ty, len, dat]
mkPDataType :: Type -> VM Type
mkPDataType ty = mkBuiltinTyConApp pdataTyCon [ty]
-- CoreExpr Construction ------------------------------------------------------
-- | Make an application of the 'PArray' data constructor.
mkPArray
:: Type -- ^ Element type
-> CoreExpr -- ^ 'Int' for the array length.
-> CoreExpr -- ^ 'PData' for the array data.
-> VM CoreExpr
mkPArray ty len dat
= do tc <- builtin parrayTyCon
let [dc] = tyConDataCons tc
return $ mkConApp dc [Type ty, len, dat]
-- Coercion Construction -----------------------------------------------------
-- | Make a coersion to some builtin type.
mkBuiltinCo :: (Builtins -> TyCon) -> VM Coercion
mkBuiltinCo get_tc
= do
tc <- builtin get_tc
= do tc <- builtin get_tc
return $ mkTyConAppCo tc []
-------------------------------------------------------------------------------
mkVScrut :: VExpr -> VM (CoreExpr, CoreExpr, TyCon, [Type])
mkVScrut (ve, le)
= do
......
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