accelerator.py Source File

accelerator.py Source File#

SDK qb Runtime Library: accelerator.py Source File
SDK qb Runtime Library v1.3
MCS001-EN
 
from typing import Sequence
 
import numpy as np
 
import qbruntime.qbruntime as _cQbRuntime
from .pinned_memory import PinnedMemory
 
 
 
class Accelerator:
    """@brief Represents an accelerator, i.e., an NPU, used for executing models."""
 
    def __init__(self, dev_no: int = 0):
        """
        @brief Creates an Accelerator object for a specific device number.
 
        The `dev_no` parameter represents the device number. For example, on Linux,
        if an ARIES NPU is attached as `/dev/aries0`, the device number is `0`.
 
        @param dev_no The device number to associate with the Accelerator.
        """
        self._accelerator = _cQbRuntime.Accelerator(dev_no)
 
    def allocate_pinned_memory(
        self, shape: Sequence[int], dtype: np.dtype
    ) -> PinnedMemory:
        """
        @brief Allocates an NPU-accessible pinned memory buffer.
 
        The returned buffer can be used as input or output for zero-copy inference
        via `Model.infer_pinned_memory`.
 
        @note This is an experimental API and is only supported on Regulus hardware.
 
        @param shape The shape of the buffer to allocate.
        @param dtype The element type of the buffer. Must be `numpy.float32` or
                     `numpy.uint8`.
        @return A PinnedMemory wrapping the allocated buffer.
        """
        cobj = self._accelerator.allocate_pinned_memory(
            [int(s) for s in shape], np.dtype(dtype)
        )
        return PinnedMemory(cobj)
 
 