OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[CMake] Add NVPL BLAS/LAPACK option (#125268)

Browse Source 
This PR add a [NVPL](https://docs.nvidia.com/nvpl/introduction.html) BLAS/LAPACK option to CMake for `aarch64` (ARM) machines.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/125268
Approved by: https://github.com/albanD
This commit is contained in:
aaitzhan 2024-05-01 17:26:28 +00:00 committed by PyTorch MergeBot
parent 39eb5d4fa4
commit e3627d05e7
2 changed files with 14 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
cmake/Dependencies.cmake
View File

@ -237,6 +237,12 @@ elseif(BLAS STREQUAL "MKL")
set(CAFFE2_USE_EIGEN_FOR_BLAS ON)
set(CAFFE2_USE_MKL OFF)
endif()
elseif(BLAS STREQUAL "NVPL")
find_package(NVPL_BLAS REQUIRED)
list(APPEND Caffe2_DEPENDENCY_LIBS nvpl::blas_lp64_omp)
set(BLAS_INFO "nvpl")
set(BLAS_FOUND 1)
set(BLAS_USE_CBLAS_DOT TRUE)
elseif(BLAS STREQUAL "vecLib")
find_package(vecLib REQUIRED)
include_directories(SYSTEM ${vecLib_INCLUDE_DIR})
@ -269,7 +275,7 @@ if(NOT INTERN_BUILD_MOBILE)
set(AT_MKL_ENABLED 0)
set(AT_MKL_SEQUENTIAL 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))
if(NOT (ATLAS_FOUND OR BLIS_FOUND OR GENERIC_BLAS_FOUND OR MKL_FOUND OR OpenBLAS_FOUND OR VECLIB_FOUND OR FlexiBLAS_FOUND OR NVPL_BLAS_FOUND))
message(WARNING "Preferred BLAS (" ${BLAS} ") cannot be found, now searching for a general BLAS library")
find_package(BLAS)
if(NOT BLAS_FOUND)

7
cmake/Modules/FindLAPACK.cmake
View File

@ -95,6 +95,13 @@ if(BLAS_FOUND)
SET(LAPACK_INFO "mkl")
ENDIF()
# NVPL
IF((NOT LAPACK_INFO) AND (BLAS_INFO STREQUAL "nvpl"))
FIND_PACKAGE(NVPL_LAPACK REQUIRED)
SET(LAPACK_LIBRARIES nvpl::lapack_lp64_omp)
SET(LAPACK_INFO "nvpl")
ENDIF()
# Accelerate
IF((NOT LAPACK_INFO) AND (BLAS_INFO STREQUAL "accelerate"))
SET(CMAKE_REQUIRED_LIBRARIES ${BLAS_LIBRARIES})