Preserve BatchNorm running state in composable quantizer

backends/arm/quantizer/arm_quantizer_utils.py

@@ -38,6 +38,8 @@
 
 logger = logging.getLogger(__name__)
 
+NodeTarget = Callable[..., Any] | str
+
 if TYPE_CHECKING:
     from executorch.backends.cortex_m.quantizer.pattern_matcher import PatternMatcher
 
@@ -255,6 +257,13 @@ class PatternQuantizer(Quantizer, QuantizerReporterUser):
         torch.ops.aten.conv_transpose2d.input,
     }
 
+    # These operands are parameters or mutable state, not activation inputs.
+    # Observing mutable state would redirect in-place updates to the observer
+    # output instead of the registered buffer.
+    UNOBSERVED_INPUT_INDICES: dict[NodeTarget, tuple[int, ...]] = {
+        torch.ops.aten.batch_norm.default: (1, 2, 3, 4),
+    }
+
     def __init__(
         self,
         quantization_config: QuantizationConfig | None,
@@ -319,13 +328,35 @@ def is_bias(self, node: Node) -> bool:
 
         return True
 
+    def is_unobserved_input(self, node: Node, input_node: Node) -> bool:
+        """Return whether an input must remain directly connected to state.
+
+        Args:
+            node: Consumer node whose positional inputs are inspected.
+            input_node: Input node being considered for quantization annotation.
+
+        Returns:
+            True when the input occupies an unobserved position for the target.
+
+        """
+        return any(
+            index < len(node.args) and input_node is node.args[index]
+            for index in self.UNOBSERVED_INPUT_INDICES.get(node.target, ())
+        )
+
     def annotate_match(
         self,
         match: list[Node],
         config: QuantizationConfig | None,
     ) -> None:
         """Annotates a matched pattern according to the given quantization
         config.
+
+        Args:
+            match: Nodes in the accepted pattern to annotate.
+            config: Quantization configuration for the pattern, or None to mark
+                its nodes as unquantized.
+
         """
 
         for node in match:
@@ -335,6 +366,8 @@ def annotate_match(
             for input_node in node.all_input_nodes:
                 if not has_float_output(input_node):
                     continue
+                if self.is_unobserved_input(node, input_node):
+                    continue
                 if self.is_weight(input_node):
                     input_qspec_map[input_node] = (
                         config.get_weight_qspec(node) if config else None

backends/arm/test/misc/test_bn_relu_folding_qat.py

@@ -16,6 +16,8 @@
 
 from executorch.backends.test.harness.tester import Quantize
 from torch import nn
+from torch.fx import Node
+from torchao.quantization.pt2e.quantize_pt2e import prepare_qat_pt2e
 
 
 input_t1 = Tuple[torch.Tensor]  # Input x
@@ -122,3 +124,52 @@ def test_bn_relu_folding_qat_tosa_INT(test_data):
         ),
     )
     pipeline.run()
+
+
+def test_qat_batch_norm_updates_registered_running_stats():
+    class ConvBatchNorm(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.conv = nn.Conv1d(2, 2, 1, bias=False)
+            self.bn = nn.BatchNorm1d(2)
+            with torch.no_grad():
+                self.conv.weight.copy_(torch.eye(2).unsqueeze(-1))
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.bn(self.conv(x))
+
+    sample = torch.stack((torch.full((4, 8), 3.0), torch.full((4, 8), -2.0)), dim=1)
+    model = torch.export.export(
+        ConvBatchNorm().train(), (sample,), strict=True
+    ).module()
+    quantizer = TOSAQuantizer(
+        TosaSpecification.create_from_string("TOSA-1.0+INT"),
+        use_composable_quantizer=True,
+    ).set_global(get_symmetric_quantization_config(is_qat=True, is_per_channel=False))
+
+    prepared = prepare_qat_pt2e(model, quantizer)
+    batch_norm_nodes = [
+        node
+        for node in prepared.graph.nodes
+        if node.target == torch.ops.aten.batch_norm.default
+    ]
+    assert len(batch_norm_nodes) == 1
+    batch_norm_node = batch_norm_nodes[0]
+
+    for state_node in batch_norm_node.args[1:5]:
+        assert isinstance(state_node, Node)
+        assert state_node.op == "get_attr"
+
+    running_mean_node = batch_norm_node.args[3]
+    running_var_node = batch_norm_node.args[4]
+    assert isinstance(running_mean_node, Node)
+    assert isinstance(running_var_node, Node)
+    running_mean = prepared.get_buffer(str(running_mean_node.target))
+    running_var = prepared.get_buffer(str(running_var_node.target))
+    initial_running_mean = running_mean.clone()
+    initial_running_var = running_var.clone()
+
+    prepared(sample)
+
+    assert not torch.equal(running_mean, initial_running_mean)
+    assert not torch.equal(running_var, initial_running_var)