Require explicit inclusion of custom op support in a minimal build (#6663)

* Remove support from custom ops from the base minimal build as they contribute too much binary growth to an Android build. Add ability to explicitly enable custom op support in a minimal build. Change one minimal build CI to test adding custom op support (unit tests are run in that build to validate)
2026-07-30 20:18:08 +00:00 · 2021-02-13 12:42:33 +10:00 · 2021-02-13 12:42:33 +10:00 · 25f7c93504
commit 25f7c93504
parent dd50c39ac6
13 changed files with 89 additions and 44 deletions
--- a/cmake/CMakeLists.txt
+++ b/cmake/CMakeLists.txt
@ -123,6 +123,7 @@ option(onnxruntime_DISABLE_RTTI "Disable RTTI" OFF)
 option(onnxruntime_DISABLE_EXCEPTIONS "Disable exception handling. Requires onnxruntime_MINIMAL_BUILD currently." OFF)
 option(onnxruntime_MINIMAL_BUILD "Exclude as much as possible from the build. Support ORT format models. No support for ONNX format models." OFF)
 option(onnxruntime_EXTENDED_MINIMAL_BUILD "onnxruntime_MINIMAL_BUILD with support for execution providers that compile kernels." OFF)
+option(onnxruntime_MINIMAL_BUILD_CUSTOM_OPS "Add custom operator kernels support to a minimal build." OFF)
 option(onnxruntime_REDUCED_OPS_BUILD "Reduced set of kernels are registered in build via modification of the kernel registration source files." OFF)
 option(onnxruntime_DISABLE_ORT_FORMAT_LOAD "Disable loading an ORT format model when onnxruntime_MINIMAL_BUILD=OFF (i.e. in a full build)." OFF)

@ -236,6 +237,10 @@ if (onnxruntime_MINIMAL_BUILD)
    add_compile_definitions(ORT_EXTENDED_MINIMAL_BUILD)
  endif()

+  if (onnxruntime_MINIMAL_BUILD_CUSTOM_OPS)
+    add_compile_definitions(ORT_MINIMAL_BUILD_CUSTOM_OPS)
+  endif()
+  
  set(onnxruntime_REDUCED_OPS_BUILD ON)

  if (NOT onnxruntime_ENABLE_PYTHON)
--- a/cmake/onnxruntime_framework.cmake
+++ b/cmake/onnxruntime_framework.cmake
@ -14,6 +14,14 @@ if (onnxruntime_MINIMAL_BUILD)
    "${ONNXRUNTIME_ROOT}/core/framework/fallback_cpu_capability.cc"
  )

+  # custom ops support must be explicitly enabled in a minimal build. exclude if not.
+  if (NOT onnxruntime_MINIMAL_BUILD_CUSTOM_OPS)
+    list(APPEND onnxruntime_framework_src_exclude
+      "${ONNXRUNTIME_INCLUDE_DIR}/core/framework/customregistry.h"
+      "${ONNXRUNTIME_ROOT}/core/framework/customregistry.cc"
+    )
+  endif()
+
  list(REMOVE_ITEM onnxruntime_framework_srcs ${onnxruntime_framework_src_exclude})
 endif()

--- a/onnxruntime/core/framework/kernel_registry_manager.cc
+++ b/onnxruntime/core/framework/kernel_registry_manager.cc
@ -45,12 +45,14 @@ Status KernelRegistryManager::RegisterKernels(const ExecutionProviders& executio
  return Status::OK();
 }

+#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
 void KernelRegistryManager::RegisterKernelRegistry(std::shared_ptr<KernelRegistry> kernel_registry) {
  if (nullptr == kernel_registry) {
    return;
  }
  custom_kernel_registries_.push_front(kernel_registry);
 }
