[Caffe] Optional NNPACK integration

The approach is pretty straightforward and very similar to the CuDNN integration:

- Add some CMake magic to detect NNPACK if it's installed on the user's
  system.
- Add some Makefile.config options
- Add some thread-local state (a threadpool for NNPACK to caffe::Caffe).
- Add an NNPACK `Engine` parameter to Convolution/Pooling/InnerProduct
  layers.
- In LayerFactory, ifdef around whether NNPACK is enabled and dispatch
  to NNPACK*Layer or *Layer depending on the Engine parameter.
- In NNPACK*Layer, ensure the NNPACK preconditions hold and call the
  appropriate nnpack_*  function, otherwise fall back to *Layer.

I'm not sure how common it is to be running CNN inference on Haswell and
beyond CPUs for Caffe users, but if it is it might be a nice boost for
them.

CMakeLists.txt

@@ -28,6 +28,7 @@ include(cmake/ConfigGen.cmake)
 # ---[ Options
 caffe_option(CPU_ONLY "Build Caffe without CUDA support" OFF) # TODO: rename to USE_CUDA
 caffe_option(USE_CUDNN "Build Caffe with cuDNN library support" ON IF NOT CPU_ONLY)
+caffe_option(USE_NNPACK "Build Caffe with NNPACK library support" OFF)
 caffe_option(BUILD_SHARED_LIBS "Build shared libraries" ON)
 caffe_option(BUILD_python "Build Python wrapper" ON)
 set(python_version "2" CACHE STRING "Specify which Python version to use")

Makefile

@@ -328,6 +328,15 @@ ifeq ($(USE_CUDNN), 1)
  COMMON_FLAGS += -DUSE_CUDNN
 endif
 
+# NNPACK acceleration configuration.
+ifeq ($(USE_NNPACK), 1)
+ LIBRARIES += nnpack
+ COMMON_FLAGS += -DUSE_NNPACK
+ INCLUDE_DIRS += $(NNPACK_INCLUDE)
+ INCLUDE_DIRS += $(NNPACK_INCLUDE)/../third-party/pthreadpool/include
+ LIBRARY_DIRS += $(NNPACK_LIB)
+endif
+
 # configure IO libraries
 ifeq ($(USE_OPENCV), 1)
  COMMON_FLAGS += -DUSE_OPENCV

Makefile.config.example

@@ -17,6 +17,11 @@
 # possibility of simultaneous read and write
 # ALLOW_LMDB_NOLOCK := 1
 
+# NNPACK acceleration switch (uncomment to build with NNPACK).
+# USE_NNPACK := 1
+# NNPACK_INCLUDE := /path/to/NNPACK/include
+# NNPACK_LIB := /path/to/nnpack/NNPACK/lib
+
 # Uncomment if you're using OpenCV 3
 # OPENCV_VERSION := 3
 

cmake/Dependencies.cmake

@@ -67,6 +67,14 @@ if(NOT HAVE_CUDA)
  add_definitions(-DCPU_ONLY)
 endif()
 
+# ---[ NNPACK
+if(USE_NNPACK)
+ find_package(NNPACK REQUIRED)
+ include_directories(SYSTEM ${NNPACK_INCLUDE_DIR})
+ include_directories(SYSTEM ${NNPACK_INCLUDE_DIR}/../third-party/pthreadpool/include)
+ list(APPEND Caffe_LINKER_LIBS ${NNPACK_LIB})
+endif()
+
 # ---[ OpenCV
 if(USE_OPENCV)
  find_package(OpenCV QUIET COMPONENTS core highgui imgproc imgcodecs)

cmake/Modules/FindNNPACK.cmake

@@ -0,0 +1,54 @@
+SET(NNPACK_INCLUDE_SEARCH_PATHS
+ /usr/include
+ /usr/local/include
+ /opt/NNPACK/include
+ $ENV{NNPACK_ROOT}
+ $ENV{NNPACK_ROOT}/include
+)
+
+SET(NNPACK_LIB_SEARCH_PATHS
+ /lib/
+ /lib64/
+ /usr/lib
+ /usr/lib64
+ /usr/local/lib
+ /usr/local/lib64
+ /opt/NNPACK/lib
+ $ENV{NNPACK_ROOT}
+ $ENV{NNPACK_ROOT}/lib
+ )
+
+FIND_PATH(NNPACK_INCLUDE_DIR NAMES nnpack.h PATHS ${NNPACK_INCLUDE_SEARCH_PATHS})
+FIND_LIBRARY(NNPACK_LIB NAMES nnpack PATHS ${NNPACK_LIB_SEARCH_PATHS})
+
+SET(NNPACK_FOUND ON)
+
+# Check include files
+IF(NOT NNPACK_INCLUDE_DIR)
+ SET(NNPACK_FOUND OFF)
+ MESSAGE(STATUS "Could not find NNPACK include. Turning NNPACK_FOUND off")
+ENDIF()
+
+# Check libraries
+IF(NOT NNPACK_LIB)
+ SET(NNPACK_FOUND OFF)
+ MESSAGE(STATUS "Could not find NNPACK lib. Turning NNPACK_FOUND off")
+ENDIF()
+
+IF (NNPACK_FOUND)
+ add_definitions(-DUSE_NNPACK)
+ IF (NOT NNPACK_FIND_QUIETLY)
+ MESSAGE(STATUS "Found NNPACK libraries: ${NNPACK_LIB}")
+ MESSAGE(STATUS "Found NNPACK include: ${NNPACK_INCLUDE_DIR}")
+ ENDIF (NOT NNPACK_FIND_QUIETLY)
+ELSE (NNPACK_FOUND)
+ IF (NNPACK_FIND_REQUIRED)
+ MESSAGE(FATAL_ERROR "Could not find NNPACK")
+ ENDIF (NNPACK_FIND_REQUIRED)
+ENDIF (NNPACK_FOUND)
+
+MARK_AS_ADVANCED(
+ NNPACK_INCLUDE_DIR
+ NNPACK_LIB
+ NNPACK
+)

