Skip to content

Add a batched gemm Kokkos implementation - #240

Draft
masterleinad wants to merge 22 commits into
masterfrom
kokkos_batchedgemm
Draft

Add a batched gemm Kokkos implementation#240
masterleinad wants to merge 22 commits into
masterfrom
kokkos_batchedgemm

Conversation

@masterleinad

Copy link
Copy Markdown
Collaborator

A first effort for a Kokkos batched gemm implementation. The current profile on Frontier looks like

BEGIN KOKKOS PROFILING REPORT:
TOTAL TIME: 118.693730 seconds
TOP-DOWN TIME TREE:
<average time> <percent of total time> <percent time in Kokkos> <percent MPI imbalance> <remainder> <kernels per second> <number of calls> <name> [type]
===================
|-> 1.01e+02 sec 85.2% 100.0% 0.0% 0.0% 1.46e+01 246 BatchedGemmKokkos::matrixVector [region]
|   |-> 9.66e+01 sec 81.4% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass2 [for]
|   |-> 4.45e+00 sec 3.7% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass1 [for]
|-> 9.83e+00 sec 8.3% 5.2% 0.0% 94.8% 9.77e+00 12 BatchedGemmKokkos::setup [region]
|   |-> 4.25e-01 sec 0.4% 100.0% 0.0% ------ 12 "d_flatBbatch"="d_flatBbatch_mirror" (HIP->HOST) [copy]
|-> 1.37e+00 sec 1.2% 36.6% 0.0% 28.1% 2.24e+04 6216 PsimagLite::kokkos_gemm [region]
    |-> 5.34e-01 sec 0.4% 21.0% 0.0% 79.0% 1.07e+04 5921 KokkosBlas::gemm[TPL_ROCBLAS,double] [region]
    |-> 2.80e-01 sec 0.2% 100.0% 0.0% ------ 18648 ""="" (HIP->HOST) [copy]

BOTTOM-UP TIME TREE:
<average time> <percent of total time> <percent time in Kokkos> <percent MPI imbalance> <number of calls> <name> [type]
===================
|-> 9.66e+01 sec 81.4% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass2 [for]
|   |-> 9.66e+01 sec 81.4% 100.0% 0.0% 0.0% 0.00e+00 246 BatchedGemmKokkos::matrixVector [region]
|-> 9.32e+00 sec 7.8% 0.0% 0.0% 0.0% 0.00e+00 12 BatchedGemmKokkos::setup [region]
|-> 4.45e+00 sec 3.7% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass1 [for]
|   |-> 4.45e+00 sec 3.7% 100.0% 0.0% 0.0% 0.00e+00 246 BatchedGemmKokkos::matrixVector [region]
|-> 4.25e-01 sec 0.4% 100.0% 0.0% ------ 12 "d_flatBbatch"="d_flatBbatch_mirror" (HIP->HOST) [copy]
|   |-> 4.25e-01 sec 0.4% 100.0% 0.0% 0.0% 0.00e+00 12 BatchedGemmKokkos::setup [region]
|-> 4.21e-01 sec 0.4% 0.0% 0.0% 0.0% 0.00e+00 5921 KokkosBlas::gemm[TPL_ROCBLAS,double] [region]
|   |-> 4.21e-01 sec 0.4% 0.0% 0.0% 0.0% 0.00e+00 5921 PsimagLite::kokkos_gemm [region]
|-> 3.84e-01 sec 0.3% 0.0% 0.0% 0.0% 0.00e+00 6216 PsimagLite::kokkos_gemm [region]
|-> 2.80e-01 sec 0.2% 100.0% 0.0% ------ 18648 ""="" (HIP->HOST) [copy]
|   |-> 2.80e-01 sec 0.2% 100.0% 0.0% 0.0% 0.00e+00 18648 PsimagLite::kokkos_gemm [region]

KOKKOS HOST SPACE:
===================
MAX MEMORY ALLOCATED: 1911406.0 kB
ALLOCATIONS AT TIME OF HIGH WATER MARK:
  99.7% BatchedGemmKokkos::setup/d_flatBbatch_mirror
  0.3% BatchedGemmKokkos::setup/d_flatAbatch_mirror

KOKKOS CUDA SPACE:
===================
MAX MEMORY ALLOCATED: 0.0 kB
ALLOCATIONS AT TIME OF HIGH WATER MARK:

