Provide an option to disable contrib ops. (#707)

cmake/CMakeLists.txt

@@ -70,6 +70,7 @@ option(onnxruntime_USE_TENSORRT "Build with TensorRT support" OFF)
 option(onnxruntime_ENABLE_LTO "Enable link time optimization, which is not stable on older GCCs" OFF)
 option(onnxruntime_CROSS_COMPILING "Cross compiling onnx runtime" OFF)
 option(onnxruntime_USE_FULL_PROTOBUF "Use full protobuf" OFF)
+option(onnxruntime_DISABLE_CONTRIB_OPS "Disable contrib ops" OFF)
 
 set(protobuf_BUILD_TESTS OFF CACHE BOOL "Build protobuf tests" FORCE)
 #nsync tests failed on Mac Build
@@ -206,7 +207,7 @@ add_subdirectory(${PROJECT_SOURCE_DIR}/external/nsync EXCLUDE_FROM_ALL)
 list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/external)
 
 #Here we support two build mode:
-#1. if ONNX_CUSTOM_PROTOC_EXECUTABLE is set, build Protobuf from source, except protoc.exe. This mode is mainly 
+#1. if ONNX_CUSTOM_PROTOC_EXECUTABLE is set, build Protobuf from source, except protoc.exe. This mode is mainly
 #   for cross-compiling
 #2. if ONNX_CUSTOM_PROTOC_EXECUTABLE is not set, Compile everything(including protoc) from source code.
 
@@ -230,6 +231,10 @@ if (onnxruntime_USE_FULL_PROTOBUF)
   add_definitions(-DUSE_FULL_PROTOBUF)
 endif()
 
+if (onnxruntime_DISABLE_CONTRIB_OPS)
+  add_definitions(-DDISABLE_CONTRIB_OPS)
+endif()
+
 if (onnxruntime_USE_CUDA AND "${onnxruntime_CUDNN_HOME}" STREQUAL "")
   message(FATAL_ERROR "onnxruntime_CUDNN_HOME required for onnxruntime_USE_CUDA")
 endif()
@@ -507,7 +512,7 @@ endif()
 
 if (onnxruntime_USE_TENSORRT)
   if (WIN32)
-    set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} /DELAYLOAD:nvinfer.dll") 
+    set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} /DELAYLOAD:nvinfer.dll")
   endif()
 endif()
 

cmake/onnxruntime_graph.cmake

@@ -2,10 +2,17 @@
 # Licensed under the MIT License.
 
 file(GLOB_RECURSE onnxruntime_graph_src
-    "${ONNXRUNTIME_INCLUDE_DIR}/core/graph/*.h"
-    "${ONNXRUNTIME_ROOT}/core/graph/*.h"
-    "${ONNXRUNTIME_ROOT}/core/graph/*.cc"
-)
+  "${ONNXRUNTIME_INCLUDE_DIR}/core/graph/*.h"
+  "${ONNXRUNTIME_ROOT}/core/graph/*.h"
+  "${ONNXRUNTIME_ROOT}/core/graph/*.cc"
+  )
+
+if (onnxruntime_DISABLE_CONTRIB_OPS)
+  list(REMOVE_ITEM onnxruntime_graph_src
+    "${ONNXRUNTIME_ROOT}/core/graph/contrib_ops/*.h"
+    "${ONNXRUNTIME_ROOT}/core/graph/contrib_ops/*.cc"
+    )
+endif()
 
 file(GLOB_RECURSE onnxruntime_ir_defs_src
     "${ONNXRUNTIME_ROOT}/core/defs/*.cc"

cmake/onnxruntime_providers.cmake

@@ -33,7 +33,14 @@ endif()
 source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_providers_common_srcs} ${onnxruntime_providers_srcs})
 # add using ONNXRUNTIME_ROOT so they show up under the 'contrib_ops' folder in Visual Studio
 source_group(TREE ${ONNXRUNTIME_ROOT} FILES ${onnxruntime_contrib_ops_srcs})
-add_library(onnxruntime_providers ${onnxruntime_providers_common_srcs} ${onnxruntime_providers_srcs} ${onnxruntime_contrib_ops_srcs})
+
+# disable contrib ops conditionally
+if(onnxruntime_DISABLE_CONTRIB_OPS)
+  add_library(onnxruntime_providers ${onnxruntime_providers_common_srcs} ${onnxruntime_providers_srcs})
+else()
+  add_library(onnxruntime_providers ${onnxruntime_providers_common_srcs} ${onnxruntime_providers_srcs} ${onnxruntime_contrib_ops_srcs})
+endif()
+
 onnxruntime_add_include_to_target(onnxruntime_providers onnxruntime_common onnxruntime_framework gsl onnx onnx_proto protobuf::libprotobuf)
 set(gemmlowp_src ${ONNXRUNTIME_ROOT}/../cmake/external/gemmlowp)
 set(re2_src ${ONNXRUNTIME_ROOT}/../cmake/external/re2)

