accelerator.py Source File

accelerator.py Source File#

SDK qb Runtime Library: accelerator.py Source File
SDK qb Runtime Library v1.4
MCS001-EN
 
from typing import Sequence, overload
 
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.
 
    @attention For the `regulus-usb` device family, do not unplug the USB device while
               the runtime is running. Disconnecting a device that is in use may crash
               the program or cause unexpected misbehavior.
    """
 
    @overload
    def __init__(self, dev_no: int = 0): ...
 
    @overload
    def __init__(self, device_name: str = "auto"): ...
 
    @overload
    def __init__(self, device_name: str = "auto", dev_no: int = 0): ...
 
    def __init__(self, device_name="auto", dev_no=0):
        """
        @brief Creates an Accelerator object for a specific target device and device number.
 
        The @p device_name parameter selects target device to open; case-insensitive.
        Supported values:
        - `"auto"`
        - `"aries"`, `"aries-rb"`
        - `"regulus"`, `"regulus-usb"`
        - `"regulus-ra"`, `"regulus-ra-usb"`
        - `"regulus-rb"`, `"regulus-rb-usb"`
 
        Automatic resolution:
        - `"auto"` resolves the target device automatically, but only when exactly one
        supported kind of target device is present in the system.
        - A base name (`"aries"`, `"regulus"`, `"regulus-usb"`, ...) is a shorthand
        for every target device whose name begins with it (e.g. `"aries"` matches
        `"aries-rb"`). It resolves to the concrete target device automatically, but
        only when exactly one such matching target device is present.
 
        The @p dev_no parameter is the device number. For example, on Linux an ARIES NPU
        attached as `/dev/aries0` has device number `0`.
 
        @param[in] device_name Device name to open (case-insensitive). See the list above.
        @param[in] dev_no The device number to associate with the Accelerator.
 
        @throws QbRuntimeError If @p device_name is unknown, if auto-resolution is ambiguous
                               (more than one candidate), or if no matching device exists.
 
        @note Supported constructor forms:
            - `Accelerator()`: auto-select the only attached NPU.
            - `Accelerator(dev_no)`: auto-detect device, open target device `dev_no`.
            - `Accelerator(device_name)`: open `device_name` at device 0.
            - `Accelerator(device_name, dev_no)`: open `device_name` at device `dev_no`.
        """
        if isinstance(device_name, int):  # Accelerator(dev_no) constructor
            dev_no, device_name = device_name, "auto"
        self._accelerator = _cQbRuntime.Accelerator(device_name, dev_no)
 
    def get_device_name(self) -> str:
        """
        @brief Returns the resolved device name of this accelerator.
 
        When created with a automatic resolve such as `"auto"`, `"aries"`,
        this returns the real target device name that was resolved
        (e.g. `"aries-rb"`, `"regulus-ra-usb"`).
 
        @return The resolved device name.
        """
        return self._accelerator.get_device_name()
 
    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 device.
 
        @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)
 
 