+#endif

 #if !defined(ORT_MINIMAL_BUILD)
 bool KernelRegistryManager::HasImplementationOf(const KernelRegistryManager& r, const Node& node, const std::string& provider_type) {
@ -84,12 +86,14 @@ Status KernelRegistryManager::SearchKernelRegistry(const onnxruntime::Node& node
    return Status(ONNXRUNTIME, FAIL, create_error_message("The node is not placed on any Execution Provider. "));
  }

+#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
  for (auto& registry : custom_kernel_registries_) {
    status = registry->TryFindKernel(node, std::string(), kernel_def_hash, kernel_create_info);
    if (status.IsOK()) {
      return status;
    }
  }
+#endif

  KernelRegistry* p = nullptr;
  auto iter = provider_type_to_registry_.find(ptype);
--- a/onnxruntime/core/framework/kernel_registry_manager.h
+++ b/onnxruntime/core/framework/kernel_registry_manager.h
@ -32,6 +32,7 @@ class KernelRegistryManager {
  // Register kernels from providers
  Status RegisterKernels(const ExecutionProviders& execution_providers) ORT_MUST_USE_RESULT;

+#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
  // The registry passed in this function has highest priority than anything already in this KernelRegistryManager,
  // and anything registered from RegisterKernels
  // For example, if you do:
@ -41,7 +42,6 @@ class KernelRegistryManager {
  // Then B > A > providers
  void RegisterKernelRegistry(std::shared_ptr<KernelRegistry> kernel_registry);

-#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
  /**
   * Search kernel registry by provider type.
   * @param type provider type string
--- a/onnxruntime/core/session/custom_ops.cc
+++ b/onnxruntime/core/session/custom_ops.cc
@ -5,7 +5,6 @@
 #pragma warning(disable : 4267)
 #endif

-#include "core/framework/customregistry.h"
 #include "core/framework/data_types.h"
 #include "core/framework/op_kernel_info.h"
 #include "core/framework/op_kernel_context_internal.h"
@ -15,8 +14,6 @@
 #include "core/session/inference_session.h"
 #include "core/session/ort_apis.h"

-// ONNXTensorElementDataType MLDataTypeToOnnxRuntimeTensorElementDataType(const onnxruntime::DataTypeImpl* cpp_type);
-
 ORT_API_STATUS_IMPL(OrtApis::KernelInfoGetAttribute_float, _In_ const OrtKernelInfo* info, _In_ const char* name, _Out_ float* out) {
  auto status = reinterpret_cast<const onnxruntime::OpKernelInfo*>(info)->GetAttr<float>(name, out);
  if (status.IsOK())
@ -69,6 +66,8 @@ ORT_API_STATUS_IMPL(OrtApis::KernelInfoGetAttribute_string, _In_ const OrtKernel
  return onnxruntime::ToOrtStatus(status);
 }

+#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
+#include "core/framework/customregistry.h"
 namespace onnxruntime {

 struct CustomOpKernel : OpKernel {
@ -178,3 +177,5 @@ common::Status CreateCustomRegistry(const std::vector<OrtCustomOpDomain*>& op_do
 }

 }  // namespace onnxruntime
+#endif  // !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
+
--- a/onnxruntime/core/session/inference_session.cc
+++ b/onnxruntime/core/session/inference_session.cc
@ -14,7 +14,6 @@
 #include "core/common/denormal.h"
 #include "core/common/logging/logging.h"
 #include "core/framework/allocatormgr.h"
-#include "core/framework/customregistry.h"
 #include "core/framework/error_code_helper.h"
 #include "core/framework/execution_frame.h"
 #include "core/framework/feeds_fetches_manager.h"
@ -44,7 +43,6 @@
 #ifdef USE_DML  // TODO: This is necessary for the workaround in TransformGraph
 #include "core/providers/dml/DmlExecutionProvider/src/GraphTransformer.h"
 #endif
-#include "core/session/custom_ops.h"
 #include "core/session/environment.h"
 #include "core/session/IOBinding.h"
 #include "core/session/inference_session_utils.h"
@ -52,6 +50,12 @@
 #include "core/util/protobuf_parsing_utils.h"
 #include "core/util/thread_utils.h"

+// custom ops are not available in a minimal build unless ORT_MINIMAL_BUILD_CUSTOM_OPS is set
+#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
+#include "core/framework/customregistry.h"
+#include "core/session/custom_ops.h"
+#endif
+
 using namespace ONNX_NAMESPACE;
 using namespace onnxruntime::experimental;
 using namespace onnxruntime::common;
@ -448,8 +452,33 @@ common::Status InferenceSession::RegisterExecutionProvider(std::unique_ptr<IExec
  return execution_providers_.Add(provider_type, std::move(p_exec_provider));
 }

-#if !defined(ORT_MINIMAL_BUILD)
+// Custom Op support
+#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
+common::Status InferenceSession::AddCustomOpDomains(const std::vector<OrtCustomOpDomain*>& op_domains) {
+  std::shared_ptr<CustomRegistry> custom_registry;
+  ORT_RETURN_IF_ERROR_SESSIONID_(CreateCustomRegistry(op_domains, custom_registry));
+  ORT_RETURN_IF_ERROR_SESSIONID_(RegisterCustomRegistry(custom_registry));
+  return Status::OK();
+}

+common::Status InferenceSession::RegisterCustomRegistry(std::shared_ptr<CustomRegistry> custom_registry) {
+  if (custom_registry == nullptr) {
+    return Status(common::ONNXRUNTIME, common::FAIL, "Received nullptr for custom registry");
+  }
+
+  custom_registries_.push_back(custom_registry);
+
+  // Insert session-level customized kernel registry.
+  kernel_registry_manager_.RegisterKernelRegistry(custom_registry->GetKernelRegistry());
+
+#if !defined(ORT_MINIMAL_BUILD)
+  custom_schema_registries_.push_back(custom_registry->GetOpschemaRegistry());
+#endif
+  return Status::OK();
+}
+#endif  // !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
+
+#if !defined(ORT_MINIMAL_BUILD)
 common::Status InferenceSession::RegisterGraphTransformer(
    std::unique_ptr<onnxruntime::GraphTransformer> p_graph_transformer, TransformerLevel level) {
  if (p_graph_transformer == nullptr) {
@ -473,32 +502,7 @@ common::Status InferenceSession::AddCustomTransformerList(const std::vector<std:

  return Status::OK();
 }
-#endif

-common::Status InferenceSession::AddCustomOpDomains(const std::vector<OrtCustomOpDomain*>& op_domains) {
-  std::shared_ptr<CustomRegistry> custom_registry;
-  ORT_RETURN_IF_ERROR_SESSIONID_(CreateCustomRegistry(op_domains, custom_registry));
-  ORT_RETURN_IF_ERROR_SESSIONID_(RegisterCustomRegistry(custom_registry));
-  return Status::OK();
-}
-
-common::Status InferenceSession::RegisterCustomRegistry(std::shared_ptr<CustomRegistry> custom_registry) {
-  if (custom_registry == nullptr) {
-    return Status(common::ONNXRUNTIME, common::FAIL, "Received nullptr for custom registry");
-  }
-
-  custom_registries_.push_back(custom_registry);
-
-  // Insert session-level customized kernel registry.
-  kernel_registry_manager_.RegisterKernelRegistry(custom_registry->GetKernelRegistry());
-
-#if !defined(ORT_MINIMAL_BUILD)
-  custom_schema_registries_.push_back(custom_registry->GetOpschemaRegistry());
-#endif
-  return Status::OK();
-}
-
-#if !defined(ORT_MINIMAL_BUILD)
 common::Status InferenceSession::SaveToOrtFormat(const std::basic_string<ORTCHAR_T>& filepath) const {
  ORT_RETURN_IF_NOT(FLATBUFFERS_LITTLEENDIAN, "ort format only supports little-edian machines");

--- a/onnxruntime/core/session/inference_session.h
+++ b/onnxruntime/core/session/inference_session.h
@ -186,6 +186,8 @@ class InferenceSession {
  common::Status AddCustomTransformerList(const std::vector<std::string>& transformers_to_enable) ORT_MUST_USE_RESULT;

 #endif  // !defined(ORT_MINIMAL_BUILD)
+
+#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
  /**
    * Add custom ops. This API is not thread safe.
    */
@ -201,6 +203,7 @@ class InferenceSession {
    * @return OK if success.
    */
  common::Status RegisterCustomRegistry(std::shared_ptr<CustomRegistry> custom_registry) ORT_MUST_USE_RESULT;
+#endif  // !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)

  /**
    * Load an ONNX or ORT format model.
@ -588,9 +591,12 @@ class InferenceSession {
  std::list<std::shared_ptr<onnxruntime::IOnnxRuntimeOpSchemaCollection>> custom_schema_registries_;
 #endif

+#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
+
  //CustomRegistry objects own the corresponding KernelRegistry and OnnxRuntimeOpSchemaRegistry objects.
  //So its lifetime should be same as its constituents. This vector is to extend the lifetime of the owner.
  std::vector<std::shared_ptr<CustomRegistry>> custom_registries_;
+#endif

  ModelMetadata model_metadata_;
  std::unordered_set<std::string> required_inputs_;
@ -654,7 +660,7 @@ class InferenceSession {
  // Longer term we may want to directly refer to offsets in this buffer for initializers so we don't need to copy
  // those into new OrtValue instances, at which point we won't free them until the InferenceSession goes away.
  std::vector<uint8_t> ort_format_model_bytes_;
-  
+
  std::shared_ptr<onnxruntime::AllocatorManager> allocator_manager_;
 };

--- a/onnxruntime/core/session/onnxruntime_c_api.cc
+++ b/onnxruntime/core/session/onnxruntime_c_api.cc
@ -439,11 +439,12 @@ static ORT_STATUS_PTR CreateSessionAndLoadModel(_In_ const OrtSessionOptions* op
        env->GetEnvironment());
  }

+#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
  // Add custom domains
-  Status status;
  if (options && !options->custom_op_domains_.empty()) {
    ORT_API_RETURN_IF_STATUS_NOT_OK(sess->AddCustomOpDomains(options->custom_op_domains_));
  }
+#endif

  // Finish load
  if (load_config_from_model) {
--- a/onnxruntime/test/shared_lib/test_ort_format_models.cc
+++ b/onnxruntime/test/shared_lib/test_ort_format_models.cc
@ -4,6 +4,9 @@
 // if we can't load an ORT format model we can't really test anything
 #if defined(ENABLE_ORT_FORMAT_LOAD)

+// custom ops are only supported in a minimal build if explicitly enabled
+#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
+
 #include "core/common/make_unique.h"
 #include "core/graph/constants.h"
 #include "core/session/onnxruntime_cxx_api.h"
@ -131,4 +134,6 @@ TEST(OrtFormatCustomOpTests, LoadOrtModel) {
 }
 #endif

+#endif  // !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
+
 #endif  // #if defined(ENABLE_ORT_FORMAT_LOAD)
--- a/tools/ci_build/build.py
+++ b/tools/ci_build/build.py
@ -439,13 +439,15 @@ def parse_arguments():
        help="Build ONNXRuntime micro-benchmarks.")

    # options to reduce binary size
-    parser.add_argument("--minimal_build", action='store',
-                        const='on', default='off', nargs='?', type=str.lower,
+    parser.add_argument("--minimal_build", default=None, nargs='*', type=str.lower,
                        help="Create a build that only supports ORT format models. "
                        "See /docs/ONNX_Runtime_Format_Model_Usage.md for more information. "
                        "RTTI is automatically disabled in a minimal build. "
                        "To enable execution providers that compile kernels at runtime (e.g. NNAPI) pass 'extended' "
-                        "as a parameter. e.g. '--minimal_build extended'.")
+                        "as a parameter. e.g. '--minimal_build extended'. "
+                        "To enable support for custom operators pass 'custom_ops' as a parameter. "
+                        "e.g. '--minimal_build custom_ops'. This can be combined with an 'extended' build by passing "
+                        "'--minimal_build extended custom_ops'")
    parser.add_argument("--include_ops_by_config", type=str,
                        help="include ops from config file. "
                        "See /docs/Reduced_Operator_Kernel_build.md for more information.")
@ -681,8 +683,11 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
        "-Donnxruntime_DISABLE_RTTI=" + ("ON" if args.disable_rtti else "OFF"),
        "-Donnxruntime_DISABLE_EXCEPTIONS=" + ("ON" if args.disable_exceptions else "OFF"),
        "-Donnxruntime_DISABLE_ORT_FORMAT_LOAD=" + ("ON" if args.disable_ort_format_load else "OFF"),
-        "-Donnxruntime_MINIMAL_BUILD=" + ("ON" if args.minimal_build != 'off' else "OFF"),
-        "-Donnxruntime_EXTENDED_MINIMAL_BUILD=" + ("ON" if args.minimal_build == 'extended' else "OFF"),
+        "-Donnxruntime_MINIMAL_BUILD=" + ("ON" if args.minimal_build is not None else "OFF"),
+        "-Donnxruntime_EXTENDED_MINIMAL_BUILD=" + ("ON" if args.minimal_build and 'extended' in args.minimal_build
+                                                   else "OFF"),
+        "-Donnxruntime_MINIMAL_BUILD_CUSTOM_OPS=" + ("ON" if args.minimal_build and 'custom_ops' in args.minimal_build
+                                                     else "OFF"),
        "-Donnxruntime_REDUCED_OPS_BUILD=" + ("ON" if args.include_ops_by_config else "OFF"),
        "-Donnxruntime_MSVC_STATIC_RUNTIME=" + (
            "ON" if args.enable_msvc_static_runtime else "OFF"),
@ -1376,7 +1381,7 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs):

            # Disable python tests in a reduced build as we don't know which ops have been included and which
            # models can run
-            if args.include_ops_by_config or args.minimal_build != 'off':
+            if args.include_ops_by_config or args.minimal_build:
                return

            if is_windows():
--- a/tools/ci_build/github/azure-pipelines/linux-cpu-minimal-build-ci-pipeline.yml
+++ b/tools/ci_build/github/azure-pipelines/linux-cpu-minimal-build-ci-pipeline.yml
@ -80,7 +80,7 @@ jobs:
      workingDirectory: $(Build.SourcesDirectory)

  - task: CmdLine@2
-    displayName: Build minimal onnxruntime [exceptions ENABLED, type reduction DISABLED] and run tests
+    displayName: Build minimal onnxruntime [exceptions ENABLED, type reduction DISABLED, custom ops ENABLED] and run tests
    inputs:
      script: |
        docker run --rm \
@ -93,7 +93,8 @@ jobs:
          onnxruntimecentoscpubuild \
            /bin/bash /onnxruntime_src/tools/ci_build/github/linux/ort_minimal/build_minimal_ort_and_run_tests.sh \
              --build-directory /build/without_type_reduction \
-              --reduced-ops-config /home/onnxruntimedev/.test_data/required_ops.ort_models.config
+              --reduced-ops-config /home/onnxruntimedev/.test_data/required_ops.ort_models.config \
+              --enable-custom-ops
      workingDirectory: $(Build.SourcesDirectory)

  - script: git checkout -- .
--- a/tools/ci_build/github/linux/ort_minimal/build_minimal_ort_and_run_tests.sh
+++ b/tools/ci_build/github/linux/ort_minimal/build_minimal_ort_and_run_tests.sh
@ -19,6 +19,7 @@ USAGE_TEXT="Usage:
 BUILD_DIR=
 REDUCED_OPS_CONFIG_FILE=
 ENABLE_TYPE_REDUCTION=
+MINIMAL_BUILD_ARGS=

 while [[ $# -gt 0 ]]
 do
@ -38,6 +39,10 @@ do
            ENABLE_TYPE_REDUCTION=1
            shift
            ;;
+        --enable-custom-ops)
+            MINIMAL_BUILD_ARGS="custom_ops"
+            shift
+            ;;
        *)
            echo "Invalid option: $1"
            echo "$USAGE_TEXT"
@ -59,7 +64,7 @@ python3 /onnxruntime_src/tools/ci_build/build.py \
    --skip_submodule_sync \
    --build_shared_lib \
    --parallel \
-    --minimal_build \
+    --minimal_build ${MINIMAL_BUILD_ARGS} \
    --disable_ml_ops \
    --include_ops_by_config ${REDUCED_OPS_CONFIG_FILE} \
    ${ENABLE_TYPE_REDUCTION:+"--enable_reduced_operator_type_support"}
--- a/tools/ci_build/github/linux/ort_minimal/readelf_utils.py
+++ b/tools/ci_build/github/linux/ort_minimal/readelf_utils.py
@ -65,7 +65,7 @@ def diff_sections_total_size(base_binary_path, binary_path, readelf_path='readel
    total = 0
    results = collections.OrderedDict()

-    for section in merged_keys:
+    for section in sorted(merged_keys):
        base_size = base_section_sizes[section] if section in base_section_sizes else 0
        size = section_sizes[section] if section in section_sizes else 0