Introduce more flexible storeChunk() syntax, use to add ADIOS2 memory selection

[franzpoeschel]

We currently don't support memory selections yet, e.g. when you have a 2D memory buffer but want to write only an arbitrary block from it. That block is then non-contiguous, but ADIOS2 has so-called memory selections to support this.

We currently don't have an API that would be able to expose this (except in Python where the native Python buffer protocol is powerful enough to express this; we currently throw an error when detecting this strides in selection are inefficient, not implemented!). Since the loadChunk()/storeChunk() API is somewhat convoluted by now anyway, I didn't want to add even further overloads.

New API Design

The new API is centered around RecordComponent::prepareLoadStore(), which returns a configuration object that can be chained with various methods before executing the actual load/store operation.

Methods return an auxiliary::DeferredComputation<T> object that encapsulates the operation without immediately executing it. The computation is only performed when get() (or operator()()) is called on the returned object:

// Returns DeferredComputation<std::shared_ptr<int>>
auto result = E_y.prepareLoadStore()
    .offset({0, 5})
    .extent({1, 5})
    .load<int>();

// Execution happens here when get() is called
auto data = result.get();

Configuration Methods

The configuration chain supports:

• offset(Offset) - Set the offset within the dataset (optional)
• extent(Extent) - Set the extent within the dataset (optional)
• memorySelection(MemorySelection) - Set memory selection for non-contiguous buffers (only available after withSharedPtr(), withRawPtr(), or withContiguousContainer())
• unsafeNoAutomaticFlush(): The returned DeferredComputation object will be a no-op object. Buffers will still be returned, but they might not yet be written/read until manually flushing.

Buffer Specification Methods

If not specifying a buffer, the following allocating load/store operations are available:

• storeSpan<T>() - Returns DynamicMemoryView<T> (a span-like view)
• load<T>() - Returns DeferredComputation<std::shared_ptr<T>>
• loadVariant() - Returns DeferredComputation<variant> for type-erased loading
• load<T>(policy) - Returns std::shared_ptr<T>>

Alternatively, you may specify a buffer using one of:

• withSharedPtr(std::shared_ptr<T>) - For shared pointer managed buffers
• withUniquePtr(UniquePtrWithLambda<T>) - For unique pointer with custom deleter
• withRawPtr(T*) - For raw pointer buffers
• withContiguousContainer(Container&) - For contiguous containers (automatically deduces size if extent is not set)

The following operations are only available after specifying a buffer. Template parameters are no longer needed as the type is given through the buffer type:

• store() / load() - Returns DeferredComputation<void>

(Load operations are only available after specifying a writable buffer, i.e. raw/shared pointer or contiguous container of non-const type)

ADIOS2 Memory Selection Support

This PR adds support for ADIOS2 memory selections, allowing writing from non-contiguous memory blocks:

std::shared_ptr<int> data;
E_y.prepareLoadStore()
    .offset({0, 5})
    .extent({1, 5})
    .withSharedPtr(data)
    .memorySelection({{1, 1}, {5, 5}})  // Write block from within the buffer
    .store();

