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

[ATen][CPU][Sparse] Use Third-Party Eigen for sparse add and addmm (#155357)

Browse Source 
This pull request adds the following ops for sparse matrices using Eigen library:
```python
    add(a_csr, b_csr)
    add(a_csc, b_csc)

    addmm(c_csr, a_csr, b_csr)
    addmm(c_csr, a_csr, b_csc)
    addmm(c_csr, a_csc, b_csc)
    addmm(c_csr, a_csc, b_csr)

    addmm(c_csc, a_csr, b_csr)
    addmm(c_csc, a_csr, b_csc)
    addmm(c_csc, a_csc, b_csc)
    addmm(c_csc, a_csc, b_csr)
```

Currently, the operations for sparse matrices on CPU are available through MKL only. The non-existence of MKL on `aarch64` causes the unavailability of these ops on any machines with ARM based CPUs, including Apple Silicon, AWS Graviton and NVIDIA Grace. This PR addresses this issue by using Eigen as a backend for the above ops.

This is a re-factored version of my previous PR #101814. The main difference with the old one, this does not enable Eigen by default.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/155357
Approved by: https://github.com/pearu, https://github.com/eqy

Co-authored-by: Eli Uriegas <eliuriegas@meta.com>
This commit is contained in:
Aidyn-A
2025-08-23 19:03:55 +00:00
committed by PyTorch MergeBot
parent 4acdbb8311
commit 3e5b021f21
14 changed files with 426 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
cmake/Dependencies.cmake
View File

@ -153,6 +153,7 @@ set(AT_MKLDNN_ACL_ENABLED 0)
set(AT_MKLDNN_ENABLED 0)
set(AT_MKL_ENABLED 0)
set(AT_KLEIDIAI_ENABLED 0)
set(AT_USE_EIGEN_SPARSE 0)
# setting default preferred BLAS options if not already present.
if(NOT INTERN_BUILD_MOBILE)
set(BLAS "MKL" CACHE STRING "Selected BLAS library")
@ -262,6 +263,15 @@ if(BLAS_LIBRARIES AND BLAS_CHECK_F2C)
include(cmake/BLAS_ABI.cmake)
endif()
if(USE_EIGEN_SPARSE AND BLAS_INFO STREQUAL "mkl")
message(WARNING "Disabling USE_EIGEN_SPARSE because MKL is enabled")
set(USE_EIGEN_SPARSE OFF)
endif()
if(USE_EIGEN_SPARSE)
set(AT_USE_EIGEN_SPARSE 1)
endif()
if(NOT INTERN_BUILD_MOBILE)
set(AT_MKL_SEQUENTIAL 0)
set(USE_BLAS 1)

1
cmake/Summary.cmake
View File

@ -135,6 +135,7 @@ function(caffe2_print_configuration_summary)
endif()
message(STATUS " BUILD_NVFUSER : ${BUILD_NVFUSER}")
message(STATUS " USE_EIGEN_FOR_BLAS : ${CAFFE2_USE_EIGEN_FOR_BLAS}")
message(STATUS " USE_EIGEN_FOR_SPARSE : ${USE_EIGEN_SPARSE}")
message(STATUS " USE_FBGEMM : ${USE_FBGEMM}")
message(STATUS " USE_KINETO : ${USE_KINETO}")
message(STATUS " USE_GFLAGS : ${USE_GFLAGS}")