dtype selective build from model API in OSS (#11760)

Given a specific model, we want to produce a binary that only includes
the minimal operators and dtypes needed to run the model. This requires
parsing the model to determine what kernels it launch, the operators
used in those kernels, and the dtypes of the tensors in the kernels.
After parsing, a header file must be generated, and the portable_kernels
lib can be rebuilt to only include the operators and dtypes specified in
the generated header.

This changes completes this E2E process. A user can now specify the
model they wish to optimize their binary for via the command line
argument `-DEXECUTORCH_SELECT_OPS_FROM_MODEL="<file path to model
pte>"`. When specified, the pte is parsed to produce a YAML file called
`seleced_operators.yaml` which describes the model's operators and
dtypes. From this YAML, a header file called `selected_op_variants.h` is
generated that selects the described operators and dtypes. When command
line argument `-DEXECUTORCH_DTYPE_SELECTIVE_BUILD=ON` is specified, the
header file is linked to the `portable_kernels` lib when it's rebuilt.
Only the model API is supported with dtype selective build, and using
other methods such as `list` or `dict` will results in a build error.

An example usage of this flow is included in
`examples/selective_build/test_selective_build.sh:test_cmake_select_ops_in_model`.
When run as `bash examples/selective_build/test_selective_build.sh
cmake`, the `cmake-out/examples/selective_build/selective_build_test`
binary is built. After stripping the binary, the following binary size
results were seen with the following models:

| Model | Default Binary Size (KB) | Dtype Selected Binary Size (KB) |
| ------- | :---:| :---:|
|add | 359 | 275|
|mul | 335 | 263 |
| add_mul| 367 | 287 |
| linear | 347 | 291 |
| softmax | 251 | 251 |
|resnet18| 643 | 515 |
|resnet50 | 643 | 515 |
|mobilebert| 707 | 415 |
|lstm| 643 | 459|
| dl3| 607 | 539|
|edsr | 543 | 371|

| Model | Default Binary Size (KB) | Dtype Selected Binary Size (KB) |
| ------- | :---:| :---:|
|emformer_transcribe| 863| 555|
|vit| 907| 687 |
|mv2| 631 | 495 |
|mv3| 843 | 535 |
|llama | 1.1M | 827|
|qwen2_5 | 1.1M | 827|

Although there is a noted reduction in the binary size, it seems that
the pte file parsing functionality from `gen_oplist.py` is incomplete.
Please see the discussion on [PR
details.

Author: BujSet

examples/selective_build/CMakeLists.txt

@@ -33,7 +33,7 @@ if(NOT CMAKE_CXX_STANDARD)
   # Can't set to 11 due to executor_runner.cpp make_unique
 endif()
 
-set(_common_compile_options -Wno-deprecated-declarations -fPIC)
+set(_common_compile_options -Wno-deprecated-declarations -fPIC -ffunction-sections -fdata-sections)
 
 # Let files say "include <executorch/path/to/header.h>".
 set(_common_include_directories ${EXECUTORCH_ROOT}/..)
@@ -123,13 +123,25 @@ gen_selected_ops(
 )
 
 generate_bindings_for_kernels(
-  LIB_NAME "select_build_lib" FUNCTIONS_YAML
-  ${EXECUTORCH_ROOT}/kernels/portable/functions.yaml CUSTOM_OPS_YAML
+  LIB_NAME 
+  "select_build_lib" 
+  FUNCTIONS_YAML
+  ${EXECUTORCH_ROOT}/kernels/portable/functions.yaml 
+  CUSTOM_OPS_YAML
   "${_custom_ops_yaml}"
+  DTYPE_SELECTIVE_BUILD
+  "${EXECUTORCH_DTYPE_SELECTIVE_BUILD}"
 )
 
 gen_operators_lib(
-  LIB_NAME "select_build_lib" KERNEL_LIBS ${_kernel_lib} DEPS executorch_core
+  LIB_NAME 
+  "select_build_lib" 
+  KERNEL_LIBS 
+  ${_kernel_lib} 
+  DEPS 
+  executorch_core
+  DTYPE_SELECTIVE_BUILD
+  "${EXECUTORCH_DTYPE_SELECTIVE_BUILD}"
 )
 
 list(TRANSFORM _executor_runner__srcs PREPEND "${EXECUTORCH_ROOT}/")

examples/selective_build/test_selective_build.sh

@@ -162,13 +162,17 @@ test_cmake_select_ops_in_yaml() {
 }
 
 test_cmake_select_ops_in_model() {
-    echo "Exporting MobilenetV2"
-    ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="mv2"
+    local model_name="add_mul"
+    local model_export_name="${model_name}.pte"
+    echo "Exporting ${model_name}"
+    ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="${model_name}"
     local example_dir=examples/selective_build
     local build_dir=cmake-out/${example_dir}
     rm -rf ${build_dir}
-    retry cmake -DCMAKE_BUILD_TYPE=$CMAKE_BUILD_TYPE \
-            -DEXECUTORCH_SELECT_OPS_FROM_MODEL="./mv2.pte" \
+    retry cmake -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \
+            -DEXECUTORCH_SELECT_OPS_FROM_MODEL="./${model_export_name}" \
+            -DEXECUTORCH_DTYPE_SELECTIVE_BUILD=ON \
+            -DEXECUTORCH_OPTIMIZE_SIZE=ON \
             -DCMAKE_INSTALL_PREFIX=cmake-out \
             -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
             -B${build_dir} \
@@ -178,10 +182,10 @@ test_cmake_select_ops_in_model() {
     cmake --build ${build_dir} -j9 --config $CMAKE_BUILD_TYPE
 
     echo 'Running selective build test'
-    ${build_dir}/selective_build_test --model_path="./mv2.pte"
+    ${build_dir}/selective_build_test --model_path="./${model_export_name}"
 
-    echo "Removing mv2.pte"
-    rm "./mv2.pte"
+    echo "Removing ${model_export_name}"
+    rm "./${model_export_name}"
 }
 
 if [[ -z $BUCK ]];

tools/cmake/Codegen.cmake

@@ -22,15 +22,20 @@ function(gen_selected_ops)
   message(STATUS "  INCLUDE_ALL_OPS: ${GEN_INCLUDE_ALL_OPS}")
   message(STATUS "  OPS_FROM_MODEL: ${GEN_OPS_FROM_MODEL}")
   message(STATUS "  DTYPE_SELECTIVE_BUILD: ${GEN_DTYPE_SELECTIVE_BUILD}")
+
+  set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME})
+
   if(GEN_DTYPE_SELECTIVE_BUILD)
-     message(STATUS "  DTYPE_SELECTIVE_BUILD is still WIP and may not be fully functional")
+    if(NOT GEN_OPS_FROM_MODEL)
+      message(FATAL_ERROR "  DTYPE_SELECTIVE_BUILD is only support with model API, please pass in a model")
+    endif()
   endif()
 
   set(_oplist_yaml
-      ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}/selected_operators.yaml
+    ${_out_dir}/selected_operators.yaml
   )
 
