Merge branch 'pytorch:main' into pr_model_improve

Author: Michiel-Olieslagers

backends/arm/_passes/cast_int64_pass.py

@@ -12,7 +12,6 @@
 from torch._export.utils import is_buffer
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
 
 
 class CastInt64BuffersToInt32Pass(ExportPass):

backends/arm/arm_backend.py

@@ -11,19 +11,13 @@
 # JIT compiler flows.
 #
 
-import logging
-
 from typing import List, Optional
 
 from executorch.backends.arm.tosa_specification import TosaSpecification
 
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
-
-
 class ArmCompileSpecBuilder:
     def __init__(self):
         self.compile_spec: List[CompileSpec] = []

backends/arm/operator_support/right_shift_support.py

@@ -17,7 +17,6 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
 
 
 @register_tosa_support_check

backends/arm/operator_support/slice_copy_support.py

@@ -16,7 +16,6 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
 
 
 @register_tosa_support_check

backends/arm/operator_support/tosa_supported_operators.py

@@ -112,6 +112,7 @@ def tosa_support_factory(
     # Negative checks: Remove nodes from partitioning
     negative_checks: list[OperatorSupportBase] = [
         CheckInt64Inputs(exported_program, reporter),
+        CheckFloat64Inputs(exported_program, reporter),
         *[
             reporter.wrap_check(check, f"Rejected by {check.__class__.__name__}")
             for check in (additional_checks if additional_checks else [])
@@ -443,3 +444,26 @@ def is_node_supported(
                         )
                         return False
         return True
+
+
+class CheckFloat64Inputs(OperatorSupportBase):
+
+    def __init__(
+        self, exported_program: ExportedProgram, reporter: WhyNoPartitionReporter
+    ):
+        self.reporter = reporter
+        super().__init__()
+
+    def is_node_supported(
+        self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
+    ) -> bool:
+
+        for input_node in node.all_input_nodes:
+            tensor = get_first_fake_tensor(input_node)
+            if tensor.dtype == torch.float64:
+                self.reporter.report_reject(
+                    node,
+                    f"Had float64 input {input_node.name} that couldn't be handled.",
+                )
+                return False
+        return True

backends/arm/operators/op_maximum.py

@@ -36,20 +36,27 @@ def define_node(
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
-        assert inputs[0].dtype == inputs[1].dtype
+        if inputs[0].dtype != inputs[1].dtype and inputs[0].dtype != output.dtype:
+            raise TypeError(
+                f"Data type of inputs and output must be the same. Got input 0 dtype: "
+                f"{inputs[0].dtype}, input 1 dtype: {inputs[1].dtype} and output "
+                f"dtype: {output.dtype}"
+            )
 
         scale_back = 1.0
         max_output = output
         if inputs[0].dtype == ts.DType.INT8:
             input_qparams = get_input_qparams(node)
-            assert (
-                len(input_qparams) == 2
-            ), f"Both inputs needs to have quantization information for {node}"
-            # insert RESCALEs to int32
-            assert (
-                input_qparams[0] == input_qparams[1]
-            ), "Both inputs must have same quantization for MAX"
+            if len(input_qparams) != 2:
+                raise ValueError(
+                    f"Both inputs need to have quantization information for {node}"
+                )
+            if input_qparams[0] != input_qparams[1]:
+                raise ValueError(
+                    "Both inputs must have the same quantization parameters for MAX"
+                )
 
+            # insert RESCALEs to int32
             operand_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(
                 tosa_graph, inputs, node
             )

backends/arm/operators/op_reciprocal.py

@@ -34,7 +34,16 @@ def define_node(
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
-        assert inputs[0].dtype == output.dtype == ts.DType.FP32
+        if len(node.all_input_nodes) != 1:
+            raise ValueError(
+                f"Expected 1 input for {self.target}, got {len(node.all_input_nodes)}"
+            )
+        if inputs[0].dtype != ts.DType.FP32 or output.dtype != ts.DType.FP32:
+            raise ValueError(
+                f"Input and output for {self.target} need to be FP32, got "
+                f"{inputs[0].dtype=} and {output.dtype=}"
+            )
+
         tosa_graph.addOperator(
             ts.TosaOp.Op().RECIPROCAL, [inputs[0].name], [output.name]
         )

backends/arm/operators/op_sub.py

@@ -40,9 +40,19 @@ def define_node(
     ) -> None:
         # Specification (0.80) states that input and output types
         # should all be the same
-        assert inputs[0].dtype == inputs[1].dtype == output.dtype
+        if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
+            raise TypeError(
+                f"All IO needs to have the same data type, got input 1: "
+                f"{inputs[0].dtype}, input 2: {inputs[1].dtype} and output: "
+                f"{output.dtype}"
+            )
+
         # Handle int8 (quantized) and int32
-        assert inputs[0].dtype in [ts.DType.INT8, ts.DType.INT32]
+        supported_dtypes = [ts.DType.INT8, ts.DType.INT32]
+        if inputs[0].dtype not in supported_dtypes:
+            raise TypeError(
+                f'IO data type needs to be {supported_dtypes}, got "{inputs[0].dtype}"'
+            )
 
         if inputs[0].dtype == ts.DType.INT8:
             rescaled_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(
@@ -97,15 +107,27 @@ def define_node(
     ) -> None:
         # Specification (0.80) states that input and output types
         # should all be the same
-        assert inputs[0].dtype == inputs[1].dtype == output.dtype
+        if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
+            raise TypeError(
+                f"All IO needs to have the same data type, got input 1: "
+                f"{inputs[0].dtype}, input 2: {inputs[1].dtype} and output: "
+                f"{output.dtype}"
+            )
 
         if inputs[0].dtype in [ts.DType.INT8, ts.DType.INT32]:
             # Call the inherited define_node for handling integers
             super().define_node(node, tosa_graph, inputs, output)
         else:
             # FP32 Sub lowering
-            assert inputs[0].dtype == ts.DType.FP32
-            assert output.dtype == ts.DType.FP32
+            if (
+                inputs[0].dtype != ts.DType.FP32
+                or inputs[1].dtype != ts.DType.FP32
+                or output.dtype != ts.DType.FP32
+            ):
+                raise TypeError(
+                    f"All IO needs to have data type fp32. Got: {inputs[0].dtype}, "
+                    f"input 2: {inputs[1].dtype} and output: {output.dtype}"
+                )
 
             # MI lowering
             tosa_graph.addOperator(

backends/arm/test/misc/test_debug_feats.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import logging
 import os
 import shutil
 import tempfile
@@ -15,9 +14,6 @@
 
 from executorch.backends.arm.test.tester.arm_tester import ArmTester
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
 
 class Linear(torch.nn.Module):
     def __init__(
@@ -205,7 +201,6 @@ def test_collate_tosa_BI_tests(self):
 
 
 def test_dump_tosa_ops(caplog):
-    caplog.set_level(logging.INFO)
     model = Linear(20, 30)
     (
         ArmTester(
@@ -222,7 +217,6 @@ def test_dump_tosa_ops(caplog):
 
 
 def test_fail_dump_tosa_ops(caplog):
-    caplog.set_level(logging.INFO)
 
     class Add(torch.nn.Module):
         def forward(self, x):

backends/arm/test/models/test_conformer.py

@@ -3,7 +3,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import logging
 import unittest
 
 import torch
@@ -14,10 +13,6 @@
 from torchaudio.models import Conformer
 
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-
 def get_test_inputs(dim, lengths, num_examples):
     return (torch.rand(num_examples, int(lengths.max()), dim), lengths)