Commit 112780e0 authored by benl@ouroborus.net's avatar benl@ouroborus.net
Browse files

Comments and formatting to vectoriser

parent 5ee7f0e6
......@@ -36,8 +36,8 @@ import Outputable
------------------------------------------------------
buildSynTyCon :: Name -> [TyVar]
-> SynTyConRhs
-> Kind -- Kind of the RHS
-> Maybe (TyCon, [Type]) -- family instance if applicable
-> Kind -- ^ Kind of the RHS
-> Maybe (TyCon, [Type]) -- ^ family instance if applicable
-> TcRnIf m n TyCon
buildSynTyCon tc_name tvs rhs@(OpenSynTyCon {}) rhs_kind _
......@@ -61,12 +61,12 @@ buildSynTyCon tc_name tvs rhs@(SynonymTyCon {}) rhs_kind mb_family
------------------------------------------------------
buildAlgTyCon :: Name -> [TyVar]
-> ThetaType -- Stupid theta
-> ThetaType -- ^ Stupid theta
-> AlgTyConRhs
-> RecFlag
-> Bool -- True <=> want generics functions
-> Bool -- True <=> was declared in GADT syntax
-> Maybe (TyCon, [Type]) -- family instance if applicable
-> Bool -- ^ True <=> want generics functions
-> Bool -- ^ True <=> was declared in GADT syntax
-> Maybe (TyCon, [Type]) -- ^ family instance if applicable
-> TcRnIf m n TyCon
buildAlgTyCon tc_name tvs stupid_theta rhs is_rec want_generics gadt_syn
......@@ -84,7 +84,7 @@ buildAlgTyCon tc_name tvs stupid_theta rhs is_rec want_generics gadt_syn
; return tycon
}
-- If a family tycon with instance types is given, the current tycon is an
-- | If a family tycon with instance types is given, the current tycon is an
-- instance of that family and we need to
--
-- (1) create a coercion that identifies the family instance type and the
......@@ -132,9 +132,9 @@ mkDataTyConRhs cons
}
mkNewTyConRhs :: Name -> TyCon -> DataCon -> TcRnIf m n AlgTyConRhs
-- Monadic because it makes a Name for the coercion TyCon
-- We pass the Name of the parent TyCon, as well as the TyCon itself,
-- because the latter is part of a knot, whereas the former is not.
-- ^ Monadic because it makes a Name for the coercion TyCon
-- We pass the Name of the parent TyCon, as well as the TyCon itself,
-- because the latter is part of a knot, whereas the former is not.
mkNewTyConRhs tycon_name tycon con
= do { co_tycon_name <- newImplicitBinder tycon_name mkNewTyCoOcc
; let co_tycon = mkNewTypeCoercion co_tycon_name tycon etad_tvs etad_rhs
......
......@@ -439,6 +439,8 @@ newTyVar fs k
u <- liftDs newUnique
return $ mkTyVar (mkSysTvName u fs) k
-- | Add a mapping between a global var and its vectorised version to the state.
defGlobalVar :: Var -> Var -> VM ()
defGlobalVar v v' = updGEnv $ \env ->
env { global_vars = extendVarEnv (global_vars env) v v'
......@@ -448,14 +450,14 @@ defGlobalVar v v' = updGEnv $ \env ->
upd env | isExportedId v = extendVarEnv env v (v, v')
| otherwise = env
-- Var ------------------------------------------------------------------------
-- | Lookup the vectorised and\/or lifted versions of this variable.
-- If it's in the global environment we get the vectorised version.
-- If it's in the local environment we get both the vectorised and lifted version.
--
lookupVar :: Var -> VM (Scope Var (Var, Var))
lookupVar v
= do
r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
= do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
case r of
Just e -> return (Local e)
Nothing -> liftM Global
......@@ -581,6 +583,8 @@ lookupFamInst tycon tys
(ppr $ mkTyConApp tycon tys)
}
-- | Run a vectorisation computation.
initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
initV pkg hsc_env guts info p
= do
......
......@@ -45,13 +45,16 @@ import Data.List ( inits, tails, zipWith4, zipWith5 )
-- ----------------------------------------------------------------------------
-- Types
-- | Vectorise a type constructor.
vectTyCon :: TyCon -> VM TyCon
vectTyCon tc
| isFunTyCon tc = builtin closureTyCon
| isBoxedTupleTyCon tc = return tc
| isUnLiftedTyCon tc = return tc
| otherwise = maybeCantVectoriseM "Tycon not vectorised:" (ppr tc)
$ lookupTyCon tc
| otherwise
= maybeCantVectoriseM "Tycon not vectorised: " (ppr tc)
$ lookupTyCon tc
vectAndLiftType :: Type -> VM (Type, Type)
vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
......@@ -67,6 +70,7 @@ vectAndLiftType ty
(tyvars, mono_ty) = splitForAllTys ty
-- | Vectorise a type.
vectType :: Type -> VM Type
vectType ty | Just ty' <- coreView ty = vectType ty'
vectType (TyVarTy tv) = return $ TyVarTy tv
......@@ -87,6 +91,7 @@ vectType ty = cantVectorise "Can't vectorise type" (ppr ty)
vectAndBoxType :: Type -> VM Type
vectAndBoxType ty = vectType ty >>= boxType
-- | Add quantified vars and dictionary parameters to the front of a type.
abstractType :: [TyVar] -> [Type] -> Type -> Type
abstractType tyvars dicts = mkForAllTys tyvars . mkFunTys dicts
......@@ -102,6 +107,7 @@ boxType ty
case r of
Just tycon' -> return $ mkTyConApp tycon' []
Nothing -> return ty
boxType ty = return ty
-- ----------------------------------------------------------------------------
......@@ -109,14 +115,21 @@ boxType ty = return ty
type TyConGroup = ([TyCon], UniqSet TyCon)
-- | Vectorise a type environment.
-- The type environment contains all the type things defined in a module.
vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
vectTypeEnv env
= do
cs <- readGEnv $ mk_map . global_tycons
-- Split the list of TyCons into the ones we have to vectorise vs the
-- ones we can pass through unchanged. We also pass through algebraic
-- types that use non Haskell98 features, as we don't handle those.
let (conv_tcs, keep_tcs) = classifyTyCons cs groups
keep_dcs = concatMap tyConDataCons keep_tcs
zipWithM_ defTyCon keep_tcs keep_tcs
zipWithM_ defDataCon keep_dcs keep_dcs
new_tcs <- vectTyConDecls conv_tcs
let orig_tcs = keep_tcs ++ conv_tcs
......@@ -151,6 +164,7 @@ vectTypeEnv env
mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
-- | Vectorise some (possibly recursively defined) type constructors.
vectTyConDecls :: [TyCon] -> VM [TyCon]
vectTyConDecls tcs = fixV $ \tcs' ->
do
......@@ -848,8 +862,8 @@ paMethods = [("dictPRepr", buildPRDict),
("fromArrPRepr", buildFromArrPRepr)]
-- | Split the given tycons into two sets depending on whether they have to be
-- converted (first list) or not (second list). The first argument contains
-- information about the conversion status of external tycons:
-- converted (first list) or not (second list). The first argument contains
-- information about the conversion status of external tycons:
--
-- * tycons which have converted versions are mapped to True
-- * tycons which are not changed by vectorisation are mapped to False
......
......@@ -281,6 +281,7 @@ combinePD ty len sel xs
where
n = length xs
-- | Like `replicatePD` but use the lifting context in the vectoriser state.
liftPD :: CoreExpr -> VM CoreExpr
liftPD x
= do
......@@ -358,6 +359,8 @@ addInlineArity DontInline _ = DontInline
inlineMe :: Inline
inlineMe = Inline 0
-- Hoising --------------------------------------------------------------------
hoistBinding :: Var -> CoreExpr -> VM ()
hoistBinding v e = updGEnv $ \env ->
env { global_bindings = (v,e) : global_bindings env }
......@@ -412,24 +415,24 @@ boxExpr ty (vexpr, lexpr)
Nothing -> return (vexpr, lexpr)
-}
-- Closures -------------------------------------------------------------------
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
= do Just 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
= 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
......@@ -471,6 +474,8 @@ buildClosure tvs vars arg_ty res_ty mk_body
mkClosure arg_ty res_ty env_ty fn env
-- Environments ---------------------------------------------------------------
buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
buildEnv [] = do
ty <- voidType
......
......@@ -42,51 +42,109 @@ vectorise backend guts = do
hsc_env <- getHscEnv
liftIO $ vectoriseIO backend hsc_env guts
-- | Vectorise a single monad, given its HscEnv (code gen environment).
vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
vectoriseIO backend hsc_env guts
= do
= do -- Get information about currently loaded external packages.
eps <- hscEPS hsc_env
-- Combine vectorisation info from the current module, and external ones.
let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
-- Run the main VM computation.
Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
return (guts' { mg_vect_info = info' })
-- | Vectorise a single module, in the VM monad.
vectModule :: ModGuts -> VM ModGuts
vectModule guts
= do
= do -- Vectorise the type environment.
-- This may add new TyCons and DataCons.
-- TODO: What new binds do we get back here?
(types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
-- TODO: What is this?
let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
updGEnv (setFamInstEnv fam_inst_env')
-- dicts <- mapM buildPADict pa_insts
-- workers <- mapM vectDataConWorkers pa_insts
-- Vectorise all the top level bindings.
binds' <- mapM vectTopBind (mg_binds guts)
return $ guts { mg_types = types'
, mg_binds = Rec tc_binds : binds'
, mg_fam_inst_env = fam_inst_env'
, mg_fam_insts = mg_fam_insts guts ++ fam_insts
}
-- | Try to vectorise a top-level binding.
-- If it doesn't vectorise then return it unharmed.
--
-- For example, for the binding
--
-- @
-- foo :: Int -> Int
-- foo = \x -> x + x
-- @
--
-- we get
-- @
-- foo :: Int -> Int
-- foo = \x -> vfoo $: x
--
-- v_foo :: Closure void vfoo lfoo
-- v_foo = closure vfoo lfoo void
--
-- vfoo :: Void -> Int -> Int
-- vfoo = ...
--
-- lfoo :: PData Void -> PData Int -> PData Int
-- lfoo = ...
-- @
--
-- @vfoo@ is the "vectorised", or scalar, version that does the same as the original
-- function foo, but takes an explicit environment.
--
-- @lfoo@ is the "lifted" version that works on arrays.
--
-- @v_foo@ combines both of these into a `Closure` that also contains the
-- environment.
--
-- The original binding @foo@ is rewritten to call the vectorised version
-- present in the closure.
--
vectTopBind :: CoreBind -> VM CoreBind
vectTopBind b@(NonRec var expr)
= do
(inline, expr') <- vectTopRhs var expr
var' <- vectTopBinder var inline expr'
hs <- takeHoisted
cexpr <- tryConvert var var' expr
= do
(inline, expr') <- vectTopRhs var expr
var' <- vectTopBinder var inline expr'
-- Vectorising the body may create other top-level bindings.
hs <- takeHoisted
-- To get the same functionality as the original body we project
-- out its vectorised version from the closure.
cexpr <- tryConvert var var' expr
return . Rec $ (var, cexpr) : (var', expr') : hs
`orElseV`
return b
vectTopBind b@(Rec bs)
= do
(vars', _, exprs') <- fixV $ \ ~(_, inlines, rhss) ->
do
vars' <- sequence [vectTopBinder var inline rhs
| (var, ~(inline, rhs))
<- zipLazy vars (zip inlines rhss)]
(inlines', exprs') <- mapAndUnzipM (uncurry vectTopRhs) bs
return (vars', inlines', exprs')
= do
(vars', _, exprs')
<- fixV $ \ ~(_, inlines, rhss) ->
do vars' <- sequence [vectTopBinder var inline rhs
| (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
(inlines', exprs')
<- mapAndUnzipM (uncurry vectTopRhs) bs
return (vars', inlines', exprs')
hs <- takeHoisted
cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
......@@ -95,11 +153,22 @@ vectTopBind b@(Rec bs)
where
(vars, exprs) = unzip bs
-- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
-- used inside of fixV in vectTopBind
vectTopBinder :: Var -> Inline -> CoreExpr -> VM Var
-- | Make the vectorised version of this top level binder, and add the mapping
-- between it and the original to the state. For some binder @foo@ the vectorised
-- version is @$v_foo@
--
-- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
-- used inside of fixV in vectTopBind
vectTopBinder
:: Var -- ^ Name of the binding.
-> Inline -- ^ Whether it should be inlined, used to annotate it.
-> CoreExpr -- ^ RHS of the binding, used to set the `Unfolding` of the returned `Var`.
-> VM Var -- ^ Name of the vectorised binding.
vectTopBinder var inline expr
= do
= do
-- Vectorise the type attached to the var.
vty <- vectType (idType var)
var' <- liftM (`setIdUnfolding` unfolding) $ cloneId mkVectOcc var vty
defGlobalVar var var'
......@@ -109,22 +178,37 @@ vectTopBinder var inline expr
Inline arity -> mkInlineRule expr (Just arity)
DontInline -> noUnfolding
vectTopRhs :: Var -> CoreExpr -> VM (Inline, CoreExpr)
-- | Vectorise the RHS of a top-level binding, in an empty local environment.
vectTopRhs
:: Var -- ^ Name of the binding.
-> CoreExpr -- ^ Body of the binding.
-> VM (Inline, CoreExpr)
vectTopRhs var expr
= closedV
$ do
(inline, vexpr) <- inBind var
$ vectPolyExpr (isLoopBreaker $ idOccInfo var)
= dtrace (vcat [text "vectTopRhs", ppr expr])
$ closedV
$ do (inline, vexpr) <- inBind var
$ vectPolyExpr (isLoopBreaker $ idOccInfo var)
(freeVars expr)
return (inline, vectorised vexpr)
tryConvert :: Var -> Var -> CoreExpr -> VM CoreExpr
-- | Project out the vectorised version of a binding from some closure,
-- or return the original body if that doesn't work.
tryConvert
:: Var -- ^ Name of the original binding (eg @foo@)
-> Var -- ^ Name of vectorised version of binding (eg @$vfoo@)
-> CoreExpr -- ^ The original body of the binding.
-> VM CoreExpr
tryConvert var vect_var rhs
= fromVect (idType var) (Var vect_var) `orElseV` return rhs
-- ----------------------------------------------------------------------------
-- Bindings
-- | Vectorise a binder variable, along with its attached type.
vectBndr :: Var -> VM VVar
vectBndr v
= do
......@@ -136,6 +220,9 @@ vectBndr v
where
mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
-- | Vectorise a binder variable, along with its attached type,
-- but give the result a new name.
vectBndrNew :: Var -> FastString -> VM VVar
vectBndrNew v fs
= do
......@@ -146,6 +233,8 @@ vectBndrNew v fs
where
upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }
-- | Vectorise a binder then run a computation with that binder in scope.
vectBndrIn :: Var -> VM a -> VM (VVar, a)
vectBndrIn v p
= localV
......@@ -154,6 +243,8 @@ vectBndrIn v p
x <- p
return (vv, x)
-- | Vectorise a binder, give it a new name, then run a computation with that binder in scope.
vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
vectBndrNewIn v fs p
= localV
......@@ -162,6 +253,7 @@ vectBndrNewIn v fs p
x <- p
return (vv, x)
-- | Vectorise some binders, then run a computation with them in scope.
vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
vectBndrsIn vs p
= localV
......@@ -170,13 +262,17 @@ vectBndrsIn vs p
x <- p
return (vvs, x)
-- ----------------------------------------------------------------------------
-- Expressions
-- | Vectorise a variable, producing the vectorised and lifted versions.
vectVar :: Var -> VM VExpr
vectVar v
= do
= do
-- lookup the variable from the environment.
r <- lookupVar v
case r of
Local (vv,lv) -> return (Var vv, Var lv)
Global vv -> do
......@@ -184,30 +280,42 @@ vectVar v
lexpr <- liftPD vexpr
return (vexpr, lexpr)
-- | Like `vectVar` but also add type applications to the variables.
vectPolyVar :: Var -> [Type] -> VM VExpr
vectPolyVar v tys
= do
vtys <- mapM vectType tys
r <- lookupVar v
vtys <- mapM vectType tys
r <- lookupVar v
case r of
Local (vv, lv) -> liftM2 (,) (polyApply (Var vv) vtys)
(polyApply (Var lv) vtys)
Global poly -> do
vexpr <- polyApply (Var poly) vtys
lexpr <- liftPD vexpr
return (vexpr, lexpr)
Local (vv, lv)
-> liftM2 (,) (polyApply (Var vv) vtys)
(polyApply (Var lv) vtys)
Global poly
-> do vexpr <- polyApply (Var poly) vtys
lexpr <- liftPD vexpr
return (vexpr, lexpr)
-- | Lifted literals are created by replicating them.
vectLiteral :: Literal -> VM VExpr
vectLiteral lit
= do
lexpr <- liftPD (Lit lit)
return (Lit lit, lexpr)
vectPolyExpr :: Bool -> CoreExprWithFVs -> VM (Inline, VExpr)
-- | Vectorise a polymorphic expression
vectPolyExpr
:: Bool -- ^ When vectorising the RHS of a binding, whether that
-- binding is a loop breaker.
-> CoreExprWithFVs
-> VM (Inline, VExpr)
vectPolyExpr loop_breaker (_, AnnNote note expr)
= do
(inline, expr') <- vectPolyExpr loop_breaker expr
= do (inline, expr') <- vectPolyExpr loop_breaker expr
return (inline, vNote note expr')
vectPolyExpr loop_breaker expr
= do
arity <- polyArity tvs
......@@ -219,13 +327,17 @@ vectPolyExpr loop_breaker expr
where
(tvs, mono) = collectAnnTypeBinders expr
-- | Vectorise a core expression.
vectExpr :: CoreExprWithFVs -> VM VExpr
vectExpr (_, AnnType ty)
= liftM vType (vectType ty)
vectExpr (_, AnnVar v) = vectVar v
vectExpr (_, AnnVar v)
= vectVar v
vectExpr (_, AnnLit lit) = vectLiteral lit
vectExpr (_, AnnLit lit)
= vectLiteral lit
vectExpr (_, AnnNote note expr)
= liftM (vNote note) (vectExpr expr)
......@@ -247,12 +359,22 @@ vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon]
-- TODO: Avoid using closure application for dictionaries.
-- vectExpr (_, AnnApp fn arg)
-- | if is application of dictionary
-- just use regular app instead of closure app.
-- for lifted version.
-- do liftPD (sub a dNumber)
-- lift the result of the selection, not sub and dNumber seprately.
vectExpr (_, AnnApp fn arg)
= do
arg_ty' <- vectType arg_ty
res_ty' <- vectType res_ty
fn' <- vectExpr fn
arg' <- vectExpr arg
mkClosureApp arg_ty' res_ty' fn' arg'
where
(arg_ty, res_ty) = splitFunTy . exprType $ deAnnotate fn
......@@ -296,7 +418,14 @@ onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
vectExpr e = cantVectorise "Can't vectorise expression" (ppr $ deAnnotate e)
vectFnExpr :: Bool -> Bool -> CoreExprWithFVs -> VM (Inline, VExpr)
-- | Vectorise an expression with an outer lambda abstraction.
vectFnExpr
:: Bool -- ^ When the RHS of a binding, whether that binding should be inlined.
-> Bool -- ^ Whether the binding is a loop breaker.
-> CoreExprWithFVs -- ^ Expression to vectorise. Must have an outer `AnnLam`.
-> VM (Inline, VExpr)
vectFnExpr inline loop_breaker e@(fvs, AnnLam bndr _)
| isId bndr = onlyIfV (isEmptyVarSet fvs)
(mark DontInline . vectScalarLam bs $ deAnnotate body)
......@@ -308,7 +437,12 @@ vectFnExpr _ _ e = mark DontInline $ vectExpr e
mark :: Inline -> VM a -> VM (Inline, a)
mark b p = do { x <- p; return (b,x) }
vectScalarLam :: [Var] -> CoreExpr -> VM VExpr
-- | Vectorise a function where are the args have scalar type, that is Int, Float or Double.
vectScalarLam
:: [Var] -- ^ Bound variables of function.
-> CoreExpr -- ^ Function body.
-> VM VExpr
vectScalarLam args body
= do
scalars <- globalScalars
......@@ -317,23 +451,24 @@ vectScalarLam args body
&& is_scalar (extendVarSetList scalars args) body
&& uses scalars body)
$ do
fn_var <- hoistExpr (fsLit "fn") (mkLams args body) DontInline
zipf <- zipScalars arg_tys res_ty
clo <- scalarClosure arg_tys res_ty (Var fn_var)
fn_var <- hoistExpr (fsLit "fn") (mkLams args body) DontInline
zipf <- zipScalars arg_tys res_ty
clo <- scalarClosure arg_tys res_ty (Var fn_var)
(zipf `App` Var fn_var)
clo_var <- hoistExpr (fsLit "clo") clo DontInline
lclo <- liftPD (Var clo_var)
lclo <- liftPD (Var clo_var)
return (Var clo_var, lclo)
where
arg_tys = map idType args
res_ty = exprType body
is_scalar_ty ty | Just (tycon, []) <- splitTyConApp_maybe ty
= tycon == intTyCon
|| tycon == floatTyCon
|| tycon == doubleTyCon
is_scalar_ty ty
| Just (tycon, []) <- splitTyConApp_maybe ty
= tycon == intTyCon
|| tycon == floatTyCon
|| tycon == doubleTyCon
| otherwise = False