Port csr_matmul_post.cpp to Kokkos

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)

PsimagLite/src/PsimagLite/BLAS.h

@@ -10,7 +10,7 @@
 #define PSIMAG_BLAS
 #include "AllocatorCpu.h"
 
-#include <PsimagLite/kokkos_gemm.h>
+#include <PsimagLite/KokkosGemm.h>
 
 #include <complex>
 

PsimagLite/src/PsimagLite/kokkos_gemm.cpp → PsimagLite/src/PsimagLite/KokkosGemm.cpp

@@ -5,24 +5,8 @@
 #include <stdexcept>
 #include <type_traits>
 
-#include <PsimagLite/kokkos_gemm.h>
-
-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 <typename T> struct KokkosType {
-	using type = T;
-};
-
-template <typename T>
-        requires(!std::is_floating_point_v<T>)
-struct KokkosType<T> {
-	using type = Kokkos::complex<typename T::value_type>;
-};
-
-}
+#include <PsimagLite/KokkosGemm.h>
+#include <PsimagLite/KokkosType.h>
 
 template <typename Scalar, typename IntegerForBlasType>
 inline void PsimagLite::kokkos_gemm(char               transa,

PsimagLite/src/PsimagLite/KokkosType.h

@@ -0,0 +1,23 @@
+#ifndef PSIMAG_KOKKOS_TYPE_H
+#define PSIMAG_KOKKOS_TYPE_H
+
+#include <Kokkos_Complex.hpp>
+
+#include <type_traits>
+
+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 <typename T> struct KokkosType {
+	using type = T;
+};
+
+template <typename T> struct KokkosType<std::complex<T>> {
+	using type = Kokkos::complex<T>;
+};
+
+}
+
+#endif // PSIMAG_KOKKOS_TYPE_H

dmrg/KronUtil/csr_matmul_post.cpp

@@ -1,4 +1,8 @@
 #include "util.h"
+#include <PsimagLite/KokkosType.h>
+
+#include <Kokkos_Core.hpp>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
 
 template <typename ComplexOrRealType>
 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<ComplexOrRealType>::type;
+
+	const int nnz = a.nonZeros();
+
+	Kokkos::View<const int*, Kokkos::HostSpace, Kokkos::MemoryUnmanaged> rowptr_host(
+	    &a.getRowPtr(0), nrow_A + 1);
+	Kokkos::View<const int*, Kokkos::HostSpace, Kokkos::MemoryUnmanaged> cols_host(&a.getCol(0),
+	                                                                               nnz);
+	Kokkos::View<const KokkosScalar*, Kokkos::HostSpace, Kokkos::MemoryUnmanaged> vals_host(
+	    reinterpret_cast<const KokkosScalar*>(&a.getValue(0)), nnz);
+
+	Kokkos::View<KokkosScalar**> x_dev_out("x_dev_out", nrow_Y, ncol_X);
+	Kokkos::View<const KokkosScalar**,
+	             Kokkos::LayoutLeft,
+	             Kokkos::HostSpace,
+	             Kokkos::MemoryUnmanaged>
+	     yin_host(reinterpret_cast<const KokkosScalar*>(&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<ExecutionSpace>;
+	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<SizeType>(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<SizeType>(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<ComplexOrRealType>(xhost(iy, jx));
 	}
 }