Note: Memory selection support depends on ADIOS2 version. For ADIOS2 versions that cannot reset memory selections (PR ornladios/ADIOS2#4169), a warning is displayed. The API gracefully handles both cases.

Migration from Old API

The old loadChunk() and storeChunk() methods are now fully implemented using the new API internally, providing full backward compatibility while benefiting from the new implementation. Existing code continues to work without changes.

Examples of Old to New Migration

Old:

std::shared_ptr<int[]> data;
E_y.loadChunk(data, {0, 5}, {1, 5});

New:

auto result = E_y.prepareLoadStore()
    .offset({0, 5})
    .extent({1, 5})
    .withSharedPtr(data)
    .load<int>();
// Later: result.get();

Old:

std::vector<int> data(100);
E_y.storeChunk(data, {0, 5});

New:

E_y.prepareLoadStore()
    .offset({0, 5})
    .extent({1, 5})
    .withContiguousContainer(data)
    .store();

Design Considerations & Open Questions

Memory Selection Limitations

• HDF5 backend throws an error when memory selection is specified (not supported by HDF5)
• ADIOS2 requires version with memory selection reset support for clean reuse of variables
• Memory selection only available after buffer specification (not on bare prepareLoadStore())

Files Changed

• include/openPMD/LoadStoreChunk.hpp - New header with the main API
• include/openPMD/LoadStoreChunk.tpp - Template implementations
• include/openPMD/auxiliary/Future.hpp - New DeferredComputation wrapper
• include/openPMD/Dataset.hpp - Added MemorySelection struct
• include/openPMD/RecordComponent.hpp - Integration of new API
• src/LoadStoreChunk.cpp - Non-template implementations
• src/auxiliary/Future.cpp - DeferredComputation implementation
• All IO backends (ADIOS2, HDF5, etc.) - Memory selection support

TODO

• What about Python? The buffer protocol in Python may already support this; consider keeping Python API unchanged for now
• Decide on deferred vs. immediate API style
• Consider introducing .noop_future() option to disable future return types for cases where you don't need to track completion
• Further binary size optimization

Related

• ADIOS2 PR ornladios/ADIOS2#4169: Memory selection reset support

TODO:

• look at hpx future for inspiration

[ax3l]

Codex:

[P1] Do not treat an arbitrary backend flush as completion of this chunk

The chunk is initially queued in RecordComponentData::m_chunks, whereas m_flushCounter advances whenever the IO handler drains its own queue, including metadata-only operations. Consequently, an unrelated query can make this guard skip the flush while the chunk is still pending. I reproduced this with the JSON backend:

int values[4] = {-7, -7, -7, -7};
auto pending = rc.prepareLoadStore().withRawPtr(values).load();
auto ranks = series.rankTable();
pending.get(); // values still contains -7
series.flush(); // values now contains the expected dataset values

This also affects allocating loads, which can return uninitialized data as though the read had completed. Please track completion of the actual chunk, or use a completion generation that advances only after the relevant frontend chunks have been flushed; draining an unrelated IO queue is insufficient.

[ax3l]

Codex:

[P2] Preserve full-selection argument expansion in loadChunkRaw

This previously delegated to loadChunk(shared_ptr, offset, extent), which expands offset {0u} to the dataset dimensionality and extent {-1u} to the remaining extent. Passing those arguments directly to the new builder makes them literal selections instead. On a 2D dataset, the documented call rc.loadChunkRaw(buffer, {0u}, {-1u}) now throws a dimensionality mismatch; on a 1D dataset, the sentinel extent is treated as an oversized chunk. The shared-pointer overload still performs the required expansion, so forwarding through it preserves the existing API contract:

Suggested change

	prepareLoadStore()
	.offset(std::move(offset))
	.extent(std::move(extent))
	.withRawPtr(ptr)
	.unsafeNoAutomaticFlush()
	.load()
	.get();
	loadChunk(auxiliary::shareRaw(ptr), std::move(offset), std::move(extent));

[ax3l]

Codex:

[P2] Apply the deleter to its argument after resetting a converted pointer

Capturing original_ptr makes the deleter ignore the pointer currently owned by the resulting UniquePtrWithLambda. This class publicly inherits std::unique_ptr, so auto q = std::move(p).static_cast_<void>(); q.reset(replacement); q.reset(); is valid usage: the first reset should delete the original allocation and the second should delete the replacement. A tracking-deleter reproducer instead receives the original address twice. With a normal deleting deleter this becomes a double free, while the replacement leaks. Please retain deletion of the supplied pointer, with the necessary type/const conversion, instead of binding deletion permanently to the original address.

[ax3l]

Claude:

Source-compatibility regression: the std::shared_ptr<T[]> overloads of loadChunk() / storeChunk() were removed. Calls without explicit template arguments still work (T deduces to double[]), but calls with explicit template arguments no longer compile:

std::shared_ptr<double[]> p(new double[10]);
rc.loadChunk<double>(p, {0}, {10});  // dev: OK, PR: no matching function
rc.storeChunk<double>(p, {0}, {10}); // dev: OK, PR: no matching function

(Verified with GCC 13 against the dev and PR headers.) Suggest keeping them as thin, header-only forwarding overloads. No extra explicit instantiations are needed, partial ordering still selects them for the implicit call, and inside the forwarder loadChunk<T[]> SFINAE-drops the T[][] candidate. With this and the storeChunk counterpart below, explicit, implicit and const[] calls compile and link again:

Suggested change

	template <typename T>
	void loadChunk(std::shared_ptr<T> data, Offset offset, Extent extent);
	template <typename T>
	void loadChunk(std::shared_ptr<T> data, Offset offset, Extent extent);
	/** Load a chunk of data into pre-allocated memory, array version.
	*
	* Kept for source compatibility with explicit template arguments,
	* e.g. loadChunk<double>(std::shared_ptr<double[]>, ...).
	*/
	template <typename T>
	void loadChunk(std::shared_ptr<T[]> data, Offset offset, Extent extent)
	{
	loadChunk<T[]>(std::move(data), std::move(offset), std::move(extent));
	}

[ax3l]

Claude:

Same source-compatibility regression for storeChunk<double>(std::shared_ptr<double[]>, ...) (see the loadChunk comment above):

Suggested change

	template <typename T>
	void storeChunk(std::shared_ptr<T> data, Offset offset, Extent extent);
	template <typename T>
	void storeChunk(std::shared_ptr<T> data, Offset offset, Extent extent);
	/** Store a chunk of data from a chunk of memory, array version.
	*
	* Kept for source compatibility with explicit template arguments,
	* e.g. storeChunk<double>(std::shared_ptr<double[]>, ...).
	*/
	template <typename T>
	void storeChunk(std::shared_ptr<T[]> data, Offset offset, Extent extent)
	{
	storeChunk<T[]>(std::move(data), std::move(offset), std::move(extent));
	}

[ax3l]

Claude:

Regression, out-of-bounds read: for N-D datasets the default extent {-1u} is now dropped, so computeExtent() falls back to the full dataset extent instead of the container size. On dev, the call below threw Dimensionality of chunk (offset=2D, extent=1D) and record component (2D) do not match. With the PR it is accepted, and the backend reads prod(dataset extent) elements from a 4-element vector:

rc.resetDataset({Datatype::INT, {100, 100}});
std::vector<int> v(4, 7);
rc.storeChunk(v, {0, 0});
series.flush(); // valgrind: Invalid read of size 4 (JSONIOHandlerImpl, to_json)

Second, smaller regression in the same block: for joined arrays, a non-default offset from the user is now silently discarded. dev threw Joined array: Must specify an empty offset (given: offset=1D, extent=1D). Now rc.storeChunk(v, {5}) on a JOINED_DIMENSION dataset is accepted. No special case for joined arrays is needed here, because computeOffset() already yields {} for them when the offset is left unset.

The suggestion below restores the dev semantics. I verified that both repros throw again with the dev error messages, and the 1D auto-size path in withContiguousContainer() is unchanged:

Suggested change

	auto joined_dim = joinedDimension();
	if (!joined_dim.has_value() && (o.size() != 1 || o.at(0) != 0u))
	{
	if (joinedDimension().has_value())
	{
	offset.clear();
	}
	else if (dim > 1u)
	{
	offset = Offset(dim, 0u);
	}
	storeChunkConfig.offset(std::move(o));
	}
	if (e.size() != 1 || e.at(0) != -1u)
	{
	storeChunkConfig.extent(std::move(e));
	}
	// default offset {0u}: leave unset, computeOffset() expands it to
	// {0, 0, ...} (or to {} for joined arrays)
	if (o.size() != 1 || o.at(0) != 0u)
	{
	storeChunkConfig.offset(std::move(o));
	}
	// default extent {-1u}: only 1D containers are auto-sized by
	// withContiguousContainer(), for N-D pass it through so that
	// verifyChunk() rejects it instead of reading past the container
	if (e.size() != 1 || e.at(0) != -1u || getDimensionality() > 1)
	{
	storeChunkConfig.extent(std::move(e));
	}

[ax3l]

Claude:

Heap buffer overflow in the new API: the container size is used only for 1D, and only as a default. It is never checked against the final extent. For N-D datasets, or when .extent() / .memorySelection() is set later in the chain (the PR's own available_chunks_test does that), nothing ensures the container is large enough. Repro (JSON, 4×4 INT dataset):

std::vector<int> small(2);
rc.prepareLoadStore().withContiguousContainer(small).load().get();
// succeeds; valgrind: Invalid write of size 4 (30 errors)

.store() has the same problem as an over-read. The final extent is only known at store() / load() time, so I suggest remembering data.size() here (e.g. std::optional<size_t> m_containerSize in ConfigureStoreChunkFromBuffer). Then store() / load() can check that prod(memorySelection ? memorySelection->extent : extent) <= *m_containerSize and throw error::WrongAPIUsage otherwise.

[ax3l]

Claude:

Missing bounds check for memory selections, which leads to out-of-bounds reads. ADIOS2's VariableBase::SetMemorySelection() only validates the rank and memoryCount[i] >= count[i]. It does not check memoryStart[i] + count[i] <= memoryCount[i]. openPMD doesn't check this either, so an inconsistent selection is silently accepted and ADIOS2 reads past the user buffer.

Repro (ADIOS2 2.12.1): 3×3 dataset, 5×5 buffer, .extent({3, 3}).memorySelection({{3, 3}, {5, 5}}). It is accepted; valgrind reports invalid reads in BP5Serializer::ProcessDeferredMinMax / GetMinMax (BP5) and Invalid read of size 4 (BP4). The valid selection from the PR's test, {{1, 1}, {5, 5}}, is clean (0 errors).

I suggest validating here, next to verifyChunk(). I verified that the out-of-range case then throws and the valid case still passes:

Suggested change

	auto [o, e, memorySelection] = std::move(cfg);
	verifyChunk(dtype, o, e);
	auto [o, e, memorySelection] = std::move(cfg);
	verifyChunk(dtype, o, e);
	if (memorySelection.has_value())
	{
	auto const &mem_offset = memorySelection->offset;
	auto const &mem_extent = memorySelection->extent;
	if (mem_offset.size() != e.size() || mem_extent.size() != e.size())
	{
	throw error::WrongAPIUsage(
	"Memory selection: dimensionality of memory offset and memory "
	"extent must match the chunk extent.");
	}
	for (size_t i = 0; i < e.size(); ++i)
	{
	if (mem_offset[i] + e[i] > mem_extent[i])
	{
	throw error::WrongAPIUsage(
	"Memory selection: memory offset + chunk extent exceeds "
	"the memory extent (dimension " +
	std::to_string(i) + ").");
	}
	}
	}

[ax3l]

Claude:

Silent data corruption with ADIOS2 2.9.x / 2.10.x, which are still supported (find_package(ADIOS2 2.9.0 ...)). The reset API only exists since ADIOS2 v2.11.0 (ornladios/ADIOS2#4169). In v2.10.0, VariableBase::SetMemorySelection() has no way to clear a selection: an empty box throws on the rank check. So after one memory-selected put, every later put of that variable keeps the stale memory selection. That includes later storeChunk() calls in the same step, later steps in group- and variable-based encoding, and the RunUniquePtrPut path. The result is wrong data or out-of-bounds reads of the user buffer.

A one-time std::cerr warning is easy to miss, especially because the damage happens in other calls than the one that triggers it. The PR's own available_chunks_test acknowledges the problem by skipping the last_row check when !CanTheMemorySelectionBeReset.

openPMD can apply the workaround from the warning text itself. A memory selection with start {0, ...} and count equal to the selection count is equivalent to contiguous memory. For example, ADIOS2File could track the names of variables that ever got a memory selection. In the write path, for those variables, when !memorySelection && !CanTheMemorySelectionBeReset, call var.SetMemorySelection({adios2::Dims(extent.size(), 0), adios2::Dims(extent.begin(), extent.end())}) before Put. Doing this only for those variables avoids forcing BP5's copy path on every put. Alternatively, reject memory selections outright for ADIOS2 < 2.11.

(Validated by reading ADIOS2 v2.10.0 VariableBase.cpp and the PR's test gating. I could only runtime-test with ADIOS2 2.12.1, where the reset path works correctly.)

[ax3l]

Claude:

Memory selections on joined arrays are accepted here but always fail at flush. For a joined array, the ADIOS2 variable has an empty Start, so SetMemorySelection() throws during the flush:

<VariableBase> <SetMemorySelection> : memoryStart size must be the same as variable /data/0/meshes/E/x start size 0

Repro: rc.resetDataset({Datatype::INT, {Dataset::JOINED_DIMENSION, 3}}), then .extent({3, 3}).withContiguousContainer(v).memorySelection({{1, 1}, {5, 5}}).store(). That fails at flush time, and as a flush-time failure it takes the rest of the IO queue with it (see the HDF5 comment). The bounds check suggested earlier doesn't catch this, because e.size() == dim holds for joined arrays. I suggest rejecting it at enqueue time; I verified that the repro then throws immediately:

Suggested change

	dWrite.memorySelection = memorySelection;
	if (memorySelection.has_value() && joinedDimension().has_value())
	{
	throw error::WrongAPIUsage(
	"Memory selections are not supported for joined arrays.");
	}
	dWrite.memorySelection = memorySelection;

[ax3l]

Claude:

Rejecting memory selections at flush time corrupts the whole output file (HDF5 and JSON). The exception is raised inside the IO task, so AbstractIOHandlerImpl::flush() prints "Clearing IO queue and passing on the exception" and drops every other pending task. That includes other components' chunks and the file's root attributes. After that, a second flush() and close() both "succeed", but the file can't be read back.

Repro (same result for .h5 and .json): in one Series, A.storeChunk(av, {0}, {4}), then a store with a memory selection on another component B, then s.flush(). The flush throws Operation unsupported in HDF5: Non-contiguous memory selections .... Then s.flush(); s.close(); both return normally. On reopening, Series r(..., READ_ONLY) fails with Read Error ... Attribute NotFound: openPMD, so even the valid data for A is lost.

Whether a backend supports memory selections is known when the chunk is enqueued. I suggest checking it in RecordComponent::storeChunk_impl() and throwing error::OperationUnsupportedInBackend there. For example, add a virtual capability query to AbstractIOHandler (like the existing fullSupportForVariableBasedEncoding()) that only ADIOS2IOHandler overrides to true. The backend-side throws can stay as a safety net.

(The queue-clearing on flush errors is existing behaviour, but this PR adds an error that users can predictably trigger and that could be caught early.)

[ax3l]

Claude:

Dangling reference: ConfigureLoadStore keeps a RecordComponent & to the handle object it was created from, not to the shared component data. RecordComponent is a cheap, copyable handle that users routinely keep in containers or temporaries, so the builder silently dangles as soon as that handle object moves or dies. Repro (BP5):

std::vector<RecordComponent> comps;
comps.push_back(series.snapshots()[0].meshes["E"]["x"]);
auto cfg = comps[0].prepareLoadStore();
for (int i = 0; i < 16; ++i)
    comps.push_back(comps[0]); // reallocation, comps[0] moves
auto p = cfg.load<int>().get(); // "works", but valgrind: 26 invalid reads

auto cfg = getComponent().prepareLoadStore(); (a temporary handle) has the same problem. Since RecordComponent is only forward-declared here, the builder can't hold it by value directly. It could hold a copy of the handle via std::shared_ptr<RecordComponent>, created in RecordComponent::prepareLoadStore() as std::make_shared<RecordComponent>(*this) (one small allocation per operation, negligible next to the IO). Then the builder's lifetime no longer depends on the caller's handle object.

[ax3l]

Claude:

The default (safe) mode hides an MPI-collective operation in get() and in ~DeferredComputation(), which can deadlock parallel writers. Series::flush() is collective (docs/source/details/mpi.rst). Here, whether get() actually flushes depends on rank-local state: the per-handler flush counter and the number of store() calls on that rank. The natural pattern for particle data, where each rank loops over its own chunks, therefore issues a different number of flushes per rank:

// rank 0 owns 2 chunks, rank 1 owns 1 chunk
for (auto &ch : local_chunks)
    rc.prepareLoadStore().offset({off++}).withContiguousContainer(ch).store().get();
series.close();

Results with mpiexec -n 2, ADIOS2 2.12.1 BP5:

• the code above: hangs (killed after 45 s);
• same code with equal chunk counts per rank: completes;
• the equivalent old API (storeChunk() per chunk, then one series.flush()): completes.

So the hang comes from the implicit flushes. The same happens if the returned object is simply dropped (the destructor flushes), including during stack unwinding on a single rank.

This is a design question, not a one-line fix. Options: make "no implicit flush" the default, so get() means "the buffer may be used after the next collective flush()"; or keep the implicit flush but document get() / the destructor as collective and add them to the table in mpi.rst. In any case, the flush in the destructor seems risky for MPI users.