diff --git a/backends/arm/quantizer/arm_quantizer_utils.py b/backends/arm/quantizer/arm_quantizer_utils.py index b15f01d11c8..633cdfa9f26 100644 --- a/backends/arm/quantizer/arm_quantizer_utils.py +++ b/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,6 +328,22 @@ 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], @@ -326,6 +351,12 @@ def annotate_match( ) -> 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 diff --git a/backends/arm/test/misc/test_bn_relu_folding_qat.py b/backends/arm/test/misc/test_bn_relu_folding_qat.py index 535ab6ea4e4..ef724737966 100644 --- a/backends/arm/test/misc/test_bn_relu_folding_qat.py +++ b/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)