dtype selective build, link codegen and scalar_type_util (#1090)

Summary:
Pull Request resolved: #1090

1. Let codegen generate `selected_mobile_ops.h` based on `selected_operators.yaml`
2. Include header in `scalar_type_utils`. Add to visibility for codegen.
3. Create test for dtype selective build in scalar_utils.
4. Add buck config for dtype selective build.

Dependencies:
1. Make scalar_type_utils depend on codegen headers (header lib only) to access `selected_mobile_ops.h`.
2. Split deps into `deps` and `kernel_deps`. Add  `kernel_deps` to executorch_generated_library. There is a circular dependency if we leave it as is, where:
- scalar_type_utils -> executorch_generated_lib (added dep)
- executorch_generated_lib -> portable_ops (portable ops are a dependency to the genrule for `selected_operators.yaml`, but they do not need to be)
- portable_ops -> kernel impls (eg. `op_add.cpp` etc.)
- kernel impls -> scalar_type_utils (for macros)

User flow:
1. User passes in the required executorch_generated_lib through buck config
2. See example in selective build examples.

To do: include another dtype selection example when we can use model.pte for selective build.

Reviewed By: larryliu0820

Differential Revision: D50561135

fbshipit-source-id: c52b6318bc1909c9c73a0ba231ecc2a6f01d1d22

Author: lucylq

codegen/targets.bzl

@@ -29,6 +29,7 @@ def define_common_targets():
         visibility = [
             "//executorch/runtime/kernel/...",
             "//executorch/kernels/...",
+            "//executorch/runtime/core/exec_aten/...",
             "@EXECUTORCH_CLIENTS",
         ],
     )

codegen/tools/gen_selected_mobile_ops.py

@@ -5,7 +5,10 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import argparse
 import os
+import sys
+from typing import Any, List
 
 import yaml
 
@@ -57,16 +60,14 @@
 }
 
 
-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:
+def write_selected_mobile_ops(yaml_file_path: str, output_dir: str) -> None:
+    with open(yaml_file_path, "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 = "true"
         body_parts = []
         for operator_name, kernel_metadata_str in et_kernel_metadata.items():
             tensor_meta = []
@@ -85,14 +86,38 @@ def write_selected_mobile_ops(output_file_path: str) -> None:
                 ]
             body_parts.append(
                 ops_and_dtypes_template.substitute(
-                    operator_name=operator_name.removeprefix("aten::"),
+                    operator_name=operator_name.replace("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"))
+        header_contents = selected_kernel_dtypes_h_template.substitute(body=body)
+        selected_mobile_ops_path = os.path.join(output_dir, "selected_mobile_ops.h")
+        with open(selected_mobile_ops_path, "wb") as out_file:
+            out_file.write(header_contents.encode("utf-8"))
+
+
+def main(argv: List[Any]) -> None:
+    parser = argparse.ArgumentParser(description="Generate operator lists")
+    parser.add_argument(
+        "--yaml-file-path",
+        "--yaml_file_path",
+        help=("The directory where selected_operators.yaml was generated)"),
+        required=True,
+    )
+    parser.add_argument(
+        "--output-dir",
+        "--output_dir",
+        help=(
+            "The directory to store the output yaml files (selected_mobile_ops.h, "
+            + "selected_kernel_dtypes.h, selected_operators.yaml)"
+        ),
+        required=True,
+    )
+
+    options = parser.parse_args(argv)
+    write_selected_mobile_ops(options.yaml_file_path, options.output_dir)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

codegen/tools/targets.bzl

@@ -123,6 +123,22 @@ def define_common_targets(is_fbcode = False):
         srcs = ["gen_selected_mobile_ops.py"],
         base_module = "executorch.codegen.tools",
         visibility = ["//executorch/..."],
+        external_deps = [
+            "gen-oplist-lib",
+        ],
+    )
+
+    runtime.python_binary(
+        name = "gen_selected_mobile_ops",
+        main_module = "executorch.codegen.tools.gen_selected_mobile_ops",
+        package_style = "inplace",
+        visibility = [
+            "PUBLIC",
+        ],
+        deps = [
+            ":gen_selected_mobile_ops_lib",
+        ],
+        _is_external_target = True,
     )
 
     runtime.python_test(
@@ -136,7 +152,6 @@ def define_common_targets(is_fbcode = False):
         ],
         deps = [
             ":gen_selected_mobile_ops_lib",
-            ":gen_oplist_lib",
             "fbsource//third-party/pypi/expecttest:expecttest",
         ],
         _is_external_target = True,

