Repository navigation
[REFACTOR][IR] Generalize staging lambdas and explicit thread reductions - #20552
Merged
Merged
Conversation
tlopex
approved these changes
Oct 5, 2026
tqchen
force-pushed
the
ir/staging-lambda-allreduce
branch
from
October 5, 2026 21:03
45bca85 to
98305c0
Compare
Move typed staging lambdas into shared IR with Expr bodies, capture-safe application, derived function signatures, and typed script construction. Represent thread reductions with explicit combiner, identity, value, predicate, destination, and axis operands. Update reduction lowering and datatype legalization to consume that contract and remove reduce_scope transport.
Keep lambda implementation with the shared expression nodes and use structural mutation for scoped substitution. Preserve generic tuple mutation behavior and let datatype legalization apply the combining lambda directly. Consolidate basic lambda usage and retain the ordinary all-reduce signature without a separate operand validator.
Apply staging lambdas through one structural map of parameter references. Derive function signatures at construction and use reflected traversal for ordinary rewrites.
Accept scalar or Tuple reduction operands and combiner results throughout lowering and datatype legalization. Construct lambda parameters from the supplied types and use ordinary callable binding. Apply promoted variables directly and infer promoted primitive result types.
Use one Lambda constructor spelling across shared and TIR script exports, printing, and callers. Keep reduction operand fields distinct from the hardware-axis map used by region-based thread launches.
tqchen
force-pushed
the
ir/staging-lambda-allreduce
branch
from
October 6, 2026 02:09
a5cb470 to
7d79464
Compare
Use the common tile-select predicate check for callable and prebuilt lambdas instead of repeating the boolean return check during callable construction.
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
2 participants
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.
Move typed staging lambdas into shared IR with Expr bodies, parameter substitution through one structural map, function signatures derived at construction, and TVMScript Lambda construction using ordinary Python callable binding.
Represent thread reductions with explicit combiner, identity, value, predicate, destination, and axis operands. Preserve scalar or Tuple operands and combiner results through reduction lowering and datatype legalization, and remove reduce_scope transport.