Class DeviceArray

tools/Polygraphy/polygraphy/cuda/cuda.py:356–485

An array on the GPU.

@mod.export()
class DeviceArray(DeviceView):
    """
    An array on the GPU.
    """

    def __init__(self, shape=None, dtype=None):
        """
        Args:
            shape (Tuple[int]): The initial shape of the buffer.
            dtype (numpy.dtype): The data type of the buffer.
        """
        super().__init__(ptr=0, shape=util.default(shape, tuple()), dtype=util.default(dtype, np.float32))
        self.allocated_nbytes = 0
        self.resize(self.shape)

    def __enter__(self):
        return self

    @staticmethod
    def raw(shape):
        """
        Creates an untyped device array of the specified shape.

        Args:
            shape (Tuple[int]):
                The initial shape of the buffer, in units of bytes.
                For example, a shape of ``(4, 4)`` would allocate a 16 byte array.

        Returns:
            DeviceArray: The raw device array.
        """
        return DeviceArray(shape=shape, dtype=np.byte)

    def resize(self, shape):
        """
        Resizes or reshapes the array to the specified shape.

        If the allocated memory region is already large enough,
        no reallocation is performed.

        Args:
            shape (Tuple[int]): The new shape.
        """
        nbytes = util.volume(shape) * self.itemsize
        if nbytes > self.allocated_nbytes:
            self.free()
            self.ptr = wrapper().malloc(nbytes)
            self.allocated_nbytes = nbytes
        self.shape = shape

    def __exit__(self, exc_type, exc_value, traceback):
        """
        Frees the underlying memory of this DeviceArray.
        """
        self.free()

    def free(self):
        """