ConjTrans support for batched TeamVector gemm by yasahi-hpc · Pull Request #2628 · kokkos/kokkos-kernels

yasahi-hpc · 2025-05-06T19:12:24Z

This PR aims at improving the implementation and testing of TeamVector Gemm. (see also #2580)

Moving implementation details into Impl namespace
Add all the specializations of the combinations of Non-Trans/Trans/ConjTrans
Add a check function of input to force A, B, C are rank 0, 1, 2 Kokkos::View
Covering all the unit-tests for all the combinations of Non-Trans/Trans/ConjTrans
Merge the unit-test for Team and TeamVector which are parameterized with Mode
Deprecated warnings for the older implementation details under KokkosBatched namespace

Signed-off-by: Yuuichi Asahi <y.asahi@nr.titech.ac.jp>

Yuuichi Asahi added 2 commits May 7, 2025 04:09

ConjTrans support TeamVector Gemm

97294dc

Signed-off-by: Yuuichi Asahi <y.asahi@nr.titech.ac.jp>

Update TeamVector Gemm usage in spmv bsrmatrix

c0a1761

Signed-off-by: Yuuichi Asahi <y.asahi@nr.titech.ac.jp>

yasahi-hpc self-assigned this May 6, 2025

Yuuichi Asahi added 2 commits May 7, 2025 04:18

fix unused variables

837ca09

Signed-off-by: Yuuichi Asahi <y.asahi@nr.titech.ac.jp>

Fix format

c029b9d

Signed-off-by: Yuuichi Asahi <y.asahi@nr.titech.ac.jp>

yasahi-hpc added CI: skip-docs Do not run the documentation checks for this pull request enhancement labels May 6, 2025

yasahi-hpc changed the title ~~Improve batched team vector gemm~~ ConjTrans support for batched TeamVector gemm May 6, 2025

Fix: call TeamVectorGemm in tests

70d4efb

Signed-off-by: Yuuichi Asahi <y.asahi@nr.titech.ac.jp>

yasahi-hpc added the AT2-CI-APPROVAL Approve CI to run at SNL label May 6, 2025

yasahi-hpc requested a review from lucbv May 6, 2025 21:04

yasahi-hpc mentioned this pull request Jun 4, 2025

Fix the behavior of extent(r) for r > rank in TeamGEMM #2667

Open

Provide feedback