Add util to find functional and/or out variant of an operator (#250)

Summary:
Pull Request resolved: #250

This is a useful API to find functional or out variant from a given operator.

It looks for all the overload of an operator name and see if the target variant exists in all the overloads.

Reviewed By: manuelcandales

Differential Revision: D49070398

fbshipit-source-id: db06e7d3be3214adc83227b3606329043e8fea08

Author: larryliu0820

exir/dialects/edge/TARGETS

@@ -17,8 +17,8 @@ python_library(
         "//caffe2:torch",
         "//caffe2/torchgen:torchgen",
         "//executorch/exir/dialects/edge/dtype:lib",
+        "//executorch/exir/dialects/edge/op:lib",
         "//executorch/exir/dialects/edge/spec:lib",
-        "//executorch/exir/operator:convert",
     ],
 )
 

exir/dialects/edge/_ops.py

@@ -4,24 +4,19 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import dataclasses
-import logging
 from typing import Any, Dict, List, Optional, Set, Union
 
 import pkg_resources
 
 import torch
 
 from executorch.exir.dialects.edge.dtype.supported import regular_tensor_str_to_dtypes
-from executorch.exir.dialects.edge.spec.utils import (
-    get_tensor_variable_names,
-    get_torch_op_overload,
-)
-from executorch.exir.operator.convert import _pybind_schema_to_native_schema
+from executorch.exir.dialects.edge.op.api import to_variant
+from executorch.exir.dialects.edge.spec.utils import get_tensor_variable_names
 
 # pyre-ignore
 from ruamel.yaml import YAML
-from torchgen.model import FunctionSchema
+from torchgen.model import SchemaKind
 
 
 class AllowedDtypeSet:
@@ -324,53 +319,9 @@ def to_out_variant(self) -> torch._ops.OpOverload:
         # return if already found
         if "_out_variant" in self.__dict__ and self._out_variant:
             return self._out_variant
-        # first check if the current operator is an out-variant
-        native_schema: Optional[FunctionSchema] = _pybind_schema_to_native_schema(
-            self._schema.schema
-        )
-        assert (
-            native_schema is not None
-        ), f"Schema: {self._schema} cannot be converted to torch.FunctionSchema"
-        if native_schema.is_out_fn():
-            out = get_torch_op_overload(
-                self.namespace,
-                self._schema.name.split("::")[1],
-                self._schema.overload_name,
-            )
-            self._out_variant = out
-            return out
-        # get all overloads
-        torch_packet = getattr(
-            getattr(torch.ops, self.namespace), self._schema.name.split("::")[1]
-        )
-        schemas: List[torch._C.FunctionSchema] = [
-            getattr(torch_packet, o)._schema
-            for o in torch._C._jit_get_operation(self._schema.name)[1]
-        ]
-        # compare the signature of out variant overload with the signature of the original overload
-        signature = dataclasses.replace(native_schema.signature(), returns=())
-        for schema in schemas:
-            # ignore self
-            if str(schema) == str(self._schema):
-                continue
-            native_s: Optional[FunctionSchema] = _pybind_schema_to_native_schema(schema)
-            if native_s is None:
-                logging.warning(
-                    f"Schema: {schema} cannot be converted to torch.FunctionSchema"
-                )
-                continue
-            if (
-                native_s.is_out_fn()
-                and dataclasses.replace(native_s.signature(), returns=()) == signature
-            ):
-                out = get_torch_op_overload(
-                    self.namespace, schema.name.split("::")[1], schema.overload_name
-                )
-                self._out_variant = out
-                return out
-        raise RuntimeError(
-            f"Out variant of operator {self.name()} can't be found. We've found the schemas of all the overloads: {[str(s) for s in schemas]}"
-        )
+        out_variant = to_variant(self._op, SchemaKind.out)
+        self._out_variant = out_variant
+        return out_variant
 
     def __getattr__(self, name):
         if name == "_schema":

exir/dialects/edge/op/TARGETS

@@ -1,4 +1,5 @@
 load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
+load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest")
 
 oncall("ai_infra_mobile_platform")
 
@@ -9,5 +10,19 @@ python_library(
     ],
     deps = [
         "//caffe2:torch",
+        "//caffe2/torchgen:torchgen",
+        "//executorch/exir/operator:convert",
+    ],
+)
+
+python_unittest(
+    name = "test_api",
+    srcs = [
+        "test/test_api.py",
+    ],
+    deps = [
+        ":lib",
+        "//caffe2:torch",
+        "//caffe2/torchgen:torchgen",
     ],
 )

exir/dialects/edge/op/api.py

@@ -7,10 +7,15 @@
 """
 APIs to help lowering edge dialect ops to other dialects.
 """
-from typing import Optional
+import dataclasses
+import logging
+from typing import List, Optional
 
 import torch
