Add FoldQATConvBNPass to fold QAT Conv-BN simulation chains into conv bias (#19315)

Summary:

Add FoldQATConvBNPass to the Cadence AOT compiler pipeline to handle QAT Conv-BN simulated fusion patterns that survive into the exported graph.

When a model is exported after QAT training, the Conv-BN simulation chain (add(var+eps) -> sqrt -> div(bn_weight) -> div(conv_out/scale) -> add(orig_bias) -> batch_norm) may not be folded by TorchAO `_fold_conv_bn_qat` due to pattern mismatch. This leaves non-quantized add/div/sqrt nodes in the graph that cause QuantFusion to crash when it tries to fuse them as quantized add ops.

The fix has three parts:

1. Add `conv1d.default` to `QuantizeFusedConvBnBiasAtenPass` conv_targets so it matches conv1d ops and can create zero biases for convs without one (mirrors the existing conv2d support).

2. Add `FoldQATConvBNPass` which matches the QAT simulation chain, computes the BN correction constant C = (orig_bias - running_mean) * bn_weight / sqrt(running_var + eps) + bn_bias, folds C into the conv quantized bias tensor, and removes the simulation chain + batch_norm nodes. No new graph nodes are created.

3. Apply these passes in the correct order in both the `get_fake_quant_model` (pre-export, on GraphModule) and `apply_pre_edge_transform_passes` (post-export, on ExportedProgram) pipelines: first `QuantizeFusedConvBnBiasAtenPass` to create zero biases for convs that lack one, then `FoldQATConvBNPass` to fold the simulation chain into those biases.

Differential Revision: D103949573

Author: Reza Sajadiany

backends/cadence/aot/BUCK

@@ -31,12 +31,23 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
+fbcode_target(_kind = runtime.python_library,
+    name = "fold_qat_conv_bn",
+    srcs = [
+        "fold_qat_conv_bn.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_library,
     name = "compiler",
     srcs = [
         "compiler.py",
     ],
     deps = [
+        ":fold_qat_conv_bn",
         ":memory_planning",
         ":ops_registrations",
         ":passes",
@@ -46,6 +57,7 @@ fbcode_target(_kind = runtime.python_library,
         "//caffe2:torch",
         "//executorch/backends/cadence/aot/quantizer:fusion_pass",
         "//executorch/backends/cadence/aot/quantizer:quantizer",
+        "//executorch/backends/transforms:quantize_fused_convbn_bias_pass",
         "//executorch/backends/transforms:decompose_sdpa",
         "//executorch/backends/transforms:remove_clone_ops",
         "//executorch/devtools:lib",
@@ -512,6 +524,23 @@ fbcode_target(_kind = python_unittest,
     ],
 )
 
+fbcode_target(_kind = python_unittest,
+    name = "test_fold_qat_conv_bn",
+    srcs = [
+        "tests/test_fold_qat_conv_bn.py",
+    ],
+    supports_static_listing = False,
+    typing = True,
+    deps = [
+        ":compiler",
+        ":fold_qat_conv_bn",
+        "//caffe2:torch",
+        "//executorch/backends/cadence/aot:ops_registrations",
+        "//executorch/backends/cadence/aot/quantizer:quantizer",
+        "//executorch/backends/transforms:quantize_fused_convbn_bias_pass",
+    ],
+)
+
 fbcode_target(_kind = python_unittest,
     name = "test_remove_ops_passes",
     srcs = [

backends/cadence/aot/compiler.py

@@ -17,6 +17,7 @@
     QuantizedInputWrapper,
     trace as trace_fn,
 )
+from executorch.backends.cadence.aot.fold_qat_conv_bn import FoldQATConvBNPass
 from executorch.backends.cadence.aot.memory_planning import (
     CadenceMemoryPlanning,
     print_memory_planning_info,
@@ -30,6 +31,9 @@
     get_default_memory_config,
     MemoryConfig,
 )
+from executorch.backends.transforms.quantize_fused_convbn_bias_pass import (
+    QuantizeFusedConvBnBiasAtenPass,
+)
 from executorch.devtools import generate_etrecord
 from executorch.exir import (
     EdgeCompileConfig,
@@ -162,6 +166,18 @@ def apply_pre_edge_transform_passes(
     which will instantiate a default quantizer for you if needed.
     Returns an ExportedProgram with the fused model.
     """
+    # Create zero biases for convs without one, quantize any float biases if exists
+    converted_program = _transform(
+        converted_program,
+        QuantizeFusedConvBnBiasAtenPass(
+            exported_program=converted_program, default_zero_bias=True
+        ),
+    )
+
+    # Fold QAT Conv-BN simulated fusion patterns
+    # Removes (div(scale) → add(bias) → batch_norm chain and absorbs the correction into the conv bias
+    FoldQATConvBNPass(converted_program)(converted_program.graph_module)
+
     # Get patterns and apply fusion of dq -> op -> q to qop
     # pyre-ignore[16]: no attribute
     patterns = [q.pattern for q in quantizer.quantizers]
@@ -205,6 +221,13 @@ def get_fake_quant_model(
 
     # Get converted graph module
     converted_gm = convert_pt2(prepared_gm, dump_graphs=dump_graphs)
+
+    # Create zero biases for convs without one, quantize any float biases
+    QuantizeFusedConvBnBiasAtenPass(default_zero_bias=True)(converted_gm)
+
+    # Fold QAT Conv-BN simulated fusion patterns (now all convs have a bias to fold into)
+    FoldQATConvBNPass()(converted_gm)
+
     return converted_gm