You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by aa...@apache.org on 2019/01/21 15:45:48 UTC
[incubator-mxnet] branch master updated: Runtime feature detection
(#13549)
This is an automated email from the ASF dual-hosted git repository.
aaronmarkham pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new 692a24a Runtime feature detection (#13549)
692a24a is described below
commit 692a24a0411e6d937bb46465f7eefd1e2df8be69
Author: Pedro Larroy <pe...@gmail.com>
AuthorDate: Mon Jan 21 16:45:07 2019 +0100
Runtime feature detection (#13549)
* Prototype for runtime feature detection
* Includes from diamond to quotes
* Add CPU feature and BLAS flavour flags
* Add BLAS flavour and CPU SSE and AVX flags
* MXNET_USE_LAPACK
* Fix C++ linting errors
* Expose runtime feature detection in the public C API and in the Python API
* Refactor Storage -> FeatureSet
* Refine documentation
* Add failure case
* Fix pylint
* Address CR comments
---
CMakeLists.txt | 2 +
Makefile | 10 ++
cmake/ChooseBlas.cmake | 4 +
include/mxnet/base.h | 50 ++-------
include/mxnet/c_api.h | 10 ++
include/mxnet/io.h | 4 +-
include/mxnet/mxfeatures.h | 188 +++++++++++++++++++++++++++++++++
python/mxnet/gluon/utils.py | 2 +-
python/mxnet/mxfeatures.py | 103 ++++++++++++++++++
src/c_api/c_api.cc | 36 ++++---
src/c_api/c_api_symbolic.cc | 12 +--
src/mxfeatures.cc | 112 ++++++++++++++++++++
tests/python/unittest/test_features.py | 40 +++++++
13 files changed, 511 insertions(+), 62 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index f61c241..2f974e8 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -423,12 +423,14 @@ if(USE_OPENMP)
endif()
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
+ add_definitions(-DMXNET_USE_OPENMP=1)
else()
if(OPENMP_FOUND)
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} ${OpenMP_EXE_LINKER_FLAGS}")
set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${OpenMP_EXE_LINKER_FLAGS}")
+ add_definitions(-DMXNET_USE_OPENMP=1)
endif()
endif()
elseif(UNIX AND NOT ANDROID)
diff --git a/Makefile b/Makefile
index 42010e4..f4769c3 100644
--- a/Makefile
+++ b/Makefile
@@ -212,6 +212,16 @@ ifeq ($(USE_CUDNN), 1)
LDFLAGS += -lcudnn
endif
+ifeq ($(use_blas), open)
+ CFLAGS += -DMXNET_USE_BLAS_OPEN=1
+else ifeq ($(use_blas), atlas)
+ CFLAGS += -DMXNET_USE_BLAS_ATLAS=1
+else ifeq ($(use_blas), mkl)
+ CFLAGS += -DMXNET_USE_BLAS_MKL=1
+else ifeq ($(use_blas), apple)
+ CFLAGS += -DMXNET_USE_BLAS_APPLE=1
+endif
+
# whether to use F16C instruction set extension for fast fp16 compute on CPU
# if cross compiling you may want to explicitly turn it off if target system does not support it
ifndef USE_F16C
diff --git a/cmake/ChooseBlas.cmake b/cmake/ChooseBlas.cmake
index 13d7083..5f4af2d 100644
--- a/cmake/ChooseBlas.cmake
+++ b/cmake/ChooseBlas.cmake
@@ -37,22 +37,26 @@ if(BLAS STREQUAL "Atlas" OR BLAS STREQUAL "atlas")
list(APPEND mshadow_LINKER_LIBS ${Atlas_LIBRARIES})
add_definitions(-DMSHADOW_USE_CBLAS=1)
add_definitions(-DMSHADOW_USE_MKL=0)
+ add_definitions(-DMXNET_USE_BLAS_ATLAS=1)
elseif(BLAS STREQUAL "Open" OR BLAS STREQUAL "open")
find_package(OpenBLAS REQUIRED)
include_directories(SYSTEM ${OpenBLAS_INCLUDE_DIR})
list(APPEND mshadow_LINKER_LIBS ${OpenBLAS_LIB})
add_definitions(-DMSHADOW_USE_CBLAS=1)
add_definitions(-DMSHADOW_USE_MKL=0)
+ add_definitions(-DMXNET_USE_BLAS_OPEN=1)
elseif(BLAS STREQUAL "MKL" OR BLAS STREQUAL "mkl")
find_package(MKL REQUIRED)
include_directories(SYSTEM ${MKL_INCLUDE_DIR})
list(APPEND mshadow_LINKER_LIBS ${MKL_LIBRARIES})
add_definitions(-DMSHADOW_USE_CBLAS=0)
add_definitions(-DMSHADOW_USE_MKL=1)
+ add_definitions(-DMXNET_USE_BLAS_MKL=1)
elseif(BLAS STREQUAL "apple")
find_package(Accelerate REQUIRED)
include_directories(SYSTEM ${Accelerate_INCLUDE_DIR})
list(APPEND mshadow_LINKER_LIBS ${Accelerate_LIBRARIES})
add_definitions(-DMSHADOW_USE_MKL=0)
add_definitions(-DMSHADOW_USE_CBLAS=1)
+ add_definitions(-DMXNET_USE_BLAS_APPLE=1)
endif()
diff --git a/include/mxnet/base.h b/include/mxnet/base.h
index 92d9c26..f88b227 100644
--- a/include/mxnet/base.h
+++ b/include/mxnet/base.h
@@ -25,47 +25,18 @@
#ifndef MXNET_BASE_H_
#define MXNET_BASE_H_
-#include <dmlc/base.h>
-#include <dmlc/io.h>
-#include <dmlc/type_traits.h>
-#include <dmlc/parameter.h>
-#include <mshadow/tensor.h>
-// nnvm headers for symbolic construction.
-#include <nnvm/op.h>
-#include <nnvm/tuple.h>
-#include <nnvm/symbolic.h>
+#include "dmlc/base.h"
#include <string>
+#include "dmlc/io.h"
+#include "dmlc/type_traits.h"
+#include "dmlc/parameter.h"
+#include "mshadow/tensor.h"
+// nnvm headers for symbolic construction.
+#include "nnvm/op.h"
+#include "nnvm/tuple.h"
+#include "nnvm/symbolic.h"
+#include "mxfeatures.h"
-/*!
- *\brief whether to use opencv support
- */
-#ifndef MXNET_USE_OPENCV
-#define MXNET_USE_OPENCV 1
-#endif
-
-/*!
- *\brief whether to use cuda support
- */
-#ifndef MXNET_USE_CUDA
-#define MXNET_USE_CUDA MSHADOW_USE_CUDA
-#endif
-
-/*!
- *\brief whether to use cudnn library for convolution
- */
-#ifndef MXNET_USE_CUDNN
-#define MXNET_USE_CUDNN MSHADOW_USE_CUDNN
-#endif
-
-/*!
- *\brief whether to use cusolver library
- */
-#ifndef MXNET_USE_CUSOLVER
-#define MXNET_USE_CUSOLVER MSHADOW_USE_CUSOLVER
-#endif
-
-/*! \brief Error message for using gpu when MXNET_USE_CUDA==0 */
-#define MXNET_GPU_NOT_ENABLED_ERROR "GPU is not enabled"
/*!
* \brief define compatible keywords in g++
@@ -412,6 +383,7 @@ inline std::ostream& operator<<(std::ostream &out, const Context &ctx) {
#define MXNET_DESCRIBE(...) describe(__VA_ARGS__ "\n\nFrom:" __FILE__ ":" STRINGIZE(__LINE__))
#define ADD_FILELINE "\n\nDefined in " __FILE__ ":L" STRINGIZE(__LINE__)
+
#if MXNET_USE_MKLDNN == 1
constexpr size_t kMKLDNNAlign = 64;
#endif
diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h
index e9f1e2d..79f2bf5 100644
--- a/include/mxnet/c_api.h
+++ b/include/mxnet/c_api.h
@@ -208,6 +208,15 @@ MXNET_DLL const char *MXGetLastError();
//-------------------------------------
// Part 0: Global State setups
//-------------------------------------
+
+/*!
+ * \brief
+ * \param feature to check mxfeatures.h
+ * \param out set to true if the feature is enabled, false otherwise
+ * \return 0 when success, -1 when failure happens.
+ */
+MXNET_DLL int MXHasFeature(const mx_uint feature, bool* out);
+
/*!
* \brief Seed all global random number generators in mxnet.
* \param seed the random number seed.
@@ -465,6 +474,7 @@ MXNET_DLL int MXGetGPUMemoryInformation64(int dev, uint64_t *free_mem, uint64_t
*/
MXNET_DLL int MXGetVersion(int *out);
+
//-------------------------------------
// Part 1: NDArray creation and deletion
//-------------------------------------
diff --git a/include/mxnet/io.h b/include/mxnet/io.h
index 3c806d8..e18f03e 100644
--- a/include/mxnet/io.h
+++ b/include/mxnet/io.h
@@ -25,12 +25,12 @@
#ifndef MXNET_IO_H_
#define MXNET_IO_H_
-#include <dmlc/data.h>
-#include <dmlc/registry.h>
#include <vector>
#include <string>
#include <utility>
#include <queue>
+#include "dmlc/data.h"
+#include "dmlc/registry.h"
#include "./base.h"
#include "./ndarray.h"
diff --git a/include/mxnet/mxfeatures.h b/include/mxnet/mxfeatures.h
new file mode 100644
index 0000000..10f9b36
--- /dev/null
+++ b/include/mxnet/mxfeatures.h
@@ -0,0 +1,188 @@
+/*
+ * 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.
+ */
+
+/*!
+ * Copyright (c) 2018 by Contributors
+ * \file mxfeatures.h
+ * \brief check MXNet features including compile time support
+ */
+
+#pragma once
+
+#include "dmlc/base.h"
+#include "mshadow/base.h"
+
+/*!
+ *\brief whether to use opencv support
+ */
+#ifndef MXNET_USE_OPENCV
+#define MXNET_USE_OPENCV 1
+#endif
+
+/*!
+ *\brief whether to use cuda support
+ */
+#ifndef MXNET_USE_CUDA
+#define MXNET_USE_CUDA MSHADOW_USE_CUDA
+#endif
+
+/*!
+ *\brief whether to use cudnn library for convolution
+ */
+#ifndef MXNET_USE_CUDNN
+#define MXNET_USE_CUDNN MSHADOW_USE_CUDNN
+#endif
+
+#ifndef MXNET_USE_NCCL
+#define MXNET_USE_NCCL 0
+#endif
+
+/*!
+ *\brief whether to use cusolver library
+ */
+#ifndef MXNET_USE_CUSOLVER
+#define MXNET_USE_CUSOLVER MSHADOW_USE_CUSOLVER
+#endif
+
+#ifndef MXNET_ENABLE_CUDA_RTC
+#define MXNET_ENABLE_CUDA_RTC 0
+#endif
+
+/*! \brief Error message for using gpu when MXNET_USE_CUDA==0 */
+#define MXNET_GPU_NOT_ENABLED_ERROR "GPU is not enabled"
+
+
+#ifndef MXNET_USE_TENSORRT
+#define MXNET_USE_TENSORRT 0
+#endif
+
+
+#ifndef MXNET_USE_BLAS_ATLAS
+#define MXNET_USE_BLAS_ATLAS 0
+#endif
+
+#ifndef MXNET_USE_BLAS_OPEN
+#define MXNET_USE_BLAS_OPEN 0
+#endif
+
+#ifndef MXNET_USE_BLAS_MKL
+#define MXNET_USE_BLAS_MKL 0
+#endif
+
+#ifndef MXNET_USE_BLAS_APPLE
+#define MXNET_USE_BLAS_APPLE 0
+#endif
+
+#ifndef MXNET_USE_LAPACK
+#define MXNET_USE_LAPACK 0
+#endif
+
+#ifndef MXNET_USE_MKLDNN
+#define MXNET_USE_MKLDNN 0
+#endif
+
+#ifndef MXNET_USE_OPENMP
+#define MXNET_USE_OPENMP 0
+#endif
+
+#ifndef MXNET_USE_F16C
+#define MXNET_USE_F16C MSHADOW_USE_F16C
+#endif
+
+#ifndef MXNET_USE_CAFFE
+#define MXNET_USE_CAFFE 0
+#endif
+
+#ifndef MXNET_USE_DIST_KVSTORE
+#define MXNET_USE_DIST_KVSTORE 0
+#endif
+
+#ifndef MXNET_USE_SIGNAL_HANDLER
+#define MXNET_USE_SIGNAL_HANDLER 0
+#endif
+
+
+
+namespace mxnet {
+namespace features {
+// Check compile flags such as CMakeLists.txt
+
+/// Compile time features
+enum : uint32_t {
+ // NVIDIA, CUDA
+ CUDA = 0,
+ CUDNN,
+ NCCL,
+ CUDA_RTC,
+ TENSORRT,
+
+ // CPU Features / optimizations
+ CPU_SSE,
+ CPU_SSE2,
+ CPU_SSE3,
+ CPU_SSE4_1,
+ CPU_SSE4_2,
+ CPU_SSE4A, // AMD extensions to SSE4
+ CPU_AVX,
+ CPU_AVX2,
+
+
+ // Multiprocessing / CPU / System
+ OPENMP,
+ SSE,
+ F16C,
+ JEMALLOC,
+
+ // Math libraries & BLAS
+ // Flavour of BLAS
+ BLAS_OPEN,
+ BLAS_ATLAS,
+ // Intel(R) Math Kernel Library
+ BLAS_MKL,
+ BLAS_APPLE,
+ // Other math libraries:
+ // Linear Algebra PACKage
+ LAPACK,
+ // Intel(R) Math Kernel Library for Deep Neural Networks
+ MKLDNN,
+
+ // Image processing
+ OPENCV,
+
+ // Misc
+ CAFFE,
+ PROFILER,
+ DIST_KVSTORE,
+ CXX14,
+ // Signal handler to print stack traces on exceptions
+ SIGNAL_HANDLER,
+ DEBUG,
+
+ // size indicator
+ MAX_FEATURES
+};
+
+
+/*!
+ * \return true if the given feature is supported
+ */
+bool is_enabled(uint32_t feat);
+
+} // namespace features
+} // namespace mxnet
diff --git a/python/mxnet/gluon/utils.py b/python/mxnet/gluon/utils.py
index 7832498..55edd95 100644
--- a/python/mxnet/gluon/utils.py
+++ b/python/mxnet/gluon/utils.py
@@ -222,7 +222,7 @@ else:
_MOVEFILE_WRITE_THROUGH = 0x8
_windows_default_flags = _MOVEFILE_WRITE_THROUGH
- text_type = unicode if sys.version_info[0] == 2 else str # noqa
+ text_type = unicode if sys.version_info[0] == 2 else str # pylint: disable=undefined-variable
def _str_to_unicode(x):
"""Handle text decoding. Internal use only"""
diff --git a/python/mxnet/mxfeatures.py b/python/mxnet/mxfeatures.py
new file mode 100644
index 0000000..c546151
--- /dev/null
+++ b/python/mxnet/mxfeatures.py
@@ -0,0 +1,103 @@
+# 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.
+
+# coding: utf-8
+# pylint: disable=not-an-iterable
+
+"""runtime detection of compile time features in the native library"""
+
+import ctypes
+import enum
+from .base import _LIB, check_call, mx_uint
+
+feature_names = [
+ "CUDA",
+ "CUDNN",
+ "NCCL",
+ "CUDA_RTC",
+ "TENSORRT",
+ "CPU_SSE",
+ "CPU_SSE2",
+ "CPU_SSE3",
+ "CPU_SSE4_1",
+ "CPU_SSE4_2",
+ "CPU_SSE4A",
+ "CPU_AVX",
+ "CPU_AVX2",
+ "OPENMP",
+ "SSE",
+ "F16C",
+ "JEMALLOC",
+ "BLAS_OPEN",
+ "BLAS_ATLAS",
+ "BLAS_MKL",
+ "BLAS_APPLE",
+ "LAPACK",
+ "MKLDNN",
+ "OPENCV",
+ "CAFFE",
+ "PROFILER",
+ "DIST_KVSTORE",
+ "CXX14",
+ "SIGNAL_HANDLER",
+ "DEBUG"
+]
+
+
+Feature = enum.Enum('Feature', {name: index for index, name in enumerate(feature_names)})
+
+
+def has_feature(feature):
+ """
+ Check the library for compile-time feature at runtime
+
+ Parameters
+ ----------
+ feature : int
+ An integer representing the feature to check
+
+ Returns
+ -------
+ boolean
+ True if the feature is enabled, false otherwise
+ """
+ res = ctypes.c_bool()
+ check_call(_LIB.MXHasFeature(mx_uint(feature), ctypes.byref(res)))
+ return res.value
+
+
+def features_enabled():
+ """
+ Returns
+ -------
+ features: list of Feature
+ list of enabled features in the back-end
+ """
+ res = []
+ for f in Feature:
+ if has_feature(f.value):
+ res.append(f)
+ return res
+
+def features_enabled_str(sep=', '):
+ """
+ Returns
+ -------
+ string with a comma separated list of enabled features in the back-end. For example:
+ "CPU_SSE, OPENMP, F16C, LAPACK, MKLDNN, OPENCV, SIGNAL_HANDLER, DEBUG"
+ """
+ return sep.join(map(lambda x: x.name, features_enabled()))
diff --git a/src/c_api/c_api.cc b/src/c_api/c_api.cc
index 80bd605..b436e8c 100644
--- a/src/c_api/c_api.cc
+++ b/src/c_api/c_api.cc
@@ -22,20 +22,6 @@
* \file c_api.cc
* \brief C API of mxnet
*/
-#include <dmlc/base.h>
-#include <dmlc/logging.h>
-#include <dmlc/io.h>
-#include <dmlc/memory_io.h>
-#include <dmlc/recordio.h>
-#include <dmlc/omp.h>
-#include <mxnet/base.h>
-#include <mxnet/ndarray.h>
-#include <mxnet/operator.h>
-#include <mxnet/io.h>
-#include <mxnet/c_api.h>
-#include <mxnet/kvstore.h>
-#include <mxnet/rtc.h>
-#include <mxnet/storage.h>
#include <vector>
#include <sstream>
#include <string>
@@ -43,6 +29,21 @@
#include <memory>
#include <functional>
#include <utility>
+#include "dmlc/base.h"
+#include "dmlc/logging.h"
+#include "dmlc/io.h"
+#include "dmlc/memory_io.h"
+#include "dmlc/recordio.h"
+#include "dmlc/omp.h"
+#include "mxnet/base.h"
+#include "mxnet/ndarray.h"
+#include "mxnet/operator.h"
+#include "mxnet/io.h"
+#include "mxnet/c_api.h"
+#include "mxnet/kvstore.h"
+#include "mxnet/rtc.h"
+#include "mxnet/storage.h"
+#include "mxnet/mxfeatures.h"
#include "./c_api_common.h"
#include "../operator/custom/custom-inl.h"
#include "../operator/tensor/matrix_op-inl.h"
@@ -85,6 +86,13 @@ inline int MXAPIGetFunctionRegInfo(const FunRegType *e,
}
// NOTE: return value is added in API_END
+
+int MXHasFeature(const mx_uint feature, bool* out) {
+ API_BEGIN();
+ *out = features::is_enabled(feature);
+ API_END();
+}
+
int MXRandomSeed(int seed) {
API_BEGIN();
mxnet::RandomSeed(seed);
diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc
index 73a8a7c..8517c9c 100644
--- a/src/c_api/c_api_symbolic.cc
+++ b/src/c_api/c_api_symbolic.cc
@@ -22,12 +22,12 @@
* \file c_api_symbolic.cc
* \brief C API of mxnet
*/
-#include <mxnet/base.h>
-#include <mxnet/c_api.h>
-#include <nnvm/c_api.h>
-#include <nnvm/pass.h>
-#include <nnvm/pass_functions.h>
-#include <nnvm/symbolic.h>
+#include "mxnet/base.h"
+#include "mxnet/c_api.h"
+#include "nnvm/c_api.h"
+#include "nnvm/pass.h"
+#include "nnvm/pass_functions.h"
+#include "nnvm/symbolic.h"
#include "./c_api_common.h"
#include "../operator/operator_common.h"
#include "../executor/exec_pass.h"
diff --git a/src/mxfeatures.cc b/src/mxfeatures.cc
new file mode 100644
index 0000000..7a435d7
--- /dev/null
+++ b/src/mxfeatures.cc
@@ -0,0 +1,112 @@
+/*
+ * 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.
+ */
+
+/*!
+ * Copyright (c) 2018 by Contributors
+ * \file mxfeatures.cc
+ * \brief check MXNet features including compile time support
+ */
+
+#include "mxnet/mxfeatures.h"
+#include <bitset>
+
+namespace mxnet {
+namespace features {
+
+class FeatureSet {
+ public:
+ FeatureSet() :
+ feature_bits() {
+ // GPU
+ feature_bits.set(CUDA, MXNET_USE_CUDA);
+ feature_bits.set(CUDNN, MXNET_USE_CUDNN);
+ feature_bits.set(NCCL, MXNET_USE_NCCL);
+ feature_bits.set(CUDA_RTC, MXNET_ENABLE_CUDA_RTC);
+ feature_bits.set(TENSORRT, MXNET_USE_TENSORRT);
+
+ // Check flags for example with gcc -msse3 -mavx2 -dM -E - < /dev/null | egrep "SSE|AVX"
+#if __SSE__
+ feature_bits.set(CPU_SSE);
+#endif
+#if __SSE2__
+ feature_bits.set(CPU_SSE2);
+#endif
+#if __SSE3__
+ feature_bits.set(CPU_SSE3);
+#endif
+#if __SSE4_1__
+ feature_bits.set(CPU_SSE4_1);
+#endif
+#if __SSE4_2__
+ feature_bits.set(CPU_SSE4_2);
+#endif
+#if __SSE4A__
+ feature_bits.set(CPU_SSE4A);
+#endif
+#if __AVX__
+ feature_bits.set(CPU_AVX);
+#endif
+#if __AVX2__
+ feature_bits.set(CPU_AVX2);
+#endif
+
+ // CPU
+ feature_bits.set(OPENMP, MXNET_USE_OPENMP);
+ feature_bits.set(F16C, MXNET_USE_F16C);
+
+ // Math
+ feature_bits.set(BLAS_OPEN, MXNET_USE_BLAS_OPEN);
+ feature_bits.set(BLAS_ATLAS, MXNET_USE_BLAS_ATLAS);
+ feature_bits.set(BLAS_MKL, MXNET_USE_BLAS_MKL);
+ feature_bits.set(BLAS_APPLE, MXNET_USE_BLAS_APPLE);
+ feature_bits.set(LAPACK, MXNET_USE_LAPACK);
+ feature_bits.set(MKLDNN, MXNET_USE_MKLDNN);
+
+ // Image
+ feature_bits.set(OPENCV, MXNET_USE_OPENCV);
+
+ // Misc
+ feature_bits.set(CAFFE, MXNET_USE_CAFFE);
+ feature_bits.set(DIST_KVSTORE, MXNET_USE_DIST_KVSTORE);
+ feature_bits.set(SIGNAL_HANDLER, MXNET_USE_SIGNAL_HANDLER);
+#ifndef NDEBUG
+ feature_bits.set(DEBUG);
+#endif
+
+#if USE_JEMALLOC == 1
+ feature_bits.set(JEMALLOC);
+#endif
+ }
+ bool is_enabled(const unsigned feat) const {
+ CHECK_LT(feat, MAX_FEATURES);
+ return feature_bits.test(feat);
+ }
+
+ private:
+ std::bitset<MAX_FEATURES> feature_bits;
+};
+
+static FeatureSet featureSet;
+
+bool is_enabled(const unsigned feat) {
+ return featureSet.is_enabled(feat);
+}
+
+} // namespace features
+} // namespace mxnet
diff --git a/tests/python/unittest/test_features.py b/tests/python/unittest/test_features.py
new file mode 100644
index 0000000..ff91181
--- /dev/null
+++ b/tests/python/unittest/test_features.py
@@ -0,0 +1,40 @@
+# 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.
+
+import mxnet as mx
+import sys
+from mxnet.mxfeatures import *
+from mxnet.base import MXNetError
+from nose.tools import *
+
+def test_runtime_features():
+ for f in Feature:
+ res = has_feature(f.value)
+ ok_(type(res) is bool)
+ for f in features_enabled():
+ ok_(type(f) is Feature)
+ ok_(type(features_enabled_str()) is str)
+ print("Features enabled: {}".format(features_enabled_str()))
+
+@raises(MXNetError)
+def test_has_feature_2large():
+ has_feature(sys.maxsize)
+
+
+if __name__ == "__main__":
+ import nose
+ nose.runmodule()