Buffer Creation and Usage

In vkdispatch, nearly all data is stored inside “buffers” (each wrapping an individual VkBuffer object and all other objects needed to manage it). These are the equivalent of torch.Tensor or wp.array() in warp-lang.

However, unlike torch.Tensor or wp.array(), vkdispatch buffers are by default multi-device. This means that when a vkdispatch buffer is allocated on a multi-device or multi-queue context, multiple VkBuffers are allocated (one for each queue on each device). This architecture has the benefit of greatly simplfying multi-GPU programs, since all buffers can be assumed to exist on all devices and all queues.

To allocate a buffer, you can either directly use the constructor of the Buffer class or the vd.asbuffer function to directly upload a numpy array to the gpu as a buffer.

Simple GPU Buffer Example

import vkdispatch as vd
import numpy as np

# Create a simple numpy array
cpu_data = np.arange(16, dtype=np.int32)
print(f"Original CPU data: {cpu_data}")

# Create a GPU buffer
gpu_buffer = vd.asbuffer(cpu_data)

# Read data back from GPU to CPU to verify
downloaded_data = gpu_buffer.read(0)
print(f"Data downloaded from GPU: {downloaded_data.flatten()}")

# Expected Output:
# Original CPU data: [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15]
# Data downloaded from GPU: [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15]

What’s happening here?

  1. We import vkdispatch and numpy (a common dependency for numerical data).

  2. We use the vd.asbuffer function to upload the numpy array to a vkdispatch buffer.

  3. read(0) retrieves data back from the GPU to the CPU. The number provided as an argument to the function is the queue index to read from. For a simple context with one device and one queue, there is only 1 queue, so we read from index 0. If the index is ommited the function returns a python list of the contents of all buffers on all queues and devices.

Buffer API Reference

class vkdispatch.Buffer(shape: Tuple[int, ...], var_type: dtype)

TODO: Docstring

__init__(shape: Tuple[int, ...], var_type: dtype) → None
_destroy() → None

Destroy the buffer and all child handles.

read(index: int | None = None) → ndarray

Read the data in the buffer at the specified device index and return it as a numpy array.

Parameters: index (int): The index to read the data from. Default is 0.

Returns: (np.ndarray): The data in the buffer as a numpy array.

write(data: bytes | ndarray, index: int = -1) → None

Given data in some numpy array, write that data to the buffer at the specified index. The default index of -1 will write to all buffers.

Parameters: data (np.ndarray): The data to write to the buffer. index (int): The index to write the data to. Default is -1 and

will write to all buffers.

Returns: None

vkdispatch.asbuffer(array: ndarray) → Buffer

Cast a numpy array to a buffer object.