xref: /aosp_15_r20/external/executorch/examples/selective_build/test_selective_build.sh (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1#!/bin/bash
2# Copyright (c) Meta Platforms, Inc. and affiliates.
3# All rights reserved.
4#
5# This source code is licensed under the BSD-style license found in the
6# LICENSE file in the root directory of this source tree.
7
8# Test the end-to-end flow of selective build, using 3 APIs:
9# 1. Select all ops
10# 2. Select from a list of ops
11# 3. Select from a yaml file
12# 4. (TODO) Select from a serialized model (.pte)
13set -e
14
15# shellcheck source=/dev/null
16source "$(dirname "${BASH_SOURCE[0]}")/../../.ci/scripts/utils.sh"
17
18
19# BUCK2 examples; test internally in fbcode/xplat
20# 1. `--config executorch.select_ops=all`: select all ops from the dependency
21#       kernel libraries, register all of them into ExecuTorch runtime.
22# 2. `--config executorch.select_ops=list`: Only select ops from `ops` kwarg
23#       in `et_operator_library` macro.
24# 3. `--config executorch.select_ops=yaml`: Only select from a yaml file from
25#       `ops_schema_yaml_target` kwarg in `et_operator_library` macro
26# 4. `--config executorch.select_ops=dict`: Only select ops from `ops_dict`
27#       kwarg in `et_operator_library` macro. Add `dtype_selective_build = True`
28#       to executorch_generated_lib to select dtypes specified in the dictionary.
29
30# Other configs:
31# - `--config executorch.max_kernel_num=N`: Only allocate memory for the
32#       required number of operators. Users can retrieve N from `selected_operators.yaml`.
33
34test_buck2_select_all_ops() {
35    echo "Exporting MobilenetV3"
36    ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="mv3"
37
38    echo "Running selective build test"
39    $BUCK run //examples/selective_build:selective_build_test \
40        --config=executorch.select_ops=all -- --model_path=./mv3.pte
41
42    echo "Removing mv3.pte"
43    rm "./mv3.pte"
44}
45
46test_buck2_select_ops_in_list() {
47    echo "Exporting add_mul"
48    ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="add_mul"
49
50    echo "Running selective build test"
51    # set max_kernel_num=22: 19 primops, add, mul
52    $BUCK run //examples/selective_build:selective_build_test \
53        --config=executorch.max_kernel_num=22 \
54        --config=executorch.select_ops=list \
55        -- --model_path=./add_mul.pte
56
57    echo "Removing add_mul.pte"
58    rm "./add_mul.pte"
59}
60
61test_buck2_select_ops_in_dict() {
62    echo "Exporting add_mul"
63    ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="add_mul"
64
65    echo "Running selective build test"
66    # select ops and their dtypes using the dictionary API.
67    $BUCK run //examples/selective_build:selective_build_test \
68        --config=executorch.select_ops=dict \
69        --config=executorch.dtype_selective_build_lib=//examples/selective_build:select_ops_in_dict_lib \
70        -- --model_path=./add_mul.pte
71
72    echo "Removing add_mul.pte"
73    rm "./add_mul.pte"
74}
75
76test_buck2_select_ops_from_yaml() {
77    echo "Exporting custom_op_1"
78    ${PYTHON_EXECUTABLE} -m examples.portable.custom_ops.custom_ops_1
79
80    echo "Running selective build test"
81    $BUCK run //examples/selective_build:selective_build_test \
82        --config=executorch.select_ops=yaml -- --model_path=./custom_ops_1.pte
83
84    echo "Removing custom_ops_1.pte"
85    rm "./custom_ops_1.pte"
86}
87
88# CMake examples; test in OSS. Check the README for more information.
89test_cmake_select_all_ops() {
90    echo "Exporting MobilenetV3"
91    ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="mv3"
92
93    local example_dir=examples/selective_build
94    local build_dir=cmake-out/${example_dir}
95    rm -rf ${build_dir}
96    retry cmake -DBUCK2="$BUCK" \
97            -DCMAKE_BUILD_TYPE=Release \
98            -DEXECUTORCH_SELECT_ALL_OPS=ON \
99            -DCMAKE_INSTALL_PREFIX=cmake-out \
100            -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
101            -B${build_dir} \
102            ${example_dir}
103
104    echo "Building ${example_dir}"
105    cmake --build ${build_dir} -j9 --config Release
106
107    echo 'Running selective build test'
108    ${build_dir}/selective_build_test --model_path="./mv3.pte"
109
110    echo "Removing mv3.pte"
111    rm "./mv3.pte"
112}
113
114test_cmake_select_ops_in_list() {
115    echo "Exporting MobilenetV2"
116    ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="mv2"
117
118    local example_dir=examples/selective_build
119    local build_dir=cmake-out/${example_dir}
120    # set MAX_KERNEL_NUM=22: 19 primops, add, mul
121    rm -rf ${build_dir}
122    retry cmake -DBUCK2="$BUCK" \
123            -DCMAKE_BUILD_TYPE=Release \
124            -DMAX_KERNEL_NUM=22 \
125            -DEXECUTORCH_SELECT_OPS_LIST="aten::convolution.out,\
126aten::_native_batch_norm_legit_no_training.out,aten::hardtanh.out,aten::add.out,\
127aten::mean.out,aten::view_copy.out,aten::permute_copy.out,aten::addmm.out,\
128aten,aten::clone.out" \
129            -DCMAKE_INSTALL_PREFIX=cmake-out \
130            -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
131            -B${build_dir} \
132            ${example_dir}
133
134    echo "Building ${example_dir}"
135    cmake --build ${build_dir} -j9 --config Release
136
137    echo 'Running selective build test'
138    ${build_dir}/selective_build_test --model_path="./mv2.pte"
139
140    echo "Removing mv2.pte"
141    rm "./mv2.pte"
142}
143
144test_cmake_select_ops_in_yaml() {
145    echo "Exporting custom_op_1"
146    ${PYTHON_EXECUTABLE} -m examples.portable.custom_ops.custom_ops_1
147    local example_dir=examples/selective_build
148    local build_dir=cmake-out/${example_dir}
149    rm -rf ${build_dir}
150    retry cmake -DBUCK2="$BUCK" \
151            -DCMAKE_BUILD_TYPE=Release \
152            -DEXECUTORCH_SELECT_OPS_YAML=ON \
153            -DCMAKE_INSTALL_PREFIX=cmake-out \
154            -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
155            -B${build_dir} \
156            ${example_dir}
157
158    echo "Building ${example_dir}"
159    cmake --build ${build_dir} -j9 --config Release
160
161    echo 'Running selective build test'
162    ${build_dir}/selective_build_test --model_path="./custom_ops_1.pte"
163
164    echo "Removing custom_ops_1.pte"
165    rm "./custom_ops_1.pte"
166}
167
168if [[ -z $BUCK ]];
169then
170  BUCK=buck2
171fi
172
173if [[ -z $PYTHON_EXECUTABLE ]];
174then
175  PYTHON_EXECUTABLE=python3
176fi
177
178if [[ $1 == "cmake" ]];
179then
180    cmake_install_executorch_lib
181    test_cmake_select_all_ops
182    test_cmake_select_ops_in_list
183    test_cmake_select_ops_in_yaml
184elif [[ $1 == "buck2" ]];
185then
186    test_buck2_select_all_ops
187    test_buck2_select_ops_in_list
188    test_buck2_select_ops_in_dict
189    test_buck2_select_ops_from_yaml
190fi
191