-  file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME})
+  file(MAKE_DIRECTORY ${_out_dir})
 
   file(GLOB_RECURSE _codegen_tools_srcs "${EXECUTORCH_ROOT}/codegen/tools/*.py")
 
@@ -64,18 +69,18 @@ function(gen_selected_ops)
 
   if(GEN_DTYPE_SELECTIVE_BUILD)
     set(_opvariant_h
-      ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}/selected_op_variants.h
+      ${_out_dir}/selected_op_variants.h
     )
     set(_gen_opvariant_command "${PYTHON_EXECUTABLE}" -m codegen.tools.gen_selected_op_variants
                           --yaml-file=${_oplist_yaml}
-                          --output-dir=${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}/
+                          --output-dir=${_out_dir}/
     )
     message("Command - ${_gen_opvariant_command}")
     add_custom_command(
-      COMMENT "Generating selected_op_variants.h for ${GEN_LIB_NAME}"
+      COMMENT "Generating ${_opvariant_h} for ${GEN_LIB_NAME}"
       OUTPUT ${_opvariant_h}
       COMMAND ${_gen_opvariant_command}
-      DEPENDS ${_oplist_yaml} ${_codegen_tools_srcs}
+      DEPENDS ${_oplist_yaml} ${GEN_OPS_SCHEMA_YAML} ${_codegen_tools_srcs}
       WORKING_DIRECTORY ${EXECUTORCH_ROOT}
     )
   endif()
