Generate 1 kernel libary with dtype selective build and used it in multiple executorch_generated_lib

Differential Revision: D75029407

Pull Request resolved: pytorch#11378

Author: larryliu0820

kernels/portable/cpu/util/targets.bzl

@@ -132,7 +132,7 @@ def define_common_targets():
             "//executorch/runtime/core/exec_aten/util:tensor_shape_to_c_string",
             "//executorch/runtime/kernel:kernel_includes",
         ],
-        visibility = ["//executorch/kernels/portable/cpu/...", "//executorch/kernels/optimized/cpu/..."],
+        visibility = ["//executorch/kernels/portable/cpu/...", "//executorch/kernels/optimized/cpu/...", "@EXECUTORCH_CLIENTS"],
     )
 
     runtime.cxx_library(

shim_et/xplat/executorch/codegen/codegen.bzl

@@ -444,10 +444,53 @@ def build_portable_header_lib(name, oplist_header_name, feature = None):
         feature = feature,
     )
 
-def build_portable_lib(name, oplist_header_name, portable_header_lib, feature = None, expose_operator_symbols = False):
-    """Build portable lib from source. We build from source so that the generated header file,
-    selected_op_variants.h, can be used to selectively build the lib for different dtypes.
+def build_portable_lib(
+    name,
+    et_operator_lib_deps = [],
+    oplist_header_name = None,
+    portable_header_lib = None,
+    feature = None,
+    expose_operator_symbols = False,
+    visibility = ["@EXECUTORCH_CLIENTS"]):
     """
