Add helper functions for setting up Nbnxm gpu object in nblib

diff --git a/api/nblib/nbnxmsetuphelpers.cpp b/api/nblib/nbnxmsetuphelpers.cpp
index 160612f8e8228fd2d170d962b7924878074c2b25..c876482b97f3b30b9072d4554e173bbd661d6170 100644 (file)
--- a/api/nblib/nbnxmsetuphelpers.cpp
+++ b/api/nblib/nbnxmsetuphelpers.cpp
@@ -129,6 +129,16 @@ Nbnxm::KernelSetup createKernelSetupCPU(const SimdKernels nbnxmSimd, const bool
     return kernelSetup;
 }
 
+Nbnxm::KernelSetup createKernelSetupGPU(const bool useTabulatedEwaldCorr)
+{
+    Nbnxm::KernelSetup kernelSetup;
+    kernelSetup.kernelType         = Nbnxm::KernelType::Gpu8x8x8;
+    kernelSetup.ewaldExclusionType = useTabulatedEwaldCorr ? Nbnxm::EwaldExclusionType::Table
+                                                           : Nbnxm::EwaldExclusionType::Analytical;
+
+    return kernelSetup;
+}
+
 std::vector<int64_t> createParticleInfoAllVdw(const size_t numParticles)
 {
     std::vector<int64_t> particleInfoAllVdw(numParticles);
@@ -178,6 +188,29 @@ gmx::StepWorkload createStepWorkload()
     return stepWorkload;
 }
 