@@ -88,14 +93,15 @@ endfunction()
 # functions_yaml CUSTOM_OPS_YAML custom_ops_yaml )
 function(generate_bindings_for_kernels)
   set(options ADD_EXCEPTION_BOUNDARY)
-  set(arg_names LIB_NAME FUNCTIONS_YAML CUSTOM_OPS_YAML)
+  set(arg_names LIB_NAME FUNCTIONS_YAML CUSTOM_OPS_YAML DTYPE_SELECTIVE_BUILD)
   cmake_parse_arguments(GEN "${options}" "${arg_names}" "" ${ARGN})
 
   message(STATUS "Generating kernel bindings:")
   message(STATUS "  LIB_NAME: ${GEN_LIB_NAME}")
   message(STATUS "  FUNCTIONS_YAML: ${GEN_FUNCTIONS_YAML}")
   message(STATUS "  CUSTOM_OPS_YAML: ${GEN_CUSTOM_OPS_YAML}")
   message(STATUS "  ADD_EXCEPTION_BOUNDARY: ${GEN_ADD_EXCEPTION_BOUNDARY}")
+  message(STATUS "  DTYPE_SELECTIVE_BUILD: ${GEN_DTYPE_SELECTIVE_BUILD}")
 
   # Command to generate selected_operators.yaml from custom_ops.yaml.
   file(GLOB_RECURSE _codegen_templates "${EXECUTORCH_ROOT}/codegen/templates/*")
@@ -104,6 +110,13 @@ function(generate_bindings_for_kernels)
   # By default selective build output is selected_operators.yaml
   set(_oplist_yaml ${_out_dir}/selected_operators.yaml)
 
+  # If dtype selective build is enable, force header file to be preserved
+  if(GEN_DTYPE_SELECTIVE_BUILD)
+    set(_opvariant_h ${_out_dir}/selected_op_variants.h)
+  else()
+    set(_opvariant_h "")
+  endif()
+
   # Command to codegen C++ wrappers to register custom ops to both PyTorch and
   # Executorch runtime.
   execute_process(
@@ -148,8 +161,9 @@ function(generate_bindings_for_kernels)
     COMMENT "Generating code for kernel registration"
     OUTPUT ${_gen_command_sources}
     COMMAND ${_gen_command}
-    DEPENDS ${_oplist_yaml} ${GEN_CUSTOM_OPS_YAML} ${GEN_FUNCTIONS_YAML}
-            ${_codegen_templates} ${_torchgen_srcs}
+    DEPENDS ${_oplist_yaml} ${_opvariant_h} ${GEN_CUSTOM_OPS_YAML}
+            ${GEN_FUNCTIONS_YAML} ${_codegen_templates}
+            ${_torchgen_srcs}
     WORKING_DIRECTORY ${EXECUTORCH_ROOT}
   )
   # Make generated file list available in parent scope
@@ -191,29 +205,85 @@ endfunction()
 
 # Generate a runtime lib for registering operators in Executorch
 function(gen_operators_lib)
-  set(multi_arg_names LIB_NAME KERNEL_LIBS DEPS)
+  set(multi_arg_names LIB_NAME KERNEL_LIBS DEPS DTYPE_SELECTIVE_BUILD)
   cmake_parse_arguments(GEN "" "" "${multi_arg_names}" ${ARGN})
 
   message(STATUS "Generating operator lib:")
   message(STATUS "  LIB_NAME: ${GEN_LIB_NAME}")
   message(STATUS "  KERNEL_LIBS: ${GEN_KERNEL_LIBS}")
   message(STATUS "  DEPS: ${GEN_DEPS}")
+  message(STATUS "  DTYPE_SELECTIVE_BUILD: ${GEN_DTYPE_SELECTIVE_BUILD}")
 
   set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME})
