DLPack Tensor Exchange#

ovstage exchanges attribute values as DLPack DLTensor s across three paths: copy-in writes (write_attribute), copy-out reads (read_attributes → fetch_read_next), and a zero-copy map/unmap path that fills ovstage-owned storage directly. Tensors can be CPU- or CUDA-resident.

The Shared Data Struct#

One data shape is reused across the API. ovstage_data_t (reads and map groups) and ovstage_write_data_t (writes) both carry:

  • tensors + tensor_count — an array of DLTensor s.

  • optional sparsity — index_map or mask, with count.

  • cuda_sync — GPU synchronization descriptor (ovstage_cuda_sync_t; described below).

tensor_count depends on the attribute kind, which is declared explicitly through is_array and never inferred:

Attribute Kind

is_array

Tensor layout

Fixed-size (scalar / fixed vector)

false

A single tensor with all prims stacked along the leading dimension (tensor_count == 1).

Array / ragged

true

One tensor per row, or a single packed tensor for all rows — tensor_count selects the transport, not the attribute kind.

DLTensor fields are the standard DLPack set: data, device ({kDLCPU, 0} or {kDLCUDA, ordinal}), ndim, dtype ({code, bits, lanes} — lanes is the tuple width, e.g. 3 for a float3), shape, strides, and byte_offset.

Copy-In Write and Copy-Out Read#

The minimal example builds a DLTensor over a CPU buffer, writes it at an ordinal, seals the ordinal, and reads the column back:

# Write one float per prim into the "temperature" column at ordinal 1,
# seal it by advancing the write floor to 1, then read it back. Tensor
# data crosses as a numpy array (CPU) via DLPack; async ops return an
# Operation whose .wait() raises OvstageError on failure. The read is
# a context manager -- block exit releases its handle even when an
# error interrupts -- and the fetched group is released in a finally.
stage.write_attribute(
    query, attr, ordinal=1, tensors=np.array([1.0, 2.0, 3.0], np.float32), is_array=False
).wait()
stage.advance_write_floor(ordinal=1).wait()

with stage.read_attributes(query, [attr], OrdinalRange.latest(1)) as read:
    read.wait()
    group = read.fetch_next()
    if group is None:
        raise SystemExit("read returned no group at ordinal 1")
    try:
        # group.array(i) is a zero-copy numpy view of tensor i (CPU).
        print("read back ordinal", group.ordinal, group.array(0))  # -> [1. 2. 3.]
    finally:
        stage.release_group(group)

// Write one float per prim into the "temperature" column (UPSERT creates
// the prims on first write), seal it by advancing the write floor to
// ordinal 1, then read the column back.
float values[] = { 1.0f, 2.0f, 3.0f };
int64_t shape[] = { 3 };
int64_t strides[] = { 1 };
DLTensor tensor{};
tensor.data = values;
tensor.device = { kDLCPU, 0 };
tensor.ndim = 1;
tensor.dtype = { kDLFloat, 32, 1 }; // {code, bits, lanes}
tensor.shape = shape;
tensor.strides = strides;

ovstage_write_data_t write{};
write.tensors = &tensor;
write.tensor_count = 1;
write.is_array = false;

ovstage_enqueue_result_t enq =
    ovstage_write_attribute(stage, query, attrArg, /*ordinal*/ 1, write, OVSTAGE_PRIM_MODE_UPSERT);
waitOp(stage, enq, "write_attribute");

ovstage_write_floor_desc_t writeFloor{};
writeFloor.ordinal = 1;
writeFloor.scope = OVSTAGE_SCOPE_ALL;
enq = ovstage_advance_write_floor(stage, &writeFloor);
waitOp(stage, enq, "advance_write_floor");

ovstage_ordinal_range_t range{};
range.end_ordinal = 1;
range.has_start_ordinal = false;
ovstage_read_handle_t read = OVSTAGE_INVALID_READ_HANDLE;
enq = ovstage_read_attributes(stage, query, &attr, 1, range, &read);
waitOp(stage, enq, "read_attributes");

ovstage_read_group_t group{};
status = ovstage_fetch_read_next(stage, read, OVSTAGE_TIMEOUT_INFINITE, &group);
check(stage, status, "fetch_read_next");
if (group.data.tensor_count != 1 || !group.data.tensors[0].data)
{
    std::fprintf(stderr, "unexpected read layout\n");
    return EXIT_FAILURE;
}
const float* out = static_cast<const float*>(group.data.tensors[0].data);
std::printf("read back ordinal %llu: %.1f %.1f %.1f\n", (unsigned long long)group.ordinal, out[0], out[1], out[2]);
ovstage_release_group(stage, &group); // the tensor data is only valid until the group is released

For a write, exactly one of tensors (client-managed; must stay valid until the op completes) or managed_tensors (storage takes ownership through the DLManagedTensorVersioned deleter) is non-NULL.

In Python, the bindings accept and return DLPack-compatible arrays, so NumPy, PyTorch, and Warp tensors interchange through from_dlpack with the same residency model.

GPU Residency and Synchronization#

cuda_sync (ovstage_cuda_sync_t) coordinates GPU producer/consumer ordering, and remains the caller’s responsibility:

typedef struct {
    uintptr_t stream;     /* 0 = none, 1 = the default stream, >1 = a specific cudaStream_t */
    uintptr_t wait_event; /* 0 = none, else a cudaEvent_t to wait on */
} ovstage_cuda_sync_t;

The two fields are independent knobs — set either, both, or neither; each non-zero field adds its own synchronization before the op:

  • stream — 0 = none; 1 = the default stream; >1 = a specific cudaStream_t. When set, all work currently queued on that stream is drained.

  • wait_event — 0 = none; otherwise a cudaEvent_t that is waited on.

  • {0, 0} — no synchronization (CPU-resident or already-synchronized data).

On a read, the synchronization is applied before your code accesses the returned data. On a write / unmap, it is what ovstage waits on before sealing your data.

Sparsity#

index_map (gather / reorder / dedup) and mask (per-element validity) are mutually exclusive. Set count (the logical element count) whenever either is present.

Zero-Copy Map/Unmap#

Note

The map/unmap path below is documented from the shipped headers; the current examples cover only the CPU, copy-in path, so this flow is not yet snippet-backed. Treat the sequence as an API reference and refer to C API Reference for exact signatures.

Instead of copying data in, a producer can fill ovstage-owned storage directly:

  1. ovstage_map_attribute with an ovstage_map_desc_t ({ attribute; dtype; semantic; prim_mode }) — creating a new column requires desc.dtype.

  2. ovstage_fetch_map_next to obtain each map group.

  3. Fill group.data.tensors[i].data in place.

  4. ovstage_unmap_group per group, then ovstage_unmap_attribute to commit the remainder and release, passing a write_done_sync (an ovstage_cuda_sync_t).

All map/unmap ops are ordinal-keyed at the session ordinal. Changing an existing prim/name’s dtype or semantic fails — delete the attribute first.

Note

This build retains only the latest committed snapshot. Returned tensor data is valid for that snapshot only — copy, retain, or transfer ownership to use it later; do not hold a borrowed pointer across further commits.

Where to Go Next#