saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-14 20:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[core][pruning][feature] cuSPARSELt build integration (#103700)

Browse source 
Summary:

This stack of PR's integrates cuSPARSELt into PyTorch.

This PR adds support for cuSPARSELt into the build process.
It adds in a new flag, USE_CUSPARSELT that defaults to false.

When USE_CUSPASRELT=1 is specified, the user can also specify
CUSPASRELT_ROOT, which defines the path to the library.

Compiling pytorch with cusparselt support can be done as follows:

``
USE_CUSPARSELT=1
CUSPARSELT_ROOT=/path/to/cusparselt

python setup.py develop
```

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/103700
Approved by: https://github.com/albanD
This commit is contained in:
Jesse Cai 2023-08-01 12:23:21 -07:00 committed by PyTorch MergeBot
parent d83b887f2a
commit f81f9093ec
11 changed files with 119 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
BUILD.bazel
View file

@ -289,6 +289,7 @@ header_template_rule(
include = "aten/src",
substitutions = {
"@AT_CUDNN_ENABLED@": "1",
"@AT_CUSPARSELT_ENABLED@": "0",
"@AT_ROCM_ENABLED@": "0",
"@AT_MAGMA_ENABLED@": "0",
"@NVCC_FLAGS_EXTRA@": "",

7
CMakeLists.txt
View file

@ -206,6 +206,9 @@ cmake_dependent_option(
cmake_dependent_option(
USE_STATIC_CUDNN "Use cuDNN static libraries" OFF
"USE_CUDNN" OFF)
cmake_dependent_option(
USE_CUSPARSELT "Use cuSPARSELt" ON
"USE_CUDA" OFF)
cmake_dependent_option(
BUILD_NVFUSER_BENCHMARK "Build C++ binaries for nvfuser benchmarks" OFF
"USE_CUDA" OFF)
@ -1152,6 +1155,10 @@ if(BUILD_SHARED_LIBS)
${PROJECT_SOURCE_DIR}/cmake/Modules/FindCUDAToolkit.cmake
DESTINATION share/cmake/Caffe2/
COMPONENT dev)
install(FILES
${PROJECT_SOURCE_DIR}/cmake/Modules/FindCUSPARSELT.cmake
DESTINATION share/cmake/Caffe2/
COMPONENT dev)
install(EXPORT Caffe2Targets DESTINATION share/cmake/Caffe2
FILE Caffe2Targets.cmake

6
aten/CMakeLists.txt
View file

@ -93,6 +93,12 @@ else()
set(AT_NNPACK_ENABLED 1)
endif()
if(NOT USE_CUSPARSELT)
set(AT_CUSPARSELT_ENABLED 0)
else()
set(AT_CUSPARSELT_ENABLED 1)
endif()
list(APPEND ATen_CPU_INCLUDE
${CMAKE_CURRENT_SOURCE_DIR}/src)
add_subdirectory(src/ATen)

1
aten/src/ATen/CMakeLists.txt
View file

@ -33,6 +33,7 @@ set_bool(AT_BLAS_USE_CBLAS_DOT BLAS_USE_CBLAS_DOT)
set_bool(AT_MAGMA_ENABLED USE_MAGMA)
set_bool(CAFFE2_STATIC_LINK_CUDA_INT CAFFE2_STATIC_LINK_CUDA)
set_bool(AT_CUDNN_ENABLED CAFFE2_USE_CUDNN)
set_bool(AT_CUSPARSELT_ENABLED CAFFE2_USE_CUSPARSELT)
configure_file(Config.h.in "${CMAKE_CURRENT_SOURCE_DIR}/Config.h")
# TODO: Do not generate CUDAConfig.h for ROCm BUILDS

2
aten/src/ATen/cuda/CUDAConfig.h.in
View file

@ -6,8 +6,8 @@
//
// NB: This header MUST NOT be included from other headers; it should
// only be included from C++ files.
#define AT_CUDNN_ENABLED() @AT_CUDNN_ENABLED@
#define AT_CUSPARSELT_ENABLED() @AT_CUSPARSELT_ENABLED@
#define AT_ROCM_ENABLED() @AT_ROCM_ENABLED@
#define AT_MAGMA_ENABLED() @AT_MAGMA_ENABLED@

5
caffe2/CMakeLists.txt
View file

@ -956,6 +956,11 @@ elseif(USE_CUDA)
set(CUDA_LINK_LIBRARIES_KEYWORD)
torch_compile_options(torch_cuda) # see cmake/public/utils.cmake
target_compile_definitions(torch_cuda PRIVATE USE_CUDA)
if(USE_CUSPARSELT)
target_link_libraries(torch_cuda PRIVATE torch::cusparselt)
target_compile_definitions(torch_cuda PRIVATE USE_CUSPARSELT)
endif()
if(USE_NCCL)
target_link_libraries(torch_cuda PRIVATE __caffe2_nccl)
target_compile_definitions(torch_cuda PRIVATE USE_NCCL)

8
cmake/Dependencies.cmake
View file

@ -38,6 +38,7 @@ if(USE_CUDA)
# public/*.cmake uses CAFFE2_USE_*
set(CAFFE2_USE_CUDA ${USE_CUDA})
set(CAFFE2_USE_CUDNN ${USE_CUDNN})
set(CAFFE2_USE_CUSPARSELT ${USE_CUSPARSELT})
set(CAFFE2_USE_NVRTC ${USE_NVRTC})
set(CAFFE2_USE_TENSORRT ${USE_TENSORRT})
include(${CMAKE_CURRENT_LIST_DIR}/public/cuda.cmake)
@ -57,6 +58,11 @@ if(USE_CUDA)
else()
caffe2_update_option(USE_CUDNN OFF)
endif()
if(CAFFE2_USE_CUSPARSELT)
list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS torch::cusparselt)
else()
caffe2_update_option(USE_CUSPARSELT OFF)
endif()
if(CAFFE2_USE_TENSORRT)
list(APPEND Caffe2_PUBLIC_CUDA_DEPENDENCY_LIBS caffe2::tensorrt)
else()
@ -68,10 +74,12 @@ if(USE_CUDA)
"-DUSE_CUDA=OFF.")
caffe2_update_option(USE_CUDA OFF)
caffe2_update_option(USE_CUDNN OFF)
caffe2_update_option(USE_CUSPARSELT OFF)
caffe2_update_option(USE_NVRTC OFF)
caffe2_update_option(USE_TENSORRT OFF)
set(CAFFE2_USE_CUDA OFF)
set(CAFFE2_USE_CUDNN OFF)
set(CAFFE2_USE_CUSPARSELT OFF)
set(CAFFE2_USE_NVRTC OFF)
set(CAFFE2_USE_TENSORRT OFF)
endif()

62
cmake/Modules/FindCUSPARSELT.cmake Normal file
View file

@ -0,0 +1,62 @@
# Find the CUSPARSELT library
#
# The following variables are optionally searched for defaults
# CUSPARSELT_ROOT: Base directory where CUSPARSELT is found
# CUSPARSELT_INCLUDE_DIR: Directory where CUSPARSELT header is searched for
# CUSPARSELT_LIBRARY: Directory where CUSPARSELT library is searched for
#
# The following are set after configuration is done:
# CUSPARSELT_FOUND
# CUSPARSELT_INCLUDE_PATH
# CUSPARSELT_LIBRARY_PATH
include(FindPackageHandleStandardArgs)
set(CUSPARSELT_ROOT $ENV{CUSPARSELT_ROOT_DIR} CACHE PATH "Folder containing NVIDIA cuSPARSELt")
if (DEFINED $ENV{CUSPARSELT_ROOT_DIR})
message(WARNING "CUSPARSELT_ROOT_DIR is deprecated. Please set CUSPARSELT_ROOT instead.")
endif()
list(APPEND CUSPARSELT_ROOT $ENV{CUSPARSELT_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR})
# Compatible layer for CMake <3.12. CUSPARSELT_ROOT will be accounted in for searching paths and libraries for CMake >=3.12.
list(APPEND CMAKE_PREFIX_PATH ${CUSPARSELT_ROOT})
set(CUSPARSELT_INCLUDE_DIR $ENV{CUSPARSELT_INCLUDE_DIR} CACHE PATH "Folder containing NVIDIA cuSPARSELt header files")
find_path(CUSPARSELT_INCLUDE_PATH cusparseLt.h
HINTS ${CUSPARSELT_INCLUDE_DIR}
PATH_SUFFIXES cuda/include cuda include)
set(CUSPARSELT_LIBRARY $ENV{CUSPARSELT_LIBRARY} CACHE PATH "Path to the cusparselt library file (e.g., libcusparseLt.so)")
find_library(CUSPARSELT_LIBRARY_PATH libcusparseLt.so
PATHS ${CUSPARSELT_LIBRARY}
PATH_SUFFIXES lib lib64 cuda/lib cuda/lib64 lib/x64)
find_package_handle_standard_args(CUSPARSELT DEFAULT_MSG CUSPARSELT_LIBRARY_PATH CUSPARSELT_INCLUDE_PATH)
if(CUSPARSELT_FOUND)
# Get cuSPARSELt version
file(READ ${CUSPARSELT_INCLUDE_PATH}/cusparseLt.h CUSPARSELT_HEADER_CONTENTS)
string(REGEX MATCH "define CUSPARSELT_VER_MAJOR * +([0-9]+)"
CUSPARSELT_VERSION_MAJOR "${CUSPARSELT_HEADER_CONTENTS}")
string(REGEX REPLACE "define CUSPARSELT_VER_MAJOR * +([0-9]+)" "\\1"
CUSPARSELT_VERSION_MAJOR "${CUSPARSELT_VERSION_MAJOR}")
string(REGEX MATCH "define CUSPARSELT_VER_MINOR * +([0-9]+)"
CUSPARSELT_VERSION_MINOR "${CUSPARSELT_HEADER_CONTENTS}")
string(REGEX REPLACE "define CUSPARSELT_VER_MINOR * +([0-9]+)" "\\1"
CUSPARSELT_VERSION_MINOR "${CUSPARSELT_VERSION_MINOR}")
string(REGEX MATCH "define CUSPARSELT_VER_PATCH * +([0-9]+)"
CUSPARSELT_VERSION_PATCH "${CUSPARSELT_HEADER_CONTENTS}")
string(REGEX REPLACE "define CUSPARSELT_VER_PATCH * +([0-9]+)" "\\1"
CUSPARSELT_VERSION_PATCH "${CUSPARSELT_VERSION_PATCH}")
# Assemble cuSPARSELt version. Use minor version since current major version is 0.
if(NOT CUSPARSELT_VERSION_MINOR)
set(CUSPARSELT_VERSION "?")
else()
set(CUSPARSELT_VERSION
"${CUSPARSELT_VERSION_MAJOR}.${CUSPARSELT_VERSION_MINOR}.${CUSPARSELT_VERSION_PATCH}")
endif()
endif()
mark_as_advanced(CUSPARSELT_ROOT CUSPARSELT_INCLUDE_DIR CUSPARSELT_LIBRARY CUSPARSELT_VERSION)

8
cmake/Summary.cmake
View file

@ -77,11 +77,15 @@ function(caffe2_print_configuration_summary)
message(STATUS " CUDA static link : ${CAFFE2_STATIC_LINK_CUDA}")
message(STATUS " USE_CUDNN : ${USE_CUDNN}")
message(STATUS " USE_EXPERIMENTAL_CUDNN_V8_API: ${USE_EXPERIMENTAL_CUDNN_V8_API}")
message(STATUS " USE_CUSPARSELT : ${USE_CUSPARSELT}")
message(STATUS " CUDA version : ${CUDA_VERSION}")
message(STATUS " USE_FLASH_ATTENTION : ${USE_FLASH_ATTENTION}")
if(${USE_CUDNN})
message(STATUS " cuDNN version : ${CUDNN_VERSION}")
endif()
if(${USE_CUSPARSELT})
message(STATUS " cuSPARSELt version : ${CUSPARSELT_VERSION}")
endif()
message(STATUS " CUDA root directory : ${CUDA_TOOLKIT_ROOT_DIR}")
message(STATUS " CUDA library : ${CUDA_cuda_driver_LIBRARY}")
message(STATUS " cudart library : ${CUDA_cudart_LIBRARY}")
@ -93,6 +97,10 @@ function(caffe2_print_configuration_summary)
get_target_property(__tmp torch::cudnn INTERFACE_LINK_LIBRARIES)
message(STATUS " cuDNN library : ${__tmp}")
endif()
if(${USE_CUSPARSELT})
get_target_property(__tmp torch::cusparselt INTERFACE_LINK_LIBRARIES)
message(STATUS " cuSPARSELt library : ${__tmp}")
endif()
message(STATUS " nvrtc : ${CUDA_nvrtc_LIBRARY}")
message(STATUS " CUDA include path : ${CUDA_INCLUDE_DIRS}")
message(STATUS " NVCC executable : ${CUDA_NVCC_EXECUTABLE}")

16
cmake/public/cuda.cmake
View file

@ -269,6 +269,22 @@ else()
message(STATUS "USE_CUDNN is set to 0. Compiling without cuDNN support")
endif()
if(CAFFE2_USE_CUSPARSELT)
find_package(CUSPARSELT)
if(NOT CUSPARSELT_FOUND)
message(WARNING
"Cannot find cuSPARSELt library. Turning the option off")
set(CAFFE2_USE_CUSPARSELT OFF)
else()
add_library(torch::cusparselt INTERFACE IMPORTED)
target_include_directories(torch::cusparselt INTERFACE ${CUSPARSELT_INCLUDE_PATH})
target_link_libraries(torch::cusparselt INTERFACE ${CUSPARSELT_LIBRARY_PATH})
endif()
else()
message(STATUS "USE_CUSPARSELT is set to 0. Compiling without cuSPARSELt support")
endif()
# curand
add_library(caffe2::curand INTERFACE IMPORTED)
if(CAFFE2_STATIC_LINK_CUDA AND NOT WIN32)

5
setup.py
View file

@ -29,10 +29,13 @@
# if used in conjunction with DEBUG or REL_WITH_DEB_INFO, will also
# build CUDA kernels with -lineinfo --source-in-ptx. Note that
# on CUDA 12 this may cause nvcc to OOM, so this is disabled by default.
#
# USE_CUDNN=0
# disables the cuDNN build
#
# USE_CUSPARSELT=0
# disables the cuSPARSELt build
#
# USE_FBGEMM=0
# disables the FBGEMM build
#