Add selective build examples for buck2

.github/workflows/pull.yml

@@ -137,3 +137,53 @@ jobs:
         # Build and test custom ops
         PYTHON_EXECUTABLE=python bash examples/custom_ops/test_custom_ops.sh "${BUILD_TOOL}"
         popd
+
+  test-selective-build-linux:
+    name: test-selective-build-linux
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    strategy:
+      matrix:
+        include:
+          - build-tool: buck2
+          # - build-tool: cmake (TODO)
+      fail-fast: false
+    with:
+      runner: linux.2xlarge
+      docker-image: executorch-ubuntu-22.04-clang12
+      submodules: 'true'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      script: |
+        PYTHON_VERSION=3.10
+        # TODO: Figure out why /opt/conda/envs/py_$PYTHON_VERSION/bin is not in the path
+        # here, as it's there in the container
+        export PATH="/opt/conda/envs/py_${PYTHON_VERSION}/bin:${PATH}"
+
+        BUILD_TOOL=${{ matrix.build-tool }}
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
+        # Test selective build
+        PYTHON_EXECUTABLE=python bash examples/selective_build/test_selective_build.sh "${BUILD_TOOL}"
+
+  test-selective-build-macos:
+    name: test-selective-build-macos
+    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
+    strategy:
+      matrix:
+        include:
+          - build-tool: buck2
+          # - build-tool: cmake (TODO)
+      fail-fast: false
+    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"
+
+        BUILD_TOOL=${{ matrix.build-tool }}
+        # Setup MacOS dependencies as there is no Docker support on MacOS atm
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-macos.sh "${BUILD_TOOL}"
+        # Build and test selective build
+        PYTHON_EXECUTABLE=python bash examples/selective_build/test_selective_build.sh "${BUILD_TOOL}"
+        popd

examples/selective_build/README.md

