[ONNX] Relax unsupported node analysis on complex dtype (pytorch#113785)

In cases like pytorch#113444, users usually stop at UnsupportedNodeAnalysis with unsupported nodes information. Although in SARIF, they can clearly see it's due to lack of COMPLEX support, in screen error message, it's only showing original FX node name, such as `aten.mul.Tensor`. ~~This PR catches the information from diagnostic messages and reveal it to users.~~

The root cause is that UnsupportedNodeAnalysis is leveraging on `onnxfunction_dispatcher.get_function_overloads()` to decide if an ATen is supported or not. However, in `onnxfunction_dispatcher.get_function_overloads()`, lacking of complex function support is considered unsupported. This PR defines Unsupported FX nodes as not in registry.
Pull Request resolved: pytorch#113785
Approved by: https://github.com/thiagocrepaldi

test/onnx/dynamo/test_registry_dispatcher.py

@@ -13,7 +13,13 @@
 from onnxscript import BFLOAT16, DOUBLE, FLOAT, FLOAT16  # type: ignore[import]
 from onnxscript.function_libs.torch_lib import ops  # type: ignore[import]
 from onnxscript.onnx_opset import opset15 as op  # type: ignore[import]
-from torch.onnx._internal.fx import diagnostics, onnxfunction_dispatcher, registration
+from torch.onnx._internal.diagnostics import infra
+from torch.onnx._internal.fx import (
+    analysis,
+    diagnostics,
+    onnxfunction_dispatcher,
+    registration,
+)
 from torch.testing._internal import common_utils
 
 # TODO: this can only be global. https://github.com/microsoft/onnxscript/issues/805
@@ -77,6 +83,60 @@ def test_custom(x, y):
             [test_original, test_custom],
         )
 
+    def test_unsupported_nodes_analysis_with_missing_aten_op(self):
+        # NOTE: simulate unsupported nodes
+        aten_mul_tensor = registration.OpName.from_name_parts(
+            namespace="aten", op_name="mul", overload="Tensor"
+        )
+        aten_mul_default = registration.OpName.from_name_parts(
+            namespace="aten", op_name="mul"
+        )
+        aten_add_tensor = registration.OpName.from_name_parts(
+            namespace="aten", op_name="add", overload="Tensor"
+        )
+        aten_add_default = registration.OpName.from_name_parts(
+            namespace="aten", op_name="add"
+        )
+
+        self.registry._registry.pop(aten_mul_tensor)
+        self.registry._registry.pop(aten_mul_default)
+        self.registry._registry.pop(aten_add_tensor)
+        self.registry._registry.pop(aten_add_default)
+
+        diagnostic_context = diagnostics.DiagnosticContext(
+            "torch.onnx.dynamo_export", torch.__version__
+        )
+        dispatcher = onnxfunction_dispatcher.OnnxFunctionDispatcher(
+            self.registry, diagnostic_context
+        )
+
+        graph: torch.fx.Graph = torch.fx.Graph()
+        x: torch.fx.Node = graph.create_node("placeholder", "x")
+        x.meta["val"] = torch.tensor(3.0)
+        b: torch.fx.Node = graph.create_node(
+            "call_function", target=torch.ops.aten.mul.Tensor, args=(x, x)
+        )
+        c: torch.fx.Node = graph.create_node(
+            "call_function", target=torch.ops.aten.add.Tensor, args=(b, b)
+        )
+        output: torch.fx.Node = graph.output(c)
+        module = torch.fx.GraphModule(torch.nn.Module(), graph)
+
+        with self.assertRaises(infra.RuntimeErrorWithDiagnostic):
+            analysis.UnsupportedFxNodesAnalysis(
+                diagnostic_context, module, dispatcher
+            ).analyze(infra.levels.ERROR)
+
+        try:
+            analysis.UnsupportedFxNodesAnalysis(
+                diagnostic_context, module, dispatcher
+            ).analyze(infra.levels.ERROR)
+        except infra.RuntimeErrorWithDiagnostic as e:
+            self.assertIn(
+                "Unsupported FX nodes: {'call_function': ['aten.mul.Tensor', 'aten.add.Tensor']}.",
+                e.diagnostic.message,
+            )
+
 
 @common_utils.instantiate_parametrized_tests
 class TestDispatcher(common_utils.TestCase):

