diff --git a/examples/selective_build/CMakeLists.txt b/examples/selective_build/CMakeLists.txt index a37b4362d78..1f6a06ea007 100644 --- a/examples/selective_build/CMakeLists.txt +++ b/examples/selective_build/CMakeLists.txt @@ -61,6 +61,10 @@ option(EXECUTORCH_SELECT_OPS_LIST "Register a list of ops, separated by comma" option(EXECUTORCH_SELECT_ALL_OPS "Whether to register all ops defined in portable kernel library." 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 +) # ------------------------------- OPTIONS END -------------------------------- # @@ -108,6 +112,8 @@ gen_selected_ops( "${EXECUTORCH_SELECT_OPS_LIST}" INCLUDE_ALL_OPS "${EXECUTORCH_SELECT_ALL_OPS}" + OPS_FROM_MODEL + "${EXECUTORCH_SELECT_OPS_FROM_MODEL}" ) generate_bindings_for_kernels( diff --git a/examples/selective_build/test_selective_build.sh b/examples/selective_build/test_selective_build.sh index 324a4fe27a5..ee5eb7f6523 100644 --- a/examples/selective_build/test_selective_build.sh +++ b/examples/selective_build/test_selective_build.sh @@ -161,6 +161,29 @@ test_cmake_select_ops_in_yaml() { rm "./custom_ops_1.pte" } +test_cmake_select_ops_in_model() { + echo "Exporting MobilenetV2" + ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="mv2" + local example_dir=examples/selective_build + local build_dir=cmake-out/${example_dir} + rm -rf ${build_dir} + retry cmake -DCMAKE_BUILD_TYPE=Release \ + -DEXECUTORCH_SELECT_OPS_FROM_MODEL="./mv2.pte" \ + -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" +} + if [[ -z $BUCK ]]; then BUCK=buck2 @@ -177,6 +200,7 @@ then test_cmake_select_all_ops test_cmake_select_ops_in_list test_cmake_select_ops_in_yaml + test_cmake_select_ops_in_model elif [[ $1 == "buck2" ]]; then test_buck2_select_all_ops diff --git a/tools/cmake/Codegen.cmake b/tools/cmake/Codegen.cmake index f1dac84de43..321ef01a60b 100644 --- a/tools/cmake/Codegen.cmake +++ b/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) + set(arg_names LIB_NAME OPS_SCHEMA_YAML ROOT_OPS INCLUDE_ALL_OPS OPS_FROM_MODEL) cmake_parse_arguments(GEN "" "" "${arg_names}" ${ARGN}) message(STATUS "Generating operator lib:") @@ -20,10 +20,12 @@ function(gen_selected_ops) message(STATUS " OPS_SCHEMA_YAML: ${GEN_OPS_SCHEMA_YAML}") 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}") set(_oplist_yaml ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}/selected_operators.yaml ) + file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${GEN_LIB_NAME}) file(GLOB_RECURSE _codegen_tools_srcs "${EXECUTORCH_ROOT}/codegen/tools/*.py") @@ -43,6 +45,9 @@ function(gen_selected_ops) if(GEN_INCLUDE_ALL_OPS) list(APPEND _gen_oplist_command --include_all_operators) endif() + if(GEN_OPS_FROM_MODEL) + list(APPEND _gen_oplist_command --model_file_path="${GEN_OPS_FROM_MODEL}") + endif() message("Command - ${_gen_oplist_command}") add_custom_command(