Skip to content

Commit

Permalink
Add loadChunkVariantAllocating
Browse files Browse the repository at this point in the history
  • Loading branch information
franzpoeschel committed Dec 4, 2023
1 parent 8469b3d commit 4c42b94
Show file tree
Hide file tree
Showing 4 changed files with 115 additions and 54 deletions.
7 changes: 7 additions & 0 deletions include/openPMD/RecordComponent.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,9 @@ class RecordComponent : public BaseRecordComponent
using shared_ptr_dataset_types = auxiliary::detail::
map_variant<auxiliary::detail::as_shared_pointer, dataset_types>::type;

using raw_ptr_dataset_types = auxiliary::detail::
map_variant<auxiliary::detail::as_raw_pointer, dataset_types>::type;

/** std::variant-based version of allocating loadChunk<T>(Offset, Extent)
*
* @return The same result that loadChunk() would return, but
Expand All @@ -215,6 +218,10 @@ class RecordComponent : public BaseRecordComponent
*/
shared_ptr_dataset_types loadChunkVariant(Offset = {0u}, Extent = {-1u});

template <typename Alloc>
raw_ptr_dataset_types
loadChunkVariantAllocating(Alloc &&alloc, Offset = {0u}, Extent = {-1u});

/** Load a chunk of data into pre-allocated memory.
*
* @param data Preallocated, contiguous buffer, large enough to load the
Expand Down
140 changes: 86 additions & 54 deletions include/openPMD/RecordComponent.tpp
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@

#pragma once

#include "openPMD/Dataset.hpp"
#include "openPMD/Error.hpp"
#include "openPMD/RecordComponent.hpp"
#include "openPMD/Span.hpp"
#include "openPMD/auxiliary/Memory.hpp"
Expand Down Expand Up @@ -54,27 +56,59 @@ inline RecordComponent &RecordComponent::makeEmpty(uint8_t dimensions)
return makeEmpty(Dataset(determineDatatype<T>(), Extent(dimensions, 0)));
}

template <typename T>
inline std::shared_ptr<T> RecordComponent::loadChunk(Offset o, Extent e)
namespace detail
{
uint8_t dim = getDimensionality();
inline void
checkOffsetExtentForLoading(RecordComponent &rc, Offset &o, Extent &e)
{
uint8_t dim = rc.getDimensionality();

// default arguments
// offset = {0u}: expand to right dim {0u, 0u, ...}
Offset offset = o;
if (o.size() == 1u && o.at(0) == 0u && dim > 1u)
offset = Offset(dim, 0u);
// default arguments
// offset = {0u}: expand to right dim {0u, 0u, ...}
Offset offset = o;
if (o.size() == 1u && o.at(0) == 0u && dim > 1u)
offset = Offset(dim, 0u);

// extent = {-1u}: take full size
Extent extent(dim, 1u);
if (e.size() == 1u && e.at(0) == -1u)
{
extent = getExtent();
for (uint8_t i = 0u; i < dim; ++i)
extent[i] -= offset[i];
// extent = {-1u}: take full size
Extent extent(dim, 1u);
if (e.size() == 1u && e.at(0) == -1u)
{
extent = rc.getExtent();
for (uint8_t i = 0u; i < dim; ++i)
extent[i] -= offset[i];
}
else
extent = e;

if (extent.size() != dim || offset.size() != dim)
{
std::ostringstream oss;
oss << "Dimensionality of chunk ("
<< "offset=" << offset.size() << "D, "
<< "extent=" << extent.size() << "D) "
<< "and record component (" << int(dim) << "D) "
<< "do not match.";
throw std::runtime_error(oss.str());
}
Extent dse = rc.getExtent();
for (uint8_t i = 0; i < dim; ++i)
if (dse[i] < offset[i] + extent[i])
throw std::runtime_error(
"Chunk does not reside inside dataset (Dimension on "
"index " +
std::to_string(i) + ". DS: " + std::to_string(dse[i]) +
" - Chunk: " + std::to_string(offset[i] + extent[i]) + ")");

o = offset;
e = extent;
}
else
extent = e;
} // namespace detail

template <typename T>
inline std::shared_ptr<T>
RecordComponent::loadChunk(Offset offset, Extent extent)
{
detail::checkOffsetExtentForLoading(*this, offset, extent);

uint64_t numPoints = 1u;
for (auto const &dimensionSize : extent)
Expand All @@ -94,8 +128,8 @@ inline std::shared_ptr<T> RecordComponent::loadChunk(Offset o, Extent e)
}

template <typename T>
inline void
RecordComponent::loadChunk(std::shared_ptr<T> data, Offset o, Extent e)
inline void RecordComponent::loadChunk(
std::shared_ptr<T> data, Offset offset, Extent extent)
{
Datatype dtype = determineDatatype(data);
if (dtype != getDatatype())
Expand All @@ -113,42 +147,8 @@ RecordComponent::loadChunk(std::shared_ptr<T> data, Offset o, Extent e)
throw std::runtime_error(err_msg);
}

uint8_t dim = getDimensionality();

// default arguments
// offset = {0u}: expand to right dim {0u, 0u, ...}
Offset offset = o;
if (o.size() == 1u && o.at(0) == 0u && dim > 1u)
offset = Offset(dim, 0u);

// extent = {-1u}: take full size
Extent extent(dim, 1u);
if (e.size() == 1u && e.at(0) == -1u)
{
extent = getExtent();
for (uint8_t i = 0u; i < dim; ++i)
extent[i] -= offset[i];
}
else
extent = e;
detail::checkOffsetExtentForLoading(*this, offset, extent);

if (extent.size() != dim || offset.size() != dim)
{
std::ostringstream oss;
oss << "Dimensionality of chunk ("
<< "offset=" << offset.size() << "D, "
<< "extent=" << extent.size() << "D) "
<< "and record component (" << int(dim) << "D) "
<< "do not match.";
throw std::runtime_error(oss.str());
}
Extent dse = getExtent();
for (uint8_t i = 0; i < dim; ++i)
if (dse[i] < offset[i] + extent[i])
throw std::runtime_error(
"Chunk does not reside inside dataset (Dimension on index " +
std::to_string(i) + ". DS: " + std::to_string(dse[i]) +
" - Chunk: " + std::to_string(offset[i] + extent[i]) + ")");
if (!data)
throw std::runtime_error(
"Unallocated pointer passed during chunk loading.");
Expand Down Expand Up @@ -192,6 +192,38 @@ inline void RecordComponent::loadChunkRaw(T *ptr, Offset offset, Extent extent)
loadChunk(auxiliary::shareRaw(ptr), std::move(offset), std::move(extent));
}

namespace detail
{
template <typename Alloc>
struct LoadChunkVariantAllocating
{
template <typename T>
static RecordComponent::raw_ptr_dataset_types
call(RecordComponent &self, Alloc &&alloc, Offset offset, Extent extent)
{
T *res;
auto dim = self.getDimensionality();
size_t flat_extent = 1;
for (uint8_t i = 0; i < dim; ++i)
{
flat_extent *= extent[i];
}
std::forward<Alloc>(alloc)(&res, flat_extent);
self.loadChunkRaw(res, std::move(offset), std::move(extent));
return res;
}
};
} // namespace detail

template <typename Alloc>
auto RecordComponent::loadChunkVariantAllocating(
Alloc &&alloc, Offset offset, Extent extent) -> raw_ptr_dataset_types
{
detail::checkOffsetExtentForLoading(*this, offset, extent);
return this->visit<detail::LoadChunkVariantAllocating<Alloc>>(
std::forward<Alloc>(alloc), std::move(offset), std::move(extent));
}

template <typename T>
inline void
RecordComponent::storeChunk(std::shared_ptr<T> data, Offset o, Extent e)
Expand Down
6 changes: 6 additions & 0 deletions include/openPMD/auxiliary/TypeTraits.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,12 @@ namespace detail
using type = std::shared_ptr<T>;
};

struct as_raw_pointer
{
template <typename T>
using type = T *;
};

template <typename...>
struct append_to_variant;

Expand Down
16 changes: 16 additions & 0 deletions test/SerialIOTest.cpp
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
// expose private and protected members for invasive testing
#include "openPMD/Datatype.hpp"
#include "openPMD/IO/Access.hpp"
#include <type_traits>
#if openPMD_USE_INVASIVE_TESTS
#define OPENPMD_private public:
#define OPENPMD_protected public:
Expand Down Expand Up @@ -1696,6 +1697,21 @@ inline void write_test(const std::string &backend)
<< '\'' << std::endl;
},
variantTypeDataset);

std::vector<char> serializationBuffer;
auto selfAllocatedVariant = rc.loadChunkVariantAllocating(
[&serializationBuffer](auto **ptr, size_t flat_extent) {
using type = std::remove_reference_t<decltype(**ptr)>;
serializationBuffer.resize(sizeof(type) * flat_extent);
*ptr = reinterpret_cast<type *>(serializationBuffer.data());
});
rc.seriesFlush();
std::visit(
[](auto *ptr) {
std::cout << "First value in loaded chunk: '" << ptr[0] << '\''
<< std::endl;
},
selfAllocatedVariant);
}

TEST_CASE("write_test", "[serial]")
Expand Down

0 comments on commit 4c42b94

Please sign in to comment.