mkBuiltinCo, voidType, mkWrapType,
mkPADictType, mkPArrayType, mkPDataType, mkPReprType, mkPArray,
+ mkBuiltinTyConApps, mkClosureTypes,
pdataReprTyCon, pdataReprDataCon, mkVScrut,
prDictOfType, prDFunOfTyCon,
paDictArgType, paDictOfType, paDFunType,
- paMethod, wrapPR, replicatePD, emptyPD, packPD, packByTagPD,
+ paMethod, wrapPR, replicatePD, emptyPD, packByTagPD,
combinePD,
liftPD,
zipScalars, scalarClosure,
polyAbstract, polyApply, polyVApply, polyArity,
Inline(..), addInlineArity, inlineMe,
hoistBinding, hoistExpr, hoistPolyVExpr, takeHoisted,
- buildClosure, buildClosures,
- mkClosureApp
) where
-import VectCore
-import VectMonad
-import MkCore ( mkCoreTup, mkCoreTupTy, mkWildCase )
+import Vectorise.Monad
+import Vectorise.Env
+import Vectorise.Vect
+import Vectorise.Builtins
+
import CoreSyn
import CoreUtils
import CoreUnfold ( mkInlineRule )
import Var
import MkId ( unwrapFamInstScrut )
import Id ( setIdUnfolding )
-import TysWiredIn
-import BasicTypes ( Boxity(..) )
+import BasicTypes
import Literal ( Literal, mkMachInt )
+
import Outputable
import FastString
mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
voidType :: VM Type
-voidType = mkBuiltinTyConApp VectMonad.voidTyCon []
+voidType = mkBuiltinTyConApp voidTyCon []
mkWrapType :: Type -> VM Type
mkWrapType ty = mkBuiltinTyConApp wrapTyCon [ty]
+
mkClosureTypes :: [Type] -> Type -> VM Type
mkClosureTypes = mkBuiltinTyConApps closureTyCon
. maybeCantVectoriseM "No PR dictionary for tycon" (ppr tycon)
$ lookupTyConPR tycon
+
paDictArgType :: TyVar -> VM (Maybe Type)
paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
where
go _ _ = return Nothing
-paDictOfType :: Type -> VM CoreExpr
-paDictOfType ty = paDictOfTyApp ty_fn ty_args
+
+-- | Get the PA dictionary for some type, or `Nothing` if there isn't one.
+paDictOfType :: Type -> VM (Maybe CoreExpr)
+paDictOfType ty
+ = paDictOfTyApp ty_fn ty_args
where
(ty_fn, ty_args) = splitAppTys ty
-paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
-paDictOfTyApp ty_fn ty_args
- | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
-paDictOfTyApp (TyVarTy tv) ty_args
- = do
- dfun <- maybeV (lookupTyVarPA tv)
- paDFunApply dfun ty_args
-paDictOfTyApp (TyConApp tc _) ty_args
- = do
- dfun <- maybeCantVectoriseM "No PA dictionary for tycon" (ppr tc)
- $ lookupTyConPA tc
- paDFunApply (Var dfun) ty_args
-paDictOfTyApp ty _
- = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
+ paDictOfTyApp :: Type -> [Type] -> VM (Maybe CoreExpr)
+ paDictOfTyApp ty_fn ty_args
+ | Just ty_fn' <- coreView ty_fn
+ = paDictOfTyApp ty_fn' ty_args
+
+ paDictOfTyApp (TyVarTy tv) ty_args
+ = do dfun <- maybeV (lookupTyVarPA tv)
+ liftM Just $ paDFunApply dfun ty_args
+
+ paDictOfTyApp (TyConApp tc _) ty_args
+ = do mdfun <- lookupTyConPA tc
+ case mdfun of
+ Nothing
+ -> pprTrace "VectUtils.paDictOfType"
+ (vcat [ text "No PA dictionary"
+ , text "for tycon: " <> ppr tc
+ , text "in type: " <> ppr ty])
+ $ return Nothing
+
+ Just dfun -> liftM Just $ paDFunApply (Var dfun) ty_args
+
+ paDictOfTyApp ty _
+ = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
+
+
paDFunType :: TyCon -> VM Type
paDFunType tc
paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
paDFunApply dfun tys
- = do
- dicts <- mapM paDictOfType tys
+ = do Just dicts <- liftM sequence $ mapM paDictOfType tys
return $ mkApps (mkTyApps dfun tys) dicts
+
paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
paMethod _ name ty
| Just tycon <- splitPrimTyCon ty
paMethod method _ ty
= do
- fn <- builtin method
- dict <- paDictOfType ty
+ fn <- builtin method
+ Just dict <- paDictOfType ty
return $ mkApps (Var fn) [Type ty, dict]
prDictOfType :: Type -> VM CoreExpr
wrapPR :: Type -> VM CoreExpr
wrapPR ty
= do
- pa_dict <- paDictOfType ty
- pr_dfun <- prDFunOfTyCon =<< builtin wrapTyCon
+ Just pa_dict <- paDictOfType ty
+ pr_dfun <- prDFunOfTyCon =<< builtin wrapTyCon
return $ mkApps pr_dfun [Type ty, pa_dict]
replicatePD :: CoreExpr -> CoreExpr -> VM CoreExpr
emptyPD :: Type -> VM CoreExpr
emptyPD = paMethod emptyPDVar "emptyPD"
-packPD :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr
-packPD ty xs len sel = liftM (`mkApps` [xs, len, sel])
- (paMethod packPDVar "packPD" ty)
-
packByTagPD :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr
-> VM CoreExpr
packByTagPD ty xs len tags t
where
n = length xs
+-- | Like `replicatePD` but use the lifting context in the vectoriser state.
liftPD :: CoreExpr -> VM CoreExpr
liftPD x
= do
scalarClosure :: [Type] -> Type -> CoreExpr -> CoreExpr -> VM CoreExpr
scalarClosure arg_tys res_ty scalar_fun array_fun
= do
- ctr <- builtin (closureCtrFun $ length arg_tys)
- pas <- mapM paDictOfType (init arg_tys)
+ ctr <- builtin (closureCtrFun $ length arg_tys)
+ Just pas <- liftM sequence $ mapM paDictOfType (init arg_tys)
return $ Var ctr `mkTyApps` (arg_tys ++ [res_ty])
`mkApps` (pas ++ [scalar_fun, array_fun])
polyApply :: CoreExpr -> [Type] -> VM CoreExpr
polyApply expr tys
- = do
- dicts <- mapM paDictOfType tys
+ = do Just dicts <- liftM sequence $ mapM paDictOfType tys
return $ expr `mkTyApps` tys `mkApps` dicts
polyVApply :: VExpr -> [Type] -> VM VExpr
polyVApply expr tys
- = do
- dicts <- mapM paDictOfType tys
- return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr
+ = do Just dicts <- liftM sequence $ mapM paDictOfType tys
+ return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr
+-- Inline ---------------------------------------------------------------------
+-- | Records whether we should inline a particular binding.
+data Inline
+ = Inline Arity
+ | DontInline
-data Inline = Inline Int -- arity
- | DontInline
-
+-- | Add to the arity contained within an `Inline`, if any.
addInlineArity :: Inline -> Int -> Inline
addInlineArity (Inline m) n = Inline (m+n)
addInlineArity DontInline _ = DontInline
+-- | Says to always inline a binding.
inlineMe :: Inline
inlineMe = Inline 0
+
+-- Hoising --------------------------------------------------------------------
hoistBinding :: Var -> CoreExpr -> VM ()
hoistBinding v e = updGEnv $ \env ->
env { global_bindings = (v,e) : global_bindings env }
where
mk_inline var = case inl of
Inline arity -> var `setIdUnfolding`
- mkInlineRule InlSat expr arity
+ mkInlineRule expr (Just arity)
DontInline -> var
hoistVExpr :: VExpr -> Inline -> VM VVar
Nothing -> return (vexpr, lexpr)
-}
-mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
-mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
- = do
- dict <- paDictOfType env_ty
- mkv <- builtin closureVar
- mkl <- builtin liftedClosureVar
- return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
- Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
-
-mkClosureApp :: Type -> Type -> VExpr -> VExpr -> VM VExpr
-mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
- = do
- vapply <- builtin applyVar
- lapply <- builtin liftedApplyVar
- lc <- builtin liftingContext
- return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
- Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [Var lc, lclo, larg])
-
-buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
-buildClosures _ _ [] _ mk_body
- = mk_body
-buildClosures tvs vars [arg_ty] res_ty mk_body
- = -- liftM vInlineMe $
- 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
- -- liftM vInlineMe
- buildClosure tvs vars arg_ty res_ty'
- . hoistPolyVExpr tvs (Inline (length vars + 1))
- $ do
- 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^>)
--- where
--- 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] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
-buildClosure tvs vars arg_ty res_ty mk_body
- = do
- (env_ty, env, bind) <- buildEnv vars
- env_bndr <- newLocalVVar (fsLit "env") env_ty
- arg_bndr <- newLocalVVar (fsLit "arg") arg_ty
-
- fn <- hoistPolyVExpr tvs (Inline 2)
- $ do
- lc <- builtin liftingContext
- body <- mk_body
- return -- . vInlineMe
- . vLams lc [env_bndr, arg_bndr]
- $ bind (vVar env_bndr)
- (vVarApps lc body (vars ++ [arg_bndr]))
-
- mkClosure arg_ty res_ty env_ty fn env
-
-buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
-buildEnv [] = do
- ty <- voidType
- void <- builtin voidVar
- pvoid <- builtin pvoidVar
- return (ty, vVar (void, pvoid), \_ body -> body)
-
-buildEnv [v] = return (vVarType v, vVar v,
- \env body -> vLet (vNonRec v env) body)
-
-buildEnv vs
- = do
-
- (lenv_tc, lenv_tyargs) <- pdataReprTyCon ty
-
- let venv_con = tupleCon Boxed (length vs)
- [lenv_con] = tyConDataCons lenv_tc
-
- venv = mkCoreTup (map Var vvs)
- lenv = Var (dataConWrapId lenv_con)
- `mkTyApps` lenv_tyargs
- `mkApps` map Var lvs
-
- vbind env body = mkWildCase env ty (exprType body)
- [(DataAlt venv_con, vvs, body)]
-
- lbind env body =
- let scrut = unwrapFamInstScrut lenv_tc lenv_tyargs env
- in
- mkWildCase scrut (exprType scrut) (exprType body)
- [(DataAlt lenv_con, lvs, body)]
-
- bind (venv, lenv) (vbody, lbody) = (vbind venv vbody,
- lbind lenv lbody)
-
- return (ty, (venv, lenv), bind)
- where
- (vvs, lvs) = unzip vs
- tys = map vVarType vs
- ty = mkCoreTupTy tys
projects / ghc-hetmet.git / blobdiff