vortex_torch.indexer.save_load

Classes

Load()

Format-aware load (copy/convert) dispatcher.

Save()

Format-aware save (copy/convert) dispatcher.
class vortex_torch.indexer.save_load.Save[source]

Bases: vOp

Format-aware save (copy/convert) dispatcher.

This operator copies or converts data from an input tensor x into a preallocated output tensor o. Both tensors are treated as rank-3 objects with logical shape

\[[S, D_0, D_1],\]

where \(S\) is a generic leading dimension (it may represent a batch, a sequence length, or a page count, depending on the surrounding runtime).

Key properties:

  • Dispatch is keyed only by the format of x (x._format).

  • No internal buffer is allocated; the provided output tensor o is validated and written into directly.

  • The output format is determined by the dispatch table and must match o._format.

_impl_map

Dispatch table keyed by x_format. Each entry maps to (callable_impl, resolved_output_format).

Type:

Dict[FORMAT, Tuple[Callable, FORMAT]]

impl

The resolved implementation selected during profile().

Type:

Optional[Callable]

output_format

The expected output format for o as determined in profile().

Type:

Optional[FORMAT]

profile(x, o, ctx)[source]

Validate inputs, resolve the implementation, and return o as the output view.

This method checks:

  • that both x and o are rank-3 vTensor instances with logical shape [S, D_0, D_1],

  • that their inner dimensions D_0 and D_1 match,

  • that a save implementation is registered for x._format, and

  • that o._format matches the format required by the dispatch.

No new buffers are allocated; the provided o will be used as the destination.

Parameters:

  • x (vTensor) – Input tensor to be copied/converted. Expected logical shape [S, D_0, D_1].

  • o (vTensor) – Preallocated output tensor. Must have logical shape compatible with x (matching D_0 and D_1) and the format expected by the selected implementation.

  • ctx (Context) – Execution context (included for API symmetry; not used for buffer allocation in this dispatcher).

Returns:

The same object as o, returned as the resolved output view.

Return type:

vTensor

Raises:

AssertionError – If types, ranks, shapes, formats, or devices are incompatible, or if no implementation is registered for x._format.

execute(x, o, ctx)[source]

Execute the resolved save operation from x into o.

The selected implementation is expected to copy or convert the contents of x into the preallocated tensor o without changing its shape or device.

Parameters:

  • x (torch.Tensor) – Input tensor to be copied/converted.

  • o (torch.Tensor) – Preallocated output tensor that will receive the data.

  • ctx (Context) – Execution context passed through to the implementation.

Returns:

The same tensor as o, after the copy/convert operation.

Return type:

torch.Tensor

Raises:

AssertionError – If profile() has not been called and no implementation is available.

class vortex_torch.indexer.save_load.Load[source]

Bases: vOp

Format-aware load (copy/convert) dispatcher.

This operator copies or converts data from an input tensor x into an internally allocated output buffer. Both input and output are treated as rank-3 tensors with logical shape

\[[S, D_0, D_1],\]

where \(S\) is a generic leading dimension (it may represent a batch size, a sequence length, or a page count, depending on the surrounding runtime).

Key properties

  • Dispatch is keyed only by the format of x (x._format).

  • The output buffer is allocated during profile() using the runtime value ctx.max_num_pages for the leading dimension.

  • The output format is determined by the dispatch table and used when wrapping the internal buffer as a vTensor.

_impl_map

Dispatch table keyed by x_format. Each entry maps to (callable_impl, resolved_output_format).

Type:

Dict[FORMAT, Tuple[Callable, FORMAT]]

impl

The resolved implementation selected during profile().

Type:

Optional[Callable]

output_format

The output tensor format as determined in profile().

Type:

Optional[FORMAT]

output_buffer

Preallocated output tensor buffer with logical shape [S_out, D_0, D_1], where S_out comes from the runtime context.

Type:

Optional[torch.Tensor]

profile(x, ctx)[source]

Validate the input, select an implementation, allocate the output buffer, and return an as_vtensor() view with the resolved format.

The input tensor is expected to have logical shape [S_in, D_0, D_1]. The output tensor is allocated with shape

\[[S_{\text{out}}, D_0, D_1],\]

where \(S_{\text{out}}\) is taken from ctx.max_num_pages to match the runtime configuration for the leading dimension.

Parameters:

  • x (vTensor) – Input tensor to be loaded from, with logical shape [S_in, D_0, D_1].

  • ctx (Context) – Execution context providing ctx.max_num_pages for the leading dimension and tracking auxiliary memory usage.

Returns:

A vTensor view wrapping the internally allocated output buffer with the resolved output format.

Return type:

vTensor

Raises:

AssertionError – If x is not a vTensor, if its rank is not 3, or if no implementation is registered for x._format.

execute(x, ctx)[source]

Run the selected load operation into the internally allocated buffer.

The selected implementation is expected to copy or convert the contents of x into the internal buffer stored in output_buffer.

Parameters:

  • x (torch.Tensor) – Input tensor to be copied/converted, on the same device as the internal output buffer.

  • ctx (Context) – Execution context passed through to the implementation.

Returns:

The internally allocated output tensor stored in output_buffer.

Return type:

torch.Tensor

Raises:

AssertionError – If profile() has not been called (no implementation or internal output buffer available).