Skip to content

Commit 22ef8e5

Browse files
ssnlezyang
ssnl
authored and
ezyang
committed
[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
1 parent d11b7fb commit 22ef8e5

File tree

20 files changed

+219
-19
lines changed

20 files changed

+219
-19
lines changed

‎.jenkins/macos-build-test.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ rm -rf $PWD/miniconda3
99
bash $PWD/miniconda3.sh -b -p $PWD/miniconda3
1010
export PATH="$PWD/miniconda3/bin:$PATH"
1111
source $PWD/miniconda3/bin/activate
12-
conda install -y numpy pyyaml setuptools cmake cffi ninja
12+
conda install -y mkl mkl-include numpy pyyaml setuptools cmake cffi ninja
1313

1414
# Build and test PyTorch
1515
git submodule update --init --recursive

‎.jenkins/win-build.sh‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,9 @@ cat >ci_scripts/build_pytorch.bat <<EOL
3232
set PATH=C:\\Program Files\\CMake\\bin;C:\\Program Files\\7-Zip;C:\\curl-7.57.0-win64-mingw\\bin;C:\\Program Files\\Git\\cmd;C:\\Program Files\\Amazon\\AWSCLI;%PATH%
3333
3434
:: Install MKL
35-
aws s3 cp s3://ossci-windows/mkl.7z mkl.7z --quiet && 7z x -aoa mkl.7z -omkl
36-
set LIB=%cd%\\mkl;%LIB%
35+
aws s3 cp s3://ossci-windows/mkl_with_headers.7z mkl.7z --quiet && 7z x -aoa mkl.7z -omkl
36+
set CMAKE_INCLUDE_PATH=%cd%\\mkl\\include
37+
set LIB=%cd%\\mkl\\lib;%LIB
3738
3839
:: Install MAGMA
3940
aws s3 cp s3://ossci-windows/magma_cuda90_release.7z magma_cuda90_release.7z --quiet && 7z x -aoa magma_cuda90_release.7z -omagma_cuda90_release
@@ -47,7 +48,7 @@ IF EXIST C:\\Jenkins\\Miniconda3 ( rd /s /q C:\\Jenkins\\Miniconda3 )
4748
curl https://repo.continuum.io/miniconda/Miniconda3-latest-Windows-x86_64.exe -O
4849
.\Miniconda3-latest-Windows-x86_64.exe /InstallationType=JustMe /RegisterPython=0 /S /AddToPath=0 /D=C:\\Jenkins\\Miniconda3
4950
call C:\\Jenkins\\Miniconda3\\Scripts\\activate.bat C:\\Jenkins\\Miniconda3
50-
call conda install -y -q numpy mkl cffi pyyaml boto3
51+
call conda install -y -q numpy cffi pyyaml boto3
5152
5253
:: Install ninja
5354
pip install ninja

‎Dockerfile‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,11 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
1414

1515
RUN curl -o ~/miniconda.sh -O https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
1616
chmod +x ~/miniconda.sh && \
17-
~/miniconda.sh -b -p /opt/conda && \
17+
~/miniconda.sh -b -p /opt/conda && \
1818
rm ~/miniconda.sh && \
19-
/opt/conda/bin/conda install numpy pyyaml scipy ipython mkl && \
19+
/opt/conda/bin/conda install numpy pyyaml scipy ipython mkl mkl-include && \
2020
/opt/conda/bin/conda install -c soumith magma-cuda90 && \
21-
/opt/conda/bin/conda clean -ya
21+
/opt/conda/bin/conda clean -ya
2222
ENV PATH /opt/conda/bin:$PATH
2323
# This must be done before pip so that requirements.txt is available
2424
WORKDIR /opt/pytorch

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ On Linux
171171
export CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" # [anaconda root directory]
172172

173173
# Install basic dependencies
174-
conda install numpy pyyaml mkl setuptools cmake cffi typing
174+
conda install numpy pyyaml mkl mkl-include setuptools cmake cffi typing
175175

176176
# Add LAPACK support for the GPU
177177
conda install -c pytorch magma-cuda80 # or magma-cuda90 if CUDA 9

‎aten/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -380,10 +380,13 @@ MACRO(Install_Required_Library ln)
380380
ENDMACRO(Install_Required_Library libname)
381381

382382
FIND_PACKAGE(BLAS)
383+
SET(AT_MKL_ENABLED 0)
383384
IF(BLAS_FOUND)
384385
SET(USE_BLAS 1)
385386
IF(BLAS_INFO STREQUAL "mkl")
386387
ADD_DEFINITIONS(-DTH_BLAS_MKL)
388+
INCLUDE_DIRECTORIES(${BLAS_INCLUDE_DIR}) # include MKL headers
389+
SET(AT_MKL_ENABLED 1)
387390
ENDIF()
388391
ENDIF(BLAS_FOUND)
389392

‎aten/src/ATen/CMakeLists.txt‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -128,12 +128,14 @@ FILE(GLOB base_cpp RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "*.cpp")
128128
FILE(GLOB native_cpp RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "native/*.cpp")
129129
FILE(GLOB native_cudnn_cpp RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "native/cudnn/*.cpp")
130130
FILE(GLOB native_cuda_cu RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "native/cuda/*.cu")
131+
FILE(GLOB native_mkl_cpp RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "native/mkl/*.cpp")
131132

132133
FILE(GLOB_RECURSE cuda_h
133134
RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
134135
"cuda/*.cuh" "cuda/*.h" "cudnn/*.cuh" "cudnn/*.h")
135136

136137
FILE(GLOB cudnn_cpp RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "cudnn/*.cpp")
138+
FILE(GLOB mkl_cpp RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "mkl/*.cpp")
137139

138140
FILE(GLOB all_python RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "*.py")
139141

@@ -179,8 +181,7 @@ ADD_CUSTOM_TARGET(aten_files_are_generated
179181
)
180182

181183

182-
SET(all_cpp ${base_cpp} ${native_cpp} ${native_cudnn_cpp} ${generated_cpp} ${ATen_CPU_SRCS} ${cpu_kernel_cpp})
183-
184+
SET(all_cpp ${base_cpp} ${native_cpp} ${native_cudnn_cpp} ${native_mkl_cpp} ${generated_cpp} ${ATen_CPU_SRCS} ${cpu_kernel_cpp})
184185

185186
INCLUDE_DIRECTORIES(${ATen_CPU_INCLUDE})
186187
IF(NOT NO_CUDA)
@@ -192,6 +193,9 @@ IF(NOT NO_CUDA)
192193
IF(CUDNN_FOUND)
193194
SET(all_cpp ${all_cpp} ${cudnn_cpp})
194195
ENDIF()
196+
IF(AT_MKL_ENABLED)
197+
SET(all_cpp ${all_cpp} ${mkl_cpp})
198+
ENDIF()
195199
endif()
196200

197201
filter_list(generated_h generated_cpp "\\.h$")
@@ -309,6 +313,7 @@ IF(CUDA_FOUND)
309313
${CUDA_cusparse_LIBRARY}
310314
${CUDA_curand_LIBRARY})
311315
CUDA_ADD_CUBLAS_TO_TARGET(ATen)
316+
CUDA_ADD_CUFFT_TO_TARGET(ATen)
312317