cmake/Summary.cmake

@@ -136,6 +136,9 @@ function(caffe_print_configuration_summary)
  if(USE_OPENCV)
  caffe_status(" OpenCV : Yes (ver. ${OpenCV_VERSION})")
  endif()
+ if(USE_NNPACK)
+ caffe_status(" NNPACK : Yes")
+ endif()
  caffe_status(" CUDA : " HAVE_CUDA THEN "Yes (ver. ${CUDA_VERSION})" ELSE "No" )
  caffe_status("")
  if(HAVE_CUDA)

include/caffe/common.hpp

@@ -18,6 +18,10 @@
 
 #include "caffe/util/device_alternate.hpp"
 
+#ifdef USE_NNPACK
+#include "caffe/nnpack_pool.h"
+#endif
+
 // Convert macro to string
 #define STRINGIFY(m) #m
 #define AS_STRING(m) STRINGIFY(m)
@@ -138,6 +142,13 @@ class Caffe {
  }
 #endif
 
+#ifdef USE_NNPACK
+ template<typename Dtype> static bool nnpack_supported();
+ inline static pthreadpool_t nnpack_threadpool() {
+ return Get().nnpack_threadpool_.pool();
+ }
+#endif
+
  // Returns the mode: running on CPU or GPU.
  inline static Brew mode() { return Get().mode_; }
  // The setters for the variables