+  if(GEN_DTYPE_SELECTIVE_BUILD)
+    set(_opvariant_h
+      ${_out_dir}/selected_op_variants.h
+    )
+  endif()
 
   add_library(${GEN_LIB_NAME})
+
+  set(_srcs_list
+    ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp
+    ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h
+  )
+  if(GEN_DTYPE_SELECTIVE_BUILD)
+    list(APPEND _srcs_list ${_opvariant_h})
+  endif()
   target_sources(
     ${GEN_LIB_NAME}
-    PRIVATE ${_out_dir}/RegisterCodegenUnboxedKernelsEverything.cpp
-            ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h
+    PRIVATE ${_srcs_list}
   )
   target_link_libraries(${GEN_LIB_NAME} PRIVATE ${GEN_DEPS})
+  set(portable_kernels_check "portable_kernels")
   if(GEN_KERNEL_LIBS)
-    target_link_libraries(${GEN_LIB_NAME} PUBLIC ${GEN_KERNEL_LIBS})
+
+    set(_common_compile_options -Wno-deprecated-declarations -ffunction-sections -fdata-sections -Os)
+
+    if(GEN_DTYPE_SELECTIVE_BUILD)
+      if("${portable_kernels_check}" IN_LIST GEN_KERNEL_LIBS)
+        list(REMOVE_ITEM GEN_KERNEL_LIBS ${portable_kernels_check})
+
+        # Build kernels_util_all_deps, since later selected_portable_kernels depends on it
+        list(TRANSFORM _kernels_util_all_deps__srcs PREPEND "${EXECUTORCH_ROOT}/")
+        add_library(selected_kernels_util_all_deps ${_kernels_util_all_deps__srcs})
+        target_link_libraries(selected_kernels_util_all_deps PRIVATE executorch_core)
+        target_include_directories(selected_kernels_util_all_deps PUBLIC ${_common_include_directories})
+        target_compile_definitions(selected_kernels_util_all_deps PUBLIC C10_USING_CUSTOM_GENERATED_MACROS)
+        target_compile_options(selected_kernels_util_all_deps PUBLIC ${_common_compile_options})
+
+        # Build selected_portable_kernels
+        list(TRANSFORM _portable_kernels__srcs PREPEND "${EXECUTORCH_ROOT}/")
+        add_library(selected_portable_kernels ${_portable_kernels__srcs})
+        target_link_libraries(selected_portable_kernels PRIVATE executorch_core selected_kernels_util_all_deps)
+        target_compile_options(selected_portable_kernels PUBLIC ${_common_compile_options})
+        target_include_directories(selected_portable_kernels PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}/)
+
+        # Make sure the header is generated before compiling the library
+        add_dependencies(selected_portable_kernels ${GEN_LIB_NAME})
+        # Create a custom target for the header to ensure proper dependency tracking
+        add_custom_target(selected_portable_kernels_header DEPENDS ${_opvariant_h})
+        add_dependencies(selected_portable_kernels selected_portable_kernels_header)
+        # Apply the compile definition for dtype selective build
+        target_compile_definitions(selected_portable_kernels PRIVATE EXECUTORCH_SELECTIVE_BUILD_DTYPE=1)
+
+        target_link_libraries(${GEN_LIB_NAME} PUBLIC selected_portable_kernels)
+      else()
+        message(FATAL_ERROR "Currently dtype selective build is only supported for portable_kernels but {${GEN_KERNEL_LIBS}} were provided!")
+      endif()
+    endif()
+
+    # After removing portable_kernels, test if there are other kernel libs provided
+    if(GEN_KERNEL_LIBS)
+      target_link_libraries(${GEN_LIB_NAME} PUBLIC ${GEN_KERNEL_LIBS})
+    endif()
   endif()
 
   target_link_options_shared_lib(${GEN_LIB_NAME})
   set(_generated_headers ${_out_dir}/Functions.h ${_out_dir}/NativeFunctions.h)
+  if(GEN_DTYPE_SELECTIVE_BUILD)
+    list(APPEND _generated_headers ${_opvariant_h})
+  endif()
   set_target_properties(
     ${GEN_LIB_NAME} PROPERTIES PUBLIC_HEADER "${_generated_headers}"
   )