Skip to main content

Crate cubecl_utils_rs

Crate cubecl_utils_rs 

Source
Expand description

Shared CubeCL helpers: GPU tensors, device-limit queries and validated dispatch geometry. No algorithms and no kernels.

§Why this exists

A CubeCL kernel dispatched with launch_unchecked that busts a device limit does not fail loudly:

  • Over-allocating shared memory makes the kernel do no work. It writes nothing, returns zeros and reports no error. Downstream code then reads uninitialised memory, which surfaces as an absurd index or a distance in an index slot rather than as anything pointing at the kernel.
  • Over-sizing a binding does the same.
  • Busting the cube-count limit is worse: the launch is rejected on the CubeCL server thread, that thread dies, and the next unrelated call on the client returns a CallError from somewhere else entirely.

So device limits are a correctness concern, not a tuning one, and they are easy to get wrong when every machine to hand reports the same numbers. Apple Silicon via wgpu reports 32 KiB of shared memory, 65535 cubes per grid dimension and a plane size pinned to exactly 32. None of that is portable.

§Design

Every limit decision is a pure function of GpuLimits. Only GpuLimits::from_client and the GpuTensor constructors touch a ComputeClient; everything else takes limits as data.

That is what makes the awkward cases testable. Asserting that a staging plan shrinks correctly on a 16 KiB device, or that a workgroup rounds to whole wave64 planes, needs no such device to be present.

use cubecl_utils_rs::prelude::*;

let limits = GpuLimits::from_client(client);
let (gx, gy) = grid_2d(n_blocks, &limits)?;
let count = checked_cube_count("my_kernel", gx, gy, 1, &limits)?;

Re-exports§

pub use crate::errors::CubeclUtilsErrors;
pub use crate::layout::pad_vectors;
pub use crate::layout::padded_dim;
pub use crate::layout::LINE_SIZE;
pub use crate::limits::checked_cube_count;
pub use crate::limits::fits_binding;
pub use crate::limits::fits_shared_memory;
pub use crate::limits::grid_2d;
pub use crate::limits::grid_2d_limited;
pub use crate::limits::plane_partitions;
pub use crate::limits::plane_uniform;
pub use crate::limits::resident_workgroups;
pub use crate::limits::resolve_workgroup_size;
pub use crate::limits::GpuLimits;
pub use crate::tensor::GpuTensor;
pub use crate::traits::CubeclFloat;

Modules§

errors
Errors in cubecl-utils-rs.
layout
Row padding to a vectorisation boundary.
limits
Device limits and the dispatch geometry derived from them.
prelude
User-facing re-exports.
tensor
GPU-resident tensor.
traits
Shared trait bounds for element-generic kernels.