From 8418a65f67a42d7c72547a152208580c0bbab6de Mon Sep 17 00:00:00 2001 From: Marty Stumpf Date: Mon, 27 Mar 2023 11:59:55 -0700 Subject: [PATCH] Address some comments. --- .../Transform/Inline/CallSiteInline.hs | 79 ++++++-------- .../src/PlutusIR/Transform/Inline/Inline.hs | 12 +- .../src/PlutusIR/Transform/Inline/Utils.hs | 103 +++++++++--------- 3 files changed, 93 insertions(+), 101 deletions(-) diff --git a/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/CallSiteInline.hs b/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/CallSiteInline.hs index 526a712d8b4..8c3b75eefee 100644 --- a/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/CallSiteInline.hs +++ b/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/CallSiteInline.hs @@ -12,9 +12,7 @@ See note [Inlining of fully applied functions]. module PlutusIR.Transform.Inline.CallSiteInline where -import Control.Lens (forMOf) import Control.Monad.State -import Data.Bifunctor (first) import PlutusIR.Core import PlutusIR.Transform.Inline.Utils @@ -26,7 +24,7 @@ each in detail below. (1) What do we mean by "fully applied"? A function is fully applied when it has been applied to all arguments as indicated by its -syntactic arity. +"syntactic arity": Consider `let v = rhs in body`, in which `body` calls `v`. @@ -39,16 +37,16 @@ term expects before it may do some work. Since we have both type and lambda abstractions, this is not a simple argument count, but rather a list of values indicating whether the next argument should be a term or a type. -Note that this is the syntactic arity, i.e. it just corresponds to the number of +Note that this is the *syntactic* arity, i.e. it just corresponds to the number of syntactic lambda and type abstractions on the outside of the term. It is thus an under-approximation of how many arguments the term may need. e.g. consider the term @let id = \x -> x in id@: the variable @id@ has syntactic arity @[]@, but does in fact need an argument before it does any work. -In the `body`, where `v` is *called*, if `VBody` is acceptable in size and cost, and +In the `body`, where `v` is *called*, if it was given the `n` term or type arguments in the correct order, then it is *fully applied*. -We inline the call of the fully applied `v` in this case, i.e., -we replace `v` in the `body` with `rhs`. E.g. +If `VBody` is acceptable in size and cost, we inline the call site of the fully applied `v` in this +case, i.e., we replace `v` in the `body` with `rhs`. E.g. let f = \x.\y -> x in @@ -72,17 +70,16 @@ in This is already a reduction in code size. However, because of this, our dead code elimination pass is able to further reduce the code to just `a`. -To find out whether a function is fully applied, -we first need to count the number of type/term lambda abstractions, -so we know how many term/type arguments they have. +Because a function can be called in the `body` multiple times and may not be fully applied for all +its calls, we cannot simply keep a substitution map like in `Inline`, +which substitute *all* occurrences of a variable. We store the function in a map, +`Utils.InScopeSet`, with all information needed for inlining saturated functions. -We pattern match the _rhs_ with `LamAbs` and `TyAbs` (lambda abstraction for terms or types), -and track the sequence of lambdas. -Then, in the _body_, we track the term and type application (`Apply` or `TyInst`) of _v_. +To find out whether a function is fully applied, +we first find the arity of a variable and put it in the `Utils.InScopeSet` map. -If _v_ is fully applied in the body, i.e., if the sequence of type and term lambda abstractions is -exactly matched by the corresponding sequence of type and term applications, then -we inline _v_, i.e., replace its occurrence with _rhs_ in the _body_. +Then, in the _body_, we track the term and type application (`Apply` or `TyInst`) of the variable +and inline right there if it is fully applied. Note that over-application of a function would also be inlined. An example of over-application: @@ -95,11 +92,6 @@ let id = \y -> y `f`'s arity is 1, but is applied to 2 arguments. We inline `f` in this case if its cost and size are acceptable. -Because a function can be called in the `body` multiple times and may not be fully applied for all -its calls, we cannot simply keep a substitution map like in `Inline`, -which substitute *all* occurrences of a variable. We store the function in a map -`Utils.CalledVarEnv` with all information needed for inlining saturated functions. - (2) How do we decide whether cost and size are acceptable? We currently reuse the heuristics 'Utils.sizeIsAcceptable' and 'Utils.costIsAcceptable' @@ -108,10 +100,10 @@ that are used in unconditional inlining. For let v = \x1.\x2...\xn.VBody in body we check `VBody` with the above "acceptable" functions. -Note that all `LamAbs` and `TyAbs` are not acceptable in terms of size, but they should have been -counted out already so we should not immediately encounter those in `VBody`. -Also, we currently reject `Constant` (has acceptable cost but not acceptable size). -We may want to check their sizes instead of just rejecting them. +The "acceptable" functions are currently quite conservative, e.g., +we currently reject `Constant` (has acceptable cost but not acceptable size). +We will work on more refined checks (e.g., checking their sizes instead of just rejecting them) to +improve the optimization. -} -- | Computes the 'Utils.Arity' of a term. Also returns the function body, for checking whether @@ -120,23 +112,25 @@ computeArity :: Term tyname name uni fun ann -> (Arity, Term tyname name uni fun ann) computeArity = \case - LamAbs _ _ _ body -> (first ((:) MkTerm) (computeArity body)) - TyAbs _ _ _ body -> (first ((:) MkType) (computeArity body)) + LamAbs _ _ _ body -> + let (nextArgs, nextBody) = computeArity body in (TermParam:nextArgs, nextBody) + TyAbs _ _ _ body -> + let (nextArgs, nextBody) = computeArity body in (TypeParam:nextArgs, nextBody) -- Whenever we encounter a body that is not a lambda or type abstraction, we are done counting tm -> ([],tm) -- | A term or type argument. -data Args tyname name uni fun ann = +data Arg tyname name uni fun ann = MkTermArg (Term tyname name uni fun ann) | MkTypeArg (Type tyname uni ann) -- | A list of type or term argument(s) that are being applied. -type ArgOrder tyname name uni fun ann = [Args tyname name uni fun ann] +type ArgOrder tyname name uni fun ann = [Arg tyname name uni fun ann] -- | A pair of argument and the annotation of the term being applied to, -- so a term that was traversed can be built back with `mkApps`. type ArgOrderWithAnn tyname name uni fun ann = - [(Args tyname name uni fun ann, ann)] + [(Arg tyname name uni fun ann, ann)] -- | Takes a term or type application expression and returns the function -- being applied and the arguments to which it is applied @@ -166,27 +160,26 @@ isFullyApplied [] [] = True isFullyApplied lamOrder argsOrder = -- start comparing from the end because there may be over-application case (last lamOrder, last argsOrder) of - (MkTerm, MkTermArg _) -> isFullyApplied (init lamOrder) (init argsOrder) - (MkType, MkTypeArg _) -> isFullyApplied (init lamOrder) (init argsOrder) - _ -> False + (TermParam, MkTermArg _) -> isFullyApplied (init lamOrder) (init argsOrder) + (TypeParam, MkTypeArg _) -> isFullyApplied (init lamOrder) (init argsOrder) + _ -> False -- | Inline fully applied functions iff the body of the function is `acceptable`. inlineSat :: forall tyname name uni fun ann. InliningConstraints tyname name uni fun => Term tyname name uni fun ann -- ^ The `body` of the `Let` term. -> InlineM tyname name uni fun ann (Term tyname name uni fun ann) -inlineSat appOrTyInstTm = - case appOrTyInstTm of +inlineSat t = + case t of -- If the term is a term or type application, check it's applying to a var that we may inline - Apply _varAnn _fun _arg -> go appOrTyInstTm - TyInst _varAnn _fun _arg -> go appOrTyInstTm + Apply _varAnn _fun _arg -> go t + TyInst _varAnn _fun _arg -> go t -- otherwise, check all subterms - _ -> forMOf termSubterms appOrTyInstTm inlineSat + _ -> pure t where go tm = do -- collect all the arguments of the term being applied to - let argsAppliedTo = fst $ collectArgs tm - args = snd $ collectArgs tm - case argsAppliedTo of + let (fun, args) = collectArgs tm + case fun of -- if it is a `Var` that is being applied to, check to see if it's fully applied Var _ann name -> do maybeVarInfo <- gets (lookupCalled name) @@ -195,11 +188,11 @@ inlineSat appOrTyInstTm = -- and it won't be in the map. ATM recursive bindings aren't inlined. Nothing -> pure tm Just varInfo -> do - isAcceptable <- acceptable (calledVarBody varInfo) + let isAcceptable = acceptable (varBody varInfo) if isAcceptable && isFullyApplied (arity varInfo) (map fst args) then do -- if the body of `Var` is `acceptable` and -- it is fully applied (over-application is allowed) then inline it - pure $ mkApps (calledVarDef varInfo) args + pure $ mkApps (varDef varInfo) args else pure tm -- if the term being applied is not a `Var`, don't inline _ -> pure tm diff --git a/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/Inline.hs b/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/Inline.hs index 804c80f8252..791fd104272 100644 --- a/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/Inline.hs +++ b/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/Inline.hs @@ -244,27 +244,27 @@ processSingleBinding -> InlineM tyname name uni fun ann (Maybe (Binding tyname name uni fun ann)) processSingleBinding body = -- when we encounter a binding that is a function, we add it to the global map - -- `Utils.CalledVarEnv`. When we check the body of the let binding we look in this map for + -- `Utils.InScopeSet`. When we check the body of the let binding we look in this map for -- call site inlining. let addVarToMap letBody v ann s n rhs = do -- we want to do unconditional inline if possible maybeRhs' <- maybeAddSubst letBody ann s n rhs case maybeRhs' of -- inline and remove binding when the conditions for unconditional inlining are met. - Nothing -> pure $ TermBind ann s v <$> maybeRhs' + Nothing -> pure Nothing Just rhsProcess -> do let varLamOrder = fst $ computeArity rhs bodyToCheck = snd $ computeArity rhs - -- add the function to `CalledVarEnv`, because we may want to inline this. + -- add the function to `InScopeSet`, because we may want to inline this. -- We don't remove the binding because we decide *at the call site* -- whether we want to inline, and it may be called more than once. - void $ modify' $ extendCalled n (MkCalledVarInfo rhs varLamOrder bodyToCheck) + void $ modify' $ extendCalled n (MkVarInfo rhs varLamOrder bodyToCheck) pure $ Just $ TermBind ann s v rhsProcess in \case -- The let binding is a function type here. - -- If it's not unconditionally inlined, we add it to the `CalledVarEnv` + -- If it's not unconditionally inlined, we add it to the `InScopeSet` -- and consider whether we want to inline at the call site. TermBind ann s v@(VarDecl _ n (TyFun _ _tyArg _tyBody)) rhs -> addVarToMap body v ann s n rhs -- The let binding is a type lambda abstraction. @@ -315,7 +315,7 @@ maybeAddSubst body ann s n rhs = do -- See Note [Inlining approach and 'Secrets of the GHC Inliner'] and [Inlining and -- purity]. This is the case where we don't know that the number of occurrences is -- exactly one, so there's no point checking if the term is immediately evaluated. - postUnconditional <- fmap ((&&) termIsPure) (acceptable rhs') + let postUnconditional = (&&) termIsPure (acceptable rhs') if postUnconditional then extendAndDrop (Done $ dupable rhs') else pure $ Just rhs' diff --git a/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/Utils.hs b/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/Utils.hs index 55449700bd5..f9f76795598 100644 --- a/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/Utils.hs +++ b/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/Utils.hs @@ -40,21 +40,21 @@ newtype InlineTerm tyname name uni fun ann = -- | Term substitution, 'Subst' in the paper. -- A map of unprocessed variable and its substitution range. -newtype TermEnv tyname name uni fun ann = - TermEnv { _unTermEnv :: UniqueMap TermUnique (InlineTerm tyname name uni fun ann) } +newtype TermSubst tyname name uni fun ann = + TermSubst { _unTermSubst :: UniqueMap TermUnique (InlineTerm tyname name uni fun ann) } deriving newtype (Semigroup, Monoid) --- | Type substitution, similar to `TermEnv` but for types. +-- | Type substitution, similar to `TermSubst` but for types. -- A map of unprocessed type variable and its substitution range. -newtype TypeEnv tyname uni ann = - TypeEnv { _unTypeEnv :: UniqueMap TypeUnique (Dupable (Type tyname uni ann)) } +newtype TypeSubst tyname uni ann = + TypeSubst { _unTypeSubst :: UniqueMap TypeUnique (Dupable (Type tyname uni ann)) } deriving newtype (Semigroup, Monoid) -- For call site inlining: -- | A mapping including all let-bindings that are functions. -newtype CalledVarEnv tyname name uni fun ann = - CalledVarEnv { _unCalledVarEnv :: UniqueMap TermUnique (CalledVarInfo tyname name uni fun ann)} +newtype InScopeSet tyname name uni fun ann = + InScopeSet { _unCalledVarEnv :: UniqueMap TermUnique (VarInfo tyname name uni fun ann)} deriving newtype (Semigroup, Monoid) {-| @@ -69,41 +69,40 @@ an under-approximation of how many arguments the term may need. e.g. consider the term @let id = \x -> x in id@: the variable @id@ has syntactic arity @[]@, but does in fact need an argument before it does any work. -} -type Arity = [TermOrType] - --- | An ordered list of the term and type arguments applied to a function. -type AppOrder = [TermOrType] +type Arity = [ParamKind] -- | Info attached to a let-binding needed for call site inlining. -data CalledVarInfo tyname name uni fun ann = - MkCalledVarInfo { - calledVarDef :: Term tyname name uni fun ann -- ^ its definition - , arity :: Arity -- ^ its sequence of term and type lambdas - , calledVarBody :: Term tyname name uni fun ann +data VarInfo tyname name uni fun ann = + MkVarInfo { + varDef :: Term tyname name uni fun ann -- ^ its definition + , arity :: Arity -- ^ its arity, storing to avoid repeated calculations. + , varBody :: Term tyname name uni fun ann -- ^ the body of the function, for checking `acceptable` or not } -- | Is the next argument a term or a type? -data TermOrType = - MkTerm | MkType +data ParamKind = + TermParam | TypeParam deriving stock (Eq, Show) -instance Pretty TermOrType where +instance Pretty ParamKind where pretty = viaShow --- | Substitution for both terms and types. --- Similar to 'Subst' in the paper, but the paper only has terms. -data Subst tyname name uni fun ann = - Subst { _termEnv :: TermEnv tyname name uni fun ann - , _typeEnv :: TypeEnv tyname uni ann - , _calledVarEnv :: CalledVarEnv tyname name uni fun ann +-- | Inliner context for both unconditional inlining and call site inlining. +-- It includes substitution for both terms and types, which is similar to 'Subst' in the paper. +-- It also includes the in-scope set for call site inlining. +data InlinerContext tyname name uni fun ann = + InlinerContext { _termSubst :: TermSubst tyname name uni fun ann + , _typeSubst :: TypeSubst tyname uni ann + , _inScopeSet :: InScopeSet tyname name uni fun ann } deriving stock (Generic) - deriving (Semigroup, Monoid) via (GenericSemigroupMonoid (Subst tyname name uni fun ann)) + deriving (Semigroup, Monoid) via + (GenericSemigroupMonoid (InlinerContext tyname name uni fun ann)) -makeLenses ''TermEnv -makeLenses ''TypeEnv -makeLenses ''CalledVarEnv -makeLenses ''Subst +makeLenses ''TermSubst +makeLenses ''TypeSubst +makeLenses ''InScopeSet +makeLenses ''InlinerContext -- Helper functions: @@ -111,56 +110,56 @@ makeLenses ''Subst lookupTerm :: (HasUnique name TermUnique) => name -- ^ The name of the variable. - -> Subst tyname name uni fun ann -- ^ The substitution. + -> InlinerContext tyname name uni fun ann -- ^ The substitution. -> Maybe (InlineTerm tyname name uni fun ann) -lookupTerm n subst = lookupName n $ subst ^. termEnv . unTermEnv +lookupTerm n subst = lookupName n $ subst ^. termSubst . unTermSubst -- | Insert the unprocessed variable into the substitution. extendTerm :: (HasUnique name TermUnique) => name -- ^ The name of the variable. -> InlineTerm tyname name uni fun ann -- ^ The substitution range. - -> Subst tyname name uni fun ann -- ^ The substitution. - -> Subst tyname name uni fun ann -extendTerm n clos subst = subst & termEnv . unTermEnv %~ insertByName n clos + -> InlinerContext tyname name uni fun ann -- ^ The substitution. + -> InlinerContext tyname name uni fun ann +extendTerm n clos subst = subst & termSubst . unTermSubst %~ insertByName n clos -- | Look up the unprocessed type variable in the substitution. lookupType :: (HasUnique tyname TypeUnique) => tyname - -> Subst tyname name uni fun ann + -> InlinerContext tyname name uni fun ann -> Maybe (Dupable (Type tyname uni ann)) -lookupType tn subst = lookupName tn $ subst ^. typeEnv . unTypeEnv +lookupType tn subst = lookupName tn $ subst ^. typeSubst . unTypeSubst -- | Check if the type substitution is empty. -isTypeSubstEmpty :: Subst tyname name uni fun ann -> Bool -isTypeSubstEmpty (Subst _ (TypeEnv tyEnv) _) = isEmpty tyEnv +isTypeSubstEmpty :: InlinerContext tyname name uni fun ann -> Bool +isTypeSubstEmpty (InlinerContext _ (TypeSubst tyEnv) _) = isEmpty tyEnv -- | Insert the unprocessed type variable into the substitution. extendType :: (HasUnique tyname TypeUnique) => tyname -- ^ The name of the type variable. -> Type tyname uni ann -- ^ Its type. - -> Subst tyname name uni fun ann -- ^ The substitution. - -> Subst tyname name uni fun ann -extendType tn ty subst = subst & typeEnv . unTypeEnv %~ insertByName tn (dupable ty) + -> InlinerContext tyname name uni fun ann -- ^ The substitution. + -> InlinerContext tyname name uni fun ann +extendType tn ty subst = subst & typeSubst . unTypeSubst %~ insertByName tn (dupable ty) -- | Look up the called variable in the substitution. lookupCalled :: (HasUnique name TermUnique) => name -- ^ The name of the variable. - -> Subst tyname name uni fun ann -- ^ The substitution. - -> Maybe (CalledVarInfo tyname name uni fun ann) -lookupCalled n subst = lookupName n $ subst ^. calledVarEnv . unCalledVarEnv + -> InlinerContext tyname name uni fun ann -- ^ The substitution. + -> Maybe (VarInfo tyname name uni fun ann) +lookupCalled n subst = lookupName n $ subst ^. inScopeSet . unCalledVarEnv -- | Insert the called variable into the substitution. extendCalled :: (HasUnique name TermUnique) => name -- ^ The name of the variable. - -> CalledVarInfo tyname name uni fun ann -- ^ The called variable's info. - -> Subst tyname name uni fun ann -- ^ The substitution. - -> Subst tyname name uni fun ann -extendCalled n info subst = subst & calledVarEnv . unCalledVarEnv %~ insertByName n info + -> VarInfo tyname name uni fun ann -- ^ The called variable's info. + -> InlinerContext tyname name uni fun ann -- ^ The substitution. + -> InlinerContext tyname name uni fun ann +extendCalled n info subst = subst & inScopeSet . unCalledVarEnv %~ insertByName n info -- General infra: @@ -195,7 +194,7 @@ makeLenses ''InlineInfo -- (determined from profiling) -- | The monad the inliner runs in. type InlineM tyname name uni fun ann = - ReaderT (InlineInfo name fun ann) (StateT (Subst tyname name uni fun ann) Quote) + ReaderT (InlineInfo name fun ann) (StateT (InlinerContext tyname name uni fun ann) Quote) -- Heuristics: @@ -251,10 +250,10 @@ effectSafe body s n purity = do -- | Should we inline? Should only inline things that won't duplicate work or code. -- See Note [Inlining approach and 'Secrets of the GHC Inliner'] -acceptable :: Term tyname name uni fun ann -> InlineM tyname name uni fun ann Bool +acceptable :: Term tyname name uni fun ann -> Bool acceptable t = -- See Note [Inlining criteria] - pure $ costIsAcceptable t && sizeIsAcceptable t + costIsAcceptable t && sizeIsAcceptable t {- Note [Inlining criteria] What gets inlined? Our goals are simple: