Update base for Update on "[ET Device Support] Schema changes: device info on Tensor and buffer-level device array"

This diff adds device placement information to the ExecuTorch schema to support representing tensor-level device type information, which will be the basic requirement for the following tensor_parser updates.

This is part of the Phase 1 implementation to make ET device type work E2E without user-specified device placement.

Design doc: https://docs.google.com/document/d/1lwd9BlohmwkN5EEvRulO_b-XnZBwv1nMb5l2K3jfuwA/edit?tab=t.0#heading=h.o6anuvkix4bu

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

[ghstack-poisoned]

Author: Gasoonjia

backends/arm/quantizer/quantization_annotator.py

@@ -13,25 +13,27 @@
 import logging
 import operator
 from dataclasses import dataclass, replace
-from typing import Callable, cast, List, Optional, Sequence
+from typing import Any, Callable, cast, Iterable, List, NamedTuple, Optional, Sequence
 
 import torch
 import torch.fx
 from executorch.backends.arm.common.debug import get_node_debug_info
 from executorch.backends.arm.common.type import ensure_type
 from executorch.backends.arm.quantizer import QuantizationConfig
-from torch._subclasses import FakeTensor
 
+from torch._subclasses import FakeTensor
 from torch.fx import Node
 from torchao.quantization.pt2e import (
     FakeQuantize,
     FusedMovingAvgObsFakeQuantize,
     MovingAveragePerChannelMinMaxObserver,
     PartialWrapper,
 )
