Add a batched gemm Kokkos implementation

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,
@@ -39,7 +23,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<int>(m);
 	int                             N      = static_cast<int>(n);
 	int                             K      = static_cast<int>(k);

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

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;
+

dmrg/CMakeLists.txt

@@ -171,6 +171,8 @@ if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") # The following tests are expensive to
   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
   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)

dmrg/Engine/MatrixVectorKron/BatchedGemmInclude.hh

@@ -1,7 +1,10 @@
 #ifndef BATCHEDGEMMINCLUDE_HH
 #define BATCHEDGEMMINCLUDE_HH
 #include "DMRGConfig.h"
-#ifdef PLUGIN_SC
+#if 1 //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 1 //KOKKOS_BATCHED
+		return "KokkosKernels";
+#elif defined(PLUGIN_SC)
 		return "PLUGIN_SC";
 #else
 		return "";