KOKKOS HIP SPACE:
===================
MAX MEMORY ALLOCATED: 1995172.2 kB
ALLOCATIONS AT TIME OF HIGH WATER MARK:
  95.5% BatchedGemmKokkos::setup/d_flatBbatch
  4.1% BatchedGemmKokkos::setup/d_flatBXbatch
  0.3% BatchedGemmKokkos::setup/d_flatAbatch

KOKKOS SYCL SPACE:
===================
MAX MEMORY ALLOCATED: 0.0 kB
ALLOCATIONS AT TIME OF HIGH WATER MARK:

KOKKOS OpenMPTarget SPACE:
===================
MAX MEMORY ALLOCATED: 0.0 kB
ALLOCATIONS AT TIME OF HIGH WATER MARK:

Host process high water mark memory consumption: 6929080 kB

END KOKKOS PROFILING REPORT.

@masterleinad masterleinad changed the title Port csr_matmul_post.cpp to Kokkos Add a batched gemm Kokkos implementation Jul 30, 2026
@masterleinad

masterleinad commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator Author

vector_size=64:

101: BEGIN KOKKOS PROFILING REPORT:
101: TOTAL TIME: 31.556133 seconds
101: TOP-DOWN TIME TREE:
101: <average time> <percent of total time> <percent time in Kokkos> <percent MPI imbalance> <remainder> <kernels per second> <number of calls> <name> [type]
101: ===================
101: |-> 2.23e+01 sec 70.6% 100.0% 0.0% 0.0% 6.63e+01 246 BatchedGemmKokkos::matrixVector [region]
101: |   |-> 1.69e+01 sec 53.7% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass2 [for]
101: |   |-> 5.29e+00 sec 16.8% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass1 [for]
101: |-> 2.52e+00 sec 8.0% 11.9% 0.0% 88.1% 3.80e+01 12 BatchedGemmKokkos::setup [region]
101: |   |-> 2.49e-01 sec 0.8% 100.0% 0.0% ------ 12 "d_flatBbatch"="d_flatBbatch_mirror" (HIP->HOST) [copy]
101: |   |-> 5.05e-02 sec 0.2% 100.0% 0.0% ------ 12 "d_flatAbatch"="d_flatAbatch_mirror" (HIP->HOST) [copy]
101: |-> 1.28e+00 sec 4.1% 36.8% 0.0% 27.9% 2.38e+04 6216 PsimagLite::kokkos_gemm [region]
101:     |-> 4.95e-01 sec 1.6% 20.9% 0.0% 79.1% 1.15e+04 5921 KokkosBlas::gemm[TPL_ROCBLAS,double] [region]
101:     |   |-> 5.28e-02 sec 0.2% 100.0% 0.0% ------ 2845 Perform Dot Product Based GEMM [for]
101:     |   |-> 5.05e-02 sec 0.2% 100.0% 0.0% ------ 2845 Initialize C for Dot Product Based GEMM [for]
101:     |-> 2.73e-01 sec 0.9% 100.0% 0.0% ------ 18648 ""="" (HIP->HOST) [copy]
101:     |-> 9.54e-02 sec 0.3% 100.0% 0.0% ------ 6216 ""="" (HOST->HIP) [copy]
101:     |-> 5.94e-02 sec 0.2% 0.0% 0.0% 100.0% 0.00e+00 295 KokkosBlas::gemv[TPL_ROCBLAS,double] [region]

vector_size=32:

101: BEGIN KOKKOS PROFILING REPORT:
101: TOTAL TIME: 31.294461 seconds
101: TOP-DOWN TIME TREE:
101: <average time> <percent of total time> <percent time in Kokkos> <percent MPI imbalance> <remainder> <kernels per second> <number of calls> <name> [type]
101: ===================
101: |-> 2.19e+01 sec 69.9% 100.0% 0.0% 0.0% 6.75e+01 246 BatchedGemmKokkos::matrixVector [region]
101: |   |-> 1.77e+01 sec 56.5% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass2 [for]
101: |   |-> 4.15e+00 sec 13.3% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass1 [for]
101: |-> 2.55e+00 sec 8.2% 11.8% 0.0% 88.2% 3.76e+01 12 BatchedGemmKokkos::setup [region]
101: |   |-> 2.49e-01 sec 0.8% 100.0% 0.0% ------ 12 "d_flatBbatch"="d_flatBbatch_mirror" (HIP->HOST) [copy]
101: |   |-> 5.06e-02 sec 0.2% 100.0% 0.0% ------ 12 "d_flatAbatch"="d_flatAbatch_mirror" (HIP->HOST) [copy]
101: |-> 1.39e+00 sec 4.4% 42.7% 0.0% 25.5% 2.20e+04 6216 PsimagLite::kokkos_gemm [region]
101:     |-> 6.02e-01 sec 1.9% 36.9% 0.0% 63.1% 9.45e+03 5921 KokkosBlas::gemm[TPL_ROCBLAS,double] [region]
101:     |   |-> 1.69e-01 sec 0.5% 100.0% 0.0% ------ 2845 Initialize C for Dot Product Based GEMM [for]
101:     |   |-> 5.30e-02 sec 0.2% 100.0% 0.0% ------ 2845 Perform Dot Product Based GEMM [for]
101:     |-> 2.78e-01 sec 0.9% 100.0% 0.0% ------ 18648 ""="" (HIP->HOST) [copy]
101:     |-> 9.48e-02 sec 0.3% 100.0% 0.0% ------ 6216 ""="" (HOST->HIP) [copy]
101:     |-> 6.06e-02 sec 0.2% 0.0% 0.0% 100.0% 0.00e+00 295 KokkosBlas::gemv[TPL_ROCBLAS,double] [region]

vector_size=16

101: BEGIN KOKKOS PROFILING REPORT:
101: TOTAL TIME: 30.381591 seconds
101: TOP-DOWN TIME TREE:
101: <average time> <percent of total time> <percent time in Kokkos> <percent MPI imbalance> <remainder> <kernels per second> <number of calls> <name> [type]
101: ===================
101: |-> 2.12e+01 sec 69.7% 100.0% 0.0% 0.0% 6.97e+01 246 BatchedGemmKokkos::matrixVector [region]
101: |   |-> 1.79e+01 sec 58.9% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass2 [for]
101: |   |-> 3.22e+00 sec 10.6% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass1 [for]
101: |-> 2.54e+00 sec 8.3% 11.8% 0.0% 88.2% 3.79e+01 12 BatchedGemmKokkos::setup [region]
101: |   |-> 2.48e-01 sec 0.8% 100.0% 0.0% ------ 12 "d_flatBbatch"="d_flatBbatch_mirror" (HIP->HOST) [copy]
101: |   |-> 5.03e-02 sec 0.2% 100.0% 0.0% ------ 12 "d_flatAbatch"="d_flatAbatch_mirror" (HIP->HOST) [copy]
101: |-> 1.30e+00 sec 4.3% 36.7% 0.0% 27.8% 2.36e+04 6216 PsimagLite::kokkos_gemm [region]
101:     |-> 5.02e-01 sec 1.7% 20.6% 0.0% 79.4% 1.13e+04 5921 KokkosBlas::gemm[TPL_ROCBLAS,double] [region]
101:     |   |-> 5.25e-02 sec 0.2% 100.0% 0.0% ------ 2845 Perform Dot Product Based GEMM [for]
101:     |   |-> 5.11e-02 sec 0.2% 100.0% 0.0% ------ 2845 Initialize C for Dot Product Based GEMM [for]
101:     |-> 2.77e-01 sec 0.9% 100.0% 0.0% ------ 18648 ""="" (HIP->HOST) [copy]
101:     |-> 9.56e-02 sec 0.3% 100.0% 0.0% ------ 6216 ""="" (HOST->HIP) [copy]
101:     |-> 6.14e-02 sec 0.2% 0.0% 0.0% 100.0% 0.00e+00 295 KokkosBlas::gemv[TPL_ROCBLAS,double] [region]

vector_size=8

