[pytorch] deprecate static dispatch (pytorch#43564)

Summary:
Pull Request resolved: pytorch#43564

Static dispatch was originally introduced for mobile selective build.

Since we have added selective build support for dynamic dispatch and
tested it in FB production for months, we can deprecate static dispatch
to reduce the complexity of the codebase.

Test Plan: Imported from OSS

Reviewed By: ezyang

Differential Revision: D23324452

Pulled By: ljk53

fbshipit-source-id: d2970257616a8c6337f90249076fca1ae93090c7

Author: ljk53

.circleci/cimodel/data/simple/mobile_definitions.py

@@ -57,11 +57,6 @@ def gen_tree(self):
         [DOCKER_REQUIREMENT_ASAN],
         ["build"]
     ),
-    MobileJob(
-        DOCKER_IMAGE_ASAN,
-        [DOCKER_REQUIREMENT_ASAN],
-        ["custom", "build", "static"]
-    ),
 
     # Use LLVM-DEV toolchain in android-ndk-r19c docker image
     MobileJob(

.circleci/config.yml

@@ -6652,13 +6652,6 @@ workflows:
           name: pytorch_linux_xenial_py3_clang5_mobile_build
           requires:
             - docker-pytorch-linux-xenial-py3-clang5-asan
-      - pytorch_linux_build:
-          build_environment: pytorch-linux-xenial-py3-clang5-mobile-custom-build-static
-          build_only: "1"
-          docker_image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-py3-clang5-asan
-          name: pytorch_linux_xenial_py3_clang5_mobile_custom_build_static
-          requires:
-            - docker-pytorch-linux-xenial-py3-clang5-asan
       - pytorch_linux_build:
           build_environment: pytorch-linux-xenial-py3-clang5-mobile-custom-build-dynamic
           build_only: "1"

.jenkins/pytorch/build-mobile.sh

@@ -22,9 +22,7 @@ retry pip install --pre torch torchvision \
 
 # Run end-to-end process of building mobile library, linking into the predictor
 # binary, and running forward pass with a real model.
-if [[ "$BUILD_ENVIRONMENT" == *-mobile-custom-build-static* ]]; then
-  TEST_CUSTOM_BUILD_STATIC=1 test/mobile/custom_build/build.sh
-elif [[ "$BUILD_ENVIRONMENT" == *-mobile-custom-build-dynamic* ]]; then
+if [[ "$BUILD_ENVIRONMENT" == *-mobile-custom-build-dynamic* ]]; then
   export LLVM_DIR="$(llvm-config-5.0 --prefix)"
   echo "LLVM_DIR: ${LLVM_DIR}"
   TEST_CUSTOM_BUILD_DYNAMIC=1 test/mobile/custom_build/build.sh

BUILD.bazel

@@ -34,7 +34,6 @@ header_template_rule(
     substitutions = {
         "cmakedefine": "define",
         "#define FEATURE_TORCH_MOBILE": "/* #undef FEATURE_TORCH_MOBILE */",
-        "#define USE_STATIC_DISPATCH": "/* #undef USE_STATIC_DISPATCH */",
         "#define C10_USE_NUMA": "/* #undef C10_USE_NUMA */",
     },
 )

CMakeLists.txt

@@ -124,7 +124,6 @@ option(BUILD_PYTHON "Build Python binaries" ON)
 option(BUILD_CAFFE2_OPS "Build Caffe2 operators" ON)
 option(BUILD_SHARED_LIBS "Build libcaffe2.so" ON)
 option(BUILD_CAFFE2_MOBILE "Build libcaffe2 for mobile (deprecating)" OFF)
-option(USE_STATIC_DISPATCH "Use static dispatch for ATen operators" OFF)
 cmake_dependent_option(
     CAFFE2_LINK_LOCAL_PROTOBUF "If set, build protobuf inside libcaffe2.so." ON
     "BUILD_SHARED_LIBS AND BUILD_CUSTOM_PROTOBUF" OFF)

aten/src/ATen/autocast_mode.cpp

@@ -202,8 +202,6 @@ Tensor binary_cross_entropy_banned(const Tensor &, const Tensor &, const c10::op
            "safe to autocast.");
 }
 
