[Caffe] Minor refactoring of NNPACK integration

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)
@@ -80,29 +88,30 @@ if(USE_OPENCV)
 endif()
 
 # ---[ BLAS
-if(NOT APPLE)
+# if(NOT APPLE)
   set(BLAS "Atlas" CACHE STRING "Selected BLAS library")
   set_property(CACHE BLAS PROPERTY STRINGS "Atlas;Open;MKL")
 
-  if(BLAS STREQUAL "Atlas" OR BLAS STREQUAL "atlas")
-    find_package(Atlas REQUIRED)
-    include_directories(SYSTEM ${Atlas_INCLUDE_DIR})
-    list(APPEND Caffe_LINKER_LIBS ${Atlas_LIBRARIES})
-  elseif(BLAS STREQUAL "Open" OR BLAS STREQUAL "open")
+  # if(BLAS STREQUAL "Atlas" OR BLAS STREQUAL "atlas")
+  #   find_package(Atlas REQUIRED)
+  #   include_directories(SYSTEM ${Atlas_INCLUDE_DIR})
+  #   list(APPEND Caffe_LINKER_LIBS ${Atlas_LIBRARIES})
+  # elseif(BLAS STREQUAL "Open" OR BLAS STREQUAL "open"
+      # )
     find_package(OpenBLAS REQUIRED)
     include_directories(SYSTEM ${OpenBLAS_INCLUDE_DIR})
     list(APPEND Caffe_LINKER_LIBS ${OpenBLAS_LIB})
-  elseif(BLAS STREQUAL "MKL" OR BLAS STREQUAL "mkl")
-    find_package(MKL REQUIRED)
-    include_directories(SYSTEM ${MKL_INCLUDE_DIR})
-    list(APPEND Caffe_LINKER_LIBS ${MKL_LIBRARIES})
-    add_definitions(-DUSE_MKL)
-  endif()
-elseif(APPLE)
-  find_package(vecLib REQUIRED)
-  include_directories(SYSTEM ${vecLib_INCLUDE_DIR})
-  list(APPEND Caffe_LINKER_LIBS ${vecLib_LINKER_LIBS})
-endif()
+  # elseif(BLAS STREQUAL "MKL" OR BLAS STREQUAL "mkl")
+  #   find_package(MKL REQUIRED)
+  #   include_directories(SYSTEM ${MKL_INCLUDE_DIR})
+  #   list(APPEND Caffe_LINKER_LIBS ${MKL_LIBRARIES})
+  #   add_definitions(-DUSE_MKL)
+  # endif()
+# elseif(APPLE)
+#   find_package(vecLib REQUIRED)
+#   include_directories(SYSTEM ${vecLib_INCLUDE_DIR})
+#   list(APPEND Caffe_LINKER_LIBS ${vecLib_LINKER_LIBS})
+# endif()
 
 # ---[ Python
 if(BUILD_python)

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

@@ -120,7 +120,8 @@ function(caffe_print_configuration_summary)
   caffe_status("  ALLOW_LMDB_NOLOCK :   ${ALLOW_LMDB_NOLOCK}")
   caffe_status("")
   caffe_status("Dependencies:")
-  caffe_status("  BLAS              : " APPLE THEN "Yes (vecLib)" ELSE "Yes (${BLAS})")
+  # caffe_status("  BLAS              : " APPLE THEN "Yes (vecLib)" ELSE "Yes (${BLAS})")
+  caffe_status("  BLAS              :   ${BLAS}")
   caffe_status("  Boost             :   Yes (ver. ${Boost_MAJOR_VERSION}.${Boost_MINOR_VERSION})")
   caffe_status("  glog              :   Yes")
   caffe_status("  gflags            :   Yes")
@@ -135,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/layers/nnpack_convolution_layer.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "caffe/layers/conv_layer.hpp"
+
+namespace caffe {
+#ifdef USE_NNPACK
+template <typename Dtype>
+class NNPackConvolutionLayer : public ConvolutionLayer<Dtype> {
+ public:
+  explicit NNPackConvolutionLayer(const LayerParameter& param)
+      : ConvolutionLayer<Dtype>(param) {}
+  virtual inline const char* type() const {
+    return "NNPackConvolution";
+  }
+  virtual void Forward_cpu(
+      const vector<Blob<Dtype>*>& bottom,
+      const vector<Blob<Dtype>*>& top);
+  virtual void Backward_cpu(
+      const vector<Blob<Dtype>*>& top,
+      const vector<bool>& propagate_down,
+      const vector<Blob<Dtype>*>& bottom);
+};
+#endif
+}

include/caffe/layers/nnpack_inner_product_layer.hpp

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "caffe/layers/inner_product_layer.hpp"
+
+namespace caffe {
+#ifdef USE_NNPACK
+template <typename Dtype>
+class NNPackInnerProductLayer : public InnerProductLayer<Dtype> {
+ public:
+  explicit NNPackInnerProductLayer(const LayerParameter& param)
+      : InnerProductLayer<Dtype>(param) {}
+  virtual inline const char* type() const {
+    return "NNPackInnerProduct";
+  }
+  virtual void Forward_cpu(
+      const vector<Blob<Dtype>*>& bottom,
+      const vector<Blob<Dtype>*>& top);
+};
+#endif
+}

include/caffe/layers/nnpack_pooling_layer.hpp

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "caffe/layers/pooling_layer.hpp"
+
+namespace caffe {
+#ifdef USE_NNPACK
+template <typename Dtype>
+class NNPackPoolingLayer : public PoolingLayer<Dtype> {
+ public:
+  explicit NNPackPoolingLayer(const LayerParameter& param)
+      : PoolingLayer<Dtype>(param) {}
+  virtual inline const char* type() const {
+    return "NNPackPooling";
+  }
+  virtual void Forward_cpu(
+      const vector<Blob<Dtype>*>& bottom,
+      const vector<Blob<Dtype>*>& top);
+};
+#endif
+}

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()