Add USE_MAGMA build flag (#55994)

Summary:
Many model pipelines/workflows don't use MAGMA even though it is included in the build by default. Leaving MAGMA kernels out of the build can save 60+MB of GPU memory when loading `libtorch_cuda.so` (tested on V100, current upstream master).

A current sharp corner of this flag is that toggling it when rebuilding requires `torch/include/THC/THCGeneral.h` to be *manually* deleted by the user, as even running `make clean` or `setup.py` with `--cmake` does not properly regenerate it with the appropriate substitution for `#cmakedefine USE_MAGMA`. Is there a way to force the regeneration of the header during a rebuild?

CC malfet ptrblck

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

Reviewed By: mruberry

Differential Revision: D27766287

Pulled By: malfet

fbshipit-source-id: 93deca57befa0febb9c5b7875ecf0015c547d421
This commit is contained in:
Eddie Yan 2021-04-15 00:41:51 -07:00 committed by Facebook GitHub Bot
parent 1995640d86
commit 81f181567a
2 changed files with 10 additions and 1 deletions

View File

@ -204,6 +204,7 @@ option(USE_GLOG "Use GLOG" OFF)
option(USE_LEVELDB "Use LEVELDB" OFF)
option(USE_LITE_PROTO "Use lite protobuf instead of full." OFF)
option(USE_LMDB "Use LMDB" OFF)
option(USE_MAGMA "Use MAGMA" ON)
option(USE_METAL "Use Metal for Caffe2 iOS build" ON)
option(USE_PYTORCH_METAL "Use Metal for PyTorch iOS build" OFF)
option(USE_NATIVE_ARCH "Use -march=native" OFF)

View File

@ -1578,7 +1578,9 @@ if(NOT INTERN_BUILD_MOBILE)
set(CUDA_ATTACH_VS_BUILD_RULE_TO_CUDA_FILE OFF)
find_package(MAGMA)
if(USE_MAGMA)
find_package(MAGMA)
endif()
if((USE_CUDA OR USE_ROCM) AND MAGMA_FOUND)
include_directories(SYSTEM ${MAGMA_INCLUDE_DIR})
if(USE_CUDA)
@ -1602,8 +1604,14 @@ if(NOT INTERN_BUILD_MOBILE)
message(STATUS "MAGMA INCLUDE DIRECTORIES: ${MAGMA_INCLUDE_DIR}")
message(STATUS "MAGMA LIBRARIES: ${MAGMA_LIBRARIES}")
message(STATUS "MAGMA V2 check: ${MAGMA_V2}")
elseif(USE_MAGMA)
message(WARNING
"Not compiling with MAGMA. Suppress this warning with "
"-DUSE_MAGMA=OFF.")
caffe2_update_option(USE_MAGMA OFF)
else()
message(STATUS "MAGMA not found. Compiling without MAGMA support")
caffe2_update_option(USE_MAGMA OFF)
endif()
# ARM specific flags