tenzr::runtime::DeviceBuffer¶

RAII owner of a PJRT device buffer that awaits in-flight work before releasing it.

#include <device_buffer.h>

Public Functions¶

Public Functions Documentation¶

function DeviceBuffer¶

inline DeviceBuffer(
    PJRT_Buffer * buffer,
    const PJRT_Api * api
)

function ~DeviceBuffer¶

inline ~DeviceBuffer()

function DeviceBuffer¶

DeviceBuffer(
    const DeviceBuffer & 
) =delete

function operator=¶

DeviceBuffer & operator=(
    const DeviceBuffer & 
) =delete

function DeviceBuffer¶

inline DeviceBuffer(
    DeviceBuffer && other
)

function operator=¶

inline DeviceBuffer & operator=(
    DeviceBuffer && other
)

function Handle¶

inline PJRT_Buffer * Handle() const

Underlying PJRT buffer handle.

function Api¶

inline const PJRT_Api * Api() const

PJRT API used to operate on the buffer.

function UnsafePointer¶

inline void * UnsafePointer() const

Platform address of the on-device buffer (for device DLPack export).

function DeviceSizeBytes¶

inline size_t DeviceSizeBytes() const

Actual bytes allocated on the device for this buffer, or 0 when the plugin does not report it.

function AwaitReady¶

inline void AwaitReady() const

Await the buffer's readiness event (does not destroy the buffer).

function DeviceId¶

inline int DeviceId() const

Local hardware id of the buffer's storage device (0 when unavailable).

function IncreaseExternalReferenceCount¶

inline void IncreaseExternalReferenceCount() const

Pin the buffer while an external framework (DLPack/CuPy/torch) shares its pointer.

function DecreaseExternalReferenceCount¶

inline void DecreaseExternalReferenceCount() const

Release an external reference previously taken by increase_external_reference_count().

function CreateView¶

static inline std::shared_ptr< DeviceBuffer > CreateView(
    const PJRT_Api * api,
    PJRT_Client * client,
    PJRT_Device * device,
    void * device_buffer_ptr,
    const std::vector< int64_t > & dims,
    PJRT_Buffer_Type element_type,
    intptr_t stream,
    void(*)(void *, void *) on_delete,
    void * on_delete_arg
)

Create a non-owned PJRT view over an external device buffer (for zero-copy import).