cubecl_utils_rs/lib.rs
1//! Shared CubeCL helpers: GPU tensors, device-limit queries and validated
2//! dispatch geometry. No algorithms and no kernels.
3//!
4//! # Why this exists
5//!
6//! A CubeCL kernel dispatched with `launch_unchecked` that busts a device limit
7//! does not fail loudly:
8//!
9//! - Over-allocating **shared memory** makes the kernel do no work. It writes
10//! nothing, returns zeros and reports no error. Downstream code then reads
11//! uninitialised memory, which surfaces as an absurd index or a distance in
12//! an index slot rather than as anything pointing at the kernel.
13//! - Over-sizing a **binding** does the same.
14//! - Busting the **cube-count** limit is worse: the launch is rejected on the
15//! CubeCL server thread, that thread dies, and the next unrelated call on the
16//! client returns a `CallError` from somewhere else entirely.
17//!
18//! So device limits are a correctness concern, not a tuning one, and they are
19//! easy to get wrong when every machine to hand reports the same numbers. Apple
20//! Silicon via wgpu reports 32 KiB of shared memory, 65535 cubes per grid
21//! dimension and a plane size pinned to exactly 32. None of that is portable.
22//!
23//! # Design
24//!
25//! Every limit decision is a pure function of [`GpuLimits`]. Only
26//! [`GpuLimits::from_client`] and the [`GpuTensor`] constructors touch a
27//! `ComputeClient`; everything else takes limits as data.
28//!
29//! That is what makes the awkward cases testable. Asserting that a staging plan
30//! shrinks correctly on a 16 KiB device, or that a workgroup rounds to whole
31//! wave64 planes, needs no such device to be present.
32//!
33//! ```no_run
34//! # use cubecl::prelude::*;
35//! use cubecl_utils_rs::prelude::*;
36//!
37//! # fn demo<R: Runtime>(client: &ComputeClient<R>, n_blocks: u32) -> Result<(), CubeclUtilsErrors> {
38//! let limits = GpuLimits::from_client(client);
39//! let (gx, gy) = grid_2d(n_blocks, &limits)?;
40//! let count = checked_cube_count("my_kernel", gx, gy, 1, &limits)?;
41//! # Ok(())
42//! # }
43//! ```
44
45#![warn(missing_docs)]
46
47pub mod errors;
48pub mod layout;
49pub mod limits;
50pub mod prelude;
51pub mod tensor;
52pub mod traits;
53
54pub use crate::errors::CubeclUtilsErrors;
55pub use crate::layout::{pad_vectors, padded_dim, LINE_SIZE};
56pub use crate::limits::{
57 checked_cube_count, fits_binding, fits_shared_memory, grid_2d, grid_2d_limited,
58 plane_partitions, plane_uniform, resident_workgroups, resolve_workgroup_size, GpuLimits,
59};
60pub use crate::tensor::GpuTensor;
61pub use crate::traits::CubeclFloat;