[TIR][Schedule]Generalize fuseReductionEpilogue to support arbitrary epilogue expressions

[kimm240]

Major Changes for Generalization

1. Pattern Matching Removal

Removed Items:

• EpilogueType enum (Bias, BiasReLU, Clipping)
• AnalyzeEpiloguePattern() function
• Pattern-specific branching logic

Current Approach:

• Directly process the entire epilogue expression without pattern matching

2. Store Entire Epilogue Expression

• Store the entire epilogue expression in epilogue_expression_
• Use the expression directly without pattern analysis

// Store the epilogue expression and reduction buffer load
epilogue_expression_ = inlined_store_->value;
reduction_buffer_load_ = loads[0];

3. Generalized Init Transformation

• Replace reduction buffer load with identity element (0)
• Apply to the entire expression to generate init value

InitSubstituter init_subst(inlined_buffer_, identity_elem);
PrimExpr init_epilogue = init_subst(epilogue_expression_);
// Simplify: 0 + C[vi, vj] -> C[vi, vj]

Examples:

• temp + C → 0 + C → C (simplify)
• max(temp + C, 0) → max(0 + C, 0) → max(C, 0)
• min(max(temp, lower), upper) → min(max(0, lower), upper)

4. Generalized Update Transformation

• Replace reduction buffer load with reduction update
• If parent is Add and the other operand is not a reduction buffer → treat as bias addend and remove
• Otherwise → apply expression as-is

class GeneralizedEpilogueApplier : public ExprMutator {
  // Replace reduction buffer load with reduction update
  // Automatically detect and remove bias addend in Add nodes
  // Automatically support other activation functions
};

Results and Verification

Existing Tests Pass

All existing tests pass, maintaining backward compatibility:

• test_fuse_reduction_epilogue_basic
• test_fuse_reduction_epilogue_fp32
• test_fuse_reduction_epilogue_numerical_correctness
• test_fuse_reduction_epilogue_multiple_epilogue
• test_matmul_bias_relu
• test_matmul_bias_relu_correctness_unified
• test_matmul_clipping
• test_matmul_clipping_correctness_unified
• Other commutative variants tests

Total: All 15 tests pass

[gemini-code-assist]

Summary of Changes

Hello @kimm240, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly refactors the fuseReductionEpilogue mechanism to support arbitrary epilogue expressions, moving away from a rigid pattern-matching approach. By storing the entire epilogue expression and introducing generalized transformation logic for both initialization and updates, the system becomes more flexible and extensible, capable of handling a wider range of post-reduction operations without needing explicit pattern definitions. This change improves the generality and maintainability of the code while ensuring backward compatibility with existing tests.

Highlights

• Pattern Matching Removal: The previous pattern-matching approach for epilogue types (Bias, BiasReLU, Clipping) has been entirely removed. This includes the deletion of the EpilogueType enum and the AnalyzeEpiloguePattern() function, simplifying the logic by directly processing the entire epilogue expression.
• Store Entire Epilogue Expression: The ReductionEpilogueFuser now stores the complete epilogue expression in epilogue_expression_ and identifies the specific reduction_buffer_load_ within it. This change eliminates the need for pattern-specific member variables and allows for a more generalized handling of post-reduction operations.
• Generalized Init Transformation: A new InitSubstituter class has been introduced to handle the initialization of the fused reduction block. It replaces the reduction buffer load within the full epilogue expression with an identity element (e.g., 0 for addition-based reductions), allowing for a generalized calculation of the initial value, which is then simplified.
• Generalized Update Transformation: The update logic now uses a GeneralizedEpilogueApplier within an UpdateSubstituter. This applier intelligently replaces the reduction buffer load with the reduction update. It also includes logic to automatically detect and remove bias addends if the reduction buffer load's parent is an Add node and the other operand is not a reduction buffer, supporting arbitrary activation functions.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request successfully generalizes the fuseReductionEpilogue primitive by removing the previous pattern-matching logic in favor of a more flexible, expression-based approach. This is a significant improvement that allows for handling arbitrary epilogue expressions. The implementation is well-aligned with the description. I have identified a couple of areas for improvement, mainly related to removing some dead code and simplifying a complex section of the implementation to enhance readability and maintainability. Overall, this is a great enhancement to the framework.

[kimm240]

@wrongtest-intellif In this PR, I implement general epilogue form. This idea is discussed in #18515.

[wrongtest-intellif]

The generalized version cover existing testcases. But since it is a "generalized" version, it would be great to add more self-contained checking or comments to prove that the transformation do not accept "false-positive" patterns. It may not be a trivial problem but we could try some improvements.

Refer to https://arxiv.org/html/2510.08726v1, which even leverage symbolic solvers for reduction fusion transformations.

[wrongtest-intellif]

Generally, we require epilogue(zero ⊕ x0 ⊕ ... ⊕ xn) == g(xn, g(xn-1, g(.....g(x0, epilogue_init(zero)))...) with g and epilogue_init deduce from epilogue expression, right?

[kimm240]

For BiasAdd, the transformation $epilogue(\sum x) = \sum x + Bias$ holds regardless of whether the bias is handled at the initialization stage or the end.

However, for non-linear operations like ReLU and Clipping, this identity does not hold if we simply substitute the identity element in the Init block, as $f(\sum x) \neq f(0) + \sum x$.

I initially aimed to generalize the logic by processing these operations through a unified substitution mechanism (essentially treating them in a 'per-iteration') to maintain the existing framework's structure-what is merged in #18515.

Should we strictly move these non-linear transformations to the final Store stage (applied once to the final sum)?
Or, since the previous 'per-iteration' behavior was already merged and used, should we keep it as an option or a specific 'fused-update' mode? I'm eager to hear about it.

[wrongtest-intellif]

The correctness matters most. The concrete strategy could be free.

[wrongtest-intellif]

Generally look good to me. Please fix lints if you'd like to check in current round of changes. cc @kimm240

[kimm240]

@wrongtest-intellif
I fix lint and push it.

If there are any specific features, improvements, or roadmap items you'd like to see developed in these areas, please let me know. I'd be more than happy to take them on!