Skip to content

Dissolve small pure workers into saturated call sites (#211) #285

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Unisay merged 3 commits into from
Jul 20, 2026
Merged

Dissolve small pure workers into saturated call sites (#211) #285

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
364c3a8
feat(optimizer): dissolve small pure workers into saturated call site…
Unisay Jul 16, 2026
5efad1e
test(golden): accept dissolved small pure workers (#211)
Unisay Jul 16, 2026
f8aee38
address copilot review on PR #285
Unisay Jul 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions changelog.d/20260716_150000_unisay_small_pure_worker_inline.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
### Changed

- Small pure workers now dissolve into their saturated call sites (#211),
generalizing the bare-primop unfolding of #281: a worker body that is a
tree of primops, equality tests and negations over parameter references,
scalar literals and cheap projection chains — possibly under a nested
lambda — is pasted at every saturated n-ary call site, bounded by the
small-inline budget. A shared helper such as `add3 x y z = x + y + z` or
`first x _ = x` no longer costs a Lua call per use: sites fold to the
inline expression and constant arguments fold further at compile time
(`add3 1 2 3` emits `6`). Workers whose bodies apply a function, branch,
or exceed the budget keep sharing, as do all value-position uses.
73 changes: 44 additions & 29 deletions lib/Language/PureScript/Backend/IR/Optimizer.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -769,9 +769,10 @@ collapse.
inlineSizeBudget ∷ Natural
inlineSizeBudget = 64

{- | The largest expression the Deref and KnownSize inlining tiers paste
(Note [Complexity and Capture gate inlining]), sized in IR nodes like
'inlineSizeBudget' but far below it: these tiers admit duplication at
{- | The largest expression the Deref and KnownSize inlining tiers and
the cheap-worker unfolding ('isCheapWorkerBody') paste (Note
[Complexity and Capture gate inlining]), sized in IR nodes like
'inlineSizeBudget' but far below it: these admissions duplicate at
every use site, so growth scales with the use count.
-}
smallInlineBudget ∷ Natural
Expand Down Expand Up @@ -1601,15 +1602,16 @@ directive pins the binding as a shared reference at partial sites.

An n-ary call — the direct worker call the uncurry split mints — is not
a unary spine, so 'unwindApp' leaves it whole and the paths above never
see it. Most workers are meant to stay shared bindings, but a worker
whose body is a bare primop over trivial operands (the residue of a
floated dictionary application like @add = Data.Semiring.add
semiringInt@, resolved through the lifted foreign) makes the call itself
the whole cost: pasting is the cheapest possible unfolding and folds
each site to the inline operator (issue #281). The n-ary 'AbsN' root is
pasted under the original 'AppN' node — never rebuilt as a unary spine,
which would leave an under-applied redex ('pasteableRoot') — so the
exact-arity 'betaReduce' consumes it in the same pass.
see it. Most workers are meant to stay shared bindings, but a small
worker whose body is cheap to duplicate ('isCheapWorkerBody' under
'smallInlineBudget' — the residue of a floated dictionary application
like @add = Data.Semiring.add semiringInt@, resolved through the lifted
foreign, or a tiny helper purs shares because its operator occurs more
than once) makes the call itself the whole cost: pasting folds each
site to the inline expression (issues #281, #211). The n-ary 'AbsN'
root is pasted under the original 'AppN' node — never rebuilt as a
unary spine, which would leave an under-applied redex ('pasteableRoot')
— so the exact-arity 'betaReduce' consumes it in the same pass.
-}
inlineSaturatedCall ∷ InlinePolicy → InlineEnv → RewriteRuleM SupplyM Ann
inlineSaturatedCall policy env expr = case expr of
Expand All @@ -1618,7 +1620,8 @@ inlineSaturatedCall policy env expr = case expr of
, Nothing ← directedArity fname
, Just rhs@(AbsN _ params body) ← Map.lookup fname env
, length args == length params
, isBarePrimOpBody body
, expSize rhs < smallInlineBudget
, isCheapWorkerBody body
, countFreeRef fname rhs == 0 →
(\rhs' → Just (AppN ann rhs' args)) <$> freshenBinders rhs
(unwindApp → (Ref _ fname, args))
Expand Down Expand Up @@ -1649,23 +1652,35 @@ inlineSaturatedCall policy env expr = case expr of
directedArity ∷ Qualified Name → Maybe Natural
directedArity fname = refQName fname >>= (`Map.lookup` policyArity policy)

{- | Whether a worker body is a bare primop — a single arithmetic,
comparison, equality or negated-equality node whose operands
'complexityOf' classifies 'Trivial': references (in practice the
worker's parameters) and scalar-sized literals, all free to re-emit.
Such a body is the cheapest possible paste: no work can be duplicated
and no allocation introduced, so 'inlineSaturatedCall' unfolds it at
every saturated n-ary call site regardless of use count (issue #281).
{- | Whether a worker body is cheap enough to duplicate at every
saturated call site: a tree of arithmetic, comparison, equality and
negation nodes over leaves 'complexityOf' classifies at most 'Deref' —
references (in practice the worker's parameters), scalar-sized
literals, and cheap-read chains over them, all free to re-evaluate —
possibly under nested abstractions. Such a body pastes to an inline
operator expression that duplicates no work and allocates nothing
beyond what the worker call already performed (an inner lambda pasted
at the site is exactly the closure the worker's own body allocated per
call), so 'inlineSaturatedCall' unfolds it at every saturated n-ary
call site regardless of use count (issues #281, #211), with code
growth bounded by 'smallInlineBudget' at the call site.

'IfThenElse' is deliberately not admitted: a decision tree pasted into
expression position lowers to a per-call IIFE — the allocation the
case-of-case rules of issue #203 exist to remove — worse than the
shared worker call it would replace. An application may hide arbitrary
work and the allocating shapes (constructors, non-empty literals,
'Let') price above 'Deref', so all of them decline through the
'complexityOf' fallback, conservative for unlisted node kinds by
construction.
-}
isBarePrimOpBody ∷ RawExp ann → Bool
isBarePrimOpBody = \case
PrimBinOp _ _ a b → isTrivialOperand a && isTrivialOperand b
Eq _ a b → isTrivialOperand a && isTrivialOperand b
PrimNot _ e → isBarePrimOpBody e
_ → False
where
isTrivialOperand ∷ RawExp ann → Bool
isTrivialOperand = (== Trivial) . complexityOf
isCheapWorkerBody ∷ RawExp ann → Bool
isCheapWorkerBody = \case
PrimBinOp _ _ a b → isCheapWorkerBody a && isCheapWorkerBody b
Eq _ a b → isCheapWorkerBody a && isCheapWorkerBody b
PrimNot _ e → isCheapWorkerBody e
AbsN _ _params body → isCheapWorkerBody body
leaf → complexityOf leaf <= Deref

-- | Re-apply the arguments 'unwindApp' peeled, as a unary spine.
rebuildSpine ∷ [Exp] → Exp → Exp
Expand Down
146 changes: 146 additions & 0 deletions test/Language/PureScript/Backend/IR/Optimizer/Spec.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ import Language.PureScript.Backend.IR.Types
, ifThenElse
, isLiteral
, lets
, listGrouping
, literalBool
, literalFloat
, literalInt
Expand Down Expand Up @@ -1615,6 +1616,151 @@ spec = describe "IR Optimizer" do
, (Name "main2", primBinOp PrimAdd n2 n1)
]

describe "dissolves small pure workers into saturated call sites (#211)" do
let mainModule = moduleNameFromString "Main"
extern = moduleNameFromString "Extern"
n1 = refImported extern (Name "n1")
n2 = refImported extern (Name "n2")
n3 = refImported extern (Name "n3")
x = Name "x"
y = Name "y"
checked = either (fail . show) pure . optimizedUberModuleChecked
uberWith name def sites =
Linker.UberModule
{ uberModuleForeigns = []
, uberModuleBindings = [Standalone (QName mainModule name, def)]
, uberModuleExports =
[(Name ("main" <> show i), site) | (i, site) ← zip [1 ∷ Int ..] sites]
}
topLevelNames uber =
fst <$> (listGrouping =<< Linker.uberModuleBindings uber)

it "folds a nested primop tree into every saturated call site" do
-- add3 = λx. λy. λz. (x + y) + z: the nested left operand is not
-- Trivial, so the bare-primop rule (#281) declined it and every
-- call paid for the shared worker.
let z = Name "z"
add3Ref = refImported mainModule (Name "add3")
add3Def =
abstraction (paramNamed x) $
abstraction (paramNamed y) $
abstraction (paramNamed z) $
primBinOp
PrimAdd
(primBinOp PrimAdd (refLocal x) (refLocal y))
(refLocal z)
saturated a b = application (application (application add3Ref a) b)
optimized ←
checked $
uberWith (Name "add3") add3Def [saturated n1 n2 n3, saturated n3 n2 n1]
Linker.uberModuleBindings optimized `shouldBe` []
Linker.uberModuleExports optimized
`shouldBe` [ (Name "main1", primBinOp PrimAdd (primBinOp PrimAdd n1 n2) n3)
, (Name "main2", primBinOp PrimAdd (primBinOp PrimAdd n3 n2) n1)
]

it "folds a parameter-selecting worker (no operator at all)" do
-- first = λx. λy. x: the body is a bare reference — not a primop,
-- so the #281 rule never saw it.
let firstRef = refImported mainModule (Name "first")
firstDef =
abstraction (paramNamed x) $
abstraction (paramNamed y) $
refLocal x
saturated a = application (application firstRef a)
optimized ←
checked $
uberWith (Name "first") firstDef [saturated n1 n2, saturated n2 n1]
Linker.uberModuleBindings optimized `shouldBe` []
Linker.uberModuleExports optimized
`shouldBe` [(Name "main1", n1), (Name "main2", n2)]

it "admits projection-chain operands (Deref leaves)" do
-- addFields = λx. λy. x.foo + y.bar: record and module tables are
-- write-once, so re-reading a field at the paste site duplicates
-- no work.
let addFieldsRef = refImported mainModule (Name "addFields")
addFieldsDef =
abstraction (paramNamed x) $
abstraction (paramNamed y) $
primBinOp
PrimAdd
(objectProp (refLocal x) (PropName "foo"))
(objectProp (refLocal y) (PropName "bar"))
saturated a = application (application addFieldsRef a)
optimized ←
checked $
uberWith
(Name "addFields")
addFieldsDef
[saturated n1 n2, saturated n2 n1]
Linker.uberModuleBindings optimized `shouldBe` []
Linker.uberModuleExports optimized
`shouldBe` [
( Name "main1"
, primBinOp
PrimAdd
(objectProp n1 (PropName "foo"))
(objectProp n2 (PropName "bar"))
)
,
( Name "main2"
, primBinOp
PrimAdd
(objectProp n2 (PropName "foo"))
(objectProp n1 (PropName "bar"))
)
]

it "keeps a worker whose body applies a function" do
-- callAdd = λx. λy. g (x + y): an application may hide arbitrary
-- work, so the call sites keep sharing the worker.
let g = refImported extern (Name "g")
callAddRef = refImported mainModule (Name "callAdd")
callAddDef =
abstraction (paramNamed x) $
abstraction (paramNamed y) $
application
g
(primBinOp PrimAdd (refLocal x) (refLocal y))
saturated a = application (application callAddRef a)
optimized ←
checked $
uberWith (Name "callAdd") callAddDef [saturated n1 n2, saturated n2 n1]
topLevelNames optimized `shouldBe` [QName mainModule (Name "callAdd$w")]

it "keeps a worker whose body branches" do
-- pickOr = λx. λy. if x then y else 0: an IfThenElse pasted into
-- expression position lowers to a per-call IIFE (issue #203) —
-- worse than the shared worker call it would replace.
let pickOrRef = refImported mainModule (Name "pickOr")
pickOrDef =
abstraction (paramNamed x) $
abstraction (paramNamed y) $
ifThenElse (refLocal x) (refLocal y) (literalInt 0)
saturated a = application (application pickOrRef a)
optimized ←
checked $
uberWith (Name "pickOr") pickOrDef [saturated n1 n2, saturated n2 n1]
topLevelNames optimized `shouldBe` [QName mainModule (Name "pickOr$w")]

it "keeps a worker whose primop tree exceeds the small budget" do
-- Same shape as add3, but grown past 'smallInlineBudget': pasting
-- it at every site trades one call for unbounded code growth.
let wideRef = refImported mainModule (Name "wide")
wideDef =
abstraction (paramNamed x) $
abstraction (paramNamed y) $
foldl'
(primBinOp PrimAdd)
(refLocal x)
(concat (replicate 8 [refLocal y, refLocal x]))
saturated a = application (application wideRef a)
optimized ←
checked $
uberWith (Name "wide") wideDef [saturated n1 n2, saturated n2 n1]
topLevelNames optimized `shouldBe` [QName mainModule (Name "wide$w")]

describe "honours @inline arity=N directives (issue #232)" do
let mainModule = moduleNameFromString "Main"
extern = moduleNameFromString "Extern"
Expand Down
33 changes: 4 additions & 29 deletions test/ps/output/Golden.Unbinding.Test/golden.ir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,37 +1,12 @@
UberModule
{ uberModuleBindings =
[ Standalone
( QName
{ qnameModuleName = ModuleName "Golden.Unbinding.Test", qnameName = Name "f$w"
}, AbsN Nothing
( ParamUnused Nothing :| [ ParamUnused Nothing ] )
( LiteralInt Nothing 3 )
)
], uberModuleForeigns = [], uberModuleExports =
{ uberModuleBindings = [], uberModuleForeigns = [], uberModuleExports =
[
( Name "a", LiteralInt Nothing 1 ),
( Name "b", LiteralInt Nothing 2 ),
( Name "f", AbsN Nothing
( ParamNamed Nothing ( Name "f$p1$0" ) :| [] )
( AbsN Nothing
( ParamNamed Nothing ( Name "f$p2$1" ) :| [] )
( AppN Nothing
( Ref Nothing ( Imported ( ModuleName "Golden.Unbinding.Test" ) ( Name "f$w" ) ) )
( Ref Nothing
( Local ( Name "f$p1$0" ) ) :|
[ Ref Nothing ( Local ( Name "f$p2$1" ) ) ]
)
)
)
( ParamUnused Nothing :| [] )
( AbsN Nothing ( ParamUnused Nothing :| [] ) ( LiteralInt Nothing 3 ) )
),
( Name "c", AppN Nothing
( Ref Nothing ( Imported ( ModuleName "Golden.Unbinding.Test" ) ( Name "f$w" ) ) )
( LiteralInt Nothing 1 :|
[ AppN Nothing
( Ref Nothing ( Imported ( ModuleName "Golden.Unbinding.Test" ) ( Name "f$w" ) ) )
( LiteralInt Nothing 2 :| [ LiteralInt Nothing 1 ] )
]
)
)
( Name "c", LiteralInt Nothing 3 )
]
}
9 changes: 2 additions & 7 deletions test/ps/output/Golden.Unbinding.Test/golden.lua
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,6 @@
local Golden_Unbinding_Test_f_S_w = function() return 3 end
return {
a = 1,
b = 2,
f = function(f_S_p1_S_0)
return function(f_S_p2_S_1)
return Golden_Unbinding_Test_f_S_w(f_S_p1_S_0, f_S_p2_S_1)
end
end,
c = Golden_Unbinding_Test_f_S_w(1, Golden_Unbinding_Test_f_S_w(2, 1))
f = function() return function() return 3 end end,
c = 3
}
35 changes: 6 additions & 29 deletions test/ps/output/Golden.UncurriedLift.Test/golden.ir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,11 +106,9 @@ UberModule
{ qnameModuleName = ModuleName "Golden.UncurriedLift.Test", qnameName = Name "addOnePlusTwoTo"
}, AbsN Nothing
( ParamNamed Nothing ( Name "c$682" ) :| [] )
( AppN Nothing
( Ref Nothing ( Imported ( ModuleName "Golden.UncurriedLift.Test" ) ( Name "add3" ) ) )
( LiteralInt Nothing 1 :|
[ LiteralInt Nothing 2, Ref Nothing ( Local ( Name "c$682" ) ) ]
)
( PrimBinOp Nothing PrimAdd
( LiteralInt Nothing 3 )
( Ref Nothing ( Local ( Name "c$682" ) ) )
)
)
], uberModuleForeigns = [], uberModuleExports =
Expand Down Expand Up @@ -155,14 +153,7 @@ UberModule
( Ref Nothing ( Imported ( ModuleName "Effect.Console" ) ( Name "log" ) ) )
( AppN Nothing
( Ref Nothing ( Imported ( ModuleName "Data.Show" ) ( Name "showIntImpl" ) ) )
( AppN Nothing
( Ref Nothing
( Imported ( ModuleName "Golden.UncurriedLift.Test" ) ( Name "add3" ) )
)
( LiteralInt Nothing 1 :|
[ LiteralInt Nothing 2, LiteralInt Nothing 3 ]
) :| []
) :| []
( LiteralInt Nothing 6 :| [] ) :| []
)
)
( Ref Nothing ( Imported ( ModuleName "Prim" ) ( Name "$magicDoRun" ) ) :| [] )
Expand Down Expand Up @@ -192,14 +183,7 @@ UberModule
( Ref Nothing ( Imported ( ModuleName "Effect.Console" ) ( Name "log" ) ) )
( AppN Nothing
( Ref Nothing ( Imported ( ModuleName "Data.Show" ) ( Name "showIntImpl" ) ) )
( AppN Nothing
( Ref Nothing
( Imported ( ModuleName "Golden.UncurriedLift.Test" ) ( Name "add3" ) )
)
( LiteralInt Nothing 1 :|
[ LiteralInt Nothing 2, LiteralInt Nothing 100 ]
) :| []
) :| []
( LiteralInt Nothing 103 :| [] ) :| []
)
)
( Ref Nothing ( Imported ( ModuleName "Prim" ) ( Name "$magicDoRun" ) ) :| [] )
Expand All @@ -209,14 +193,7 @@ UberModule
( Ref Nothing ( Imported ( ModuleName "Effect.Console" ) ( Name "log" ) ) )
( AppN Nothing
( Ref Nothing ( Imported ( ModuleName "Data.Show" ) ( Name "showIntImpl" ) ) )
( AppN Nothing
( Ref Nothing
( Imported ( ModuleName "Golden.UncurriedLift.Test" ) ( Name "add3" ) )
)
( LiteralInt Nothing 1 :|
[ LiteralInt Nothing 2, LiteralInt Nothing 200 ]
) :| []
) :| []
( LiteralInt Nothing 203 :| [] ) :| []
)
)
( Ref Nothing ( Imported ( ModuleName "Prim" ) ( Name "$magicDoRun" ) ) :| [] )
Expand Down
Loading