frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

AT_MKL_SEQUENTIAL and build changes (#70259)

Browse Source 
Summary:
Re-land of  https://github.com/pytorch/pytorch/pull/69419

Pull Request resolved: https://github.com/pytorch/pytorch/pull/70259

Test Plan: Imported from OSS

Reviewed By: malfet

Differential Revision: D33246757

Pulled By: ngimel

fbshipit-source-id: 738f8558d4cad6752be14108f9931ec3514f6682
This commit is contained in:
Peter Bell
2021-12-22 13:50:47 -08:00
committed by Facebook GitHub Bot
parent b37de0a4bb
commit c34aa715fa
4 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
BUILD.bazel
View File

@ -318,6 +318,7 @@ header_template_rule(
substitutions = {
"@AT_MKLDNN_ENABLED@": "1",
"@AT_MKL_ENABLED@": "1",
"@AT_MKL_SEQUENTIAL@": "0",
"@AT_FFTW_ENABLED@": "0",
"@AT_POCKETFFT_ENABLED@": "0",
"@AT_NNPACK_ENABLED@": "0",

1
aten/src/ATen/Config.h.in
View File

@ -8,6 +8,7 @@
#define AT_MKLDNN_ENABLED() @AT_MKLDNN_ENABLED@
#define AT_MKL_ENABLED() @AT_MKL_ENABLED@
#define AT_MKL_SEQUENTIAL() @AT_MKL_SEQUENTIAL@
#define AT_FFTW_ENABLED() @AT_FFTW_ENABLED@
#define AT_POCKETFFT_ENABLED() @AT_POCKETFFT_ENABLED@
#define AT_NNPACK_ENABLED() @AT_NNPACK_ENABLED@

2
aten/src/ATen/ParallelOpenMP.cpp
View File

@ -31,7 +31,7 @@ void init_num_threads() {
if (nthreads > 0) {
set_num_threads(nthreads);
} else {
#if defined(_OPENMP) && AT_MKL_ENABLED() && !defined(TH_BLAS_MKL_SEQ)
#if defined(_OPENMP) && AT_MKL_ENABLED() && !AT_MKL_SEQUENTIAL()
// If we are using MKL an OpenMP make sure the number of threads match.
// Otherwise, MKL and our OpenMP-enabled functions will keep changing the
// size of the OpenMP thread pool, resulting in worse performance (and memory

4
cmake/Dependencies.cmake
View File

@ -247,6 +247,7 @@ endif()
if(NOT INTERN_BUILD_MOBILE)
set(AT_MKL_ENABLED 0)
set(AT_MKL_SEQUENTIAL 0)
set(AT_MKL_MT 0)
set(USE_BLAS 1)
if(NOT (ATLAS_FOUND OR BLIS_FOUND OR GENERIC_BLAS_FOUND OR MKL_FOUND OR OpenBLAS_FOUND OR VECLIB_FOUND OR FlexiBLAS_FOUND))
@ -258,9 +259,8 @@ if(NOT INTERN_BUILD_MOBILE)
endif()
if(MKL_FOUND)
add_definitions(-DTH_BLAS_MKL)
if("${MKL_THREADING}" STREQUAL "SEQ")
add_definitions(-DTH_BLAS_MKL_SEQ=1)
set(AT_MKL_SEQUENTIAL 1)
endif()
if(MSVC AND MKL_LIBRARIES MATCHES ".*libiomp5md\\.lib.*")
add_definitions(-D_OPENMP_NOFORCE_MANIFEST)