101: BEGIN KOKKOS PROFILING REPORT:
101: TOTAL TIME: 34.869436 seconds
101: TOP-DOWN TIME TREE:
101: <average time> <percent of total time> <percent time in Kokkos> <percent MPI imbalance> <remainder> <kernels per second> <number of calls> <name> [type]
101: ===================
101: |-> 2.53e+01 sec 72.7% 100.0% 0.0% 0.0% 5.82e+01 246 BatchedGemmKokkos::matrixVector [region]
101: |   |-> 2.24e+01 sec 64.1% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass2 [for]
101: |   |-> 2.95e+00 sec 8.5% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass1 [for]
101: |-> 2.53e+00 sec 7.3% 12.1% 0.0% 87.9% 3.79e+01 12 BatchedGemmKokkos::setup [region]
101: |   |-> 2.54e-01 sec 0.7% 100.0% 0.0% ------ 12 "d_flatBbatch"="d_flatBbatch_mirror" (HIP->HOST) [copy]
101: |   |-> 5.05e-02 sec 0.1% 100.0% 0.0% ------ 12 "d_flatAbatch"="d_flatAbatch_mirror" (HIP->HOST) [copy]
101: |-> 1.28e+00 sec 3.7% 37.5% 0.0% 28.1% 2.38e+04 6216 PsimagLite::kokkos_gemm [region]
101:     |-> 4.87e-01 sec 1.4% 21.6% 0.0% 78.4% 1.17e+04 5921 KokkosBlas::gemm[TPL_ROCBLAS,double] [region]
101:     |   |-> 5.30e-02 sec 0.2% 100.0% 0.0% ------ 2845 Perform Dot Product Based GEMM [for]
101:     |   |-> 5.20e-02 sec 0.1% 100.0% 0.0% ------ 2845 Initialize C for Dot Product Based GEMM [for]
101:     |-> 2.79e-01 sec 0.8% 100.0% 0.0% ------ 18648 ""="" (HIP->HOST) [copy]
101:     |-> 9.71e-02 sec 0.3% 100.0% 0.0% ------ 6216 ""="" (HOST->HIP) [copy]
101:     |-> 5.99e-02 sec 0.2% 0.0% 0.0% 100.0% 0.00e+00 295 KokkosBlas::gemv[TPL_ROCBLAS,double] [region]

vector_size=4

101: BEGIN KOKKOS PROFILING REPORT:
101: TOTAL TIME: 34.999324 seconds
101: TOP-DOWN TIME TREE:
101: <average time> <percent of total time> <percent time in Kokkos> <percent MPI imbalance> <remainder> <kernels per second> <number of calls> <name> [type]
101: ===================
101: |-> 2.57e+01 sec 73.3% 100.0% 0.0% 0.0% 5.75e+01 246 BatchedGemmKokkos::matrixVector [region]
101: |   |-> 2.27e+01 sec 64.8% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass2 [for]
101: |   |-> 2.94e+00 sec 8.4% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass1 [for]
101: |-> 2.57e+00 sec 7.4% 11.7% 0.0% 88.3% 3.73e+01 12 BatchedGemmKokkos::setup [region]
101: |   |-> 2.49e-01 sec 0.7% 100.0% 0.0% ------ 12 "d_flatBbatch"="d_flatBbatch_mirror" (HIP->HOST) [copy]
101: |   |-> 5.04e-02 sec 0.1% 100.0% 0.0% ------ 12 "d_flatAbatch"="d_flatAbatch_mirror" (HIP->HOST) [copy]
101: |-> 1.28e+00 sec 3.6% 37.4% 0.0% 28.0% 2.39e+04 6216 PsimagLite::kokkos_gemm [region]
101:     |-> 4.84e-01 sec 1.4% 21.3% 0.0% 78.7% 1.18e+04 5921 KokkosBlas::gemm[TPL_ROCBLAS,double] [region]
101:     |   |-> 5.26e-02 sec 0.2% 100.0% 0.0% ------ 2845 Perform Dot Product Based GEMM [for]
101:     |   |-> 5.07e-02 sec 0.1% 100.0% 0.0% ------ 2845 Initialize C for Dot Product Based GEMM [for]
101:     |-> 2.78e-01 sec 0.8% 100.0% 0.0% ------ 18648 ""="" (HIP->HOST) [copy]
101:     |-> 9.61e-02 sec 0.3% 100.0% 0.0% ------ 6216 ""="" (HOST->HIP) [copy]
101:     |-> 6.11e-02 sec 0.2% 0.0% 0.0% 100.0% 0.00e+00 295 KokkosBlas::gemv[TPL_ROCBLAS,double] [region]

vector_size=2

