Skip to content

Add eager_inline() directive for schedule-time inlining - #9256

Open
alexreinking wants to merge 3 commits into
mainfrom
alexreinking/eager-inline
Open

Add eager_inline() directive for schedule-time inlining#9256
alexreinking wants to merge 3 commits into
mainfrom
alexreinking/eager-inline

Conversation

@alexreinking

@alexreinking alexreinking commented Jul 29, 2026

Copy link
Copy Markdown
Member

This PR adds a scheduling directive, g.eager_inline({f1, ..., fN}) that immediately and destructively inlines the bodies of f1, ..., fN into g. This happens in sequence so that direct calls to f_j by f_i, i < j can be inlined without needing to issue the directive multiple times.

I'm placing this on the bottom of the stack because it is the easiest to review.

Breaking changes

None. This is a new directive.

Checklist

  • Tests added or updated (not required for docs, CI config, or typo fixes)
  • Documentation updated (if public API changed)
  • Python bindings updated (if public API changed)
  • Commits include AI attribution where applicable (see Code of Conduct)

Stack created with GitHub Stacks CLIGive Feedback 💬

@alexreinking
alexreinking force-pushed the alexreinking/eager-inline branch from c226ea0 to de722b4 Compare July 29, 2026 19:33
@alexreinking
alexreinking marked this pull request as ready for review July 29, 2026 19:33
@codecov

codecov Bot commented Jul 30, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 64.70588% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.24%. Comparing base (ceea694) to head (fc9374c).
⚠️ Report is 4 commits behind head on main.

Files with missing lines Patch % Lines
src/Func.cpp 57.14% 4 Missing and 2 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #9256      +/-   ##
==========================================
- Coverage   70.26%   70.24%   -0.03%     
==========================================
  Files         257      257              
  Lines       79106    79122      +16     
  Branches    18954    18957       +3     
==========================================
- Hits        55583    55578       -5     
- Misses      17885    17895      +10     
- Partials     5638     5649      +11     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@mcourteaux

Copy link
Copy Markdown
Contributor

I'd like to see an actual example (code or description) of when this is useful. I don't really grasp what it means when you say:

This is useful to surface structure that other schedule-time directives need to see.

What's now possible that otherwise is impossible?

@alexreinking

Copy link
Copy Markdown
Member Author

What's now possible that otherwise is impossible?

That's in #9190 (top of the stack) — the idea is that if you have a function that looks like f(x) += g(x, r) * h(x, r) then you need to be able to look inside the definitions of g and h to find invariant factors. Say g(x, r) looks like sc_g(x) * Wg(x, r) and h has a similar form. You'd like to be able to pull the sc_g(x) * sc_h(x) factor into a write-back, but it's just not visible without inlining at a point in time that hoist_invariants can see it.

@alexreinking
alexreinking force-pushed the alexreinking/eager-inline branch 3 times, most recently from dd530c1 to 2f709b3 Compare July 31, 2026 21:30
@abadams

abadams commented Aug 3, 2026

Copy link
Copy Markdown
Member

Let's discuss in person how this interacts with .in()/.clone_in() and how it would interact with proposed changes to .in()/.clone_in(). I think it's fine but it's worth working through.

@alexreinking
alexreinking force-pushed the alexreinking/eager-inline branch 2 times, most recently from a7a483c to cbfbb1c Compare August 4, 2026 21:37
@alexreinking

Copy link
Copy Markdown
Member Author

Let's discuss in person how this interacts with .in()/.clone_in() and how it would interact with proposed changes to .in()/.clone_in(). I think it's fine but it's worth working through.

I followed up on the potential for strange interactions with .in() and .clone_in(). In short, it's fine especially since we're planning to rework .in() and .clone_in() as eager DAG rewrites soon, anyway.

If someone writes f.in().eager_inline(f), assuming f is inlinable (no update definitions), they get a normal/correct call to f, but then lose the schedule that was on f. It's a weird thing to do, but it's not pathological.

Writing f.clone_in().eager_inline(f) is a no-op because there are no calls to f.

Still, it might be worth warning or erroring if you .eager_inline(f) when there are no direct calls to f. Thoughts?

alexreinking and others added 3 commits August 4, 2026 18:13
Add Func::eager_inline({f1..fN}), which inlines direct calls to each given
Func into this Func's definitions immediately, at schedule time, processed
left to right so that inlining an earlier Func exposes direct calls to a later
one. Unlike compute_inline(), which only marks a Func to be inlined during
lowering, eager_inline() rewrites the caller's definitions in place, surfacing
structure that other schedule-time directives can then act on -- e.g. exposing
an invariant factor buried in a call so hoist_invariants() can hoist it from
h(x) += f(x) * g(x).

Built on Internal::inline_function(Function, Function). Removes the redundant
validate_schedule_inlined_function() call from the Inliner constructor: it is
also (and properly) called at lowering time in ScheduleFunctions with loop
levels locked, and calling it in the constructor inspects unlocked loop
levels, which breaks schedule-time inlining. Adds a Python binding and
correctness/eager_inline.cpp.

This Func predates hoist_invariants() in commit order, so
eager_inline.cpp doesn't yet cover the composed case; a TODO marks it.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
eager_inline() previously lived only on Func and inlined into every one of the
Func's definitions at once. But the structure it surfaces is stage-scoped -- the
schedule-time directives that consume it (e.g. rfactor()) operate on a single
Stage -- so inlining into all definitions is both too coarse and conceptually
mismatched.

Make Stage::eager_inline() the primary API: it rewrites only that stage's
definition in place, via a new inline_function(Definition &, Function) overload.
Func::eager_inline() is kept as a shorthand targeting the initial (pure)
definition, mirroring how Func::vectorize() and friends delegate to the initial
Stage; inline into an update with f.update(n).eager_inline(...).

Add stage-scoping coverage to correctness/eager_inline (inlining into one stage
leaves the others untouched) and an error test for the undefined-Func assert.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@alexreinking
alexreinking force-pushed the alexreinking/eager-inline branch from cbfbb1c to fc9374c Compare August 4, 2026 22:13
@alexreinking alexreinking added the release_notes For changes that may warrant a note in README for official releases. label Aug 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

release_notes For changes that may warrant a note in README for official releases.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants