Add custom ops test to CI job

.github/workflows/pull.yml

@@ -31,7 +31,7 @@ jobs:
         bash .ci/scripts/test.sh
 
         # Test custom ops
-        bash examples/custom_ops/test_custom_ops.sh
+        bash examples/custom_ops/test_custom_ops.sh buck2
 
   buck-build-test-macos:
     name: buck-build-test-macos
@@ -49,6 +49,9 @@ jobs:
 
         # Build and test Executorch
         bash .ci/scripts/test.sh
+
+        # Test custom ops
+        bash examples/custom_ops/test_custom_ops.sh buck2
         popd
 
   cmake-build-test-linux:
@@ -68,3 +71,24 @@ jobs:
 
         # Build and test Executorch
         bash .ci/scripts/test-cmake.sh
+
+        # Build and test custom ops
+        bash examples/custom_ops/test_custom_ops.sh cmake
+
+  cmake-build-test-macos:
+    name: cmake-build-test-macos
+    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
+    with:
+      runner: macos-m1-12
+      submodules: 'true'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      script: |
+        WORKSPACE=$(pwd)
+
+        pushd "${WORKSPACE}/pytorch/executorch"
+        # Setup MacOS dependencies as there is no Docker support on MacOS atm
+        bash .ci/scripts/setup-macos.sh
+
+        # Build and test custom ops
+        bash examples/custom_ops/test_custom_ops.sh cmake
+        popd

examples/custom_ops/test_custom_ops.sh

@@ -14,7 +14,7 @@ set -e
 test_buck2_custom_op_1() {
   local model_name='custom_ops_1'
   echo "Exporting ${model_name}.pte"
-  python3 -m "examples.custom_ops.${model_name}"
+  python -m "examples.custom_ops.${model_name}"
   # should save file custom_ops_1.pte
 
   echo 'Running executor_runner'
@@ -29,7 +29,7 @@ test_buck2_custom_op_1() {
 test_cmake_custom_op_1() {
   local model_name='custom_ops_1'
   echo "Exporting ${model_name}.pte"
-  python3 -m "examples.custom_ops.${model_name}"
+  python -m "examples.custom_ops.${model_name}"
   # should save file custom_ops_1.pte
   (rm -rf cmake-out \
     && mkdir cmake-out \
@@ -50,7 +50,7 @@ test_buck2_custom_op_2() {
   SO_LIB=$(buck2 build //examples/custom_ops:custom_ops_aot_lib_2 --show-output | grep "buck-out" | cut -d" " -f2)
 
   echo "Exporting ${model_name}.pte"
-  python3 -m "examples.custom_ops.${model_name}" --so_library="$SO_LIB"
+  python -m "examples.custom_ops.${model_name}" --so_library="$SO_LIB"
   # should save file custom_ops_2.pte
 
   buck2 run //examples/executor_runner:executor_runner \
@@ -92,14 +92,24 @@ test_cmake_custom_op_2() {
 
   EXT=$(get_shared_lib_ext)
   echo "Exporting ${model_name}.pte"
-  python3 -m "examples.custom_ops.${model_name}" --so_library="cmake-out/examples/custom_ops/libcustom_ops_aot_lib$EXT"
+  python -m "examples.custom_ops.${model_name}" --so_library="cmake-out/examples/custom_ops/libcustom_ops_aot_lib$EXT"
   # should save file custom_ops_2.pte
 
   echo 'Running executor_runner'
   cmake-out/executor_runner "--model_path=./${model_name}.pte"
 }
 
-test_buck2_custom_op_1
-test_cmake_custom_op_1
-test_buck2_custom_op_2
-test_cmake_custom_op_2
+if [[ $1 == "cmake" ]];
+then
+  test_cmake_custom_op_1
+  test_cmake_custom_op_2
+elif [[ $1 == "buck2" ]];
+then
+  test_buck2_custom_op_1
+  test_buck2_custom_op_2
+else
+  test_cmake_custom_op_1
+  test_cmake_custom_op_2
+  test_buck2_custom_op_1
+  test_buck2_custom_op_2
+fi

third-party/link_torch.sh

@@ -11,14 +11,39 @@
 
 set -e
 
-LIB=$(python3 -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')
+# Named argument:
+# -o: output directory/file
+# -f: a list of files/directories that we want to link to the output directory, separated by comma ",".
+# These paths need to be in relative path format to the python library path.
+while getopts ":o:f:" opt; do
+  case $opt in
+    o) OUT="$OPTARG"
+    ;;
+    f) FILES="$OPTARG"
+    ;;
+    \?) echo "Invalid option -$OPTARG" >&2
+    exit 1
+    ;;
+  esac
 
-SUBPATH=$1
-OUT=$2
-if [[ -f "$LIB/$SUBPATH" ]] || [[ -d "$LIB/$SUBPATH" ]];
-then
-    ln -s "$LIB/$SUBPATH" "$OUT"
-else
-    echo "Error: $LIB/$SUBPATH doesn't exist"
+  case $OPTARG in
+    -*) echo "Option $opt needs to be one of -o, -f."
     exit 1
-fi
+    ;;
+  esac
+done
+
+LIB=$(python3 -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')
+
+# delimiter ,
+export IFS=","
+
+for SUBPATH in $FILES; do
+  if [[ -f "$LIB/$SUBPATH" ]] || [[ -d "$LIB/$SUBPATH" ]];
+  then
+      ln -s "$LIB/$SUBPATH" "$OUT"
+  else
+      echo "Error: $LIB/$SUBPATH doesn't exist"
+      exit 1
+  fi
+done