[ONNX] Update decomposition logic to loop over onnx registry (pytorch#153168)

* [ONNX] Update decomposition logic to loop over onnx registry (pytorch#151826)

Fixes pytorch#150367

This PR makes decomposition table from onnx registry, which includes registered ops not only ATen and prim. This will help to keep the custom ops that are specified in the custom_translation table from decomposition during ONNX export.
Pull Request resolved: pytorch#151826
Approved by: https://github.com/justinchuby

(cherry picked from commit 6cd1741)

* [ONNX] Add test for decomp_table update (pytorch#153671)

Added a test to strengthen the case for cherry-picking pytorch#153168. The original PR didn’t include this test since the fix for decomp_table and the registry was already covered by existing tests. However, it's reasonable to include a dedicated test for the specific issue (pytorch#150367 ) when considering the cherry-pick.
Pull Request resolved: pytorch#153671
Approved by: https://github.com/justinchuby

---------

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

Author: pytorchbot

test/onnx/exporter/test_api.py

@@ -359,6 +359,51 @@ def onnx_add(self: FLOAT, other: FLOAT) -> FLOAT:
         self.assertIn("Sub", all_nodes)
         self.assertNotIn("Add", all_nodes)
 
+    def test_custom_translation_table_supports_custom_op_with_its_decomp(self):
+        with torch.library._scoped_library("mylib", "FRAGMENT") as lib:
+            torch.library.define(
+                "mylib::foo",
+                "(Tensor a, Tensor b) -> Tensor",
+                tags=torch.Tag.pt2_compliant_tag,
+                lib=lib,
+            )
+
+            @torch.library.impl("mylib::foo", "CompositeImplicitAutograd", lib=lib)
+            @torch.library.register_fake("mylib::foo")
+            def foo_impl(a, b):
+                return a + b
+
+            class M(torch.nn.Module):
+                def forward(self, x, y):
+                    return torch.ops.mylib.foo(x, y)
+
+            def onnx_add(self: FLOAT, other: FLOAT) -> FLOAT:
+                # Replace add with Sub
+                return op.Sub(self, other)
+
+            # With the custom op defined, we can use it in the model
+            # and replace it with a custom translation table
+            custom_translation_table = {
+                torch.ops.mylib.foo.default: onnx_add,
+            }
+            onnx_program = torch.onnx.export(
+                M(),
+                (torch.ones(3, 3), torch.ones(3, 3)),
+                custom_translation_table=custom_translation_table,
+                dynamo=True,
+            )
+            all_nodes = [n.op_type for n in onnx_program.model.graph]
+            self.assertIn("Sub", all_nodes)
+            self.assertNotIn("Add", all_nodes)
+
+            # Without the custom op defined, it's going to be decomposed
+            onnx_program_decomp = torch.onnx.export(
+                M(), (torch.ones(3, 3), torch.ones(3, 3)), dynamo=True
+            )
+            all_nodes_decomp = [n.op_type for n in onnx_program_decomp.model.graph]
+            self.assertIn("Add", all_nodes_decomp)
+            self.assertNotIn("Sub", all_nodes_decomp)
+
 
 class TestFakeTensorExport(common_utils.TestCase):
     """Test exporting in fake mode."""

torch/onnx/_internal/exporter/_decomp.py

@@ -14,7 +14,7 @@
 
 def get_onnx_implemented_overloads(
     registry: _registration.ONNXRegistry,
-) -> list[torch._ops.OperatorBase]:
+) -> list[_registration.TorchOp]:
     """
     Creates a set of OperatorBase and Callable objects that represent ONNX-supported PyTorch operations.
 
@@ -24,24 +24,19 @@ def get_onnx_implemented_overloads(
     Returns:
         A collection of OperatorBase and Callable objects representing ONNX-supported PyTorch operations.
     """
-    registered_ops: list[torch._ops.OperatorBase] = []
-    for op_namespace in (torch.ops.aten, torch.ops.prims):
-        op_names = dir(op_namespace)
-        for op_name in op_names:
-            op_overload_packet = getattr(op_namespace, op_name)
-            if not isinstance(op_overload_packet, torch._ops.OpOverloadPacket):
-                continue
-
-            for overload_name in op_overload_packet.overloads():
-                op_overload = getattr(op_overload_packet, overload_name)
-                if registry.is_registered(op_overload):
-                    registered_ops.append(op_overload)
+    registered_ops: list[_registration.TorchOp] = []
+    for onnx_decomp_meta in registry.functions.values():
+        assert len(onnx_decomp_meta) > 0
+        # Different OnnxDecompMeta for the same TorchOp should
+        # have the same fx_target.
+        fx_target = onnx_decomp_meta[0].fx_target
+        registered_ops.append(fx_target)
     return registered_ops
 
 
 def create_onnx_friendly_decomposition_table(
-    onnx_registered_ops: set[torch._ops.OperatorBase],
-) -> dict[torch._ops.OperatorBase, Callable]:
+    onnx_registered_ops: set[_registration.TorchOp],
+) -> dict[_registration.TorchOp, Callable]:
     """
     This function creates a dictionary of op overloads and their decomposition functions
     for ops that do not have ONNX symbolic functions. If an op already has an ONNX symbolic function,
@@ -55,7 +50,7 @@ def create_onnx_friendly_decomposition_table(
         Dict[torch._ops.OperatorBase, Callable]: A dictionary that maps op overloads to their corresponding
         decomposition functions.
     """
-    decomposition_table: dict[torch._ops.OperatorBase, Callable] = {}
+    decomposition_table: dict[_registration.TorchOp, Callable] = {}
 
     for op_overload, decomp_fn in itertools.chain(
         torch.export.default_decompositions().items(),  # type: ignore[attr-defined]