Copy unit tests from torchgen to ET codegen

Differential Revision: D75236020

Pull Request resolved: #11074

Author: larryliu0820

codegen/test/TARGETS

@@ -0,0 +1,8 @@
+# Any targets that should be shared between fbcode and xplat must be defined in
+# targets.bzl. This file can contain xplat-only targets.
+
+load(":targets.bzl", "define_common_targets")
+
+oncall("executorch")
+
+define_common_targets()

codegen/test/targets.bzl

@@ -0,0 +1,21 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_oss_build_kwargs", "runtime")
+
+def define_common_targets():
+  """Defines targets that should be shared between fbcode and xplat.
+
+  The directory containing this targets.bzl file should also contain both
+  TARGETS and BUCK files that call this function.
+  """
+
+  runtime.python_test(
+      name = "test_gen",
+      srcs = glob(["test_*.py"]),
+      package_style = "inplace",
+      deps = [
+          "//executorch/codegen:gen_lib",
+          "fbsource//third-party/pypi/expecttest:expecttest",
+      ],
+      external_deps = [
+          "torchgen",
+      ],
+  )

codegen/test/test_executorch_custom_ops.py

@@ -0,0 +1,153 @@
+# 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.
+
+from __future__ import annotations
+
+import tempfile
+import unittest
+from typing import Any
+from unittest.mock import ANY, Mock, patch
+
+import expecttest
+
+import torchgen
+from executorch.codegen.api.custom_ops import ComputeNativeFunctionStub
+from executorch.codegen.model import ETKernelIndex
+from torchgen.gen_executorch import gen_headers
+from torchgen.model import Location, NativeFunction
+from torchgen.selective_build.selector import SelectiveBuilder
+from torchgen.utils import FileManager
+
+
+SPACES = "    "
+
+
+def _get_native_function_from_yaml(yaml_obj: dict[str, object]) -> NativeFunction:
+    native_function, _ = NativeFunction.from_yaml(
+        yaml_obj,
+        loc=Location(__file__, 1),
+        valid_tags=set(),
+    )
+    return native_function
+
+
+class TestComputeNativeFunctionStub(expecttest.TestCase):
+    """
+    Could use torch.testing._internal.common_utils to reduce boilerplate.
+    GH CI job doesn't build torch before running tools unit tests, hence
+    manually adding these parametrized tests.
+    """
+
+    def _test_function_schema_generates_correct_kernel(
+        self, obj: dict[str, Any], expected: str
+    ) -> None:
+        func = _get_native_function_from_yaml(obj)
+
+        gen = ComputeNativeFunctionStub()
+        res = gen(func)
+        self.assertIsNotNone(res)
+        self.assertExpectedInline(
+            str(res),
+            expected,
+        )
+
+    def test_function_schema_generates_correct_kernel_tensor_out(self) -> None:
+        obj = {"func": "custom::foo.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)"}
+        expected = """
+at::Tensor & wrapper_CPU_out_foo_out(const at::Tensor & self, at::Tensor & out) {
+    return out;
+}
+    """
+        self._test_function_schema_generates_correct_kernel(obj, expected)
+
+    def test_function_schema_generates_correct_kernel_no_out(self) -> None:
+        obj = {"func": "custom::foo.Tensor(Tensor self) -> Tensor"}
+        expected = """
+at::Tensor wrapper_CPU_Tensor_foo(const at::Tensor & self) {
+    return self;
+}
+    """
+        self._test_function_schema_generates_correct_kernel(obj, expected)
+
+    def test_function_schema_generates_correct_kernel_no_return(self) -> None:
+        obj = {"func": "custom::foo.out(Tensor self, *, Tensor(a!)[] out) -> ()"}
+        expected = f"""
+void wrapper_CPU_out_foo_out(const at::Tensor & self, at::TensorList out) {{
+{SPACES}
+}}
+    """
+        self._test_function_schema_generates_correct_kernel(obj, expected)
+
+    def test_function_schema_generates_correct_kernel_3_returns(self) -> None:
+        obj = {
+            "func": "custom::foo(Tensor self, Tensor[] other) -> (Tensor, Tensor, Tensor)"
+        }
+        expected = """
+::std::tuple<at::Tensor,at::Tensor,at::Tensor> wrapper_CPU__foo(const at::Tensor & self, at::TensorList other) {
+    return ::std::tuple<at::Tensor, at::Tensor, at::Tensor>(
+                at::Tensor(), at::Tensor(), at::Tensor()
+            );
+}
+    """
+        self._test_function_schema_generates_correct_kernel(obj, expected)
+
+    def test_function_schema_generates_correct_kernel_1_return_no_out(self) -> None:
+        obj = {"func": "custom::foo(Tensor[] a) -> Tensor"}
+        expected = """
+at::Tensor wrapper_CPU__foo(at::TensorList a) {
+    return at::Tensor();
+}
+    """
+        self._test_function_schema_generates_correct_kernel(obj, expected)
+
+    def test_schema_has_no_return_type_argument_throws(self) -> None:
+        func = _get_native_function_from_yaml(
+            {"func": "custom::foo.bool(Tensor self) -> bool"}
+        )
+
+        gen = ComputeNativeFunctionStub()
+        with self.assertRaisesRegex(Exception, "Can't handle this return type"):
+            gen(func)
+
+
+class TestGenCustomOpsHeader(unittest.TestCase):
+    @patch.object(torchgen.utils.FileManager, "write_with_template")
+    @patch.object(torchgen.utils.FileManager, "write")
+    def test_fm_writes_custom_ops_header_when_boolean_is_true(
+        self, unused: Mock, mock_method: Mock
+    ) -> None:
+        with tempfile.TemporaryDirectory() as tempdir:
+            fm = FileManager(tempdir, tempdir, False)
+            gen_headers(
+                native_functions=[],
+                gen_custom_ops_header=True,
+                custom_ops_native_functions=[],
+                selector=SelectiveBuilder.get_nop_selector(),
+                kernel_index=ETKernelIndex(index={}),  # type: ignore[arg-type]
+                cpu_fm=fm,
+                use_aten_lib=False,
+            )
+            mock_method.assert_called_once_with(
+                "CustomOpsNativeFunctions.h", "NativeFunctions.h", ANY
+            )
+
+    @patch.object(torchgen.utils.FileManager, "write_with_template")
+    @patch.object(torchgen.utils.FileManager, "write")
+    def test_fm_doesnot_writes_custom_ops_header_when_boolean_is_false(
+        self, unused: Mock, mock_method: Mock
+    ) -> None:
+        with tempfile.TemporaryDirectory() as tempdir:
+            fm = FileManager(tempdir, tempdir, False)
+            gen_headers(
+                native_functions=[],
+                gen_custom_ops_header=False,
+                custom_ops_native_functions=[],
+                selector=SelectiveBuilder.get_nop_selector(),
+                kernel_index=ETKernelIndex(index={}),  # type: ignore[arg-type]
+                cpu_fm=fm,
+                use_aten_lib=False,
+            )
+            mock_method.assert_not_called()