xref: /aosp_15_r20/external/executorch/examples/portable/custom_ops/test_custom_ops.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 using custom operator in a PyTorch model and use
9# EXIR to capture and export a model file. Then use `executor_runner` demo C++
10# binary to run the model.
11
12set -e
13
14# shellcheck source=/dev/null
15source "$(dirname "${BASH_SOURCE[0]}")/../../../.ci/scripts/utils.sh"
16
17test_cmake_custom_op_1() {
18  local model_name='custom_ops_1'
19  echo "Exporting ${model_name}.pte"
20  ${PYTHON_EXECUTABLE} -m "examples.portable.custom_ops.${model_name}"
21  # should save file custom_ops_1.pte
22  local example_dir=examples/portable/custom_ops
23  local build_dir=cmake-out/${example_dir}
24  rm -rf ${build_dir}
25  retry cmake \
26        -DREGISTER_EXAMPLE_CUSTOM_OP=1 \
27        -DCMAKE_INSTALL_PREFIX=cmake-out \
28        -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
29        -B${build_dir} \
30        ${example_dir}
31
32  echo "Building ${example_dir}"
33  cmake --build ${build_dir} -j9 --config Release
34
35  echo 'Running custom_ops_executor_runner'
36  ${build_dir}/custom_ops_executor_runner --model_path="./${model_name}.pte"
37}
38
39get_shared_lib_ext() {
40  UNAME=$(uname)
41  if [[ $UNAME == "Darwin" ]];
42  then
43    EXT=".dylib"
44  elif [[ $UNAME == "Linux" ]];
45  then
46    EXT=".so"
47  else
48    echo "Unsupported platform $UNAME"
49    exit 1
50  fi
51  echo $EXT
52}
53
54test_cmake_custom_op_2() {
55  local model_name='custom_ops_2'
56  SITE_PACKAGES="$(${PYTHON_EXECUTABLE} -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')"
57  CMAKE_PREFIX_PATH="$PWD/cmake-out/lib/cmake/ExecuTorch;${SITE_PACKAGES}/torch"
58
59  local example_dir=examples/portable/custom_ops
60  local build_dir=cmake-out/${example_dir}
61  rm -rf ${build_dir}
62  retry cmake \
63        -DREGISTER_EXAMPLE_CUSTOM_OP=2 \
64        -DCMAKE_PREFIX_PATH="$CMAKE_PREFIX_PATH" \
65        -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
66        -B${build_dir} \
67        ${example_dir}
68
69  echo "Building ${example_dir}"
70  cmake --build ${build_dir} -j9 --config Release
71
72  EXT=$(get_shared_lib_ext)
73  echo "Exporting ${model_name}.pte"
74  ${PYTHON_EXECUTABLE} -m "examples.portable.custom_ops.${model_name}" --so_library="cmake-out/examples/portable/custom_ops/libcustom_ops_aot_lib$EXT"
75  # should save file custom_ops_2.pte
76
77  echo 'Running custom_ops_executor_runner'
78  ${build_dir}/custom_ops_executor_runner --model_path="./${model_name}.pte"
79}
80
81if [[ -z $PYTHON_EXECUTABLE ]];
82then
83  PYTHON_EXECUTABLE=python3
84fi
85
86cmake_install_executorch_lib
87test_cmake_custom_op_1
88test_cmake_custom_op_2
89