codegen/tools/test/test_gen_selected_mobile_ops.py

@@ -15,6 +15,7 @@
 class TestGenSelectedMobileOpsHeader(expecttest.TestCase):
     def setUp(self):
         self.temp_dir = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp_dir.cleanup)
         self.selected_ops_yaml = os.path.join(
             self.temp_dir.name, "selected_operators.yaml"
         )
@@ -54,7 +55,10 @@ 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)
+        gen_selected_mobile_ops.write_selected_mobile_ops(
+            os.path.join(self.temp_dir.name, "selected_operators.yaml"),
+            self.temp_dir.name,
+        )
         with open(
             os.path.join(self.temp_dir.name, "selected_mobile_ops.h"), "r"
         ) as result:
@@ -78,3 +82,51 @@ def test_generates_correct_header(self) -> None:
 }
 """,
             )
+
+
+class TestGenSelectedMobileOpsHeader_Empty(expecttest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp_dir.cleanup)
+        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(
+                """
+build_features: []
+custom_classes: []
+et_kernel_metadata: {}
+include_all_non_op_selectives: false
+include_all_operators: true
+kernel_metadata: {}
+operators: {}
+                """
+            )
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_generates_correct_header(self) -> None:
+        gen_selected_mobile_ops.write_selected_mobile_ops(
+            os.path.join(self.temp_dir.name, "selected_operators.yaml"),
+            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 true;
+}
+""",
+            )

examples/selective_build/README.md

@@ -21,6 +21,7 @@ Check out `targets.bzl` for demo of 3 selective build APIs:
 
 Other configs:
 - `--config executorch.max_kernel_num=N`: Only allocate memory for the required number of operators. Take this result from `selected_operators.yaml`.
+- `--config executorch.dtype_selective_build_lib=<executorch_generated_lib_name>`: Use dtype selective build. For each op, we register the dtypes that are used. Eg. if the model only uses the float implementation of add, then only the float add will be registered. Pass in the executorch_generated_lib name (see buck2 list example).
 
 ## CMake examples
 

examples/selective_build/targets.bzl

