Merge pull request #4046 from rupertnash/operator-plugin-sizes

Allow plugin operators to take advantage of the estimated size API
ornladios · Mar 2, 2024 · 6eff222 · 6eff222
2 parents e38e07f + 2a8d390
commit 6eff222
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 0 deletions.
diff --git a/plugins/operators/EncryptionOperator.cpp b/plugins/operators/EncryptionOperator.cpp
@@ -166,6 +166,16 @@ EncryptionOperator::InverseOperate(const char *bufferIn, const size_t sizeIn, ch
 
 bool EncryptionOperator::IsDataTypeValid(const DataType type) const { return true; }
 
+size_t EncryptionOperator::GetEstimatedSize(const size_t ElemCount, const size_t ElemSize,
+                                            const size_t ndims, const size_t *dims) const
+{
+    size_t sizeIn = ElemCount * ElemSize;
+    return (sizeof(size_t)                // Data size
+            + crypto_secretbox_NONCEBYTES // Nonce
+            + sizeIn                      // Data
+            + crypto_secretbox_MACBYTES   // MAC
+    );
+}
 } // end namespace plugin
 } // end namespace adios2
 

diff --git a/plugins/operators/EncryptionOperator.h b/plugins/operators/EncryptionOperator.h
@@ -41,6 +41,9 @@ class EncryptionOperator : public PluginOperatorInterface
 
     bool IsDataTypeValid(const DataType type) const override;
 
+    size_t GetEstimatedSize(const size_t ElemCount, const size_t ElemSize, const size_t ndims,
+                            const size_t *dims) const override;
+
 private:
     struct EncryptImpl;
     std::unique_ptr<EncryptImpl> Impl;

diff --git a/source/adios2/operator/plugin/PluginOperator.cpp b/source/adios2/operator/plugin/PluginOperator.cpp
@@ -79,6 +79,28 @@ void PluginOperator::PluginInit(const std::string &pluginName, const std::string
     m_Impl->m_Plugin = m_Impl->m_HandleCreate(m_Parameters);
 }
 
+size_t PluginOperator::GetEstimatedSize(const size_t ElemCount, const size_t ElemSize,
+                                        const size_t ndims, const size_t *dims) const
+{
+    // Need to calculate the size of the header written by Operate and then add our plugin's size.
+    constexpr size_t commonHeaderSize =
+        sizeof(m_TypeEnum) + sizeof(std::uint8_t) + sizeof(std::uint16_t);
+
+    auto &pp = m_Impl->m_PluginParams;
+    // Want to use std::transform_reduce but C++11
+    size_t paramsSize = 1; // for the number of parameters
+    for (auto &&p : pp)
+    {
+        // Need length and string for key and values.
+        paramsSize += p.first.size() + p.second.size() + 2;
+    }
+
+    // Plugin's estimate of size so it doesn't need to know about headers.
+    auto implSize = m_Impl->m_Plugin->GetEstimatedSize(ElemCount, ElemSize, ndims, dims);
+
+    return commonHeaderSize + paramsSize + implSize;
+}
+
 size_t PluginOperator::Operate(const char *dataIn, const Dims &blockStart, const Dims &blockCount,
                                const DataType type, char *bufferOut)
 {

diff --git a/source/adios2/operator/plugin/PluginOperator.h b/source/adios2/operator/plugin/PluginOperator.h
@@ -37,6 +37,9 @@ class PluginOperator : public core::Operator
     PluginOperator(const Params &parameters);
     virtual ~PluginOperator();
 
+    size_t GetEstimatedSize(const size_t ElemCount, const size_t ElemSize, const size_t ndims,
+                            const size_t *dims) const override;
+
     size_t Operate(const char *dataIn, const Dims &blockStart, const Dims &blockCount,
                    const DataType type, char *bufferOut) override;