Update on "Dtype selective build: fail if not xplat, if portable/optimized not in kernel_deps"

#10985

Try to make user error harder for dtype selective build.

Emit warning for now, as too many failures when set to failure :(
For example:
```
buck2 build //xplat/sgr/resources/tests/handwriting:pkg
buck2 build fbsource//xplat/sgr/resources/mwa:main_pkg_libAndroid

Differential Revision: [D75027794](https://our.internmc.facebook.com/intern/diff/D75027794/)

[ghstack-poisoned]

Author: lucylq

backends/arm/operators/op_avg_pool2d.py

@@ -17,6 +17,7 @@
     register_node_visitor,
 )
 from executorch.backends.arm.operators.operator_validation_utils import (
+    adjust_pooling_pad_if_needed,
     validate_num_inputs,
     validate_same_dtype,
 )
@@ -63,6 +64,20 @@ def _build_generic_avgpool2d(
         except IndexError:
             pad_size_list = [0, 0, 0, 0]
 
+        # Adjust the padding as necessary
+        pad_size_list[1] = adjust_pooling_pad_if_needed(
+            input_tensor.shape[2],
+            kernel_size_list[0],
+            stride_size_list[0],
+            pad_size_list[1],
+        )
+        pad_size_list[3] = adjust_pooling_pad_if_needed(
+            input_tensor.shape[3],
+            kernel_size_list[1],
+            stride_size_list[1],
+            pad_size_list[3],
+        )
+
         attr = ts.TosaSerializerAttribute()
         attr.PoolAttribute(
             kernel=kernel_size_list,
@@ -192,6 +207,20 @@ def _build_generic_avgpool2d(
         except IndexError:
             pad_size_list = [0, 0, 0, 0]
 
+        # Adjust the padding as necessary
+        pad_size_list[1] = adjust_pooling_pad_if_needed(
+            input_tensor.shape[2],
+            kernel_size_list[0],
+            stride_size_list[0],
+            pad_size_list[1],
+        )
+        pad_size_list[3] = adjust_pooling_pad_if_needed(
+            input_tensor.shape[3],
+            kernel_size_list[1],
+            stride_size_list[1],
+            pad_size_list[3],
+        )
+
         attr = ts.TosaSerializerAttribute()
         attr.AvgPool2dAttribute(
             kernel=kernel_size_list,

backends/arm/operators/op_max_pool2d.py

@@ -17,31 +17,14 @@
     register_node_visitor,
 )
 from executorch.backends.arm.operators.operator_validation_utils import (
+    adjust_pooling_pad_if_needed,
     validate_num_inputs,
     validate_same_dtype,
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
 
 
-# Similarly to Conv2d, the TOSA spec requires that following is exactly divisible:
-# `(input + 2 * pad - kernel_size) / stride`
-# PyTorch however, does not require this, so as needed, we must adjust the padding.
-def adjust_pad_if_needed(
-    input_size: int, kernel_size: int, stride: int, pad: int
-) -> int:
-    if pad == 0:
-        return pad
-
-    mod_remainder = (input_size + 2 * pad - kernel_size) % stride
-
-    # No need to adjust
-    if mod_remainder == 0:
-        return pad
-
-    return pad - mod_remainder
-
-
 @register_node_visitor
 class MaxPool2dVisitor_0_80(NodeVisitor):
     target = "aten.max_pool2d.default"
@@ -82,13 +65,13 @@ def define_node(
             pad_size_list = [0, 0, 0, 0]
 
         # Adjust the padding as necessary
-        pad_size_list[1] = adjust_pad_if_needed(
+        pad_size_list[1] = adjust_pooling_pad_if_needed(
             input_tensor.shape[2],
             kernel_size[0],
             stride[0],
             pad_size_list[1],
         )
-        pad_size_list[3] = adjust_pad_if_needed(
+        pad_size_list[3] = adjust_pooling_pad_if_needed(
             input_tensor.shape[3],
             kernel_size[1],
             stride[1],
@@ -167,13 +150,13 @@ def define_node(
             pad_size_list = [0, 0, 0, 0]
 
         # Adjust the padding as necessary
-        pad_size_list[1] = adjust_pad_if_needed(
+        pad_size_list[1] = adjust_pooling_pad_if_needed(
             input_tensor.shape[2],
             kernel_size[0],
             stride[0],
             pad_size_list[1],
         )
-        pad_size_list[3] = adjust_pad_if_needed(
+        pad_size_list[3] = adjust_pooling_pad_if_needed(
             input_tensor.shape[3],
             kernel_size[1],
             stride[1],

backends/arm/operators/operator_validation_utils.py

@@ -99,3 +99,40 @@ def validate_same_dtype(op_name: str, tensors: List[Any]):
                 f"{op_name}: Expected all tensors to have dtype {reference_dtype}, but "
                 f"found inconsistent dtype {tensor.dtype}."
             )
+
+
+def adjust_pooling_pad_if_needed(
+    input_size: int, kernel_size: int, stride: int, pad: int
+) -> int:
+    """
+    Calculates the padding that needs to be removed to a pooling window to make it
+    divisible by the kernels stride. All inputs should correspond to the same dimension.
+
+    Parameters:
+    -----------
+    input_size : int
+        The size of the input to the operator.
+
+    kernel_size : int
+        The size of the kernel.
+
+    stride : int
+        The size of the stride.
+
+    pad : int
+        The amount of padding.
+
+    Output:
+    -------
+    An int, representing the padding to remove to make the window divisible.
+    """
+    if pad == 0:
+        return pad
+
+    mod_remainder = (input_size + 2 * pad - kernel_size) % stride
+
+    # No need to adjust
+    if mod_remainder == 0:
+        return pad
+
+    return pad - mod_remainder

backends/arm/quantizer/quantization_annotator.py

@@ -10,6 +10,7 @@
 
 import torch
 import torch.fx
+import torch.nn.functional as F
 from executorch.backends.arm.quantizer import QuantizationConfig
 from executorch.backends.arm.tosa_utils import get_node_debug_info
 from torch.ao.quantization.quantizer import QuantizationSpecBase, SharedQuantizationSpec
@@ -142,29 +143,33 @@ def _match_pattern(
 
     Each 'pattern' element is composed of a list of disjunctive nodes types.
     """
-    assert len(pattern) == 2, "Only two-nodes patterns supported currently"
-
-    if node.target in pattern[0]:
-        assert len(node.users) != 0
-        parent = node
-        child = next(iter(node.users))
-    elif node.target in pattern[1]:
-        assert len(node.args) != 0
-        parent = node.args[0]  # type: ignore[assignment]
-        child = node
-    else:
-        return False
-
-    if len(parent.users) != 1:
-        return False
-
-    if parent.target not in pattern[0] or child.target not in pattern[1]:
-        return False
-
+    assert len(pattern) > 0, "No pattern provided"
     if filter_fn is not None:
-        return filter_fn(parent) and filter_fn(child)
-
-    return True
+        if not filter_fn(node):
+            return False
+    if len(pattern) == 1:
+        # Base case where it has passed the filter_fn. Simply look if node.target is in pattern.
+        return node.target in pattern[0]
+    if node.target not in [op for sub_pattern in pattern for op in sub_pattern]:
+        # node.target not in pattern. No need to look at the rest of the pattern.
+        return False
+    # Find the index of this node's target in pattern
+    idx = [node.target in sub_pattern for sub_pattern in pattern].index(True)
+    left_pattern = pattern[:idx]
+    # Exclude idx as this contains node.target which we have already matched
+    right_pattern = pattern[idx + 1 :]
+    left_condition = True
+    right_condition = True
+    # Recursively look at the rest of the pattern by calling this function for
+    # node's input and user node with updated patterns.
+    if len(left_pattern) > 0:
+        parent = node.all_input_nodes[0]
+        if len(parent.users) != 1:
+            return False
+        left_condition = _match_pattern(parent, left_pattern, filter_fn)
+    if len(right_pattern) > 0:
+        right_condition = _match_pattern(list(node.users)[0], right_pattern, filter_fn)
+    return left_condition and right_condition
 
 
 _one_to_one = [
@@ -274,6 +279,58 @@ def any_or_hardtanh_min_zero(n: Node):
         return n.target != torch.ops.aten.hardtanh.default or n.args[1] == 0
 
     if _match_pattern(
+        node,
+        [
+            [
+                torch.ops.aten.conv1d.default,
+                torch.ops.aten.conv2d.default,
+                torch.ops.aten.conv2d.padding,
+            ],
+            [torch.ops.aten.batch_norm.default, F.batch_norm],
+            [torch.ops.aten.relu.default, torch.ops.aten.hardtanh.default],
+        ],
+        filter_fn=any_or_hardtanh_min_zero,
+    ):
+        if node.target in (
+            torch.ops.aten.conv1d.default,
+            torch.ops.aten.conv2d.default,
+            torch.ops.aten.conv2d.padding,
+        ):
+            quant_properties.quant_inputs = [
+                _QuantProperty(0, input_act_qspec),
+                _QuantProperty(1, weight_qspec, mark_annotated=True),
+                _QuantProperty(2, bias_qspec, optional=True, mark_annotated=True),
+            ]
+        elif node.target in (
+            torch.ops.aten.relu.default,
+            torch.ops.aten.hardtanh.default,
+        ):
+            quant_properties.quant_output = _QuantProperty(0, output_act_qspec)
+
+    elif _match_pattern(
+        node,
+        [
+            [
+                torch.ops.aten.conv1d.default,
+                torch.ops.aten.conv2d.default,
+                torch.ops.aten.conv2d.padding,
+            ],
+            [torch.ops.aten.batch_norm.default, F.batch_norm],
+        ],
+    ):
+        if node.target in (
+            torch.ops.aten.conv1d.default,
+            torch.ops.aten.conv2d.default,
+            torch.ops.aten.conv2d.padding,
+        ):
+            quant_properties.quant_inputs = [
+                _QuantProperty(0, input_act_qspec),
+                _QuantProperty(1, weight_qspec, mark_annotated=True),
+                _QuantProperty(2, bias_qspec, optional=True, mark_annotated=True),
+            ]
+        elif node.target in [torch.ops.aten.batch_norm.default, F.batch_norm]:
+            quant_properties.quant_output = _QuantProperty(0, output_act_qspec)
+    elif _match_pattern(
         node,
         [
             [

backends/arm/test/misc/test_bn_relu_folding_qat.py

@@ -0,0 +1,66 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Tuple
+
+import torch
+import torch.nn.functional as F
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_quantization_config,
+    TOSAQuantizer,
+)
+from executorch.backends.arm.test import common, conftest
+from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineBI
+
+from executorch.backends.xnnpack.test.tester.tester import Quantize
+from torch import nn
+
+
+input_t1 = Tuple[torch.Tensor]  # Input x
+
+
+class ConvModule(torch.nn.Module):
+    input_shape = (1, 28, 28)
+    batch_size = 64
+    test_data: input_t1 = (torch.randn(batch_size, *input_shape),)
+
+    def __init__(self, batch_norm: bool = True) -> None:
+        super().__init__()
+        self.conv = torch.nn.Conv2d(1, 16, 3, stride=2)
+        self.bn = nn.BatchNorm2d(num_features=16) if batch_norm else nn.Identity()
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = F.relu(x)
+
+        return x
+
+
+models = {
+    "conv_bn_relu": ConvModule(batch_norm=True),
+    "conv_relu": ConvModule(batch_norm=False),
+}
+
+
+@common.parametrize("model", models)
+def test_qat_tosa_BI(model: torch.nn.Module):
+    pipeline = TosaPipelineBI[input_t1](model, model.test_data, [], [], qtol=1)
+    tosa_version = conftest.get_option("tosa_version")
+    tosa_profiles = {
+        "0.80": common.TosaSpecification.create_from_string("TOSA-0.80+BI"),
+        "1.0": common.TosaSpecification.create_from_string("TOSA-1.0+INT"),
+    }
+    tosa_spec = tosa_profiles[tosa_version]
+    quantizer = TOSAQuantizer(tosa_spec)
+    pipeline.change_args(
+        "quantize",
+        Quantize(
+            quantizer=quantizer,
+            quantization_config=get_symmetric_quantization_config(is_qat=True),
+            is_qat=True,
+        ),
+    )
+    pipeline.run()