@@ -17,8 +17,10 @@ def define_common_targets():
     executorch_generated_lib(
         name = "select_all_lib",
         functions_yaml_target = "//executorch/kernels/portable:functions.yaml",
-        deps = [
+        kernel_deps = [
             "//executorch/kernels/portable:operators",
+        ],
+        deps = [
             ":select_all_ops",
         ],
     )
@@ -35,10 +37,13 @@ def define_common_targets():
     executorch_generated_lib(
         name = "select_ops_in_list_lib",
         functions_yaml_target = "//executorch/kernels/portable:functions.yaml",
-        deps = [
+        kernel_deps = [
             "//executorch/kernels/portable:operators",
+        ],
+        deps = [
             ":select_ops_in_list",
         ],
+        visibility = ["//executorch/runtime/core/..."],
     )
 
     # Select all ops from a yaml file
@@ -50,9 +55,11 @@ def define_common_targets():
     executorch_generated_lib(
         name = "select_ops_from_yaml_lib",
         custom_ops_yaml_target = "//executorch/examples/portable/custom_ops:custom_ops.yaml",
-        deps = [
+        kernel_deps = [
             "//executorch/examples/portable/custom_ops:custom_ops_1",
             "//executorch/examples/portable/custom_ops:custom_ops_2",
+        ],
+        deps = [
             ":select_ops_from_yaml",
         ],
     )

examples/selective_build/test_selective_build.sh

@@ -35,7 +35,9 @@ test_buck2_select_ops_in_list() {
     # set max_kernel_num=17: 14 primops, add, mul
     $BUCK run //examples/selective_build:selective_build_test \
         --config=executorch.max_kernel_num=17 \
-        --config=executorch.select_ops=list -- --model_path=./add_mul.pte
+        --config=executorch.select_ops=list \
+        --config=executorch.dtype_selective_build_lib="//examples/selective_build:select_ops_in_list_lib" \
+        -- --model_path=./add_mul.pte
 
     echo "Removing add_mul.pte"
     rm "./add_mul.pte"

runtime/core/exec_aten/util/targets.bzl

@@ -1,4 +1,5 @@
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+load("@fbsource//xplat/executorch/codegen:codegen.bzl", "et_operator_library", "executorch_generated_lib")
 
 def define_common_targets():
     """Defines targets that should be shared between fbcode and xplat.
@@ -7,9 +8,68 @@ def define_common_targets():
     TARGETS and BUCK files that call this function.
     """
 
+    et_operator_library(
+        name = "two_ops",
+        ops = [
+            "aten::add.out",
+            "aten::mm.out",
+        ],
+    )
+
+    executorch_generated_lib(
+        name = "two_ops_lib",
+        functions_yaml_target = "//executorch/kernels/portable:functions.yaml",
+        kernel_deps = ["//executorch/kernels/portable:operators"],
+        deps = [":two_ops"],
+    )
+
+    runtime.cxx_library(
+        name = "scalar_type_util_TEST_ONLY",
+        srcs = [],
+        exported_headers = [
+            "scalar_type_util.h",
+        ],
+        visibility = [
+            "//executorch/...",
+            "@EXECUTORCH_CLIENTS",
+        ],
+        exported_preprocessor_flags = ["-DEXECUTORCH_SELECTIVE_BUILD_DTYPE"],
+        exported_deps = [
+            "//executorch/runtime/core/portable_type:scalar_type",
+            ":two_ops_lib_headers",
+        ],
+    )
+
+    dtype_selective_build_lib = native.read_config("executorch", "dtype_selective_build_lib", None)
+    if dtype_selective_build_lib != None:
+        # retrieve selected_mobile_ops.h from codegen
+        genrule_name = dtype_selective_build_lib + "_et_op_dtype_gen[selected_mobile_ops]"
+        runtime.cxx_library(
+            name = "dtype_headers",
+            srcs = [],
+            exported_headers = {
+                "selected_mobile_ops.h": genrule_name,
+            },
+            visibility = [
+                "//executorch/...",
+                "@EXECUTORCH_CLIENTS",
+            ],
+        )
+
     for aten_mode in (True, False):
         aten_suffix = "_aten" if aten_mode else ""
 
+        exported_preprocessor_flags_ = []
+        exported_deps_ = ["//executorch/runtime/core:core"]
+        if aten_mode:
+            exported_preprocessor_flags_ += ["-DUSE_ATEN_LIB"]
+        else:
+            exported_deps_ += ["//executorch/runtime/core/portable_type:scalar_type"]
+
+        if dtype_selective_build_lib != None:
+            exported_preprocessor_flags_ += ["-DEXECUTORCH_SELECTIVE_BUILD_DTYPE"]
+            exported_deps_ += [":dtype_headers"]
+
         runtime.cxx_library(
             name = "scalar_type_util" + aten_suffix,
             srcs = [],
@@ -20,10 +80,8 @@ def define_common_targets():
                 "//executorch/...",
                 "@EXECUTORCH_CLIENTS",
             ],
-            exported_preprocessor_flags = ["-DUSE_ATEN_LIB"] if aten_mode else [],
-            exported_deps = [
-                "//executorch/runtime/core:core",
-            ] + [] if aten_mode else ["//executorch/runtime/core/portable_type:scalar_type"],
+            exported_preprocessor_flags = exported_preprocessor_flags_,
+            exported_deps = exported_deps_,
             exported_external_deps = ["libtorch"] if aten_mode else [],
         )
 

runtime/core/exec_aten/util/test/selected_mobile_ops_header_test.cpp

@@ -0,0 +1,53 @@
+/*
+ * 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.
+ */
+
+#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/test/utils/DeathTest.h>
+#include <gtest/gtest.h>
+
+using namespace ::testing;
+using exec_aten::ScalarType;
+using torch::executor::ScalarTypeToCppType;
+
+// Add test for specific dtypes when we can run on model file.
+
+TEST(SelectedMobileOpsHeaderTest, UnknownOp) {
+  ET_EXPECT_DEATH(
+      ET_SWITCH_TWO_TYPES(
+          Float,
+          Int,
+          exec_aten::ScalarType::Float,
+          ctx,
+          "addmm.out",
+          CTYPE_OUT,
+          [&] { return true; }),
+      "");
+}
+
+TEST(SelectedMobileOpsHeaderTest, OpWithDtype) {
+  ASSERT_EQ(
+      ET_SWITCH_TWO_TYPES(
+          Float,
+          Int,
+          exec_aten::ScalarType::Float,
+          ctx,
+          "add.out",
+          CTYPE_OUT,
+          [&] { return true; }),
+      true);
+  ASSERT_EQ(
+      ET_SWITCH_TWO_TYPES(
+          Float,
+          Int,
+          exec_aten::ScalarType::Int,
+          ctx,
+          "mm.out",
+          CTYPE_OUT,
+          [&] { return true; }),
+      true);
+}