From a0ba7f8c757a26b8fec01801a56b9d6b715d2968 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Wed, 15 Jul 2026 14:47:58 -0400 Subject: [PATCH 01/41] Port csr_matmul_post.cpp to Kokkos --- PsimagLite/src/CMakeLists.txt | 2 +- PsimagLite/src/PsimagLite/BLAS.h | 2 +- .../{kokkos_gemm.cpp => KokkosGemm.cpp} | 20 +-- .../{kokkos_gemm.h => KokkosGemm.h} | 0 PsimagLite/src/PsimagLite/KokkosType.h | 23 +++ dmrg/KronUtil/csr_matmul_post.cpp | 155 +++++++++++++----- 6 files changed, 139 insertions(+), 63 deletions(-) rename PsimagLite/src/PsimagLite/{kokkos_gemm.cpp => KokkosGemm.cpp} (91%) rename PsimagLite/src/PsimagLite/{kokkos_gemm.h => KokkosGemm.h} (100%) create mode 100644 PsimagLite/src/PsimagLite/KokkosType.h diff --git a/PsimagLite/src/CMakeLists.txt b/PsimagLite/src/CMakeLists.txt index 8b12b0ed8..87223cd45 100644 --- a/PsimagLite/src/CMakeLists.txt +++ b/PsimagLite/src/CMakeLists.txt @@ -37,7 +37,7 @@ FetchContent_MakeAvailable(KokkosKernels) list(POP_BACK CMAKE_MESSAGE_INDENT) target_link_libraries(psimaglite PUBLIC Kokkos::kokkoskernels Kokkos::kokkos) -target_sources(psimaglite PRIVATE PsimagLite/kokkos_gemm.cpp) +target_sources(psimaglite PRIVATE PsimagLite/KokkosGemm.cpp) target_include_directories(psimaglite PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/Io ${CMAKE_CURRENT_SOURCE_DIR}/Ainur) diff --git a/PsimagLite/src/PsimagLite/BLAS.h b/PsimagLite/src/PsimagLite/BLAS.h index fa448d0f1..ea2d1b6ce 100644 --- a/PsimagLite/src/PsimagLite/BLAS.h +++ b/PsimagLite/src/PsimagLite/BLAS.h @@ -10,7 +10,7 @@ #define PSIMAG_BLAS #include "AllocatorCpu.h" -#include +#include #include diff --git a/PsimagLite/src/PsimagLite/kokkos_gemm.cpp b/PsimagLite/src/PsimagLite/KokkosGemm.cpp similarity index 91% rename from PsimagLite/src/PsimagLite/kokkos_gemm.cpp rename to PsimagLite/src/PsimagLite/KokkosGemm.cpp index 2ffc79381..b66efc2da 100644 --- a/PsimagLite/src/PsimagLite/kokkos_gemm.cpp +++ b/PsimagLite/src/PsimagLite/KokkosGemm.cpp @@ -5,24 +5,8 @@ #include #include -#include - -namespace { - -// The scalar types that are floating point types and their corresponding std::complex types. -// We need to map std::complex to Kokkos::complex while keeping all the other types which is what -// KokkosType does. -template struct KokkosType { - using type = T; -}; - -template - requires(!std::is_floating_point_v) -struct KokkosType { - using type = Kokkos::complex; -}; - -} +#include +#include template inline void PsimagLite::kokkos_gemm(char transa, diff --git a/PsimagLite/src/PsimagLite/kokkos_gemm.h b/PsimagLite/src/PsimagLite/KokkosGemm.h similarity index 100% rename from PsimagLite/src/PsimagLite/kokkos_gemm.h rename to PsimagLite/src/PsimagLite/KokkosGemm.h diff --git a/PsimagLite/src/PsimagLite/KokkosType.h b/PsimagLite/src/PsimagLite/KokkosType.h new file mode 100644 index 000000000..502ce8f04 --- /dev/null +++ b/PsimagLite/src/PsimagLite/KokkosType.h @@ -0,0 +1,23 @@ +#ifndef PSIMAG_KOKKOS_TYPE_H +#define PSIMAG_KOKKOS_TYPE_H + +#include + +#include + +namespace PsimagLite { + +// The scalar types that are floating point types and their corresponding std::complex types. +// We need to map std::complex to Kokkos::complex while keeping all the other types which is what +// KokkosType does. +template struct KokkosType { + using type = T; +}; + +template struct KokkosType> { + using type = Kokkos::complex; +}; + +} + +#endif // PSIMAG_KOKKOS_TYPE_H diff --git a/dmrg/KronUtil/csr_matmul_post.cpp b/dmrg/KronUtil/csr_matmul_post.cpp index 937c359b2..6664e5c2a 100644 --- a/dmrg/KronUtil/csr_matmul_post.cpp +++ b/dmrg/KronUtil/csr_matmul_post.cpp @@ -1,4 +1,8 @@ #include "util.h" +#include + +#include +#include template void csr_matmul_post(char trans_A, @@ -33,6 +37,41 @@ void csr_matmul_post(char int isConjTranspose = (trans_A == 'C') || (trans_A == 'c'); int isConj = (trans_A == 'Z') || (trans_A == 'z'); + Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_matmul_post"); + + using ExecutionSpace = Kokkos::DefaultExecutionSpace; + ExecutionSpace exec; + + using KokkosScalar = typename PsimagLite::KokkosType::type; + + const int nnz = a.nonZeros(); + + Kokkos::View rowptr_host( + &a.getRowPtr(0), nrow_A + 1); + Kokkos::View cols_host(&a.getCol(0), + nnz); + Kokkos::View vals_host( + reinterpret_cast(&a.getValue(0)), nnz); + + Kokkos::View x_dev_out("x_dev_out", nrow_Y, ncol_X); + Kokkos::View + yin_host(reinterpret_cast(&yin(0, 0)), nrow_Y, ncol_Y); + auto y_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, yin_host); + + // Copy CSR arrays to device so the TeamPolicy kernel can access them + auto d_rowptr = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, rowptr_host); + auto d_cols = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, cols_host); + auto d_vals = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, vals_host); + + // FIXME We can try using KokkosKernels directly later + using team_policy = Kokkos::TeamPolicy; + using member_type = team_policy::member_type; + + // One team per output row iy; let Kokkos pick team/vector sizes + team_policy policy(nrow_Y, Kokkos::AUTO); if (isTranspose || isConjTranspose) { /* * ---------------------------------------------------------- @@ -45,28 +84,41 @@ void csr_matmul_post(char assert(nrow_X == nrow_Y); assert(static_cast(ncol_Y) == a.cols() && (ncol_X == nrow_A)); - int ia = 0; - for (ia = 0; ia < nrow_A; ia++) { - int istart = a.getRowPtr(ia); - int iend = a.getRowPtr(ia + 1); - int k = 0; - for (k = istart; k < iend; k++) { - int ja = a.getCol(k); - ComplexOrRealType aij = a.getValue(k); - ComplexOrRealType atji = aij; - if (is_complex && isConjTranspose) { - atji = PsimagLite::conj(atji); - }; - - int iy = 0; - for (iy = 0; iy < nrow_Y; iy++) { - int ix = iy; - int jx = ia; - - xout(ix, jx) += (yin(iy, ja) * atji); - } - } - } + Kokkos::parallel_for( + "csr_matmul_post::team_transpose", + policy, + KOKKOS_LAMBDA(const member_type& team) { + const int iy = team.league_rank(); + + // create a subview for the current Y row for faster access + auto yrow = Kokkos::subview(y_dev, iy, Kokkos::ALL); + + // For transpose: X(iy,ia) += sum_j Y(iy,ja)*A(ia,ja) + Kokkos::parallel_for( + Kokkos::TeamThreadRange(team, nrow_A), + [&](int ia) + { + int istart = d_rowptr(ia); + int iend = d_rowptr(ia + 1); + KokkosScalar local_sum; + // reduce contributions across the nonzeros of row + // ia + Kokkos::parallel_reduce( + Kokkos::ThreadVectorRange(team, istart, iend), + [&](int k, KokkosScalar& lsum) + { + int ja = d_cols(k); + KokkosScalar aij = d_vals(k); + if constexpr (is_complex) + if (isConjTranspose) + aij = Kokkos::conj(aij); + lsum += yrow(ja) * aij; + }, + local_sum); + Kokkos::single(Kokkos::PerThread(team), + [&]() { x_dev_out(iy, ia) += local_sum; }); + }); + }); } else { /* * --------------------------------------------- @@ -77,26 +129,43 @@ void csr_matmul_post(char assert(nrow_X == nrow_Y); assert(ncol_Y == nrow_A && static_cast(ncol_X) == a.cols()); - int ia = 0; - for (ia = 0; ia < nrow_A; ia++) { - int istart = a.getRowPtr(ia); - int iend = a.getRowPtr(ia + 1); - int k = 0; - for (k = istart; k < iend; k++) { - int ja = a.getCol(k); - ComplexOrRealType aij = a.getValue(k); - if (is_complex && isConj) { - aij = PsimagLite::conj(aij); - }; - - int iy = 0; - for (iy = 0; iy < nrow_Y; iy++) { - int ix = iy; - int jx = ja; - - xout(ix, jx) += (yin(iy, ia) * aij); - } - } - } + Kokkos::parallel_for( + "csr_matmul_post::team_no_transpose", + policy, + KOKKOS_LAMBDA(const member_type& team) { + const int iy = team.league_rank(); + + // create a subview for the current Y row for faster access + auto yrow = Kokkos::subview(y_dev, iy, Kokkos::ALL); + + // Non-transpose: X(iy,ja) += sum_ia Y(iy,ia)*A(ia,ja) + // Parallelize over ia (TeamThreadRange), then vectorize over + // nonzeros and use atomics for updates + Kokkos::parallel_for( + Kokkos::TeamThreadRange(team, nrow_A), + [&](int ia) + { + int istart = d_rowptr(ia); + int iend = d_rowptr(ia + 1); + KokkosScalar yval = yrow(ia); + Kokkos::parallel_for( + Kokkos::ThreadVectorRange(team, istart, iend), + [&](int k) + { + int ja = d_cols(k); + KokkosScalar aij = d_vals(k); + if constexpr (is_complex) + if (isConj) + aij = Kokkos::conj(aij); + KokkosScalar prod = yval * aij; + Kokkos::atomic_add(&x_dev_out(iy, ja), prod); + }); + }); + }); + } + auto xhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace {}, x_dev_out); + for (int iy = 0; iy < nrow_Y; ++iy) { + for (int jx = 0; jx < ncol_X; ++jx) + xout(iy, jx) += static_cast(xhost(iy, jx)); } } From e3955148cbff96d8b376ce7b6eaba3c3761cae1f Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Thu, 16 Jul 2026 14:48:51 -0400 Subject: [PATCH 02/41] Add more Kokkos profiling annotations --- PsimagLite/src/PsimagLite/kokkos_gemm.cpp | 2 +- dmrg/KronUtil/csr_den_kron_mult.cpp | 12 ++++++++++++ dmrg/KronUtil/csr_kron_mult.cpp | 9 +++++++++ dmrg/KronUtil/csr_matmul_post.cpp | 3 +++ dmrg/KronUtil/csr_matmul_pre.cpp | 3 +++ dmrg/KronUtil/den_csr_kron_mult.cpp | 12 ++++++++++++ dmrg/KronUtil/den_kron_mult.cpp | 12 ++++++++++++ dmrg/KronUtil/den_matmul_post.cpp | 10 +++++++--- dmrg/KronUtil/den_matmul_pre.cpp | 4 ++++ 9 files changed, 63 insertions(+), 4 deletions(-) diff --git a/PsimagLite/src/PsimagLite/kokkos_gemm.cpp b/PsimagLite/src/PsimagLite/kokkos_gemm.cpp index 2ffc79381..dbd551a2a 100644 --- a/PsimagLite/src/PsimagLite/kokkos_gemm.cpp +++ b/PsimagLite/src/PsimagLite/kokkos_gemm.cpp @@ -39,7 +39,7 @@ inline void PsimagLite::kokkos_gemm(char transa, Scalar* C, IntegerForBlasType ldc) { - Kokkos::Profiling::ScopedRegion scoped_region("kokkos_gemm"); + Kokkos::Profiling::ScopedRegion scoped_region("PsimagLite::kokkos_gemm"); int M = static_cast(m); int N = static_cast(n); int K = static_cast(k); diff --git a/dmrg/KronUtil/csr_den_kron_mult.cpp b/dmrg/KronUtil/csr_den_kron_mult.cpp index 56170dadb..633fda1cb 100644 --- a/dmrg/KronUtil/csr_den_kron_mult.cpp +++ b/dmrg/KronUtil/csr_den_kron_mult.cpp @@ -1,5 +1,7 @@ #include "util.h" +#include + template void csr_den_kron_mult_method(const int imethod, const char transA, @@ -12,6 +14,8 @@ void csr_den_kron_mult_method(const int SizeType offsetX, PsimagLite::GemmR& gemmR) { + Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_den_kron_mult_method"); + const bool is_complex = PsimagLite::IsComplexNumber::True; const int isTransA = (transA == 'T') || (transA == 't'); const int isConjTransA = (transA == 'C') || (transA == 'c'); @@ -85,6 +89,8 @@ void csr_den_kron_mult_method(const int */ if (imethod == 1) { + Kokkos::Profiling::ScopedRegion region( + "PsimagLite::csr_den_kron_mult_method::imethod1"); /* * -------------------------------------------- @@ -167,6 +173,9 @@ void csr_den_kron_mult_method(const int xout); } } else if (imethod == 2) { + Kokkos::Profiling::ScopedRegion region( + "PsimagLite::csr_den_kron_mult_method::imethod2"); + /* * --------------------- * YAt(jb,ia) = Y(jb,ja) * tranpose(A(ia,ja)) @@ -246,6 +255,9 @@ void csr_den_kron_mult_method(const int gemmR); } } else if (imethod == 3) { + Kokkos::Profiling::ScopedRegion region( + "PsimagLite::csr_den_kron_mult_method::imethod3"); + /* * --------------------------------------------- * C = kron(A,B) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index badc6dad6..ab480c2b7 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -1,5 +1,7 @@ #include "util.h" +#include + template void csr_to_den(const PsimagLite::CrsMatrix& a, PsimagLite::Matrix& a_) @@ -40,6 +42,8 @@ void csr_kron_mult_method(const int imethod, const PsimagLite::MatrixNonOwned& yin, PsimagLite::MatrixNonOwned& xout) { + Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_kron_mult_method"); + const bool is_complex = PsimagLite::IsComplexNumber::True; const int isTransA = (transA == 'T') || (transA == 't'); const int isTransB = (transB == 'T') || (transB == 't'); @@ -109,6 +113,7 @@ void csr_kron_mult_method(const int imethod, */ if (imethod == 1) { + Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod1"); /* * -------------------------------------------- @@ -185,6 +190,8 @@ void csr_kron_mult_method(const int imethod, xout); } } else if (imethod == 2) { + Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod2"); + /* * --------------------- * YAt(jb,ia) = Y(jb,ja) * tranpose(A(ia,ja)) @@ -262,6 +269,8 @@ void csr_kron_mult_method(const int imethod, xout); } } else if (imethod == 3) { + Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3"); + /* * --------------------------------------------- * C = kron(A,B) diff --git a/dmrg/KronUtil/csr_matmul_post.cpp b/dmrg/KronUtil/csr_matmul_post.cpp index 937c359b2..0c458e35c 100644 --- a/dmrg/KronUtil/csr_matmul_post.cpp +++ b/dmrg/KronUtil/csr_matmul_post.cpp @@ -1,5 +1,7 @@ #include "util.h" +#include + template void csr_matmul_post(char trans_A, const PsimagLite::CrsMatrix& a, @@ -10,6 +12,7 @@ void csr_matmul_post(char const int ncol_X, PsimagLite::MatrixNonOwned& xout) { + Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_matmul_post"); /* * ------------------------------------------------------- * A in compressed sparse ROW format diff --git a/dmrg/KronUtil/csr_matmul_pre.cpp b/dmrg/KronUtil/csr_matmul_pre.cpp index bb0d03b0b..1d9603441 100644 --- a/dmrg/KronUtil/csr_matmul_pre.cpp +++ b/dmrg/KronUtil/csr_matmul_pre.cpp @@ -1,5 +1,7 @@ #include "util.h" +#include + template void csr_matmul_pre(char trans_A, const PsimagLite::CrsMatrix& a, @@ -10,6 +12,7 @@ void csr_matmul_pre(char t const int ncol_X, PsimagLite::MatrixNonOwned& xout) { + Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_matmul_pre"); /* * ------------------------------------------------------- * A in compressed sparse ROW format diff --git a/dmrg/KronUtil/den_csr_kron_mult.cpp b/dmrg/KronUtil/den_csr_kron_mult.cpp index 9bcb0e7d7..b66d05d2e 100644 --- a/dmrg/KronUtil/den_csr_kron_mult.cpp +++ b/dmrg/KronUtil/den_csr_kron_mult.cpp @@ -1,5 +1,7 @@ #include "util.h" +#include + template void den_csr_kron_mult_method(const int imethod, const char transA, @@ -12,6 +14,8 @@ void den_csr_kron_mult_method(const int SizeType offsetX, PsimagLite::GemmR& gemmR) { + Kokkos::Profiling::ScopedRegion region("PsimagLite::den_csr_kron_mult_method"); + const bool is_complex = PsimagLite::IsComplexNumber::True; const int isTransA = (transA == 'T') || (transA == 't'); @@ -87,6 +91,8 @@ void den_csr_kron_mult_method(const int */ if (imethod == 1) { + Kokkos::Profiling::ScopedRegion region( + "PsimagLite::den_csr_kron_mult_method::imethod1"); /* * -------------------------------------------- @@ -192,6 +198,9 @@ void den_csr_kron_mult_method(const int }; } } else if (imethod == 2) { + Kokkos::Profiling::ScopedRegion region( + "PsimagLite::den_csr_kron_mult_method::imethod2"); + /* * --------------------- * YAt(jb,ia) = Y(jb,ja) * tranpose(A(ia,ja)) @@ -294,6 +303,9 @@ void den_csr_kron_mult_method(const int xout); } } else if (imethod == 3) { + Kokkos::Profiling::ScopedRegion region( + "PsimagLite::den_csr_kron_mult_method::imethod3"); + /* * --------------------------------------------- * C = kron(A,B) diff --git a/dmrg/KronUtil/den_kron_mult.cpp b/dmrg/KronUtil/den_kron_mult.cpp index 02220f21c..9cd8f7979 100644 --- a/dmrg/KronUtil/den_kron_mult.cpp +++ b/dmrg/KronUtil/den_kron_mult.cpp @@ -1,5 +1,7 @@ #include "util.h" +#include + template void den_kron_mult_method(const int imethod, const char transA, @@ -12,6 +14,8 @@ void den_kron_mult_method(const int SizeType offsetX, PsimagLite::GemmR& gemmR) { + Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_den_kron_mult_method"); + const bool is_complex = PsimagLite::IsComplexNumber::True; const int nrow_A = a_.n_row(); const int ncol_A = a_.n_col(); @@ -69,6 +73,8 @@ void den_kron_mult_method(const int */ if (imethod == 1) { + Kokkos::Profiling::ScopedRegion region( + "PsimagLite::csr_den_kron_mult_method::imethod1"); /* * -------------------------------------------- @@ -173,6 +179,9 @@ void den_kron_mult_method(const int }; } } else if (imethod == 2) { + Kokkos::Profiling::ScopedRegion region( + "PsimagLite::csr_den_kron_mult_method::imethod2"); + /* * --------------------- * YAt(jb,ia) = Y(jb,ja) * tranpose(A(ia,ja)) @@ -276,6 +285,9 @@ void den_kron_mult_method(const int gemmR); } } else if (imethod == 3) { + Kokkos::Profiling::ScopedRegion region( + "PsimagLite::csr_den_kron_mult_method::imethod3"); + /* * --------------------------------------------- * C = kron(A,B) diff --git a/dmrg/KronUtil/den_matmul_post.cpp b/dmrg/KronUtil/den_matmul_post.cpp index 53bbba390..7736086ad 100644 --- a/dmrg/KronUtil/den_matmul_post.cpp +++ b/dmrg/KronUtil/den_matmul_post.cpp @@ -1,5 +1,7 @@ #include "util.h" +#include + template void den_matmul_post(const char trans_A, const int nrow_A, @@ -13,7 +15,8 @@ void den_matmul_post(const char PsimagLite::MatrixNonOwned& xout, PsimagLite::GemmR& gemmR) { - const bool is_complex = PsimagLite::IsComplexNumber::True; + Kokkos::Profiling::ScopedRegion region("PsimagLite::den_matmul_post"); + /* * ------------------------------------------------------- * A in dense matrix format @@ -33,8 +36,9 @@ void den_matmul_post(const char * ------------------------------------------------------- */ - int isTranspose = (trans_A == 'T') || (trans_A == 't'); - int isConjTranspose = (trans_A == 'C') || (trans_A == 'c'); + const bool is_complex = PsimagLite::IsComplexNumber::True; + int isTranspose = (trans_A == 'T') || (trans_A == 't'); + int isConjTranspose = (trans_A == 'C') || (trans_A == 'c'); const bool use_blas = true; diff --git a/dmrg/KronUtil/den_matmul_pre.cpp b/dmrg/KronUtil/den_matmul_pre.cpp index 909b2c637..6637a2ce4 100644 --- a/dmrg/KronUtil/den_matmul_pre.cpp +++ b/dmrg/KronUtil/den_matmul_pre.cpp @@ -1,5 +1,7 @@ #include "util.h" +#include + template void den_matmul_pre(const char trans_A, const int nrow_A, @@ -13,6 +15,8 @@ void den_matmul_pre(const char t PsimagLite::MatrixNonOwned& xout, PsimagLite::GemmR& gemmR) { + Kokkos::Profiling::ScopedRegion region("PsimagLite::den_matmul_pre"); + /* * ------------------------------------------------------- * A in dense matrix format From 92a8d617314ae36ac260c669de617eee7ecad42e Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Fri, 17 Jul 2026 09:23:01 -0400 Subject: [PATCH 03/41] Print sizes csr matrix --- dmrg/KronUtil/csr_kron_mult.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index ab480c2b7..2dcea72ee 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -270,7 +270,6 @@ void csr_kron_mult_method(const int imethod, } } else if (imethod == 3) { Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3"); - /* * --------------------------------------------- * C = kron(A,B) @@ -279,6 +278,11 @@ void csr_kron_mult_method(const int imethod, * --------------------------------------------- */ + std::cerr << "nrow_A: " << nrow_A << ' ' + << "nnz_A: " << csr_nnz(a) << ' ' + << "nrow_B: " << nrow_B << ' ' + << "nnz_B: " << csr_nnz(b) << '\n'; + int ia = 0; int ka = 0; int ib = 0; From aeddf1c97efafda95e5a28768a5bc61571280a05 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Fri, 17 Jul 2026 10:03:20 -0400 Subject: [PATCH 04/41] Joint A and B --- dmrg/KronUtil/csr_kron_mult.cpp | 155 +++++++++++++++++++++++------- dmrg/KronUtil/csr_matmul_post.cpp | 2 - 2 files changed, 121 insertions(+), 36 deletions(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 2dcea72ee..93f6a02a2 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -1,6 +1,7 @@ #include "util.h" - +#include #include +#include template void csr_to_den(const PsimagLite::CrsMatrix& a, @@ -278,48 +279,134 @@ void csr_kron_mult_method(const int imethod, * --------------------------------------------- */ +#if 0 std::cerr << "nrow_A: " << nrow_A << ' ' << "nnz_A: " << csr_nnz(a) << ' ' << "nrow_B: " << nrow_B << ' ' << "nnz_B: " << csr_nnz(b) << '\n'; - - int ia = 0; - int ka = 0; - int ib = 0; - int kb = 0; - for (ia = 0; ia < nrow_A; ia++) { - int istarta = a.getRowPtr(ia); - int ienda = a.getRowPtr(ia + 1); - for (ka = istarta; ka < ienda; ka++) { - int ja = a.getCol(ka); - ComplexOrRealType aij = a.getValue(ka); - if (is_complex && isConjTransA) { - aij = PsimagLite::conj(aij); - }; +#endif - for (ib = 0; ib < nrow_B; ib++) { - int istartb = b.getRowPtr(ib); - int iendb = b.getRowPtr(ib + 1); + // Build flat lists of nonzeros for A and B on the host, then copy to device + using ExecutionSpace = Kokkos::DefaultExecutionSpace; + using KokkosScalar = typename PsimagLite::KokkosType::type; - for (kb = istartb; kb < iendb; kb++) { - int jb = b.getCol(kb); - ComplexOrRealType bij = b.getValue(kb); - if (is_complex && isConjTransB) { - bij = PsimagLite::conj(bij); - }; + int nnzA = csr_nnz(a); + int nnzB = csr_nnz(b); - ComplexOrRealType cij = aij * bij; + // host-side temporary arrays + std::vector A_row(nnzA); + std::vector A_col(nnzA); + std::vector A_val(nnzA); + { + int idx = 0; + for (int ia = 0; ia < nrow_A; ++ia) { + int istart = a.getRowPtr(ia); + int iend = a.getRowPtr(ia + 1); + for (int ka = istart; ka < iend; ++ka) { + A_row[idx] = ia; + A_col[idx] = a.getCol(ka); + ComplexOrRealType aval = a.getValue(ka); + if (is_complex && isConjTransA) + aval = PsimagLite::conj(aval); + A_val[idx] = static_cast(aval); + ++idx; + } + } + } - int ix = (isTransB || isConjTransB) ? jb : ib; - int jx = (isTransA || isConjTransA) ? ja : ia; - int iy = (isTransB || isConjTransB) ? ib : jb; - int jy = (isTransA || isConjTransA) ? ia : ja; + std::vector B_row(nnzB); + std::vector B_col(nnzB); + std::vector B_val(nnzB); + { + int idx = 0; + for (int ib = 0; ib < nrow_B; ++ib) { + int istart = b.getRowPtr(ib); + int iend = b.getRowPtr(ib + 1); + for (int kb = istart; kb < iend; ++kb) { + B_row[idx] = ib; + B_col[idx] = b.getCol(kb); + ComplexOrRealType bval = b.getValue(kb); + if (is_complex && isConjTransB) + bval = PsimagLite::conj(bval); + B_val[idx] = static_cast(bval); + ++idx; + } + } + } - xout(ix, jx) += cij * yin(iy, jy); - }; - }; - }; - }; + // create device views + Kokkos::View A_row_h("A_row_h", nnzA); + Kokkos::View A_col_h("A_col_h", nnzA); + Kokkos::View A_val_h("A_val_h", nnzA); + Kokkos::View B_row_h("B_row_h", nnzB); + Kokkos::View B_col_h("B_col_h", nnzB); + Kokkos::View B_val_h("B_val_h", nnzB); + + for (int i = 0; i < nnzA; ++i) { + A_row_h(i) = A_row[i]; + A_col_h(i) = A_col[i]; + A_val_h(i) = A_val[i]; + } + for (int i = 0; i < nnzB; ++i) { + B_row_h(i) = B_row[i]; + B_col_h(i) = B_col[i]; + B_val_h(i) = B_val[i]; + } + + auto A_row_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_row_h); + auto A_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_col_h); + auto A_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_val_h); + auto B_row_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_row_h); + auto B_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_col_h); + auto B_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_val_h); + // device yin and xout + + auto yin_host = Kokkos::View( + reinterpret_cast(&yin(0, 0)), nrow_Y, ncol_Y); + + auto y_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, yin_host); + + auto x_dev = Kokkos::View("x_dev", nrow_X, ncol_X); + Kokkos::deep_copy(x_dev, KokkosScalar(0)); + + const size_t totalPairs = static_cast(nnzA) * static_cast(nnzB); + + Kokkos::parallel_for( + "csr_kron_mult::imethod3_pairs", + Kokkos::RangePolicy(0, totalPairs), + KOKKOS_LAMBDA(const size_t idx) { + const int ia_idx = static_cast(idx / nnzB); + const int ib_idx = static_cast(idx % nnzB); + + int ia = A_row_dev(ia_idx); + int ja = A_col_dev(ia_idx); + KokkosScalar aij = A_val_dev(ia_idx); + + int ib = B_row_dev(ib_idx); + int jb = B_col_dev(ib_idx); + KokkosScalar bij = B_val_dev(ib_idx); + + KokkosScalar cij = aij * bij; + + int ix = (isTransB || isConjTransB) ? jb : ib; + int jx = (isTransA || isConjTransA) ? ja : ia; + int iy = (isTransB || isConjTransB) ? ib : jb; + int jy = (isTransA || isConjTransA) ? ia : ja; + + KokkosScalar prod = cij * y_dev(iy, jy); + Kokkos::atomic_add(&x_dev(ix, jx), prod); + }); + + // copy back and accumulate into xout + + auto xhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace {}, x_dev); + for (int ix = 0; ix < nrow_X; ++ix) { + for (int jx = 0; jx < ncol_X; ++jx) + xout(ix, jx) += static_cast(xhost(ix, jx)); + } }; } diff --git a/dmrg/KronUtil/csr_matmul_post.cpp b/dmrg/KronUtil/csr_matmul_post.cpp index fc895d61f..557677980 100644 --- a/dmrg/KronUtil/csr_matmul_post.cpp +++ b/dmrg/KronUtil/csr_matmul_post.cpp @@ -40,8 +40,6 @@ void csr_matmul_post(char int isConjTranspose = (trans_A == 'C') || (trans_A == 'c'); int isConj = (trans_A == 'Z') || (trans_A == 'z'); - Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_matmul_post"); - using ExecutionSpace = Kokkos::DefaultExecutionSpace; ExecutionSpace exec; From 0cb7206be55dee2d69cde7fdcbbc0256e8b52026 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Fri, 17 Jul 2026 10:18:22 -0400 Subject: [PATCH 05/41] Revert "Joint A and B" This reverts commit aeddf1c97efafda95e5a28768a5bc61571280a05. --- dmrg/KronUtil/csr_kron_mult.cpp | 155 +++++++----------------------- dmrg/KronUtil/csr_matmul_post.cpp | 2 + 2 files changed, 36 insertions(+), 121 deletions(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 93f6a02a2..2dcea72ee 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -1,7 +1,6 @@ #include "util.h" -#include + #include -#include template void csr_to_den(const PsimagLite::CrsMatrix& a, @@ -279,134 +278,48 @@ void csr_kron_mult_method(const int imethod, * --------------------------------------------- */ -#if 0 std::cerr << "nrow_A: " << nrow_A << ' ' << "nnz_A: " << csr_nnz(a) << ' ' << "nrow_B: " << nrow_B << ' ' << "nnz_B: " << csr_nnz(b) << '\n'; -#endif - - // Build flat lists of nonzeros for A and B on the host, then copy to device - using ExecutionSpace = Kokkos::DefaultExecutionSpace; - using KokkosScalar = typename PsimagLite::KokkosType::type; - - int nnzA = csr_nnz(a); - int nnzB = csr_nnz(b); - - // host-side temporary arrays - std::vector A_row(nnzA); - std::vector A_col(nnzA); - std::vector A_val(nnzA); - { - int idx = 0; - for (int ia = 0; ia < nrow_A; ++ia) { - int istart = a.getRowPtr(ia); - int iend = a.getRowPtr(ia + 1); - for (int ka = istart; ka < iend; ++ka) { - A_row[idx] = ia; - A_col[idx] = a.getCol(ka); - ComplexOrRealType aval = a.getValue(ka); - if (is_complex && isConjTransA) - aval = PsimagLite::conj(aval); - A_val[idx] = static_cast(aval); - ++idx; - } - } - } - - std::vector B_row(nnzB); - std::vector B_col(nnzB); - std::vector B_val(nnzB); - { - int idx = 0; - for (int ib = 0; ib < nrow_B; ++ib) { - int istart = b.getRowPtr(ib); - int iend = b.getRowPtr(ib + 1); - for (int kb = istart; kb < iend; ++kb) { - B_row[idx] = ib; - B_col[idx] = b.getCol(kb); - ComplexOrRealType bval = b.getValue(kb); - if (is_complex && isConjTransB) - bval = PsimagLite::conj(bval); - B_val[idx] = static_cast(bval); - ++idx; - } - } - } - - // create device views - Kokkos::View A_row_h("A_row_h", nnzA); - Kokkos::View A_col_h("A_col_h", nnzA); - Kokkos::View A_val_h("A_val_h", nnzA); - Kokkos::View B_row_h("B_row_h", nnzB); - Kokkos::View B_col_h("B_col_h", nnzB); - Kokkos::View B_val_h("B_val_h", nnzB); - - for (int i = 0; i < nnzA; ++i) { - A_row_h(i) = A_row[i]; - A_col_h(i) = A_col[i]; - A_val_h(i) = A_val[i]; - } - for (int i = 0; i < nnzB; ++i) { - B_row_h(i) = B_row[i]; - B_col_h(i) = B_col[i]; - B_val_h(i) = B_val[i]; - } - - auto A_row_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_row_h); - auto A_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_col_h); - auto A_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_val_h); - auto B_row_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_row_h); - auto B_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_col_h); - auto B_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_val_h); - // device yin and xout - - auto yin_host = Kokkos::View( - reinterpret_cast(&yin(0, 0)), nrow_Y, ncol_Y); - - auto y_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, yin_host); - - auto x_dev = Kokkos::View("x_dev", nrow_X, ncol_X); - Kokkos::deep_copy(x_dev, KokkosScalar(0)); - - const size_t totalPairs = static_cast(nnzA) * static_cast(nnzB); - - Kokkos::parallel_for( - "csr_kron_mult::imethod3_pairs", - Kokkos::RangePolicy(0, totalPairs), - KOKKOS_LAMBDA(const size_t idx) { - const int ia_idx = static_cast(idx / nnzB); - const int ib_idx = static_cast(idx % nnzB); - - int ia = A_row_dev(ia_idx); - int ja = A_col_dev(ia_idx); - KokkosScalar aij = A_val_dev(ia_idx); - - int ib = B_row_dev(ib_idx); - int jb = B_col_dev(ib_idx); - KokkosScalar bij = B_val_dev(ib_idx); + + int ia = 0; + int ka = 0; + int ib = 0; + int kb = 0; + for (ia = 0; ia < nrow_A; ia++) { + int istarta = a.getRowPtr(ia); + int ienda = a.getRowPtr(ia + 1); + for (ka = istarta; ka < ienda; ka++) { + int ja = a.getCol(ka); + ComplexOrRealType aij = a.getValue(ka); + if (is_complex && isConjTransA) { + aij = PsimagLite::conj(aij); + }; - KokkosScalar cij = aij * bij; + for (ib = 0; ib < nrow_B; ib++) { + int istartb = b.getRowPtr(ib); + int iendb = b.getRowPtr(ib + 1); - int ix = (isTransB || isConjTransB) ? jb : ib; - int jx = (isTransA || isConjTransA) ? ja : ia; - int iy = (isTransB || isConjTransB) ? ib : jb; - int jy = (isTransA || isConjTransA) ? ia : ja; + for (kb = istartb; kb < iendb; kb++) { + int jb = b.getCol(kb); + ComplexOrRealType bij = b.getValue(kb); + if (is_complex && isConjTransB) { + bij = PsimagLite::conj(bij); + }; - KokkosScalar prod = cij * y_dev(iy, jy); - Kokkos::atomic_add(&x_dev(ix, jx), prod); - }); + ComplexOrRealType cij = aij * bij; - // copy back and accumulate into xout + int ix = (isTransB || isConjTransB) ? jb : ib; + int jx = (isTransA || isConjTransA) ? ja : ia; + int iy = (isTransB || isConjTransB) ? ib : jb; + int jy = (isTransA || isConjTransA) ? ia : ja; - auto xhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace {}, x_dev); - for (int ix = 0; ix < nrow_X; ++ix) { - for (int jx = 0; jx < ncol_X; ++jx) - xout(ix, jx) += static_cast(xhost(ix, jx)); - } + xout(ix, jx) += cij * yin(iy, jy); + }; + }; + }; + }; }; } diff --git a/dmrg/KronUtil/csr_matmul_post.cpp b/dmrg/KronUtil/csr_matmul_post.cpp index 557677980..fc895d61f 100644 --- a/dmrg/KronUtil/csr_matmul_post.cpp +++ b/dmrg/KronUtil/csr_matmul_post.cpp @@ -40,6 +40,8 @@ void csr_matmul_post(char int isConjTranspose = (trans_A == 'C') || (trans_A == 'c'); int isConj = (trans_A == 'Z') || (trans_A == 'z'); + Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_matmul_post"); + using ExecutionSpace = Kokkos::DefaultExecutionSpace; ExecutionSpace exec; From 9818a77c50ec9f60a46b5df95bb8cbca12e3207f Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 20 Jul 2026 13:08:08 -0400 Subject: [PATCH 06/41] Fix --- dmrg/KronUtil/csr_kron_mult.cpp | 6 +- dmrg/KronUtil/csr_matmul_post.cpp | 158 ++++++++---------------------- 2 files changed, 47 insertions(+), 117 deletions(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 2dcea72ee..34c8e28da 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -278,11 +278,13 @@ void csr_kron_mult_method(const int imethod, * --------------------------------------------- */ +#if 0 std::cerr << "nrow_A: " << nrow_A << ' ' << "nnz_A: " << csr_nnz(a) << ' ' << "nrow_B: " << nrow_B << ' ' - << "nnz_B: " << csr_nnz(b) << '\n'; - + << "nnz_B: " << csr_nnz(b) << '\n'; +#endif + int ia = 0; int ka = 0; int ib = 0; diff --git a/dmrg/KronUtil/csr_matmul_post.cpp b/dmrg/KronUtil/csr_matmul_post.cpp index fc895d61f..937c359b2 100644 --- a/dmrg/KronUtil/csr_matmul_post.cpp +++ b/dmrg/KronUtil/csr_matmul_post.cpp @@ -1,10 +1,4 @@ #include "util.h" -#include - -#include -#include - -#include template void csr_matmul_post(char trans_A, @@ -16,7 +10,6 @@ void csr_matmul_post(char const int ncol_X, PsimagLite::MatrixNonOwned& xout) { - Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_matmul_post"); /* * ------------------------------------------------------- * A in compressed sparse ROW format @@ -40,41 +33,6 @@ void csr_matmul_post(char int isConjTranspose = (trans_A == 'C') || (trans_A == 'c'); int isConj = (trans_A == 'Z') || (trans_A == 'z'); - Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_matmul_post"); - - using ExecutionSpace = Kokkos::DefaultExecutionSpace; - ExecutionSpace exec; - - using KokkosScalar = typename PsimagLite::KokkosType::type; - - const int nnz = a.nonZeros(); - - Kokkos::View rowptr_host( - &a.getRowPtr(0), nrow_A + 1); - Kokkos::View cols_host(&a.getCol(0), - nnz); - Kokkos::View vals_host( - reinterpret_cast(&a.getValue(0)), nnz); - - Kokkos::View x_dev_out("x_dev_out", nrow_Y, ncol_X); - Kokkos::View - yin_host(reinterpret_cast(&yin(0, 0)), nrow_Y, ncol_Y); - auto y_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, yin_host); - - // Copy CSR arrays to device so the TeamPolicy kernel can access them - auto d_rowptr = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, rowptr_host); - auto d_cols = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, cols_host); - auto d_vals = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, vals_host); - - // FIXME We can try using KokkosKernels directly later - using team_policy = Kokkos::TeamPolicy; - using member_type = team_policy::member_type; - - // One team per output row iy; let Kokkos pick team/vector sizes - team_policy policy(nrow_Y, Kokkos::AUTO); if (isTranspose || isConjTranspose) { /* * ---------------------------------------------------------- @@ -87,41 +45,28 @@ void csr_matmul_post(char assert(nrow_X == nrow_Y); assert(static_cast(ncol_Y) == a.cols() && (ncol_X == nrow_A)); - Kokkos::parallel_for( - "csr_matmul_post::team_transpose", - policy, - KOKKOS_LAMBDA(const member_type& team) { - const int iy = team.league_rank(); - - // create a subview for the current Y row for faster access - auto yrow = Kokkos::subview(y_dev, iy, Kokkos::ALL); - - // For transpose: X(iy,ia) += sum_j Y(iy,ja)*A(ia,ja) - Kokkos::parallel_for( - Kokkos::TeamThreadRange(team, nrow_A), - [&](int ia) - { - int istart = d_rowptr(ia); - int iend = d_rowptr(ia + 1); - KokkosScalar local_sum; - // reduce contributions across the nonzeros of row - // ia - Kokkos::parallel_reduce( - Kokkos::ThreadVectorRange(team, istart, iend), - [&](int k, KokkosScalar& lsum) - { - int ja = d_cols(k); - KokkosScalar aij = d_vals(k); - if constexpr (is_complex) - if (isConjTranspose) - aij = Kokkos::conj(aij); - lsum += yrow(ja) * aij; - }, - local_sum); - Kokkos::single(Kokkos::PerThread(team), - [&]() { x_dev_out(iy, ia) += local_sum; }); - }); - }); + int ia = 0; + for (ia = 0; ia < nrow_A; ia++) { + int istart = a.getRowPtr(ia); + int iend = a.getRowPtr(ia + 1); + int k = 0; + for (k = istart; k < iend; k++) { + int ja = a.getCol(k); + ComplexOrRealType aij = a.getValue(k); + ComplexOrRealType atji = aij; + if (is_complex && isConjTranspose) { + atji = PsimagLite::conj(atji); + }; + + int iy = 0; + for (iy = 0; iy < nrow_Y; iy++) { + int ix = iy; + int jx = ia; + + xout(ix, jx) += (yin(iy, ja) * atji); + } + } + } } else { /* * --------------------------------------------- @@ -132,43 +77,26 @@ void csr_matmul_post(char assert(nrow_X == nrow_Y); assert(ncol_Y == nrow_A && static_cast(ncol_X) == a.cols()); - Kokkos::parallel_for( - "csr_matmul_post::team_no_transpose", - policy, - KOKKOS_LAMBDA(const member_type& team) { - const int iy = team.league_rank(); - - // create a subview for the current Y row for faster access - auto yrow = Kokkos::subview(y_dev, iy, Kokkos::ALL); - - // Non-transpose: X(iy,ja) += sum_ia Y(iy,ia)*A(ia,ja) - // Parallelize over ia (TeamThreadRange), then vectorize over - // nonzeros and use atomics for updates - Kokkos::parallel_for( - Kokkos::TeamThreadRange(team, nrow_A), - [&](int ia) - { - int istart = d_rowptr(ia); - int iend = d_rowptr(ia + 1); - KokkosScalar yval = yrow(ia); - Kokkos::parallel_for( - Kokkos::ThreadVectorRange(team, istart, iend), - [&](int k) - { - int ja = d_cols(k); - KokkosScalar aij = d_vals(k); - if constexpr (is_complex) - if (isConj) - aij = Kokkos::conj(aij); - KokkosScalar prod = yval * aij; - Kokkos::atomic_add(&x_dev_out(iy, ja), prod); - }); - }); - }); - } - auto xhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace {}, x_dev_out); - for (int iy = 0; iy < nrow_Y; ++iy) { - for (int jx = 0; jx < ncol_X; ++jx) - xout(iy, jx) += static_cast(xhost(iy, jx)); + int ia = 0; + for (ia = 0; ia < nrow_A; ia++) { + int istart = a.getRowPtr(ia); + int iend = a.getRowPtr(ia + 1); + int k = 0; + for (k = istart; k < iend; k++) { + int ja = a.getCol(k); + ComplexOrRealType aij = a.getValue(k); + if (is_complex && isConj) { + aij = PsimagLite::conj(aij); + }; + + int iy = 0; + for (iy = 0; iy < nrow_Y; iy++) { + int ix = iy; + int jx = ja; + + xout(ix, jx) += (yin(iy, ia) * aij); + } + } + } } } From fbdc2b229cb092879a6f4ab7249a8f42d382a1c4 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 20 Jul 2026 13:08:39 -0400 Subject: [PATCH 07/41] Revert "Fix" This reverts commit 9818a77c50ec9f60a46b5df95bb8cbca12e3207f. --- dmrg/KronUtil/csr_kron_mult.cpp | 6 +- dmrg/KronUtil/csr_matmul_post.cpp | 158 ++++++++++++++++++++++-------- 2 files changed, 117 insertions(+), 47 deletions(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 34c8e28da..2dcea72ee 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -278,13 +278,11 @@ void csr_kron_mult_method(const int imethod, * --------------------------------------------- */ -#if 0 std::cerr << "nrow_A: " << nrow_A << ' ' << "nnz_A: " << csr_nnz(a) << ' ' << "nrow_B: " << nrow_B << ' ' - << "nnz_B: " << csr_nnz(b) << '\n'; -#endif - + << "nnz_B: " << csr_nnz(b) << '\n'; + int ia = 0; int ka = 0; int ib = 0; diff --git a/dmrg/KronUtil/csr_matmul_post.cpp b/dmrg/KronUtil/csr_matmul_post.cpp index 937c359b2..fc895d61f 100644 --- a/dmrg/KronUtil/csr_matmul_post.cpp +++ b/dmrg/KronUtil/csr_matmul_post.cpp @@ -1,4 +1,10 @@ #include "util.h" +#include + +#include +#include + +#include template void csr_matmul_post(char trans_A, @@ -10,6 +16,7 @@ void csr_matmul_post(char const int ncol_X, PsimagLite::MatrixNonOwned& xout) { + Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_matmul_post"); /* * ------------------------------------------------------- * A in compressed sparse ROW format @@ -33,6 +40,41 @@ void csr_matmul_post(char int isConjTranspose = (trans_A == 'C') || (trans_A == 'c'); int isConj = (trans_A == 'Z') || (trans_A == 'z'); + Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_matmul_post"); + + using ExecutionSpace = Kokkos::DefaultExecutionSpace; + ExecutionSpace exec; + + using KokkosScalar = typename PsimagLite::KokkosType::type; + + const int nnz = a.nonZeros(); + + Kokkos::View rowptr_host( + &a.getRowPtr(0), nrow_A + 1); + Kokkos::View cols_host(&a.getCol(0), + nnz); + Kokkos::View vals_host( + reinterpret_cast(&a.getValue(0)), nnz); + + Kokkos::View x_dev_out("x_dev_out", nrow_Y, ncol_X); + Kokkos::View + yin_host(reinterpret_cast(&yin(0, 0)), nrow_Y, ncol_Y); + auto y_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, yin_host); + + // Copy CSR arrays to device so the TeamPolicy kernel can access them + auto d_rowptr = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, rowptr_host); + auto d_cols = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, cols_host); + auto d_vals = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, vals_host); + + // FIXME We can try using KokkosKernels directly later + using team_policy = Kokkos::TeamPolicy; + using member_type = team_policy::member_type; + + // One team per output row iy; let Kokkos pick team/vector sizes + team_policy policy(nrow_Y, Kokkos::AUTO); if (isTranspose || isConjTranspose) { /* * ---------------------------------------------------------- @@ -45,28 +87,41 @@ void csr_matmul_post(char assert(nrow_X == nrow_Y); assert(static_cast(ncol_Y) == a.cols() && (ncol_X == nrow_A)); - int ia = 0; - for (ia = 0; ia < nrow_A; ia++) { - int istart = a.getRowPtr(ia); - int iend = a.getRowPtr(ia + 1); - int k = 0; - for (k = istart; k < iend; k++) { - int ja = a.getCol(k); - ComplexOrRealType aij = a.getValue(k); - ComplexOrRealType atji = aij; - if (is_complex && isConjTranspose) { - atji = PsimagLite::conj(atji); - }; - - int iy = 0; - for (iy = 0; iy < nrow_Y; iy++) { - int ix = iy; - int jx = ia; - - xout(ix, jx) += (yin(iy, ja) * atji); - } - } - } + Kokkos::parallel_for( + "csr_matmul_post::team_transpose", + policy, + KOKKOS_LAMBDA(const member_type& team) { + const int iy = team.league_rank(); + + // create a subview for the current Y row for faster access + auto yrow = Kokkos::subview(y_dev, iy, Kokkos::ALL); + + // For transpose: X(iy,ia) += sum_j Y(iy,ja)*A(ia,ja) + Kokkos::parallel_for( + Kokkos::TeamThreadRange(team, nrow_A), + [&](int ia) + { + int istart = d_rowptr(ia); + int iend = d_rowptr(ia + 1); + KokkosScalar local_sum; + // reduce contributions across the nonzeros of row + // ia + Kokkos::parallel_reduce( + Kokkos::ThreadVectorRange(team, istart, iend), + [&](int k, KokkosScalar& lsum) + { + int ja = d_cols(k); + KokkosScalar aij = d_vals(k); + if constexpr (is_complex) + if (isConjTranspose) + aij = Kokkos::conj(aij); + lsum += yrow(ja) * aij; + }, + local_sum); + Kokkos::single(Kokkos::PerThread(team), + [&]() { x_dev_out(iy, ia) += local_sum; }); + }); + }); } else { /* * --------------------------------------------- @@ -77,26 +132,43 @@ void csr_matmul_post(char assert(nrow_X == nrow_Y); assert(ncol_Y == nrow_A && static_cast(ncol_X) == a.cols()); - int ia = 0; - for (ia = 0; ia < nrow_A; ia++) { - int istart = a.getRowPtr(ia); - int iend = a.getRowPtr(ia + 1); - int k = 0; - for (k = istart; k < iend; k++) { - int ja = a.getCol(k); - ComplexOrRealType aij = a.getValue(k); - if (is_complex && isConj) { - aij = PsimagLite::conj(aij); - }; - - int iy = 0; - for (iy = 0; iy < nrow_Y; iy++) { - int ix = iy; - int jx = ja; - - xout(ix, jx) += (yin(iy, ia) * aij); - } - } - } + Kokkos::parallel_for( + "csr_matmul_post::team_no_transpose", + policy, + KOKKOS_LAMBDA(const member_type& team) { + const int iy = team.league_rank(); + + // create a subview for the current Y row for faster access + auto yrow = Kokkos::subview(y_dev, iy, Kokkos::ALL); + + // Non-transpose: X(iy,ja) += sum_ia Y(iy,ia)*A(ia,ja) + // Parallelize over ia (TeamThreadRange), then vectorize over + // nonzeros and use atomics for updates + Kokkos::parallel_for( + Kokkos::TeamThreadRange(team, nrow_A), + [&](int ia) + { + int istart = d_rowptr(ia); + int iend = d_rowptr(ia + 1); + KokkosScalar yval = yrow(ia); + Kokkos::parallel_for( + Kokkos::ThreadVectorRange(team, istart, iend), + [&](int k) + { + int ja = d_cols(k); + KokkosScalar aij = d_vals(k); + if constexpr (is_complex) + if (isConj) + aij = Kokkos::conj(aij); + KokkosScalar prod = yval * aij; + Kokkos::atomic_add(&x_dev_out(iy, ja), prod); + }); + }); + }); + } + auto xhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace {}, x_dev_out); + for (int iy = 0; iy < nrow_Y; ++iy) { + for (int jx = 0; jx < ncol_X; ++jx) + xout(iy, jx) += static_cast(xhost(iy, jx)); } } From 9519958533c0c8a32fa22f19b9bf9686000f47bb Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 20 Jul 2026 13:08:41 -0400 Subject: [PATCH 08/41] Reapply "Joint A and B" This reverts commit 0cb7206be55dee2d69cde7fdcbbc0256e8b52026. --- dmrg/KronUtil/csr_kron_mult.cpp | 155 +++++++++++++++++++++++------- dmrg/KronUtil/csr_matmul_post.cpp | 2 - 2 files changed, 121 insertions(+), 36 deletions(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 2dcea72ee..93f6a02a2 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -1,6 +1,7 @@ #include "util.h" - +#include #include +#include template void csr_to_den(const PsimagLite::CrsMatrix& a, @@ -278,48 +279,134 @@ void csr_kron_mult_method(const int imethod, * --------------------------------------------- */ +#if 0 std::cerr << "nrow_A: " << nrow_A << ' ' << "nnz_A: " << csr_nnz(a) << ' ' << "nrow_B: " << nrow_B << ' ' << "nnz_B: " << csr_nnz(b) << '\n'; - - int ia = 0; - int ka = 0; - int ib = 0; - int kb = 0; - for (ia = 0; ia < nrow_A; ia++) { - int istarta = a.getRowPtr(ia); - int ienda = a.getRowPtr(ia + 1); - for (ka = istarta; ka < ienda; ka++) { - int ja = a.getCol(ka); - ComplexOrRealType aij = a.getValue(ka); - if (is_complex && isConjTransA) { - aij = PsimagLite::conj(aij); - }; +#endif - for (ib = 0; ib < nrow_B; ib++) { - int istartb = b.getRowPtr(ib); - int iendb = b.getRowPtr(ib + 1); + // Build flat lists of nonzeros for A and B on the host, then copy to device + using ExecutionSpace = Kokkos::DefaultExecutionSpace; + using KokkosScalar = typename PsimagLite::KokkosType::type; - for (kb = istartb; kb < iendb; kb++) { - int jb = b.getCol(kb); - ComplexOrRealType bij = b.getValue(kb); - if (is_complex && isConjTransB) { - bij = PsimagLite::conj(bij); - }; + int nnzA = csr_nnz(a); + int nnzB = csr_nnz(b); - ComplexOrRealType cij = aij * bij; + // host-side temporary arrays + std::vector A_row(nnzA); + std::vector A_col(nnzA); + std::vector A_val(nnzA); + { + int idx = 0; + for (int ia = 0; ia < nrow_A; ++ia) { + int istart = a.getRowPtr(ia); + int iend = a.getRowPtr(ia + 1); + for (int ka = istart; ka < iend; ++ka) { + A_row[idx] = ia; + A_col[idx] = a.getCol(ka); + ComplexOrRealType aval = a.getValue(ka); + if (is_complex && isConjTransA) + aval = PsimagLite::conj(aval); + A_val[idx] = static_cast(aval); + ++idx; + } + } + } - int ix = (isTransB || isConjTransB) ? jb : ib; - int jx = (isTransA || isConjTransA) ? ja : ia; - int iy = (isTransB || isConjTransB) ? ib : jb; - int jy = (isTransA || isConjTransA) ? ia : ja; + std::vector B_row(nnzB); + std::vector B_col(nnzB); + std::vector B_val(nnzB); + { + int idx = 0; + for (int ib = 0; ib < nrow_B; ++ib) { + int istart = b.getRowPtr(ib); + int iend = b.getRowPtr(ib + 1); + for (int kb = istart; kb < iend; ++kb) { + B_row[idx] = ib; + B_col[idx] = b.getCol(kb); + ComplexOrRealType bval = b.getValue(kb); + if (is_complex && isConjTransB) + bval = PsimagLite::conj(bval); + B_val[idx] = static_cast(bval); + ++idx; + } + } + } - xout(ix, jx) += cij * yin(iy, jy); - }; - }; - }; - }; + // create device views + Kokkos::View A_row_h("A_row_h", nnzA); + Kokkos::View A_col_h("A_col_h", nnzA); + Kokkos::View A_val_h("A_val_h", nnzA); + Kokkos::View B_row_h("B_row_h", nnzB); + Kokkos::View B_col_h("B_col_h", nnzB); + Kokkos::View B_val_h("B_val_h", nnzB); + + for (int i = 0; i < nnzA; ++i) { + A_row_h(i) = A_row[i]; + A_col_h(i) = A_col[i]; + A_val_h(i) = A_val[i]; + } + for (int i = 0; i < nnzB; ++i) { + B_row_h(i) = B_row[i]; + B_col_h(i) = B_col[i]; + B_val_h(i) = B_val[i]; + } + + auto A_row_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_row_h); + auto A_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_col_h); + auto A_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_val_h); + auto B_row_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_row_h); + auto B_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_col_h); + auto B_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_val_h); + // device yin and xout + + auto yin_host = Kokkos::View( + reinterpret_cast(&yin(0, 0)), nrow_Y, ncol_Y); + + auto y_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, yin_host); + + auto x_dev = Kokkos::View("x_dev", nrow_X, ncol_X); + Kokkos::deep_copy(x_dev, KokkosScalar(0)); + + const size_t totalPairs = static_cast(nnzA) * static_cast(nnzB); + + Kokkos::parallel_for( + "csr_kron_mult::imethod3_pairs", + Kokkos::RangePolicy(0, totalPairs), + KOKKOS_LAMBDA(const size_t idx) { + const int ia_idx = static_cast(idx / nnzB); + const int ib_idx = static_cast(idx % nnzB); + + int ia = A_row_dev(ia_idx); + int ja = A_col_dev(ia_idx); + KokkosScalar aij = A_val_dev(ia_idx); + + int ib = B_row_dev(ib_idx); + int jb = B_col_dev(ib_idx); + KokkosScalar bij = B_val_dev(ib_idx); + + KokkosScalar cij = aij * bij; + + int ix = (isTransB || isConjTransB) ? jb : ib; + int jx = (isTransA || isConjTransA) ? ja : ia; + int iy = (isTransB || isConjTransB) ? ib : jb; + int jy = (isTransA || isConjTransA) ? ia : ja; + + KokkosScalar prod = cij * y_dev(iy, jy); + Kokkos::atomic_add(&x_dev(ix, jx), prod); + }); + + // copy back and accumulate into xout + + auto xhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace {}, x_dev); + for (int ix = 0; ix < nrow_X; ++ix) { + for (int jx = 0; jx < ncol_X; ++jx) + xout(ix, jx) += static_cast(xhost(ix, jx)); + } }; } diff --git a/dmrg/KronUtil/csr_matmul_post.cpp b/dmrg/KronUtil/csr_matmul_post.cpp index fc895d61f..557677980 100644 --- a/dmrg/KronUtil/csr_matmul_post.cpp +++ b/dmrg/KronUtil/csr_matmul_post.cpp @@ -40,8 +40,6 @@ void csr_matmul_post(char int isConjTranspose = (trans_A == 'C') || (trans_A == 'c'); int isConj = (trans_A == 'Z') || (trans_A == 'z'); - Kokkos::Profiling::ScopedRegion region("PsimagLite::csr_matmul_post"); - using ExecutionSpace = Kokkos::DefaultExecutionSpace; ExecutionSpace exec; From e5f07822eee0ba936e2572e8994a932c45b35774 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 20 Jul 2026 13:56:52 -0400 Subject: [PATCH 09/41] Use MDRangePolicy --- dmrg/KronUtil/csr_kron_mult.cpp | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 93f6a02a2..86d9b6bdf 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -376,11 +376,8 @@ void csr_kron_mult_method(const int imethod, Kokkos::parallel_for( "csr_kron_mult::imethod3_pairs", - Kokkos::RangePolicy(0, totalPairs), - KOKKOS_LAMBDA(const size_t idx) { - const int ia_idx = static_cast(idx / nnzB); - const int ib_idx = static_cast(idx % nnzB); - + Kokkos::MDRangePolicy>({0,0}, {nnzB, nnzA}), + KOKKOS_LAMBDA(const size_t ib_idx, const size_t ia_idx) { int ia = A_row_dev(ia_idx); int ja = A_col_dev(ia_idx); KokkosScalar aij = A_val_dev(ia_idx); From 21bf56c22827271e84f478bd774f66d3b6dcd15b Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 20 Jul 2026 14:26:37 -0400 Subject: [PATCH 10/41] no atomic --- dmrg/KronUtil/csr_kron_mult.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 86d9b6bdf..117ffe548 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -394,7 +394,7 @@ void csr_kron_mult_method(const int imethod, int jy = (isTransA || isConjTransA) ? ia : ja; KokkosScalar prod = cij * y_dev(iy, jy); - Kokkos::atomic_add(&x_dev(ix, jx), prod); + x_dev(ix, jx) = prod; }); // copy back and accumulate into xout From 403302660bb7dcecbb9beb22601d4874148e4dcd Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 20 Jul 2026 14:26:40 -0400 Subject: [PATCH 11/41] Revert "no atomic" This reverts commit 21bf56c22827271e84f478bd774f66d3b6dcd15b. --- dmrg/KronUtil/csr_kron_mult.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 117ffe548..86d9b6bdf 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -394,7 +394,7 @@ void csr_kron_mult_method(const int imethod, int jy = (isTransA || isConjTransA) ? ia : ja; KokkosScalar prod = cij * y_dev(iy, jy); - x_dev(ix, jx) = prod; + Kokkos::atomic_add(&x_dev(ix, jx), prod); }); // copy back and accumulate into xout From ba52627bf0596a030af86177b535ee3d156ba0a3 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 20 Jul 2026 16:36:43 -0400 Subject: [PATCH 12/41] Swap around conj --- dmrg/KronUtil/csr_kron_mult.cpp | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 86d9b6bdf..981e6ad76 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -290,14 +290,15 @@ void csr_kron_mult_method(const int imethod, using ExecutionSpace = Kokkos::DefaultExecutionSpace; using KokkosScalar = typename PsimagLite::KokkosType::type; - int nnzA = csr_nnz(a); - int nnzB = csr_nnz(b); + int nnzA = a.nonZeros(); + int nnzB = b.nonZeros(); // host-side temporary arrays std::vector A_row(nnzA); std::vector A_col(nnzA); std::vector A_val(nnzA); { + Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3::fill_A"); int idx = 0; for (int ia = 0; ia < nrow_A; ++ia) { int istart = a.getRowPtr(ia); @@ -306,8 +307,6 @@ void csr_kron_mult_method(const int imethod, A_row[idx] = ia; A_col[idx] = a.getCol(ka); ComplexOrRealType aval = a.getValue(ka); - if (is_complex && isConjTransA) - aval = PsimagLite::conj(aval); A_val[idx] = static_cast(aval); ++idx; } @@ -318,6 +317,7 @@ void csr_kron_mult_method(const int imethod, std::vector B_col(nnzB); std::vector B_val(nnzB); { + Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3::fill_B"); int idx = 0; for (int ib = 0; ib < nrow_B; ++ib) { int istart = b.getRowPtr(ib); @@ -326,8 +326,6 @@ void csr_kron_mult_method(const int imethod, B_row[idx] = ib; B_col[idx] = b.getCol(kb); ComplexOrRealType bval = b.getValue(kb); - if (is_complex && isConjTransB) - bval = PsimagLite::conj(bval); B_val[idx] = static_cast(bval); ++idx; } @@ -335,13 +333,15 @@ void csr_kron_mult_method(const int imethod, } // create device views - Kokkos::View A_row_h("A_row_h", nnzA); - Kokkos::View A_col_h("A_col_h", nnzA); - Kokkos::View A_val_h("A_val_h", nnzA); - Kokkos::View B_row_h("B_row_h", nnzB); - Kokkos::View B_col_h("B_col_h", nnzB); - Kokkos::View B_val_h("B_val_h", nnzB); + Kokkos::View A_row_h(Kokkos::view_alloc(Kokkos::WithoutInitializing, "A_row_h"), nnzA); + Kokkos::View A_col_h(Kokkos::view_alloc(Kokkos::WithoutInitializing,"A_col_h"), nnzA); + Kokkos::View A_val_h(Kokkos::view_alloc(Kokkos::WithoutInitializing,"A_val_h"), nnzA); + Kokkos::View B_row_h(Kokkos::view_alloc(Kokkos::WithoutInitializing,"B_row_h"), nnzB); + Kokkos::View B_col_h(Kokkos::view_alloc(Kokkos::WithoutInitializing,"B_col_h"), nnzB); + Kokkos::View B_val_h(Kokkos::view_alloc(Kokkos::WithoutInitializing, "B_val_h"), nnzB); +{ + Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3::fill_AB"); for (int i = 0; i < nnzA; ++i) { A_row_h(i) = A_row[i]; A_col_h(i) = A_col[i]; @@ -352,6 +352,7 @@ void csr_kron_mult_method(const int imethod, B_col_h(i) = B_col[i]; B_val_h(i) = B_val[i]; } +} auto A_row_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_row_h); auto A_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_col_h); @@ -381,10 +382,14 @@ void csr_kron_mult_method(const int imethod, int ia = A_row_dev(ia_idx); int ja = A_col_dev(ia_idx); KokkosScalar aij = A_val_dev(ia_idx); + if constexpr(is_complex) if (isConjTransA) + aij = Kokkos::conj(aij); int ib = B_row_dev(ib_idx); int jb = B_col_dev(ib_idx); KokkosScalar bij = B_val_dev(ib_idx); + if constexpr(is_complex) if (isConjTransA) + bij = Kokkos::conj(bij); KokkosScalar cij = aij * bij; From 5aae57a5a5277cdd0b70bdd5f5e56a2cd88e0157 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 20 Jul 2026 17:15:14 -0400 Subject: [PATCH 13/41] Avoid some allocations --- dmrg/KronUtil/csr_kron_mult.cpp | 22 ++++------------------ 1 file changed, 4 insertions(+), 18 deletions(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 981e6ad76..e120c8333 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -295,8 +295,6 @@ void csr_kron_mult_method(const int imethod, // host-side temporary arrays std::vector A_row(nnzA); - std::vector A_col(nnzA); - std::vector A_val(nnzA); { Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3::fill_A"); int idx = 0; @@ -305,17 +303,12 @@ void csr_kron_mult_method(const int imethod, int iend = a.getRowPtr(ia + 1); for (int ka = istart; ka < iend; ++ka) { A_row[idx] = ia; - A_col[idx] = a.getCol(ka); - ComplexOrRealType aval = a.getValue(ka); - A_val[idx] = static_cast(aval); ++idx; } } } std::vector B_row(nnzB); - std::vector B_col(nnzB); - std::vector B_val(nnzB); { Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3::fill_B"); int idx = 0; @@ -324,9 +317,6 @@ void csr_kron_mult_method(const int imethod, int iend = b.getRowPtr(ib + 1); for (int kb = istart; kb < iend; ++kb) { B_row[idx] = ib; - B_col[idx] = b.getCol(kb); - ComplexOrRealType bval = b.getValue(kb); - B_val[idx] = static_cast(bval); ++idx; } } @@ -334,23 +324,19 @@ void csr_kron_mult_method(const int imethod, // create device views Kokkos::View A_row_h(Kokkos::view_alloc(Kokkos::WithoutInitializing, "A_row_h"), nnzA); - Kokkos::View A_col_h(Kokkos::view_alloc(Kokkos::WithoutInitializing,"A_col_h"), nnzA); - Kokkos::View A_val_h(Kokkos::view_alloc(Kokkos::WithoutInitializing,"A_val_h"), nnzA); + Kokkos::View A_col_h(&a.getCol(0), nnzA); + Kokkos::View A_val_h(reinterpret_cast(&a.getValue(0)), nnzA); Kokkos::View B_row_h(Kokkos::view_alloc(Kokkos::WithoutInitializing,"B_row_h"), nnzB); - Kokkos::View B_col_h(Kokkos::view_alloc(Kokkos::WithoutInitializing,"B_col_h"), nnzB); - Kokkos::View B_val_h(Kokkos::view_alloc(Kokkos::WithoutInitializing, "B_val_h"), nnzB); + Kokkos::View B_col_h(&b.getCol(0), nnzB); + Kokkos::View B_val_h(reinterpret_cast(&b.getValue(0)), nnzB); { Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3::fill_AB"); for (int i = 0; i < nnzA; ++i) { A_row_h(i) = A_row[i]; - A_col_h(i) = A_col[i]; - A_val_h(i) = A_val[i]; } for (int i = 0; i < nnzB; ++i) { B_row_h(i) = B_row[i]; - B_col_h(i) = B_col[i]; - B_val_h(i) = B_val[i]; } } From f5b40902ef58d78825811b099f22098b2e8e34cf Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Tue, 21 Jul 2026 11:04:19 -0400 Subject: [PATCH 14/41] Intermediate with TeamPolicy --- dmrg/KronUtil/csr_kron_mult.cpp | 246 ++++++++++++++++++++------------ 1 file changed, 152 insertions(+), 94 deletions(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index e120c8333..9ea72bf43 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -279,122 +279,180 @@ void csr_kron_mult_method(const int imethod, * --------------------------------------------- */ -#if 0 - std::cerr << "nrow_A: " << nrow_A << ' ' - << "nnz_A: " << csr_nnz(a) << ' ' - << "nrow_B: " << nrow_B << ' ' - << "nnz_B: " << csr_nnz(b) << '\n'; -#endif - - // Build flat lists of nonzeros for A and B on the host, then copy to device using ExecutionSpace = Kokkos::DefaultExecutionSpace; using KokkosScalar = typename PsimagLite::KokkosType::type; - int nnzA = a.nonZeros(); - int nnzB = b.nonZeros(); + const int nnzA = a.nonZeros(); + const int nnzB = b.nonZeros(); - // host-side temporary arrays - std::vector A_row(nnzA); - { - Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3::fill_A"); - int idx = 0; + /* + * CPU fast-path: for small problems, direct loops avoid GPU transfer + * overhead (~400+ µs/call on discrete GPUs). The threshold below (tunable) + * targets the crossover where GPU kernel savings outweigh copy latency. + */ + const size_t totalPairs = static_cast(nnzA) * static_cast(nnzB); + static constexpr size_t kGpuThreshold = 100000; + if (totalPairs < kGpuThreshold) { + Kokkos::Profiling::ScopedRegion cpuRegion("PsimgLite::csr_kron_mult_method::imethod3::cpu"); for (int ia = 0; ia < nrow_A; ++ia) { - int istart = a.getRowPtr(ia); - int iend = a.getRowPtr(ia + 1); - for (int ka = istart; ka < iend; ++ka) { - A_row[idx] = ia; - ++idx; - } - } - } - - std::vector B_row(nnzB); - { - Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3::fill_B"); - int idx = 0; - for (int ib = 0; ib < nrow_B; ++ib) { - int istart = b.getRowPtr(ib); - int iend = b.getRowPtr(ib + 1); - for (int kb = istart; kb < iend; ++kb) { - B_row[idx] = ib; - ++idx; + const int istart_a = a.getRowPtr(ia); + const int iend_a = a.getRowPtr(ia + 1); + for (int ka = istart_a; ka < iend_a; ++ka) { + const int ja = a.getCol(ka); + ComplexOrRealType aij = a.getValue(ka); + if constexpr (is_complex) + if (isConjTransA) + aij = PsimagLite::conj(aij); + for (int ib = 0; ib < nrow_B; ++ib) { + const int istart_b = b.getRowPtr(ib); + const int iend_b = b.getRowPtr(ib + 1); + for (int kb = istart_b; kb < iend_b; ++kb) { + const int jb = b.getCol(kb); + ComplexOrRealType bij = b.getValue(kb); + if constexpr (is_complex) + if (isConjTransB) + bij = PsimagLite::conj(bij); + + const int ix = (isTransB || isConjTransB) ? jb : ib; + const int jx = (isTransA || isConjTransA) ? ja : ia; + const int iy = (isTransB || isConjTransB) ? ib : jb; + const int jy = (isTransA || isConjTransA) ? ia : ja; + + xout(ix, jx) += aij * bij * yin(iy, jy); + } + } } } + return; } - // create device views - Kokkos::View A_row_h(Kokkos::view_alloc(Kokkos::WithoutInitializing, "A_row_h"), nnzA); - Kokkos::View A_col_h(&a.getCol(0), nnzA); - Kokkos::View A_val_h(reinterpret_cast(&a.getValue(0)), nnzA); - Kokkos::View B_row_h(Kokkos::view_alloc(Kokkos::WithoutInitializing,"B_row_h"), nnzB); - Kokkos::View B_col_h(&b.getCol(0), nnzB); - Kokkos::View B_val_h(reinterpret_cast(&b.getValue(0)), nnzB); - -{ - Kokkos::Profiling::ScopedRegion region("PsimgLite::csr_kron_mult_method::imethod3::fill_AB"); - for (int i = 0; i < nnzA; ++i) { - A_row_h(i) = A_row[i]; - } - for (int i = 0; i < nnzB; ++i) { - B_row_h(i) = B_row[i]; - } -} - - auto A_row_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_row_h); - auto A_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_col_h); - auto A_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, A_val_h); - auto B_row_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_row_h); - auto B_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_col_h); - auto B_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, B_val_h); - // device yin and xout + /* + * GPU path for large problems. + * + * Use compact rowptr arrays instead of expanded row-index arrays. + * All four are unmanaged host views wrapping the existing CrsMatrix + * storage (zero extra allocation on the host side). + * + * Kernel structure: TeamPolicy(nrow_A, AUTO) + * league rank → A row ia (one team per A row) + * team threads → B row ib (one thread per B row) + * serial inner → A / B nonzeros within (ia, ib) + * + * For the non-transpose case this layout eliminates atomic_add: + * jx = ia — unique per team → no inter-team write conflict + * ix = ib — unique per thread → no intra-team write conflict + * x_dev(ib, ia) is written by exactly one (team, thread) pair. + * + * With LayoutLeft both the x_dev writes and the y_dev reads are + * stride-1 (ib / jb vary across threads, ia / ja fixed per team). + */ + Kokkos::View + A_rowptr_h(&a.getRowPtr(0), nrow_A + 1); + Kokkos::View + A_col_h(&a.getCol(0), nnzA); + Kokkos::View + A_val_h(reinterpret_cast(&a.getValue(0)), nnzA); + + Kokkos::View + B_rowptr_h(&b.getRowPtr(0), nrow_B + 1); + Kokkos::View + B_col_h(&b.getCol(0), nnzB); + Kokkos::View + B_val_h(reinterpret_cast(&b.getValue(0)), nnzB); + + auto A_rowptr_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, A_rowptr_h); + auto A_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, A_col_h); + auto A_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, A_val_h); + auto B_rowptr_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, B_rowptr_h); + auto B_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, B_col_h); + auto B_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, B_val_h); auto yin_host = Kokkos::View( reinterpret_cast(&yin(0, 0)), nrow_Y, ncol_Y); + auto y_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, yin_host); - auto y_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace {}, yin_host); - - auto x_dev = Kokkos::View("x_dev", nrow_X, ncol_X); + auto x_dev = Kokkos::View("x_dev", nrow_X, ncol_X); Kokkos::deep_copy(x_dev, KokkosScalar(0)); - const size_t totalPairs = static_cast(nnzA) * static_cast(nnzB); - - Kokkos::parallel_for( - "csr_kron_mult::imethod3_pairs", - Kokkos::MDRangePolicy>({0,0}, {nnzB, nnzA}), - KOKKOS_LAMBDA(const size_t ib_idx, const size_t ia_idx) { - int ia = A_row_dev(ia_idx); - int ja = A_col_dev(ia_idx); - KokkosScalar aij = A_val_dev(ia_idx); - if constexpr(is_complex) if (isConjTransA) - aij = Kokkos::conj(aij); - - int ib = B_row_dev(ib_idx); - int jb = B_col_dev(ib_idx); - KokkosScalar bij = B_val_dev(ib_idx); - if constexpr(is_complex) if (isConjTransA) - bij = Kokkos::conj(bij); - - KokkosScalar cij = aij * bij; + using TeamPolicy = Kokkos::TeamPolicy; + using TeamMember = typename TeamPolicy::member_type; - int ix = (isTransB || isConjTransB) ? jb : ib; - int jx = (isTransA || isConjTransA) ? ja : ia; - int iy = (isTransB || isConjTransB) ? ib : jb; - int jy = (isTransA || isConjTransA) ? ia : ja; - - KokkosScalar prod = cij * y_dev(iy, jy); - Kokkos::atomic_add(&x_dev(ix, jx), prod); - }); - - // copy back and accumulate into xout + if (!isTransA && !isConjTransA && !isTransB && !isConjTransB) { + /* + * Non-transpose path: ix = ib (unique per thread), jx = ia (unique + * per team). Each (ia, ib) cell of x_dev is touched by exactly one + * (team, thread) pair → plain += without atomic_add is correct. + */ + Kokkos::parallel_for( + "csr_kron_mult::imethod3_nn", + TeamPolicy(nrow_A, Kokkos::AUTO), + KOKKOS_LAMBDA(const TeamMember& team) { + const int ia = team.league_rank(); + const int ka_begin = A_rowptr_dev(ia); + const int ka_end = A_rowptr_dev(ia + 1); + Kokkos::parallel_for( + Kokkos::TeamThreadRange(team, nrow_B), + [=](int ib) { + const int kb_begin = B_rowptr_dev(ib); + const int kb_end = B_rowptr_dev(ib + 1); + KokkosScalar acc = 0; + for (int ka = ka_begin; ka < ka_end; ++ka) { + const int ja = A_col_dev(ka); + KokkosScalar aij = A_val_dev(ka); + for (int kb = kb_begin; kb < kb_end; ++kb) + acc += aij * B_val_dev(kb) * y_dev(B_col_dev(kb), ja); + } + x_dev(ib, ia) += acc; + }); + }); + } else { + /* + * General (transpose / conjugate) path: ix or jx may not be unique + * per (team, thread) pair → atomic_add required. + */ + Kokkos::parallel_for( + "csr_kron_mult::imethod3_gen", + TeamPolicy(nrow_A, Kokkos::AUTO), + KOKKOS_LAMBDA(const TeamMember& team) { + const int ia = team.league_rank(); + const int ka_begin = A_rowptr_dev(ia); + const int ka_end = A_rowptr_dev(ia + 1); + Kokkos::parallel_for( + Kokkos::TeamThreadRange(team, nrow_B), + [=](int ib) { + const int kb_begin = B_rowptr_dev(ib); + const int kb_end = B_rowptr_dev(ib + 1); + for (int ka = ka_begin; ka < ka_end; ++ka) { + const int ja = A_col_dev(ka); + KokkosScalar aij = A_val_dev(ka); + if constexpr (is_complex) + if (isConjTransA) + aij = Kokkos::conj(aij); + for (int kb = kb_begin; kb < kb_end; ++kb) { + const int jb = B_col_dev(kb); + KokkosScalar bij = B_val_dev(kb); + if constexpr (is_complex) + if (isConjTransB) + bij = Kokkos::conj(bij); + const int ix = (isTransB || isConjTransB) ? jb : ib; + const int jx = (isTransA || isConjTransA) ? ja : ia; + const int iy = (isTransB || isConjTransB) ? ib : jb; + const int jy = (isTransA || isConjTransA) ? ia : ja; + Kokkos::atomic_add(&x_dev(ix, jx), aij * bij * y_dev(iy, jy)); + } + } + }); + }); + } - auto xhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace {}, x_dev); - for (int ix = 0; ix < nrow_X; ++ix) { + // copy result back and accumulate into xout + auto xhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, x_dev); + for (int ix = 0; ix < nrow_X; ++ix) for (int jx = 0; jx < ncol_X; ++jx) xout(ix, jx) += static_cast(xhost(ix, jx)); - } }; } From 6d0ba06680e3fe106ea13b1348746ff5a59d035c Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Tue, 21 Jul 2026 11:14:27 -0400 Subject: [PATCH 15/41] Don't use CPU --- dmrg/KronUtil/csr_kron_mult.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 9ea72bf43..2f3db5167 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -292,7 +292,7 @@ void csr_kron_mult_method(const int imethod, */ const size_t totalPairs = static_cast(nnzA) * static_cast(nnzB); static constexpr size_t kGpuThreshold = 100000; - if (totalPairs < kGpuThreshold) { + if (false) { Kokkos::Profiling::ScopedRegion cpuRegion("PsimgLite::csr_kron_mult_method::imethod3::cpu"); for (int ia = 0; ia < nrow_A; ++ia) { const int istart_a = a.getRowPtr(ia); From 757e9f53654d463b706df559feb2f1b40a36d02b Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Thu, 23 Jul 2026 14:08:23 -0400 Subject: [PATCH 16/41] Kokkos batched implementation --- .../MatrixVectorKron/BatchedGemmInclude.hh | 13 +- .../MatrixVectorKron/BatchedGemmKokkos.h | 455 ++++++++++++++++++ 2 files changed, 464 insertions(+), 4 deletions(-) create mode 100644 dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh index 1fa802098..51b6298ab 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh @@ -1,7 +1,10 @@ #ifndef BATCHEDGEMMINCLUDE_HH #define BATCHEDGEMMINCLUDE_HH #include "DMRGConfig.h" -#ifdef PLUGIN_SC +#ifdef KOKKOS_BATCHED +#include "BatchedGemmKokkos.h" +#define BATCHED_GEMM BatchedGemmKokkos +#elif defined(PLUGIN_SC) #include "BatchedGemmPluginSc.h" #define BATCHED_GEMM BatchedGemmPluginSc #else @@ -18,15 +21,17 @@ public: static void failIfNotSupported() { -#ifdef PLUGIN_SC +#if defined(KOKKOS_BATCHED) || defined(PLUGIN_SC) return; #endif - err("BatchedGemm needs -DPLUGIN_SC in Config.make\n"); + err("BatchedGemm needs DMRG_BUILD_BATCHED_KOKKOS=ON or -DPLUGIN_SC\n"); } static std::string info() { -#ifdef PLUGIN_SC +#ifdef KOKKOS_BATCHED + return "KokkosKernels"; +#elif defined(PLUGIN_SC) return "PLUGIN_SC"; #else return ""; diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h new file mode 100644 index 000000000..012096e48 --- /dev/null +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h @@ -0,0 +1,455 @@ +#ifndef BATCHED_GEMM_KOKKOS_H +#define BATCHED_GEMM_KOKKOS_H +// Don't include this file directly; use BatchedGemmInclude.hh + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace Dmrg { + +/* + * Batched GEMM using KokkosKernels with a sparse-batch layout. + * + * Algorithm (mirrors apply_Htarget_sparse_not_ready.cpp): + * + * Pass 1: For each non-zero triple (ipatch, jpatch, ioperator): + * BXbatch[ip][:, colBX:colBX+lps[jp]] = + * Bbatch[ip][:, colB:colB+rps[jp]] * X[jp] + * + * Pass 2: For each ipatch: + * Y[ip] = BXbatch[ip] * Abatch[ip]^T + * + * Abatch[ip] = left-operator (xc) matrices packed column-by-column (ldA padded). + * Bbatch[ip] = right-operator (yc) matrices packed column-by-column (ldB padded). + * BXbatch[ip] is the intermediate work buffer (ldB padded, same col-width as Abatch). + * + * All Abatch and Bbatch data live on the device for the lifetime of the object. + * BXbatch, vin, vout are (re-)written on each matrixVector call. + */ + +template +class BatchedGemmKokkos { + + using ArrayOfMatStructType = typename InitKronType::ArrayOfMatStructType; + using GenIjPatchType = typename InitKronType::GenIjPatchType; + using MatrixDenseOrSparseType = typename ArrayOfMatStructType::MatrixDenseOrSparseType; + using VectorType = typename MatrixDenseOrSparseType::VectorType; + using ComplexOrRealType = typename VectorType::value_type; + using MatrixType = PsimagLite::Matrix; + using VectorMatrixType = typename PsimagLite::Vector::Type; + using VectorSizeType = PsimagLite::Vector::Type; + + // Kokkos types + using KS = typename PsimagLite::KokkosType::type; + using DevExecSpace = Kokkos::DefaultExecutionSpace; + using DevMemSpace = typename DevExecSpace::memory_space; + using DevScalView = Kokkos::View; + + // Compact struct holding all parameters for one batched GEMM call. + // Stored in device memory and accessed inside the kernel. + struct GemmArgs { + int m, n, k; // GEMM dimensions + int lda, ldb, ldc; // leading dimensions (column-major strides) + long long a_off, b_off, c_off; // element offsets into device flat arrays + }; + + using DevArgsView = Kokkos::View; + + static const int ialign_ = 32; + +public: + + BatchedGemmKokkos(const InitKronType& initKron) + : initKron_(initKron) + , progress_("BatchedGemmKokkos") + { + if (!enabled()) + return; + setup_(); + } + + ~BatchedGemmKokkos() + { + for (SizeType i = 0; i < garbage_.size(); ++i) { + delete garbage_[i]; + garbage_[i] = nullptr; + } + } + + bool enabled() const { return initKron_.params().options.isSet("BatchedGemm"); } + + void matrixVector(VectorType& vout, const VectorType& vin) const + { + assert(enabled()); + + Kokkos::Profiling::ScopedRegion region("BatchedGemmKokkos::matrixVector"); + + const SizeType totalXY = vin.size(); + assert(vout.size() == totalXY); + assert(d_vin_.extent(0) == totalXY); + + // --- H2D: copy vin to device ----------------------------------------------- + { + using HV = Kokkos::View; + HV hv(reinterpret_cast(vin.data()), totalXY); + Kokkos::deep_copy(d_vin_, hv); + } + + // --- Zero output and work buffers ------------------------------------------- + Kokkos::deep_copy(d_vout_, KS(0)); + Kokkos::deep_copy(d_flatBXbatch_, KS(0)); + + DevExecSpace exec; + + // --- Pass 1: BXbatch[ip] = Bbatch[ip] * X[jp] (NoT x NoT) ---------------- + { + const DevScalView flatBbatch = d_flatBbatch_; + const DevScalView vin_dev = d_vin_; + const DevScalView flatBXbatch = d_flatBXbatch_; + const DevArgsView args = d_pass1_; + + using MemberType = typename Kokkos::TeamPolicy::member_type; + Kokkos::parallel_for( + "BatchedGemmKokkos_Pass1", + Kokkos::TeamPolicy(exec, static_cast(nbatch1_), + Kokkos::AUTO, Kokkos::AUTO), + KOKKOS_LAMBDA(const MemberType& member) { + const int i = member.league_rank(); + const GemmArgs& ag = args(i); + + using UV = Kokkos::View; + + // A = Bbatch[ip] block: (lda, k) padded -> subview (m, k) active rows + UV A_full(flatBbatch.data() + ag.a_off, ag.lda, ag.k); + auto A = Kokkos::subview(A_full, + Kokkos::make_pair(0, ag.m), Kokkos::ALL()); + + // B = X[jp]: (k, n) exact -- no padding for X slice + UV B(vin_dev.data() + ag.b_off, ag.ldb, ag.n); + + // C = BXbatch[ip] block: (ldc, n) padded -> subview (m, n) active rows + UV C_full(flatBXbatch.data() + ag.c_off, ag.ldc, ag.n); + auto C = Kokkos::subview(C_full, + Kokkos::make_pair(0, ag.m), Kokkos::ALL()); + + KokkosBatched::TeamVectorGemm< + MemberType, + KokkosBatched::Trans::NoTranspose, + KokkosBatched::Trans::NoTranspose, + KokkosBatched::Algo::Gemm::Unblocked>:: + invoke(member, KS(1), A, B, KS(0), C); + }); + } + + exec.fence(); + + // --- Pass 2: Y[ip] = BXbatch[ip] * Abatch[ip]^T (NoT x T) --------------- + { + const DevScalView flatBXbatch = d_flatBXbatch_; + const DevScalView flatAbatch = d_flatAbatch_; + const DevScalView vout_dev = d_vout_; + const DevArgsView args = d_pass2_; + + using MemberType = typename Kokkos::TeamPolicy::member_type; + Kokkos::parallel_for( + "BatchedGemmKokkos_Pass2", + Kokkos::TeamPolicy(exec, static_cast(nbatch2_), + Kokkos::AUTO, Kokkos::AUTO), + KOKKOS_LAMBDA(const MemberType& member) { + const int i = member.league_rank(); + const GemmArgs& ag = args(i); + + if (ag.k == 0) + return; // no connections for this ipatch; Y[ip] already zero + + using UV = Kokkos::View; + + // A = BXbatch[ip]: (lda, k) padded -> subview (m, k) + UV A_full(flatBXbatch.data() + ag.a_off, ag.lda, ag.k); + auto A = Kokkos::subview(A_full, + Kokkos::make_pair(0, ag.m), Kokkos::ALL()); + + // B = Abatch[ip]: (ldb, k) padded -> subview (n, k); will be transposed + UV B_full(flatAbatch.data() + ag.b_off, ag.ldb, ag.k); + auto B = Kokkos::subview(B_full, + Kokkos::make_pair(0, ag.n), Kokkos::ALL()); + + // C = Y[ip]: (ldc, n) exact -- no padding for output + UV C(vout_dev.data() + ag.c_off, ag.ldc, ag.n); + + KokkosBatched::TeamVectorGemm< + MemberType, + KokkosBatched::Trans::NoTranspose, + KokkosBatched::Trans::Transpose, + KokkosBatched::Algo::Gemm::Unblocked>:: + invoke(member, KS(1), A, B, KS(0), C); + }); + } + + exec.fence(); + + // --- D2H: copy vout back to host ------------------------------------------- + { + using HV = Kokkos::View; + HV hv(reinterpret_cast(vout.data()), totalXY); + Kokkos::deep_copy(hv, d_vout_); + } + } + +private: + + static int iceil(int x, int n) { return (x + n - 1) / n; } + + // Dense matrix reference: pointer + source dimensions (no padding). + // For sparse matrices, the expansion is heap-allocated and owned by garbage_. + struct MatRef { + ComplexOrRealType* ptr; + int rows; // = source leading dimension + int cols; + }; + + MatRef getMatRef(const MatrixDenseOrSparseType& mat) + { + if (mat.isZero()) + return { nullptr, 0, 0 }; + if (!mat.isDense()) { + MatrixType* dense = new MatrixType(); + crsMatrixToFullMatrix(*dense, mat.sparse()); + garbage_.push_back(dense); + return { &(*dense)(0, 0), + static_cast(dense->rows()), + static_cast(dense->cols()) }; + } + const MatrixType& d = mat.dense(); + return { const_cast(&d(0, 0)), + static_cast(d.rows()), + static_cast(d.cols()) }; + } + + // Column-major copy: src(m, n) with lds -> dst(m, n) with ldd. + static void lacpy(const ComplexOrRealType* src, int m, int n, int lds, + ComplexOrRealType* dst, int ldd) + { + for (int j = 0; j < n; ++j) + for (int i = 0; i < m; ++i) + dst[i + static_cast(j) * ldd] + = src[i + static_cast(j) * lds]; + } + + void setup_() + { + Kokkos::Profiling::ScopedRegion region("BatchedGemmKokkos::setup"); + + const SizeType npatches = initKron_.numberOfPatches(InitKronType::OLD); + const SizeType noperator = initKron_.connections(); + + // Per-patch sizes and padded leading dimensions. + VectorSizeType lps(npatches, 0); // left patch size + VectorSizeType rps(npatches, 0); // right patch size + VectorSizeType ldA(npatches, 0); // padded lps -> Abatch leading dim + VectorSizeType ldB(npatches, 0); // padded rps -> Bbatch / BXbatch leading dim + VectorSizeType xyStart(npatches + 1, 0); // patch start offset in vin/vout + + for (SizeType ip = 0; ip < npatches; ++ip) { + const SizeType lg = initKron_.patch(InitKronType::NEW, GenIjPatchType::LEFT)[ip]; + const SizeType rg = initKron_.patch(InitKronType::NEW, GenIjPatchType::RIGHT)[ip]; + const int L1 = initKron_.lrs(InitKronType::NEW).left().partition(lg); + const int L2 = initKron_.lrs(InitKronType::NEW).left().partition(lg + 1); + const int R1 = initKron_.lrs(InitKronType::NEW).right().partition(rg); + const int R2 = initKron_.lrs(InitKronType::NEW).right().partition(rg + 1); + lps[ip] = static_cast(L2 - L1); + rps[ip] = static_cast(R2 - R1); + ldA[ip] = static_cast(ialign_ * iceil(static_cast(lps[ip]), ialign_)); + ldB[ip] = static_cast(ialign_ * iceil(static_cast(rps[ip]), ialign_)); + } + + xyStart[0] = 0; + for (SizeType ip = 0; ip < npatches; ++ip) + xyStart[ip + 1] = xyStart[ip] + lps[ip] * rps[ip]; + + // Column widths: + // AbatchCols[ip] = sum of lps[jp] over all non-zero (jp, k) connections for ip + // BbatchCols[ip] = sum of rps[jp] over all non-zero (jp, k) connections for ip + VectorSizeType AbatchCols(npatches, 0); + VectorSizeType BbatchCols(npatches, 0); + + for (SizeType ip = 0; ip < npatches; ++ip) { + for (SizeType jp = 0; jp < npatches; ++jp) { + for (SizeType k = 0; k < noperator; ++k) { + const MatrixDenseOrSparseType* Ap = initKron_.xc(k)(ip, jp); + const MatrixDenseOrSparseType* Bp = initKron_.yc(k)(ip, jp); + if (!Ap || !Bp) continue; + if (Ap->isZero() || Bp->isZero()) continue; + AbatchCols[ip] += lps[jp]; + BbatchCols[ip] += rps[jp]; + } + } + } + + // Flat offsets: each patch ip occupies a contiguous slice. + VectorSizeType AbatchOff(npatches + 1, 0); + VectorSizeType BbatchOff(npatches + 1, 0); + VectorSizeType BXbatchOff(npatches + 1, 0); // ldB rows x AbatchCols columns + + for (SizeType ip = 0; ip < npatches; ++ip) { + AbatchOff[ip + 1] = AbatchOff[ip] + ldA[ip] * AbatchCols[ip]; + BbatchOff[ip + 1] = BbatchOff[ip] + ldB[ip] * BbatchCols[ip]; + BXbatchOff[ip + 1] = BXbatchOff[ip] + ldB[ip] * AbatchCols[ip]; + } + + const SizeType totalAbatch = AbatchOff[npatches]; + const SizeType totalBbatch = BbatchOff[npatches]; + const SizeType totalBXbatch = BXbatchOff[npatches]; + + // Host packing buffers (zeroed). + std::vector h_flatAbatch(totalAbatch, ComplexOrRealType(0)); + std::vector h_flatBbatch(totalBbatch, ComplexOrRealType(0)); + + // Build GEMM arg lists while packing matrices. + std::vector pass1_args; + std::vector pass2_args; + pass2_args.reserve(npatches); + + for (SizeType ip = 0; ip < npatches; ++ip) { + long long colA = 0; // column cursor in Abatch[ip] (= BXbatch[ip]) + long long colB = 0; // column cursor in Bbatch[ip] + + for (SizeType jp = 0; jp < npatches; ++jp) { + for (SizeType k = 0; k < noperator; ++k) { + const MatrixDenseOrSparseType* Amat = initKron_.xc(k)(ip, jp); + const MatrixDenseOrSparseType* Bmat = initKron_.yc(k)(ip, jp); + if (!Amat || !Bmat) continue; + if (Amat->isZero() || Bmat->isZero()) continue; + + MatRef Aref = getMatRef(*Amat); // left operator (xc) + MatRef Bref = getMatRef(*Bmat); // right operator (yc) + assert(Aref.ptr && Bref.ptr); + + const int mA = static_cast(lps[ip]); + const int nAk = static_cast(lps[jp]); // A cols = BX cols per op + const int mB = static_cast(rps[ip]); + const int nBk = static_cast(rps[jp]); // B cols = X rows + + // Pack A (left operator) into Abatch[ip] at column colA + lacpy(Aref.ptr, mA, nAk, Aref.rows, + h_flatAbatch.data() + + AbatchOff[ip] + colA * static_cast(ldA[ip]), + static_cast(ldA[ip])); + + // Pack B (right operator) into Bbatch[ip] at column colB + lacpy(Bref.ptr, mB, nBk, Bref.rows, + h_flatBbatch.data() + + BbatchOff[ip] + colB * static_cast(ldB[ip]), + static_cast(ldB[ip])); + + // Pass 1 GEMM: C(mB, nAk) = Bbatch_block(mB, nBk) * X_jp(nBk, nAk) + GemmArgs a1; + a1.m = mB; + a1.n = nAk; + a1.k = nBk; + a1.lda = static_cast(ldB[ip]); // Bbatch leading dim + a1.ldb = nBk; // X[jp] no padding: ld = rps[jp] + a1.ldc = static_cast(ldB[ip]); // BXbatch leading dim + a1.a_off = static_cast(BbatchOff[ip]) + + colB * static_cast(ldB[ip]); + a1.b_off = static_cast(xyStart[jp]); + a1.c_off = static_cast(BXbatchOff[ip]) + + colA * static_cast(ldB[ip]); + pass1_args.push_back(a1); + + colA += nAk; + colB += nBk; + } + } + + // Pass 2 GEMM: Y[ip](mB, mA) = BXbatch[ip](mB, k) * Abatch[ip](mA, k)^T + const int totalCols = static_cast(colA); // = AbatchCols[ip] + GemmArgs a2; + a2.m = static_cast(rps[ip]); + a2.n = static_cast(lps[ip]); + a2.k = totalCols; + a2.lda = static_cast(ldB[ip]); // BXbatch leading dim + a2.ldb = static_cast(ldA[ip]); // Abatch leading dim + a2.ldc = static_cast(rps[ip]); // Y[ip] no padding: ld = rps[ip] + a2.a_off = static_cast(BXbatchOff[ip]); + a2.b_off = static_cast(AbatchOff[ip]); + a2.c_off = static_cast(xyStart[ip]); + pass2_args.push_back(a2); + } + + nbatch1_ = pass1_args.size(); + nbatch2_ = pass2_args.size(); // == npatches + + // Allocate device arrays and upload. + d_flatAbatch_ = DevScalView("d_flatAbatch", totalAbatch ? totalAbatch : 1); + d_flatBbatch_ = DevScalView("d_flatBbatch", totalBbatch ? totalBbatch : 1); + d_flatBXbatch_ = DevScalView("d_flatBXbatch", totalBXbatch ? totalBXbatch : 1); + d_vin_ = DevScalView("d_vin", xyStart[npatches] ? xyStart[npatches] : 1); + d_vout_ = DevScalView("d_vout", xyStart[npatches] ? xyStart[npatches] : 1); + + { + auto hA = Kokkos::create_mirror_view(d_flatAbatch_); + auto hB = Kokkos::create_mirror_view(d_flatBbatch_); + for (SizeType i = 0; i < totalAbatch; ++i) + hA(i) = *reinterpret_cast(&h_flatAbatch[i]); + for (SizeType i = 0; i < totalBbatch; ++i) + hB(i) = *reinterpret_cast(&h_flatBbatch[i]); + Kokkos::deep_copy(d_flatAbatch_, hA); + Kokkos::deep_copy(d_flatBbatch_, hB); + } + + d_pass1_ = DevArgsView("d_pass1", nbatch1_ ? nbatch1_ : 1); + d_pass2_ = DevArgsView("d_pass2", nbatch2_ ? nbatch2_ : 1); + { + auto h1 = Kokkos::create_mirror_view(d_pass1_); + auto h2 = Kokkos::create_mirror_view(d_pass2_); + for (SizeType i = 0; i < nbatch1_; ++i) h1(i) = pass1_args[i]; + for (SizeType i = 0; i < nbatch2_; ++i) h2(i) = pass2_args[i]; + Kokkos::deep_copy(d_pass1_, h1); + Kokkos::deep_copy(d_pass2_, h2); + } + + { + PsimagLite::OstringStream msg(std::cout.precision()); + msg() << "setup done: npatches=" << npatches + << " noperator=" << noperator + << " pass1_batches=" << nbatch1_ + << " pass2_batches=" << nbatch2_ + << " Abatch=" << totalAbatch << "elems" + << " Bbatch=" << totalBbatch << "elems" + << " BXbatch=" << totalBXbatch << "elems"; + progress_.printline(msg, std::cout); + } + } + + // ----------------------------------------------------------------------- + + const InitKronType& initKron_; + PsimagLite::ProgressIndicator progress_; + mutable VectorMatrixType garbage_; // owns sparse->dense expansions (freed in dtor) + + SizeType nbatch1_ = 0; // number of pass-1 GEMMs + SizeType nbatch2_ = 0; // number of pass-2 GEMMs (== npatches) + + DevScalView d_flatAbatch_; // left-operator matrices, device, persistent + DevScalView d_flatBbatch_; // right-operator matrices, device, persistent + mutable DevScalView d_flatBXbatch_; // intermediate BX work buffer, per-call + mutable DevScalView d_vin_; // input vector on device, per-call + mutable DevScalView d_vout_; // output vector on device, per-call + DevArgsView d_pass1_; // pass-1 GEMM parameters, persistent + DevArgsView d_pass2_; // pass-2 GEMM parameters, persistent +}; + +} // namespace Dmrg +#endif // BATCHED_GEMM_KOKKOS_H From 46212595f5df83f66842e9b19960a4e87087a3d9 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Thu, 23 Jul 2026 14:52:49 -0400 Subject: [PATCH 17/41] Optimixe more --- CMakeLists.txt | 5 + TestSuite/inputs/input342.ain | 3 +- dmrg/CMakeLists.txt | 26 +- .../MatrixVectorKron/BatchedGemmKokkos.h | 256 ++++++------------ dmrg/GPUPlugin/DMRGConfig.h.in | 1 + dmrg/KronUtil/csr_kron_mult.cpp | 69 ++--- 6 files changed, 140 insertions(+), 220 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 3747e5fa2..fda339d21 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -63,6 +63,11 @@ option(DMRG_BUILD_BATCHED "Enable building with Batched support" ON) include(CMakeDependentOption) cmake_dependent_option(DMRG_BUILD_BATCHED_MAGMA "Enable building with Magma" OFF "DMRG_BUILD_BATCHED" OFF) +# KokkosKernels batched GEMM: portable GPU path, no separate GPU plugin needed. +# Kokkos and KokkosKernels are already pulled in via PsimagLite. +cmake_dependent_option(DMRG_BUILD_BATCHED_KOKKOS + "Use KokkosKernels for batched GEMM (GPU-portable, no GPU plugin required)" + OFF "DMRG_BUILD_BATCHED" OFF) add_subdirectory(util) diff --git a/TestSuite/inputs/input342.ain b/TestSuite/inputs/input342.ain index 247ddffb6..8a2b4c2de 100644 --- a/TestSuite/inputs/input342.ain +++ b/TestSuite/inputs/input342.ain @@ -47,7 +47,8 @@ potentialV=[0.,...x32]; Model=FeAsBasedScExtended; FeAsMode=INT_PAPER33; Orbitals=2; -SolverOptions=twositedmrg,minimizedisk; +SolverOptions=BatchedGemm,twositedmrg,minimizedisk; +integer DenseSparseThreshold=0; Version=61289987cdd32b8485213ac0415baf4e9cd16432; OutputFile=data342; InfiniteLoopKeptStates=20; diff --git a/dmrg/CMakeLists.txt b/dmrg/CMakeLists.txt index 4c8035dea..1a3c0f765 100644 --- a/dmrg/CMakeLists.txt +++ b/dmrg/CMakeLists.txt @@ -6,11 +6,21 @@ add_custom_command( VERBATIM) if(DMRG_BUILD_BATCHED) - add_subdirectory(GPUPlugin) - set(PLUGIN_SC ON) - set(FpType "double") - if(DMRG_BUILD_BATCHED_MAGMA) - set(USE_MAGMA ON) + if(DMRG_BUILD_BATCHED_KOKKOS) + # KokkosKernels path: portable GPU implementation, no separate GPU plugin. + # Kokkos and KokkosKernels are already a dependency of psimaglite. + set(KOKKOS_BATCHED ON) + set(FpType "double") + message(STATUS "BatchedGemm: using KokkosKernels (DMRG_BUILD_BATCHED_KOKKOS=ON)") + else() + # Original GPU plugin path (optionally with MAGMA). + add_subdirectory(GPUPlugin) + set(PLUGIN_SC ON) + set(FpType "double") + if(DMRG_BUILD_BATCHED_MAGMA) + set(USE_MAGMA ON) + endif() + message(STATUS "BatchedGemm: using GPU plugin (PLUGIN_SC)") endif() endif() @@ -25,14 +35,14 @@ add_library(dmrgpp_utils ProgramGlobals.cpp Provenance.cpp Utils.cpp Su2Related. target_sources(dmrgpp_utils PUBLIC ${CMAKE_CURRENT_BINARY_DIR}/GitRevision.h) target_include_directories(dmrgpp_utils PUBLIC ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR}/GPUPlugin Engine) target_link_libraries(dmrgpp_utils PUBLIC psimaglite::psimaglite) -if(DMRG_BUILD_BATCHED) +if(DMRG_BUILD_BATCHED AND NOT DMRG_BUILD_BATCHED_KOKKOS) target_link_libraries(dmrgpp_utils PUBLIC gpuplugin) endif() # DMRG runner add_library(dmrg_runner Engine/DmrgRunner.cpp) target_link_libraries(dmrg_runner PUBLIC dmrgpp_utils kronutil) -if(DMRG_BUILD_BATCHED) +if(DMRG_BUILD_BATCHED AND NOT DMRG_BUILD_BATCHED_KOKKOS) target_link_libraries(dmrg_runner PUBLIC gpuplugin) endif() @@ -170,6 +180,8 @@ if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") # The following tests are expensive to add_test(NAME input341 COMMAND ./dmrg -f ${PATH_TO_INPUTS}/input341.ain) # energy -0.093321 add_lowest_eigenvalue_check_test(output341 input341 -0.093321 runForinput341.cout) add_test(NAME input342 COMMAND ./dmrg -f ${PATH_TO_INPUTS}/input342.ain) # energy -0.0933211 + set_tests_properties(input342 PROPERTIES + ENVIRONMENT "KOKKOS_TOOLS_LIBS=$ENV{HOME}/kokkos-tools/build/profiling/space-time-stack/libkp_space_time_stack.so") add_lowest_eigenvalue_check_test(output342 input342 -0.0933211 runForinput342.cout) endif() add_test(NAME input351 COMMAND ./dmrg -f ${PATH_TO_INPUTS}/input351.ain) # energy -0.408019 diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h index 012096e48..dfafc476e 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h @@ -11,8 +11,7 @@ #include #include -#include -#include +#include namespace Dmrg { @@ -25,12 +24,16 @@ namespace Dmrg { * BXbatch[ip][:, colBX:colBX+lps[jp]] = * Bbatch[ip][:, colB:colB+rps[jp]] * X[jp] * - * Pass 2: For each ipatch: + * Pass 2: For each ipatch (with connections): * Y[ip] = BXbatch[ip] * Abatch[ip]^T * - * Abatch[ip] = left-operator (xc) matrices packed column-by-column (ldA padded). - * Bbatch[ip] = right-operator (yc) matrices packed column-by-column (ldB padded). - * BXbatch[ip] is the intermediate work buffer (ldB padded, same col-width as Abatch). + * Abatch[ip] = left-operator (xc) matrices packed column-by-column (no padding). + * Bbatch[ip] = right-operator (yc) matrices packed column-by-column (no padding). + * BXbatch[ip] is the intermediate work buffer (no padding, same col-width as Abatch). + * + * No padding is used so that unmanaged Kokkos::LayoutLeft views can be created + * directly from device pointers, enabling dispatch to rocBLAS/cuBLAS via + * KokkosBlas::gemm. * * All Abatch and Bbatch data live on the device for the lifetime of the object. * BXbatch, vin, vout are (re-)written on each matrixVector call. @@ -54,18 +57,13 @@ class BatchedGemmKokkos { using DevMemSpace = typename DevExecSpace::memory_space; using DevScalView = Kokkos::View; - // Compact struct holding all parameters for one batched GEMM call. - // Stored in device memory and accessed inside the kernel. + // Parameters for one KokkosBlas::gemm call. + // With no padding: leading dims = m (for A,C) or k (for B in pass1) or n (for B in pass2). struct GemmArgs { int m, n, k; // GEMM dimensions - int lda, ldb, ldc; // leading dimensions (column-major strides) long long a_off, b_off, c_off; // element offsets into device flat arrays }; - using DevArgsView = Kokkos::View; - - static const int ialign_ = 32; - public: BatchedGemmKokkos(const InitKronType& initKron) @@ -104,97 +102,40 @@ class BatchedGemmKokkos { Kokkos::deep_copy(d_vin_, hv); } - // --- Zero output and work buffers ------------------------------------------- + // Zero d_vout_ so patches with no connections produce zero output. + // d_flatBXbatch_ is NOT zeroed: pass1 uses beta=0 (overwrites all used blocks). Kokkos::deep_copy(d_vout_, KS(0)); - Kokkos::deep_copy(d_flatBXbatch_, KS(0)); DevExecSpace exec; - // --- Pass 1: BXbatch[ip] = Bbatch[ip] * X[jp] (NoT x NoT) ---------------- - { - const DevScalView flatBbatch = d_flatBbatch_; - const DevScalView vin_dev = d_vin_; - const DevScalView flatBXbatch = d_flatBXbatch_; - const DevArgsView args = d_pass1_; - - using MemberType = typename Kokkos::TeamPolicy::member_type; - Kokkos::parallel_for( - "BatchedGemmKokkos_Pass1", - Kokkos::TeamPolicy(exec, static_cast(nbatch1_), - Kokkos::AUTO, Kokkos::AUTO), - KOKKOS_LAMBDA(const MemberType& member) { - const int i = member.league_rank(); - const GemmArgs& ag = args(i); - - using UV = Kokkos::View; - - // A = Bbatch[ip] block: (lda, k) padded -> subview (m, k) active rows - UV A_full(flatBbatch.data() + ag.a_off, ag.lda, ag.k); - auto A = Kokkos::subview(A_full, - Kokkos::make_pair(0, ag.m), Kokkos::ALL()); - - // B = X[jp]: (k, n) exact -- no padding for X slice - UV B(vin_dev.data() + ag.b_off, ag.ldb, ag.n); - - // C = BXbatch[ip] block: (ldc, n) padded -> subview (m, n) active rows - UV C_full(flatBXbatch.data() + ag.c_off, ag.ldc, ag.n); - auto C = Kokkos::subview(C_full, - Kokkos::make_pair(0, ag.m), Kokkos::ALL()); - - KokkosBatched::TeamVectorGemm< - MemberType, - KokkosBatched::Trans::NoTranspose, - KokkosBatched::Trans::NoTranspose, - KokkosBatched::Algo::Gemm::Unblocked>:: - invoke(member, KS(1), A, B, KS(0), C); - }); + using UV = Kokkos::View; + + // --- Pass 1: BXbatch[ip] block = Bbatch[ip] block * X[jp] (N x N) --------- + // Each call: C(m, n) = A(m, k) * B(k, n), beta=0 overwrites C. + // A = Bbatch block (m=rps[ip], k=rps[jp]), ld = rps[ip] + // B = X[jp] (k=rps[jp], n=lps[jp]), ld = rps[jp] + // C = BXbatch block (m=rps[ip], n=lps[jp]), ld = rps[ip] + for (const GemmArgs& ag : h_pass1_) { + UV A(d_flatBbatch_.data() + ag.a_off, ag.m, ag.k); + UV B(d_vin_.data() + ag.b_off, ag.k, ag.n); + UV C(d_flatBXbatch_.data() + ag.c_off, ag.m, ag.n); + KokkosBlas::gemm(exec, "N", "N", KS(1), A, B, KS(0), C); } exec.fence(); - // --- Pass 2: Y[ip] = BXbatch[ip] * Abatch[ip]^T (NoT x T) --------------- - { - const DevScalView flatBXbatch = d_flatBXbatch_; - const DevScalView flatAbatch = d_flatAbatch_; - const DevScalView vout_dev = d_vout_; - const DevArgsView args = d_pass2_; - - using MemberType = typename Kokkos::TeamPolicy::member_type; - Kokkos::parallel_for( - "BatchedGemmKokkos_Pass2", - Kokkos::TeamPolicy(exec, static_cast(nbatch2_), - Kokkos::AUTO, Kokkos::AUTO), - KOKKOS_LAMBDA(const MemberType& member) { - const int i = member.league_rank(); - const GemmArgs& ag = args(i); - - if (ag.k == 0) - return; // no connections for this ipatch; Y[ip] already zero - - using UV = Kokkos::View; - - // A = BXbatch[ip]: (lda, k) padded -> subview (m, k) - UV A_full(flatBXbatch.data() + ag.a_off, ag.lda, ag.k); - auto A = Kokkos::subview(A_full, - Kokkos::make_pair(0, ag.m), Kokkos::ALL()); - - // B = Abatch[ip]: (ldb, k) padded -> subview (n, k); will be transposed - UV B_full(flatAbatch.data() + ag.b_off, ag.ldb, ag.k); - auto B = Kokkos::subview(B_full, - Kokkos::make_pair(0, ag.n), Kokkos::ALL()); - - // C = Y[ip]: (ldc, n) exact -- no padding for output - UV C(vout_dev.data() + ag.c_off, ag.ldc, ag.n); - - KokkosBatched::TeamVectorGemm< - MemberType, - KokkosBatched::Trans::NoTranspose, - KokkosBatched::Trans::Transpose, - KokkosBatched::Algo::Gemm::Unblocked>:: - invoke(member, KS(1), A, B, KS(0), C); - }); + // --- Pass 2: Y[ip] = BXbatch[ip] * Abatch[ip]^T (N x T) ------------------ + // Each call: C(m, n) = A(m, k) * B(n, k)^T, beta=0 overwrites C. + // A = BXbatch[ip] (m=rps[ip], k=AbatchCols[ip]), ld = rps[ip] + // B = Abatch[ip] (n=lps[ip], k=AbatchCols[ip]), ld = lps[ip] (transposed) + // C = Y[ip] (m=rps[ip], n=lps[ip]), ld = rps[ip] + for (const GemmArgs& ag : h_pass2_) { + if (ag.k == 0) + continue; // no connections for this ipatch; Y[ip] stays zero + UV A(d_flatBXbatch_.data() + ag.a_off, ag.m, ag.k); + UV B(d_flatAbatch_.data() + ag.b_off, ag.n, ag.k); + UV C(d_vout_.data() + ag.c_off, ag.m, ag.n); + KokkosBlas::gemm(exec, "N", "T", KS(1), A, B, KS(0), C); } exec.fence(); @@ -209,8 +150,6 @@ class BatchedGemmKokkos { private: - static int iceil(int x, int n) { return (x + n - 1) / n; } - // Dense matrix reference: pointer + source dimensions (no padding). // For sparse matrices, the expansion is heap-allocated and owned by garbage_. struct MatRef { @@ -254,11 +193,9 @@ class BatchedGemmKokkos { const SizeType npatches = initKron_.numberOfPatches(InitKronType::OLD); const SizeType noperator = initKron_.connections(); - // Per-patch sizes and padded leading dimensions. - VectorSizeType lps(npatches, 0); // left patch size - VectorSizeType rps(npatches, 0); // right patch size - VectorSizeType ldA(npatches, 0); // padded lps -> Abatch leading dim - VectorSizeType ldB(npatches, 0); // padded rps -> Bbatch / BXbatch leading dim + // Per-patch sizes — NO padding: leading dim == patch size. + VectorSizeType lps(npatches, 0); // left patch size + VectorSizeType rps(npatches, 0); // right patch size VectorSizeType xyStart(npatches + 1, 0); // patch start offset in vin/vout for (SizeType ip = 0; ip < npatches; ++ip) { @@ -270,15 +207,13 @@ class BatchedGemmKokkos { const int R2 = initKron_.lrs(InitKronType::NEW).right().partition(rg + 1); lps[ip] = static_cast(L2 - L1); rps[ip] = static_cast(R2 - R1); - ldA[ip] = static_cast(ialign_ * iceil(static_cast(lps[ip]), ialign_)); - ldB[ip] = static_cast(ialign_ * iceil(static_cast(rps[ip]), ialign_)); } xyStart[0] = 0; for (SizeType ip = 0; ip < npatches; ++ip) xyStart[ip + 1] = xyStart[ip] + lps[ip] * rps[ip]; - // Column widths: + // Column widths for each patch's batch block. // AbatchCols[ip] = sum of lps[jp] over all non-zero (jp, k) connections for ip // BbatchCols[ip] = sum of rps[jp] over all non-zero (jp, k) connections for ip VectorSizeType AbatchCols(npatches, 0); @@ -297,15 +232,18 @@ class BatchedGemmKokkos { } } - // Flat offsets: each patch ip occupies a contiguous slice. + // Flat offsets: no padding, ld[ip] = rps[ip] or lps[ip]. + // Abatch[ip]: lps[ip] rows x AbatchCols[ip] cols + // Bbatch[ip]: rps[ip] rows x BbatchCols[ip] cols + // BXbatch[ip]: rps[ip] rows x AbatchCols[ip] cols VectorSizeType AbatchOff(npatches + 1, 0); VectorSizeType BbatchOff(npatches + 1, 0); - VectorSizeType BXbatchOff(npatches + 1, 0); // ldB rows x AbatchCols columns + VectorSizeType BXbatchOff(npatches + 1, 0); for (SizeType ip = 0; ip < npatches; ++ip) { - AbatchOff[ip + 1] = AbatchOff[ip] + ldA[ip] * AbatchCols[ip]; - BbatchOff[ip + 1] = BbatchOff[ip] + ldB[ip] * BbatchCols[ip]; - BXbatchOff[ip + 1] = BXbatchOff[ip] + ldB[ip] * AbatchCols[ip]; + AbatchOff[ip + 1] = AbatchOff[ip] + lps[ip] * AbatchCols[ip]; + BbatchOff[ip + 1] = BbatchOff[ip] + rps[ip] * BbatchCols[ip]; + BXbatchOff[ip + 1] = BXbatchOff[ip] + rps[ip] * AbatchCols[ip]; } const SizeType totalAbatch = AbatchOff[npatches]; @@ -316,13 +254,13 @@ class BatchedGemmKokkos { std::vector h_flatAbatch(totalAbatch, ComplexOrRealType(0)); std::vector h_flatBbatch(totalBbatch, ComplexOrRealType(0)); - // Build GEMM arg lists while packing matrices. - std::vector pass1_args; - std::vector pass2_args; - pass2_args.reserve(npatches); + // Build host GEMM arg lists while packing operator matrices. + h_pass1_.clear(); + h_pass2_.clear(); + h_pass2_.reserve(npatches); for (SizeType ip = 0; ip < npatches; ++ip) { - long long colA = 0; // column cursor in Abatch[ip] (= BXbatch[ip]) + long long colA = 0; // column cursor in Abatch[ip] and BXbatch[ip] long long colB = 0; // column cursor in Bbatch[ip] for (SizeType jp = 0; jp < npatches; ++jp) { @@ -336,37 +274,39 @@ class BatchedGemmKokkos { MatRef Bref = getMatRef(*Bmat); // right operator (yc) assert(Aref.ptr && Bref.ptr); - const int mA = static_cast(lps[ip]); - const int nAk = static_cast(lps[jp]); // A cols = BX cols per op - const int mB = static_cast(rps[ip]); - const int nBk = static_cast(rps[jp]); // B cols = X rows + const int mA = static_cast(lps[ip]); // Abatch rows + const int nAk = static_cast(lps[jp]); // Abatch cols per op + const int mB = static_cast(rps[ip]); // Bbatch / BXbatch rows + const int nBk = static_cast(rps[jp]); // Bbatch cols per op - // Pack A (left operator) into Abatch[ip] at column colA + // Pack A (left operator) into Abatch[ip] at column colA. + // Abatch[ip]: lps[ip] rows, ld = lps[ip] (no padding). lacpy(Aref.ptr, mA, nAk, Aref.rows, h_flatAbatch.data() - + AbatchOff[ip] + colA * static_cast(ldA[ip]), - static_cast(ldA[ip])); + + AbatchOff[ip] + colA * static_cast(lps[ip]), + mA); - // Pack B (right operator) into Bbatch[ip] at column colB + // Pack B (right operator) into Bbatch[ip] at column colB. + // Bbatch[ip]: rps[ip] rows, ld = rps[ip] (no padding). lacpy(Bref.ptr, mB, nBk, Bref.rows, h_flatBbatch.data() - + BbatchOff[ip] + colB * static_cast(ldB[ip]), - static_cast(ldB[ip])); + + BbatchOff[ip] + colB * static_cast(rps[ip]), + mB); - // Pass 1 GEMM: C(mB, nAk) = Bbatch_block(mB, nBk) * X_jp(nBk, nAk) + // Pass 1 GEMM: C(mB, nAk) = A(mB, nBk) * B(nBk, nAk) + // A = Bbatch block, ld = rps[ip] = mB + // B = X[jp], ld = rps[jp] = nBk + // C = BXbatch block, ld = rps[ip] = mB GemmArgs a1; - a1.m = mB; - a1.n = nAk; - a1.k = nBk; - a1.lda = static_cast(ldB[ip]); // Bbatch leading dim - a1.ldb = nBk; // X[jp] no padding: ld = rps[jp] - a1.ldc = static_cast(ldB[ip]); // BXbatch leading dim + a1.m = mB; + a1.n = nAk; + a1.k = nBk; a1.a_off = static_cast(BbatchOff[ip]) - + colB * static_cast(ldB[ip]); + + colB * static_cast(rps[ip]); a1.b_off = static_cast(xyStart[jp]); a1.c_off = static_cast(BXbatchOff[ip]) - + colA * static_cast(ldB[ip]); - pass1_args.push_back(a1); + + colA * static_cast(rps[ip]); + h_pass1_.push_back(a1); colA += nAk; colB += nBk; @@ -374,29 +314,26 @@ class BatchedGemmKokkos { } // Pass 2 GEMM: Y[ip](mB, mA) = BXbatch[ip](mB, k) * Abatch[ip](mA, k)^T + // A = BXbatch[ip], ld = rps[ip] = mB + // B = Abatch[ip], ld = lps[ip] = mA (transposed) + // C = Y[ip], ld = rps[ip] = mB const int totalCols = static_cast(colA); // = AbatchCols[ip] GemmArgs a2; - a2.m = static_cast(rps[ip]); - a2.n = static_cast(lps[ip]); - a2.k = totalCols; - a2.lda = static_cast(ldB[ip]); // BXbatch leading dim - a2.ldb = static_cast(ldA[ip]); // Abatch leading dim - a2.ldc = static_cast(rps[ip]); // Y[ip] no padding: ld = rps[ip] + a2.m = static_cast(rps[ip]); + a2.n = static_cast(lps[ip]); + a2.k = totalCols; a2.a_off = static_cast(BXbatchOff[ip]); a2.b_off = static_cast(AbatchOff[ip]); a2.c_off = static_cast(xyStart[ip]); - pass2_args.push_back(a2); + h_pass2_.push_back(a2); } - nbatch1_ = pass1_args.size(); - nbatch2_ = pass2_args.size(); // == npatches - - // Allocate device arrays and upload. + // Allocate device arrays and upload operator matrices. d_flatAbatch_ = DevScalView("d_flatAbatch", totalAbatch ? totalAbatch : 1); d_flatBbatch_ = DevScalView("d_flatBbatch", totalBbatch ? totalBbatch : 1); d_flatBXbatch_ = DevScalView("d_flatBXbatch", totalBXbatch ? totalBXbatch : 1); - d_vin_ = DevScalView("d_vin", xyStart[npatches] ? xyStart[npatches] : 1); - d_vout_ = DevScalView("d_vout", xyStart[npatches] ? xyStart[npatches] : 1); + d_vin_ = DevScalView("d_vin", xyStart[npatches] ? xyStart[npatches] : 1); + d_vout_ = DevScalView("d_vout", xyStart[npatches] ? xyStart[npatches] : 1); { auto hA = Kokkos::create_mirror_view(d_flatAbatch_); @@ -409,23 +346,12 @@ class BatchedGemmKokkos { Kokkos::deep_copy(d_flatBbatch_, hB); } - d_pass1_ = DevArgsView("d_pass1", nbatch1_ ? nbatch1_ : 1); - d_pass2_ = DevArgsView("d_pass2", nbatch2_ ? nbatch2_ : 1); - { - auto h1 = Kokkos::create_mirror_view(d_pass1_); - auto h2 = Kokkos::create_mirror_view(d_pass2_); - for (SizeType i = 0; i < nbatch1_; ++i) h1(i) = pass1_args[i]; - for (SizeType i = 0; i < nbatch2_; ++i) h2(i) = pass2_args[i]; - Kokkos::deep_copy(d_pass1_, h1); - Kokkos::deep_copy(d_pass2_, h2); - } - { PsimagLite::OstringStream msg(std::cout.precision()); msg() << "setup done: npatches=" << npatches << " noperator=" << noperator - << " pass1_batches=" << nbatch1_ - << " pass2_batches=" << nbatch2_ + << " pass1_batches=" << h_pass1_.size() + << " pass2_batches=" << h_pass2_.size() << " Abatch=" << totalAbatch << "elems" << " Bbatch=" << totalBbatch << "elems" << " BXbatch=" << totalBXbatch << "elems"; @@ -439,16 +365,14 @@ class BatchedGemmKokkos { PsimagLite::ProgressIndicator progress_; mutable VectorMatrixType garbage_; // owns sparse->dense expansions (freed in dtor) - SizeType nbatch1_ = 0; // number of pass-1 GEMMs - SizeType nbatch2_ = 0; // number of pass-2 GEMMs (== npatches) + std::vector h_pass1_; // pass-1 GEMM parameters (host, persistent) + std::vector h_pass2_; // pass-2 GEMM parameters (host, persistent) DevScalView d_flatAbatch_; // left-operator matrices, device, persistent DevScalView d_flatBbatch_; // right-operator matrices, device, persistent mutable DevScalView d_flatBXbatch_; // intermediate BX work buffer, per-call mutable DevScalView d_vin_; // input vector on device, per-call mutable DevScalView d_vout_; // output vector on device, per-call - DevArgsView d_pass1_; // pass-1 GEMM parameters, persistent - DevArgsView d_pass2_; // pass-2 GEMM parameters, persistent }; } // namespace Dmrg diff --git a/dmrg/GPUPlugin/DMRGConfig.h.in b/dmrg/GPUPlugin/DMRGConfig.h.in index b10e5c17f..d936ca38f 100644 --- a/dmrg/GPUPlugin/DMRGConfig.h.in +++ b/dmrg/GPUPlugin/DMRGConfig.h.in @@ -3,6 +3,7 @@ #define DMRG_CONFIG_H #cmakedefine PLUGIN_SC +#cmakedefine KOKKOS_BATCHED #cmakedefine FpType @FpType@ #cmakedefine USE_MAGMA #cmakedefine USE_COMPLEX_Z diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 2f3db5167..52958151d 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -280,6 +280,7 @@ void csr_kron_mult_method(const int imethod, */ using ExecutionSpace = Kokkos::DefaultExecutionSpace; + using MemorySpace = ExecutionSpace::memory_space; using KokkosScalar = typename PsimagLite::KokkosType::type; const int nnzA = a.nonZeros(); @@ -292,39 +293,6 @@ void csr_kron_mult_method(const int imethod, */ const size_t totalPairs = static_cast(nnzA) * static_cast(nnzB); static constexpr size_t kGpuThreshold = 100000; - if (false) { - Kokkos::Profiling::ScopedRegion cpuRegion("PsimgLite::csr_kron_mult_method::imethod3::cpu"); - for (int ia = 0; ia < nrow_A; ++ia) { - const int istart_a = a.getRowPtr(ia); - const int iend_a = a.getRowPtr(ia + 1); - for (int ka = istart_a; ka < iend_a; ++ka) { - const int ja = a.getCol(ka); - ComplexOrRealType aij = a.getValue(ka); - if constexpr (is_complex) - if (isConjTransA) - aij = PsimagLite::conj(aij); - for (int ib = 0; ib < nrow_B; ++ib) { - const int istart_b = b.getRowPtr(ib); - const int iend_b = b.getRowPtr(ib + 1); - for (int kb = istart_b; kb < iend_b; ++kb) { - const int jb = b.getCol(kb); - ComplexOrRealType bij = b.getValue(kb); - if constexpr (is_complex) - if (isConjTransB) - bij = PsimagLite::conj(bij); - - const int ix = (isTransB || isConjTransB) ? jb : ib; - const int jx = (isTransA || isConjTransA) ? ja : ia; - const int iy = (isTransB || isConjTransB) ? ib : jb; - const int jy = (isTransA || isConjTransA) ? ia : ja; - - xout(ix, jx) += aij * bij * yin(iy, jy); - } - } - } - } - return; - } /* * GPU path for large problems. @@ -346,6 +314,8 @@ void csr_kron_mult_method(const int imethod, * With LayoutLeft both the x_dev writes and the y_dev reads are * stride-1 (ib / jb vary across threads, ia / ja fixed per team). */ + Kokkos::Profiling::pushRegion("PsimgLite::csr_kron_mult_method::imethod3::view_init"); + Kokkos::View A_rowptr_h(&a.getRowPtr(0), nrow_A + 1); Kokkos::View @@ -360,22 +330,23 @@ void csr_kron_mult_method(const int imethod, Kokkos::View B_val_h(reinterpret_cast(&b.getValue(0)), nnzB); - auto A_rowptr_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, A_rowptr_h); - auto A_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, A_col_h); - auto A_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, A_val_h); - auto B_rowptr_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, B_rowptr_h); - auto B_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, B_col_h); - auto B_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, B_val_h); + auto A_rowptr_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), A_rowptr_h); + auto A_col_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), A_col_h); + auto A_val_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), A_val_h); + auto B_rowptr_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), B_rowptr_h); + auto B_col_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), B_col_h); + auto B_val_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), B_val_h); auto yin_host = Kokkos::View( reinterpret_cast(&yin(0, 0)), nrow_Y, ncol_Y); - auto y_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, yin_host); + auto y_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), yin_host); auto x_dev = Kokkos::View("x_dev", nrow_X, ncol_X); - Kokkos::deep_copy(x_dev, KokkosScalar(0)); + + Kokkos::Profiling::popRegion(); using TeamPolicy = Kokkos::TeamPolicy; using TeamMember = typename TeamPolicy::member_type; @@ -388,7 +359,7 @@ void csr_kron_mult_method(const int imethod, */ Kokkos::parallel_for( "csr_kron_mult::imethod3_nn", - TeamPolicy(nrow_A, Kokkos::AUTO), + TeamPolicy(nrow_A, Kokkos::AUTO, 32), KOKKOS_LAMBDA(const TeamMember& team) { const int ia = team.league_rank(); const int ka_begin = A_rowptr_dev(ia); @@ -399,13 +370,15 @@ void csr_kron_mult_method(const int imethod, const int kb_begin = B_rowptr_dev(ib); const int kb_end = B_rowptr_dev(ib + 1); KokkosScalar acc = 0; - for (int ka = ka_begin; ka < ka_end; ++ka) { + Kokkos::parallel_reduce(Kokkos::ThreadVectorRange(team, ka_begin, ka_end), [&](int ka, KokkosScalar& update) { const int ja = A_col_dev(ka); KokkosScalar aij = A_val_dev(ka); for (int kb = kb_begin; kb < kb_end; ++kb) - acc += aij * B_val_dev(kb) * y_dev(B_col_dev(kb), ja); - } - x_dev(ib, ia) += acc; + update += aij * B_val_dev(kb) * y_dev(B_col_dev(kb), ja); + }, acc); +Kokkos::single(Kokkos::PerThread(team), + [&]() { + x_dev(ib, ia) += acc; }); }); }); } else { @@ -448,11 +421,15 @@ void csr_kron_mult_method(const int imethod, }); } +{ + Kokkos::Profiling::ScopedRegion cpuRegion("PsimgLite::csr_kron_mult_method::imethod3::copy_back"); + // copy result back and accumulate into xout auto xhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, x_dev); for (int ix = 0; ix < nrow_X; ++ix) for (int jx = 0; jx < ncol_X; ++jx) xout(ix, jx) += static_cast(xhost(ix, jx)); +} }; } From 188dc3925f06af3e748f3fd1f1903c96c884ed31 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 27 Jul 2026 14:44:56 -0400 Subject: [PATCH 18/41] Minimize allocation overhead --- dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h index dfafc476e..30f043e14 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h @@ -329,15 +329,15 @@ class BatchedGemmKokkos { } // Allocate device arrays and upload operator matrices. - d_flatAbatch_ = DevScalView("d_flatAbatch", totalAbatch ? totalAbatch : 1); - d_flatBbatch_ = DevScalView("d_flatBbatch", totalBbatch ? totalBbatch : 1); - d_flatBXbatch_ = DevScalView("d_flatBXbatch", totalBXbatch ? totalBXbatch : 1); - d_vin_ = DevScalView("d_vin", xyStart[npatches] ? xyStart[npatches] : 1); - d_vout_ = DevScalView("d_vout", xyStart[npatches] ? xyStart[npatches] : 1); + d_flatAbatch_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_flatAbatch"), totalAbatch ? totalAbatch : 1); + d_flatBbatch_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_flatBbatch"), totalBbatch ? totalBbatch : 1); + d_flatBXbatch_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_flatBXbatch"), totalBXbatch ? totalBXbatch : 1); + d_vin_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_vin"), xyStart[npatches] ? xyStart[npatches] : 1); + d_vout_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_vout"), xyStart[npatches] ? xyStart[npatches] : 1); { - auto hA = Kokkos::create_mirror_view(d_flatAbatch_); - auto hB = Kokkos::create_mirror_view(d_flatBbatch_); + auto hA = Kokkos::create_mirror_view(Kokkos::view_alloc(Kokkos::WithoutInitializing), d_flatAbatch_); + auto hB = Kokkos::create_mirror_view(Kokkos::view_alloc(Kokkos::WithoutInitializing), d_flatBbatch_); for (SizeType i = 0; i < totalAbatch; ++i) hA(i) = *reinterpret_cast(&h_flatAbatch[i]); for (SizeType i = 0; i < totalBbatch; ++i) From 296278c20c8c4df6a4ad9711bc1f1944e4df93f0 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 27 Jul 2026 15:48:40 -0400 Subject: [PATCH 19/41] Add data342_batchedgemm --- TestSuite/inputs/input342.ain | 3 +- TestSuite/inputs/input342_batchedgemm.ain | 60 +++++++++++++++++++++++ dmrg/CMakeLists.txt | 4 +- 3 files changed, 63 insertions(+), 4 deletions(-) create mode 100644 TestSuite/inputs/input342_batchedgemm.ain diff --git a/TestSuite/inputs/input342.ain b/TestSuite/inputs/input342.ain index 8a2b4c2de..247ddffb6 100644 --- a/TestSuite/inputs/input342.ain +++ b/TestSuite/inputs/input342.ain @@ -47,8 +47,7 @@ potentialV=[0.,...x32]; Model=FeAsBasedScExtended; FeAsMode=INT_PAPER33; Orbitals=2; -SolverOptions=BatchedGemm,twositedmrg,minimizedisk; -integer DenseSparseThreshold=0; +SolverOptions=twositedmrg,minimizedisk; Version=61289987cdd32b8485213ac0415baf4e9cd16432; OutputFile=data342; InfiniteLoopKeptStates=20; diff --git a/TestSuite/inputs/input342_batchedgemm.ain b/TestSuite/inputs/input342_batchedgemm.ain new file mode 100644 index 000000000..89dd2361b --- /dev/null +++ b/TestSuite/inputs/input342_batchedgemm.ain @@ -0,0 +1,60 @@ +##Ainur1.0 + +TotalNumberOfSites=8; +NumberOfTerms=3; + +gt0:DegreesOfFreedom=2; +gt0:GeometryKind=ladderx; +gt0:GeometryOptions=ConstantValues; +gt0:LadderLeg=2; +gt0:dir0:Connectors=[ +[-0.058, 0], +[0, -0.2196]]; + +gt0:dir1:Connectors=[ +[-0.2196, 0], +[0, -0.058]]; + +gt0:dir2:Connectors=[ +[0.20828, 0.079], +[0.079, 0.20828]]; + +gt0:dir3:Connectors=[ +[0.20828, -0.079], +[-0.079, 0.20828]]; + +gt1:DegreesOfFreedom=1; +gt1:GeometryKind=ladderx; +gt1:GeometryOptions=ConstantValues; +gt1:LadderLeg=2; +gt1:dir0:Connectors= [0.1]; +gt1:dir1:Connectors= [0.1]; +gt1:dir2:Connectors= [0.02]; +gt1:dir3:Connectors= [0.02]; + +gt2:DegreesOfFreedom=1; +gt2:GeometryKind=ladderx; +gt2:GeometryOptions=ConstantValues; +gt2:LadderLeg=2; +gt2:dir0:Connectors= [0.1]; +gt2:dir1:Connectors= [0.1]; +gt2:dir2:Connectors= [0.02]; +gt2:dir3:Connectors= [0.02]; + +hubbardU=[4.,...x4]; +# The syntax below implies 32 zeroes for the vector +potentialV=[0.,...x32]; +Model=FeAsBasedScExtended; +FeAsMode=INT_PAPER33; +Orbitals=2; +SolverOptions=BatchedGemm,twositedmrg,minimizedisk; +integer DenseSparseThreshold=0; +Version=61289987cdd32b8485213ac0415baf4e9cd16432; +OutputFile=data342_batchedgemm; +InfiniteLoopKeptStates=20; +FiniteLoops=[ + [@auto, 250, 0], + [@auto, 500, 0]]; +TargetElectronsUp=8; +TargetElectronsDown=8; + diff --git a/dmrg/CMakeLists.txt b/dmrg/CMakeLists.txt index 1a3c0f765..69ca3fb4c 100644 --- a/dmrg/CMakeLists.txt +++ b/dmrg/CMakeLists.txt @@ -180,9 +180,9 @@ if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") # The following tests are expensive to add_test(NAME input341 COMMAND ./dmrg -f ${PATH_TO_INPUTS}/input341.ain) # energy -0.093321 add_lowest_eigenvalue_check_test(output341 input341 -0.093321 runForinput341.cout) add_test(NAME input342 COMMAND ./dmrg -f ${PATH_TO_INPUTS}/input342.ain) # energy -0.0933211 - set_tests_properties(input342 PROPERTIES - ENVIRONMENT "KOKKOS_TOOLS_LIBS=$ENV{HOME}/kokkos-tools/build/profiling/space-time-stack/libkp_space_time_stack.so") add_lowest_eigenvalue_check_test(output342 input342 -0.0933211 runForinput342.cout) + add_test(NAME input342_batchedgemm COMMAND ./dmrg -f ${PATH_TO_INPUTS}/input342_batchedgemm.ain) # energy -0.0933211 + add_lowest_eigenvalue_check_test(output342_batchedgemm input342_batchedgemm -0.0933211 runForinput342_batchedgemm.cout) endif() add_test(NAME input351 COMMAND ./dmrg -f ${PATH_TO_INPUTS}/input351.ain) # energy -0.408019 add_lowest_eigenvalue_check_test(output351 input351 -0.408019 runForinput351.cout) From d271d67a9c9e95e60257c244c4bc6115ff2ec72e Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Wed, 29 Jul 2026 14:41:51 -0400 Subject: [PATCH 20/41] Revert "Optimixe more" This reverts commit 46212595f5df83f66842e9b19960a4e87087a3d9. --- CMakeLists.txt | 5 - dmrg/CMakeLists.txt | 24 +- .../MatrixVectorKron/BatchedGemmKokkos.h | 250 ++++++++++++------ dmrg/GPUPlugin/DMRGConfig.h.in | 1 - dmrg/KronUtil/csr_kron_mult.cpp | 69 +++-- 5 files changed, 216 insertions(+), 133 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index fda339d21..3747e5fa2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -63,11 +63,6 @@ option(DMRG_BUILD_BATCHED "Enable building with Batched support" ON) include(CMakeDependentOption) cmake_dependent_option(DMRG_BUILD_BATCHED_MAGMA "Enable building with Magma" OFF "DMRG_BUILD_BATCHED" OFF) -# KokkosKernels batched GEMM: portable GPU path, no separate GPU plugin needed. -# Kokkos and KokkosKernels are already pulled in via PsimagLite. -cmake_dependent_option(DMRG_BUILD_BATCHED_KOKKOS - "Use KokkosKernels for batched GEMM (GPU-portable, no GPU plugin required)" - OFF "DMRG_BUILD_BATCHED" OFF) add_subdirectory(util) diff --git a/dmrg/CMakeLists.txt b/dmrg/CMakeLists.txt index 69ca3fb4c..71d66bce6 100644 --- a/dmrg/CMakeLists.txt +++ b/dmrg/CMakeLists.txt @@ -6,21 +6,11 @@ add_custom_command( VERBATIM) if(DMRG_BUILD_BATCHED) - if(DMRG_BUILD_BATCHED_KOKKOS) - # KokkosKernels path: portable GPU implementation, no separate GPU plugin. - # Kokkos and KokkosKernels are already a dependency of psimaglite. - set(KOKKOS_BATCHED ON) - set(FpType "double") - message(STATUS "BatchedGemm: using KokkosKernels (DMRG_BUILD_BATCHED_KOKKOS=ON)") - else() - # Original GPU plugin path (optionally with MAGMA). - add_subdirectory(GPUPlugin) - set(PLUGIN_SC ON) - set(FpType "double") - if(DMRG_BUILD_BATCHED_MAGMA) - set(USE_MAGMA ON) - endif() - message(STATUS "BatchedGemm: using GPU plugin (PLUGIN_SC)") + add_subdirectory(GPUPlugin) + set(PLUGIN_SC ON) + set(FpType "double") + if(DMRG_BUILD_BATCHED_MAGMA) + set(USE_MAGMA ON) endif() endif() @@ -35,14 +25,14 @@ add_library(dmrgpp_utils ProgramGlobals.cpp Provenance.cpp Utils.cpp Su2Related. target_sources(dmrgpp_utils PUBLIC ${CMAKE_CURRENT_BINARY_DIR}/GitRevision.h) target_include_directories(dmrgpp_utils PUBLIC ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR}/GPUPlugin Engine) target_link_libraries(dmrgpp_utils PUBLIC psimaglite::psimaglite) -if(DMRG_BUILD_BATCHED AND NOT DMRG_BUILD_BATCHED_KOKKOS) +if(DMRG_BUILD_BATCHED) target_link_libraries(dmrgpp_utils PUBLIC gpuplugin) endif() # DMRG runner add_library(dmrg_runner Engine/DmrgRunner.cpp) target_link_libraries(dmrg_runner PUBLIC dmrgpp_utils kronutil) -if(DMRG_BUILD_BATCHED AND NOT DMRG_BUILD_BATCHED_KOKKOS) +if(DMRG_BUILD_BATCHED) target_link_libraries(dmrg_runner PUBLIC gpuplugin) endif() diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h index 30f043e14..0137b4560 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h @@ -11,7 +11,8 @@ #include #include -#include +#include +#include namespace Dmrg { @@ -24,16 +25,12 @@ namespace Dmrg { * BXbatch[ip][:, colBX:colBX+lps[jp]] = * Bbatch[ip][:, colB:colB+rps[jp]] * X[jp] * - * Pass 2: For each ipatch (with connections): + * Pass 2: For each ipatch: * Y[ip] = BXbatch[ip] * Abatch[ip]^T * - * Abatch[ip] = left-operator (xc) matrices packed column-by-column (no padding). - * Bbatch[ip] = right-operator (yc) matrices packed column-by-column (no padding). - * BXbatch[ip] is the intermediate work buffer (no padding, same col-width as Abatch). - * - * No padding is used so that unmanaged Kokkos::LayoutLeft views can be created - * directly from device pointers, enabling dispatch to rocBLAS/cuBLAS via - * KokkosBlas::gemm. + * Abatch[ip] = left-operator (xc) matrices packed column-by-column (ldA padded). + * Bbatch[ip] = right-operator (yc) matrices packed column-by-column (ldB padded). + * BXbatch[ip] is the intermediate work buffer (ldB padded, same col-width as Abatch). * * All Abatch and Bbatch data live on the device for the lifetime of the object. * BXbatch, vin, vout are (re-)written on each matrixVector call. @@ -57,13 +54,18 @@ class BatchedGemmKokkos { using DevMemSpace = typename DevExecSpace::memory_space; using DevScalView = Kokkos::View; - // Parameters for one KokkosBlas::gemm call. - // With no padding: leading dims = m (for A,C) or k (for B in pass1) or n (for B in pass2). + // Compact struct holding all parameters for one batched GEMM call. + // Stored in device memory and accessed inside the kernel. struct GemmArgs { int m, n, k; // GEMM dimensions + int lda, ldb, ldc; // leading dimensions (column-major strides) long long a_off, b_off, c_off; // element offsets into device flat arrays }; + using DevArgsView = Kokkos::View; + + static const int ialign_ = 32; + public: BatchedGemmKokkos(const InitKronType& initKron) @@ -102,40 +104,97 @@ class BatchedGemmKokkos { Kokkos::deep_copy(d_vin_, hv); } - // Zero d_vout_ so patches with no connections produce zero output. - // d_flatBXbatch_ is NOT zeroed: pass1 uses beta=0 (overwrites all used blocks). + // --- Zero output and work buffers ------------------------------------------- Kokkos::deep_copy(d_vout_, KS(0)); + Kokkos::deep_copy(d_flatBXbatch_, KS(0)); DevExecSpace exec; - using UV = Kokkos::View; - - // --- Pass 1: BXbatch[ip] block = Bbatch[ip] block * X[jp] (N x N) --------- - // Each call: C(m, n) = A(m, k) * B(k, n), beta=0 overwrites C. - // A = Bbatch block (m=rps[ip], k=rps[jp]), ld = rps[ip] - // B = X[jp] (k=rps[jp], n=lps[jp]), ld = rps[jp] - // C = BXbatch block (m=rps[ip], n=lps[jp]), ld = rps[ip] - for (const GemmArgs& ag : h_pass1_) { - UV A(d_flatBbatch_.data() + ag.a_off, ag.m, ag.k); - UV B(d_vin_.data() + ag.b_off, ag.k, ag.n); - UV C(d_flatBXbatch_.data() + ag.c_off, ag.m, ag.n); - KokkosBlas::gemm(exec, "N", "N", KS(1), A, B, KS(0), C); + // --- Pass 1: BXbatch[ip] = Bbatch[ip] * X[jp] (NoT x NoT) ---------------- + { + const DevScalView flatBbatch = d_flatBbatch_; + const DevScalView vin_dev = d_vin_; + const DevScalView flatBXbatch = d_flatBXbatch_; + const DevArgsView args = d_pass1_; + + using MemberType = typename Kokkos::TeamPolicy::member_type; + Kokkos::parallel_for( + "BatchedGemmKokkos_Pass1", + Kokkos::TeamPolicy(exec, static_cast(nbatch1_), + Kokkos::AUTO, Kokkos::AUTO), + KOKKOS_LAMBDA(const MemberType& member) { + const int i = member.league_rank(); + const GemmArgs& ag = args(i); + + using UV = Kokkos::View; + + // A = Bbatch[ip] block: (lda, k) padded -> subview (m, k) active rows + UV A_full(flatBbatch.data() + ag.a_off, ag.lda, ag.k); + auto A = Kokkos::subview(A_full, + Kokkos::make_pair(0, ag.m), Kokkos::ALL()); + + // B = X[jp]: (k, n) exact -- no padding for X slice + UV B(vin_dev.data() + ag.b_off, ag.ldb, ag.n); + + // C = BXbatch[ip] block: (ldc, n) padded -> subview (m, n) active rows + UV C_full(flatBXbatch.data() + ag.c_off, ag.ldc, ag.n); + auto C = Kokkos::subview(C_full, + Kokkos::make_pair(0, ag.m), Kokkos::ALL()); + + KokkosBatched::TeamVectorGemm< + MemberType, + KokkosBatched::Trans::NoTranspose, + KokkosBatched::Trans::NoTranspose, + KokkosBatched::Algo::Gemm::Unblocked>:: + invoke(member, KS(1), A, B, KS(0), C); + }); } exec.fence(); - // --- Pass 2: Y[ip] = BXbatch[ip] * Abatch[ip]^T (N x T) ------------------ - // Each call: C(m, n) = A(m, k) * B(n, k)^T, beta=0 overwrites C. - // A = BXbatch[ip] (m=rps[ip], k=AbatchCols[ip]), ld = rps[ip] - // B = Abatch[ip] (n=lps[ip], k=AbatchCols[ip]), ld = lps[ip] (transposed) - // C = Y[ip] (m=rps[ip], n=lps[ip]), ld = rps[ip] - for (const GemmArgs& ag : h_pass2_) { - if (ag.k == 0) - continue; // no connections for this ipatch; Y[ip] stays zero - UV A(d_flatBXbatch_.data() + ag.a_off, ag.m, ag.k); - UV B(d_flatAbatch_.data() + ag.b_off, ag.n, ag.k); - UV C(d_vout_.data() + ag.c_off, ag.m, ag.n); - KokkosBlas::gemm(exec, "N", "T", KS(1), A, B, KS(0), C); + // --- Pass 2: Y[ip] = BXbatch[ip] * Abatch[ip]^T (NoT x T) --------------- + { + const DevScalView flatBXbatch = d_flatBXbatch_; + const DevScalView flatAbatch = d_flatAbatch_; + const DevScalView vout_dev = d_vout_; + const DevArgsView args = d_pass2_; + + using MemberType = typename Kokkos::TeamPolicy::member_type; + Kokkos::parallel_for( + "BatchedGemmKokkos_Pass2", + Kokkos::TeamPolicy(exec, static_cast(nbatch2_), + Kokkos::AUTO, Kokkos::AUTO), + KOKKOS_LAMBDA(const MemberType& member) { + const int i = member.league_rank(); + const GemmArgs& ag = args(i); + + if (ag.k == 0) + return; // no connections for this ipatch; Y[ip] already zero + + using UV = Kokkos::View; + + // A = BXbatch[ip]: (lda, k) padded -> subview (m, k) + UV A_full(flatBXbatch.data() + ag.a_off, ag.lda, ag.k); + auto A = Kokkos::subview(A_full, + Kokkos::make_pair(0, ag.m), Kokkos::ALL()); + + // B = Abatch[ip]: (ldb, k) padded -> subview (n, k); will be transposed + UV B_full(flatAbatch.data() + ag.b_off, ag.ldb, ag.k); + auto B = Kokkos::subview(B_full, + Kokkos::make_pair(0, ag.n), Kokkos::ALL()); + + // C = Y[ip]: (ldc, n) exact -- no padding for output + UV C(vout_dev.data() + ag.c_off, ag.ldc, ag.n); + + KokkosBatched::TeamVectorGemm< + MemberType, + KokkosBatched::Trans::NoTranspose, + KokkosBatched::Trans::Transpose, + KokkosBatched::Algo::Gemm::Unblocked>:: + invoke(member, KS(1), A, B, KS(0), C); + }); } exec.fence(); @@ -150,6 +209,8 @@ class BatchedGemmKokkos { private: + static int iceil(int x, int n) { return (x + n - 1) / n; } + // Dense matrix reference: pointer + source dimensions (no padding). // For sparse matrices, the expansion is heap-allocated and owned by garbage_. struct MatRef { @@ -193,9 +254,11 @@ class BatchedGemmKokkos { const SizeType npatches = initKron_.numberOfPatches(InitKronType::OLD); const SizeType noperator = initKron_.connections(); - // Per-patch sizes — NO padding: leading dim == patch size. - VectorSizeType lps(npatches, 0); // left patch size - VectorSizeType rps(npatches, 0); // right patch size + // Per-patch sizes and padded leading dimensions. + VectorSizeType lps(npatches, 0); // left patch size + VectorSizeType rps(npatches, 0); // right patch size + VectorSizeType ldA(npatches, 0); // padded lps -> Abatch leading dim + VectorSizeType ldB(npatches, 0); // padded rps -> Bbatch / BXbatch leading dim VectorSizeType xyStart(npatches + 1, 0); // patch start offset in vin/vout for (SizeType ip = 0; ip < npatches; ++ip) { @@ -207,13 +270,15 @@ class BatchedGemmKokkos { const int R2 = initKron_.lrs(InitKronType::NEW).right().partition(rg + 1); lps[ip] = static_cast(L2 - L1); rps[ip] = static_cast(R2 - R1); + ldA[ip] = static_cast(ialign_ * iceil(static_cast(lps[ip]), ialign_)); + ldB[ip] = static_cast(ialign_ * iceil(static_cast(rps[ip]), ialign_)); } xyStart[0] = 0; for (SizeType ip = 0; ip < npatches; ++ip) xyStart[ip + 1] = xyStart[ip] + lps[ip] * rps[ip]; - // Column widths for each patch's batch block. + // Column widths: // AbatchCols[ip] = sum of lps[jp] over all non-zero (jp, k) connections for ip // BbatchCols[ip] = sum of rps[jp] over all non-zero (jp, k) connections for ip VectorSizeType AbatchCols(npatches, 0); @@ -232,18 +297,15 @@ class BatchedGemmKokkos { } } - // Flat offsets: no padding, ld[ip] = rps[ip] or lps[ip]. - // Abatch[ip]: lps[ip] rows x AbatchCols[ip] cols - // Bbatch[ip]: rps[ip] rows x BbatchCols[ip] cols - // BXbatch[ip]: rps[ip] rows x AbatchCols[ip] cols + // Flat offsets: each patch ip occupies a contiguous slice. VectorSizeType AbatchOff(npatches + 1, 0); VectorSizeType BbatchOff(npatches + 1, 0); - VectorSizeType BXbatchOff(npatches + 1, 0); + VectorSizeType BXbatchOff(npatches + 1, 0); // ldB rows x AbatchCols columns for (SizeType ip = 0; ip < npatches; ++ip) { - AbatchOff[ip + 1] = AbatchOff[ip] + lps[ip] * AbatchCols[ip]; - BbatchOff[ip + 1] = BbatchOff[ip] + rps[ip] * BbatchCols[ip]; - BXbatchOff[ip + 1] = BXbatchOff[ip] + rps[ip] * AbatchCols[ip]; + AbatchOff[ip + 1] = AbatchOff[ip] + ldA[ip] * AbatchCols[ip]; + BbatchOff[ip + 1] = BbatchOff[ip] + ldB[ip] * BbatchCols[ip]; + BXbatchOff[ip + 1] = BXbatchOff[ip] + ldB[ip] * AbatchCols[ip]; } const SizeType totalAbatch = AbatchOff[npatches]; @@ -254,13 +316,13 @@ class BatchedGemmKokkos { std::vector h_flatAbatch(totalAbatch, ComplexOrRealType(0)); std::vector h_flatBbatch(totalBbatch, ComplexOrRealType(0)); - // Build host GEMM arg lists while packing operator matrices. - h_pass1_.clear(); - h_pass2_.clear(); - h_pass2_.reserve(npatches); + // Build GEMM arg lists while packing matrices. + std::vector pass1_args; + std::vector pass2_args; + pass2_args.reserve(npatches); for (SizeType ip = 0; ip < npatches; ++ip) { - long long colA = 0; // column cursor in Abatch[ip] and BXbatch[ip] + long long colA = 0; // column cursor in Abatch[ip] (= BXbatch[ip]) long long colB = 0; // column cursor in Bbatch[ip] for (SizeType jp = 0; jp < npatches; ++jp) { @@ -274,39 +336,37 @@ class BatchedGemmKokkos { MatRef Bref = getMatRef(*Bmat); // right operator (yc) assert(Aref.ptr && Bref.ptr); - const int mA = static_cast(lps[ip]); // Abatch rows - const int nAk = static_cast(lps[jp]); // Abatch cols per op - const int mB = static_cast(rps[ip]); // Bbatch / BXbatch rows - const int nBk = static_cast(rps[jp]); // Bbatch cols per op + const int mA = static_cast(lps[ip]); + const int nAk = static_cast(lps[jp]); // A cols = BX cols per op + const int mB = static_cast(rps[ip]); + const int nBk = static_cast(rps[jp]); // B cols = X rows - // Pack A (left operator) into Abatch[ip] at column colA. - // Abatch[ip]: lps[ip] rows, ld = lps[ip] (no padding). + // Pack A (left operator) into Abatch[ip] at column colA lacpy(Aref.ptr, mA, nAk, Aref.rows, h_flatAbatch.data() - + AbatchOff[ip] + colA * static_cast(lps[ip]), - mA); + + AbatchOff[ip] + colA * static_cast(ldA[ip]), + static_cast(ldA[ip])); - // Pack B (right operator) into Bbatch[ip] at column colB. - // Bbatch[ip]: rps[ip] rows, ld = rps[ip] (no padding). + // Pack B (right operator) into Bbatch[ip] at column colB lacpy(Bref.ptr, mB, nBk, Bref.rows, h_flatBbatch.data() - + BbatchOff[ip] + colB * static_cast(rps[ip]), - mB); + + BbatchOff[ip] + colB * static_cast(ldB[ip]), + static_cast(ldB[ip])); - // Pass 1 GEMM: C(mB, nAk) = A(mB, nBk) * B(nBk, nAk) - // A = Bbatch block, ld = rps[ip] = mB - // B = X[jp], ld = rps[jp] = nBk - // C = BXbatch block, ld = rps[ip] = mB + // Pass 1 GEMM: C(mB, nAk) = Bbatch_block(mB, nBk) * X_jp(nBk, nAk) GemmArgs a1; - a1.m = mB; - a1.n = nAk; - a1.k = nBk; + a1.m = mB; + a1.n = nAk; + a1.k = nBk; + a1.lda = static_cast(ldB[ip]); // Bbatch leading dim + a1.ldb = nBk; // X[jp] no padding: ld = rps[jp] + a1.ldc = static_cast(ldB[ip]); // BXbatch leading dim a1.a_off = static_cast(BbatchOff[ip]) - + colB * static_cast(rps[ip]); + + colB * static_cast(ldB[ip]); a1.b_off = static_cast(xyStart[jp]); a1.c_off = static_cast(BXbatchOff[ip]) - + colA * static_cast(rps[ip]); - h_pass1_.push_back(a1); + + colA * static_cast(ldB[ip]); + pass1_args.push_back(a1); colA += nAk; colB += nBk; @@ -314,18 +374,18 @@ class BatchedGemmKokkos { } // Pass 2 GEMM: Y[ip](mB, mA) = BXbatch[ip](mB, k) * Abatch[ip](mA, k)^T - // A = BXbatch[ip], ld = rps[ip] = mB - // B = Abatch[ip], ld = lps[ip] = mA (transposed) - // C = Y[ip], ld = rps[ip] = mB const int totalCols = static_cast(colA); // = AbatchCols[ip] GemmArgs a2; - a2.m = static_cast(rps[ip]); - a2.n = static_cast(lps[ip]); - a2.k = totalCols; + a2.m = static_cast(rps[ip]); + a2.n = static_cast(lps[ip]); + a2.k = totalCols; + a2.lda = static_cast(ldB[ip]); // BXbatch leading dim + a2.ldb = static_cast(ldA[ip]); // Abatch leading dim + a2.ldc = static_cast(rps[ip]); // Y[ip] no padding: ld = rps[ip] a2.a_off = static_cast(BXbatchOff[ip]); a2.b_off = static_cast(AbatchOff[ip]); a2.c_off = static_cast(xyStart[ip]); - h_pass2_.push_back(a2); + pass2_args.push_back(a2); } // Allocate device arrays and upload operator matrices. @@ -335,6 +395,9 @@ class BatchedGemmKokkos { d_vin_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_vin"), xyStart[npatches] ? xyStart[npatches] : 1); d_vout_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_vout"), xyStart[npatches] ? xyStart[npatches] : 1); + nbatch1_ = pass1_args.size(); + nbatch2_ = pass2_args.size(); // == npatches + { auto hA = Kokkos::create_mirror_view(Kokkos::view_alloc(Kokkos::WithoutInitializing), d_flatAbatch_); auto hB = Kokkos::create_mirror_view(Kokkos::view_alloc(Kokkos::WithoutInitializing), d_flatBbatch_); @@ -346,12 +409,23 @@ class BatchedGemmKokkos { Kokkos::deep_copy(d_flatBbatch_, hB); } + d_pass1_ = DevArgsView("d_pass1", nbatch1_ ? nbatch1_ : 1); + d_pass2_ = DevArgsView("d_pass2", nbatch2_ ? nbatch2_ : 1); + { + auto h1 = Kokkos::create_mirror_view(d_pass1_); + auto h2 = Kokkos::create_mirror_view(d_pass2_); + for (SizeType i = 0; i < nbatch1_; ++i) h1(i) = pass1_args[i]; + for (SizeType i = 0; i < nbatch2_; ++i) h2(i) = pass2_args[i]; + Kokkos::deep_copy(d_pass1_, h1); + Kokkos::deep_copy(d_pass2_, h2); + } + { PsimagLite::OstringStream msg(std::cout.precision()); msg() << "setup done: npatches=" << npatches << " noperator=" << noperator - << " pass1_batches=" << h_pass1_.size() - << " pass2_batches=" << h_pass2_.size() + << " pass1_batches=" << nbatch1_ + << " pass2_batches=" << nbatch2_ << " Abatch=" << totalAbatch << "elems" << " Bbatch=" << totalBbatch << "elems" << " BXbatch=" << totalBXbatch << "elems"; @@ -365,14 +439,16 @@ class BatchedGemmKokkos { PsimagLite::ProgressIndicator progress_; mutable VectorMatrixType garbage_; // owns sparse->dense expansions (freed in dtor) - std::vector h_pass1_; // pass-1 GEMM parameters (host, persistent) - std::vector h_pass2_; // pass-2 GEMM parameters (host, persistent) + SizeType nbatch1_ = 0; // number of pass-1 GEMMs + SizeType nbatch2_ = 0; // number of pass-2 GEMMs (== npatches) DevScalView d_flatAbatch_; // left-operator matrices, device, persistent DevScalView d_flatBbatch_; // right-operator matrices, device, persistent mutable DevScalView d_flatBXbatch_; // intermediate BX work buffer, per-call mutable DevScalView d_vin_; // input vector on device, per-call mutable DevScalView d_vout_; // output vector on device, per-call + DevArgsView d_pass1_; // pass-1 GEMM parameters, persistent + DevArgsView d_pass2_; // pass-2 GEMM parameters, persistent }; } // namespace Dmrg diff --git a/dmrg/GPUPlugin/DMRGConfig.h.in b/dmrg/GPUPlugin/DMRGConfig.h.in index d936ca38f..b10e5c17f 100644 --- a/dmrg/GPUPlugin/DMRGConfig.h.in +++ b/dmrg/GPUPlugin/DMRGConfig.h.in @@ -3,7 +3,6 @@ #define DMRG_CONFIG_H #cmakedefine PLUGIN_SC -#cmakedefine KOKKOS_BATCHED #cmakedefine FpType @FpType@ #cmakedefine USE_MAGMA #cmakedefine USE_COMPLEX_Z diff --git a/dmrg/KronUtil/csr_kron_mult.cpp b/dmrg/KronUtil/csr_kron_mult.cpp index 52958151d..2f3db5167 100644 --- a/dmrg/KronUtil/csr_kron_mult.cpp +++ b/dmrg/KronUtil/csr_kron_mult.cpp @@ -280,7 +280,6 @@ void csr_kron_mult_method(const int imethod, */ using ExecutionSpace = Kokkos::DefaultExecutionSpace; - using MemorySpace = ExecutionSpace::memory_space; using KokkosScalar = typename PsimagLite::KokkosType::type; const int nnzA = a.nonZeros(); @@ -293,6 +292,39 @@ void csr_kron_mult_method(const int imethod, */ const size_t totalPairs = static_cast(nnzA) * static_cast(nnzB); static constexpr size_t kGpuThreshold = 100000; + if (false) { + Kokkos::Profiling::ScopedRegion cpuRegion("PsimgLite::csr_kron_mult_method::imethod3::cpu"); + for (int ia = 0; ia < nrow_A; ++ia) { + const int istart_a = a.getRowPtr(ia); + const int iend_a = a.getRowPtr(ia + 1); + for (int ka = istart_a; ka < iend_a; ++ka) { + const int ja = a.getCol(ka); + ComplexOrRealType aij = a.getValue(ka); + if constexpr (is_complex) + if (isConjTransA) + aij = PsimagLite::conj(aij); + for (int ib = 0; ib < nrow_B; ++ib) { + const int istart_b = b.getRowPtr(ib); + const int iend_b = b.getRowPtr(ib + 1); + for (int kb = istart_b; kb < iend_b; ++kb) { + const int jb = b.getCol(kb); + ComplexOrRealType bij = b.getValue(kb); + if constexpr (is_complex) + if (isConjTransB) + bij = PsimagLite::conj(bij); + + const int ix = (isTransB || isConjTransB) ? jb : ib; + const int jx = (isTransA || isConjTransA) ? ja : ia; + const int iy = (isTransB || isConjTransB) ? ib : jb; + const int jy = (isTransA || isConjTransA) ? ia : ja; + + xout(ix, jx) += aij * bij * yin(iy, jy); + } + } + } + } + return; + } /* * GPU path for large problems. @@ -314,8 +346,6 @@ void csr_kron_mult_method(const int imethod, * With LayoutLeft both the x_dev writes and the y_dev reads are * stride-1 (ib / jb vary across threads, ia / ja fixed per team). */ - Kokkos::Profiling::pushRegion("PsimgLite::csr_kron_mult_method::imethod3::view_init"); - Kokkos::View A_rowptr_h(&a.getRowPtr(0), nrow_A + 1); Kokkos::View @@ -330,23 +360,22 @@ void csr_kron_mult_method(const int imethod, Kokkos::View B_val_h(reinterpret_cast(&b.getValue(0)), nnzB); - auto A_rowptr_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), A_rowptr_h); - auto A_col_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), A_col_h); - auto A_val_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), A_val_h); - auto B_rowptr_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), B_rowptr_h); - auto B_col_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), B_col_h); - auto B_val_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), B_val_h); + auto A_rowptr_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, A_rowptr_h); + auto A_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, A_col_h); + auto A_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, A_val_h); + auto B_rowptr_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, B_rowptr_h); + auto B_col_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, B_col_h); + auto B_val_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, B_val_h); auto yin_host = Kokkos::View( reinterpret_cast(&yin(0, 0)), nrow_Y, ncol_Y); - auto y_dev = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(ExecutionSpace {}, MemorySpace{}), yin_host); + auto y_dev = Kokkos::create_mirror_view_and_copy(ExecutionSpace{}, yin_host); auto x_dev = Kokkos::View("x_dev", nrow_X, ncol_X); - - Kokkos::Profiling::popRegion(); + Kokkos::deep_copy(x_dev, KokkosScalar(0)); using TeamPolicy = Kokkos::TeamPolicy; using TeamMember = typename TeamPolicy::member_type; @@ -359,7 +388,7 @@ void csr_kron_mult_method(const int imethod, */ Kokkos::parallel_for( "csr_kron_mult::imethod3_nn", - TeamPolicy(nrow_A, Kokkos::AUTO, 32), + TeamPolicy(nrow_A, Kokkos::AUTO), KOKKOS_LAMBDA(const TeamMember& team) { const int ia = team.league_rank(); const int ka_begin = A_rowptr_dev(ia); @@ -370,15 +399,13 @@ void csr_kron_mult_method(const int imethod, const int kb_begin = B_rowptr_dev(ib); const int kb_end = B_rowptr_dev(ib + 1); KokkosScalar acc = 0; - Kokkos::parallel_reduce(Kokkos::ThreadVectorRange(team, ka_begin, ka_end), [&](int ka, KokkosScalar& update) { + for (int ka = ka_begin; ka < ka_end; ++ka) { const int ja = A_col_dev(ka); KokkosScalar aij = A_val_dev(ka); for (int kb = kb_begin; kb < kb_end; ++kb) - update += aij * B_val_dev(kb) * y_dev(B_col_dev(kb), ja); - }, acc); -Kokkos::single(Kokkos::PerThread(team), - [&]() { - x_dev(ib, ia) += acc; }); + acc += aij * B_val_dev(kb) * y_dev(B_col_dev(kb), ja); + } + x_dev(ib, ia) += acc; }); }); } else { @@ -421,15 +448,11 @@ Kokkos::single(Kokkos::PerThread(team), }); } -{ - Kokkos::Profiling::ScopedRegion cpuRegion("PsimgLite::csr_kron_mult_method::imethod3::copy_back"); - // copy result back and accumulate into xout auto xhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, x_dev); for (int ix = 0; ix < nrow_X; ++ix) for (int jx = 0; jx < ncol_X; ++jx) xout(ix, jx) += static_cast(xhost(ix, jx)); -} }; } From d3db8e360afb3980f8841d23844dc958b1832b4d Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Wed, 29 Jul 2026 15:47:51 -0400 Subject: [PATCH 21/41] Enable Kokkos implementation --- dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh index 51b6298ab..e9e802f61 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh @@ -1,7 +1,7 @@ #ifndef BATCHEDGEMMINCLUDE_HH #define BATCHEDGEMMINCLUDE_HH #include "DMRGConfig.h" -#ifdef KOKKOS_BATCHED +#if 1 //KOKKOS_BATCHED #include "BatchedGemmKokkos.h" #define BATCHED_GEMM BatchedGemmKokkos #elif defined(PLUGIN_SC) @@ -29,7 +29,7 @@ public: static std::string info() { -#ifdef KOKKOS_BATCHED +#ifdef 1 //KOKKOS_BATCHED return "KokkosKernels"; #elif defined(PLUGIN_SC) return "PLUGIN_SC"; From c3d8302563c33cb94106f460c3cf05d1f065945d Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Fri, 31 Jul 2026 12:08:08 -0400 Subject: [PATCH 22/41] Fused 17s --- .../MatrixVectorKron/BatchedGemmInclude.hh | 2 +- .../MatrixVectorKron/BatchedGemmKokkos.h | 236 ++++++++++++++++-- 2 files changed, 223 insertions(+), 15 deletions(-) diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh index e9e802f61..06e3e6fa6 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh @@ -29,7 +29,7 @@ public: static std::string info() { -#ifdef 1 //KOKKOS_BATCHED +#if 1 //KOKKOS_BATCHED return "KokkosKernels"; #elif defined(PLUGIN_SC) return "PLUGIN_SC"; diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h index 0137b4560..e459592ac 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h @@ -64,6 +64,24 @@ class BatchedGemmKokkos { using DevArgsView = Kokkos::View; + // Extended struct for fused Pass1+Pass2 kernel with scratch memory optimization + struct FusedGemmArgs { + // Pass 1 (Bbatch * X -> BXbatch) + int m1, n1, k1; + int lda1, ldb1, ldc1; + long long a_off1, b_off1, c_off1; + + // Pass 2 (BXbatch * Abatch^T -> Y) + int m2, n2, k2; + int lda2, ldb2, ldc2; + long long a_off2, b_off2, c_off2; + + // Flags + int hasPass1, hasPass2; + }; + + using DevFusedArgsView = Kokkos::View; + static const int ialign_ = 32; public: @@ -104,8 +122,116 @@ class BatchedGemmKokkos { Kokkos::deep_copy(d_vin_, hv); } - // --- Zero output and work buffers ------------------------------------------- + // --- Zero output buffer ------------------------------------------- Kokkos::deep_copy(d_vout_, KS(0)); + + // Use fused kernel if enabled, otherwise fall back to two-pass + if (useFusedKernel_ && nbatch_fused_ > 0) { + matrixVectorFused_(); + } else { + matrixVectorTwoPass_(); + } + + // --- D2H: copy vout back to host ------------------------------------------- + { + using HV = Kokkos::View; + HV hv(reinterpret_cast(vout.data()), totalXY); + Kokkos::deep_copy(hv, d_vout_); + } + } + +private: + + // Optimized fused kernel combining Pass1 and Pass2 with scratch memory + void matrixVectorFused_() const + { + Kokkos::Profiling::ScopedRegion region("BatchedGemmKokkos::matrixVectorFused"); + + const DevScalView flatBbatch = d_flatBbatch_; + const DevScalView flatAbatch = d_flatAbatch_; + const DevScalView vin_dev = d_vin_; + const DevScalView vout_dev = d_vout_; + const DevFusedArgsView args = d_fused_args_; + const size_t max_scratch_size = max_scratch_size_; + + DevExecSpace exec; + + using MemberType = typename Kokkos::TeamPolicy::member_type; + Kokkos::TeamPolicy policy(exec, static_cast(nbatch_fused_), + Kokkos::AUTO, 32); + + if (max_scratch_size > 0) { + policy.set_scratch_size(1, Kokkos::PerTeam(max_scratch_size)); + } + + Kokkos::parallel_for( + "BatchedGemmKokkos_FusedPass1Pass2", + policy, + KOKKOS_LAMBDA(const MemberType& member) { + const int i = member.league_rank(); + const FusedGemmArgs& ag = args(i); + + using UV = Kokkos::View; + using ScratchView = Kokkos::View; + + // --- Pass 1: BXbatch[ip] = Bbatch[ip] * X[jp] --- + ScratchView BXbatch_scratch; + if (ag.hasPass1 && ag.k1 > 0) { + // Allocate BXbatch in scratch memory for this team + BXbatch_scratch = ScratchView(member.team_scratch(1), ag.m1, ag.n1); + + // A = Bbatch[ip] block + UV A_full(flatBbatch.data() + ag.a_off1, ag.lda1, ag.k1); + auto A = Kokkos::subview(A_full, + Kokkos::make_pair(0, ag.m1), Kokkos::ALL()); + + // B = X[jp] + UV B(vin_dev.data() + ag.b_off1, ag.ldb1, ag.n1); + + // C = BXbatch (in scratch memory) + UV C(BXbatch_scratch.data(), ag.m1, ag.n1); + + KokkosBatched::TeamVectorGemm< + MemberType, + KokkosBatched::Trans::NoTranspose, + KokkosBatched::Trans::NoTranspose, + KokkosBatched::Algo::Gemm::Blocked>:: + invoke(member, KS(1), A, B, KS(0), C); + + member.team_barrier(); + } + + // --- Pass 2: Y[ip] = BXbatch[ip] * Abatch[ip]^T --- + if (ag.hasPass2 && ag.k2 > 0 && ag.hasPass1) { + // A = BXbatch (from scratch) + UV A(BXbatch_scratch.data(), ag.m1, ag.n1); + + // B = Abatch[ip] + UV B_full(flatAbatch.data() + ag.b_off2, ag.ldb2, ag.k2); + auto B = Kokkos::subview(B_full, + Kokkos::make_pair(0, ag.n2), Kokkos::ALL()); + + // C = Y[ip] (global output) + UV C_out(vout_dev.data() + ag.c_off2, ag.ldc2, ag.n2); + + KokkosBatched::TeamVectorGemm< + MemberType, + KokkosBatched::Trans::NoTranspose, + KokkosBatched::Trans::Transpose, + KokkosBatched::Algo::Gemm::Blocked>:: + invoke(member, KS(1), A, B, KS(0), C_out); + } + }); + + exec.fence(); + } + + // Fallback two-pass kernel (original implementation) + void matrixVectorTwoPass_() const + { + Kokkos::Profiling::ScopedRegion region("BatchedGemmKokkos::matrixVectorTwoPass"); + Kokkos::deep_copy(d_flatBXbatch_, KS(0)); DevExecSpace exec; @@ -198,17 +324,8 @@ class BatchedGemmKokkos { } exec.fence(); - - // --- D2H: copy vout back to host ------------------------------------------- - { - using HV = Kokkos::View; - HV hv(reinterpret_cast(vout.data()), totalXY); - Kokkos::deep_copy(hv, d_vout_); - } } -private: - static int iceil(int x, int n) { return (x + n - 1) / n; } // Dense matrix reference: pointer + source dimensions (no padding). @@ -247,6 +364,73 @@ class BatchedGemmKokkos { = src[i + static_cast(j) * lds]; } + // Build fused GEMM arguments combining Pass1 and Pass2 per ipatch + void buildFusedArgs_(std::vector& fused_args, + const std::vector& pass1_args, + const std::vector& pass2_args) + { + const SizeType npatches = pass2_args.size(); + + // Group Pass1 GEMMs by their output ipatch (destination in BXbatch) + std::vector> pass1_by_ipatch(npatches); + for (int idx = 0; idx < static_cast(pass1_args.size()); ++idx) { + // Find which ipatch this Pass1 GEMM belongs to based on output offset + // Note: We assume Pass2 args correspond to ipatch order + for (SizeType ip = 0; ip < npatches; ++ip) { + // Simple heuristic: group by checking offset ranges + // In practice, the pass2_args[ip].a_off indicates BXbatch range for ipatch ip + if (pass1_args[idx].c_off == pass2_args[ip].a_off) { + pass1_by_ipatch[ip].push_back(idx); + break; + } + } + } + + // Create one fused GEMM per ipatch + for (SizeType ip = 0; ip < npatches; ++ip) { + FusedGemmArgs fused = {}; + fused.hasPass1 = pass1_by_ipatch[ip].empty() ? 0 : 1; + fused.hasPass2 = (pass2_args[ip].k == 0) ? 0 : 1; + + if (fused.hasPass1 && !pass1_by_ipatch[ip].empty()) { + // Use the first Pass1 GEMM's parameters (they should all have same m,k) + const GemmArgs& ag1 = pass1_args[pass1_by_ipatch[ip][0]]; + fused.m1 = ag1.m; + fused.k1 = ag1.k; + fused.lda1 = ag1.lda; + fused.ldb1 = ag1.ldb; + fused.ldc1 = ag1.ldc; + fused.a_off1 = ag1.a_off; + fused.b_off1 = ag1.b_off; + fused.c_off1 = ag1.c_off; + + // Sum up all n1 values from Pass1 GEMMs for this ipatch + int n1_total = 0; + for (int idx : pass1_by_ipatch[ip]) { + n1_total += pass1_args[idx].n; + } + fused.n1 = n1_total; + } else { + fused.m1 = fused.n1 = fused.k1 = 0; + fused.lda1 = fused.ldb1 = fused.ldc1 = 0; + fused.a_off1 = fused.b_off1 = fused.c_off1 = 0; + } + + // Copy Pass2 GEMM args + fused.m2 = pass2_args[ip].m; + fused.n2 = pass2_args[ip].n; + fused.k2 = pass2_args[ip].k; + fused.lda2 = pass2_args[ip].lda; + fused.ldb2 = pass2_args[ip].ldb; + fused.ldc2 = pass2_args[ip].ldc; + fused.a_off2 = pass2_args[ip].a_off; + fused.b_off2 = pass2_args[ip].b_off; + fused.c_off2 = pass2_args[ip].c_off; + + fused_args.push_back(fused); + } + } + void setup_() { Kokkos::Profiling::ScopedRegion region("BatchedGemmKokkos::setup"); @@ -388,6 +572,19 @@ class BatchedGemmKokkos { pass2_args.push_back(a2); } + // Build fused GEMM arguments combining Pass1 and Pass2 per ipatch + std::vector fused_args; + buildFusedArgs_(fused_args, pass1_args, pass2_args); + + // Calculate max scratch memory needed for fused kernel + max_scratch_size_ = 0; + for (const auto& fg : fused_args) { + if (fg.hasPass1 && fg.m1 > 0 && fg.n1 > 0) { + size_t bx_size = static_cast(fg.m1) * static_cast(fg.n1) * sizeof(KS); + max_scratch_size_ = (bx_size > max_scratch_size_) ? bx_size : max_scratch_size_; + } + } + // Allocate device arrays and upload operator matrices. d_flatAbatch_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_flatAbatch"), totalAbatch ? totalAbatch : 1); d_flatBbatch_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_flatBbatch"), totalBbatch ? totalBbatch : 1); @@ -397,6 +594,7 @@ class BatchedGemmKokkos { nbatch1_ = pass1_args.size(); nbatch2_ = pass2_args.size(); // == npatches + nbatch_fused_ = fused_args.size(); { auto hA = Kokkos::create_mirror_view(Kokkos::view_alloc(Kokkos::WithoutInitializing), d_flatAbatch_); @@ -411,13 +609,17 @@ class BatchedGemmKokkos { d_pass1_ = DevArgsView("d_pass1", nbatch1_ ? nbatch1_ : 1); d_pass2_ = DevArgsView("d_pass2", nbatch2_ ? nbatch2_ : 1); + d_fused_args_ = DevFusedArgsView("d_fused_args", nbatch_fused_ ? nbatch_fused_ : 1); { auto h1 = Kokkos::create_mirror_view(d_pass1_); auto h2 = Kokkos::create_mirror_view(d_pass2_); + auto hf = Kokkos::create_mirror_view(d_fused_args_); for (SizeType i = 0; i < nbatch1_; ++i) h1(i) = pass1_args[i]; for (SizeType i = 0; i < nbatch2_; ++i) h2(i) = pass2_args[i]; + for (SizeType i = 0; i < nbatch_fused_; ++i) hf(i) = fused_args[i]; Kokkos::deep_copy(d_pass1_, h1); Kokkos::deep_copy(d_pass2_, h2); + Kokkos::deep_copy(d_fused_args_, hf); } { @@ -426,9 +628,10 @@ class BatchedGemmKokkos { << " noperator=" << noperator << " pass1_batches=" << nbatch1_ << " pass2_batches=" << nbatch2_ + << " fused_batches=" << nbatch_fused_ << " Abatch=" << totalAbatch << "elems" << " Bbatch=" << totalBbatch << "elems" - << " BXbatch=" << totalBXbatch << "elems"; + << " BXbatch=" << totalBXbatch << "elems (can be on-device scratch)"; progress_.printline(msg, std::cout); } } @@ -439,16 +642,21 @@ class BatchedGemmKokkos { PsimagLite::ProgressIndicator progress_; mutable VectorMatrixType garbage_; // owns sparse->dense expansions (freed in dtor) - SizeType nbatch1_ = 0; // number of pass-1 GEMMs - SizeType nbatch2_ = 0; // number of pass-2 GEMMs (== npatches) + SizeType nbatch1_ = 0; // number of pass-1 GEMMs + SizeType nbatch2_ = 0; // number of pass-2 GEMMs (== npatches) + SizeType nbatch_fused_ = 0; // number of fused GEMM groups + size_t max_scratch_size_ = 0; // max scratch memory needed for fused kernel + + bool useFusedKernel_ = true; // Enable fused kernel optimization DevScalView d_flatAbatch_; // left-operator matrices, device, persistent DevScalView d_flatBbatch_; // right-operator matrices, device, persistent - mutable DevScalView d_flatBXbatch_; // intermediate BX work buffer, per-call + mutable DevScalView d_flatBXbatch_; // intermediate BX work buffer (for two-pass mode) mutable DevScalView d_vin_; // input vector on device, per-call mutable DevScalView d_vout_; // output vector on device, per-call DevArgsView d_pass1_; // pass-1 GEMM parameters, persistent DevArgsView d_pass2_; // pass-2 GEMM parameters, persistent + DevFusedArgsView d_fused_args_; // fused GEMM parameters, persistent }; } // namespace Dmrg From 682b0ce31c25cb30859e9a3863b3a115a1239898 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Fri, 31 Jul 2026 12:21:44 -0400 Subject: [PATCH 23/41] Kokkos::AUTO also works --- dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h index e459592ac..8d814b314 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h @@ -158,7 +158,7 @@ class BatchedGemmKokkos { using MemberType = typename Kokkos::TeamPolicy::member_type; Kokkos::TeamPolicy policy(exec, static_cast(nbatch_fused_), - Kokkos::AUTO, 32); + Kokkos::AUTO, Kokkos::AUTO); if (max_scratch_size > 0) { policy.set_scratch_size(1, Kokkos::PerTeam(max_scratch_size)); From 076a81da90ba31b55e7e6e895cca36625fd9d14c Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Fri, 31 Jul 2026 12:22:24 -0400 Subject: [PATCH 24/41] Revert "Kokkos::AUTO also works" This reverts commit 682b0ce31c25cb30859e9a3863b3a115a1239898. --- dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h index 8d814b314..e459592ac 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h @@ -158,7 +158,7 @@ class BatchedGemmKokkos { using MemberType = typename Kokkos::TeamPolicy::member_type; Kokkos::TeamPolicy policy(exec, static_cast(nbatch_fused_), - Kokkos::AUTO, Kokkos::AUTO); + Kokkos::AUTO, 32); if (max_scratch_size > 0) { policy.set_scratch_size(1, Kokkos::PerTeam(max_scratch_size)); From 2fc457cbee47860f440a9e7b2d61c60510d91d3c Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Fri, 31 Jul 2026 12:22:37 -0400 Subject: [PATCH 25/41] Revert "Fused 17s" This reverts commit c3d8302563c33cb94106f460c3cf05d1f065945d. --- .../MatrixVectorKron/BatchedGemmInclude.hh | 2 +- .../MatrixVectorKron/BatchedGemmKokkos.h | 236 ++---------------- 2 files changed, 15 insertions(+), 223 deletions(-) diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh index 06e3e6fa6..e9e802f61 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh @@ -29,7 +29,7 @@ public: static std::string info() { -#if 1 //KOKKOS_BATCHED +#ifdef 1 //KOKKOS_BATCHED return "KokkosKernels"; #elif defined(PLUGIN_SC) return "PLUGIN_SC"; diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h index e459592ac..0137b4560 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h @@ -64,24 +64,6 @@ class BatchedGemmKokkos { using DevArgsView = Kokkos::View; - // Extended struct for fused Pass1+Pass2 kernel with scratch memory optimization - struct FusedGemmArgs { - // Pass 1 (Bbatch * X -> BXbatch) - int m1, n1, k1; - int lda1, ldb1, ldc1; - long long a_off1, b_off1, c_off1; - - // Pass 2 (BXbatch * Abatch^T -> Y) - int m2, n2, k2; - int lda2, ldb2, ldc2; - long long a_off2, b_off2, c_off2; - - // Flags - int hasPass1, hasPass2; - }; - - using DevFusedArgsView = Kokkos::View; - static const int ialign_ = 32; public: @@ -122,116 +104,8 @@ class BatchedGemmKokkos { Kokkos::deep_copy(d_vin_, hv); } - // --- Zero output buffer ------------------------------------------- + // --- Zero output and work buffers ------------------------------------------- Kokkos::deep_copy(d_vout_, KS(0)); - - // Use fused kernel if enabled, otherwise fall back to two-pass - if (useFusedKernel_ && nbatch_fused_ > 0) { - matrixVectorFused_(); - } else { - matrixVectorTwoPass_(); - } - - // --- D2H: copy vout back to host ------------------------------------------- - { - using HV = Kokkos::View; - HV hv(reinterpret_cast(vout.data()), totalXY); - Kokkos::deep_copy(hv, d_vout_); - } - } - -private: - - // Optimized fused kernel combining Pass1 and Pass2 with scratch memory - void matrixVectorFused_() const - { - Kokkos::Profiling::ScopedRegion region("BatchedGemmKokkos::matrixVectorFused"); - - const DevScalView flatBbatch = d_flatBbatch_; - const DevScalView flatAbatch = d_flatAbatch_; - const DevScalView vin_dev = d_vin_; - const DevScalView vout_dev = d_vout_; - const DevFusedArgsView args = d_fused_args_; - const size_t max_scratch_size = max_scratch_size_; - - DevExecSpace exec; - - using MemberType = typename Kokkos::TeamPolicy::member_type; - Kokkos::TeamPolicy policy(exec, static_cast(nbatch_fused_), - Kokkos::AUTO, 32); - - if (max_scratch_size > 0) { - policy.set_scratch_size(1, Kokkos::PerTeam(max_scratch_size)); - } - - Kokkos::parallel_for( - "BatchedGemmKokkos_FusedPass1Pass2", - policy, - KOKKOS_LAMBDA(const MemberType& member) { - const int i = member.league_rank(); - const FusedGemmArgs& ag = args(i); - - using UV = Kokkos::View; - using ScratchView = Kokkos::View; - - // --- Pass 1: BXbatch[ip] = Bbatch[ip] * X[jp] --- - ScratchView BXbatch_scratch; - if (ag.hasPass1 && ag.k1 > 0) { - // Allocate BXbatch in scratch memory for this team - BXbatch_scratch = ScratchView(member.team_scratch(1), ag.m1, ag.n1); - - // A = Bbatch[ip] block - UV A_full(flatBbatch.data() + ag.a_off1, ag.lda1, ag.k1); - auto A = Kokkos::subview(A_full, - Kokkos::make_pair(0, ag.m1), Kokkos::ALL()); - - // B = X[jp] - UV B(vin_dev.data() + ag.b_off1, ag.ldb1, ag.n1); - - // C = BXbatch (in scratch memory) - UV C(BXbatch_scratch.data(), ag.m1, ag.n1); - - KokkosBatched::TeamVectorGemm< - MemberType, - KokkosBatched::Trans::NoTranspose, - KokkosBatched::Trans::NoTranspose, - KokkosBatched::Algo::Gemm::Blocked>:: - invoke(member, KS(1), A, B, KS(0), C); - - member.team_barrier(); - } - - // --- Pass 2: Y[ip] = BXbatch[ip] * Abatch[ip]^T --- - if (ag.hasPass2 && ag.k2 > 0 && ag.hasPass1) { - // A = BXbatch (from scratch) - UV A(BXbatch_scratch.data(), ag.m1, ag.n1); - - // B = Abatch[ip] - UV B_full(flatAbatch.data() + ag.b_off2, ag.ldb2, ag.k2); - auto B = Kokkos::subview(B_full, - Kokkos::make_pair(0, ag.n2), Kokkos::ALL()); - - // C = Y[ip] (global output) - UV C_out(vout_dev.data() + ag.c_off2, ag.ldc2, ag.n2); - - KokkosBatched::TeamVectorGemm< - MemberType, - KokkosBatched::Trans::NoTranspose, - KokkosBatched::Trans::Transpose, - KokkosBatched::Algo::Gemm::Blocked>:: - invoke(member, KS(1), A, B, KS(0), C_out); - } - }); - - exec.fence(); - } - - // Fallback two-pass kernel (original implementation) - void matrixVectorTwoPass_() const - { - Kokkos::Profiling::ScopedRegion region("BatchedGemmKokkos::matrixVectorTwoPass"); - Kokkos::deep_copy(d_flatBXbatch_, KS(0)); DevExecSpace exec; @@ -324,8 +198,17 @@ class BatchedGemmKokkos { } exec.fence(); + + // --- D2H: copy vout back to host ------------------------------------------- + { + using HV = Kokkos::View; + HV hv(reinterpret_cast(vout.data()), totalXY); + Kokkos::deep_copy(hv, d_vout_); + } } +private: + static int iceil(int x, int n) { return (x + n - 1) / n; } // Dense matrix reference: pointer + source dimensions (no padding). @@ -364,73 +247,6 @@ class BatchedGemmKokkos { = src[i + static_cast(j) * lds]; } - // Build fused GEMM arguments combining Pass1 and Pass2 per ipatch - void buildFusedArgs_(std::vector& fused_args, - const std::vector& pass1_args, - const std::vector& pass2_args) - { - const SizeType npatches = pass2_args.size(); - - // Group Pass1 GEMMs by their output ipatch (destination in BXbatch) - std::vector> pass1_by_ipatch(npatches); - for (int idx = 0; idx < static_cast(pass1_args.size()); ++idx) { - // Find which ipatch this Pass1 GEMM belongs to based on output offset - // Note: We assume Pass2 args correspond to ipatch order - for (SizeType ip = 0; ip < npatches; ++ip) { - // Simple heuristic: group by checking offset ranges - // In practice, the pass2_args[ip].a_off indicates BXbatch range for ipatch ip - if (pass1_args[idx].c_off == pass2_args[ip].a_off) { - pass1_by_ipatch[ip].push_back(idx); - break; - } - } - } - - // Create one fused GEMM per ipatch - for (SizeType ip = 0; ip < npatches; ++ip) { - FusedGemmArgs fused = {}; - fused.hasPass1 = pass1_by_ipatch[ip].empty() ? 0 : 1; - fused.hasPass2 = (pass2_args[ip].k == 0) ? 0 : 1; - - if (fused.hasPass1 && !pass1_by_ipatch[ip].empty()) { - // Use the first Pass1 GEMM's parameters (they should all have same m,k) - const GemmArgs& ag1 = pass1_args[pass1_by_ipatch[ip][0]]; - fused.m1 = ag1.m; - fused.k1 = ag1.k; - fused.lda1 = ag1.lda; - fused.ldb1 = ag1.ldb; - fused.ldc1 = ag1.ldc; - fused.a_off1 = ag1.a_off; - fused.b_off1 = ag1.b_off; - fused.c_off1 = ag1.c_off; - - // Sum up all n1 values from Pass1 GEMMs for this ipatch - int n1_total = 0; - for (int idx : pass1_by_ipatch[ip]) { - n1_total += pass1_args[idx].n; - } - fused.n1 = n1_total; - } else { - fused.m1 = fused.n1 = fused.k1 = 0; - fused.lda1 = fused.ldb1 = fused.ldc1 = 0; - fused.a_off1 = fused.b_off1 = fused.c_off1 = 0; - } - - // Copy Pass2 GEMM args - fused.m2 = pass2_args[ip].m; - fused.n2 = pass2_args[ip].n; - fused.k2 = pass2_args[ip].k; - fused.lda2 = pass2_args[ip].lda; - fused.ldb2 = pass2_args[ip].ldb; - fused.ldc2 = pass2_args[ip].ldc; - fused.a_off2 = pass2_args[ip].a_off; - fused.b_off2 = pass2_args[ip].b_off; - fused.c_off2 = pass2_args[ip].c_off; - - fused_args.push_back(fused); - } - } - void setup_() { Kokkos::Profiling::ScopedRegion region("BatchedGemmKokkos::setup"); @@ -572,19 +388,6 @@ class BatchedGemmKokkos { pass2_args.push_back(a2); } - // Build fused GEMM arguments combining Pass1 and Pass2 per ipatch - std::vector fused_args; - buildFusedArgs_(fused_args, pass1_args, pass2_args); - - // Calculate max scratch memory needed for fused kernel - max_scratch_size_ = 0; - for (const auto& fg : fused_args) { - if (fg.hasPass1 && fg.m1 > 0 && fg.n1 > 0) { - size_t bx_size = static_cast(fg.m1) * static_cast(fg.n1) * sizeof(KS); - max_scratch_size_ = (bx_size > max_scratch_size_) ? bx_size : max_scratch_size_; - } - } - // Allocate device arrays and upload operator matrices. d_flatAbatch_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_flatAbatch"), totalAbatch ? totalAbatch : 1); d_flatBbatch_ = DevScalView(Kokkos::view_alloc(Kokkos::WithoutInitializing, "d_flatBbatch"), totalBbatch ? totalBbatch : 1); @@ -594,7 +397,6 @@ class BatchedGemmKokkos { nbatch1_ = pass1_args.size(); nbatch2_ = pass2_args.size(); // == npatches - nbatch_fused_ = fused_args.size(); { auto hA = Kokkos::create_mirror_view(Kokkos::view_alloc(Kokkos::WithoutInitializing), d_flatAbatch_); @@ -609,17 +411,13 @@ class BatchedGemmKokkos { d_pass1_ = DevArgsView("d_pass1", nbatch1_ ? nbatch1_ : 1); d_pass2_ = DevArgsView("d_pass2", nbatch2_ ? nbatch2_ : 1); - d_fused_args_ = DevFusedArgsView("d_fused_args", nbatch_fused_ ? nbatch_fused_ : 1); { auto h1 = Kokkos::create_mirror_view(d_pass1_); auto h2 = Kokkos::create_mirror_view(d_pass2_); - auto hf = Kokkos::create_mirror_view(d_fused_args_); for (SizeType i = 0; i < nbatch1_; ++i) h1(i) = pass1_args[i]; for (SizeType i = 0; i < nbatch2_; ++i) h2(i) = pass2_args[i]; - for (SizeType i = 0; i < nbatch_fused_; ++i) hf(i) = fused_args[i]; Kokkos::deep_copy(d_pass1_, h1); Kokkos::deep_copy(d_pass2_, h2); - Kokkos::deep_copy(d_fused_args_, hf); } { @@ -628,10 +426,9 @@ class BatchedGemmKokkos { << " noperator=" << noperator << " pass1_batches=" << nbatch1_ << " pass2_batches=" << nbatch2_ - << " fused_batches=" << nbatch_fused_ << " Abatch=" << totalAbatch << "elems" << " Bbatch=" << totalBbatch << "elems" - << " BXbatch=" << totalBXbatch << "elems (can be on-device scratch)"; + << " BXbatch=" << totalBXbatch << "elems"; progress_.printline(msg, std::cout); } } @@ -642,21 +439,16 @@ class BatchedGemmKokkos { PsimagLite::ProgressIndicator progress_; mutable VectorMatrixType garbage_; // owns sparse->dense expansions (freed in dtor) - SizeType nbatch1_ = 0; // number of pass-1 GEMMs - SizeType nbatch2_ = 0; // number of pass-2 GEMMs (== npatches) - SizeType nbatch_fused_ = 0; // number of fused GEMM groups - size_t max_scratch_size_ = 0; // max scratch memory needed for fused kernel - - bool useFusedKernel_ = true; // Enable fused kernel optimization + SizeType nbatch1_ = 0; // number of pass-1 GEMMs + SizeType nbatch2_ = 0; // number of pass-2 GEMMs (== npatches) DevScalView d_flatAbatch_; // left-operator matrices, device, persistent DevScalView d_flatBbatch_; // right-operator matrices, device, persistent - mutable DevScalView d_flatBXbatch_; // intermediate BX work buffer (for two-pass mode) + mutable DevScalView d_flatBXbatch_; // intermediate BX work buffer, per-call mutable DevScalView d_vin_; // input vector on device, per-call mutable DevScalView d_vout_; // output vector on device, per-call DevArgsView d_pass1_; // pass-1 GEMM parameters, persistent DevArgsView d_pass2_; // pass-2 GEMM parameters, persistent - DevFusedArgsView d_fused_args_; // fused GEMM parameters, persistent }; } // namespace Dmrg From e1eb9018f9f9b4d9eade4340d55e4fb8345e2608 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Fri, 31 Jul 2026 13:32:52 -0400 Subject: [PATCH 26/41] 16s with two pass --- dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh | 2 +- dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh index e9e802f61..06e3e6fa6 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh @@ -29,7 +29,7 @@ public: static std::string info() { -#ifdef 1 //KOKKOS_BATCHED +#if 1 //KOKKOS_BATCHED return "KokkosKernels"; #elif defined(PLUGIN_SC) return "PLUGIN_SC"; diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h index 0137b4560..a508d5727 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h @@ -146,7 +146,7 @@ class BatchedGemmKokkos { MemberType, KokkosBatched::Trans::NoTranspose, KokkosBatched::Trans::NoTranspose, - KokkosBatched::Algo::Gemm::Unblocked>:: + KokkosBatched::Algo::Gemm::Blocked>:: invoke(member, KS(1), A, B, KS(0), C); }); } @@ -192,7 +192,7 @@ class BatchedGemmKokkos { MemberType, KokkosBatched::Trans::NoTranspose, KokkosBatched::Trans::Transpose, - KokkosBatched::Algo::Gemm::Unblocked>:: + KokkosBatched::Algo::Gemm::Blocked>:: invoke(member, KS(1), A, B, KS(0), C); }); } From a3fc791a92352248f2b0a0c9d2ab19b8ea7f1c2b Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Fri, 31 Jul 2026 16:00:40 -0400 Subject: [PATCH 27/41] Fix fences --- dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h index a508d5727..f387b71d4 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmKokkos.h @@ -151,8 +151,6 @@ class BatchedGemmKokkos { }); } - exec.fence(); - // --- Pass 2: Y[ip] = BXbatch[ip] * Abatch[ip]^T (NoT x T) --------------- { const DevScalView flatBXbatch = d_flatBXbatch_; @@ -197,13 +195,12 @@ class BatchedGemmKokkos { }); } - exec.fence(); - // --- D2H: copy vout back to host ------------------------------------------- { using HV = Kokkos::View; HV hv(reinterpret_cast(vout.data()), totalXY); - Kokkos::deep_copy(hv, d_vout_); + Kokkos::deep_copy(exec, hv, d_vout_); + exec.fence(); } } From 2036da8d4f453e9ac41b584487dbbf748447e805 Mon Sep 17 00:00:00 2001 From: Daniel Arndt Date: Mon, 3 Aug 2026 10:29:11 -0400 Subject: [PATCH 28/41] Test BatchedGemm implementations --- PsimagLite/src/PsimagLite/KokkosGemm.cpp | 34 +- .../MatrixVectorKron/BatchedGemmInclude.hh | 17 +- dmrg/Engine/tests/CMakeLists.txt | 10 + dmrg/Engine/tests/test_BatchedGemmKokkos.cpp | 341 ++++++++++++++++++ 4 files changed, 385 insertions(+), 17 deletions(-) create mode 100644 dmrg/Engine/tests/test_BatchedGemmKokkos.cpp diff --git a/PsimagLite/src/PsimagLite/KokkosGemm.cpp b/PsimagLite/src/PsimagLite/KokkosGemm.cpp index 26b43349a..03143c174 100644 --- a/PsimagLite/src/PsimagLite/KokkosGemm.cpp +++ b/PsimagLite/src/PsimagLite/KokkosGemm.cpp @@ -52,35 +52,47 @@ inline void PsimagLite::kokkos_gemm(char transa, Kokkos::DefaultExecutionSpace exec; decltype(exec)::memory_space mem; - KOKKOS_ASSERT(ldaVal == (ta == 'N' ? M : K)); + // allow padded leading dimensions (ldaVal/ldbVal/ldcVal >= required) + if (ldaVal < req_lda || ldbVal < req_ldb || ldcVal < req_ldc) { + throw std::runtime_error("kokkos_gemm: invalid leading dimension"); + } + + // Create host unmanaged views that reflect the actual storage (use lda/ldb as the first + // extent) Kokkos::View - Aview_op( - reinterpret_cast(A), ta == 'N' ? M : K, ta == 'N' ? K : M); + Aview_op(reinterpret_cast(A), ldaVal, (ta == 'N' ? K : M)); auto Aview_op_device = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(exec, mem), Aview_op); - KOKKOS_ASSERT(ldbVal == (tb == 'N' ? K : N)); Kokkos::View - Bview_op( - reinterpret_cast(B), tb == 'N' ? K : N, tb == 'N' ? N : K); + Bview_op(reinterpret_cast(B), ldbVal, (tb == 'N' ? N : K)); auto Bview_op_device = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(exec, mem), Bview_op); + // Create C view that reflects storage with possible padding (ldcVal >= M) Kokkos::View - Cview(reinterpret_cast(C), M, N); + Cview(reinterpret_cast(C), ldcVal, N); auto Cview_device = Kokkos::create_mirror_view_and_copy(Kokkos::view_alloc(exec, mem), Cview); - const char transA[2] = { ta, '\0' }; - const char transB[2] = { tb, '\0' }; - KokkosBlas::gemm( - exec, transA, transB, alpha, Aview_op_device, Bview_op_device, beta, Cview_device); + // Create subviews that present the logical matrix sizes (m x k for A if not transposed, + // etc.) + using Pair = Kokkos::pair; + auto Aop = (ta == 'N') ? Kokkos::subview(Aview_op_device, Pair(0, M), Pair(0, K)) + : Kokkos::subview(Aview_op_device, Pair(0, K), Pair(0, M)); + auto Bop = (tb == 'N') ? Kokkos::subview(Bview_op_device, Pair(0, K), Pair(0, N)) + : Kokkos::subview(Bview_op_device, Pair(0, N), Pair(0, K)); + auto Cop = Kokkos::subview(Cview_device, Pair(0, M), Pair(0, N)); + + const char transA2[2] = { ta, '\0' }; + const char transB2[2] = { tb, '\0' }; + KokkosBlas::gemm(exec, transA2, transB2, alpha, Aop, Bop, beta, Cop); Kokkos::deep_copy(exec, Cview, Cview_device); exec.fence(); } diff --git a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh index 1fa802098..28ac46a0a 100644 --- a/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh +++ b/dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh @@ -1,12 +1,15 @@ #ifndef BATCHEDGEMMINCLUDE_HH #define BATCHEDGEMMINCLUDE_HH #include "DMRGConfig.h" -#ifdef PLUGIN_SC +#if 0 // KOKKOS_BATCHED +#include "BatchedGemmKokkos.h" +#define BATCHED_GEMM ::Dmrg::BatchedGemmKokkos +#elif defined(PLUGIN_SC) #include "BatchedGemmPluginSc.h" -#define BATCHED_GEMM BatchedGemmPluginSc +#define BATCHED_GEMM ::Dmrg::BatchedGemmPluginSc #else #include "BatchedGemmCpu.h" -#define BATCHED_GEMM BatchedGemmCpu +#define BATCHED_GEMM ::Dmrg::BatchedGemmCpu #endif #include @@ -18,15 +21,17 @@ public: static void failIfNotSupported() { -#ifdef PLUGIN_SC +#if defined(KOKKOS_BATCHED) || defined(PLUGIN_SC) return; #endif - err("BatchedGemm needs -DPLUGIN_SC in Config.make\n"); + err("BatchedGemm needs DMRG_BUILD_BATCHED_KOKKOS=ON or -DPLUGIN_SC\n"); } static std::string info() { -#ifdef PLUGIN_SC +#if 1 // KOKKOS_BATCHED + return "KokkosKernels"; +#elif defined(PLUGIN_SC) return "PLUGIN_SC"; #else return ""; diff --git a/dmrg/Engine/tests/CMakeLists.txt b/dmrg/Engine/tests/CMakeLists.txt index 35b69c010..68427de12 100644 --- a/dmrg/Engine/tests/CMakeLists.txt +++ b/dmrg/Engine/tests/CMakeLists.txt @@ -2,3 +2,13 @@ add_executable(test_LastKrylovSlots test_LastKrylovSlots.cpp) target_link_libraries(test_LastKrylovSlots PRIVATE Catch2::Catch2WithMain psimaglite) target_include_directories(test_LastKrylovSlots PRIVATE ${CMAKE_SOURCE_DIR}/dmrg/Engine) catch_discover_tests(test_LastKrylovSlots) + +# BatchedGemmKokkos test +add_executable(test_BatchedGemmKokkos test_BatchedGemmKokkos.cpp) +# Link Kokkos and KokkosKernels since the test initializes Kokkos and uses KokkosBatched +target_link_libraries(test_BatchedGemmKokkos PRIVATE Catch2::Catch2WithMain dmrgpp_utils psimaglite Kokkos::kokkos + KokkosKernels::kokkoskernels gpuplugin) +# Ensure headers from Engine and the MatrixVectorKron subdir are visible +target_include_directories(test_BatchedGemmKokkos PRIVATE ${CMAKE_SOURCE_DIR}/dmrg/Engine + ${CMAKE_SOURCE_DIR}/dmrg/Engine/MatrixVectorKron) +catch_discover_tests(test_BatchedGemmKokkos) diff --git a/dmrg/Engine/tests/test_BatchedGemmKokkos.cpp b/dmrg/Engine/tests/test_BatchedGemmKokkos.cpp new file mode 100644 index 000000000..111637434 --- /dev/null +++ b/dmrg/Engine/tests/test_BatchedGemmKokkos.cpp @@ -0,0 +1,341 @@ +// #include "BatchedGemmKokkos.h" +#include "BatchedGemmCpu.h" +#include "BatchedGemmPluginSc.h" +#include +#include +#include +#include +#include +#include +#include + +using SizeType = unsigned long; + +// Minimal fake types to instantiate BatchedGemmKokkos for a 1-patch, 1-operator case +struct FakeParams { + struct Opt { + bool isSet(const std::string& s) const { return s == "BatchedGemm"; } + } options; +}; + +struct FakeLR { + std::vector parts; // partition points, length = groups+1 + int partition(size_t i) const { return parts[i]; } + int size() const { return parts.back(); } +}; + +struct FakeLrs { + FakeLR l, r; + const FakeLR& left() const { return l; } + const FakeLR& right() const { return r; } +}; + +// Matrix wrapper used by BatchedGemmKokkos via ArrayOfMatStructType. +struct FakeMatrixDenseOrSparse { + using value_type = double; + using VectorType = PsimagLite::Vector::Type; + using MatrixType = PsimagLite::Matrix; + using CrsType = PsimagLite::CrsMatrix; + MatrixType mat; + CrsType crs; + bool zero = false; + FakeMatrixDenseOrSparse() + : mat() + , crs(0, 0) + , zero(true) + { } + FakeMatrixDenseOrSparse(const MatrixType& m) + : mat(m) + , crs(0, 0) + , zero(false) + { } + bool isZero() const { return zero; } + bool isDense() const { return true; } + const MatrixType& dense() const { return mat; } + const CrsType& sparse() const { return crs; } + int rows() const { return mat.rows(); } + int cols() const { return mat.cols(); } +}; + +// ArrayOfMatStructType: callable (ip,jp) -> pointer to FakeMatrixDenseOrSparse +struct FakeArrayOfMatStructType { + using MatrixDenseOrSparseType = FakeMatrixDenseOrSparse; + + std::vector> storage; // [ip][jp] + FakeArrayOfMatStructType(size_t npatches = 0) + { + storage.resize(npatches); + for (size_t i = 0; i < npatches; ++i) + storage[i].resize(npatches); + } + const FakeMatrixDenseOrSparse* operator()(size_t ip, size_t jp) const + { + // If matrix is default-constructed, treat as zero -> return nullptr + if (storage[ip][jp].isZero()) + return nullptr; + return &storage[ip][jp]; + } + FakeMatrixDenseOrSparse* operator()(size_t ip, size_t jp) + { + if (storage[ip][jp].isZero()) + return nullptr; + return &storage[ip][jp]; + } + void set(size_t ip, size_t jp, const FakeMatrixDenseOrSparse& m) { storage[ip][jp] = m; } +}; + +struct GenIjPatch { + enum LeftOrRightEnumType + { + LEFT = 0, + RIGHT = 1 + }; + using BasisType = int; +}; + +struct FakeInitKron { + enum WhatBasisEnum + { + OLD = 0, + NEW = 1 + }; + + using ArrayOfMatStructType = FakeArrayOfMatStructType; + using GenIjPatchType = GenIjPatch; + using SparseMatrixType = PsimagLite::CrsMatrix; + + FakeParams p; + FakeArrayOfMatStructType xc0, yc0; + FakeLrs lrs_; // used for both NEW/OLD in this fake + std::vector patchLeft, patchRight; + size_t npatches_ = 0; + size_t noperator_ = 0; + + FakeInitKron(size_t npatches, size_t noperator) + : p() + , xc0(npatches) + , yc0(npatches) + , npatches_(npatches) + , noperator_(noperator) + { + // default partitions: single group covering full size; user will set lrs_.l.parts + // and lrs_.r.parts + patchLeft.resize(npatches_, 0); + patchRight.resize(npatches_, 0); + } + + const FakeParams& params() const { return p; } + + size_t numberOfPatches(int /*which*/) const { return npatches_; } + size_t connections() const { return noperator_; } + + // Return patch group indices + const std::vector& patch(WhatBasisEnum /*which*/, + GenIjPatch::LeftOrRightEnumType side) const + { + if (side == GenIjPatch::LEFT) + return patchLeft; + return patchRight; + } + + const FakeLrs& lrs(WhatBasisEnum /*which*/) const { return lrs_; } + + // offsetForPatches: cumulative offsets in vin/vout per patch + SizeType offsetForPatches(WhatBasisEnum /*which*/, size_t ipatch) const + { + SizeType off = 0; + for (size_t p = 0; p < ipatch; ++p) { + size_t igroup = patchLeft[p]; + size_t jgroup = patchRight[p]; + int L1 = lrs_.left().partition(igroup); + int L2 = lrs_.left().partition(igroup + 1); + int R1 = lrs_.right().partition(jgroup); + int R2 = lrs_.right().partition(jgroup + 1); + off += static_cast((L2 - L1) * (R2 - R1)); + } + return off; + } + + // Accessors for matrices: expect xc(k) and yc(k) to provide (ip,jp) access + const FakeArrayOfMatStructType& xc(size_t /*k*/) const { return xc0; } + const FakeArrayOfMatStructType& yc(size_t /*k*/) const { return yc0; } + + void + checks(const FakeMatrixDenseOrSparse&, const FakeMatrixDenseOrSparse&, size_t, size_t) const + { } +}; + +template