pinned_memory.py Source File

pinned_memory.py Source File#

SDK qb Runtime Library: pinned_memory.py Source File
SDK qb Runtime Library v1.3
MCS001-EN
 
from typing import Tuple, Union
 
import numpy as np
 
import qbruntime.qbruntime as _cQbRuntime
 
_Shape = Tuple[int, ...]
 
__all__ = ["PinnedMemory"]
 
 
 
class PinnedMemory:
    """@brief An NPU-accessible pinned (physically contiguous) memory buffer.
 
    A PinnedMemory wraps a buffer that is pinned in device memory so the NPU can
    read from / write to it directly, enabling zero-copy inference via
    `Model.infer_pinned_memory`. Instances are created through
    `Accelerator.allocate_pinned_memory`, not constructed directly.
 
    @note This is an experimental API and is only supported on Regulus hardware.
    """
 
    def __init__(
        self,
        _pinned_memory: Union[
            _cQbRuntime.PinnedMemoryFloat, _cQbRuntime.PinnedMemoryUint8
        ],
    ):
        """@brief Internal constructor.
 
        @param _pinned_memory The underlying C++ pinned-memory object. Use
                    `Accelerator.allocate_pinned_memory` to create instances.
        """
        self._pinned_memory = _pinned_memory
        # Zero-copy numpy view over the pinned buffer. Built once here so that
        # indexing does not rebuild it. Shares memory with the pinned buffer and
        # is only valid while this PinnedMemory is alive and available.
        self._array = np.asarray(_pinned_memory)
 
    def __getitem__(self, key):
        """@brief Reads element(s) from the pinned buffer via indexing.
 
        Supports the same indexing as numpy (integers, slices, tuples, etc.),
        e.g. `pm[0]`, `pm[1, 2]`, `pm[:10]`. The data is read directly from the
        pinned buffer without an intermediate copy.
 
        @param key The index or slice.
        @return The selected element or a view into the pinned buffer.
        """
        return self._array[key]
 
    def __setitem__(self, key, value):
        """@brief Writes value(s) into the pinned buffer via indexing.
 
        Writes directly into the pinned buffer, e.g. `pm[0] = 1.0`,
        `pm[:10] = src`. This is the recommended way to fill an input buffer
        before `Model.infer_pinned_memory`.
 
        @param key The index or slice.
        @param value The value(s) to write.
        """
        self._array[key] = value
 
    def __len__(self) -> int:
        """@brief Returns the length of the buffer's first dimension."""
        return self.shape[0]
 
    @property
    def shape(self) -> _Shape:
        """@brief Returns the shape of the buffer."""
        return tuple(self._pinned_memory.shape())
 
    @property
    def dtype(self) -> np.dtype:
        """@brief Returns the numpy dtype of the buffer's elements."""
        return self._array.dtype
 
    @property
    def size(self) -> int:
        """@brief Returns the size of the buffer in bytes."""
        return self._pinned_memory.size()
 
    @property
    def is_available(self) -> bool:
        """@brief Returns whether the buffer is allocated and still usable."""
        return self._pinned_memory.is_available()
 
 