Refactor codegen components to prepare for benchmark generation (pytorch#5560)

Summary:
Pull Request resolved: pytorch#5560

## Context

Refactor operator test code generation scripts, such that components can be re-used to generate operator benchmarks.

In broad strokes, the refactors implemented by this diff are as follows:

* Improve granularity of Python modules
* Replace `test` with `correctness_test`, to make it clear that we are generating correctness tests. **Note that I haven't changed the top level target name `compute_graph_op_tests_bin` since I believe it would be too verbose.
ghstack-source-id: 244283559
exported-using-ghexport

Reviewed By: nathanaelsee

Differential Revision: D63286131

fbshipit-source-id: 1177ea381e6381045f1c97491dd7ec006690f574

Author: SS-JIA

backends/vulkan/test/op_tests/cases.py

@@ -8,7 +8,7 @@
 from collections import namedtuple
 from typing import Callable
 
-from executorch.backends.vulkan.test.op_tests.utils.codegen import VkTestSuite
+from executorch.backends.vulkan.test.op_tests.utils.test_suite import VkTestSuite
 
 
 # Prime numbers dim sizes for testing

backends/vulkan/test/op_tests/generate_op_tests.py renamed to backends/vulkan/test/op_tests/generate_op_correctness_tests.py

@@ -10,12 +10,14 @@
 from typing import Dict
 
 from executorch.backends.vulkan.test.op_tests.cases import test_suites
+from executorch.backends.vulkan.test.op_tests.utils.gen_computegraph import (
+    ComputeGraphGen,
+)
 
-from executorch.backends.vulkan.test.op_tests.utils.codegen import VkCppTestFileGen
-from executorch.backends.vulkan.test.op_tests.utils.codegen_base import (
-    TestSuite,
-    TestSuiteGen,
+from executorch.backends.vulkan.test.op_tests.utils.gen_correctness_vk import (
+    VkCorrectnessTestFileGen,
 )
+from executorch.backends.vulkan.test.op_tests.utils.test_suite import TestSuite
 from torchgen import local
 
 from torchgen.gen import parse_native_yaml, ParsedYaml
@@ -37,7 +39,7 @@ def construct_f_map(parsed_yaml: ParsedYaml) -> Dict[str, NativeFunction]:
 
 
 def process_test_suites(
-    cpp_generator: VkCppTestFileGen,
+    cpp_generator: VkCorrectnessTestFileGen,
     f_map: Dict[str, NativeFunction],
     test_suites: Dict[str, TestSuite],
 ) -> None:
@@ -53,12 +55,12 @@ def generate_cpp(
     native_functions_yaml_path: str, tags_path: str, output_dir: str
 ) -> None:
     output_file = os.path.join(output_dir, "op_tests.cpp")
-    cpp_generator = VkCppTestFileGen(output_file)
+    cpp_generator = VkCorrectnessTestFileGen(output_file)
 
     parsed_yaml = parse_native_yaml(native_functions_yaml_path, tags_path)
     f_map = construct_f_map(parsed_yaml)
 
-    TestSuiteGen.backend_key = parsed_yaml.backend_indices[DispatchKey.CPU]
+    ComputeGraphGen.backend_key = parsed_yaml.backend_indices[DispatchKey.CPU]
 
     process_test_suites(cpp_generator, f_map, test_suites)
 
@@ -67,16 +69,14 @@ def generate_cpp(
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Generate a simple Hello World C++ program."
-    )
+    parser = argparse.ArgumentParser()
     parser.add_argument(
         "--aten-yaml-path",
         help="path to native_functions.yaml file.",
     )
     parser.add_argument(
         "--tags-path",
-        help="Path to tags.yaml. Required by yaml parsing in codegen system.",
+        help="Path to tags.yaml. Required by yaml parsing in gen_correctness_vk system.",
     )
     parser.add_argument("-o", "--output", help="Output directory", required=True)
     args = parser.parse_args()

backends/vulkan/test/op_tests/targets.bzl

@@ -8,9 +8,9 @@ def define_common_targets(is_fbcode = False):
         return
 
     runtime.python_library(
-        name = "generate_op_tests_lib",
+        name = "generate_op_correctness_tests_lib",
         srcs = native.glob(["utils/*.py"]) + [
-            "generate_op_tests.py",
+            "generate_op_correctness_tests.py",
             "cases.py",
         ],
         base_module = "executorch.backends.vulkan.test.op_tests",
@@ -21,23 +21,23 @@ def define_common_targets(is_fbcode = False):
     )
 
     runtime.python_binary(
-        name = "generate_op_tests",
-        main_module = "executorch.backends.vulkan.test.op_tests.generate_op_tests",
+        name = "generate_op_correctness_tests",
+        main_module = "executorch.backends.vulkan.test.op_tests.generate_op_correctness_tests",
         deps = [
-            ":generate_op_tests_lib",
+            ":generate_op_correctness_tests_lib",
         ],
     )
 
     aten_src_path = runtime.external_dep_location("aten-src-path")
     genrule_cmd = [
-        "$(exe :generate_op_tests)",
+        "$(exe :generate_op_correctness_tests)",
         "--tags-path $(location {})/aten/src/ATen/native/tags.yaml".format(aten_src_path),
         "--aten-yaml-path $(location {})/aten/src/ATen/native/native_functions.yaml".format(aten_src_path),
         "-o $OUT",
     ]
 
     runtime.genrule(
-        name = "generated_op_tests_cpp",
+        name = "generated_op_correctness_tests_cpp",
         outs = {
             "op_tests.cpp": ["op_tests.cpp"],
         },
@@ -66,7 +66,7 @@ def define_common_targets(is_fbcode = False):
     runtime.cxx_binary(
         name = "compute_graph_op_tests_bin",
         srcs = [
-            ":generated_op_tests_cpp[op_tests.cpp]",
+            ":generated_op_correctness_tests_cpp[op_tests.cpp]",
         ],
         define_static_target = False,
         deps = [
@@ -79,7 +79,7 @@ def define_common_targets(is_fbcode = False):
     runtime.cxx_test(
         name = "compute_graph_op_tests",
         srcs = [
-            ":generated_op_tests_cpp[op_tests.cpp]",
+            ":generated_op_correctness_tests_cpp[op_tests.cpp]",
         ],
         contacts = ["[email protected]"],
         fbandroid_additional_loaded_sonames = [

backends/vulkan/test/op_tests/utils/aten_types.py

@@ -0,0 +1,30 @@
+# 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.
+
+####################
+## ATen C++ Types ##
+####################
+
+AT_INT_ARRAY_REF = "at::IntArrayRef"
+AT_SCALAR = "at::Scalar"
+AT_TENSOR = "at::Tensor"
+AT_TENSOR_LIST = "at::TensorList"
+BOOL = "bool"
+DOUBLE = "double"
+INT = "int64_t"
+OPT_AT_DOUBLE_ARRAY_REF = "::std::optional<at::ArrayRef<double>>"
+OPT_AT_INT_ARRAY_REF = "at::OptionalIntArrayRef"
+OPT_AT_TENSOR = "::std::optional<at::Tensor>"
+OPT_BOOL = "::std::optional<bool>"
+OPT_INT64 = "::std::optional<int64_t>"
+OPT_DEVICE = "::std::optional<at::Device>"
+OPT_LAYOUT = "::std::optional<at::Layout>"
+OPT_MEMORY_FORMAT = "::std::optional<at::MemoryFormat>"
+OPT_SCALAR_TYPE = "::std::optional<at::ScalarType>"
+STRING = "c10::string_view"
+TWO_TENSOR_TUPLE = "::std::tuple<at::Tensor,at::Tensor>"
+THREE_TENSOR_TUPLE = "::std::tuple<at::Tensor,at::Tensor,at::Tensor>"
+TENSOR_VECTOR = "::std::vector<at::Tensor>"

backends/vulkan/test/op_tests/utils/codegen.py renamed to backends/vulkan/test/op_tests/utils/gen_computegraph.py

@@ -6,15 +6,14 @@
 
 import re
 from dataclasses import dataclass
-from typing import Any, List, Optional, Union
+from typing import List, Optional, Union
 
-from executorch.backends.vulkan.test.op_tests.utils.codegen_base import (
+from executorch.backends.vulkan.test.op_tests.utils.aten_types import (
     AT_INT_ARRAY_REF,
     AT_SCALAR,
     AT_TENSOR,
     AT_TENSOR_LIST,
     BOOL,
-    CppTestFileGen,
     DOUBLE,
     INT,
     OPT_AT_DOUBLE_ARRAY_REF,
@@ -28,37 +27,20 @@
     OPT_SCALAR_TYPE,
     STRING,
     TENSOR_VECTOR,
-    TestSuite,
-    TestSuiteGen,
     THREE_TENSOR_TUPLE,
     TWO_TENSOR_TUPLE,
 )
+from executorch.backends.vulkan.test.op_tests.utils.test_suite import TestSuite
 
 from torchgen.api import cpp
 from torchgen.api.types import CppSignatureGroup
-
 from torchgen.gen import generate_static_dispatch_backend_call, translate_args
-
 from torchgen.gen_aoti_c_shim import gen_static_dispatch_backend_call_signature
 from torchgen.model import NativeFunction, Variant
 
-##################################
-## Custom Test Suite Definition ##
-##################################
-
-
-@dataclass
-class VkTestSuite(TestSuite):
-    def __init__(self, input_cases: List[Any]):
-        super().__init__(input_cases)
-        self.storage_types: List[str] = ["utils::kTexture3D"]
-        self.layouts: List[str] = ["utils::kChannelsPacked"]
-        self.data_gen: str = "make_rand_tensor"
-
-
-##########################
-## Code Generator Class ##
-##########################
+###################################
+## Compute Graph Code Generation ##
+###################################
 
 
 @dataclass
@@ -105,6 +87,8 @@ def vk_out(self):
 
 
 class ComputeGraphGen:
+    backend_key = None
+
     def __init__(self, op_reg_name: str, f: NativeFunction, suite_def: TestSuite):
         self.op_reg_name = op_reg_name
         self.f = f
@@ -230,7 +214,7 @@ def gen_decl(self, fn_name: str, ret_type: str = "void") -> str:
 
     def create_aten_fn_call(self) -> str:
         func_call = generate_static_dispatch_backend_call(
-            self.f_sig, self.f, TestSuiteGen.backend_key
+            self.f_sig, self.f, ComputeGraphGen.backend_key
         )[7:].replace("::cpu", "")
 
         return func_call
@@ -611,147 +595,3 @@ def gen_op_check_fn(self) -> str:
         op_check_fn += "\n  }"
 
         return op_check_fn
-
-
-##################################
-## Test Fixture Code Generation ##
-##################################
-
-test_fixture_template = """
-class GeneratedOpsTest_{op_name} : public ::testing::TestWithParam< ::std::tuple<at::ScalarType, utils::StorageType, utils::GPUMemoryLayout>> {{
- protected:
-  ComputeGraph* graph;
-  at::ScalarType test_dtype = at::kFloat;
-  float rtol = {rtol};
-  float atol = {atol};
-
-  void SetUp() override {{
-    GraphConfig config;
-    utils::StorageType default_storage_type;
-    utils::GPUMemoryLayout default_memory_layout;
-    std::tie(test_dtype, default_storage_type, default_memory_layout) = GetParam();
-    config.set_storage_type_override(default_storage_type);
-    config.set_memory_layout_override(default_memory_layout);
-    graph = new ComputeGraph(config);
-
-    if (test_dtype == at::kHalf) {{
-      rtol = 1e-2;
-      atol = 1e-2;
-    }}
-  }}
-
-  void TearDown() override {{
-    delete graph;
-    graph = nullptr;
-  }}
-
-  {check_fn}
-}};
-"""
-
-
-class VkTestSuiteGen(TestSuiteGen):
-    def __init__(self, op_reg_name: str, f: NativeFunction, inputs: VkTestSuite):
-        super().__init__(f, inputs)
-        self.op_reg_name = op_reg_name
-        self.generator = ComputeGraphGen(self.op_reg_name, self.f, self.suite_def)
-
-    def generate_fixture_cpp(self) -> str:
-        check_fn = ""
-        if not self.suite_def.requires_prepack:
-            check_fn = self.generator.gen_op_check_fn()
-
-        prepacked_check_fn = ""
-        if self.suite_def.supports_prepack():
-            self.generator.should_prepack = True
-            prepacked_check_fn = self.generator.gen_op_check_fn()
-            check_fn += "\n\n  "
-            check_fn += prepacked_check_fn
-
-        return test_fixture_template.format(
-            op_name=self.op_name,
-            check_fn=check_fn,
-            rtol=self.suite_def.rtol,
-            atol=self.suite_def.atol,
-        )
-
-    def gen_parameterization(self) -> str:
-        dtypes = self.suite_def.dtypes
-        storage_types = self.suite_def.storage_types
-        layouts = self.suite_def.layouts
-
-        return f"""
-INSTANTIATE_TEST_SUITE_P(
-  Combos_{self.op_name},
-  GeneratedOpsTest_{self.op_name},
-    ::testing::Combine(
-      ::testing::Values({', '.join(dtypes)}),
-      ::testing::Values({', '.join(storage_types)}),
-      ::testing::Values({', '.join(layouts)})));
-        """
-
-
-##############################
-## Test File Code Generation ##
-###############################
-
-preamble_str = """
-#include <executorch/backends/vulkan/runtime/api/api.h>
-#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
-#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
-
-#include <tuple>
-
-using namespace vkcompute;
-using TensorOptions = at::TensorOptions;
-
-vkapi::ScalarType from_at_scalartype(c10::ScalarType at_scalartype) {
-  switch (at_scalartype) {
-    case c10::kFloat:
-      return vkapi::kFloat;
-    case c10::kHalf:
-      return vkapi::kHalf;
-    case c10::kInt:
-      return vkapi::kInt;
-    case c10::kLong:
-      return vkapi::kInt;
-    case c10::kChar:
-      return vkapi::kChar;
-    default:
-      VK_THROW("Unsupported at::ScalarType!");
-  }
-}
-
-#ifdef USE_VULKAN_FP16_INFERENCE
-bool check_close(at::Tensor& t1, at::Tensor& t2, float rtol=1e-2, float atol=1e-2) {
-#else
-bool check_close(at::Tensor& t1, at::Tensor& t2, float rtol=1e-5, float atol=1e-5) {
-#endif
-  // Skip checking index tensors
-  if (t1.scalar_type() == at::kLong || t2.scalar_type() == at::kLong) {
-    return true;
-  }
-  bool is_close = at::allclose(t1, t2, rtol, atol);
-  if (!is_close && t1.numel() < 500) {
-    std::cout << "reference: " << std::endl;
-    print(t1, 150);
-    std::cout << std::endl;
-    std::cout << "vulkan: " << std::endl;
-    print(t2, 150);
-    std::cout << std::endl;
-  }
-  return is_close;
-}
-"""
-
-
-class VkCppTestFileGen(CppTestFileGen):
-    def __init__(self, out_path: str):
-        super().__init__(out_path)
-
-    def generate_preamble(self) -> str:
-        return preamble_str
-
-    def add_suite(self, op_reg_name: str, f: NativeFunction, all_input_cases) -> None:
-        suites_gen = VkTestSuiteGen(op_reg_name, f, all_input_cases)
-        self.suites_gens.append(suites_gen)