+static gmx::SimulationWorkload createSimulationWorkload()
+{
+    gmx::SimulationWorkload simulationWork;
+    simulationWork.computeNonbonded = true;
+    return simulationWork;
+}
+
+gmx::SimulationWorkload createSimulationWorkloadGpu()
+{
+    gmx::SimulationWorkload simulationWork = createSimulationWorkload();
+
+    simulationWork.useGpuNonbonded = true;
+    simulationWork.useGpuUpdate    = false;
+
+    return simulationWork;
+}
+
+std::shared_ptr<gmx::DeviceStreamManager> createDeviceStreamManager(const DeviceInformation& deviceInfo,
+                                                                    const gmx::SimulationWorkload& simulationWorkload)
+{
+    return std::make_shared<gmx::DeviceStreamManager>(deviceInfo, false, simulationWorkload, false);
+}
+
 real ewaldCoeff(const real ewald_rtol, const real pairlistCutoff)
 {
     return calc_ewaldcoeff_q(pairlistCutoff, ewald_rtol);
@@ -279,6 +312,53 @@ std::unique_ptr<nonbonded_verlet_t> createNbnxmCPU(const size_t              num
     return nbv;
 }
 
+std::unique_ptr<nonbonded_verlet_t> createNbnxmGPU(const size_t               numParticleTypes,
+                                                   const NBKernelOptions&     options,
+                                                   const std::vector<real>&   nonbondedParameters,
+                                                   const interaction_const_t& interactionConst,
+                                                   const gmx::DeviceStreamManager& deviceStreamManager)
+{
+    const auto pinPolicy       = gmx::PinningPolicy::PinnedIfSupported;
+    const int  combinationRule = static_cast<int>(options.ljCombinationRule);
+
+    Nbnxm::KernelSetup kernelSetup = createKernelSetupGPU(options.useTabulatedEwaldCorr);
+
+    PairlistParams pairlistParams(kernelSetup.kernelType, false, options.pairlistCutoff, false);
+
+
+    // nbnxn_atomdata is always initialized with 1 thread if the GPU is used
+    constexpr int numThreadsInit = 1;
+    // multiple energy groups are not supported on the GPU
+    constexpr int numEnergyGroups = 1;
+    auto          atomData        = std::make_unique<nbnxn_atomdata_t>(pinPolicy,
+                                                       gmx::MDLogger(),
+                                                       kernelSetup.kernelType,
+                                                       combinationRule,
+                                                       numParticleTypes,
+                                                       nonbondedParameters,
+                                                       numEnergyGroups,
+                                                       numThreadsInit);
+
+    NbnxmGpu* nbnxmGpu = Nbnxm::gpu_init(
+            deviceStreamManager, &interactionConst, pairlistParams, atomData.get(), false);
+
+    // minimum iList count for GPU balancing
+    int iListCount = Nbnxm::gpu_min_ci_balanced(nbnxmGpu);
+
+    auto pairlistSets = std::make_unique<PairlistSets>(pairlistParams, false, iListCount);
+    auto pairSearch   = std::make_unique<PairSearch>(
+            PbcType::Xyz, false, nullptr, nullptr, pairlistParams.pairlistType, false, options.numOpenMPThreads, pinPolicy);
+
+    // Put everything together
+    auto nbv = std::make_unique<nonbonded_verlet_t>(
+            std::move(pairlistSets), std::move(pairSearch), std::move(atomData), kernelSetup, nbnxmGpu, nullptr);
+
+    // Some paramters must be copied to NbnxmGpu to have a fully constructed nonbonded_verlet_t
+    Nbnxm::gpu_init_atomdata(nbv->gpu_nbv, nbv->nbat.get());
+
+    return nbv;
+}
+
 void setGmxNonBondedNThreads(int numThreads)
 {
     gmx_omp_nthreads_set(ModuleMultiThread::Pairsearch, numThreads);

diff --git a/api/nblib/nbnxmsetuphelpers.h b/api/nblib/nbnxmsetuphelpers.h
index 2f22bd825b80699aca222bce185f8bdb4b94c6b0..5c70811e5de134aef1527b829987204ec2af6d4f 100644 (file)
--- a/api/nblib/nbnxmsetuphelpers.h
+++ b/api/nblib/nbnxmsetuphelpers.h
@@ -91,6 +91,9 @@ void checkKernelSetupSimd(SimdKernels nbnxmSimd);
 //! Creates and returns the kernel setup for CPU
 Nbnxm::KernelSetup createKernelSetupCPU(const SimdKernels nbnxmSimd, const bool useTabulatedEwaldCorr);
 
+//! Creates and returns the kernel setup for GPU
+Nbnxm::KernelSetup createKernelSetupGPU(const bool useTabulatedEwaldCorr);
+
 //! Create Particle info array to mark those that undergo VdV interaction
 std::vector<int64_t> createParticleInfoAllVdw(size_t numParticles);
 
@@ -101,6 +104,13 @@ std::vector<real> createNonBondedParameters(const std::vector<ParticleType>& par
 //! Create a step work object
 gmx::StepWorkload createStepWorkload();
 
+//! Create a SimulationWorkload object for use with createDeviceStreamManager
+gmx::SimulationWorkload createSimulationWorkloadGpu();
+
+//! Create a DeviceStreamManager; could be shared among multiple force calculators
+std::shared_ptr<gmx::DeviceStreamManager> createDeviceStreamManager(const DeviceInformation& deviceInfo,
+                                                                    const gmx::SimulationWorkload& simulationWorkload);
+
 //! Computes the Ewald splitting coefficient for Coulomb
 real ewaldCoeff(real ewald_rtol, real pairlistCutoff);
 
@@ -113,6 +123,13 @@ std::unique_ptr<nonbonded_verlet_t> createNbnxmCPU(size_t                    num
                                                    int                       numEnergyGroups,
                                                    gmx::ArrayRef<const real> nonbondedParameters);
 
+//! Create nonbonded_verlet_gpu object
+std::unique_ptr<nonbonded_verlet_t> createNbnxmGPU(size_t                     numParticleTypes,
+                                                   const NBKernelOptions&     options,
+                                                   const std::vector<real>&   nonbondedParameters,
+                                                   const interaction_const_t& interactionConst,
+                                                   const gmx::DeviceStreamManager& deviceStreamManager);
+
 //! Set number of OpenMP threads in the GROMACS backend
 void setGmxNonBondedNThreads(int numThreads);
 

diff --git a/api/nblib/tests/CMakeLists.txt b/api/nblib/tests/CMakeLists.txt
index ef2ace7d19a30b4bd0a653694a8d6648cef10fd3..b0e80d50faedc8bb7cafe2518331ddbd991d9129 100644 (file)
--- a/api/nblib/tests/CMakeLists.txt
+++ b/api/nblib/tests/CMakeLists.txt
@@ -52,6 +52,7 @@ set(exename "nblib-setup-test")
 
 gmx_add_gtest_executable(
     ${exename}
+    HARDWARE_DETECTION
     CPP_SOURCE_FILES
     # files with code for tests
         box.cpp
@@ -61,7 +62,6 @@ gmx_add_gtest_executable(
         molecules.cpp
         nbnxmsetup.cpp
         topology.cpp
-        tpr.cpp
         virials.cpp
     )
 target_link_libraries(${exename} PRIVATE mdrun_test_infrastructure)
@@ -70,6 +70,19 @@ target_include_directories(${exename} PRIVATE ${PROJECT_SOURCE_DIR}/api)
 gmx_register_gtest_test(${testname} ${exename} INTEGRATION_TEST)
 add_dependencies(nblib-tests ${exename})
 
+set(exename "nblib-tpr-test")
+
+gmx_add_unit_test(
+        NbLibTprTests
+        ${exename}
+        CPP_SOURCE_FILES
+        # files with code for tests
+        tpr.cpp
+)
+target_link_libraries(${exename} PRIVATE mdrun_test_infrastructure)
+target_link_libraries(${exename} PRIVATE nblib_test_infrastructure nblib)
+add_dependencies(nblib-tests ${exename})
+
 set(testname "NbLibIntegrationTests")
 set(exename "nblib-integration-test")
 

diff --git a/api/nblib/tests/nbnxmsetup.cpp b/api/nblib/tests/nbnxmsetup.cpp
index a7bff00fdf72a013374580a90cdd89b8cc24c95c..369173467e5338b0037a43741eb187d40489477b 100644 (file)
--- a/api/nblib/tests/nbnxmsetup.cpp
+++ b/api/nblib/tests/nbnxmsetup.cpp
@@ -43,6 +43,7 @@
  */
 #include <cmath>
 
+#include "gromacs/hardware/device_management.h"
 #include "gromacs/mdtypes/forcerec.h"
 #include "gromacs/mdtypes/interaction_const.h"
 #include "gromacs/mdtypes/simulation_workload.h"
@@ -52,6 +53,7 @@
 #include "nblib/nbnxmsetuphelpers.h"
 
 #include "testutils/testasserts.h"
+#include "testutils/test_hardware_environment.h"
 
 namespace nblib
 {
@@ -191,6 +193,48 @@ TEST(NbnxmSetupTest, CanCreateNbnxmCPU)
     EXPECT_NO_THROW(createNbnxmCPU(numParticles, nbKernelOptions, numEnergyGroups, nonbondedParameters));
 }
 
+#if GMX_GPU_CUDA
+TEST(NbnxmSetupTest, canCreateKernelSetupGPU)
+{
+    NBKernelOptions    nbKernelOptions;
+    Nbnxm::KernelSetup kernelSetup = createKernelSetupGPU(nbKernelOptions.useTabulatedEwaldCorr);
+    EXPECT_EQ(kernelSetup.kernelType, Nbnxm::KernelType::Gpu8x8x8);
+    EXPECT_EQ(kernelSetup.ewaldExclusionType, Nbnxm::EwaldExclusionType::Analytical);
+}
+
+TEST(NbnxmSetupTest, CanCreateDeviceStreamManager)
+{
+    const auto& testDeviceList = gmx::test::getTestHardwareEnvironment()->getTestDeviceList();
+    for (const auto& testDevice : testDeviceList)
+    {
+        const DeviceInformation& deviceInfo = testDevice->deviceInfo();
+        setActiveDevice(deviceInfo);
+        gmx::SimulationWorkload simulationWork = createSimulationWorkloadGpu();
+        EXPECT_NO_THROW(createDeviceStreamManager(deviceInfo, simulationWork));
+    }
+}
+
+TEST(NbnxmSetupTest, CanCreateNbnxmGPU)
+{
+    const auto& testDeviceList = gmx::test::getTestHardwareEnvironment()->getTestDeviceList();
+    for (const auto& testDevice : testDeviceList)
+    {
+        const DeviceInformation& deviceInfo = testDevice->deviceInfo();
+        setActiveDevice(deviceInfo);
+        size_t                  numParticles = 1;
+        NBKernelOptions         nbKernelOptions;
+        std::vector<real>       nonbondedParameters = { 1, 1 };
+        gmx::SimulationWorkload simulationWork      = createSimulationWorkloadGpu();
+        interaction_const_t     interactionConst    = createInteractionConst(nbKernelOptions);
+        // set DeviceInformation and create the DeviceStreamManager
+        auto deviceStreamManager = createDeviceStreamManager(deviceInfo, simulationWork);
+        EXPECT_NO_THROW(createNbnxmGPU(
+                numParticles, nbKernelOptions, nonbondedParameters, interactionConst, *deviceStreamManager));
+    }
+}
+
+#endif
+
 } // namespace
 } // namespace test
 } // namespace nblib