-
-#ifndef USE_STATIC_DISPATCH
 namespace {
 /*****************************************************************************************************************
 This section performs load-time registration for autocast wrappers.
@@ -378,7 +376,6 @@ TORCH_LIBRARY_IMPL(aten, Autocast, m) {
 }
 
 }
-#endif
 
 } // namespace autocast
 } // namespace at

aten/src/ATen/cudnn/AutocastRNN.cpp

@@ -104,14 +104,12 @@ _cudnn_rnn_cast_reflatten(const Tensor & input,
 #endif // AT_CUDNN_ENABLED()
 }
 
-#ifndef USE_STATIC_DISPATCH
 namespace {
 TORCH_LIBRARY_IMPL(aten, Autocast, m) {
   m.impl("_cudnn_rnn",
          TORCH_FN((&at::autocast::_cudnn_rnn_cast_reflatten)));
 }
 } // anonymous namespace
-#endif
 
 } // namespace autocast
 } // namespace at

aten/src/ATen/function_wrapper.py

@@ -146,14 +146,10 @@ def TypedDict(name, attrs, total=True):  # type: ignore
 
 // ${schema_string}
 ${return_type} Tensor::${api_name}(${method_formals}) const {
-#ifdef USE_STATIC_DISPATCH
-    ${static_dispatch_method_body}
-#else
     static auto op = c10::Dispatcher::singleton()
         .findSchemaOrThrow("aten::${operator_name}", "${overload_name}")
         .typed<${tensor_method_cpp_signature}>();
     return op.call(${tensor_method_actuals});
-#endif
 }
 """)
 
@@ -172,45 +168,13 @@ def TypedDict(name, attrs, total=True):  # type: ignore
 
 // ${schema_string}
 ${return_type} ${api_name}(${formals}) {
-#ifdef USE_STATIC_DISPATCH
-    ${static_dispatch_function_body}
-#else
     static auto op = c10::Dispatcher::singleton()
         .findSchemaOrThrow("aten::${operator_name}", "${overload_name}")
         .typed<${function_cpp_signature}>();
     return op.call(${function_actuals});
-#endif
-}
-""")
-
-# In order to rely on the linker to strip unused ops, it requires us to dispatch statically
-# in Functions.h and TensorMethods.cpp.
-#
-# NB: The default body also needs to apply a variable guard, as in some
-# situations what we think is a default body actually does have an
-# explicit derivative, and thereby would have gotten unwrapped by
-# the time you get to the implementation.
-STATIC_DISPATCH_FUNCTION_DEFAULT_BODY = CodeTemplate("""\
-at::AutoNonVariableTypeMode _var_guard(true);
-${return_call} TypeDefault::${type_wrapper_name}(${actuals});
-""")
-
-STATIC_DISPATCH_FUNCTION_SWITCH_BODY = CodeTemplate("""\
-at::AutoNonVariableTypeMode _var_guard(true);
-${dispatch_key_init}
-switch (dispatchKeyToBackend(${dispatch_key_var_name})) {
-    ${static_dispatch_function_cases}
-    default:
-        AT_ERROR("${api_name} not implemented for ", at::toString(${dispatch_key_var_name}));
 }
 """)
 
-STATIC_DISPATCH_FUNCTION_SWITCH_CASE = CodeTemplate("""\
-case Backend::${backend}:
-    ${return_call} ${backend}Type::${type_wrapper_name}(${actuals});
-    break;
-""")
-
 IFDEF_BLOCK = CodeTemplate("""\
 #ifdef ${ifdef_guard}
 ${content}
@@ -246,10 +210,6 @@ def TypedDict(name, attrs, total=True):  # type: ignore
     ('ComplexDouble', 'ComplexDouble', 'ComplexDouble', False),
 ]
 
-static_dispatch_backends = ['CPU', 'QuantizedCPU', 'Vulkan']
-static_dispatch_backends_ifdef_guard = {'Vulkan' : 'USE_VULKAN'}
-
-
 class NYIError(Exception):
     """Indicates we don't support this declaration yet"""
 
@@ -1136,44 +1096,6 @@ def swizzle_self(f):  # blegh
 
             method_actuals = maybe_unwrap_optional_tensors(option, formals, option['method_actuals'])
 
-            if isinstance(type_method_dispatch, dict):
-                static_dispatch_function_cases = []
-                # NB: As this code is currently written, there will NEVER be
-                # a backend generated for variable dispatch.  There is nothing
-                # stopping us from actually implementing this, however, if you
-                # really wanted variable on mobile, there's nothing stopping
-                # you from implementing this (however, you would have an
-                # annoying phase problem, since code generation for variable
-                # happens in tools/ which happens later than here.)
-                #
-                # If you pass in a variable to the dispatch, and variable is
-                # enabled, this switch will fail.  This is intentional: you
-                # probably need to disable variable globally in the mobile
-                # calling code.
-                for backend in static_dispatch_backends:
-                    if backend in type_method_dispatch:
-                        static_dispatch_function_case = STATIC_DISPATCH_FUNCTION_SWITCH_CASE.substitute(
-                            option,
-                            backend=backend,
-                            backend_function=type_method_dispatch[backend],
-                            actuals=method_actuals)
-                        if (backend in static_dispatch_backends_ifdef_guard):
-                            static_dispatch_function_cases.append(IFDEF_BLOCK.substitute(
-                                option,
-                                ifdef_guard=static_dispatch_backends_ifdef_guard[backend],
-                                content=static_dispatch_function_case))
-                        else:
-                            static_dispatch_function_cases.append(static_dispatch_function_case)
-
-                static_dispatch_method_body = STATIC_DISPATCH_FUNCTION_SWITCH_BODY.substitute(
-                    option,
-                    dispatch_key_var_name=dispatch_key_var_name,
-                    dispatch_key_init=dispatch_key_init,
-                    static_dispatch_function_cases=static_dispatch_function_cases)
-            else:
-                static_dispatch_method_body = STATIC_DISPATCH_FUNCTION_DEFAULT_BODY.substitute(
-                    option, actuals=method_actuals)
-
             # See NOTE[UnboxedOnly]
             if option['use_c10_dispatcher'] == 'full':
                 tensor_method_actuals = option['schema_order_method_actuals']
@@ -1184,13 +1106,12 @@ def swizzle_self(f):  # blegh
                 tensor_method_cpp_signature = option['cpp_signature']
 
             method_definition = TENSOR_METHOD_DEFINITION.substitute(
-                option, static_dispatch_method_body=static_dispatch_method_body,
+                option,
                 tensor_method_actuals=tensor_method_actuals,
                 tensor_method_cpp_signature=tensor_method_cpp_signature
             )
             return FunctionCode(
-                declaration=TENSOR_METHOD_DECLARATION.substitute(
-                    option, static_dispatch_method_body=static_dispatch_method_body),
+                declaration=TENSOR_METHOD_DECLARATION.substitute(option),
                 definition=method_definition)
 
         def gen_namespace_function(option, multidispatch_formals):
@@ -1204,31 +1125,6 @@ def gen_namespace_function(option, multidispatch_formals):
 
             actuals = maybe_unwrap_optional_tensors(option, formals, option['actuals'])
 
-            if isinstance(type_method_dispatch, dict):
-                static_dispatch_function_cases = []
-                for backend in static_dispatch_backends:
-                    if backend in type_method_dispatch:
-                        static_dispatch_function_case = STATIC_DISPATCH_FUNCTION_SWITCH_CASE.substitute(
-                            option,
-                            backend=backend,
-                            backend_function=type_method_dispatch[backend],
-                            actuals=actuals)
-                        if (backend in static_dispatch_backends_ifdef_guard):
-                            static_dispatch_function_cases.append(IFDEF_BLOCK.substitute(
-                                option,
-                                ifdef_guard=static_dispatch_backends_ifdef_guard[backend],
-                                content=static_dispatch_function_case))
-                        else:
-                            static_dispatch_function_cases.append(static_dispatch_function_case)
-                static_dispatch_function_body = STATIC_DISPATCH_FUNCTION_SWITCH_BODY.substitute(
-                    option,
-                    dispatch_key_var_name=dispatch_key_var_name,
-                    dispatch_key_init=dispatch_key_init,
-                    static_dispatch_function_cases=static_dispatch_function_cases)
-            else:
-                static_dispatch_function_body = STATIC_DISPATCH_FUNCTION_DEFAULT_BODY.substitute(
-                    option, actuals=actuals)
-
             # See NOTE[UnboxedOnly]
             if option['use_c10_dispatcher'] == 'full':
                 function_actuals = option['schema_order_actuals']
@@ -1239,7 +1135,7 @@ def gen_namespace_function(option, multidispatch_formals):
                 function_cpp_signature = option['cpp_signature']
 
             fn_definition = FUNCTION_DEFINITION.substitute(
-                option, static_dispatch_function_body=static_dispatch_function_body,
+                option,
                 function_actuals=function_actuals,
                 function_cpp_signature=function_cpp_signature)
 

aten/src/ATen/native/TensorProperties.cpp

@@ -53,20 +53,14 @@ bool cudnn_is_acceptable(const Tensor& self) {
 }
 
 Tensor detach(const Tensor& self) {
-#ifndef USE_STATIC_DISPATCH
   // this just exists to give us a hook in VariableType and an entry in Declarations.yaml
   //AT_ERROR("detach is not implemented for Tensor");
-#endif
-  // this is no-op for USE_STATIC_DISPATCH mode
   return self;
 }
 
 Tensor & detach_(Tensor & self) {
-#ifndef USE_STATIC_DISPATCH
   // this just exists to give us a hook in VariableType and an entry in Declarations.yaml
   //AT_ERROR("detach_ is not implemented for Tensor");
-#endif
-  // this is no-op for USE_STATIC_DISPATCH mode
   return self;
 }
 

aten/src/ATen/templates/TensorMethods.cpp

@@ -12,15 +12,6 @@
 #include <ATen/quantized/Quantizer.h>
 #include <torch/csrc/WindowsTorchApiMacro.h>
 
-#ifdef USE_STATIC_DISPATCH
-#include <ATen/TypeDefault.h>
-#include <ATen/CPUType.h>
-#include <ATen/QuantizedCPUType.h>
-#ifdef USE_VULKAN
-#include <ATen/VulkanType.h>
-#endif
-#endif
-
 namespace at {
 
 Tensor Tensor::cpu() const {

c10/macros/cmake_macros.h.in

@@ -14,9 +14,4 @@
 // to converging libtorch and caffe2 mobile builds and removing it eventually.
 #cmakedefine FEATURE_TORCH_MOBILE
 
-// If defined it will use static dispatch for ATen operators.
-// Should expose this macro for projects including ATen headers to inherient
-// the same option.
-#cmakedefine USE_STATIC_DISPATCH
-
 #endif // C10_MACROS_CMAKE_MACROS_H_

caffe2/core/macros.h.in

@@ -79,5 +79,4 @@ static_assert(
   {"USE_MKLDNN", "${CAFFE2_USE_MKLDNN}"}, \
   {"USE_NVTX", "${CAFFE2_USE_NVTX}"}, \
   {"USE_TRT", "${CAFFE2_USE_TRT}"}, \
-  {"USE_STATIC_DISPATCH", "${USE_STATIC_DISPATCH}"},   \
 }

cmake/Codegen.cmake

@@ -167,7 +167,7 @@ if(INTERN_BUILD_ATEN_OPS)
   endif()
 
   if(SELECTED_OP_LIST)
-    if(NOT USE_STATIC_DISPATCH AND NOT OP_DEPENDENCY)
+    if(NOT OP_DEPENDENCY)
       message(INFO "Use default op dependency graph .yaml file for custom build with dynamic dispatch.")
       set(OP_DEPENDENCY ${CMAKE_CURRENT_LIST_DIR}/../tools/code_analyzer/default_op_deps.yaml)
     endif()

cmake/Summary.cmake

@@ -21,7 +21,6 @@ function(caffe2_print_configuration_summary)
   message(STATUS "  TORCH_VERSION         : ${TORCH_VERSION}")
   message(STATUS "  CAFFE2_VERSION        : ${CAFFE2_VERSION}")
   message(STATUS "  BUILD_CAFFE2_MOBILE   : ${BUILD_CAFFE2_MOBILE}")
-  message(STATUS "  USE_STATIC_DISPATCH   : ${USE_STATIC_DISPATCH}")
   message(STATUS "  BUILD_BINARY          : ${BUILD_BINARY}")
   message(STATUS "  BUILD_CUSTOM_PROTOBUF : ${BUILD_CUSTOM_PROTOBUF}")
   if(${CAFFE2_LINK_LOCAL_PROTOBUF})

scripts/build_android.sh

@@ -61,7 +61,6 @@ CMAKE_ARGS=()
 
 if [ -z "${BUILD_CAFFE2_MOBILE:-}" ]; then
   # Build PyTorch mobile
-  CMAKE_ARGS+=("-DUSE_STATIC_DISPATCH=ON")
   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")

scripts/build_ios.sh

@@ -13,7 +13,6 @@ CMAKE_ARGS=()
 
 if [ -z "${BUILD_CAFFE2_MOBILE:-}" ]; then
   # Build PyTorch mobile
-  CMAKE_ARGS+=("-DUSE_STATIC_DISPATCH=ON")
   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")
@@ -62,7 +61,7 @@ fi
 
 # IOS_PLATFORM controls type of iOS platform (see ios-cmake)
 if [ -n "${IOS_PLATFORM:-}" ]; then
-  CMAKE_ARGS+=("-DIOS_PLATFORM=${IOS_PLATFORM}")    
+  CMAKE_ARGS+=("-DIOS_PLATFORM=${IOS_PLATFORM}")
   if [ "${IOS_PLATFORM}" == "WATCHOS" ]; then
       # enable bitcode by default for watchos
       CMAKE_ARGS+=("-DCMAKE_C_FLAGS=-fembed-bitcode")

scripts/build_mobile.sh

@@ -15,7 +15,6 @@ echo "Bash: $(/bin/bash --version | head -1)"
 echo "Caffe2 path: $CAFFE2_ROOT"
 
 CMAKE_ARGS=()
-CMAKE_ARGS+=("-DUSE_STATIC_DISPATCH=ON")
 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")

test/mobile/custom_build/build.sh

@@ -4,23 +4,18 @@
 # size for mobile devices and the flow to integrate it with a simple predictor
 # in c++.
 #
-# There are three custom build types:
+# Supported custom build types:
 #
 # 1. `TEST_DEFAULT_BUILD=1 ./build.sh` - it is similar to the prebuilt libtorch
 # libraries released for Android and iOS (same CMake build options + host
 # toolchain), which doesn't contain autograd function nor backward ops thus is
 # smaller than full LibTorch.
 #
-# 2. `TEST_CUSTOM_BUILD_STATIC=1 ./build.sh` - it further optimizes libtorch
+# 2. `TEST_CUSTOM_BUILD_DYNAMIC=1 ./build.sh` - it further optimizes libtorch
 # size by only including ops used by a specific model.
-#
-# 3. `TEST_CUSTOM_BUILD_DYNAMIC=1 ./build.sh` - similar as 2) except that it
-# relies on the op dependency graph (instead of static dispatch) to calculate
-# and keep all transitively dependent ops by the model.
 # Note that LLVM_DIR environment variable should be set to the location of
 # LLVM-dev toolchain.
 #
-# Type 2) will be deprecated by type 3) in the future.
 ###############################################################################
 
 set -ex -o pipefail
@@ -59,25 +54,13 @@ run_default_build() {
     "${SRC_ROOT}/scripts/build_mobile.sh"
 }
 
-run_custom_build_with_static_dispatch() {
-  LIBTORCH_BUILD_ROOT="${BUILD_ROOT}/build_custom_libtorch_static"
-  LIBTORCH_INSTALL_PREFIX="${LIBTORCH_BUILD_ROOT}/install"
-
-  BUILD_ROOT="${LIBTORCH_BUILD_ROOT}" \
-    "${SRC_ROOT}/scripts/build_mobile.sh" \
-    -DCMAKE_CXX_FLAGS="-DSTRIP_ERROR_MESSAGES" \
-    -DUSE_STATIC_DISPATCH=ON \
-    -DSELECTED_OP_LIST="${ROOT_OPS}"
-}
-
 run_custom_build_with_dynamic_dispatch() {
   LIBTORCH_BUILD_ROOT="${BUILD_ROOT}/build_custom_libtorch_dynamic"
   LIBTORCH_INSTALL_PREFIX="${LIBTORCH_BUILD_ROOT}/install"
 
   BUILD_ROOT="${LIBTORCH_BUILD_ROOT}" \
     "${SRC_ROOT}/scripts/build_mobile.sh" \
     -DCMAKE_CXX_FLAGS="-DSTRIP_ERROR_MESSAGES" \
-    -DUSE_STATIC_DISPATCH=OFF \
     -DSELECTED_OP_LIST="${ROOT_OPS}" \
     -DOP_DEPENDENCY="${OP_DEPENDENCY}"
 }
@@ -115,13 +98,6 @@ test_default_build() {
   run_predictor
 }
 
-test_custom_build_with_static_dispatch() {
-  prepare_model_and_dump_root_ops
-  run_custom_build_with_static_dispatch
-  build_predictor
-  run_predictor
-}
-
 test_custom_build_with_dynamic_dispatch() {
   prepare_model_and_dump_root_ops
   generate_op_dependency_graph
@@ -134,10 +110,6 @@ if [ -n "${TEST_DEFAULT_BUILD}" ]; then
   test_default_build
 fi
 
-if [ -n "${TEST_CUSTOM_BUILD_STATIC}" ]; then
-  test_custom_build_with_static_dispatch
-fi
-
 if [ -n "${TEST_CUSTOM_BUILD_DYNAMIC}" ]; then
   test_custom_build_with_dynamic_dispatch
 fi

tools/code_analyzer/build.sh

@@ -62,7 +62,6 @@ build_torch_mobile() {
 
   BUILD_ROOT="${TORCH_BUILD_ROOT}" "${SRC_ROOT}/scripts/build_mobile.sh" \
     -DCMAKE_CXX_FLAGS="-S -emit-llvm -DSTRIP_ERROR_MESSAGES" \
-    -DUSE_STATIC_DISPATCH=OFF \
     ${MOBILE_BUILD_FLAGS}
 }
 

tools/jit/gen_unboxing_wrappers.py

@@ -166,11 +166,7 @@ def from_ivalue(arg, value):
         .layout(${layout})
         .device(${device})
         .pinned_memory(${pin_memory});
-#ifdef USE_STATIC_DISPATCH
-    auto result_ = at::${name}(${args_with_tensor_options});
-#else
     auto result_ = torch::${name}(${args_with_tensor_options});
-#endif
 """)
 CALL_METHOD_WITH_TENSOR_OPTIONS = CodeTemplate("""\
 const auto options = TensorOptions()

torch/csrc/jit/runtime/register_c10_ops.cpp

@@ -122,14 +122,7 @@ Operator createOperatorFromC10_withTracingHandledHere(
       jit::tracer::setTracingState(nullptr);
     }
 
-#ifdef USE_STATIC_DISPATCH
-    {
-      at::AutoNonVariableTypeMode non_var_type_mode(true);
-      op.callBoxed(stack);
-    }
-#else
     op.callBoxed(stack);
-#endif // USE_STATIC_DISPATCH
 
     if (tracer_state) {
       jit::tracer::setTracingState(std::move(tracer_state));