+    WARNING: Before using this, please consider using executorch_generated_lib instead. This
+    function is only for special cases where you need to build a portable kernel library with
+    dtype selective build enabled and also wants to share it across more than one executorch_generated_lib.
+    Any other use case is likely wrong and you should use executorch_generated_lib instead.
+
+    Create a new portable kernel library based on `portable_header_lib`. `portable_header_lib`
+    should contain the header `selected_op_variants.h` generated by `dtype_header_genrule`.
+
+    Notice that this is giving a library that is different than //executorch/kernels/portable/cpu:cpu,
+    because of the generated header `selected_op_variants.h`. The original portable kernel library
+    doesn't have that header and thus include all the dtypes possible.
+
+    If no `portable_header_lib` is provided, try to create one based on the deps. In this case
+    we require `deps` to be present. Notice that this way we are always enabling dtype selective
+    build.
+
+    Args:
+        name: name of the new portable kernel library.
+        et_operator_lib_deps: list of deps to use to create the portable header library.
+        oplist_header_name: the name of the header genrule (dtype_header_genrule)
+        portable_header_lib: the name of the header library (build_portable_header_lib)
+        feature: feature to use for the new portable kernel library.
+        expose_operator_symbols: expose operator symbols to library users. This only works in xplat.
+        visibility: visibility of the new portable kernel library.
+    """
+
+    if not portable_header_lib:
+        if not oplist_header_name:
+            if not et_operator_lib_deps:
+                fail("Either et_operator_lib_deps or oplist_header_name must be provided.")
+            oplist_header_name = name + "_header"
+            dtype_header_genrule(
+                name = oplist_header_name,
+                deps = et_operator_lib_deps,
+                visibility = visibility,
+            )
+        portable_header_lib = name + "_portable_header_lib"
+        build_portable_header_lib(portable_header_lib, oplist_header_name, feature)
 
     # Copy portable cpp files.
     portable_source_files = []
@@ -546,6 +589,72 @@ def build_optimized_lib(name, oplist_header_name, portable_header_lib, feature =
         feature = feature,
     )
 
+def selected_operators_genrule(
+    name,
+    deps,
+    platforms = get_default_executorch_platforms(),
+):
+    """Generates selected_operators.yaml from the list of deps. We look into the trasitive closure of all the deps,
+    and look for macros `et_operator_library`.
+
+    `gen_all_oplist` is the python binary we use to aggregate all the `et_operator_library`s into single
+    `selected_operators.yaml` file.
+
+    This file can be furthur used to generate `selected_op_variants.h` (see dtype_header_genrule) for dtype
+    selective build work.
+    """
+    runtime.genrule(
+        name = name,
+        macros_only = False,
+        cmd = ("$(exe fbsource//xplat/executorch/codegen/tools:gen_all_oplist) " +
+               "--model_file_list_path $(@query_outputs \'attrfilter(labels, et_operator_library, deps(set({deps})))\') " +
+               "--allow_include_all_overloads " +
+               "--output_dir $OUT ").format(deps = " ".join(["\"{}\"".format(d) for d in deps])),
+        outs = {"selected_operators.yaml": ["selected_operators.yaml"]},
+        default_outs = ["."],
+        platforms = platforms,
+    )
+
+def dtype_header_genrule(
+    name,
+    visibility,
+    deps = [],
+    selected_operators_genrule_name = None,
+    platforms = get_default_executorch_platforms(),
+):
+    """Generate selected_op_variants.h from selected_operators.yaml.
+
+    Given a `selected_operators.yaml` (passed in as selected_operators_genrule_name), we should be able to determine
+    what dtypes to be enabled for kernels in the kernel library. For example, `add.out` kernel needs to support
+    both float16 and float32 etc.
+
+    This information is recorded in `selected_op_variants.h` and it should be used to compile a new kernel library.
+
+    Notice that until this stage we are kernel library agnostic, meaning the header should be applicable to any
+    kernel library that includes it.
+    """
+    if not selected_operators_genrule_name:
+        if not deps:
+            fail("Either deps or selected_operators_genrule_name must be provided.")
+        selected_operators_genrule_name = name + "_selected_operators"
+        selected_operators_genrule(
+            name = selected_operators_genrule_name,
+            deps = deps,
+        )
+
+    runtime.genrule(
+        name = name,
+        macros_only = False,
+        cmd = ("$(exe //executorch/codegen/tools:gen_selected_op_variants) " +
+               "--yaml_file_path $(location :{}[selected_operators.yaml]) " +
+               "--output_dir $OUT").format(selected_operators_genrule_name),
+        outs = {"selected_op_variants": ["selected_op_variants.h"]},
+        default_outs = ["."],
+        platforms = platforms,
+        visibility = visibility,
+        _is_external_target = True,
+    )
+
 def executorch_generated_lib(
         name,
         functions_yaml_target = None,
@@ -721,32 +830,11 @@ def executorch_generated_lib(
 
     # genrule for selective build from static operator list
     oplist_dir_name = name + "_et_oplist"
-    runtime.genrule(
-        name = oplist_dir_name,
-        macros_only = False,
-        cmd = ("$(exe fbsource//xplat/executorch/codegen/tools:gen_all_oplist) " +
-               "--model_file_list_path $(@query_outputs \'attrfilter(labels, et_operator_library, deps(set({deps})))\') " +
-               "--allow_include_all_overloads " +
-               "--output_dir $OUT ").format(deps = " ".join(["\"{}\"".format(d) for d in deps])),
-        outs = {"selected_operators.yaml": ["selected_operators.yaml"]},
-        default_outs = ["."],
-        platforms = platforms,
-    )
+    selected_operators_genrule(name = oplist_dir_name, deps = deps, platforms = platforms)
 
     # genrule to generate selected_op_variants.h from selected_operators.yaml above
     oplist_header_name = name + "_et_op_dtype_gen"
-    runtime.genrule(
-        name = oplist_header_name,
-        macros_only = False,
-        cmd = ("$(exe //executorch/codegen/tools:gen_selected_op_variants) " +
-               "--yaml_file_path $(location :{}[selected_operators.yaml]) " +
-               "--output_dir $OUT").format(oplist_dir_name),
-        outs = {"selected_op_variants": ["selected_op_variants.h"]},
-        default_outs = ["."],
-        platforms = platforms,
-        visibility = visibility,
-        _is_external_target = True,
-    )
+    dtype_header_genrule(name = oplist_header_name, selected_operators_genrule_name = oplist_dir_name, platforms = platforms, visibility = visibility)
 
     # codegen genrule(s). For ATen mode we expect two genrules, one for ATen ops one for custom ops.
     for genrule_name in genrules:
@@ -777,7 +865,7 @@ def executorch_generated_lib(
 
             # Build portable lib.
             portable_lib_name = name + "_portable_lib"
-            build_portable_lib(portable_lib_name, oplist_header_name, portable_header_lib, feature, expose_operator_symbols)
+            build_portable_lib(name = portable_lib_name, portable_header_lib = portable_header_lib, feature = feature, expose_operator_symbols = expose_operator_symbols)
             kernel_deps.append(":{}".format(portable_lib_name))
 
         if "//executorch/kernels/optimized:optimized_operators" in kernel_deps: