Allow hoisting invariant factors in rfactor#9190
Draft
alexreinking wants to merge 4 commits into
Draft
Conversation
alexreinking
added
the
release_notes
alexreinking
changed the title
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #9190 +/- ##
=======================================
Coverage ? 70.11%
=======================================
Files ? 256
Lines ? 79298
Branches ? 18992
=======================================
Hits ? 55603
Misses ? 17943
Partials ? 5752 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
alexreinking
marked this pull request as draft
July 6, 2026 14:22
alexreinking
force-pushed
the
alexreinking/rfactor-hoisting
branch
from
July 20, 2026 04:07
9ad9a14 to
9e244ff
Compare
Introduce Stage::hoist_invariants(), which applies the distributive law of a semiring to hoist a loop-invariant factor out of an associative reduction. Called on an update definition, it splits the update into an intermediate that accumulates the factor-free reduction over all of the update's RVars and a write-back that applies the hoisted factor once, returning the intermediate. It handles the +/*, min/+, max/+, or/&&, and and/|| laws, finds factors nested at any depth of an associative chain, and errors when no distributable invariant factor is found. It does not change any types (see change_type()). rfactor() and hoist_invariants() share a private rfactor_impl() helper, and factor extraction sees through Let nodes so hoist_invariants() composes on an rfactor-produced intermediate. Adds general as_binary_operands() and make_binary_op() helpers to IROperator, a Python binding, and correctness tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add Func::change_type(Type, unsafe), which changes the type at which a Func computes and stores its values. It works eagerly at schedule time by splitting the Func in two: a returned intermediate that copies the Func's definitions but accumulates at the new type (inserting casts, preferring integer forms like widening_mul over float round-trips), and the original Func, rewritten in place into an inline wrapper that casts the intermediate's result back to the original type so every existing consumer is unaffected. Composed after hoist_invariants() and rfactor(), this exposes e.g. an Int(32) dot-product accumulator from a Float(32) reduction. Safety is validated with the bounds machinery: for an integer target, change_type() bounds the accumulator by combining the per-term value range (constant_integer_bounds) with the reduction extent. Statically-safe cases pass silently; a case provable only under a runtime precondition (symbolic RDom extent) records that condition, which a new lowering pass (add_type_change_checks, modeled on add_split_factor_checks) injects into the pipeline's assertion block and no_asserts strips. Otherwise change_type() errors unless unsafe=true. change_type() can also be chained. Supporting changes: Function::clear_definition() to redefine a Func in place as the wrapper; FuncSchedule carries the injected type_change_checks; get_associative_identity() for retyped reduction identities; StrictifyFloat treats int<->float casts as strict so change_type() won't strip a user's strict_cast. Adds a Python binding and correctness/change_type.cpp. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add a quantized (int8 x int8 -> f32) mat-vec performance test targeting ARM's SDOT instruction, scheduled by composing hoist_invariants(), rfactor(), and change_type(). hoist_invariants() lifts the invariant scale product out of the reduction, rfactor() splits the scale-free reduction by block, and change_type(Int(32)) retypes the innermost per-block dot product so CodeGen_ARM matches it to SDOT. The generated assembly contains sdot instructions, the Hoisted/PlainRfactor numerical cross-check passes, and the composed schedule is several times faster than the non-hoisting variant. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
alexreinking
force-pushed
the
alexreinking/rfactor-hoisting
branch
from
July 20, 2026 15:28
9e244ff to
328fd0a
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds a mode to
rfactorthat hoists invariant factors according to a detected distributive (semiring) law into the write-back part of the accumulation step. This enables writing straightforward quantized kernels in the algorithm language and factoring them into efficient kernels.On my system (Apple M3 Max, 1 core), the schedule in test/performance/tiled_matmul_arm_neon.cpp outcompetes existing schedules:
The difference in performance is in mostly due to hoisted rfactor's ability to target SDOT instructions.
Breaking changes
No breaking changes. The mode is 100% opt-in.
Just like rfactor, it is possible to write equivalent algorithms that can be scheduled equally well. However, doing so is relatively finicky and more complicated.
Checklist