@@ -175,6 +186,10 @@ class Caffe {
  int solver_count_;
  bool root_solver_;
 
+#ifdef USE_NNPACK
+ NNPACKPool nnpack_threadpool_;
+#endif
+
  private:
  // The private constructor to avoid duplicate instantiation.
  Caffe();

include/caffe/nnpack_pool.h

@@ -0,0 +1,37 @@
+#pragma once
+
+#include <boost/noncopyable.hpp>
+
+#include "nnpack.h"
+
+namespace caffe {
+class NNPACKPool : public boost::noncopyable {
+ public:
+ NNPACKPool() {
+#ifdef USE_MKL
+ const size_t num_mkl_threads = mkl_get_max_threads();
+#else
+ // Can we do better here?
+ const size_t num_mkl_threads = 1;
+#endif
+ if (num_mkl_threads > 1) {
+ pool_ = pthreadpool_create(num_mkl_threads);
+ } else {
+ pool_ = NULL;
+ }
+
+ }
+ ~NNPACKPool() {
+ if (pool_) {
+ pthreadpool_destroy(pool_);
+ }
+ pool_ = NULL;
+ }
+
+ pthreadpool_t pool() { return pool_; };
+
+ private:
+ pthreadpool_t pool_;
+};
+
+}

src/caffe/common.cpp

@@ -49,6 +49,19 @@ void GlobalInit(int* pargc, char*** pargv) {
  ::google::InstallFailureSignalHandler();
 }
 
+#ifdef USE_NNPACK
+ template<>
+ bool Caffe::nnpack_supported<double>() {
+ return false;
+ }
+
+ template<>
+ bool Caffe::nnpack_supported<float>() {
+ static enum nnp_status nnpack_status = nnp_initialize();
+ return nnpack_status == nnp_status_success;
+ }
+#endif
+
 #ifdef CPU_ONLY // CPU-only Caffe.
 
 Caffe::Caffe()

src/caffe/layer_factory.cpp

@@ -16,6 +16,8 @@
 #include "caffe/layers/tanh_layer.hpp"
 #include "caffe/proto/caffe.pb.h"
 
+#include "caffe/layers/inner_product_layer.hpp"
+
 #ifdef USE_CUDNN
 #include "caffe/layers/cudnn_conv_layer.hpp"
 #include "caffe/layers/cudnn_lcn_layer.hpp"
@@ -27,6 +29,12 @@
 #include "caffe/layers/cudnn_tanh_layer.hpp"
 #endif
 
+#ifdef USE_NNPACK
+#include "caffe/layers/nnpack_convolution_layer.hpp"
+#include "caffe/layers/nnpack_pooling_layer.hpp"
+#include "caffe/layers/nnpack_inner_product_layer.hpp"
+#endif
+
 #ifdef WITH_PYTHON_LAYER
 #include "caffe/layers/python_layer.hpp"
 #endif
@@ -57,6 +65,17 @@ shared_ptr<Layer<Dtype> > GetConvolutionLayer(
  }
  if (engine == ConvolutionParameter_Engine_CAFFE) {
  return shared_ptr<Layer<Dtype> >(new ConvolutionLayer<Dtype>(param));
+#ifdef USE_NNPACK
+ } else if (engine == ConvolutionParameter_Engine_NNPACK) {
+ // If we're in CPU mode and on supported processor, we can use NNPACK.
+ // Otherwise, we can't fall-through (since we'll get an unknown
+ // layer, so just return the default ConvolutionLayer
+ if ((Caffe::mode() == Caffe::CPU) && Caffe::nnpack_supported<Dtype>()) {
+ return shared_ptr<Layer<Dtype> >(
+ new NNPackConvolutionLayer<Dtype>(param));
+ }
+ return shared_ptr<Layer<Dtype> >(new ConvolutionLayer<Dtype>(param));
+#endif
 #ifdef USE_CUDNN
  } else if (engine == ConvolutionParameter_Engine_CUDNN) {
  if (use_dilation) {
@@ -84,6 +103,16 @@ shared_ptr<Layer<Dtype> > GetPoolingLayer(const LayerParameter& param) {
  }
  if (engine == PoolingParameter_Engine_CAFFE) {
  return shared_ptr<Layer<Dtype> >(new PoolingLayer<Dtype>(param));
+#ifdef USE_NNPACK
+ } else if (engine == PoolingParameter_Engine_NNPACK) {
+ // If we're in CPU mode and on supported processor, we can use NNPACK.
+ // Otherwise, we can't fall-through (since we'll get an unknown
+ // layer, so just return the default ConvolutionLayer
+ if ((Caffe::mode() == Caffe::CPU) && Caffe::nnpack_supported<Dtype>()) {
+ return shared_ptr<Layer<Dtype> >(new NNPackPoolingLayer<Dtype>(param));
+ }
+ return shared_ptr<Layer<Dtype> >(new PoolingLayer<Dtype>(param));
+#endif
 #ifdef USE_CUDNN
  } else if (engine == PoolingParameter_Engine_CUDNN) {
  if (param.top_size() > 1) {
@@ -109,6 +138,34 @@ shared_ptr<Layer<Dtype> > GetPoolingLayer(const LayerParameter& param) {
 
 REGISTER_LAYER_CREATOR(Pooling, GetPoolingLayer);
 
+// Get pooling layer according to engine.
+template <typename Dtype>
+shared_ptr<Layer<Dtype> > GetInnerProductLayer(const LayerParameter& param) {
+ InnerProductParameter_Engine engine = param.inner_product_param().engine();
+ if (engine == InnerProductParameter_Engine_DEFAULT) {
+ engine = InnerProductParameter_Engine_CAFFE;
+ }
+ if (engine == InnerProductParameter_Engine_CAFFE) {
+ return shared_ptr<Layer<Dtype> >(new InnerProductLayer<Dtype>(param));
+#ifdef USE_NNPACK
+ } else if (engine == InnerProductParameter_Engine_NNPACK) {
+ // If we're in CPU mode and on supported processor, we can use NNPACK.
+ // Otherwise, we can't fall-through (since we'll get an unknown
+ // layer, so just return the default InnerProductLayer
+ if ((Caffe::mode() == Caffe::CPU) && Caffe::nnpack_supported<Dtype>()) {
+ return shared_ptr<Layer<Dtype> >(
+ new NNPackInnerProductLayer<Dtype>(param));
+ }
+ return shared_ptr<Layer<Dtype> >(new InnerProductLayer<Dtype>(param));
+#endif
+ } else {
+ LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
+ }
+ return shared_ptr<Layer<Dtype> >();
+}
+
+REGISTER_LAYER_CREATOR(InnerProduct, GetInnerProductLayer);
+
 // Get LRN layer according to engine
 template <typename Dtype>
 shared_ptr<Layer<Dtype> > GetLRNLayer(const LayerParameter& param) {

src/caffe/layers/inner_product_layer.cpp

@@ -145,6 +145,5 @@ STUB_GPU(InnerProductLayer);
 #endif
 
 INSTANTIATE_CLASS(InnerProductLayer);
-REGISTER_LAYER_CLASS(InnerProduct);
 
 } // namespace caffe