Dependency refactor, add tests (pytorch#1165)

Summary:
Pull Request resolved: pytorch#1165

Refactor dependency chain following:
https://docs.google.com/document/d/1vuU9UdQA5bq3WU7YnEgJowq9ey6vkF3KUKA0rpIxxIY/edit#heading=h.8raqyft9y50

**New dep chain:**
op_<op>.cpp -> selected_op_variants.h
selected_op_variants.h -> scalar_type_util

**Old dep chain:**
op_<op>.cpp -> scalar_type_util
scalar_type_util -> selected_op_variants.h

**Testing**
Add e2e test with addmul model
Update unit test to use addmul model

Gated to only run with fbcode, as there is dependency on exir which is only in fbcode.

Reviewed By: larryliu0820

Differential Revision: D51036567

fbshipit-source-id: ba76b3c61fc9c044232b0ae66f4dba788254bfab

Author: lucylq

codegen/targets.bzl

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

examples/portable/scripts/export.py

@@ -26,6 +26,7 @@ def main() -> None:
         required=True,
         help=f"provide a model name. Valid ones: {list(MODEL_NAME_TO_MODEL.keys())}",
     )
+    parser.add_argument("-o", "--output_dir", default=".", help="output directory")
 
     args = parser.parse_args()
 
@@ -40,7 +41,7 @@ def main() -> None:
     )
 
     prog = export_to_exec_prog(model, example_inputs)
-    save_pte_program(prog.buffer, args.model_name)
+    save_pte_program(prog.buffer, args.model_name, args.output_dir)
 
 
 if __name__ == "__main__":

examples/portable/utils.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
+import os
 
 from typing import Tuple, Union
 
