From 39eb871ddfe4118293f1211b8e3427505b1cf333 Mon Sep 17 00:00:00 2001
From: Liu Yiqun <liuyiqun01@baidu.com>
Date: Mon, 21 May 2018 02:38:37 +0000
Subject: [PATCH] Add an interface to set the number of threads for math
 function, and set the default value to 1 for inference.

---
 cmake/external/openblas.cmake               |  2 ++
 paddle/fluid/inference/io.cc                |  4 ++++
 paddle/fluid/operators/math/blas.h          | 15 +++++++++++++++
 paddle/fluid/operators/math/math_function.h |  2 ++
 paddle/math/MathFunctions.h                 |  9 ++++-----
 5 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/cmake/external/openblas.cmake b/cmake/external/openblas.cmake
index 8af2765f587..4a49a92f2b1 100644
--- a/cmake/external/openblas.cmake
+++ b/cmake/external/openblas.cmake
@@ -29,6 +29,8 @@ IF(NOT ${CBLAS_FOUND})
         "${CBLAS_INSTALL_DIR}/lib/${CMAKE_STATIC_LIBRARY_PREFIX}openblas${CMAKE_STATIC_LIBRARY_SUFFIX}"
         CACHE FILEPATH "openblas library." FORCE)
 
+    ADD_DEFINITIONS(-DPADDLE_USE_OPENBLAS)
+
     SET(OPENBLAS_CC "${CMAKE_C_COMPILER} -Wno-unused-but-set-variable -Wno-unused-variable")
     SET(OPENBLAS_COMMIT "v0.2.20")
 
diff --git a/paddle/fluid/inference/io.cc b/paddle/fluid/inference/io.cc
index 65db7c7b500..6b03ac7119b 100644
--- a/paddle/fluid/inference/io.cc
+++ b/paddle/fluid/inference/io.cc
@@ -20,16 +20,20 @@ limitations under the License. */
 #include "paddle/fluid/framework/block_desc.h"
 #include "paddle/fluid/framework/feed_fetch_type.h"
 #include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/operators/math/blas.h"
 #include "paddle/fluid/pybind/pybind.h"
 
 DEFINE_string(devices, "", "The devices to be used which is joined by comma.");
 DEFINE_bool(init_p2p, false, "Whether to init p2p.");
+DEFINE_int32(math_num_threads, 1,
+             "Number of threads used to run math functions.");
 
 namespace paddle {
 namespace inference {
 
 void Init(const std::vector<std::string> argv) {
   framework::InitGflags(argv);
+  operators::math::SetNumThreads(FLAGS_math_num_threads);
   // init devices
   std::vector<int> devices;
   std::string token;
diff --git a/paddle/fluid/operators/math/blas.h b/paddle/fluid/operators/math/blas.h
index dabde43850d..b393139acd0 100644
--- a/paddle/fluid/operators/math/blas.h
+++ b/paddle/fluid/operators/math/blas.h
@@ -20,13 +20,16 @@
 #ifdef PADDLE_WITH_MKLML
 #include <mkl_cblas.h>
 #include <mkl_lapacke.h>
+#include <mkl_service.h>
 #include <mkl_vml_functions.h>
 #endif
 
 #ifdef PADDLE_USE_OPENBLAS
 #include <cblas.h>
+#ifdef LAPACK_FOUND
 #include <lapacke.h>
 #endif
+#endif
 
 #ifndef LAPACK_FOUND
 extern "C" {
@@ -46,6 +49,18 @@ namespace paddle {
 namespace operators {
 namespace math {
 
+static void SetNumThreads(int num_threads) {
+#ifdef PADDLE_USE_OPENBLAS
+  int real_num_threads = num_threads > 1 ? num_threads : 1;
+  openblas_set_num_threads(real_num_threads);
+#elif defined(PADDLE_WITH_MKLML)
+  int real_num_threads = num_threads > 1 ? num_threads : 1;
+  mkl_set_num_threads(real_num_threads);
+#else
+  PADDLE_ENFORCE(false, "To be implemented.");
+#endif
+}
+
 /**
  * Matrix Descriptor of a memory buffer.
  *
diff --git a/paddle/fluid/operators/math/math_function.h b/paddle/fluid/operators/math/math_function.h
index d4b0e17ed44..8b296b6a07c 100644
--- a/paddle/fluid/operators/math/math_function.h
+++ b/paddle/fluid/operators/math/math_function.h
@@ -21,8 +21,10 @@ limitations under the License. */
 
 #ifdef PADDLE_USE_OPENBLAS
 #include <cblas.h>
+#ifdef LAPACK_FOUND
 #include <lapacke.h>
 #endif
+#endif
 
 #ifndef LAPACK_FOUND
 extern "C" {
diff --git a/paddle/math/MathFunctions.h b/paddle/math/MathFunctions.h
index f3d8b1a39e8..854e4baa398 100644
--- a/paddle/math/MathFunctions.h
+++ b/paddle/math/MathFunctions.h
@@ -12,8 +12,7 @@ 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. */
 
-#ifndef MATHFUNCTIONS_H_
-#define MATHFUNCTIONS_H_
+#pragma once
 
 #ifdef PADDLE_WITH_MKLML
 #include <mkl_cblas.h>
@@ -21,7 +20,7 @@ limitations under the License. */
 #include <mkl_vml_functions.h>
 #endif
 
-#if defined(PADDLE_USE_VECLIB)
+#ifdef PADDLE_USE_VECLIB
 extern "C" {
 #include <cblas.h>
 #include <clapack.h>
@@ -30,8 +29,10 @@ extern "C" {
 
 #ifdef PADDLE_USE_OPENBLAS
 #include <cblas.h>
+#ifdef LAPACK_FOUND
 #include <lapacke.h>
 #endif
+#endif
 
 #ifndef LAPACK_FOUND
 extern "C" {
@@ -126,5 +127,3 @@ template <class T>
 void vTanh(const int n, const T* a, T* r);
 
 }  // namespace paddle
-
-#endif  // MATHFUNCTIONS_H_
-- 
GitLab