kokkos-kernels
kokkos-kernels copied to clipboard
kokkos
Adding more support for SpGEMM via rocsparse
Changes:
-
Adding Jacobi SpGEMM via rocsparse
-
Enabling rocsparse testing for SpGEMM
-
Enabling complex datatypes with SpGEMM
-
Adding option for non-int size types (e.g. size_t and int64_t) for SpGEMM
@lucbv are you up for reviewing this? It adds some feature support for SpGEMM with rocsparse.
Status Flag 'Pre-Test Inspection' - - This Pull Request Requires Inspection... The code must be inspected by a member of the Team before Testing/Merging NO INSPECTION HAS BEEN PERFORMED ON THIS PULL REQUEST! - This PR must be inspected by setting label 'AT: PRE-TEST INSPECTED'.
Status Flag 'Pre-Test Inspection' - - This Pull Request Requires Inspection... The code must be inspected by a member of the Team before Testing/Merging NO INSPECTION HAS BEEN PERFORMED ON THIS PULL REQUEST! - This PR must be inspected by setting label 'AT: PRE-TEST INSPECTED'.
@seanofthemillers @brian-kelley Is it fair to say that this PR has been superseded by recent work that was merged in Kokkos Kernels on integrating TPLs and redesigning our interfaces for SPGEMM? In which case I would suggest just closing this.
@lucbv Hmm... the testing and complex support are superseded now, but I haven't supported any TPLs for SpGEMM-Jacobi. I also haven't converted non-int ordinals/offsets to int to call TPLs.
- For SpGEMM-Jacobi, ideally any TPLs would be supported through the unification layer (KokkosSparse_spgemm_jacobi_spec.hpp does exist now), instead of the old SpGEMM way of having rocSparse be an "algorithm" choice in the handle. In fact, the algorithm enums corresponding to TPLs are now deprecated
- On converting ordinals/offsets to int, I think it's a good idea in the sense that the time to convert is probably a lot smaller than the savings from using TPLs instead of native. But reusing the symbolic phase would require keeping a second copy of the product matrix around
@lucbv @brian-kelley Sorry, I've been a bit sidetracked. I can rewrite this to work with the new backend for spgemm-jacobi, and remove my changes for testing/complex. I'll can try to get back on this next week.
I am a little worried about memory consumption for duplicating some of these arrays (index arrays and the D^{-1} * A allocation). Hiding those allocations inside a handle may not always be advisable. With the new spgemm-jacobi backend, is there a way to chose not to use rocsparse? Outside of recompiling the code with rocsparse disabled.
@seanofthemillers When TPL support is done through the unification layer, the decision to call a TPL is done at compile-time and normally there isn't a way to opt out of the TPL at runtime. This design assumes that the TPL will always outperform the implementation that we write. In a few places we have alternate runtime paths for cases where TPLs perform poorly (for example, KokkosBlas::gemm where A is short-and-fat, and B is tall-and-skinny).
I think in this case, the only decision to make is, should non-int index types be supported (requiring a copy of int32-typed indices in the handle)? I don't have a strong opinion either way.
Status Flag 'Pre-Test Inspection' - - This Pull Request Requires Inspection... The code must be inspected by a member of the Team before Testing/Merging NO INSPECTION HAS BEEN PERFORMED ON THIS PULL REQUEST! - This PR must be inspected by setting label 'AT: PRE-TEST INSPECTED'.