test/onnx/test_fx_to_onnx.py

@@ -266,6 +266,48 @@ def forward(self, input):
             expected_node="aten.clone.default",
         )
 
+    def test_missing_complex_onnx_variant_raises_errors_in_dispatcher(self):
+        registry = torch.onnx.OnnxRegistry()
+
+        # NOTE: simulate unsupported nodes
+        aten_mul_tensor = registration.OpName.from_name_parts(
+            namespace="aten", op_name="mul", overload="Tensor"
+        )
+
+        # Only keep real aten.mul to test missing complex aten.mul
+        registry._registry[aten_mul_tensor] = [
+            onnx_func
+            for onnx_func in registry._registry[aten_mul_tensor]
+            if not onnx_func.is_complex
+        ]
+
+        class TraceModel(torch.nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.mul.Tensor(input, input)
+
+        x = torch.tensor([1 + 2j, 3 + 4j], dtype=torch.complex64)
+
+        with self.assertRaises(torch.onnx.OnnxExporterError) as e:
+            torch.onnx.dynamo_export(
+                TraceModel(),
+                x,
+                export_options=torch.onnx.ExportOptions(onnx_registry=registry),
+            )
+
+        try:
+            torch.onnx.dynamo_export(
+                TraceModel(),
+                x,
+                export_options=torch.onnx.ExportOptions(onnx_registry=registry),
+            )
+        except torch.onnx.OnnxExporterError as e:
+            assert_has_diagnostics(
+                e.onnx_program.diagnostic_context,
+                diagnostics.rules.no_symbolic_function_for_call_function,
+                diagnostics.levels.ERROR,
+                expected_node="aten.mul.Tensor",
+            )
+
     def test_dynamo_export_retains_readable_parameter_and_buffer_names(self):
         class SubModule(torch.nn.Module):
             def __init__(self):

torch/onnx/_internal/fx/analysis/unsupported_nodes.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Dict, List
+from typing import Dict
 
-import torch
-from torch.onnx._internal.fx import _pass, diagnostics
+from torch.onnx._internal.fx import _pass, diagnostics, registration
 
 
 @dataclasses.dataclass
@@ -52,23 +51,35 @@ def analyze(
             RuntimeErrorWithDiagnostic: If diagnostics are emitted and the diagnostic
                 level is `ERROR`.
         """
-        unsupported_nodes: List[torch.fx.Node] = []
+
+        op_to_target_mapping: Dict[str, Dict[str, None]] = {}
         for node in self.module.graph.nodes:
             if node.op == "call_function":
-                try:
-                    # NOTE: OPSchema matcher is not in this analysis scope.
-                    self.onnxfunction_dispatcher.get_function_overloads(
-                        node, self.diagnostic_context
+                # NOTE: OPSchema matcher is not in this analysis scope.
+                internal_opname: registration.OpName = (
+                    self.onnxfunction_dispatcher._get_aten_name(
+                        node=node, diagnostic_context=self.diagnostic_context
+                    )
+                )
+                overload_registration = (
+                    self.onnxfunction_dispatcher.onnx_registry.is_registered_op(
+                        namespace=internal_opname.namespace,
+                        op_name=internal_opname.op_name,
+                        overload=internal_opname.overload,
+                    )
+                )
+                # NOTE: Fall back to default overload if the ONNX registry doesn't have the overload.
+                default_registration = (
+                    self.onnxfunction_dispatcher.onnx_registry.is_registered_op(
+                        namespace=internal_opname.namespace,
+                        op_name=internal_opname.op_name,
+                        overload=None,
+                    )
+                )
+                if not overload_registration and not default_registration:
+                    op_to_target_mapping.setdefault(node.op, {}).setdefault(
+                        str(node.target), None
                     )
-                except diagnostics.RuntimeErrorWithDiagnostic as e:
-                    unsupported_nodes.append(node)
-
-        op_to_target_mapping: Dict[str, Dict[str, None]] = {}
-
-        for node in unsupported_nodes:
-            op = node.op
-            target = node.target
-            op_to_target_mapping.setdefault(op, {}).setdefault(str(target), None)
 
         analysis_result = UnsupportedFxNodesAnalysisResult(op_to_target_mapping)
         self._lint(analysis_result, diagnostic_level)