@@ -0,0 +1,25 @@
+# Selective Build Examples
+To optimize binary size of ExecuTorch runtime, selective build can be used. This folder contains examples to select only the operators needed for Executorch build. We provide APIs for both CMake build and buck2 build. This example will demonstrate both. You can find more information on how to use buck2 macros in [wiki](https://github.com/pytorch/executorch/blob/main/docs/website/docs/tutorials/selective_build.md).
+
+## How to run
+
+Prerequisite: finish the [setting up wiki](https://github.com/pytorch/executorch/blob/main/docs/website/docs/tutorials/00_setting_up_executorch.md).
+
+Run:
+
+```bash
+bash test_selective_build.sh
+```
+
+## BUCK2 examples
+
+Check out `targets.bzl` for demo of 3 selective build APIs:
+1. Select all ops from the dependency kernel libraries, register all of them into Executorch runtime.
+2. Only select ops from `ops` kwarg in `et_operator_library` macro.
+3. Only select from a yaml file from `ops_schema_yaml_target` kwarg in `et_operator_library` macro.
+
+We have one more API incoming: only select from an exported model file (.pte).
+
+## CMake examples
+
+TODO

examples/selective_build/TARGETS

@@ -0,0 +1,8 @@
+# Any targets that should be shared between fbcode and xplat must be defined in
+# targets.bzl. This file can contain fbcode-only targets.
+
+load(":targets.bzl", "define_common_targets")
+
+oncall("executorch")
+
+define_common_targets()

examples/selective_build/targets.bzl

@@ -0,0 +1,80 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_oss_build_kwargs", "runtime")
+load("@fbsource//xplat/executorch/codegen:codegen.bzl", "et_operator_library", "executorch_generated_lib")
+
+def define_common_targets():
+    """Defines targets that should be shared between fbcode and xplat.
+
+    The directory containing this targets.bzl file should also contain both
+    TARGETS and BUCK files that call this function.
+    """
+
+    # Select all ops: register all the ops in portable/functions.yaml
+    et_operator_library(
+        name = "select_all_ops",
+        include_all_operators = True,
+    )
+
+    executorch_generated_lib(
+        name = "select_all_lib",
+        functions_yaml_target = "//executorch/kernels/portable:functions.yaml",
+        deps = [
+            "//executorch/kernels/portable:operators",
+            ":select_all_ops",
+        ],
+    )
+
+    # Select a list of operators: defined in `ops`
+    et_operator_library(
+        name = "select_ops_in_list",
+        ops = [
+            "aten::add.out",
+            "aten::mm.out",
+        ],
+    )
+
+    executorch_generated_lib(
+        name = "select_ops_in_list_lib",
+        functions_yaml_target = "//executorch/kernels/portable:functions.yaml",
+        deps = [
+            "//executorch/kernels/portable:operators",
+            ":select_ops_in_list",
+        ],
+    )
+
+    # Select all ops from a yaml file
+    et_operator_library(
+        name = "select_ops_from_yaml",
+        ops_schema_yaml_target = "//executorch/examples/custom_ops:custom_ops.yaml",
+    )
+
+    executorch_generated_lib(
+        name = "select_ops_from_yaml_lib",
+        custom_ops_yaml_target = "//executorch/examples/custom_ops:custom_ops.yaml",
+        deps = [
+            "//executorch/examples/custom_ops:custom_ops_1",
+            "//executorch/examples/custom_ops:custom_ops_2",
+            ":select_ops_from_yaml",
+        ],
+    )
+
+    # Select all ops from a given model
+    # TODO(larryliu0820): Add this
+
+    # ~~~ Test binary for selective build ~~~
+    select_ops = native.read_config("executorch", "select_ops", None)
+    lib = []
+    if select_ops == "all":
+        lib.append(":select_all_lib")
+    elif select_ops == "list":
+        lib.append(":select_ops_in_list_lib")
+    elif select_ops == "yaml":
+        lib.append(":select_ops_from_yaml_lib")
+    runtime.cxx_binary(
+        name = "selective_build_test",
+        srcs = [],
+        deps = [
+            "//executorch/examples/executor_runner:executor_runner_lib",
+        ] + lib,
+        define_static_target = True,
+        **get_oss_build_kwargs()
+    )

examples/selective_build/test_selective_build.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Test the end-to-end flow of selective build, using 3 APIs:
+# 1. Select all ops
+# 2. Select from a list of ops
+# 3. Select from a yaml file
+# 4. (TODO) Select from a serialized model (.pte)
+set -e
+
+test_buck2_select_all_ops() {
+    echo "Exporting MobilenetV3"
+    ${PYTHON_EXECUTABLE} -m examples.export.export_example --model_name="mv3"
+
+    echo "Running executor_runner"
+    buck2 run //examples/selective_build:selective_build_test \
+        --config=executorch.select_ops=all -- --model_path=./mv3.pte
+
+    echo "Removing mv3.pte"
+    rm "./mv3.pte"
+}
+
+test_buck2_select_ops_in_list() {
+    echo "Exporting add_mul"
+    ${PYTHON_EXECUTABLE} -m examples.export.export_example --model_name="add_mul"
+
+    echo "Running executor_runner"
+    buck2 run //examples/selective_build:selective_build_test \
+        --config=executorch.select_ops=list -- --model_path=./add_mul.pte
+
+    echo "Removing add_mul.pte"
+    rm "./add_mul.pte"
+}
+
+test_buck2_select_ops_from_yaml() {
+    echo "Exporting custom_op_1"
+    ${PYTHON_EXECUTABLE} -m examples.custom_ops.custom_ops_1
+
+    echo "Running executor_runner"
+    buck2 run //examples/selective_build:selective_build_test \
+        --config=executorch.select_ops=yaml -- --model_path=./custom_ops_1.pte
+
+    echo "Removing custom_ops_1.pte"
+    rm "./custom_ops_1.pte"
+}
+
+# TODO(larryliu0820) Add example to select ops from model.
+
+if [[ -z $PYTHON_EXECUTABLE ]];
+then
+  PYTHON_EXECUTABLE=python3
+fi
+
+
+test_buck2_select_all_ops
+test_buck2_select_ops_in_list
+test_buck2_select_ops_from_yaml