Rewriter: Fold Batchnorm nodes

[AyoubMDL]

Fuses BatchNormalization nodes into the following nodes (Conv, ConvTranspose, Gemm)
(#2301)

[Copilot]

Pull Request Overview

This PR adds functionality to fold BatchNormalization nodes into the preceding Conv, ConvTranspose, and Gemm nodes, simplifying the graph and improving runtime efficiency.

• Implements rewrite rules for Gemm, Conv, and ConvTranspose to absorb batchnorm parameters into weights and biases.
• Adds comprehensive unit tests covering folding behavior with/without bias and a non-initializer scenario.
• Registers the new fold_batchnorm rule in the default rewrite-rule set.

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
onnxscript/rewriter/fold_batchnorm.py	Implements rewrite rules and helper functions to fold BatchNormalization.
onnxscript/rewriter/fold_batchnorm_test.py	Adds tests for folding BatchNorm into ConvTranspose, Conv, Gemm, and non-init case.
onnxscript/rewriter/init.py	Registers fold_batchnorm in the default rewrite-rule set.

Comments suppressed due to low confidence (2)

onnxscript/rewriter/fold_batchnorm.py:70

• [nitpick] Using x.name + "_bias" can produce non-obvious or conflicting initializer names. Consider deriving the bias name from the inbound node's output or weight name for clarity and uniqueness.

            bias_name = x.name + "_bias"

onnxscript/rewriter/fold_batchnorm_test.py:38

• Add a test case where the BatchNormalization node has a non-default epsilon attribute to ensure folding logic correctly respects custom epsilon values.

                Y = BatchNormalization(X1, gamma, beta, input_mean, input_var)

[codecov]

Codecov Report

Attention: Patch coverage is 98.75000% with 2 lines in your changes missing coverage. Please review.

Project coverage is 73.81%. Comparing base (a3ce145) to head (a0b7bee).
Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/rewriter/fuse_batchnorm_test.py	97.75%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2312      +/-   ##
==========================================
+ Coverage   73.68%   73.81%   +0.13%     
==========================================
  Files         240      242       +2     
  Lines       30978    31138     +160     
  Branches     3517     3530      +13     
==========================================
+ Hits        22825    22984     +159     
  Misses       6932     6932              
- Partials     1221     1222       +1     

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

[justinchuby]

Thank you!

1. Could you follow

   onnxscript/onnxscript/rewriter/llama_rule_sets.py

   Lines 35 to 48 in 2ae13be

   	class CastIdentity(orp.RewriteRuleClassBase):
   	"""Replaces ``Cast(., to=to)`` by ``Identity`` if possible."""
   	def pattern(self, op, x, to):
   	return op.Cast(x, to=to)
   	def rewrite(self, op, x: ir.Value, to: ir.Attr):
   	return op.Identity(x)
   	def check(self, context, x, to) -> orp.MatchResult:
   	check_result = orp.MatchResult()
   	if x.dtype != to.as_int():
   	return check_result.fail("Input and output types are not the same")
   	return check_result

   to implement the rules as an pattern.RewriteRuleClassBase? You may create additional subclasses for the different variants for the rules.
2. In the subclasses I would also create a docstring to describe what the rules do in plain language.
3. You may also follow the readme (https://github.com/microsoft/onnxscript#coding-style) to format the code with lintrunner.

[AyoubMDL]

1. Could you follow

   onnxscript/onnxscript/rewriter/llama_rule_sets.py

   Lines 35 to 48 in 2ae13be

   	class CastIdentity(orp.RewriteRuleClassBase):
   	"""Replaces ``Cast(., to=to)`` by ``Identity`` if possible."""
   	def pattern(self, op, x, to):
   	return op.Cast(x, to=to)
   	def rewrite(self, op, x: ir.Value, to: ir.Attr):
   	return op.Identity(x)
   	def check(self, context, x, to) -> orp.MatchResult:
   	check_result = orp.MatchResult()
   	if x.dtype != to.as_int():
   	return check_result.fail("Input and output types are not the same")
   	return check_result

   
   to implement the rules as an pattern.RewriteRuleClassBase? You may create additional subclasses for the different variants for the rules.

@justinchuby Added the suggestion—please let me know if it aligns with what you had in mind.

[justinchuby]

The rules look good to me, although I am not sure it should be part of the default rule set? @gramalingam do you have any recommendations? Should we enable multiple levels of optimization, or is this fine?

[justinchuby]

Thanks a lot!