[WIP] CMake NNPack support

.gitignore

@@ -127,6 +127,11 @@ CMakeFiles
 cmake_install.cmake
 lib
 
+# Kate / Kdevelop files
+*.kate-swp
+*.kdev4
+
+
 # Visual Studio Code
 .vscode
 

.gitmodules

@@ -26,3 +26,6 @@
 [submodule "3rdparty/cub"]
 	path = 3rdparty/cub
 	url = https://github.com/dmlc/cub
+[submodule "3rdparty/nnpack/NNPACK"]
+	path = 3rdparty/NNPACK
+	url = https://github.com/Maratyszcza/NNPACK

CMakeLists.txt

@@ -34,6 +34,7 @@ mxnet_option(USE_VTUNE            "Enable use of Intel Amplifier XE (VTune)" OFF
 mxnet_option(ENABLE_CUDA_RTC      "Build with CUDA runtime compilation support" ON)
 mxnet_option(INSTALL_EXAMPLES     "Install the example source files." OFF)
 mxnet_option(USE_SIGNAL_HANDLER   "Print stack traces on segfaults." OFF)
+mxnet_option(USE_NNPACK           "Build with NNPack support." OFF)
 
 if(USE_CUDA AND NOT USE_OLDCMAKECUDA)
   message(STATUS "CMake version '${CMAKE_VERSION}' using generator '${CMAKE_GENERATOR}'")
@@ -551,6 +552,52 @@ if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/nnvm/CMakeLists.txt")
   list(APPEND mxnet_LINKER_LIBS ${nnvm_LINKER_LIBS})
 endif()
 
+# ---[ NNPack
+if(USE_NNPACK)
+  if (USE_MKLDNN)
+    message(FATAL_ERROR "Either MKLDNN or NNPack can be enabled but not both.")
+  endif()
+  # Add in NNPack and its dependencies
+  set(NNPACK_SOURCE_DIR "${CMAKE_SOURCE_DIR}/3rdparty/NNPACK")
+  if (EXISTS "${NNPACK_SOURCE_DIR}")
+    if (GOOGLETEST_SOURCE_DIR AND EXISTS "${GOOGLETEST_SOURCE_DIR}")
+      set(GOOGLETEST_SOURCE_DIR "${GTEST_ROOT}" CACHE STRING "Google Test source directory")
+    endif()
+
+    # Disable NNPack internal testing
+    set(NNPACK_BUILD_TESTS OFF CACHE BOOL "")
+    set(NNPACK_BUILD_BENCHMARKS OFF CACHE BOOL "")
+
+    # Compile statically
+    set(NNPACK_LIBRARY_TYPE "static" CACHE STRING "")
+    set(PTHREADPOOL_LIBRARY_TYPE "static" CACHE STRING "")
+    set(CPUINFO_LIBRARY_TYPE "static" CACHE STRING "")
+
+    # put NNPack dependencies in appropriate folders, see NNPack
+    # for other options
+    set(CONFU_DEPENDENCIES_SOURCE_DIR "${NNPACK_SOURCE_DIR}/deps"
+        CACHE PATH "Confu-style dependencies source directory")
+    set(CONFU_DEPENDENCIES_BINARY_DIR "${CMAKE_BINARY_DIR}/3rdparty/NNPACK/deps"
+        CACHE PATH "Confu-style dependencies source directory")
+
+    add_subdirectory("${NNPACK_SOURCE_DIR}")
+
+    # compile with -fPIC
+    set_property(TARGET nnpack PROPERTY POSITION_INDEPENDENT_CODE ON)
+    set_property(TARGET pthreadpool PROPERTY POSITION_INDEPENDENT_CODE ON)
+    set_property(TARGET cpuinfo PROPERTY POSITION_INDEPENDENT_CODE ON)
+
+    include_directories(${NNPACK_SOURCE_DIR}/include)
+    #PTHREADPOOL_SOURCE_DIR is set in cache by NNPack
+    include_directories(${PTHREADPOOL_SOURCE_DIR}/include)
+    add_definitions(-DMXNET_USE_NNPACK=1)
+    set(NNPack_LINKER_LIBS nnpack)
+    list(APPEND mxnet_LINKER_LIBS ${NNPack_LINKER_LIBS})
+  else()
+    message("NNPack submodule not found.")
+  endif()
+endif()
+
 if(NOT MSVC)
   # Only add c++11 flags and definitions after cuda compiling
   add_definitions(-DDMLC_USE_CXX11)

src/operator/convolution_v1.cc

@@ -25,9 +25,6 @@
 */
 
 #include "./convolution_v1-inl.h"
-#if MXNET_USE_NNPACK == 1
-#include "./nnpack/nnpack_convolution-inl.h"
-#endif  // MXNET_USE_NNPACK
 
 namespace mxnet {
 namespace op {

src/operator/nn/convolution.cc

@@ -28,9 +28,7 @@
 #include "../elemwise_op_common.h"
 #include "./mkldnn/mkldnn_ops-inl.h"
 #include "./mkldnn/mkldnn_base-inl.h"
-#if MXNET_USE_NNPACK == 1
-#include "./nnpack/nnpack_convolution-inl.h"
-#endif  // MXNET_USE_NNPACK
+#include "./nnpack/nnpack_ops-inl.h"
 
 namespace mxnet {
 namespace op {

src/operator/nn/fully_connected.cc

@@ -25,9 +25,7 @@
 #include "./fully_connected-inl.h"
 #include "./mkldnn/mkldnn_ops-inl.h"
 #include "./mkldnn/mkldnn_base-inl.h"
-#if MXNET_USE_NNPACK == 1
-#include "./nnpack/nnpack_fully_connected-inl.h"
-#endif  // MXNET_USE_NNPACK
+#include "./nnpack/nnpack_ops-inl.h"
 
 namespace mxnet {
 namespace op {

src/operator/nn/nnpack/nnpack_ops-inl.h

@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file nnpack_ops-inl.h
+ * \brief
+ * \author David Braude
+*/
+
+#ifndef MXNET_OPERATOR_NN_NNPACK_NNPACK_OPS_INL_H_
+#define MXNET_OPERATOR_NN_NNPACK_NNPACK_OPS_INL_H_
+
+#if MXNET_USE_NNPACK == 1
+
+#include <mxnet/io.h>
+#include <mxnet/base.h>
+#include <mxnet/ndarray.h>
+#include <mxnet/operator.h>
+#include <mxnet/operator_util.h>
+#include <dmlc/logging.h>
+#include <dmlc/optional.h>
+#include <vector>
+#include <nnpack.h>
+
+// TODO:
+// Convolutional layer
+//     Inference-optimized forward propagation (nnp_convolution_inference)
+//     Training-optimized forward propagation (nnp_convolution_output)
+//     Training-optimized backward input gradient update (nnp_convolution_input_gradient)
+//     Training-optimized backward kernel gradient update (nnp_convolution_kernel_gradient)
+// Fully-connected layer
+//     Inference-optimized forward propagation (nnp_fully_connected_inference and nnp_fully_connected_inference_f16f32 version for FP16 weights)
+//     Training-optimized forward propagation (nnp_fully_connected_output)
+// Max pooling layer
+//     Forward propagation, both for training and inference, (nnp_max_pooling_output)
+// ReLU layer (with parametrized negative slope)
+//     Forward propagation, both for training and inference, optionally in-place, (nnp_relu_output)
+//     Backward input gradient update (nnp_relu_input_gradient)
+
+namespace mxnet {
+namespace op {
+
+/* For softmax */
+void NNPACKSoftmaxForward(const nnvm::NodeAttrs& attrs, const OpContext &ctx,
+                          const NDArray &in_data, const OpReqType &req,
+                          const NDArray &out_data);
+
+}  // namespace op
+}  // namespace mxnet
+#endif  // MXNET_USE_MKLDNN == 1
+
+#endif  // MXNET_OPERATOR_NN_MKLDNN_MKLDNN_OPS_INL_H_

src/operator/nnpack/nnpack_util.cc → src/operator/nn/nnpack/nnpack_softmax.cc

@@ -18,20 +18,31 @@
  */
 
 /*!
- * Copyright (c) 2016 by Contributors
- * \file nnpack_util.cc
+ * \file nnpack_softmax.cc
  * \brief
- * \author Wei Wu
+ * \author David Braude
 */
 
-#if MXNET_USE_NNPACK == 1
-#include "nnpack_util.h"
+#include "../softmax-inl.h"
+#include "./nnpack_ops-inl.h"
+
 
+#if MXNET_USE_NNPACK == 1
 namespace mxnet {
 namespace op {
 
-NNPACKInitialize nnpackinitialize;
+void NNPACKSoftmaxForward(const nnvm::NodeAttrs& attrs, const OpContext &ctx,
+                          const NDArray &in_data, const OpReqType &req,
+                          const NDArray &out_data) {
+  const SoftmaxParam& param = nnvm::get<SoftmaxParam>(attrs.parsed);
+//     enum nnp_status nnp_softmax_output(
+//     size_t batch_size,
+//     size_t channels,
+//     const float input[],
+//     float output[],
+//     pthreadpool_t threadpool);
+}
 
-}  // namespace op
-}  // namespace mxnet
-#endif  // MXNET_USE_NNPACK
+}   // namespace op
+}   // namespace mxnet
+#endif

src/operator/nn/pooling.cc

@@ -25,12 +25,8 @@
 */
 #include "../elemwise_op_common.h"
 #include "./pooling-inl.h"
-#if MXNET_USE_NNPACK == 1
-#include "./nnpack/nnpack_pooling-inl.h"
-#endif  // MXNET_USE_NNPACK
-#if MXNET_USE_MKLDNN == 1
+// #include "./nnpack/nnpack_pooling-inl.h"
 #include "./mkldnn/mkldnn_pooling-inl.h"
-#endif  // MXNET_USE_MKLDNN
 
 namespace mxnet {
 namespace op {