cmake/onnxruntime_unittests.cmake

@@ -111,13 +111,16 @@ if(onnxruntime_USE_CUDA)
 endif()
 
 set(onnxruntime_test_providers_src_patterns
-  "${TEST_SRC_DIR}/contrib_ops/*.h"
-  "${TEST_SRC_DIR}/contrib_ops/*.cc"
   "${TEST_SRC_DIR}/providers/*.h"
   "${TEST_SRC_DIR}/providers/*.cc"
   "${TEST_SRC_DIR}/framework/TestAllocatorManager.cc"
   "${TEST_SRC_DIR}/framework/TestAllocatorManager.h"
   )
+if(NOT onnxruntime_DISABLE_CONTRIB_OPS)
+  list(APPEND onnxruntime_test_providers_src_patterns
+    "${TEST_SRC_DIR}/contrib_ops/*.h"
+    "${TEST_SRC_DIR}/contrib_ops/*.cc")
+endif()
 
 file(GLOB onnxruntime_test_providers_src ${onnxruntime_test_providers_src_patterns})
 file(GLOB_RECURSE onnxruntime_test_providers_cpu_src

onnxruntime/core/framework/environment.cc

@@ -4,10 +4,12 @@
 #include "core/framework/environment.h"
 #include "core/framework/allocatormgr.h"
 #include "core/graph/constants.h"
-#include "core/graph/contrib_ops/contrib_defs.h"
 #include "core/graph/op.h"
 #include "onnx/defs/operator_sets.h"
 #include "onnx/defs/operator_sets-ml.h"
+#ifndef DISABLE_CONTRIB_OPS
+#include "core/graph/contrib_ops/contrib_defs.h"
+#endif
 
 namespace onnxruntime {
 using namespace ::onnxruntime::common;
@@ -32,7 +34,9 @@ Status Environment::Initialize() {
       ONNX_NAMESPACE::OpSchemaRegistry::DomainToVersionRange::Instance().AddDomainToVersion(onnxruntime::kMSDomain, 1, 1);
       // Register contributed schemas.
       // The corresponding kernels are registered inside the appropriate execution provider.
+#ifndef DISABLE_CONTRIB_OPS
       contrib::RegisterContribSchemas();
+#endif
       RegisterOnnxOperatorSetSchema();
       RegisterOnnxMLOperatorSetSchema();
       RegisterOnnxFunctionBuilder();

onnxruntime/core/providers/cpu/cpu_execution_provider.cc

@@ -4,7 +4,11 @@
 #include "core/providers/cpu/cpu_execution_provider.h"
 #include "core/framework/op_kernel.h"
 #include "core/framework/kernel_registry.h"
+
+#ifndef DISABLE_CONTRIB_OPS
 #include "contrib_ops/contrib_kernels.h"
+#endif
+
 #include "core/framework/compute_capability.h"
 
 namespace onnxruntime {
@@ -605,7 +609,9 @@ void RegisterOnnxMLOperatorKernels(KernelRegistry& kernel_registry) {
 static void RegisterCPUKernels(KernelRegistry& kernel_registry) {
   RegisterOnnxOperatorKernels(kernel_registry);
   ::onnxruntime::ml::RegisterOnnxMLOperatorKernels(kernel_registry);
+#ifndef DISABLE_CONTRIB_OPS
   ::onnxruntime::contrib::RegisterContribKernels(kernel_registry);
+#endif
 }
 
 std::shared_ptr<KernelRegistry> GetCpuKernelRegistry() {

onnxruntime/test/ir/op_reg_test.cc

@@ -11,6 +11,7 @@ using namespace ONNX_NAMESPACE;
 
 namespace onnxruntime {
 namespace test {
+#ifndef DISABLE_CONTRIB_OPS
 TEST(OpRegistrationTest, AffineOp) {
   auto op = OpSchemaRegistry::Schema("Affine");
   EXPECT_TRUE(nullptr != op);
@@ -26,6 +27,7 @@ TEST(OpRegistrationTest, AffineOp) {
   EXPECT_EQ(attr_beta.name, "beta");
   EXPECT_EQ(attr_beta.type, AttrType::AttributeProto_AttributeType_FLOAT);
 }
+#endif
 
 TEST(FeatureVectorizerTest, TraditionalMlOpTest) {
   Model model("traditionalMl");

onnxruntime/test/providers/cpu/activation/activation_op_test.cc

@@ -183,6 +183,7 @@ TEST(ActivationOpTest, PRelu_MultiChannel) {
   test.Run();
 }
 
+#ifndef DISABLE_CONTRIB_OPS
 TEST(ActivationOpTest, ParametricSoftplus) {
   static constexpr float alpha = 2.0f;
   static constexpr float beta = 1.5f;
@@ -198,6 +199,7 @@ TEST(ActivationOpTest, ParametricSoftplus) {
                          },
                          {{"alpha", alpha}, {"beta", beta}});
 }
+#endif
 
 TEST(ActivationOpTest, Softplus) {
   TestUnaryElementwiseOp("Softplus",

onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc

@@ -760,6 +760,7 @@ TEST(MathOpTest, Mean_8) {
   test.Run();
 }
 
+#ifndef DISABLE_CONTRIB_OPS
 TEST(MathOpTest, AffineDefaultAttributes) {
   OpTester test("Affine");
   std::vector<int64_t> dims{2, 2};
@@ -777,6 +778,7 @@ TEST(MathOpTest, Affine) {
   test.AddOutput<float>("B", dims, {1.0f, 3.0f, 5.0f, 7.0f});
   test.Run();
 }
+#endif
 
 template <float (&op)(float value)>
 void TrigTest(OpTester& test, std::initializer_list<float> input) {

onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc

@@ -36,6 +36,7 @@ void Quantize(float scale, uint8_t zero_point,
   }
 }
 
+#ifndef DISABLE_CONTRIB_OPS
 TEST(ConvTest, QLinearConv2DTest) {
   OpTester test("QLinearConv", 1, onnxruntime::kMSDomain);
 
@@ -173,8 +174,9 @@ TEST(ConvTest, QLinearConv3DTest) {
 
   test.AddOutput<uint8_t>("y", Y_shape, result_quantized);
 
-  test.Run();      
+  test.Run();
 }
+#endif
 }  // namespace
 }  // namespace test
 }  // namespace onnxruntime

onnxruntime/test/providers/cpu/tensor/tensor_op_test.cc

@@ -302,6 +302,7 @@ TEST(TensorOpTest, CastToString) {
   TestCastOp(int_16_input, int_string_data, shape, TensorProto::STRING);
 }
 
+#ifndef DISABLE_CONTRIB_OPS
 TEST(TensorOpTest, CropBorderOnly) {
   const int N = 2, C = 1, H = 3, W = 4;
   std::vector<float> X = {1.0f, 2.0f, 3.0f, 4.0f,
@@ -352,6 +353,7 @@ TEST(TensorOpTest, CropBorderAndScale) {
   test.AddOutput<float>("output", {N, C, scale[0], scale[1]}, output);
   test.Run();
 }
+#endif
 
 std::pair<float, float> MeanStdev(std::vector<float>& v) {
   float sum = std::accumulate(v.begin(), v.end(), 0.0f);
@@ -579,6 +581,7 @@ TEST(TensorOpTest, MeanVarianceNormalizationCPUTest) {
   MeanVarianceNormalizationFunctionDefaultPerChannel();
 }
 
+#ifndef DISABLE_CONTRIB_OPS
 TEST(TensorOpTest, ImageScalerTest) {
   const int64_t N = 1, C = 2, H = 2, W = 2;
   std::vector<float> X = {
@@ -605,6 +608,6 @@ TEST(TensorOpTest, ImageScalerTest) {
   test.AddOutput<float>("output", {N, C, H, W}, result);
   test.Run();
 }
-
+#endif
 }  // namespace test
 }  // namespace onnxruntime

tools/ci_build/build.py

@@ -130,6 +130,7 @@ Use the individual flags to only run the specified stages.
     parser.add_argument("--use_tensorrt", action='store_true', help="Build with TensorRT")
     parser.add_argument("--tensorrt_home", help="Path to TensorRT installation dir")
     parser.add_argument("--use_full_protobuf", action='store_true', help="Use the full protobuf library")
+    parser.add_argument("--disable_contrib_ops", action='store_true', help="Disable contrib ops (reduces binary size)")
     return parser.parse_args()
 
 def resolve_executable_path(command_or_path):
@@ -316,6 +317,7 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
                  "-Donnxruntime_CROSS_COMPILING=" + ("ON" if args.arm64 or args.arm else "OFF"),
                  "-Donnxruntime_BUILD_x86=" + ("ON" if args.x86 else "OFF"),
                  "-Donnxruntime_USE_FULL_PROTOBUF=" + ("ON" if args.use_full_protobuf else "OFF"),
+                 "-Donnxruntime_DISABLE_CONTRIB_OPS=" + ("ON" if args.disable_contrib_ops else "OFF"),
                  ]
     if args.use_brainslice:
         bs_pkg_name = args.brain_slice_package_name.split('.', 1)