101: BEGIN KOKKOS PROFILING REPORT:
101: TOTAL TIME: 54.220531 seconds
101: TOP-DOWN TIME TREE:
101: <average time> <percent of total time> <percent time in Kokkos> <percent MPI imbalance> <remainder> <kernels per second> <number of calls> <name> [type]
101: ===================
101: |-> 4.49e+01 sec 82.8% 100.0% 0.0% 0.0% 3.29e+01 246 BatchedGemmKokkos::matrixVector [region]
101: |   |-> 4.20e+01 sec 77.5% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass2 [for]
101: |   |-> 2.82e+00 sec 5.2% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass1 [for]
101: |-> 2.57e+00 sec 4.7% 11.7% 0.0% 88.3% 3.73e+01 12 BatchedGemmKokkos::setup [region]
101: |   |-> 2.49e-01 sec 0.5% 100.0% 0.0% ------ 12 "d_flatBbatch"="d_flatBbatch_mirror" (HIP->HOST) [copy]
101: |-> 1.28e+00 sec 2.4% 37.5% 0.0% 27.7% 2.38e+04 6216 PsimagLite::kokkos_gemm [region]
101:     |-> 4.91e-01 sec 0.9% 21.4% 0.0% 78.6% 1.16e+04 5921 KokkosBlas::gemm[TPL_ROCBLAS,double] [region]
101:     |-> 2.80e-01 sec 0.5% 100.0% 0.0% ------ 18648 ""="" (HIP->HOST) [copy]
101:     |-> 9.49e-02 sec 0.2% 100.0% 0.0% ------ 6216 ""="" (HOST->HIP) [copy]
101:     |-> 6.00e-02 sec 0.1% 0.0% 0.0% 100.0% 0.00e+00 295 KokkosBlas::gemv[TPL_ROCBLAS,double] [region]

With vector_size=2/64

101: BEGIN KOKKOS PROFILING REPORT:
101: TOTAL TIME: 29.206763 seconds
101: TOP-DOWN TIME TREE:
101: <average time> <percent of total time> <percent time in Kokkos> <percent MPI imbalance> <remainder> <kernels per second> <number of calls> <name> [type]
101: ===================
101: |-> 1.96e+01 sec 67.3% 100.0% 0.0% 0.0% 7.51e+01 246 BatchedGemmKokkos::matrixVector [region]
101: |   |-> 1.67e+01 sec 57.1% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass2 [for]
101: |   |-> 2.92e+00 sec 10.0% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass1 [for]
101: |-> 2.63e+00 sec 9.0% 11.4% 0.0% 88.6% 3.66e+01 12 BatchedGemmKokkos::setup [region]
101: |   |-> 2.48e-01 sec 0.9% 100.0% 0.0% ------ 12 "d_flatBbatch"="d_flatBbatch_mirror" (HIP->HOST) [copy]
101: |   |-> 5.09e-02 sec 0.2% 100.0% 0.0% ------ 12 "d_flatAbatch"="d_flatAbatch_mirror" (HIP->HOST) [copy]
101: |-> 1.26e+00 sec 4.3% 37.1% 0.0% 28.4% 2.43e+04 6216 PsimagLite::kokkos_gemm [region]
101:     |-> 4.77e-01 sec 1.6% 21.6% 0.0% 78.4% 1.19e+04 5921 KokkosBlas::gemm[TPL_ROCBLAS,double] [region]
101:     |   |-> 5.24e-02 sec 0.2% 100.0% 0.0% ------ 2845 Perform Dot Product Based GEMM [for]
101:     |   |-> 5.05e-02 sec 0.2% 100.0% 0.0% ------ 2845 Initialize C for Dot Product Based GEMM [for]
101:     |-> 2.70e-01 sec 0.9% 100.0% 0.0% ------ 18648 ""="" (HIP->HOST) [copy]
101:     |-> 9.34e-02 sec 0.3% 100.0% 0.0% ------ 6216 ""="" (HOST->HIP) [copy]
101:     |-> 5.94e-02 sec 0.2% 0.0% 0.0% 100.0% 0.00e+00 295 KokkosBlas::gemv[TPL_ROCBLAS,double] [region]

@masterleinad

Copy link
Copy Markdown
Collaborator Author

Just switching to KokkosBatched::Algo::Gemm::Blocked gives me

101: |-> 6.00e+00 sec 23.3% 100.0% 0.0% 0.0% 2.46e+02 246 BatchedGemmKokkos::matrixVector [region]
101: |   |-> 4.66e+00 sec 18.1% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass2 [for]
101: |   |-> 1.30e+00 sec 5.0% 100.0% 0.0% ------ 246 BatchedGemmKokkos_Pass1 [for]

for a total runtime ~16s.


using DevArgsView = Kokkos::View<GemmArgs*, DevMemSpace>;

static const int ialign_ = 32;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where does 32 come from? Is it related to the warpsize?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants