Struct GpuTensor

Source

pub struct GpuTensor<R: Runtime, F: CubeElement + Numeric> { /* private fields */ }

Expand description

GPU-resident tensor for use with CubeCL kernels.

Implementations§

Source §

impl<R: Runtime, F: Numeric + CubeElement> GpuTensor<R, F>

Source

pub fn from_slice( data: &[F], shape: Vec<usize>, client: &ComputeClient<R>, ) -> Result<Self, CubeclUtilsErrors>

Create a tensor from CPU data.

§Params

data - Slice of values to upload
shape - Dimensions of the tensor
client - GPU compute client for memory allocation

§Returns

A new tensor with the data copied to GPU memory, or BindingTooLarge when the allocation exceeds what this device binds in one go.

Source

pub fn empty( shape: Vec<usize>, client: &ComputeClient<R>, ) -> Result<Self, CubeclUtilsErrors>

Create an uninitialised tensor.

§Params

shape - Dimensions of the tensor
client - GPU compute client for memory allocation

§Returns

A new tensor with allocated but uninitialised GPU memory, or BindingTooLarge when the allocation exceeds what this device binds in one go.

§Note

The allocation returns quickly but its pages are not backed until something writes them. On a large buffer the first kernel write pays that fault and can cost more than the kernel itself, so prefer reusing one scratch tensor over allocating per call. See Self::reshaped_view.

Source

pub fn into_tensor_arg(&self) -> TensorArg<R>

Convert to a TensorArg for kernel launches.

§Returns

A TensorArg suitable for passing to CubeCL kernels. Vectorisation width is not set per tensor; it is passed once at launch as the argument for the kernel’s N: Size generic.

Source

pub fn read( self, client: &ComputeClient<R>, ) -> Result<Vec<F>, CubeclUtilsErrors>

Read tensor data back to CPU.

Consumes the tensor and transfers data from GPU to CPU memory.

§Params

client - GPU compute client for memory transfer

§Returns

Vector of exactly Self::len elements.

§Note

The read is truncated to the shape. That only matters for a tensor produced by Self::reshaped_view, where the underlying allocation is larger than the view: the runtime hands back the whole binding, and returning that would silently include another view’s data.

Source

pub fn shape(&self) -> &[usize]

Dimensions of the tensor.

§Returns

Slice of the per-dimension extents, outermost first.

Source

pub fn len(&self) -> usize

Number of elements the underlying allocation holds.

§Returns

Product of the shape dimensions.

Source

pub fn is_empty(&self) -> bool

Whether the tensor holds no elements.

§Returns

True if the element count is zero.

Source

pub fn reshaped_view(&self, shape: Vec<usize>) -> Self

Reinterpret an existing allocation under a smaller shape.

Shares the underlying buffer rather than allocating, so callers can keep one scratch tensor alive across several differently shaped uses. The first kernel write to a fresh allocation faults its pages in, which for a large buffer costs more than the kernel itself, so reuse is worth the sharp edge.

§Params

shape - New shape; its element count must not exceed the current one

§Returns

A tensor sharing this one’s buffer, with row-major strides for shape.

§Note

The returned tensor aliases self. Writing through both concurrently is a data race the type system does not prevent here.

Source

pub fn vram_bytes(&self) -> usize

Size of the tensor on the GPU, in bytes.

§Returns

Element count multiplied by the element size.

Source

pub fn handle(&self) -> &Handle

Return the handle of the tensor.

Escape hatch for crates that need to hand the raw buffer to their own kernels or to a library matmul without going through GpuTensor::into_tensor_arg.