From 22ef8e5654c45d1f5404e3add6ad19678c0b80a9 Mon Sep 17 00:00:00 2001 From: Tongzhou Wang Date: Mon, 19 Mar 2018 15:43:14 -0400 Subject: [PATCH] [fft][1 of 3] build system and helpers to support cuFFT and MKL (#5855) This is the first of three PRs that #5537 will be split into. This PR adds mkl headers to included files, and provides helper functions for MKL fft and cuFFT. In particular, on POSIX, headers are using mkl-include from conda, and on Windows, it is from a new file @yf225 and I made and uploaded to s3. * add mkl-include to required packages * include MKL headers; add AT_MKL_ENABLED flag; add a method to query MKL availability * Add MKL and CUFFT helpers --- .jenkins/macos-build-test.sh | 2 +- .jenkins/win-build.sh | 7 +- Dockerfile | 6 +- README.md | 2 +- aten/CMakeLists.txt | 3 + aten/src/ATen/CMakeLists.txt | 9 ++- aten/src/ATen/Config.h.in | 3 +- aten/src/ATen/Context.cpp | 8 ++ aten/src/ATen/Context.h | 7 +- aten/src/ATen/mkl/Descriptors.h | 44 ++++++++++ aten/src/ATen/mkl/Exceptions.h | 19 +++++ aten/src/ATen/mkl/Limits.h | 11 +++ aten/src/ATen/mkl/README.md | 4 + aten/src/ATen/native/cuda/CuFFTUtils.h | 84 ++++++++++++++++++++ aten/src/ATen/test/verify_api_visibility.cpp | 4 + aten/src/TH/cmake/FindLAPACK.cmake | 2 +- aten/src/TH/cmake/FindMKL.cmake | 12 +-- test/common.py | 3 + torch/backends/mkl/__init__.py | 6 ++ torch/csrc/Module.cpp | 2 + 20 files changed, 219 insertions(+), 19 deletions(-) create mode 100644 aten/src/ATen/mkl/Descriptors.h create mode 100644 aten/src/ATen/mkl/Exceptions.h create mode 100644 aten/src/ATen/mkl/Limits.h create mode 100644 aten/src/ATen/mkl/README.md create mode 100644 aten/src/ATen/native/cuda/CuFFTUtils.h create mode 100644 torch/backends/mkl/__init__.py diff --git a/.jenkins/macos-build-test.sh b/.jenkins/macos-build-test.sh index 6f9cd071953f..4d1e8bd129ef 100755 --- a/.jenkins/macos-build-test.sh +++ b/.jenkins/macos-build-test.sh @@ -9,7 +9,7 @@ rm -rf $PWD/miniconda3 bash $PWD/miniconda3.sh -b -p $PWD/miniconda3 export PATH="$PWD/miniconda3/bin:$PATH" source $PWD/miniconda3/bin/activate -conda install -y numpy pyyaml setuptools cmake cffi ninja +conda install -y mkl mkl-include numpy pyyaml setuptools cmake cffi ninja # Build and test PyTorch git submodule update --init --recursive diff --git a/.jenkins/win-build.sh b/.jenkins/win-build.sh index f1b056a88b05..eb0131f1a946 100755 --- a/.jenkins/win-build.sh +++ b/.jenkins/win-build.sh @@ -32,8 +32,9 @@ cat >ci_scripts/build_pytorch.bat < +#include + +namespace at { namespace native { + +struct DftiDescriptorDeleter { + void operator()(DFTI_DESCRIPTOR* desc) { + if (desc != nullptr) { + MKL_DFTI_CHECK(DftiFreeDescriptor(&desc)); + } + } +}; + +class DftiDescriptor { +public: + void init(DFTI_CONFIG_VALUE precision, DFTI_CONFIG_VALUE signal_type, MKL_LONG signal_ndim, MKL_LONG* sizes) { + if (desc_ != nullptr) { + throw std::runtime_error("DFTI DESCRIPTOR can only be initialized once"); + } + DFTI_DESCRIPTOR *raw_desc; + if (signal_ndim == 1) { + MKL_DFTI_CHECK(DftiCreateDescriptor(&raw_desc, precision, signal_type, 1, sizes[0])); + } else { + MKL_DFTI_CHECK(DftiCreateDescriptor(&raw_desc, precision, signal_type, signal_ndim, sizes)); + } + desc_.reset(raw_desc); + } + + DFTI_DESCRIPTOR *get() const { + if (desc_ == nullptr) { + throw std::runtime_error("DFTI DESCRIPTOR has not been initialized"); + } + return desc_.get(); + } + +private: + std::unique_ptr desc_; +}; + + +}} // at::native diff --git a/aten/src/ATen/mkl/Exceptions.h b/aten/src/ATen/mkl/Exceptions.h new file mode 100644 index 000000000000..e954a074cdbb --- /dev/null +++ b/aten/src/ATen/mkl/Exceptions.h @@ -0,0 +1,19 @@ +#pragma once + +#include +#include +#include +#include + +namespace at { namespace native { + +static inline void MKL_DFTI_CHECK(MKL_INT status) +{ + if (status && !DftiErrorClass(status, DFTI_NO_ERROR)) { + std::ostringstream ss; + ss << "MKL FFT error: " << DftiErrorMessage(status); + throw std::runtime_error(ss.str()); + } +} + +}} // namespace at::native diff --git a/aten/src/ATen/mkl/Limits.h b/aten/src/ATen/mkl/Limits.h new file mode 100644 index 000000000000..b0d3829e458e --- /dev/null +++ b/aten/src/ATen/mkl/Limits.h @@ -0,0 +1,11 @@ +#pragma once + +#include + +namespace at { namespace native { + + // Since size of MKL_LONG varies on different platforms (linux 64 bit, windows + // 32 bit), we need to programmatically calculate the max. + static int64_t MKL_LONG_MAX = ((1LL << (sizeof(MKL_LONG) * 8 - 2)) - 1) * 2 + 1; + +}} // namespace diff --git a/aten/src/ATen/mkl/README.md b/aten/src/ATen/mkl/README.md new file mode 100644 index 000000000000..2916246ec841 --- /dev/null +++ b/aten/src/ATen/mkl/README.md @@ -0,0 +1,4 @@ +All files living in this directory are written with the assumption that MKL is available, +which means that these code are not guarded by `#if AT_MKL_ENABLED()`. Therefore, whenever +you need to use definitions from here, please guard the `#include` and +definition usages with `#if AT_MKL_ENABLED()` macro, e.g. [SpectralOps.cpp](native/mkl/SpectralOps.cpp). \ No newline at end of file diff --git a/aten/src/ATen/native/cuda/CuFFTUtils.h b/aten/src/ATen/native/cuda/CuFFTUtils.h new file mode 100644 index 000000000000..9048734efaa5 --- /dev/null +++ b/aten/src/ATen/native/cuda/CuFFTUtils.h @@ -0,0 +1,84 @@ +#pragma once + +#include "ATen/ATen.h" +#include "ATen/Config.h" + +#include +#include +#include +#include +#include + + +namespace at { namespace native { + +static inline std::string _cudaGetErrorEnum(cufftResult error) +{ + switch (error) + { + case CUFFT_SUCCESS: + return "CUFFT_SUCCESS"; + case CUFFT_INVALID_PLAN: + return "CUFFT_INVALID_PLAN"; + case CUFFT_ALLOC_FAILED: + return "CUFFT_ALLOC_FAILED"; + case CUFFT_INVALID_TYPE: + return "CUFFT_INVALID_TYPE"; + case CUFFT_INVALID_VALUE: + return "CUFFT_INVALID_VALUE"; + case CUFFT_INTERNAL_ERROR: + return "CUFFT_INTERNAL_ERROR"; + case CUFFT_EXEC_FAILED: + return "CUFFT_EXEC_FAILED"; + case CUFFT_SETUP_FAILED: + return "CUFFT_SETUP_FAILED"; + case CUFFT_INVALID_SIZE: + return "CUFFT_INVALID_SIZE"; + case CUFFT_UNALIGNED_DATA: + return "CUFFT_UNALIGNED_DATA"; + case CUFFT_INCOMPLETE_PARAMETER_LIST: + return "CUFFT_INCOMPLETE_PARAMETER_LIST"; + case CUFFT_INVALID_DEVICE: + return "CUFFT_INVALID_DEVICE"; + case CUFFT_PARSE_ERROR: + return "CUFFT_PARSE_ERROR"; + case CUFFT_NO_WORKSPACE: + return "CUFFT_NO_WORKSPACE"; + case CUFFT_NOT_IMPLEMENTED: + return "CUFFT_NOT_IMPLEMENTED"; + case CUFFT_LICENSE_ERROR: + return "CUFFT_LICENSE_ERROR"; + case CUFFT_NOT_SUPPORTED: + return "CUFFT_NOT_SUPPORTED"; + default: + std::ostringstream ss; + ss << "unknown error " << error; + return ss.str(); + } +} + +static inline void CUFFT_CHECK(cufftResult error) +{ + if (error != CUFFT_SUCCESS) { + std::ostringstream ss; + ss << "cuFFT error: " << _cudaGetErrorEnum(error); + throw std::runtime_error(ss.str()); + } +} + +class CufftHandle { +public: + explicit CufftHandle() { + CUFFT_CHECK(cufftCreate(&raw_plan)); + } + + const cufftHandle &get() const { return raw_plan; } + + ~CufftHandle() { + CUFFT_CHECK(cufftDestroy(raw_plan)); + } +private: + cufftHandle raw_plan; +}; + +}} // at::native diff --git a/aten/src/ATen/test/verify_api_visibility.cpp b/aten/src/ATen/test/verify_api_visibility.cpp index ed5f88a983b3..ba01240ee8a4 100644 --- a/aten/src/ATen/test/verify_api_visibility.cpp +++ b/aten/src/ATen/test/verify_api_visibility.cpp @@ -8,4 +8,8 @@ #error "AT_CUDNN_ENABLED should not be visible in public headers" #endif +#ifdef AT_MKL_ENABLED +#error "AT_MKL_ENABLED should not be visible in public headers" +#endif + auto main() -> int {} diff --git a/aten/src/TH/cmake/FindLAPACK.cmake b/aten/src/TH/cmake/FindLAPACK.cmake index 9eca0730fc26..38cb9d04b3b9 100644 --- a/aten/src/TH/cmake/FindLAPACK.cmake +++ b/aten/src/TH/cmake/FindLAPACK.cmake @@ -180,7 +180,7 @@ IF (NOT LAPACK_FOUND AND LAPACK_FIND_REQUIRED) ENDIF (NOT LAPACK_FOUND AND LAPACK_FIND_REQUIRED) IF(NOT LAPACK_FIND_QUIETLY) IF(LAPACK_FOUND) - MESSAGE(STATUS "Found a library with LAPACK API. (${LAPACK_INFO})") + MESSAGE(STATUS "Found a library with LAPACK API (${LAPACK_INFO}).") ELSE(LAPACK_FOUND) MESSAGE(STATUS "Cannot find a library with LAPACK API. Not using LAPACK.") ENDIF(LAPACK_FOUND) diff --git a/aten/src/TH/cmake/FindMKL.cmake b/aten/src/TH/cmake/FindMKL.cmake index 88f0aa3281d4..49f841ab81d5 100644 --- a/aten/src/TH/cmake/FindMKL.cmake +++ b/aten/src/TH/cmake/FindMKL.cmake @@ -175,7 +175,7 @@ FOREACH(mklrtl ${mklrtls} "") IF (NOT MKL_LIBRARIES AND NOT INTEL_MKL_SEQUENTIAL) CHECK_ALL_LIBRARIES(MKL_LIBRARIES cblas_sgemm "mkl_${mkliface}${mkl64};${mklthread};mkl_core;${mklrtl};pthread;${mkl_m};${mkl_dl}" "") - ENDIF (NOT MKL_LIBRARIES AND NOT INTEL_MKL_SEQUENTIAL) + ENDIF (NOT MKL_LIBRARIES AND NOT INTEL_MKL_SEQUENTIAL) ENDFOREACH(mklthread) ENDFOREACH(mkl64) ENDFOREACH(mkliface) @@ -200,7 +200,7 @@ FOREACH(mklrtl ${mklrtls} "") IF (NOT MKL_LIBRARIES) CHECK_ALL_LIBRARIES(MKL_LIBRARIES cblas_sgemm "mkl_${mkliface}${mkl64};${mklthread};mkl_core;${mklrtl};pthread;${mkl_m};${mkl_dl}" "") - ENDIF (NOT MKL_LIBRARIES) + ENDIF (NOT MKL_LIBRARIES) ENDFOREACH(mklthread) ENDFOREACH(mkl64) ENDFOREACH(mkliface) @@ -211,7 +211,7 @@ IF (NOT MKL_LIBRARIES) SET(MKL_VERSION 900) CHECK_ALL_LIBRARIES(MKL_LIBRARIES cblas_sgemm "mkl;guide;pthread;m" "") -ENDIF (NOT MKL_LIBRARIES) +ENDIF (NOT MKL_LIBRARIES) # Include files IF (MKL_LIBRARIES) @@ -228,7 +228,7 @@ IF (MKL_LIBRARIES) MARK_AS_ADVANCED(MKL_LAPACK_LIBRARIES) ENDIF (NOT MKL_LAPACK_LIBRARIES) IF (NOT MKL_SCALAPACK_LIBRARIES) - FIND_LIBRARY(MKL_SCALAPACK_LIBRARIES NAMES "mkl_scalapack${mkl64}${mkls}") + FIND_LIBRARY(MKL_SCALAPACK_LIBRARIES NAMES "mkl_scalapack${mkl64}${mkls}") MARK_AS_ADVANCED(MKL_SCALAPACK_LIBRARIES) ENDIF (NOT MKL_SCALAPACK_LIBRARIES) IF (NOT MKL_SOLVER_LIBRARIES) @@ -243,7 +243,7 @@ IF (MKL_LIBRARIES) ENDFOREACH(mkl64) ENDIF (MKL_LIBRARIES) -# LibIRC: intel compiler always links this; +# LibIRC: intel compiler always links this; # gcc does not; but mkl kernels sometimes need it. IF (MKL_LIBRARIES) IF (CMAKE_COMPILER_IS_GNUCC) @@ -269,7 +269,7 @@ ENDIF (MKL_LIBRARIES) # Standard termination IF(NOT MKL_FOUND AND MKL_FIND_REQUIRED) - MESSAGE(FATAL_ERROR "MKL library not found. Please specify library location") + MESSAGE(FATAL_ERROR "MKL library not found. Please specify library location") ENDIF(NOT MKL_FOUND AND MKL_FIND_REQUIRED) IF(NOT MKL_FIND_QUIETLY) IF(MKL_FOUND) diff --git a/test/common.py b/test/common.py index 2980cd5a0a9b..07b4f20b224a 100644 --- a/test/common.py +++ b/test/common.py @@ -19,6 +19,7 @@ import torch.cuda from torch.autograd import Variable from torch._six import string_classes import torch.backends.cudnn +import torch.backends.mkl torch.set_default_tensor_type('torch.DoubleTensor') @@ -54,6 +55,8 @@ try: except ImportError: TEST_SCIPY = False +TEST_MKL = torch.backends.mkl.is_available() + def skipIfNoLapack(fn): @wraps(fn) diff --git a/torch/backends/mkl/__init__.py b/torch/backends/mkl/__init__.py new file mode 100644 index 000000000000..f3d27d1fa08c --- /dev/null +++ b/torch/backends/mkl/__init__.py @@ -0,0 +1,6 @@ +import torch + + +def is_available(): + r"""Returns whether PyTorch is built with MKL support.""" + return torch._C.has_mkl diff --git a/torch/csrc/Module.cpp b/torch/csrc/Module.cpp index 611175ef70d0..8e9b185a96f8 100644 --- a/torch/csrc/Module.cpp +++ b/torch/csrc/Module.cpp @@ -505,6 +505,8 @@ static PyObject* initModule() { // setting up TH Errors so that they throw C++ exceptions at::init(); + ASSERT_TRUE(PyModule_AddObject(module, "has_mkl", at::hasMKL() ? Py_True : Py_False) == 0); + auto& defaultGenerator = at::globalContext().defaultGenerator(at::kCPU); THPDefaultGenerator = (THPGenerator*)THPGenerator_NewWithGenerator( defaultGenerator);