Flags for LMDB, LevelDB, and Caffe2 ops (#11462)

Summary:
Add flags for LMDB and LevelDB, default `OFF`. These can be enabled with

```
USE_LMDB=1 USE_LEVELDB=1 python setup.py build_deps
```

Also add a flag to build Caffe2 ops, which is default `ON`. Disable with

```
NO_CAFFE2_OPS=1 python setup.py build_deps
```

cc Yangqing soumith pjh5 mingzhe09088
Pull Request resolved: https://github.com/pytorch/pytorch/pull/11462

Reviewed By: soumith

Differential Revision: D9758156

Pulled By: orionr

fbshipit-source-id: 95fd206d72fdf44df54fc5d0aeab598bff900c63

.jenkins/caffe2/build.sh

@@ -226,7 +226,7 @@ else
     export MAX_JOBS=`expr $(nproc) - 1`
   fi
 
-  USE_OPENCV=1 BUILD_BINARY=1 python setup.py install --user
+  USE_LEVELDB=1 USE_LMDB=1 USE_OPENCV=1 BUILD_BINARY=1 python setup.py install --user
 
   # This is to save test binaries for testing
   cp -r torch/lib/tmp_install $INSTALL_PREFIX

CMakeLists.txt

@@ -60,6 +60,7 @@ option(BUILD_BINARY "Build C++ binaries" OFF)
 option(BUILD_DOCS "Build Caffe2 documentation" OFF)
 option(BUILD_CUSTOM_PROTOBUF "Build and use Caffe2's own protobuf under third_party" ON)
 option(BUILD_PYTHON "Build Python binaries" ON)
+option(BUILD_CAFFE2_OPS "Build Caffe2 operators" ON)
 option(BUILD_SHARED_LIBS "Build libcaffe2.so" ON)
 cmake_dependent_option(
     CAFFE2_LINK_LOCAL_PROTOBUF "If set, build protobuf inside libcaffe2.so." ON

caffe2/CMakeLists.txt

@@ -86,8 +86,10 @@ add_subdirectory(mobile)
 add_subdirectory(mpi)
 add_subdirectory(observers)
 add_subdirectory(onnx)
-add_subdirectory(operators)
-add_subdirectory(operators/rnn)
+if (BUILD_CAFFE2_OPS)
+  add_subdirectory(operators)
+  add_subdirectory(operators/rnn)
+endif()
 add_subdirectory(opt)
 add_subdirectory(perfkernels)
 add_subdirectory(python)

caffe2/contrib/aten/CMakeLists.txt

@@ -1,4 +1,4 @@
-if(NOT BUILD_ATEN_MOBILE)
+if(NOT BUILD_ATEN_MOBILE AND BUILD_CAFFE2_OPS)
   # Add source generated by Codegen.cmake and pass to parent
   list(APPEND Caffe2_CPU_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/aten_op.cc)
   list(APPEND Caffe2_GPU_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/aten_op_cuda.cc)

caffe2/share/contrib/CMakeLists.txt

@@ -1,4 +1,4 @@
-if (USE_NNPACK)
+if (USE_NNPACK AND BUILD_CAFFE2_OPS)
   add_subdirectory(nnpack)
 endif()
 if (USE_ZSTD)

cmake/Dependencies.cmake

@@ -770,7 +770,7 @@ if (USE_NNAPI AND NOT ANDROID)
   caffe2_update_option(USE_NNAPI OFF)
 endif()
 
-if (NOT BUILD_ATEN_MOBILE)
+if (NOT BUILD_ATEN_MOBILE AND BUILD_CAFFE2_OPS)
   if (CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO)
     list(APPEND Caffe2_DEPENDENCY_LIBS aten_op_header_gen)
     if (USE_CUDA)

cmake/Summary.cmake

@@ -38,6 +38,7 @@ function (caffe2_print_configuration_summary)
     message(STATUS "    Python includes     : ${PYTHON_INCLUDE_DIRS}")
     message(STATUS "    Python site-packages: ${PYTHON_SITE_PACKAGES}")
   endif()
+  message(STATUS "  BUILD_CAFFE2_OPS      : ${BUILD_CAFFE2_OPS}")
   message(STATUS "  BUILD_SHARED_LIBS     : ${BUILD_SHARED_LIBS}")
   message(STATUS "  BUILD_TEST            : ${BUILD_TEST}")
 

scripts/onnx/install-develop.sh

@@ -15,4 +15,4 @@ pip install -e "$tp2_dir/onnx"
 # Install caffe2 and pytorch
 pip install -r "$top_dir/caffe2/requirements.txt"
 pip install -r "$top_dir/requirements.txt"
-USE_OPENCV=1 BUILD_BINARY=1 python setup.py build_deps develop
+USE_LEVELDB=1 USE_LMDB=1 USE_OPENCV=1 BUILD_BINARY=1 python setup.py build_deps develop

scripts/onnx/install.sh

@@ -35,4 +35,4 @@ _pip_install -b "$BUILD_DIR/onnx" "file://$tp2_dir/onnx#egg=onnx"
 # Install caffe2 and pytorch
 pip install -r "$top_dir/caffe2/requirements.txt"
 pip install -r "$top_dir/requirements.txt"
-USE_OPENCV=1 BUILD_BINARY=1 python setup.py install
+USE_LEVELDB=1 USE_LMDB=1 USE_OPENCV=1 BUILD_BINARY=1 python setup.py install

setup.py

@@ -46,12 +46,21 @@
 #     disables use of system-wide nccl (we will use our submoduled
 #     copy in third_party/nccl)
 #
+#   NO_CAFFE2_OPS
+#     disable Caffe2 operators build
+#
 #   USE_GLOO_IBVERBS
 #     toggle features related to distributed support
 #
 #   USE_OPENCV
 #     enables use of OpenCV for additional operators
 #
+#   USE_LEVELDB
+#     enables use of LevelDB for storage
+#
+#   USE_LMBD
+#     enables use of LMDB for storage
+#
 #   BUILD_BINARY
 #     enables the additional binaries/ build
 #
@@ -144,11 +153,13 @@ use_env_vars = ['CUDA', 'CUDNN', 'MIOPEN', 'MKLDNN', 'NNPACK', 'DISTRIBUTED',
 list(map(hotpatch_var, use_env_vars))
 
 # Also hotpatch a few with BUILD_* equivalent
-build_env_vars = ['BINARY', 'TEST']
+build_env_vars = ['BINARY', 'TEST', 'CAFFE2_OPS']
 [hotpatch_var(v, 'BUILD_') for v in build_env_vars]
 
 from tools.setup_helpers.cuda import USE_CUDA, CUDA_HOME, CUDA_VERSION
-from tools.setup_helpers.build import BUILD_BINARY, BUILD_TEST, USE_OPENCV
+from tools.setup_helpers.build import (BUILD_BINARY, BUILD_TEST,
+                                       BUILD_CAFFE2_OPS, USE_LEVELDB,
+                                       USE_LMDB, USE_OPENCV)
 from tools.setup_helpers.rocm import USE_ROCM, ROCM_HOME, ROCM_VERSION
 from tools.setup_helpers.cudnn import (USE_CUDNN, CUDNN_LIBRARY,
                                        CUDNN_LIB_DIR, CUDNN_INCLUDE_DIR)
@@ -376,7 +387,10 @@ def build_libs(libs):
     my_env["BUILD_PYTHON"] = "ON"
     my_env["BUILD_BINARY"] = "ON" if BUILD_BINARY else "OFF"
     my_env["BUILD_TEST"] = "ON" if BUILD_TEST else "OFF"
+    my_env["BUILD_CAFFE2_OPS"] = "ON" if BUILD_CAFFE2_OPS else "OFF"
     my_env["INSTALL_TEST"] = "ON" if BUILD_TEST else "OFF"
+    my_env["USE_LEVELDB"] = "ON" if USE_LEVELDB else "OFF"
+    my_env["USE_LMDB"] = "ON" if USE_LMDB else "OFF"
     my_env["USE_OPENCV"] = "ON" if USE_OPENCV else "OFF"
 
     try:

tools/build_pytorch_libs.bat

@@ -180,10 +180,13 @@ goto:eof
                   -DBUILD_BINARY=%BUILD_BINARY% ^
                   -DBUILD_TEST=OFF ^
                   -DINSTALL_TEST=%INSTALL_TEST% ^
+                  -DBUILD_CAFFE2_OPS=%BUILD_CAFFE2_OPS% ^
                   -DONNX_NAMESPACE=%ONNX_NAMESPACE% ^
                   -DUSE_CUDA=%USE_CUDA% ^
                   -DUSE_CUDNN=OFF ^
                   -DUSE_NNPACK=%USE_NNPACK% ^
+                  -DUSE_LEVELDB=%USE_LEVELDB% ^
+                  -DUSE_LMDB=%USE_LMDB% ^
                   -DUSE_OPENCV=%USE_OPENCV% ^
                   -DUSE_GLOG=OFF ^
                   -DUSE_GFLAGS=OFF ^

tools/build_pytorch_libs.sh

@@ -263,11 +263,14 @@ function build_caffe2() {
       -DBUILD_BINARY=$BUILD_BINARY \
       -DBUILD_TEST=$BUILD_TEST \
       -DINSTALL_TEST=$INSTALL_TEST \
+      -DBUILD_CAFFE2_OPS=$BUILD_CAFFE2_OPS \
       -DONNX_NAMESPACE=$ONNX_NAMESPACE \
       -DUSE_CUDA=$USE_CUDA \
       -DCAFFE2_STATIC_LINK_CUDA=$CAFFE2_STATIC_LINK_CUDA \
       -DUSE_ROCM=$USE_ROCM \
       -DUSE_NNPACK=$USE_NNPACK \
+      -DUSE_LEVELDB=$USE_LEVELDB \
+      -DUSE_LMDB=$USE_LMDB \
       -DUSE_OPENCV=$USE_OPENCV \
       -DUSE_GLOG=OFF \
       -DUSE_GFLAGS=OFF \

tools/setup_helpers/build.py

@@ -2,4 +2,7 @@ from .env import check_env_flag, check_negative_env_flag
 
 BUILD_BINARY = check_env_flag('BUILD_BINARY')
 BUILD_TEST = not check_negative_env_flag('BUILD_TEST')
+BUILD_CAFFE2_OPS = not check_negative_env_flag('BUILD_CAFFE2_OPS')
+USE_LEVELDB = check_env_flag('USE_LEVELDB')
+USE_LMDB = check_env_flag('USE_LMDB')
 USE_OPENCV = check_env_flag('USE_OPENCV')