BP4 engine capable of using device buffers with Put

ornladios · Oct 8, 2021 · a310850 · a310850
1 parent 39c8c99
commit a310850
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 0 deletions.
diff --git a/source/adios2/helper/adiosMemory.h b/source/adios2/helper/adiosMemory.h
@@ -39,6 +39,15 @@ template <class T>
 void InsertToBuffer(std::vector<char> &buffer, const T *source,
                     const size_t elements = 1) noexcept;
 
+/*
+ * Copies data from a GPU buffer to a specific location in the adios buffer
+ */
+#ifdef ADIOS2_HAVE_CUDA
+template <class T>
+void CopyFromGPUToBuffer(std::vector<char> &buffer, size_t &position,
+                         const T *source, const size_t elements = 1) noexcept;
+#endif
+
 /**
  * Copies data to a specific location in the buffer updating position
  * Does not update vec.size().

diff --git a/source/adios2/helper/adiosMemory.inl b/source/adios2/helper/adiosMemory.inl
@@ -20,6 +20,10 @@
 #include <iostream>
 #include <thread>
 /// \endcond
+#ifdef ADIOS2_HAVE_CUDA
+  #include <cuda.h>
+  #include <cuda_runtime.h>
+#endif
 
 #include "adios2/helper/adiosMath.h"
 #include "adios2/helper/adiosSystem.h"
@@ -74,6 +78,18 @@ void InsertToBuffer(std::vector<char> &buffer, const T *source,
     buffer.insert(buffer.end(), src, src + elements * sizeof(T));
 }
 
+#ifdef ADIOS2_HAVE_CUDA
+template <class T>
+void CopyFromGPUToBuffer(std::vector<char> &buffer, size_t &position,
+                         const T *source, const size_t elements) noexcept
+{
+    const char *src = reinterpret_cast<const char *>(source);
+    cudaMemcpy(buffer.data() + position, src, elements * sizeof(T),
+               cudaMemcpyDeviceToHost);
+    position += elements * sizeof(T);
+}
+#endif
+
 template <class T>
 void CopyToBuffer(std::vector<char> &buffer, size_t &position, const T *source,
                   const size_t elements) noexcept

diff --git a/source/adios2/toolkit/format/bp/BPSerializer.tcc b/source/adios2/toolkit/format/bp/BPSerializer.tcc
@@ -72,6 +72,18 @@ inline void BPSerializer::PutPayloadInBuffer(
 {
     const size_t blockSize = helper::GetTotalSize(blockInfo.Count);
     m_Profiler.Start("memcpy");
+
+#ifdef ADIOS2_HAVE_CUDA
+    if (blockInfo.IsGPU)
+    {
+        helper::CopyFromGPUToBuffer(m_Data.m_Buffer, m_Data.m_Position,
+                                    blockInfo.Data, blockSize);
+        m_Profiler.Stop("memcpy");
+        m_Data.m_AbsolutePosition += blockSize * sizeof(T);
+        return;
+    }
+#endif
+
     if (!blockInfo.MemoryStart.empty())
     {
         helper::CopyMemoryBlock(

diff --git a/source/adios2/toolkit/format/bp/bp4/BP4Serializer.tcc b/source/adios2/toolkit/format/bp/bp4/BP4Serializer.tcc
@@ -334,6 +334,15 @@ BP4Serializer::GetBPStats(const bool singleValue,
     stats.Step = m_MetadataSet.TimeStep;
     stats.FileIndex = GetFileIndex();
 
+#ifdef ADIOS2_HAVE_CUDA
+    if (blockInfo.IsGPU)
+    {
+        const size_t size = helper::GetTotalSize(blockInfo.Count);
+        helper::CUDAMinMax(blockInfo.Data, size, stats.Min, stats.Max);
+        return stats;
+    }
+#endif
+
     // support span
     if (blockInfo.Data == nullptr && m_Parameters.StatsLevel > 0)
     {