Add FuseQATConvBN to fuse_ops (#19442)

Summary:

Adds a FuseQATConvBN which folds the QAT Conv-BN simulation chain (`conv → q → dq → div(scale) → add(orig_bias) → batch_norm`) inserted by `prepare_qat_pt2e` into the conv's quantized bias and removes the chain.

The pass runs in two steps inside a single `call()`:
  1. Bias prep — for each conv, create a zero-filled quantized bias if missing, or quantize a float bias as per-tensor int32. Required so step 2 has a quantized bias slot to write the BN correction into.
  2. Fold — for each matched chain, compute the BN correction
       C = (orig_bias - running_mean) * bn_weight / sqrt(running_var + eps) + bn_bias
     and absorb it into the conv's quantized bias in place. Erase the chain + batch_norm.

Differential Revision: D104497938

Author: ethansfng

backends/cadence/aot/BUCK

@@ -45,6 +45,7 @@ fbcode_target(_kind = runtime.python_library,
         ":utils",
         "//caffe2:torch",
         "//executorch/backends/cadence/aot/quantizer:fusion_pass",
+        "//executorch/backends/cadence/aot/quantizer/passes:fuse_ops",
         "//executorch/backends/cadence/aot/quantizer:quantizer",
         "//executorch/backends/transforms:decompose_sdpa",
         "//executorch/backends/transforms:remove_clone_ops",

backends/cadence/aot/compiler.py

@@ -22,6 +22,7 @@
     print_memory_planning_info,
 )
 from executorch.backends.cadence.aot.quantizer.fusion_pass import QuantFusion
+from executorch.backends.cadence.aot.quantizer.passes.fuse_ops import FuseQATConvBN
 from executorch.backends.cadence.aot.quantizer.quantizer import (
     CadenceDefaultQuantizer,
     CadenceQuantizer,
@@ -37,9 +38,10 @@
     ExecutorchBackendConfig,
     ExecutorchProgramManager,
 )
+from executorch.exir.pass_manager import PassManager
 from executorch.exir.passes import ToOutVarPass
 from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
-from executorch.exir.program._program import _transform, to_edge
+from executorch.exir.program._program import to_edge
 from torch.export.exported_program import ExportedProgram
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e
 
@@ -162,13 +164,17 @@ def apply_pre_edge_transform_passes(
     which will instantiate a default quantizer for you if needed.
     Returns an ExportedProgram with the fused model.
     """
-    # Get patterns and apply fusion of dq -> op -> q to qop
     # pyre-ignore[16]: no attribute
     patterns = [q.pattern for q in quantizer.quantizers]
-    fused_program = _transform(converted_program, QuantFusion(patterns))
+    PassManager(
+        [
+            FuseQATConvBN(converted_program),
+            QuantFusion(patterns),
+        ]
+    )(converted_program.graph_module)
 
     # Apply torch ops passes (e.g., ReplaceMulTensorWithMulAndFullOpsPass)
-    fused_program = apply_torch_ops_passes(fused_program)
+    fused_program = apply_torch_ops_passes(converted_program)
 
     return fused_program
 

backends/cadence/aot/quantizer/passes/BUCK

@@ -0,0 +1,16 @@
+load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("odai_jarvis")
+
+fbcode_target(_kind = runtime.python_library,
+    name = "fuse_ops",
+    srcs = [
+        "fuse_ops.py",
+    ],
+    typing = True,
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/transforms:quantize_fused_convbn_bias_pass",
+    ],
+)