Generate selective_mobile_ops.h file for executorch (pytorch#1021)

Summary:
Pull Request resolved: pytorch#1021

Generate `selective_mobile_ops_executorch.h`

Test using an expecttest. Check cases where we have multiple dtypes and default case (return true).

Reviewed By: JacobSzwejbka

Differential Revision: D49236461

fbshipit-source-id: 8e229313cc0d3e7175d99b42b9f036203a7808ab

codegen/tools/gen_selected_mobile_ops.py

@@ -0,0 +1,98 @@
+#!/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 os
+
+import yaml
+
+from torchgen.code_template import CodeTemplate
+
+
+ops_and_dtypes_template_str = """((string_view(operator_name).compare("$operator_name") == 0)\n        && ($dtype_checks))"""
+ops_and_dtypes_template = CodeTemplate(ops_and_dtypes_template_str)
+
+selected_kernel_dtypes_h_template_str = """#pragma once
+/**
+ * Generated by executorch/codegen/tools/gen_selected_mobile_ops_header.py
+ */
+
+inline constexpr bool should_include_kernel_dtype(
+  const char *operator_name,
+  exec_aten::ScalarType scalar_type
+) {
+  return $body;
+}
+"""
+selected_kernel_dtypes_h_template = CodeTemplate(selected_kernel_dtypes_h_template_str)
+
+# enum from: https://github.com/pytorch/executorch/blob/main/runtime/core/portable_type/scalar_type.h
+dtype_enum_to_type = {
+    "0": "Byte",
+    "1": "Char",
+    "2": "Short",
+    "3": "Int",
+    "4": "Long",
+    "5": "Half",
+    "6": "Float",
+    "7": "Double",
+    "8": "ComplexHalf",
+    "9": "ComplexFloat",
+    "10": "ComplexDouble",
+    "11": "Bool",
+    "12": "QInt8",
+    "13": "QUInt8",
+    "14": "QInt32",
+    "15": "BFloat16",
+    "16": "QUInt4x2",
+    "17": "QUInt2x4",
+    "18": "Bits1x8",
+    "19": "Bits2x4",
+    "20": "Bits4x2",
+    "21": "Bits8",
+    "22": "Bits16",
+}
+
+
+def write_selected_mobile_ops(output_file_path: str) -> None:
+    with open(
+        os.path.join(output_file_path, "selected_operators.yaml"), "r"
+    ) as selected_operators_file:
+        # Collect et_kernel_metadata from selected_operators.yaml and extract dtypes
+        # Example format: v1/6;0,1|6;0,1|6;0,1|6;0,1  # Float, 0, 1
+        selected_operators_dict = yaml.safe_load(selected_operators_file)
+        et_kernel_metadata = selected_operators_dict.get("et_kernel_metadata", {})
+        assert isinstance(et_kernel_metadata, dict)
+        body = "return true;"
+        body_parts = []
+        for operator_name, kernel_metadata_str in et_kernel_metadata.items():
+            tensor_meta = []
+            for kernel_metadata in kernel_metadata_str:
+                if kernel_metadata == "default":
+                    break
+                else:
+                    x = kernel_metadata.split("/")[1]
+                    tensor_meta.extend(x.split("|"))
+            conditions = ["true"]
+            if len(tensor_meta) > 0:
+                dtype_set = set([x.split(";")[0] for x in tensor_meta])
+                dtype_list = sorted([dtype_enum_to_type[x] for x in dtype_set])
+                conditions = [
+                    "scalar_type == exec_aten::ScalarType::" + x for x in dtype_list
+                ]
+            body_parts.append(
+                ops_and_dtypes_template.substitute(
+                    operator_name=operator_name.removeprefix("aten::"),
+                    dtype_checks=" || ".join(conditions),
+                ),
+            )
+            body = "\n || ".join(body_parts)
+            header_contents = selected_kernel_dtypes_h_template.substitute(body=body)
+            selected_mobile_ops_path = os.path.join(
+                output_file_path, "selected_mobile_ops.h"
+            )
+            with open(selected_mobile_ops_path, "wb") as out_file:
+                out_file.write(header_contents.encode("utf-8"))

codegen/tools/targets.bzl

@@ -118,6 +118,30 @@ def define_common_targets(is_fbcode = False):
         _is_external_target = True,
     )
 
