mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

[sparse] Update the guard of cusparse SpMM and SpMv algorithms to cusparse version 11.7.1 onwards.

Browse Source 
PiperOrigin-RevId: 486051658
This commit is contained in:
Tianjian Lu 2022-11-03 21:39:10 -07:00 committed by jax authors
parent 974134fe31
commit 46368e4e73
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
jaxlib/gpu/vendor.h
View File

@ -224,9 +224,9 @@ typedef cusparseDnVecDescr_t gpusparseDnVecDescr_t;
// Use CUSPARSE_SPMV_COO_ALG2 and CUSPARSE_SPMV_CSR_ALG2 for SPMV and
// use CUSPARSE_SPMM_COO_ALG2 and CUSPARSE_SPMM_CSR_ALG3 for SPMM, which
// provide deterministic (bit-wise) results for each run. These indexing modes
// are available in CUSPARSE 11.4 and newer (which was released as part of
// CUDA 11.2.1)
#if CUSPARSE_VERSION >= 11400
// are fully supported (both row- and column-major inputs) in CUSPARSE 11.7.1
// and newer (which was released as part of CUDA 11.8)
#if CUSPARSE_VERSION > 11700
#define GPUSPARSE_SPMV_COO_ALG CUSPARSE_SPMV_COO_ALG2
#define GPUSPARSE_SPMV_CSR_ALG CUSPARSE_SPMV_CSR_ALG2
#define GPUSPARSE_SPMM_COO_ALG CUSPARSE_SPMM_COO_ALG2