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)); } }