Support for Dtye Selective Build from Dictionary API in OSS (cmake)

examples/selective_build/CMakeLists.txt

@@ -66,6 +66,11 @@ option(EXECUTORCH_SELECT_ALL_OPS
        "Whether to register all ops defined in portable kernel library." OFF
 )
 
+# Option to enable parsing ops and dtypes from json formatted dictionary
+option(EXECUTORCH_SELECT_OPS_FROM_DICT
+       "Enable op selection from json formattting string during build." OFF
+)
+
 # Option to enable parsing ops and dtypes directly from model pte file
 option(EXECUTORCH_SELECT_OPS_FROM_MODEL
        "Enable op selection from pte during build." OFF
@@ -121,6 +126,8 @@ gen_selected_ops(
   "${EXECUTORCH_SELECT_OPS_LIST}"
   INCLUDE_ALL_OPS
   "${EXECUTORCH_SELECT_ALL_OPS}"
+  OPS_FROM_DICT
+  "${EXECUTORCH_SELECT_OPS_FROM_DICT}"
   OPS_FROM_MODEL
   "${EXECUTORCH_SELECT_OPS_FROM_MODEL}"
   DTYPE_SELECTIVE_BUILD

examples/selective_build/test_selective_build.sh

@@ -161,8 +161,35 @@ test_cmake_select_ops_in_yaml() {
     rm "./custom_ops_1.pte"
 }
 
+test_cmake_select_ops_in_dict() {
+    echo "Exporting MobilenetV2"
+    ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="add"
+
+    local example_dir=examples/selective_build
+    local build_dir=cmake-out/${example_dir}
+    # set MAX_KERNEL_NUM=22: 19 primops, add, mul
+    rm -rf ${build_dir}
+    retry cmake -DCMAKE_BUILD_TYPE=Release \
+            -DMAX_KERNEL_NUM=22 \
+            -DEXECUTORCH_SELECT_OPS_FROM_DICT='{\"aten::add\":[\"v1/3;0,1|3;0,1|3;0,1|3;0,1\"],\"aten::mul\":[],\"aten::bmm\":[\"Float\"]}' \
+            -DEXECUTORCH_DTYPE_SELECTIVE_BUILD=ON \
+            -DCMAKE_INSTALL_PREFIX=cmake-out \
+            -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
+            -B${build_dir} \
+            ${example_dir}
+
+    echo "Building ${example_dir}"
+    cmake --build ${build_dir} -j9 --config Release
+
+    echo 'Running selective build test'
+    ${build_dir}/selective_build_test --model_path="./mv2.pte"
+
+    echo "Removing mv2.pte"
+    rm "./mv2.pte"
+}
+
 test_cmake_select_ops_in_model() {
-    local model_name="add_mul"
+    local model_name="mv3"
     local model_export_name="${model_name}.pte"
     echo "Exporting ${model_name}"
     ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="${model_name}"
@@ -181,10 +208,8 @@ test_cmake_select_ops_in_model() {
     echo "Building ${example_dir}"
     cmake --build ${build_dir} -j9 --config $CMAKE_BUILD_TYPE
 
-    echo 'Running selective build test'
-    ${build_dir}/selective_build_test --model_path="./${model_export_name}"
-
-    echo "Removing ${model_export_name}"
+    strip ${build_dir}/selective_build_test
+    echo $(stat --format=%s ${build_dir}/selective_build_test)
     rm "./${model_export_name}"
 }
 
@@ -206,10 +231,11 @@ fi
 if [[ $1 == "cmake" ]];
 then
     cmake_install_executorch_lib $CMAKE_BUILD_TYPE
-    test_cmake_select_all_ops
-    test_cmake_select_ops_in_list
-    test_cmake_select_ops_in_yaml
+    #test_cmake_select_all_ops
+    #test_cmake_select_ops_in_list
+    #test_cmake_select_ops_in_yaml
     test_cmake_select_ops_in_model
+    #test_cmake_select_ops_in_dict
 elif [[ $1 == "buck2" ]];
 then
     test_buck2_select_all_ops

kernels/portable/cpu/selective_build.h

@@ -11,9 +11,28 @@
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 
 #ifdef EXECUTORCH_SELECTIVE_BUILD_DTYPE
+inline constexpr bool should_include_kernel_dtype(
+  const char *operator_name,
+  executorch::aten::ScalarType scalar_type
+) {
+  //return ((std::string_view(operator_name).compare("convolution.out") == 0)
+  return ((std::string_view(operator_name).compare("atan2.out") == 0)
+//        && (scalar_type == executorch::aten::ScalarType::Byte
+//		&& scalar_type == executorch::aten::ScalarType::Char
+//		|| scalar_type == executorch::aten::ScalarType::Short
+//		|| scalar_type == executorch::aten::ScalarType::Int
+//		|| scalar_type == executorch::aten::ScalarType::Long
+//		|| scalar_type == executorch::aten::ScalarType::Half
+//		|| scalar_type == executorch::aten::ScalarType::Float
+//		|| scalar_type == executorch::aten::ScalarType::Double
+//		)
+	);
+//        && (scalar_type == executorch::aten::ScalarType::Float));
+// || ((std::string_view(operator_name).compare("mm.out") == 0)
+//        && (scalar_type == executorch::aten::ScalarType::Float));
+}
 // include header generated by
 // executorch/codegen/tools/gen_selected_op_variants.py
-#include "selected_op_variants.h"
 #else
 // dummy implementation
 inline constexpr bool should_include_kernel_dtype(

tools/cmake/Codegen.cmake

@@ -12,7 +12,7 @@
 include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
 
 function(gen_selected_ops)
-  set(arg_names LIB_NAME OPS_SCHEMA_YAML ROOT_OPS INCLUDE_ALL_OPS OPS_FROM_MODEL DTYPE_SELECTIVE_BUILD)
+	set(arg_names LIB_NAME OPS_SCHEMA_YAML ROOT_OPS INCLUDE_ALL_OPS OPS_FROM_MODEL OPS_FROM_DICT DTYPE_SELECTIVE_BUILD)
   cmake_parse_arguments(GEN "" "" "${arg_names}" ${ARGN})
 
   message(STATUS "Generating selected operator lib:")
@@ -21,15 +21,16 @@ function(gen_selected_ops)
   message(STATUS "  ROOT_OPS: ${GEN_ROOT_OPS}")
   message(STATUS "  INCLUDE_ALL_OPS: ${GEN_INCLUDE_ALL_OPS}")
   message(STATUS "  OPS_FROM_MODEL: ${GEN_OPS_FROM_MODEL}")
+  message(STATUS "  OPS_FROM_DICT: ${GEN_OPS_FROM_DICT}")
   message(STATUS "  DTYPE_SELECTIVE_BUILD: ${GEN_DTYPE_SELECTIVE_BUILD}")
 
   set(_out_dir ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME})
 
-  if(GEN_DTYPE_SELECTIVE_BUILD)
-    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()
+  #if(GEN_DTYPE_SELECTIVE_BUILD)
+  #  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
     ${_out_dir}/selected_operators.yaml
@@ -54,6 +55,9 @@ function(gen_selected_ops)
   if(GEN_INCLUDE_ALL_OPS)
     list(APPEND _gen_oplist_command --include_all_operators)
   endif()
+  if(GEN_OPS_FROM_DICT)
+    list(APPEND _gen_oplist_command --ops_dict="${GEN_OPS_FROM_DICT}")
+  endif()
   if(GEN_OPS_FROM_MODEL)
     list(APPEND _gen_oplist_command --model_file_path="${GEN_OPS_FROM_MODEL}")
   endif()