313318
if(CUDNN_FOUND)
314319
target_link_libraries(ATen ${CUDNN_LIBRARIES})

‎aten/src/ATen/Config.h.in‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
11
#pragma once
22

3-
// Test these using #if AT_CUDA_ENABLED()(), not #ifdef, so that it's
3+
// Test these using #if AT_CUDA_ENABLED(), not #ifdef, so that it's
44
// obvious if you forgot to include Config.h
55
// c.f. https://stackoverflow.com/questions/33759787/generating-an-error-if-checked-boolean-macro-is-not-defined
66

77
#define AT_CUDA_ENABLED() @AT_CUDA_ENABLED@
88
#define AT_CUDNN_ENABLED() @AT_CUDNN_ENABLED@
99
#define AT_NNPACK_ENABLED() @AT_NNPACK_ENABLED@
10+
#define AT_MKL_ENABLED() @AT_MKL_ENABLED@
1011

1112
#if !AT_CUDA_ENABLED() && AT_CUDNN_ENABLED()
1213
#error "Cannot enable CuDNN without CUDA"

‎aten/src/ATen/Context.cpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,14 @@ void Context::setBenchmarkCuDNN(bool b) {
8989
benchmark_cudnn = b;
9090
}
9191

92+
bool Context::hasMKL() const {
93+
#if AT_MKL_ENABLED()
94+
return true;
95+
#else
96+
return false;
97+
#endif
98+
}
99+
92100
bool Context::hasCUDA() const {
93101
#if AT_CUDA_ENABLED()
94102
int count;

‎aten/src/ATen/Context.h‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ class AT_API Context {
4343
runtime_error("%s backend type not enabled.",toString(p));
4444
return *generator;
4545
}
46+
bool hasMKL() const;
4647
bool hasCUDA() const;
4748
int64_t current_device() const;
4849
// defined in header so that getType has ability to inline
@@ -103,7 +104,7 @@ static inline void init() {
103104
}
104105

105106
static inline Type& getType(Backend p, ScalarType s) {
106-
return globalContext().getType(p,s);
107+
return globalContext().getType(p, s);
107108
}
108109

109110
static inline Type& CPU(ScalarType s) {
@@ -118,6 +119,10 @@ static inline bool hasCUDA() {
118119
return globalContext().hasCUDA();
119120
}
120121

122+
static inline bool hasMKL() {
123+
return globalContext().hasMKL();
124+
}
125+
121126
static inline int64_t current_device() {
122127
return globalContext().current_device();
123128
}

‎aten/src/ATen/mkl/Descriptors.h‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
#pragma once
2+
3+
#include "Exceptions.h"
4+
#include <mkl_dfti.h>
5+
#include <ATen/Tensor.h>
6+
7+
namespace at { namespace native {
8+
9+
struct DftiDescriptorDeleter {
10+
void operator()(DFTI_DESCRIPTOR* desc) {
11+
if (desc != nullptr) {
12+
MKL_DFTI_CHECK(DftiFreeDescriptor(&desc));
13+
}
14+
}
15+
};
16+
17+
class DftiDescriptor {
18+
public:
19+
void init(DFTI_CONFIG_VALUE precision, DFTI_CONFIG_VALUE signal_type, MKL_LONG signal_ndim, MKL_LONG* sizes) {
20+
if (desc_ != nullptr) {
21+
throw std::runtime_error("DFTI DESCRIPTOR can only be initialized once");
22+
}
23+
DFTI_DESCRIPTOR *raw_desc;
24+
if (signal_ndim == 1) {
25+
MKL_DFTI_CHECK(DftiCreateDescriptor(&raw_desc, precision, signal_type, 1, sizes[0]));
26+
} else {
27+
MKL_DFTI_CHECK(DftiCreateDescriptor(&raw_desc, precision, signal_type, signal_ndim, sizes));
28+
}
29+
desc_.reset(raw_desc);
30+
}
31+
32+
DFTI_DESCRIPTOR *get() const {
33+
if (desc_ == nullptr) {
34+
throw std::runtime_error("DFTI DESCRIPTOR has not been initialized");
35+
}
36+
return desc_.get();
37+
}
38+
39+
private:
40+
std::unique_ptr<DFTI_DESCRIPTOR, DftiDescriptorDeleter> desc_;
41+
};
42+
43+
44+
}} // at::native

0 commit comments

Comments
 (0)