-from torch._ops import OpOverloadPacket
+
+from executorch.exir.operator.convert import _pybind_schema_to_native_schema
+from torch._ops import OpOverload, OpOverloadPacket
+from torchgen.model import FunctionSchema, SchemaKind
 
 
 def get_torch_op_overload(
@@ -26,3 +31,56 @@ def get_torch_op_overload(
 def get_callable(name):
     main, suffix = name.split(".")
     return get_torch_op_overload("aten", main, suffix)
+
+
+def to_variant(op: OpOverload, variant: SchemaKind) -> OpOverload:
+    """Given an operator overload, return its corresponding variant. Currently
+    only supports functional variant and out variant.
+    Argument:
+        op (OpOverload): operator overload instance.
+        variant (SchemaKind): the variant we are looking for.
+    Returns:
+        OpOverload: The matched variant operator.
+    Example:
+        torch.ops.aten.add.Tensor, SchemaKind.out -> torch.ops.aten.add.out
+        torch.ops.aten.add.out, SchemaKind.functional -> torch.ops.aten.add.Tensor
+    """
+    assert (
+        variant == SchemaKind.functional or variant == SchemaKind.out
+    ), f"Only support out variant and functional variant, got {variant}"
+    # first check if the current operator is the target variant
+    native_schema: Optional[FunctionSchema] = _pybind_schema_to_native_schema(
+        op._schema
+    )
+    assert (
+        native_schema is not None
+    ), f"Schema: {op._schema} cannot be converted to torch.FunctionSchema"
+
+    # get all overloads
+    torch_packet = getattr(
+        getattr(torch.ops, op.namespace), op._schema.name.split("::")[1]
+    )
+    schemas: List[torch._C.FunctionSchema] = [
+        getattr(torch_packet, o)._schema
+        for o in torch._C._jit_get_operation(op._schema.name)[1]
+    ]
+    # compare the signature of out variant overload with the signature of the original overload
+    signature = dataclasses.replace(native_schema.signature(), returns=())
+    for schema in schemas:
+        native_s: Optional[FunctionSchema] = _pybind_schema_to_native_schema(schema)
+        if native_s is None:
+            logging.warning(
+                f"Schema: {schema} cannot be converted to torch.FunctionSchema"
+            )
+            continue
+        if (
+            native_s.kind() == variant
+            and dataclasses.replace(native_s.signature(), returns=()) == signature
+        ):
+            op_variant = get_torch_op_overload(
+                op.namespace, schema.name.split("::")[1], schema.overload_name
+            )
+            return op_variant
+    raise RuntimeError(
+        f"{variant} variant of operator {op.name()} can't be found. We've found the schemas of all the overloads: {[str(s) for s in schemas]}"
+    )

exir/dialects/edge/op/test/test_api.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env fbpython
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+
+from executorch.exir.dialects.edge.op.api import to_variant
+from torchgen.model import SchemaKind
+
+aten = torch.ops.aten
+
+OPS_TO_FUNCTIONAL = {
+    aten.add.out: aten.add.Tensor,
+    aten._native_batch_norm_legit_no_training.out: aten._native_batch_norm_legit_no_training.default,
+    aten.addmm.out: aten.addmm.default,
+    aten.view_copy.out: aten.view_copy.default,
+}
+
+
+class TestApi(unittest.TestCase):
+    """Test api.py"""
+
+    def test_to_out_variant_returns_self_when_given_out_variant(self) -> None:
+        op = aten.add.out
+        variant = to_variant(op, SchemaKind.out)
+        self.assertEqual(variant, op)
+
+    def test_to_functional_variant_returns_self_when_given_functional(self) -> None:
+        op = aten.leaky_relu.default
+        variant = to_variant(op, SchemaKind.functional)
+        self.assertEqual(variant, op)
+
+    def test_to_functional_variant_returns_correct_op(
+        self,
+    ) -> None:
+        for op in OPS_TO_FUNCTIONAL:
+            variant = to_variant(op, SchemaKind.functional)
+            self.assertEqual(variant, OPS_TO_FUNCTIONAL[op])
+
+    def test_to_out_variant_returns_correct_op(
+        self,
+    ) -> None:
+        inv_map = {v: k for k, v in OPS_TO_FUNCTIONAL.items()}
+        for op in inv_map:
+            variant = to_variant(op, SchemaKind.out)
+            self.assertEqual(variant, inv_map[op])

exir/dialects/edge/test/test_edge_ops.py

@@ -417,7 +417,8 @@ def test_to_out_variant_returns_correct_op(self) -> None:
     def test_to_out_variant_raises_exception_when_no_out_variant(self) -> None:
         view_op = ops.edge.aten.view.default
         with self.assertRaisesRegex(
-            RuntimeError, "Out variant of operator aten::view can't be found"
+            RuntimeError,
+            "SchemaKind.out variant of operator aten::view can't be found.",
         ):
             view_op.to_out_variant()