+
 from torchao.quantization.pt2e.quantizer import (
     annotate_input_qspec_map,
     annotate_output_qspec,
+    FixedQParamsQuantizationSpec,
     QuantizationSpec,
     QuantizationSpecBase,
     SharedQuantizationSpec,
@@ -78,6 +80,11 @@ def __init__(self):
         self.quant_output: Optional[_QuantProperty] = None
 
 
+class _QParams(NamedTuple):
+    scale: float
+    zero_point: int
+
+
 def _as_list(x):
     """Return ``x`` wrapped as a list if needed.
 
@@ -391,14 +398,16 @@ def _annotate_output(node: Node, quant_property: _QuantProperty):
 
 
 def _match_pattern(
-    node: Node, pattern: List[List], filter_fn: Optional[Callable[[Node], bool]] = None
+    node: Node,
+    pattern: Sequence[Iterable[object]],
+    filter_fn: Optional[Callable[[Node], bool]] = None,
 ) -> bool:
     """Check whether a node chain matches a pattern.
 
     Verify a chain of ancestors -> node -> descendants matches the provided
     ``pattern``. If ``filter_fn`` is provided, require all nodes in the chain
-    to pass the filter. Each pattern element is a list of disjunctive node
-    targets.
+    to pass the filter. Each pattern element is an iterable of disjunctive
+    node targets.
 
     """
     if len(pattern) < 1:
@@ -432,16 +441,39 @@ def _match_pattern(
     return left_condition and right_condition
 
 
-_conv_ops = [
+_conv_ops = {
     torch.ops.aten.conv1d.default,
     torch.ops.aten.conv2d.default,
     torch.ops.aten.conv2d.padding,
     torch.ops.aten.conv_transpose2d.input,
     torch.ops.aten.conv3d.default,
     torch.ops.aten.conv3d.padding,
-]
+}
 
-_one_to_one = [
+# For these ops, we use fixed qspecs, meaning that quantization params for
+# these are statically defined. This is to prevent issues with out-of-range
+# values when using dynamic quantization.
+#
+# Dict of operator to a dict of num_bits to qparams for that operator.
+_fixed_input_qspec_ops: dict[Any, dict[int, _QParams]] = {
+    # acos has a valid range of [-1, 1]
+    torch.ops.aten.acos.default: {
+        8: _QParams((1.0 - (-1.0)) / (1 << 8), 0),
+        16: _QParams((1.0 - (-1.0)) / (1 << 16), 0),
+    },
+    # asin has a valid range of [-1, 1]
+    torch.ops.aten.asin.default: {
+        8: _QParams((1.0 - (-1.0)) / (1 << 8), 0),
+        16: _QParams((1.0 - (-1.0)) / (1 << 16), 0),
+    },
+    # atanh has a valid range of (-1, 1) (excluding -1 and 1).
+    torch.ops.aten.atanh.default: {
+        8: _QParams((0.999 - (-0.999)) / (1 << 8), 0),
+        16: _QParams((0.99999 - (-0.99999)) / (1 << 16), 0),
+    },
+}
+
+_one_to_one = {
     torch.ops.aten.abs.default,
     torch.ops.aten.ceil.default,
     torch.ops.aten.erf.default,
@@ -472,16 +504,13 @@ def _match_pattern(
     torch.ops.aten.log1p.default,
     torch.ops.aten.acosh.default,
     torch.ops.aten.sign.default,
-    torch.ops.aten.asin.default,
-    torch.ops.aten.atanh.default,
     torch.ops.aten.asinh.default,
     torch.ops.aten.cosh.default,
-    torch.ops.aten.acos.default,
     torch.ops.aten.cumsum.default,
     torch.ops.aten.tan.default,
-]
+}
 
-_one_to_one_shared_input_qspec = [
+_one_to_one_shared_input_qspec = {
     torch.ops.aten.squeeze.default,
     torch.ops.aten.squeeze_copy.default,
     torch.ops.aten.squeeze_copy.dim,
@@ -539,9 +568,9 @@ def _match_pattern(
     # dequant -> neg -> requant chain.
     torch.ops.aten.neg.default,
     torch.ops.aten.detach_copy.default,
-]
+}
 
-_one_to_one_shared_input_or_input_act_qspec = [
+_one_to_one_shared_input_or_input_act_qspec = {
     torch.ops.aten.alias.default,
     torch.ops.aten.clone.default,
     torch.ops.aten.hardtanh.default,
@@ -562,7 +591,7 @@ def _match_pattern(
     torch.ops.aten.alias_copy.default,
     torch.ops.aten.pixel_shuffle.default,
     torch.ops.aten.pixel_unshuffle.default,
-]
+}
 
 
 def get_quant_properties(  # noqa: C901
@@ -615,13 +644,13 @@ def any_or_hardtanh_min_zero(n: Node):
         node,
         [
             _conv_ops,
-            [torch.ops.aten.batch_norm.default],
-            [
+            {torch.ops.aten.batch_norm.default},
+            {
                 torch.ops.aten.relu.default,
                 torch.ops.aten.relu_.default,
                 torch.ops.aten.hardtanh.default,
                 torch.ops.aten.hardtanh_.default,
-            ],
+            },
         ],
         filter_fn=any_or_hardtanh_min_zero,
     ):
@@ -644,7 +673,7 @@ def any_or_hardtanh_min_zero(n: Node):
         node,
         [
             _conv_ops,
-            [torch.ops.aten.batch_norm.default],
+            {torch.ops.aten.batch_norm.default},
         ],
     ):
         if node.target in _conv_ops:
@@ -654,23 +683,21 @@ def any_or_hardtanh_min_zero(n: Node):
                 _QuantProperty(1, conv_weight_qspec, mark_annotated=True),
                 _QuantProperty(2, bias_qspec, optional=True, mark_annotated=True),
             ]
-        elif node.target in [
-            torch.ops.aten.batch_norm.default,
-        ]:
+        elif node.target in {torch.ops.aten.batch_norm.default}:
             quant_properties.quant_output = _QuantProperty(0, output_act_qspec)
     elif not is_symmetric and _match_pattern(
         node,
         [
-            [
+            {
                 *_conv_ops,
                 torch.ops.aten.linear.default,
-            ],
-            [
+            },
+            {
                 torch.ops.aten.relu.default,
                 torch.ops.aten.relu_.default,
                 torch.ops.aten.hardtanh.default,
                 torch.ops.aten.hardtanh_.default,
-            ],
+            },
         ],
         any_or_hardtanh_min_zero,
     ):
@@ -784,6 +811,25 @@ def any_or_hardtanh_min_zero(n: Node):
     elif node.target in _one_to_one:
         quant_properties.quant_inputs = [_QuantProperty(0, input_act_qspec)]
         quant_properties.quant_output = _QuantProperty(0, output_act_qspec)
+    elif node.target in _fixed_input_qspec_ops:
+        num_bits = torch.iinfo(input_act_qspec.dtype).bits
+        qparams = _fixed_input_qspec_ops[node.target][num_bits]
+
+        quant_properties.quant_inputs = [
+            _QuantProperty(
+                0,
+                FixedQParamsQuantizationSpec(
+                    dtype=input_act_qspec.dtype,
+                    scale=qparams.scale,
+                    zero_point=qparams.zero_point,
+                    quant_min=input_act_qspec.quant_min,
+                    quant_max=input_act_qspec.quant_max,
+                    qscheme=input_act_qspec.qscheme,
+                    is_dynamic=input_act_qspec.is_dynamic,
+                ),
+            )
+        ]
+        quant_properties.quant_output = _QuantProperty(0, output_act_qspec)
     elif node.target in _one_to_one_shared_input_qspec:
         input_node = ensure_type(Node, node.args[0])
         quant_properties.quant_inputs = [_QuantProperty(0, input_act_qspec)]

backends/arm/test/ops/test_acos.py

@@ -65,7 +65,6 @@ def test_acos_tosa_INT(test_data: Tuple):
         (test_data(),),
         aten_op=aten_op,
         exir_op=exir_op,
-        frobenius_threshold=0.5,  # MLETORCH-1709
     )
     pipeline.run()
 

backends/arm/test/ops/test_asin.py

@@ -55,8 +55,6 @@ def test_asin_tosa_INT(test_data: Tuple):
         (test_data(),),
         aten_op=[],
         exir_op=[],
-        frobenius_threshold=0.6,  # MLETORCH-1709
-        cosine_threshold=0.8,  # MLETORCH-1709
     )
     pipeline.run()
 

backends/arm/test/ops/test_atanh.py

@@ -26,11 +26,10 @@
 test_data_suite = {
     "zeros": torch.zeros(1, 10, 10, 10),
     "zeros_alt_shape": torch.zeros(1, 10, 3, 5),
-    "ones": torch.ones(10, 10, 10),
     "rand": torch.rand(10, 10) - 0.5,
     "rand_alt_shape": torch.rand(1, 10, 3, 5) - 0.5,
     "ramp": torch.arange(-1, 1, 0.2),
-    "near_bounds": torch.tensor([-0.999999, -0.999, -0.9, 0.9, 0.999, 0.999999]),
+    "near_bounds": torch.tensor([-0.99, -0.9, 0.9, 0.99]),
     "on_bounds": torch.tensor([-1.0, 1.0]),
 }
 
@@ -58,9 +57,11 @@ def test_atanh_tosa_INT(test_data: Tuple):
         (test_data,),
         aten_op=aten_op,
         exir_op=exir_op,
-        frobenius_threshold=None,  # MLETORCH-1709
-        cosine_threshold=0.7,
     )
+    if torch.any(test_data >= 1) or torch.any(test_data <= -1):
+        # The quantized model will saturate to max/min values while the
+        # original model will return inf/-inf, so comparison wont be valid here.
+        pipeline.pop_stage("run_method_and_compare_outputs.original_model")
     pipeline.run()
 
 

backends/arm/test/tester/analyze_output_utils.py

@@ -337,22 +337,6 @@ def dump_error_output(
     logger.error(f"{atol=}, {rtol=}, {qtol=}")
 
 
-if __name__ == "__main__":
-    """This is expected to produce the example output of print_diff."""
-    torch.manual_seed(0)
-    a = torch.rand(3, 3, 2, 2) * 0.01
-    b = a.clone().detach()
-    logger.info(b)
-
-    # Errors in all channels in element (1,1)
-    a[1, :, 1, 1] = 0
-    # Errors in (0,0) and (1,1) in channel 1
-    a[2, 1, 1, 1] = 0
-    a[2, 1, 0, 0] = 0
-
-    print_error_diffs(a, b)
-
-
 def compare_rel_frobenius_and_cosine_similarity(
     reference_output: torch.Tensor,
     test_output: torch.Tensor,
@@ -452,3 +436,19 @@ def compare_rel_frobenius_and_cosine_similarity(
             f"Tensor-wise comparison failed: Cosine similarity {cosine_similarity} is below threshold {cosine_threshold}."
             f" (Relative frobenius error: {relative_frobenius_error}, threshold {frobenius_threshold})."
         )
+
+
+if __name__ == "__main__":
+    """This is expected to produce the example output of print_diff."""
+    torch.manual_seed(0)
+    a = torch.rand(3, 3, 2, 2) * 0.01
+    b = a.clone().detach()
+    logger.info(b)
+
+    # Errors in all channels in element (1,1)
+    a[1, :, 1, 1] = 0
+    # Errors in (0,0) and (1,1) in channel 1
+    a[2, 1, 1, 1] = 0
+    a[2, 1, 0, 0] = 0
+
+    print_error_diffs(a, b)

backends/cadence/aot/BUCK

@@ -300,6 +300,7 @@ fbcode_target(_kind = runtime.python_library,
     ],
     typing = True,
     deps = [
+        ":fuse_ops",
         ":ops_registrations",
         "//caffe2:torch",
         "//executorch/backends/cadence/aot:pass_utils",

backends/cadence/aot/decompose_ops.py

@@ -23,10 +23,12 @@
 from torch.fx.node import Argument
 
 
-@register_cadence_pass(CadencePassAttribute(opt_level=0))
+@register_cadence_pass(CadencePassAttribute(opt_level=1))
 class DecomposeAtenApproxGeluPass(ExportPass):
     """
-    Decompose the aten gelu op with an approximate arg to a series of simpler ops
+    Decompose the aten gelu op with an approximate arg to a series of simpler ops.
+    This is an optimization - gelu has a portable kernel fallback, but decomposing
+    may be more efficient on some backends.
     """
 
     def call_operator(

backends/cadence/aot/functions.yaml

@@ -309,10 +309,15 @@
     - arg_meta: null
       kernel_name: impl::generic::quantized_relu_asym8u_asym8u_per_tensor_out
 
-- func: cadence::quantized_max_pool2d.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode, *, Tensor(a!) out) -> Tensor(a!)
+- func: cadence::quantized_max_pool2d_nchw.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode, *, Tensor(a!) out) -> Tensor(a!)
   kernels:
     - arg_meta: null
-      kernel_name: impl::generic::quantized_max_pool2d_out
+      kernel_name: impl::generic::quantized_max_pool2d_nchw_out
+
+- func: cadence::quantized_max_pool2d_nhwc.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: impl::generic::quantized_max_pool2d_nhwc_out
 
 - func: cadence::quantized_add.per_tensor_out(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, int Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!)
   kernels:

backends/cadence/aot/fuse_ops.py

@@ -1170,7 +1170,10 @@ def can_fuse_for_chain(
             return False
 
         # checking that permut2(permut1(identity)) == identity, modulo unitary dimensions
-        input_shape = cast(torch.fx.Node, producer.args[0]).meta["val"].shape
+        producer_input = cast(torch.fx.Node, producer.args[0])
+        if "val" not in producer_input.meta:
+            return False
+        input_shape = producer_input.meta["val"].shape
         ident_dims = list(range(len(input_shape)))
         # this mapping helps to handle both transpose and permutations
         f: dict[Any, Callable] = {