+    runtime.python_library(
+        name = "gen_selected_mobile_ops_lib",
+        srcs = ["gen_selected_mobile_ops.py"],
+        base_module = "executorch.codegen.tools",
+        visibility = ["//executorch/..."],
+    )
+
+    runtime.python_test(
+        name = "test_gen_selected_mobile_ops",
+        srcs = [
+            "test/test_gen_selected_mobile_ops.py",
+        ],
+        package_style = "inplace",
+        visibility = [
+            "PUBLIC",
+        ],
+        deps = [
+            ":gen_selected_mobile_ops_lib",
+            ":gen_oplist_lib",
+            "fbsource//third-party/pypi/expecttest:expecttest",
+        ],
+        _is_external_target = True,
+    )
+
     # TODO(larryliu0820): This is a hack to only run these two on fbcode. These targets depends on exir which is only available in fbcode.
     if not runtime.is_oss and is_fbcode:
         runtime.python_binary(

codegen/tools/test/test_gen_selected_mobile_ops.py

@@ -0,0 +1,80 @@
+#!/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 os
+import tempfile
+
+import executorch.codegen.tools.gen_selected_mobile_ops as gen_selected_mobile_ops
+import expecttest
+
+
+class TestGenSelectedMobileOpsHeader(expecttest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.selected_ops_yaml = os.path.join(
+            self.temp_dir.name, "selected_operators.yaml"
+        )
+        with open(self.selected_ops_yaml, "w") as f:
+            f.write(
+                """
+include_all_non_op_selectives: False
+include_all_operators: False
+debug_info:
+  - model1@v100
+  - model2@v50
+operators:
+  aten::add:
+    is_root_operator: Yes
+    is_used_for_training: Yes
+    include_all_overloads: No
+  aten::add.int:
+    is_root_operator: No
+    is_used_for_training: No
+    include_all_overloads: Yes
+kernel_metadata: {}
+et_kernel_metadata:
+  aten::add.out:
+    # A list of different kernel keys (tensors with dtype-enum/dim-order) combinations used in model
+      - v1/6;0,1|6;0,1|6;0,1|6;0,1  # Float, 0, 1
+      - v1/3;0,1|3;0,1|3;0,1|3;0,1  # Int, 0, 1
+  aten::mul.out:
+      - v1/6;0,1|6;0,1|6;0,1|6;0,1  # Float, 0, 1
+  aten::sub.out:
+      - default
+build_features: []
+custom_classes: []
+            """
+            )
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_generates_correct_header(self) -> None:
+        gen_selected_mobile_ops.write_selected_mobile_ops(self.temp_dir.name)
+        with open(
+            os.path.join(self.temp_dir.name, "selected_mobile_ops.h"), "r"
+        ) as result:
+            self.assertExpectedInline(
+                result.read(),
+                """#pragma once
+/**
+ * Generated by executorch/codegen/tools/gen_selected_mobile_ops_header.py
+ */
+
+inline constexpr bool should_include_kernel_dtype(
+  const char *operator_name,
+  exec_aten::ScalarType scalar_type
+) {
+  return ((string_view(operator_name).compare("add.out") == 0)
+        && (scalar_type == exec_aten::ScalarType::Float || scalar_type == exec_aten::ScalarType::Int))
+ || ((string_view(operator_name).compare("mul.out") == 0)
+        && (scalar_type == exec_aten::ScalarType::Float))
+ || ((string_view(operator_name).compare("sub.out") == 0)
+        && (true));
+}
+""",
+            )