ZFP: Enable cuda execution policy

ornladios · Oct 21, 2021 · a740238 · a740238
1 parent 9bd5718
commit a740238
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 0 deletions.
diff --git a/source/adios2/common/ADIOSTypes.h b/source/adios2/common/ADIOSTypes.h
@@ -305,9 +305,17 @@ namespace zfp
 namespace key
 {
 constexpr char accuracy[] = "accuracy";
+constexpr char backend[] = "backend";
 constexpr char rate[] = "rate";
 constexpr char precision[] = "precision";
 }
+
+namespace value
+{
+constexpr char backend_cuda[] = "cuda";
+constexpr char backend_omp[] = "omp";
+constexpr char backend_serial[] = "serial";
+}
 }
 #endif
 

diff --git a/source/adios2/operator/compress/CompressZFP.cpp b/source/adios2/operator/compress/CompressZFP.cpp
@@ -11,6 +11,11 @@
 #include "adios2/helper/adiosFunctions.h"
 #include <sstream>
 
+/* ZFP will default to SERIAL if CUDA is not available */
+#ifndef ZFP_DEFAULT_EXECUTION_POLICY
+#define ZFP_DEFAULT_EXECUTION_POLICY zfp_exec_cuda
+#endif
+
 namespace adios2
 {
 namespace core
@@ -256,6 +261,7 @@ zfp_stream *CompressZFP::GetZFPStream(const Dims &dimensions, DataType type,
                                       const Params &parameters) const
 {
     zfp_stream *stream = zfp_stream_open(NULL);
+    zfp_stream_set_execution(stream, ZFP_DEFAULT_EXECUTION_POLICY);
 
     auto itAccuracy = parameters.find("accuracy");
     const bool hasAccuracy = itAccuracy != parameters.end();
@@ -266,6 +272,30 @@ zfp_stream *CompressZFP::GetZFPStream(const Dims &dimensions, DataType type,
     auto itPrecision = parameters.find("precision");
     const bool hasPrecision = itPrecision != parameters.end();
 
+    auto itBackend = parameters.find("backend");
+    const bool hasBackend = itBackend != parameters.end();
+
+    if (hasBackend)
+    {
+        auto policy = ZFP_DEFAULT_EXECUTION_POLICY;
+        const auto backend = itBackend->second;
+
+        if (backend == "cuda")
+        {
+            policy = zfp_exec_cuda;
+        }
+        else if (backend == "omp")
+        {
+            policy = zfp_exec_omp;
+        }
+        else if (backend == "serial")
+        {
+            policy = zfp_exec_serial;
+        }
+
+        zfp_stream_set_execution(stream, policy);
+    }
+
     if ((hasAccuracy && hasRate) || (hasAccuracy && hasPrecision) ||
         (hasRate && hasPrecision) || !(hasAccuracy || hasRate || hasPrecision))
     {

diff --git a/testing/adios2/engine/bp/operations/TestBPWriteReadZfp.cpp b/testing/adios2/engine/bp/operations/TestBPWriteReadZfp.cpp
@@ -866,6 +866,7 @@ void ZFPRate2DSmallSel(const std::string rate)
         // add operations
         adios2::Operator szOp =
             adios.DefineOperator("ZFPCompressor", adios2::ops::LossyZFP);
+        szOp.SetParameter("backend", "serial");
 
         var_r32.AddOperation(szOp, {{adios2::ops::zfp::key::rate, rate}});
         var_r64.AddOperation(szOp, {{adios2::ops::zfp::key::rate, rate}});