npu_data.py Source File

npu_data.py Source File#

SDK qb Runtime Library: npu_data.py Source File
SDK qb Runtime Library v1.5
MCS001-KR
 
from typing import Tuple
 
import numpy as np
 
import qbruntime.qbruntime as _cQbRuntime
from .accelerator import Accelerator
 
_Shape = Tuple[int, ...]
 
__all__ = ["NPUData"]
 
# Maps the C++ DataType enum to the corresponding numpy dtype.
_DTYPE_MAP = {
    _cQbRuntime.DataType.Float32: np.float32,
    _cQbRuntime.DataType.Float16: np.float16,
    _cQbRuntime.DataType.Int8: np.int8,
    _cQbRuntime.DataType.Uint8: np.uint8,
}
 
 
 
class NPUData:
    """@brief A model input or output tensor that can reside on the host (CPU) or NPU.
 
    An NPUData backs a single model input/output tensor: it owns its host buffer,
    holds NPU memory that belongs to the `Accelerator` while the data is on the NPU,
    and lets the data be moved between CPU and NPU. Instances are created through
    `Model.acquire_input_npu_data` / `Model.acquire_output_npu_data`, not constructed
    directly.
 
    The data is accessible as a numpy array (via indexing, e.g. `nd[...]`) only while
    it resides on the CPU. Bring it back with `cpu` before reading results that were
    produced on the NPU.
 
    An NPU-resident NPUData may be passed to any model that expects a tensor of the
    same shape and dtype on the same accelerator, so several models can share one
    NPU-resident tensor without copying it back to the host in between.
 
    An NPUData is not owned by the `Model` it came from and may outlive it. Its NPU
    memory belongs to the `Accelerator` it resides on: if that accelerator is
    destroyed while the data is still on the NPU, the memory is reclaimed and the
    NPUData becomes invalid (`cpu`, `launch`, `copy_from` and indexing raise, and
    `dev_no` returns -1). A CPU-resident NPUData is unaffected.
 
    @note This is an advanced API rather than a typical usage.
    @warning This API is in beta: it may still contain bugs, and its behavior may
             change in a future release.
    @warning NPUData objects are not fully thread-safe; the client must serialize
             access to an NPUData from different threads. In particular, `launch` and
             `cpu` change where the data resides and must not run concurrently with
             each other.
    """
 
    def __init__(self, _npu_data: _cQbRuntime.NPUData):
        """@brief Internal constructor.
 
        @param _npu_data The underlying C++ NPUData object. Use
                    `Model.acquire_input_npu_data` / `Model.acquire_output_npu_data`
                    to create instances.
        """
        self._npu_data = _npu_data
        # Zero-copy numpy view over the host buffer, built once here (like PinnedMemory)
        # so indexing does not rebuild it. Unlike PinnedMemory, an NPUData is only
        # host-accessible while CPU-resident, so this is None while the data lives on
        # the NPU (acquired with upload=True or after launch()) and is (re)built when
        # the data is brought back via cpu(). Shares memory with the host buffer and is
        # only valid while this NPUData is alive and CPU-resident.
        self._array = self._make_array()
 
    def _make_array(self):
        """@brief Builds the numpy view, or returns None if not CPU-resident.
 
        The residency is checked first (dev_no() == -1 means the data is on the host):
        requesting the buffer protocol while the data is on the NPU would propagate a
        C++ exception out of pybind's `extern "C"` buffer slot and call std::terminate,
        so it must not be attempted.
        """
        if self._npu_data.dev_no() != -1:
            return None
        return np.asarray(self._npu_data)
 
    def launch(self, acc: Accelerator) -> None:
        """@brief Uploads the data to the NPU memory of the given accelerator.
 
        @warning Not thread-safe: see the class-level note.
 
        @param acc The accelerator on which to place the data.
        """
        self._npu_data.launch(acc._accelerator)
        # No longer host-accessible.
        self._array = None
 
    def cpu(self) -> None:
        """@brief Brings the data back to host (CPU) memory.
 
        @warning Not thread-safe: see the class-level note.
        """
        self._npu_data.cpu()
        # Host buffer is available again; rebuild the view.
        self._array = self._make_array()
 
    def copy_from(self, src: "NPUData") -> None:
        """@brief Copies the data from another NPUData.
 
        Both tensors must have the same shape and dtype; the source may reside on
        either side (a CPU-to-NPU or NPU-to-CPU copy is staged through the host).
 
        @param src The source NPUData to copy from.
        """
        self._npu_data.copy_from(src._npu_data)
 
    def __getitem__(self, key):
        """@brief Reads element(s) from the host buffer via numpy indexing."""
        if self._array is None:
            raise ValueError("NPUData is not on the CPU; call cpu() first.")
        return self._array[key]
 
    def __setitem__(self, key, value):
        """@brief Writes value(s) into the host buffer via numpy indexing."""
        if self._array is None:
            raise ValueError("NPUData is not on the CPU; call cpu() first.")
        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 data."""
        return tuple(self._npu_data.shape())
 
    @property
    def dtype(self) -> np.dtype:
        """@brief Returns the numpy dtype of the data's elements."""
        data_type = self._npu_data.data_type()
        if data_type not in _DTYPE_MAP:
            # NPUData rejects 16-bit float types at acquisition, so this only happens
            # for a corrupted object.
            raise ValueError(f"NPUData has no numpy equivalent for {data_type}.")
        return np.dtype(_DTYPE_MAP[data_type])
 
    @property
    def dev_no(self) -> int:
        """@brief Device number of the accelerator the data resides on.
 
        @return The device number when the data is on the NPU, or -1 when it is on the
                host (CPU) or has been invalidated (see the class-level note on
                accelerator destruction).
        """
        return self._npu_data.dev_no()
 
    @property
    def hardware_name(self) -> str:
        """@brief Hardware name (device type) of the accelerator the data resides on.
 
        Device numbers are assigned per hardware type (e.g. an "aries-rb" and a
        "regulus-ra" can both be device #0), so ``dev_no`` alone does not identify a
        physical device; pair it with this name to distinguish devices.
 
        @return The hardware name (e.g. "aries-rb") when the data is on the NPU, or an
                empty string when it is on the host (CPU) or has been invalidated (see
                the class-level note on accelerator destruction).
        """
        return self._npu_data.hardware_name()
 
 