[PyTorch] Add API to Query if a PyTorch Op is Supported in Core ML (#2081)

add API to query if a PyTorch op is supported in Core ML

---------

Co-authored-by: yifan_shen3 <[email protected]>

Author: YifanShenSZ

coremltools/converters/mil/frontend/torch/__init__.py

@@ -11,4 +11,4 @@
     from . import ops, quantization_ops
     from .dialect_ops import (torch_tensor_assign, torch_upsample_bilinear,
                               torch_upsample_nearest_neighbor)
-    from .torch_op_registry import register_torch_op
+    from .torch_op_registry import register_torch_op, is_torch_fx_node_supported

coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py

@@ -5,6 +5,7 @@
 
 import itertools
 import os
+from unittest.mock import patch
 
 import numpy as np
 import pytest
@@ -19,6 +20,7 @@
 )
 from coremltools.converters.mil.frontend.torch.test.testing_utils import _copy_input_data
 from coremltools.converters.mil.frontend.torch.torch_op_registry import (
+    TorchOpsRegistry,
     _TORCH_OPS_REGISTRY,
     register_torch_op,
 )
@@ -264,6 +266,71 @@ def test_func_dummy(context, inputs):
         # Cleanup the test
         del _TORCH_OPS_REGISTRY.name_to_func_mapping["test_func_dummy"]
 
+
+@pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND)
+class TestFxNodeSupport:
+    """
+    The API ``ct.converters.mil.frontend.torch.is_torch_fx_node_supported`` is used
+    by 3rd-party code ExecuTorch: https://github.com/pytorch/executorch/pull/1415,
+    so we cannot break it
+    """
+
+    @staticmethod
+    def test_simple_case():
+        class Model(torch.nn.Module):
+            def forward(self, a, x, b):
+                y = torch.mm(a, x)
+                z = y + b
+                a.sub_(z)
+                y = torch.mm(a, x)
+                z = y + b
+                return z
+
+        model = Model()
+        model.eval()
+        symbolic_traced = torch.fx.symbolic_trace(model)
+
+        for node in symbolic_traced.graph.nodes:
+            # There are many types of torch fx node,
+            # we only support "call_function" node for now
+            if node.op == "call_function":
+                # All PyTorch ops in the example model are supported, so they should all return true
+                assert ct.converters.mil.frontend.torch.is_torch_fx_node_supported(node)
+            # Other types of torch fx node are not supported
+            else:
+                assert not ct.converters.mil.frontend.torch.is_torch_fx_node_supported(node)
+
+    @staticmethod
+    def test_unsupported_op():
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                z = x + y
+                return torch.nn.functional.softmax(z)
+
+        model = Model()
+        model.eval()
+        symbolic_traced = torch.fx.symbolic_trace(model)
+
+        # Mock our torch ops registry, pretending that only "add" is supported
+        with patch.object(
+            TorchOpsRegistry,
+            "__contains__",
+            side_effect=(lambda op_name: op_name == "add"),
+        ):
+            for node in symbolic_traced.graph.nodes:
+                # There are many types of torch fx node,
+                # we only support "call_function" node for now
+                if node.op == "call_function":
+                    # Only "add" is supported
+                    assert (
+                        (node.target.__name__.lower() == "add")
+                        == ct.converters.mil.frontend.torch.is_torch_fx_node_supported(node)
+                    )
+                # Other types of torch fx node are not supported
+                else:
+                    assert not ct.converters.mil.frontend.torch.is_torch_fx_node_supported(node)
+
+
 #################################################################################
 # Note: Starting from here, all of the following tests are also used as examples
 # in https://coremltools.readme.io/docs as a reference.

coremltools/converters/mil/frontend/torch/torch_op_registry.py

@@ -3,8 +3,11 @@
 #  Use of this source code is governed by a BSD-3-clause license that can be
 #  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
 
-from typing import Callable
+from typing import Callable, List
 
+import torch
+
+from coremltools import _logger as logger
 from coremltools.models._deprecation import deprecated as _deprecated
 
 
@@ -127,3 +130,45 @@ def func_wrapper(func):
         # decorator called without argument
         return func_wrapper
     return func_wrapper(_func)
+
+
+def is_torch_fx_node_supported(torch_fx_node: torch.fx.Node) -> bool:
+    # There are many types of torch fx node:
+    #     1. call_function
+    #     2. call_module
+    #     3. call_method
+    #     4. get_attr
+    #     5. placeholder
+    #     6. output
+    #     ...
+    # Only "call_*" nodes contain PyTorch ops,
+    # among them we only support "call_function" node for now
+    if torch_fx_node.op != "call_function":
+        logger.warning(
+            "For now, among all types of torch fx nodes, CoreML only supports call_function node"
+        )
+        return False
+
+    # Get the target in torch fx node, and canonicalize it to lower-case string
+    torch_fx_node_target = torch_fx_node.target
+    if isinstance(torch_fx_node_target, str):
+        torch_fx_node_target_name = torch_fx_node_target.lower()
+    else:
+        torch_fx_node_target_name = torch_fx_node.target.__name__.lower()
+    # Since we are only dealing with "call_function" node,
+    # the contained PyTorch op must be functional, i.e. not in-place
+    assert (
+        not torch_fx_node_target_name.endswith("_")
+    ), (
+        "For now, since CoreML only supports call_function torch fx node, "
+        "all ops should be functional, i.e. there should not be any in-place op"
+    )
+    # Target name may or may not contain prefix "aten.":
+    #     1. For usual fx node, target is a PyTorch function, i.e. no prefix
+    #     2. For executorch exported fx node, target is executorch.exir.dialects.edge._ops.EdgeOp,
+    #        whose name has format "aten.xx.yy"
+    _ATEN_NODE_PREFIX = "aten."
+    if torch_fx_node_target_name.startswith(_ATEN_NODE_PREFIX):
+        torch_fx_node_target_name = torch_fx_node_target_name[len(_ATEN_NODE_PREFIX):]
+
+    return torch_fx_node_target_name in _TORCH_OPS_REGISTRY