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 transported data rows 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.

Fixed-Size Canonical Layout#

Fixed-size tensors returned by raw reads and maps use a lane-canonical layout:

  • tensor_count == 1 and ndim == 1.

  • shape[0] == N, where N is the transported data-row count.

  • dtype.lanes is the complete tuple width.

Logical prims select those data rows one-to-one when index_map is NULL. If index_map is present, logical prim i selects row index_map[i] and N may differ from count: it can be smaller when rows are shared, or larger when a query touches only part of a transported bucket.

Copy-in writes may use either the canonical form or compact convenience dimensions. Convenience dimensions are folded into dtype.lanes and are not preserved as schema.

Without index_map, the leading dimension must equal the logical element count. A flat shape = [N * L], lanes = 1 tensor is not a convenience encoding of N rows of width L; use shape = [N, L], lanes = 1 or the canonical shape = [N], lanes = L form.

Value per data row

Canonical write / raw read / raw map

Accepted convenience write

Python DLPack export

Scalar

shape = [N], lanes = 1

Same

(N,)

point3f

shape = [N], lanes = 3

shape = [N, 3], lanes = 1

(N, 3)

matrix4d

shape = [N], lanes = 16

shape = [N, 4, 4], lanes = 1 (or [N, 4] with 4 lanes)

(N, 16)

Python DLPack export expands a multi-lane dtype by exactly one trailing axis; it does not reconstruct the convenience input shape. Array/ragged attributes are outside this fixed-size layout rule.

Important

In raw OVStage fixed-size transport, dtype.lanes is the canonical encoding of the complete component width; the attribute semantic separately supplies the geometric role. For example, three point3f data rows use shape = [3], lanes = 3, and OVSTAGE_SEMANTIC_POINT.

shape = [3, 3], lanes = 1 is accepted as a convenience write, but OVStage normalizes it. Raw reads and maps return shape = [3], lanes = 3 — not shape = [9], lanes = 1. The original trailing shape is not preserved. Python’s ReadGroup.array(0) is a separate flat base-element view of length 9; use ReadGroup.dlpack(0) when a consumer should see the expanded (3, 3) lane axis.

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.

Some producers expose library vector types as scalar DLPack tensors with a trailing component axis. For example, Warp exports vec3f as (N, 3) with lanes = 1. For an array-valued point3f[] row, first create float3 = DLDataType(code=DLDataTypeCode.kDLFloat, bits=32, lanes=3), then call make_dltensor(warp_points, dtype=float3). The adapter aliases the same allocation while folding only complete compact trailing axes into dtype.lanes. It validates the unchanged base type, a positive bit width, and byte extent, and requires byte-aligned source elements; no automatic shape-based inference is performed, so a scalar array with shape (N, 3) keeps its original meaning unless the caller requests the lane fold. When folding consumes every source axis, such as (3,) into a three-lane dtype, the adapter normalizes the mathematically rank-zero result to shape=(1,), ndim=1: a one-element tensor view compatible with ovstage’s leading-dimension transport. Explicit overrides for this view must use shape=[1], ndim=1, and compact strides=[1].

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.

They address different axes. index_map picks the source row each logical element reads (index_map[i] is a row of the payload, not a prim), so it gathers, reorders, or broadcasts data. mask picks the target elements to write, leaving the rest untouched. To write a subset of a query’s prims, use mask.

Where the payload declares a transported row count — shape[0] for a fixed-size payload, tensor_count for per-row array transport — every index_map entry must be below it, and an out-of-range entry is rejected with OVSTAGE_ERROR_INVALID_ARGUMENT. Unreferenced rows are left unused, so the payload may be wider than the query. Packed array transport carries no row count of its own, so the map defines the partition instead: the payload is cut into max(index_map) + 1 uniform rows, and a partition the payload cannot support (one that does not divide evenly, or whose rows are not a whole number of dtype elements) is rejected with OVSTAGE_ERROR_INVALID_ARGUMENT. A mask leaves the row partition alone, so a masked payload must still carry a row for every logical element, including the unselected ones.

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#