@@ -78,8 +79,8 @@ def export_to_exec_prog(
     return exec_prog
 
 
-def save_pte_program(buffer: bytes, model_name: str) -> None:
-    filename = f"{model_name}.pte"
+def save_pte_program(buffer: bytes, model_name: str, output_dir: str = "") -> None:
+    filename = os.path.join(output_dir, f"{model_name}.pte")
     try:
         with open(filename, "wb") as file:
             file.write(buffer)

examples/selective_build/README.md

@@ -21,7 +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).
+- `--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 model example in targets.bzl).
 
 ## CMake examples
 

examples/selective_build/targets.bzl

@@ -43,7 +43,6 @@ def define_common_targets():
         deps = [
             ":select_ops_in_list",
         ],
-        visibility = ["//executorch/runtime/core/..."],
     )
 
     # Select all ops from a yaml file
@@ -67,6 +66,28 @@ def define_common_targets():
     # Select all ops from a given model
     # TODO(larryliu0820): Add this
 
+    if not runtime.is_oss:
+        runtime.genrule(
+            name = "add_mul_model",
+            outs = {"add_mul": ["add_mul.pte"]},
+            cmd = "$(exe fbcode//executorch/examples/portable/scripts:export) --model_name add_mul --output_dir $OUT",
+            macros_only = False,
+            visibility = ["//executorch/..."],
+        )
+
+        et_operator_library(
+            name = "select_ops_from_model",
+            model = ":add_mul_model[add_mul]",
+        )
+
+        executorch_generated_lib(
+            name = "select_ops_from_model_lib",
+            functions_yaml_target = "//executorch/kernels/portable:functions.yaml",
+            kernel_deps = ["//executorch/kernels/portable:operators"],
+            deps = [":select_ops_from_model"],
+            visibility = ["//executorch/kernels/..."],
+        )
+
     # ~~~ Test binary for selective build ~~~
     select_ops = native.read_config("executorch", "select_ops", None)
     lib = []
@@ -76,6 +97,8 @@ def define_common_targets():
         lib.append(":select_ops_in_list_lib")
     elif select_ops == "yaml":
         lib.append(":select_ops_from_yaml_lib")
+    elif select_ops == "model":
+        lib.append(":select_ops_from_model_lib")
     runtime.cxx_binary(
         name = "selective_build_test",
         srcs = [],

examples/selective_build/test_selective_build.sh

@@ -47,7 +47,6 @@ test_buck2_select_ops_in_list() {
     $BUCK run //examples/selective_build:selective_build_test \
         --config=executorch.max_kernel_num=17 \
         --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"

kernels/portable/cpu/TARGETS

@@ -3,4 +3,4 @@
 
 load(":targets.bzl", "define_common_targets")
 
-define_common_targets()
+define_common_targets(is_fbcode = True)

kernels/portable/cpu/scalar_utils.h

@@ -12,6 +12,7 @@
 #include <cmath>
 #include <limits>
 
+#include <executorch/kernels/portable/cpu/selective_build.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/core/portable_type/scalar.h>

kernels/portable/cpu/selective_build.h

@@ -0,0 +1,44 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+
+#ifdef EXECUTORCH_SELECTIVE_BUILD_DTYPE
+// include header generated by
+// executorch/codegen/tools/gen_selected_op_variants.py
+#include <executorch/kernels/portable/cpu/selected_op_variants.h>
+#else
+// dummy implementation
+inline constexpr bool should_include_kernel_dtype(
+    const char* /*operator_name*/,
+    exec_aten::ScalarType /*scalar_type*/
+) {
+  return true;
+}
+#endif
+
+namespace torch {
+namespace executor {
+#define ET_INTERNAL_CHECK_SELECTIVE_BUILD(enum_type)               \
+  do {                                                             \
+    if (!should_include_kernel_dtype(et_switch_name, enum_type)) { \
+      ET_LOG(                                                      \
+          Error,                                                   \
+          "dtype '%" PRId8 "' not selected for operator %s",       \
+          static_cast<int8_t>(enum_type),                          \
+          et_switch_name);                                         \
+      torch::executor::runtime_abort();                            \
+    }                                                              \
+  } while (0)
+
+} // namespace executor
+} // namespace torch
+
+#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>

kernels/portable/cpu/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")
 load("@fbsource//xplat/executorch/kernels/portable:op_registration_util.bzl", "define_op_target", "op_target")
 
 # Operators that are listed in `functions.yaml`, and are thus compatible with
@@ -821,7 +822,7 @@ _CUSTOM_OPS = (
     ),
 )
 
-def define_common_targets():
+def define_common_targets(is_fbcode = False):
     """Defines targets that should be shared between fbcode and xplat.
 
     The directory containing this targets.bzl file should also contain both
@@ -878,11 +879,30 @@ def define_common_targets():
         ],
     )
 
+    dtype_selective_build_lib = native.read_config("executorch", "dtype_selective_build_lib", None)
+    if dtype_selective_build_lib != None:
+        # retrieve selected_op_variants.h from codegen
+        genrule_name = dtype_selective_build_lib + "_et_op_dtype_gen[selected_op_variants]"
+        runtime.cxx_library(
+            name = "dtype_headers",
+            srcs = [],
+            exported_headers = {
+                "selected_op_variants.h": genrule_name,
+            },
+            visibility = [
+                "//executorch/...",
+                "@EXECUTORCH_CLIENTS",
+            ],
+        )
+
     # Only for use by targets in this directory.
     runtime.cxx_library(
         name = "scalar_utils",
         srcs = [],
-        exported_headers = ["scalar_utils.h"],
+        # include dtype selective build flag and header
+        exported_preprocessor_flags = ["-DEXECUTORCH_SELECTIVE_BUILD_DTYPE"] if dtype_selective_build_lib != None else [],
+        exported_headers = ["scalar_utils.h", "selective_build.h"],
+        exported_deps = [":dtype_headers"] if dtype_selective_build_lib != None else [],
         visibility = [
             "//executorch/kernels/portable/cpu/...",
             "//executorch/kernels/optimized/cpu/...",
@@ -893,3 +913,45 @@ def define_common_targets():
             "//executorch/runtime/core/exec_aten/util:scalar_type_util",
         ],
     )
+
+    # dtype selective build test artifacts
+    if is_fbcode:
+        et_operator_library(
+            name = "add_model",
+            model = "fbcode//executorch/test/models:exported_programs[ModuleAdd.pte]",
+        )
+
+        executorch_generated_lib(
+            name = "add_model_lib",
+            functions_yaml_target = "//executorch/kernels/portable:functions.yaml",
+            kernel_deps = ["//executorch/kernels/portable:operators"],
+            deps = [":add_model"],
+            visibility = ["//executorch/kernels/..."],
+        )
+
+        runtime.cxx_library(
+            name = "dtype_headers_TEST_ONLY",
+            srcs = [],
+            exported_headers = {
+                "selected_op_variants.h": ":add_model_lib_et_op_dtype_gen[selected_op_variants]",
+            },
+            visibility = [
+                "//executorch/...",
+                "@EXECUTORCH_CLIENTS",
+            ],
+        )
+
+        runtime.cxx_library(
+            name = "scalar_utils_TEST_ONLY",
+            srcs = [],
+            exported_preprocessor_flags = ["-DEXECUTORCH_SELECTIVE_BUILD_DTYPE"],
+            exported_headers = ["scalar_utils.h", "selective_build.h"],
+            exported_deps = [":dtype_headers_TEST_ONLY"],
+            visibility = [
+                "//executorch/kernels/...",
+                "@EXECUTORCH_CLIENTS",
+            ],
+            deps = [
+                "//executorch/runtime/core/exec_aten/util:scalar_type_util",
+            ],
+        )

kernels/portable/targets.bzl

@@ -97,8 +97,8 @@ def define_common_targets():
         deps = [
             ":executorch_aten_ops",
             ":executorch_custom_ops",
-            "//executorch/kernels/portable:operators",
         ],
+        kernel_deps = ["//executorch/kernels/portable:operators"],
         **generated_lib_common_args
     )
 

kernels/test/TARGETS

@@ -7,7 +7,7 @@ load(":targets.bzl", "define_common_targets")
 
 oncall("executorch")
 
-define_common_targets()
+define_common_targets(is_fbcode = True)
 
 python_unittest(
     name = "gen_supported_features_test",

runtime/core/exec_aten/util/test/selected_op_variants_test.cpp renamed to kernels/test/dtype_selective_build_test.cpp

@@ -6,47 +6,48 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/kernels/portable/cpu/scalar_utils.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) {
+TEST(DtypeSelectiveBuildTest, UnknownOp) {
   ET_EXPECT_DEATH(
       ET_SWITCH_TWO_TYPES(
           Float,
           Int,
           exec_aten::ScalarType::Float,
           ctx,
-          "addmm.out",
+          "unknown.out",
           CTYPE_OUT,
           [&] { return true; }),
       "");
 }
 
-TEST(SelectedMobileOpsHeaderTest, OpWithDtype) {
-  ASSERT_EQ(
+TEST(DtypeSelectiveBuildTest, OpWithoutDtype) {
+  ET_EXPECT_DEATH(
       ET_SWITCH_TWO_TYPES(
           Float,
           Int,
-          exec_aten::ScalarType::Float,
+          exec_aten::ScalarType::Int,
           ctx,
           "add.out",
           CTYPE_OUT,
           [&] { return true; }),
-      true);
+      "");
+}
+
+TEST(DtypeSelectiveBuildTest, OpWithDtype) {
   ASSERT_EQ(
       ET_SWITCH_TWO_TYPES(
           Float,
           Int,
-          exec_aten::ScalarType::Int,
+          exec_aten::ScalarType::Float,
           ctx,
-          "mm.out",
+          "add.out",
           CTYPE_OUT,
           [&] { return true; }),
       true);

kernels/test/targets.bzl

@@ -28,7 +28,7 @@ def make_example_generated_op_test_target():
             "//executorch/kernels/test:function_header_wrapper_portable",
         )
 
-def define_common_targets():
+def define_common_targets(is_fbcode = False):
     """Defines targets that should be shared between fbcode and xplat.
 
     The directory containing this targets.bzl file should also contain both
@@ -280,3 +280,15 @@ def define_common_targets():
     _common_op_test("op_zeros_test", ["aten", "portable"])
 
     make_example_generated_op_test_target()
+
+    # dtype selective build test
+    if is_fbcode:
+        runtime.cxx_test(
+            name = "dtype_selective_build_test",
+            srcs = ["dtype_selective_build_test.cpp"],
+            deps = [
+                "//executorch/kernels/portable/cpu:scalar_utils_TEST_ONLY",
+                "//executorch/runtime/core/exec_aten:lib",
+            ],
+            preprocessor_flags = ["-DEXECUTORCH_SELECTIVE_BUILD_DTYPE"],
+        )