custom build script (#30144)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/30144

Create script to produce libtorch that only contains ops needed by specific
models. Developers can use this workflow to further optimize mobile build size.

Need keep a dummy stub for unused (stripped) ops because some JIT side
logic requires certain function schemas to be existed in the JIT op
registry.

Test Steps:
1. Build "dump_operator_names" binary and use it to dump root ops needed
by a specific model:
```
build/bin/dump_operator_names --model=mobilenetv2.pk --output=mobilenetv2.yaml
```

2. The MobileNetV2 model should use the following ops:
```
- aten::t
- aten::dropout
- aten::mean.dim
- aten::add.Tensor
- prim::ListConstruct
- aten::addmm
- aten::_convolution
- aten::batch_norm
- aten::hardtanh_
- aten::mm
```
NOTE that for some reason it outputs "aten::addmm" but actually uses "aten::mm".
You need fix it manually for now.

3. Run custom build script locally (use Android as an example):
```
SELECTED_OP_LIST=mobilenetv2.yaml scripts/build_pytorch_android.sh armeabi-v7a
```

4. Checkout demo app that uses locally built library instead of
downloading from jcenter repo:
```
git clone --single-branch --branch custom_build git@github.com:ljk53/android-demo-app.git
```

5. Copy locally built libraries to demo app folder:
```
find ${HOME}/src/pytorch/android -name '*.aar' -exec cp {} ${HOME}/src/android-demo-app/HelloWorldApp/app/libs/ \;
```

6. Build demo app with locally built libtorch:
```
cd ${HOME}/src/android-demo-app/HelloWorldApp
./gradlew clean && ./gradlew assembleDebug
```

7. Install and run the demo app.

In-APK arm-v7 libpytorch_jni.so build size reduced from 5.5M to 2.9M.

Test Plan: Imported from OSS

Differential Revision: D18612127

Pulled By: ljk53

fbshipit-source-id: fa8d5e1d3259143c7346abd1c862773be8c7e29a

CMakeLists.txt

@@ -205,6 +205,8 @@ cmake_dependent_option(
   "MSVC" OFF)
 
 set(ONNX_NAMESPACE "onnx_torch" CACHE STRING "A namespace for ONNX; needed to build with other frameworks that share ONNX.")
+set(SELECTED_OP_LIST "" CACHE STRING
+    "Path to the yaml file that contains the list of operators to include for custom build. Include all operators by default.")
 
 # This is a fix for a rare build issue on Ubuntu:
 # symbol lookup error: miniconda3/envs/pytorch-py3.7/lib/libmkl_intel_lp64.so: undefined symbol: mkl_blas_dsyrk

caffe2/CMakeLists.txt

@@ -296,6 +296,7 @@ if (NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
       --declarations-path "${CMAKE_BINARY_DIR}/aten/src/ATen/Declarations.yaml"
       --nn-path "aten/src"
       $<$<BOOL:${INTERN_DISABLE_AUTOGRAD}>:--disable-autograd>
+      $<$<BOOL:${SELECTED_OP_LIST}>:--selected-op-list-path="${SELECTED_OP_LIST}">
     DEPENDS
     "${CMAKE_BINARY_DIR}/aten/src/ATen/Declarations.yaml"
     "${CMAKE_CURRENT_LIST_DIR}/../aten/src/THNN/generic/THNN.h"

cmake/Summary.cmake

@@ -130,7 +130,9 @@ function (caffe2_print_configuration_summary)
     message(STATUS "    USE_GLOO            : ${USE_GLOO}")
   endif()
   message(STATUS "  BUILD_NAMEDTENSOR   : ${BUILD_NAMEDTENSOR}")
-
+  if(NOT "${SELECTED_OP_LIST}" STREQUAL "")
+    message(STATUS "  SELECTED_OP_LIST    : ${SELECTED_OP_LIST}")
+  endif()
   message(STATUS "  Public Dependencies  : ${Caffe2_PUBLIC_DEPENDENCY_LIBS}")
   message(STATUS "  Private Dependencies : ${Caffe2_DEPENDENCY_LIBS}")
 endfunction()

scripts/build_android.sh

@@ -61,6 +61,10 @@ if [ -n "${BUILD_PYTORCH_MOBILE:-}" ]; then
   CMAKE_ARGS+=("-DCMAKE_PREFIX_PATH=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')")
   CMAKE_ARGS+=("-DPYTHON_EXECUTABLE=$(python -c 'import sys; print(sys.executable)')")
   CMAKE_ARGS+=("-DBUILD_CUSTOM_PROTOBUF=OFF")
+  # custom build with selected ops
+  if [ -n "${SELECTED_OP_LIST}" ]; then
+    CMAKE_ARGS+=("-DSELECTED_OP_LIST=${SELECTED_OP_LIST}")
+  fi
 else
   # Build protobuf from third_party so we have a host protoc binary.
   echo "Building protoc"

scripts/build_ios.sh

@@ -23,6 +23,10 @@ if [ -n "${BUILD_PYTORCH_MOBILE:-}" ]; then
   CMAKE_ARGS+=("-DCMAKE_PREFIX_PATH=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')")
   CMAKE_ARGS+=("-DPYTHON_EXECUTABLE=$(python -c 'import sys; print(sys.executable)')")
   CMAKE_ARGS+=("-DBUILD_CUSTOM_PROTOBUF=OFF")
+  # custom build with selected ops
+  if [ -n "${SELECTED_OP_LIST}" ]; then
+    CMAKE_ARGS+=("-DSELECTED_OP_LIST=${SELECTED_OP_LIST}")
+  fi
   # bitcode
   if [ "${ENABLE_BITCODE:-}" == '1' ]; then
     CMAKE_ARGS+=("-DCMAKE_C_FLAGS=-fembed-bitcode")

scripts/build_mobile.sh

@@ -27,6 +27,10 @@ CMAKE_ARGS+=("-DCMAKE_PREFIX_PATH=$(python -c 'from distutils.sysconfig import g
 CMAKE_ARGS+=("-DPYTHON_EXECUTABLE=$(python -c 'import sys; print(sys.executable)')")
 CMAKE_ARGS+=("-DBUILD_CUSTOM_PROTOBUF=OFF")
 CMAKE_ARGS+=("-DBUILD_SHARED_LIBS=OFF")
+# custom build with selected ops
+if [ -n "${SELECTED_OP_LIST}" ]; then
+  CMAKE_ARGS+=("-DSELECTED_OP_LIST=${SELECTED_OP_LIST}")
+fi
 
 # If Ninja is installed, prefer it to Make
 if [ -x "$(command -v ninja)" ]; then

tools/jit/gen_jit_dispatch.py

@@ -15,8 +15,9 @@ torch/csrc/jit/generated/
 import argparse
 import copy
 import re
+import yaml
 from itertools import groupby
-from ..autograd.utils import CodeTemplate, write
+from ..autograd.utils import CodeTemplate, YamlLoader, write
 from ..autograd.gen_autograd import load_aten_declarations
 from ..autograd.gen_autograd import RETURNS_VIEWS_OF_INPUT
 
@@ -263,7 +264,13 @@ def argument_order(decl):
     return decl.get('jit_argument_order') or list(range(len(decl['arguments'])))
 
 
-def gen_jit_dispatch(declarations, out, template_path, disable_autograd=False):
+def load_op_list(path):
+    with open(path, 'r') as f:
+        op_list = yaml.load(f, Loader=YamlLoader)
+    return op_list
+
+
+def gen_jit_dispatch(declarations, out, template_path, disable_autograd=False, selected_op_list_path=None):
     REGISTER_ATEN_OPS_CPP = CodeTemplate.from_file(template_path + '/register_aten_ops.cpp')
 
     ops = []
@@ -307,6 +314,8 @@ def gen_jit_dispatch(declarations, out, template_path, disable_autograd=False):
         return 'jit_type' in arg and arg['jit_type'] in {"Tensor!", "Tensor(a!)"}
 
     def emit_decl_variant(decl):
+        if ('emit_dummy_placeholder' in decl):
+            return "DUMMY_OPERATION"
         kw_assignments = []
 
         # mutable arguments in aten are passed as non const references
@@ -337,11 +346,13 @@ def gen_jit_dispatch(declarations, out, template_path, disable_autograd=False):
                                              lvalues=lvalues)
         return constructor
 
-    def filter_decls(jit_decls, disable_autograd):
+    def filter_decls(jit_decls, disable_autograd, selected_op_list):
         result = []
         for decl in jit_decls:
             if disable_autograd and is_backward_op(decl):
                 continue
+            if selected_op_list and signature_without_args(decl) not in selected_op_list:
+                decl['emit_dummy_placeholder'] = True
             result.append(decl)
         return result
 
@@ -431,7 +442,8 @@ def gen_jit_dispatch(declarations, out, template_path, disable_autograd=False):
                 additional_jit_decls.append(decl_copy)
 
     jit_decls.extend(additional_jit_decls)
-    jit_decls = filter_decls(jit_decls, disable_autograd)
+    selected_op_list = load_op_list(selected_op_list_path) if selected_op_list_path else None
+    jit_decls = filter_decls(jit_decls, disable_autograd, selected_op_list)
 
     # Group and sort the generated snippets to ensure that the
     # generation is deterministic
@@ -563,6 +575,12 @@ def signature(decl, should_match_schema=True):
     return match_signature(decl, constructed_string, should_match_schema)
 
 
+def signature_without_args(decl):
+    name = decl['name'] if not is_out_variant(decl) else decl['name'][:-4]
+    overload_name = '.' + decl['overload_name'] if not decl['overload_name'] == '' else ''
+    return 'aten::{}{}'.format(name, overload_name)
+
+
 def main():
     parser = argparse.ArgumentParser(
         description='Generate JIT op dispatch')

tools/jit/templates/register_aten_ops.cpp

@@ -86,6 +86,11 @@ c10::OperatorOptions atenOperatorOptions() {
   return result;
 }
 
+int (*DUMMY_OPERATION)(Stack&) = [](Stack& stack) -> int {
+  TORCH_CHECK(false, "Operator has been stripped in the custom build.")
+  return 0;
+};
+
 RegisterOperators reg(
     {Operator(
          "aten::get_device(Tensor self) -> int",

tools/setup_helpers/generate_code.py

@@ -24,7 +24,8 @@ def generate_code(ninja_global=None,
                   nn_path=None,
                   install_dir=None,
                   subset=None,
-                  disable_autograd=False):
+                  disable_autograd=False,
+                  selected_op_list_path=None):
     # cwrap depends on pyyaml, so we can't import it earlier
     root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
     sys.path.insert(0, root)
@@ -52,7 +53,8 @@ def generate_code(ninja_global=None,
             declarations_path or DECLARATIONS_PATH,
             jit_gen_dir,
             'tools/jit/templates',
-            disable_autograd=disable_autograd)
+            disable_autograd=disable_autograd,
+            selected_op_list_path=selected_op_list_path)
 
 
 def main():
@@ -71,6 +73,10 @@ def main():
         action='store_true',
         help='It can skip generating autograd related code when the flag is set',
     )
+    parser.add_argument(
+        '--selected-op-list-path',
+        help='Path to the yaml file that contains the list of operators to include for custom build.',
+    )
     options = parser.parse_args()
     generate_code(
         options.ninja_global,
@@ -79,6 +85,7 @@ def main():
         options.install_dir,
         options.subset,
         options.disable_autograd,
+        options.selected_op_list_path,
     )