diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 70f7f6412..3760c94e5 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -57,6 +57,9 @@ cubek-fft = { path = "../crates/cubek-fft", version = "=0.3.0-pre.1", features = cubek-interpolate = { path = "../crates/cubek-interpolate", version = "=0.3.0-pre.1", features = [ "benchmarks", ] } +cubek-linalg = { path = "../crates/cubek-linalg", version = "=0.3.0-pre.1", features = [ + "benchmarks", +] } cubek-matmul = { path = "../crates/cubek-matmul", version = "=0.3.0-pre.1", features = [ "benchmarks", ] } @@ -139,3 +142,7 @@ name = "fft" [[bench]] harness = false name = "quantized_matmul" + +[[bench]] +harness = false +name = "qr" diff --git a/benchmarks/benches/qr.rs b/benchmarks/benches/qr.rs new file mode 100644 index 000000000..aec409450 --- /dev/null +++ b/benchmarks/benches/qr.rs @@ -0,0 +1 @@ +benchmarks::run_bench!(qr); diff --git a/benchmarks/src/lib.rs b/benchmarks/src/lib.rs index a4c62d93c..d733040b1 100644 --- a/benchmarks/src/lib.rs +++ b/benchmarks/src/lib.rs @@ -5,6 +5,7 @@ pub use cubek_attention::eval::forward::benchmarks as attention; pub use cubek_convolution::eval::benchmarks as conv2d; pub use cubek_fft::eval::benchmarks as fft; pub use cubek_interpolate::eval::benchmarks as interpolate; +pub use cubek_linalg::eval::benchmarks as qr; pub use cubek_matmul::eval::benchmarks::gemm; pub use cubek_matmul::eval::benchmarks::gemm_cpu; pub use cubek_matmul::eval::benchmarks::gemm_cpu_tiled; @@ -38,6 +39,7 @@ pub fn all() -> &'static [&'static dyn BenchmarkCategory] { &crate::interpolate::Category, &crate::memcpy_async::Category, &crate::pool::Category, + &crate::qr::Category, &crate::quantized_matmul::Category, &crate::reduce::Category, &crate::split_k::Category, diff --git a/crates/cubek-linalg/Cargo.toml b/crates/cubek-linalg/Cargo.toml new file mode 100644 index 000000000..e7928b1c1 --- /dev/null +++ b/crates/cubek-linalg/Cargo.toml @@ -0,0 +1,36 @@ +[package] +authors = [ + "Jorge Perez Burgos ", +] +categories = ["science", "mathematics", "algorithms"] +description = "CubeK: Linear Algebra Kernels" +edition.workspace = true +keywords = [] +license.workspace = true +name = "cubek-linalg" +readme.workspace = true +repository = "https://github.com/tracel-ai/cubek/tree/main/crates/cubek-linalg" +version.workspace = true + +[features] +default = ["std", "cubecl/default"] +std = ["cubecl/std", "thiserror/std"] +# Enables `pub mod eval::benchmarks` — catalogue of QR benchmark problems / +# strategies consumed by the benchmark registry. +benchmarks = ["dep:cubek-test-utils", "cubecl/test-runtime"] + +[dependencies] +cubecl = { workspace = true, features = ["stdlib"] } +cubecl-common = { workspace = true } +cubek-std = { path = "../cubek-std", version = "=0.3.0-pre.1", default-features = false } +cubek-matmul = { path = "../cubek-matmul", version = "=0.3.0-pre.1", default-features = false } +cubek-test-utils = { path = "../cubek-test-utils", version = "=0.3.0-pre.1", default-features = false, optional = true } + +thiserror = { workspace = true } + +[dev-dependencies] +cubecl = { workspace = true, features = ["test-runtime"] } +cubecl-common = { workspace = true } +cubek-test-utils = { path = "../cubek-test-utils", version = "=0.3.0-pre.1", default-features = false } +num-traits = { workspace = true } +paste = "1.0" diff --git a/crates/cubek-linalg/examples/profile_qr.rs b/crates/cubek-linalg/examples/profile_qr.rs new file mode 100644 index 000000000..39645ec12 --- /dev/null +++ b/crates/cubek-linalg/examples/profile_qr.rs @@ -0,0 +1,75 @@ +//! Minimal driver for profiling the QR decomposition under nsys/ncu. +//! +//! Run with: +//! ```sh +//! cargo build --release -p cubek-linalg --example profile_qr --features cubecl/cuda +//! nsys profile --stats=true target/release/examples/profile_qr [m] [n] [iters] [dtype] +//! ``` +//! +//! `dtype` is `f32` (default) or `f64`; pass a 5th arg `tf32` to opt the +//! trailing-update GEMMs into tensor cores (see `BahtTsqrStrategy`). + +use cubecl::prelude::*; +use cubecl::std::tensor::TensorHandle; +use cubecl::{Runtime, TestRuntime, future}; + +fn run(m: usize, n: usize, iters: usize, label: &str, allow_tf32: bool) { + let client = TestRuntime::client(&Default::default()); + + // Col-major pseudo-random data, same layout the tests use. + let mut state = 0x2545F4914F6CDD1Du64; + let data: Vec = (0..m * n) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + let v = (state as f64 / u64::MAX as f64) - 0.5; + ::from(v).unwrap() + }) + .collect(); + let handle = client.create_from_slice(E::as_bytes(&data)); + let a = TensorHandle::::new( + handle, + vec![m, n], + vec![1, m], + E::as_type_native_unchecked(), + ); + + let strategy = cubek_linalg::routines::BahtTsqrStrategy { allow_tf32 }; + + // Warmup (JIT compilation). + cubek_linalg::qr_with_strategy::(&client, &a, strategy).unwrap(); + future::block_on(client.sync()).unwrap(); + + let mut times_ms = Vec::with_capacity(iters); + for i in 0..iters { + let start = std::time::Instant::now(); + cubek_linalg::qr_with_strategy::(&client, &a, strategy).unwrap(); + future::block_on(client.sync()).unwrap(); + let elapsed = start.elapsed(); + times_ms.push(elapsed.as_secs_f64() * 1e3); + println!("iter {i}: {elapsed:?}"); + } + + times_ms.sort_by(|x, y| x.partial_cmp(y).unwrap()); + let mean: f64 = times_ms.iter().sum::() / times_ms.len() as f64; + let median = times_ms[times_ms.len() / 2]; + println!("{m}x{n} {label}: mean {mean:.3} ms, median {median:.3} ms"); +} + +fn main() { + let mut args = std::env::args().skip(1); + let m: usize = args.next().and_then(|a| a.parse().ok()).unwrap_or(2048); + let n: usize = args.next().and_then(|a| a.parse().ok()).unwrap_or(2048); + let iters: usize = args.next().and_then(|a| a.parse().ok()).unwrap_or(3); + let dtype = args.next().unwrap_or_else(|| "f32".to_string()); + // Optional 5th arg: `tf32` opts the trailing-update GEMMs into tensor cores. + let allow_tf32 = args.next().as_deref() == Some("tf32"); + let label = if allow_tf32 { "tf32" } else { "full" }; + + match dtype.as_str() { + "f64" => run::(m, n, iters, &format!("f64/{label}"), allow_tf32), + "f32" => run::(m, n, iters, &format!("f32/{label}"), allow_tf32), + other => panic!("unsupported dtype `{other}`; use f32 or f64"), + } +} diff --git a/crates/cubek-linalg/src/components/baht_tsqr.rs b/crates/cubek-linalg/src/components/baht_tsqr.rs new file mode 100644 index 000000000..73dca0f95 --- /dev/null +++ b/crates/cubek-linalg/src/components/baht_tsqr.rs @@ -0,0 +1,577 @@ +//! # TSQR-inspired blocked Householder QR +//! +//! Factorizes the whole panel with two kernel dispatches per column — a +//! reflector build and a panel update, both parallel across a full cube with +//! shared-memory tree reductions — builds the block reflector's T matrix from +//! a V^T·V Gram GEMM, and applies the trailing updates through GEMMs. + +use cubecl::calculate_cube_count_elemwise; +use cubecl::prelude::*; +use cubecl::std::tensor::TensorHandle; +use cubek_matmul::definition::MatmulElems; +use cubek_matmul::strategy::Strategy; +use cubek_std::InputBinding; + +use crate::definition::QRSetupError; +use crate::routines::{BahtTsqrBlueprint, BahtTsqrLaunchSettings}; + +#[cube(launch_unchecked)] +fn clear_buffer_kernel(buf: &mut [F], n: u32) { + let idx = ABSOLUTE_POS_X; + if idx < n { + buf[idx as usize] = F::cast_from(0.0); + } +} + +/// Sum-reduce each unit's `local` partial across the cube; every unit +/// returns the total and the shared buffer is free for reuse on return. +/// +/// With `use_plane_reduce` (blueprint-gated on plane-op support and a uniform +/// plane size) the partials are folded by the hardware `plane_sum`, one value +/// per plane is merged through `reduce`, and only two barriers are needed. +/// Otherwise a power-of-two shared-memory tree (`log2(cube_dim)` barriers) +/// does the same job on any hardware. +#[cube] +fn cube_reduce_sum( + reduce: &mut Shared<[F]>, + local: F, + tdx: u32, + cube_dim_x: u32, + #[comptime] use_plane_reduce: bool, +) -> F { + let total = if comptime![use_plane_reduce] { + let plane_total = plane_sum(local); + if UNIT_POS_PLANE == 0 { + reduce[(tdx / PLANE_DIM) as usize] = plane_total; + } + sync_cube(); + let num_planes = cube_dim_x.div_ceil(PLANE_DIM); + let mut acc = F::cast_from(0.0); + for p in 0..num_planes { + acc += reduce[p as usize]; + } + acc + } else { + reduce[tdx as usize] = local; + sync_cube(); + let mut pow2 = 1u32; + while pow2 < cube_dim_x { + if (tdx as usize).is_multiple_of((pow2 as usize) * 2) && (tdx + pow2 < cube_dim_x) { + let val = reduce[(tdx + pow2) as usize]; + reduce[tdx as usize] += val; + } + pow2 *= 2; + sync_cube(); + } + reduce[0] + }; + // Every unit finishes reading before the caller may reuse the buffer. + sync_cube(); + total +} + +/// Compute the Householder reflector for one panel column and write it +/// directly into column `j` of `v_buf` (col-major `[rows, tile]`, cleared at +/// the start of the tile) with the conventional implicit-1 head: +/// `v_buf[j*rows + col] = 1`, tail `= r[k]/v0`. +/// +/// One cube; the `sigma = ||r_tail||²` reduction runs through +/// [`cube_reduce_sum`] and the normalization is parallel over the cube. All flat offsets +/// are computed in `usize` so `col*rows` cannot wrap on huge matrices. +#[cube(launch_unchecked)] +fn householder_kernel( + r: &[F], + rows: u32, + col: u32, + v_buf: &mut [F], + beta_vec: &mut [F], + j: u32, + #[comptime] shared_size: usize, + #[comptime] use_plane_reduce: bool, +) { + let tdx = UNIT_POS_X; + let cube_dim_x = CUBE_DIM_X; + let zero = F::cast_from(0.0); + let one = F::cast_from(1.0); + let two = F::cast_from(2.0); + + let dim = rows - col; + let r_offset = col as usize * rows as usize + col as usize; + + let mut reduce = Shared::<[F]>::new_slice(shared_size); + let mut local = zero; + let mut i = tdx + 1; + while i < dim { + let v = r[r_offset + i as usize]; + local = fma(v, v, local); + i += cube_dim_x; + } + + let sigma = cube_reduce_sum::(&mut reduce, local, tdx, cube_dim_x, use_plane_reduce); + + if tdx == 0 { + if sigma == zero { + beta_vec[j as usize] = zero; + reduce[0] = one; + } else { + let r0 = r[r_offset]; + let mu = F::sqrt(fma(r0, r0, sigma)); + let v0 = if r0 <= zero { + r0 - mu + } else { + -sigma / (r0 + mu) + }; + let v0sq = v0 * v0; + beta_vec[j as usize] = -two * v0sq / (v0sq + sigma); + reduce[0] = v0; + } + } + sync_cube(); + let v0 = reduce[0]; + + let v_base = j as usize * rows as usize + col as usize; + let mut k = tdx; + while k < dim { + if k == 0 { + v_buf[v_base] = one; + } else { + v_buf[v_base + k as usize] = r[r_offset + k as usize] / v0; + } + k += cube_dim_x; + } +} + +/// Apply the reflector held in `v_buf` column `j` to the remaining panel +/// columns of R. One cube per target column: the `v·r` dot product is +/// reduced through [`cube_reduce_sum`], then the rank-1 update runs parallel +/// over the cube. The reflector head is the explicit 1 written by +/// [`householder_kernel`], so the dot spans the full `dim` range uniformly. +#[cube(launch_unchecked)] +fn apply_householder_kernel( + rows: u32, + col: u32, + r: &mut Tensor, + v_buf: &[F], + beta_vec: &[F], + j: u32, + #[comptime] shared_size: usize, + #[comptime] use_plane_reduce: bool, +) { + let tdx = UNIT_POS_X; + let cube_dim_x = CUBE_DIM_X; + let target_col = col + CUBE_POS_X; + let dim = rows - col; + let r_base = target_col as usize * rows as usize + col as usize; + let v_base = j as usize * rows as usize + col as usize; + + let mut reduce = Shared::<[F]>::new_slice(shared_size); + let mut local = F::cast_from(0.0); + let mut k = tdx; + while k < dim { + local = fma(v_buf[v_base + k as usize], r[r_base + k as usize], local); + k += cube_dim_x; + } + + let dot_f = cube_reduce_sum::(&mut reduce, local, tdx, cube_dim_x, use_plane_reduce) + * beta_vec[j as usize]; + + let mut k2 = tdx; + while k2 < dim { + r[r_base + k2 as usize] += v_buf[v_base + k2 as usize] * dot_f; + k2 += cube_dim_x; + } +} + +/// Build the block reflector's T matrix from the Gram matrix. Columns of T +/// depend on all previous columns, so `j` advances sequentially with a cube +/// sync per step, but the rows `i < j` within a column are independent and +/// are computed in parallel across the cube. +/// +/// T and the Gram matrix are both staged in shared memory: each of the `tile` +/// serialized steps reads back columns written by earlier steps, so building +/// T directly in global memory puts a global round trip inside every step of +/// the dependency chain. The whole working set is `tile * tile` elements +/// (4 KiB at f32/tile=32), so it fits comfortably. +#[cube(launch_unchecked)] +fn build_t_tsqr_kernel( + current_tile: u32, + gram: &[F], + beta_vec: &[F], + t_mat: &mut [F], + #[comptime] tile: u32, +) { + let tdx = UNIT_POS_X; + let cube_dim_x = CUBE_DIM_X; + let zero = F::cast_from(0.0); + let total = comptime!(tile * tile); + + let mut t_sh = Shared::<[F]>::new_slice(comptime!((tile * tile) as usize)); + let mut g_sh = Shared::<[F]>::new_slice(comptime!((tile * tile) as usize)); + + let mut idx = tdx; + while idx < total { + t_sh[idx as usize] = zero; + g_sh[idx as usize] = gram[idx as usize]; + idx += cube_dim_x; + } + sync_cube(); + + for j in 0u32..current_tile { + if tdx == 0 { + t_sh[(j * tile + j) as usize] = beta_vec[j as usize]; + } + let mut i = tdx; + while i < j { + let mut sum = zero; + for k in i..j { + sum = fma( + t_sh[(k * tile + i) as usize], + g_sh[(k * tile + j) as usize], + sum, + ); + } + t_sh[(j * tile + i) as usize] = beta_vec[j as usize] * sum; + i += cube_dim_x; + } + sync_cube(); + } + + let mut o = tdx; + while o < total { + t_mat[o as usize] = t_sh[o as usize]; + o += cube_dim_x; + } +} + +/// R's trailing block is col-major while Z is row-major, so this update is a +/// transposing add and cannot be vectorized on both operands at once. Making Z +/// col-major to fix that measurably de-vectorizes the matmul that writes it +/// (its output write drops from `acc_size_4` to `acc_size_1`), which costs +/// more than this kernel saves — so it stays scalar and 2D. +#[cube(launch_unchecked)] +fn update_trailing_r_kernel( + rows: u32, + cols: u32, + col_start_trailing: u32, + r: &mut Tensor, + z_buf: &[F], +) { + let row = ABSOLUTE_POS_X; + let col = ABSOLUTE_POS_Y; + let trailing_cols = cols - col_start_trailing; + if row < rows && col < trailing_cols { + let r_idx = (col_start_trailing + col) as usize * rows as usize + row as usize; + let z_idx = row as usize * trailing_cols as usize + col as usize; + r[r_idx] += z_buf[z_idx]; + } +} + +/// `dst += z` over a tight contiguous region, launched 1D over `Vector` lanes +/// so the loads/stores are widened. Used for the Q^T update, whose two +/// operands are tight col-major `[rows, rows]` buffers with identical flat +/// indexing. +#[cube(launch_unchecked)] +fn add_assign_kernel( + n_vecs: u32, + dst: &mut [Vector], + z_buf: &[Vector], +) { + let idx = ABSOLUTE_POS_X; + if idx < n_vecs { + dst[idx as usize] += z_buf[idx as usize]; + } +} + +/// Launch QR decomposition using the TSQR-inspired kernels and GEMM updates. +pub fn launch( + client: &ComputeClient, + q_handle: &TensorHandle, + r_handle: &TensorHandle, + blueprint: BahtTsqrBlueprint, + settings: BahtTsqrLaunchSettings, +) -> Result<(), QRSetupError> { + let rows = r_handle.shape()[0] as u32; + let cols = r_handle.shape()[1] as u32; + let use_plane_reduce = blueprint.use_plane_reduce; + + let BahtTsqrLaunchSettings { + tile, + max_cube_dim, + thread_block_size, + cube_dim_2d, + strategy_gram, + strategy_w, + strategy_tall, + } = settings; + + let num_tiles = cols.div_ceil(tile); + let dtype = E::as_type_native_unchecked(); + let storage_dtype = dtype.storage_type(); + let elem_size = storage_dtype.size(); + + let rows_us = rows as usize; + let cols_us = cols as usize; + let tile_us = tile as usize; + + // Raw scratch allocations: every consumer below re-declares them with + // hand-computed tight strides (or indexes them flat), and each region is + // fully written before it is read, so no zero-fill dispatches are needed + // and no allocator stride policy can interfere. + let beta_vec = client.empty(tile_us * elem_size); + let v_buf_global = client.empty(rows_us * tile_us * elem_size); + let w_buf_global = client.empty(rows_us * tile_us * elem_size); + let gram_buf_global = client.empty(tile_us * tile_us * elem_size); + let t_buf_global = client.empty(tile_us * tile_us * elem_size); + let s_buf_global = client.empty(tile_us * cols_us * elem_size); + let s_tile_global = client.empty(tile_us * rows_us * elem_size); + let z_buf_global = client.empty(rows_us * rows_us * elem_size); + + let mut matmul_dtypes = MatmulElems::from_single_dtype(dtype); + + // Widest supported vector size dividing a contiguous element count, plus + // the resulting lane count. Both operands of every `add_assign_kernel` + // launch are tight, so divisibility of the count is the only constraint. + let vec_split = |n: usize| -> (VectorSize, usize) { + let v = client + .io_optimized_vector_sizes(elem_size) + .filter(|v| n.is_multiple_of(*v)) + .max() + .unwrap_or(1); + (v, n / v) + }; + + let launch_matmul = |strategy: &Strategy, + lhs: InputBinding, + rhs: InputBinding, + out: TensorBinding, + dtypes: &mut MatmulElems| + -> Result<(), QRSetupError> { + cubek_matmul::launch::launch_ref(strategy, client, lhs, rhs, out, dtypes) + .map_err(|e| QRSetupError::Matmul(e.to_string())) + }; + + for k in 0..num_tiles { + let col_start = k * tile; + let current_tile = tile.min(cols - col_start); + let current_tile_us = current_tile as usize; + + let v_buf = TensorHandle::::new( + v_buf_global.clone(), + vec![rows_us, current_tile_us], + vec![1, rows_us], + dtype, + ); + let w_buf = TensorHandle::::new( + w_buf_global.clone(), + vec![rows_us, current_tile_us], + vec![1, rows_us], + dtype, + ); + + let n_clear = rows * current_tile; + let cd_clear = CubeDim::new_1d(max_cube_dim); + let cc_clear = calculate_cube_count_elemwise(client, n_clear as usize, cd_clear); + unsafe { + clear_buffer_kernel::launch_unchecked::( + client, + cc_clear.clone(), + cd_clear, + BufferArg::from_raw_parts(v_buf_global.clone(), n_clear as usize), + n_clear, + ); + } + + let cd_panel = CubeDim::new_1d(max_cube_dim); + let shared_size = max_cube_dim as usize; + + for j in 0..current_tile { + let col = col_start + j; + + unsafe { + householder_kernel::launch_unchecked::( + client, + CubeCount::new_1d(1), + cd_panel, + BufferArg::from_raw_parts(r_handle.handle.clone(), rows_us * cols_us), + rows, + col, + BufferArg::from_raw_parts(v_buf.handle.clone(), rows_us * current_tile_us), + BufferArg::from_raw_parts(beta_vec.clone(), tile_us), + j, + shared_size, + use_plane_reduce, + ); + } + + // One cube per remaining panel column, cube-wide reduction inside. + let n_upd = current_tile - j; + unsafe { + apply_householder_kernel::launch_unchecked::( + client, + CubeCount::new_1d(n_upd), + cd_panel, + rows, + col, + r_handle.clone().into_arg(), + BufferArg::from_raw_parts(v_buf.handle.clone(), rows_us * current_tile_us), + BufferArg::from_raw_parts(beta_vec.clone(), tile_us), + j, + shared_size, + use_plane_reduce, + ); + } + } + + let mut v_t_gram = InputBinding::Normal(v_buf.clone().binding(), storage_dtype); + v_t_gram.swap_dims(0, 1); + launch_matmul( + &strategy_gram, + v_t_gram, + InputBinding::Normal(v_buf.clone().binding(), storage_dtype), + TensorHandle::::new( + gram_buf_global.clone(), + vec![current_tile_us, current_tile_us], + vec![tile_us, 1], + dtype, + ) + .binding(), + &mut matmul_dtypes, + )?; + + unsafe { + build_t_tsqr_kernel::launch_unchecked::( + client, + CubeCount::new_1d(1), + CubeDim::new_1d(tile), + current_tile, + BufferArg::from_raw_parts(gram_buf_global.clone(), tile_us * tile_us), + BufferArg::from_raw_parts(beta_vec.clone(), tile_us), + BufferArg::from_raw_parts(t_buf_global.clone(), tile_us * tile_us), + tile, + ); + } + + let t_buf = TensorHandle::::new( + t_buf_global.clone(), + vec![current_tile_us, current_tile_us], + vec![1, tile_us], + dtype, + ); + launch_matmul( + &strategy_w, + InputBinding::Normal(v_buf.clone().binding(), storage_dtype), + InputBinding::Normal(t_buf.clone().binding(), storage_dtype), + w_buf.clone().binding(), + &mut matmul_dtypes, + )?; + + let has_trailing = col_start + current_tile < cols; + let trailing_cols = if has_trailing { + cols - (col_start + current_tile) + } else { + 0 + }; + let r_trail_offset = + (col_start + current_tile) as u64 * rows as u64 * core::mem::size_of::() as u64; + + if has_trailing { + let trailing_us = trailing_cols as usize; + let r_trailing = TensorHandle::::new( + r_handle.handle.clone().offset_start(r_trail_offset), + vec![rows_us, trailing_us], + vec![1, rows_us], + dtype, + ); + let s_r = TensorHandle::::new( + s_buf_global.clone(), + vec![current_tile_us, trailing_us], + vec![trailing_us, 1], + dtype, + ); + let z_r = TensorHandle::::new( + z_buf_global.clone(), + vec![rows_us, trailing_us], + vec![trailing_us, 1], + dtype, + ); + let mut w_t = InputBinding::Normal(w_buf.clone().binding(), storage_dtype); + w_t.swap_dims(0, 1); + launch_matmul( + &strategy_tall, + w_t, + InputBinding::Normal(r_trailing.clone().binding(), storage_dtype), + s_r.clone().binding(), + &mut matmul_dtypes, + )?; + launch_matmul( + &strategy_tall, + InputBinding::Normal(v_buf.clone().binding(), storage_dtype), + InputBinding::Normal(s_r.clone().binding(), storage_dtype), + z_r.clone().binding(), + &mut matmul_dtypes, + )?; + let cc_r = CubeCount::new_2d( + rows.div_ceil(thread_block_size), + trailing_cols.div_ceil(thread_block_size), + ); + unsafe { + update_trailing_r_kernel::launch_unchecked::( + client, + cc_r, + cube_dim_2d, + rows, + cols, + col_start + current_tile, + r_handle.clone().into_arg(), + BufferArg::from_raw_parts(z_buf_global.clone(), rows_us * trailing_us), + ); + } + } + + let s_tile_qt = TensorHandle::::new( + s_tile_global.clone(), + vec![current_tile_us, rows_us], + vec![rows_us, 1], + dtype, + ); + let mut w_t2 = InputBinding::Normal(w_buf.clone().binding(), storage_dtype); + w_t2.swap_dims(0, 1); + launch_matmul( + &strategy_tall, + w_t2, + InputBinding::Normal(q_handle.clone().binding(), storage_dtype), + s_tile_qt.clone().binding(), + &mut matmul_dtypes, + )?; + let z_qt = TensorHandle::::new( + z_buf_global.clone(), + vec![rows_us, rows_us], + vec![1, rows_us], + dtype, + ); + launch_matmul( + &strategy_tall, + InputBinding::Normal(v_buf.clone().binding(), storage_dtype), + InputBinding::Normal(s_tile_qt.clone().binding(), storage_dtype), + z_qt.clone().binding(), + &mut matmul_dtypes, + )?; + let n_q = rows_us * rows_us; + let (q_vec, n_q_vecs) = vec_split(n_q); + let cd_q = CubeDim::new_1d(max_cube_dim); + let cc_q = calculate_cube_count_elemwise(client, n_q_vecs, cd_q); + unsafe { + add_assign_kernel::launch_unchecked::( + client, + cc_q, + cd_q, + q_vec, + n_q_vecs as u32, + BufferArg::from_raw_parts(q_handle.handle.clone(), n_q_vecs), + BufferArg::from_raw_parts(z_buf_global.clone(), n_q_vecs), + ); + } + } + + Ok(()) +} diff --git a/crates/cubek-linalg/src/components/mod.rs b/crates/cubek-linalg/src/components/mod.rs new file mode 100644 index 000000000..37de0289b --- /dev/null +++ b/crates/cubek-linalg/src/components/mod.rs @@ -0,0 +1,2 @@ +pub mod baht_tsqr; +pub mod solve; diff --git a/crates/cubek-linalg/src/components/solve.rs b/crates/cubek-linalg/src/components/solve.rs new file mode 100644 index 000000000..5f5b1f386 --- /dev/null +++ b/crates/cubek-linalg/src/components/solve.rs @@ -0,0 +1,54 @@ +//! Kernels used to solve `Ax = b` from a QR decomposition: +//! `y = Q^T·b` followed by back substitution on `Rx = y`. + +use cubecl::prelude::*; + +/// Kernel to compute y = Q_already_t * b +/// Q_already_t is m x m stored col-major as produced by the QR kernels +/// (`Q^T[i, j]` at flat index `j*m + i`). b is m x 1. +/// y_i = sum_j (Q_already_t)_ij * b_j +#[cube(launch_unchecked)] +pub fn q_already_t_b_kernel(m: u32, q_t: &Tensor, b: &Tensor, y: &mut Tensor) { + let i = ABSOLUTE_POS_X; + if i < m { + let mut sum = 0.0f64; + for j in 0..m { + sum = fma( + f64::cast_from(q_t[j as usize * m as usize + i as usize]), + f64::cast_from(b[j as usize]), + sum, + ); + } + y[i as usize] = F::cast_from(sum); + } +} + +/// Back substitution for Rx = y where R is upper triangular, stored tight +/// col-major `[rows, n]` as produced by the QR kernels. +/// Accumulates in f64 to minimise rounding in the triangular solve. +#[cube(launch_unchecked)] +pub fn back_substitution_kernel( + n: u32, + rows: u32, + r: &Tensor, + y: &Tensor, + x: &mut Tensor, +) { + if ABSOLUTE_POS == 0 { + let mut i = n; + while i > 0 { + i -= 1; + let mut sum = 0.0f64; + for j in (i + 1)..n { + sum = fma( + f64::cast_from(r[j as usize * rows as usize + i as usize]), + f64::cast_from(x[j as usize]), + sum, + ); + } + let diag_idx = i as usize * rows as usize + i as usize; + x[i as usize] = + F::cast_from((f64::cast_from(y[i as usize]) - sum) / f64::cast_from(r[diag_idx])); + } + } +} diff --git a/crates/cubek-linalg/src/definition/error.rs b/crates/cubek-linalg/src/definition/error.rs new file mode 100644 index 000000000..7f945ae27 --- /dev/null +++ b/crates/cubek-linalg/src/definition/error.rs @@ -0,0 +1,34 @@ +use cubecl::ir::StorageType; +use thiserror::Error; + +/// Errors that can occur when trying to launch QR decomposition. +#[derive(Debug, Error, PartialEq, Eq, Clone, Hash)] +pub enum QRSetupError { + /// The input should be a non-empty matrix where m should be greater or equal to n. + #[error("The input should be a non-empty matrix where m should be greater or equal to n.")] + InvalidShape, + /// The element type the kernels were launched with does not match the tensor's. + #[error("Element type mismatch: launched as {launched:?} but the tensor holds {actual:?}.")] + TypeMismatch { + /// Element type the generic entry point was instantiated with. + launched: StorageType, + /// Element type of the tensor handle. + actual: StorageType, + }, + /// An internal matmul refused the requested launch. + #[error("Internal matmul setup failed: {0}")] + Matmul(String), + /// The routine requires more shared memory than the hardware provides. + #[error( + "The routine requires {requested} bytes of shared memory but only {available} are available." + )] + SharedMemoryLimitExceeded { + /// Number of shared memory bytes the routine asked for. + requested: usize, + /// Number of shared memory bytes the hardware provides. + available: usize, + }, + /// A forced blueprint is inconsistent with the problem or the hardware. + #[error("Invalid blueprint: {0}")] + InvalidBlueprint(String), +} diff --git a/crates/cubek-linalg/src/definition/mod.rs b/crates/cubek-linalg/src/definition/mod.rs new file mode 100644 index 000000000..3e670c92a --- /dev/null +++ b/crates/cubek-linalg/src/definition/mod.rs @@ -0,0 +1,5 @@ +mod error; +mod problem; + +pub use error::*; +pub use problem::*; diff --git a/crates/cubek-linalg/src/definition/problem.rs b/crates/cubek-linalg/src/definition/problem.rs new file mode 100644 index 000000000..dc5a1f761 --- /dev/null +++ b/crates/cubek-linalg/src/definition/problem.rs @@ -0,0 +1,32 @@ +use cubecl::ir::StorageType; + +use crate::definition::QRSetupError; + +/// Runtime description of a QR factorization problem: an `m x n` matrix with `m >= n`. +/// +/// Only runtime information belongs here (shapes, element type); anything that +/// changes the generated kernel code goes in a routine blueprint instead. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct QRProblem { + /// Number of rows (`m`) of the input matrix. + pub rows: usize, + /// Number of columns (`n`) of the input matrix. + pub cols: usize, + /// Element type of the input matrix. + pub dtype: StorageType, +} + +impl QRProblem { + /// Build a problem description from a tensor shape, validating that it + /// describes a supported (`m >= n`) matrix. + pub fn from_shape(shape: &[usize], dtype: StorageType) -> Result { + if shape.len() != 2 || shape[0] < shape[1] || shape[1] == 0 { + return Err(QRSetupError::InvalidShape); + } + Ok(Self { + rows: shape[0], + cols: shape[1], + dtype, + }) + } +} diff --git a/crates/cubek-linalg/src/eval/benchmarks/benchmark.rs b/crates/cubek-linalg/src/eval/benchmarks/benchmark.rs new file mode 100644 index 000000000..20a505153 --- /dev/null +++ b/crates/cubek-linalg/src/eval/benchmarks/benchmark.rs @@ -0,0 +1,88 @@ +use std::marker::PhantomData; + +use cubecl::{ + Runtime, TestRuntime, + benchmark::{Benchmark, TimingMethod}, + client::ComputeClient, + future, + prelude::*, + std::tensor::TensorHandle, + zspace::Shape, +}; +use cubek_test_utils::{RunSamples, StridedLayout, TestInput}; + +use crate::eval::benchmarks::problem::QrProblem; +use crate::eval::benchmarks::strategy::QrStrategy; + +pub fn bench( + _strategy: &QrStrategy, + problem: &QrProblem, + num_samples: usize, +) -> Result { + let device = ::Device::default(); + let client = ::client(&device); + + let bench = QrBench:: { + m: problem.m, + n: problem.n, + client, + samples: num_samples, + _e: PhantomData, + }; + + let durations = bench + .run(TimingMethod::System) + .map_err(|e| format!("benchmark failed: {e}"))? + .durations; + + Ok(RunSamples::new(durations)) +} + +struct QrBench { + m: usize, + n: usize, + client: ComputeClient, + samples: usize, + _e: PhantomData, +} + +impl Benchmark for QrBench { + type Input = TensorHandle; + type Output = (); + + fn prepare(&self) -> Self::Input { + let storage = E::as_type_native_unchecked().storage_type(); + + // Col-major input: the layout the QR kernels work in, so the timed + // run measures the factorization rather than a layout conversion. + TestInput::builder(self.client.clone(), Shape::from(vec![self.m, self.n])) + .layout(StridedLayout::ColMajor) + .dtype(storage) + .uniform(0, -1., 1.) + .generate_without_host_data() + } + + fn execute(&self, a: Self::Input) -> Result<(), String> { + crate::launch::qr::(&self.client, &a) + .map(|_| ()) + .map_err(|err| format!("{err:?}")) + } + + fn num_samples(&self) -> usize { + self.samples + } + + fn name(&self) -> String { + format!( + "qr-{}-baht_tsqr-{}x{}", + E::as_type_native_unchecked(), + self.m, + self.n, + ) + .to_lowercase() + } + + fn sync(&self) { + future::block_on(self.client.sync()).unwrap() + } +} diff --git a/crates/cubek-linalg/src/eval/benchmarks/mod.rs b/crates/cubek-linalg/src/eval/benchmarks/mod.rs new file mode 100644 index 000000000..70d0c5217 --- /dev/null +++ b/crates/cubek-linalg/src/eval/benchmarks/mod.rs @@ -0,0 +1,43 @@ +//! Benchmark catalogue for `cubek-linalg`. + +mod benchmark; +mod problem; +mod strategy; + +pub use benchmark::bench; +pub use problem::{QrProblem, problems}; +pub use strategy::{QrStrategy, strategies}; + +use cubek_test_utils::{CatalogEntry, RunSamples}; + +pub struct Category; + +impl cubek_test_utils::Category for Category { + type Problem = QrProblem; + type Strategy = QrStrategy; + + fn id(&self) -> &'static str { + "qr" + } + + fn label(&self) -> &'static str { + "QR decomposition" + } + + fn problems(&self) -> Vec> { + problems() + } + + fn strategies(&self) -> Vec> { + strategies() + } + + fn bench( + &self, + strategy: &QrStrategy, + problem: &QrProblem, + num_samples: usize, + ) -> Result { + bench(strategy, problem, num_samples) + } +} diff --git a/crates/cubek-linalg/src/eval/benchmarks/problem.rs b/crates/cubek-linalg/src/eval/benchmarks/problem.rs new file mode 100644 index 000000000..8e0248fe3 --- /dev/null +++ b/crates/cubek-linalg/src/eval/benchmarks/problem.rs @@ -0,0 +1,22 @@ +use cubek_test_utils::CatalogEntry; + +pub struct QrProblem { + /// Number of rows (`m`) of the input matrix; must be >= `n`. + pub m: usize, + /// Number of columns (`n`) of the input matrix. + pub n: usize, +} + +pub fn problems() -> Vec> { + vec![ + CatalogEntry::new("128x128", "128x128", QrProblem { m: 128, n: 128 }), + CatalogEntry::new("512x512", "512x512", QrProblem { m: 512, n: 512 }), + CatalogEntry::new("1024x1024", "1024x1024", QrProblem { m: 1024, n: 1024 }), + // Tall-skinny case: the shape TSQR-style panel factorizations target. + CatalogEntry::new( + "2048x512", + "2048x512 (tall-skinny)", + QrProblem { m: 2048, n: 512 }, + ), + ] +} diff --git a/crates/cubek-linalg/src/eval/benchmarks/strategy.rs b/crates/cubek-linalg/src/eval/benchmarks/strategy.rs new file mode 100644 index 000000000..a41588d10 --- /dev/null +++ b/crates/cubek-linalg/src/eval/benchmarks/strategy.rs @@ -0,0 +1,8 @@ +use cubek_test_utils::CatalogEntry; + +/// Marker type: `cubek-linalg` only implements the `baht_tsqr` QR strategy. +pub struct QrStrategy; + +pub fn strategies() -> Vec> { + vec![CatalogEntry::new("baht_tsqr", "BahtTsqr", QrStrategy)] +} diff --git a/crates/cubek-linalg/src/eval/mod.rs b/crates/cubek-linalg/src/eval/mod.rs new file mode 100644 index 000000000..d4b49f112 --- /dev/null +++ b/crates/cubek-linalg/src/eval/mod.rs @@ -0,0 +1,2 @@ +#[cfg(feature = "benchmarks")] +pub mod benchmarks; diff --git a/crates/cubek-linalg/src/launch/base.rs b/crates/cubek-linalg/src/launch/base.rs new file mode 100644 index 000000000..4bc6f7ec1 --- /dev/null +++ b/crates/cubek-linalg/src/launch/base.rs @@ -0,0 +1,99 @@ +use cubecl::prelude::*; +use cubecl::std::tensor::{TensorHandle, identity, into_contiguous}; + +use crate::{ + components, + definition::{QRProblem, QRSetupError}, + routines::{BahtTsqrRoutine, BahtTsqrStrategy, BlueprintStrategy, QRRoutine}, +}; + +/// The `(Q, R)` pair produced by a QR decomposition. +pub type QRTuple = (TensorHandle, TensorHandle); + +/// Allocate and seed the Q (identity) and R (copy of A) buffers the QR +/// kernels work in-place on. +fn initialize( + client: &ComputeClient, + a: &TensorHandle, + problem: &QRProblem, +) -> QRTuple { + let m = problem.rows; + let n = problem.cols; + let dtype = problem.dtype; + + // Allocate Q as identity (col-major [rows, rows]). + // IMPORTANT: TensorHandle::zeros / ::empty may return a GPU-padded buffer + // (e.g. pitch 4 for a 3-row matrix). The BAHT kernels index Q^T with flat + // `col*rows+row` arithmetic that assumes NO pitch padding. We therefore + // reserve a tight (non-padded) buffer via `client.empty` and fill it on the + // GPU with cubecl's `identity` kernel — no host round-trip, any float type. + // + // The identity kernel derives the diagonal stride from `strides[0]`, so we + // hand it standard row-major strides `[m, 1]`. The identity matrix is + // symmetric, so the resulting tight buffer is simultaneously the col-major + // `[1, m]` identity the QR kernels consume. + let q_shape = vec![m, m]; + let elem_size = dtype.size(); + let q_handle = client.empty(m * m * elem_size); + let q_contig = TensorHandle::::new(q_handle.clone(), q_shape.clone(), vec![m, 1], dtype); + identity::launch::(client, &q_contig); + let q = TensorHandle::::new(q_handle, q_shape, vec![1, m], dtype); + + // Build R as a tight col-major copy of A, entirely on-device and accepting + // any input layout (row-major, col-major, pitch-padded): a tight col-major + // `[m, n]` buffer is byte-identical to a tight row-major `[n, m]` buffer of + // the transpose, so run `into_contiguous` (tight `client.empty` output, + // reads the source through its strides) on a transposed *view* of A and + // re-declare the result with col-major strides `[1, m]`. + let a_strides = a.strides(); + let a_t = TensorHandle::::new( + a.handle.clone(), + vec![n, m], + vec![a_strides[1], a_strides[0]], + dtype, + ); + let r_t = into_contiguous::(client, a_t.binding(), dtype); + let r = TensorHandle::::new(r_t.handle, vec![m, n], vec![1, m], dtype); + + (q, r) +} + +/// It launches a QR decomposition over a m x n matrix a using the TSQR-inspired +/// blocked Householder routine ([`BahtTsqrRoutine`]). +/// +/// Specify the client and the matrix a to decompose. In case of success it +/// will return a tuple with the matrix Q and the matrix R in this order. +pub fn qr( + client: &ComputeClient, + a: &TensorHandle, +) -> Result, QRSetupError> { + qr_with_strategy::(client, a, Default::default()) +} + +/// [`qr`] with explicit routine knobs. +/// +/// The only knob today is [`BahtTsqrStrategy::allow_tf32`], which lets the +/// trailing-update GEMMs run on tensor cores: about 2x faster end to end, at a +/// reconstruction error near `1e-2` instead of f32 accuracy. [`qr`] leaves it +/// off, so its accuracy is unchanged. +pub fn qr_with_strategy( + client: &ComputeClient, + a: &TensorHandle, + strategy: BahtTsqrStrategy, +) -> Result, QRSetupError> { + let problem = QRProblem::from_shape(a.shape(), a.dtype)?; + let launched = EG::as_type_native_unchecked().storage_type(); + if problem.dtype != launched { + return Err(QRSetupError::TypeMismatch { + launched, + actual: problem.dtype, + }); + } + let (q, r) = initialize::(client, a, &problem); + + let (blueprint, settings) = + BahtTsqrRoutine::prepare(client, &problem, BlueprintStrategy::Inferred(strategy))?; + components::baht_tsqr::launch::(client, &q, &r, blueprint, settings)?; + + Ok((q, r)) +} diff --git a/crates/cubek-linalg/src/launch/mod.rs b/crates/cubek-linalg/src/launch/mod.rs new file mode 100644 index 000000000..dd6e1f100 --- /dev/null +++ b/crates/cubek-linalg/src/launch/mod.rs @@ -0,0 +1,5 @@ +mod base; +mod solve; + +pub use base::*; +pub use solve::*; diff --git a/crates/cubek-linalg/src/launch/solve.rs b/crates/cubek-linalg/src/launch/solve.rs new file mode 100644 index 000000000..8536af697 --- /dev/null +++ b/crates/cubek-linalg/src/launch/solve.rs @@ -0,0 +1,80 @@ +use cubecl::calculate_cube_count_elemwise; +use cubecl::prelude::*; +use cubecl::std::tensor::TensorHandle; + +use crate::{ + components::solve::{back_substitution_kernel, q_already_t_b_kernel}, + definition::QRSetupError, +}; + +/// Solve Ax = b using QR decomposition. +pub fn solve( + client: &ComputeClient, + a: &TensorHandle, + b: &TensorHandle, +) -> Result, QRSetupError> { + let shape_a = a.shape(); + let shape_b = b.shape(); + + if shape_a.len() != 2 || shape_b.len() != 1 || shape_a[0] != shape_b[0] { + return Err(QRSetupError::InvalidShape); + } + + let m = shape_a[0]; + let n = shape_a[1]; + + if m < n { + return Err(QRSetupError::InvalidShape); + } + + // `qr` checks `a` against `E`; `b` must match as well since the solve + // kernels read it as `E` too. + let launched = E::as_type_native_unchecked().storage_type(); + if b.dtype != launched { + return Err(QRSetupError::TypeMismatch { + launched, + actual: b.dtype, + }); + } + + // 1. A = QR + let (q, r) = crate::launch::qr::(client, a)?; + + // 2. y = Q^T * b + let y = TensorHandle::zeros(client, vec![m], a.dtype); + let max_cube_dim = client.properties().hardware.max_cube_dim.0; + let cd_q = CubeDim::new_1d(max_cube_dim.min(m as u32)); + let cc_q = calculate_cube_count_elemwise(client, m, cd_q); + + // All strategies return Q^T + unsafe { + q_already_t_b_kernel::launch_unchecked::( + client, + cc_q, + cd_q, + m as u32, + q.clone().into_arg(), + b.clone().into_arg(), + y.clone().into_arg(), + ); + } + + // 3. Rx = y (first n elements of y if m > n) + let x = TensorHandle::zeros(client, vec![n], a.dtype); + + // For back substitution, we use a single cube since it's sequential + unsafe { + back_substitution_kernel::launch_unchecked::( + client, + CubeCount::new_1d(1), + CubeDim::new_1d(1), + n as u32, + m as u32, + r.clone().into_arg(), + y.clone().into_arg(), + x.clone().into_arg(), + ); + } + + Ok(x) +} diff --git a/crates/cubek-linalg/src/lib.rs b/crates/cubek-linalg/src/lib.rs new file mode 100644 index 000000000..40c9bf782 --- /dev/null +++ b/crates/cubek-linalg/src/lib.rs @@ -0,0 +1,21 @@ +//! Linear algebra kernels for cubek. +//! +//! QR decomposition is implemented with the TSQR-inspired blocked Householder +//! routine ([`routines::BahtTsqrRoutine`]), the fastest of the strategies +//! benchmarked for this crate. The crate follows the blueprint-routine +//! architecture: +//! - [`definition`] holds the runtime problem descriptions and errors. +//! - [`routines`] adapts the algorithm to the hardware, producing a minimal +//! comptime `Blueprint` plus runtime launch settings. +//! - [`components`] holds the kernels, specialized by their blueprint. +//! - [`launch`](mod@launch) validates the input and dispatches to the +//! matching routine and component. + +pub mod components; +pub mod definition; +pub mod eval; +pub mod launch; +pub mod routines; + +pub use definition::*; +pub use launch::*; diff --git a/crates/cubek-linalg/src/routines/baht_tsqr.rs b/crates/cubek-linalg/src/routines/baht_tsqr.rs new file mode 100644 index 000000000..408057714 --- /dev/null +++ b/crates/cubek-linalg/src/routines/baht_tsqr.rs @@ -0,0 +1,152 @@ +use cubecl::features::Plane; +use cubecl::ir::{ElemType, FloatKind}; +use cubecl::prelude::*; +use cubecl::tf32; +use cubek_matmul::strategy::Strategy; + +use crate::{ + definition::{QRProblem, QRSetupError}, + routines::{BlueprintStrategy, QRRoutine}, +}; + +/// Whether the element type is f64, for which the specialized unit matmul +/// routines are not supported; fall back to auto-selection in that case. +fn is_f64(problem: &QRProblem) -> bool { + problem.dtype == StorageType::Scalar(ElemType::Float(FloatKind::F64)) +} + +/// TSQR-inspired blocked Householder QR: the whole panel is factorized with +/// minimal dispatches, then the block reflector is applied through GEMMs. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct BahtTsqrRoutine; + +/// Tunable knobs for [`BahtTsqrRoutine`]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct BahtTsqrStrategy { + /// Trade accuracy for speed on the trailing-update GEMMs. See + /// [`BahtTsqrBlueprint::allow_tf32`]. Off by default. + pub allow_tf32: bool, +} + +/// Comptime specialization settings for the TSQR kernels. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub struct BahtTsqrBlueprint { + /// Fold the panel reductions with the hardware `plane_sum` (one shared + /// slot per plane, two barriers) instead of the generic shared-memory + /// tree. Requires plane-op support and a uniform plane size. + pub use_plane_reduce: bool, + /// Let the trailing-update GEMMs pick a tensor-core matmul, whose f32 + /// path rounds its inputs to tf32 (~10 mantissa bits instead of 24). + /// + /// Roughly **2x faster end to end** — those GEMMs are essentially the + /// entire matmul cost — at a reconstruction error around `1e-2` rather + /// than f32 accuracy. Off by default; enable it only when the caller's + /// tolerance permits. No effect for f64, which has no tensor-core path. + pub allow_tf32: bool, +} + +/// Runtime launch parameters for the TSQR kernels and GEMM updates. +#[derive(Clone)] +pub struct BahtTsqrLaunchSettings { + /// Number of reflectors batched per panel. + pub tile: u32, + /// Cube dim used for the 1D panel kernels. + pub max_cube_dim: u32, + /// Side of the square cube dim used for the 2D update kernels. + pub thread_block_size: u32, + pub cube_dim_2d: CubeDim, + /// Matmul strategy for the V^T·V Gram matrix. + pub strategy_gram: Strategy, + /// Matmul strategy for W = V·T. + pub strategy_w: Strategy, + /// Matmul strategy for the tall trailing-R and Q^T updates. + pub strategy_tall: Strategy, +} + +impl QRRoutine for BahtTsqrRoutine { + type Strategy = BahtTsqrStrategy; + type Blueprint = BahtTsqrBlueprint; + type LaunchSettings = BahtTsqrLaunchSettings; + + fn prepare( + client: &ComputeClient, + problem: &QRProblem, + strategy: BlueprintStrategy, + ) -> Result<(Self::Blueprint, Self::LaunchSettings), QRSetupError> { + let hardware = &client.properties().hardware; + + let plane_reduce_supported = client.properties().features.plane.contains(Plane::Ops) + && hardware.plane_size_min == hardware.plane_size_max + && hardware.plane_size_min > 0; + + let tf32_supported = client + .properties() + .supports_type(tf32::as_type_native_unchecked().storage_type()); + + let blueprint = match strategy { + BlueprintStrategy::Forced(blueprint) => { + if blueprint.use_plane_reduce && !plane_reduce_supported { + return Err(QRSetupError::InvalidBlueprint( + "use_plane_reduce requires plane-op support and a uniform plane size" + .to_string(), + )); + } + blueprint + } + BlueprintStrategy::Inferred(strategy) => BahtTsqrBlueprint { + use_plane_reduce: plane_reduce_supported, + // Honour the opt-in only where tf32 actually exists. Without a + // tensor-core path `Strategy::Auto` resolves to something + // *slower* than the full-precision unit routine (measured on + // wgpu: 24.1ms -> 30.5ms at 1024x1024), so taking the accuracy + // hit there would buy nothing. + allow_tf32: strategy.allow_tf32 && tf32_supported, + }, + }; + let thread_block_size = (hardware.max_cube_dim.0 as f64).sqrt() as u32; + let max_cube_dim = hardware.max_cube_dim.0.min(256); + let tile = 32u32.min(problem.cols as u32).min(max_cube_dim); + + // `Strategy::Auto` resolves to a tensor-core (CMMA) matmul whose f32 + // path silently downgrades the stage/register types to tf32 (see + // `cubek_matmul::definition::adjust_dtypes`), losing ~13 mantissa bits + // and pushing the QR reconstruction error to ~1e-2 — over the test + // tolerance — so f32 defaults to full-precision strategies. + // + // The unit routines don't support f64, so f64 uses Auto — safe there + // because no f64 CMMA exists and Auto falls back to the + // full-precision SimpleUnit. + // + // Only `strategy_tall` is worth trading accuracy for: measured at + // 1024x1024, forcing it to Auto takes the whole decomposition from + // 23.5ms to 12.0ms, while forcing gram or W to Auto changes nothing + // end to end. So the tf32 opt-in touches the trailing updates alone + // and gram/W stay full-precision unconditionally. + let (strategy_gram, strategy_w, strategy_tall) = if is_f64(problem) { + (Strategy::Auto, Strategy::Auto, Strategy::Auto) + } else { + let tall = if blueprint.allow_tf32 { + Strategy::Auto + } else { + Strategy::DoubleUnit(Default::default()) + }; + ( + Strategy::SimpleVecMat(Default::default()), + Strategy::DoubleUnit(Default::default()), + tall, + ) + }; + + let settings = BahtTsqrLaunchSettings { + tile, + max_cube_dim, + thread_block_size, + cube_dim_2d: CubeDim::new_2d(thread_block_size, thread_block_size), + strategy_gram, + strategy_w, + strategy_tall, + }; + + Ok((blueprint, settings)) + } +} diff --git a/crates/cubek-linalg/src/routines/base.rs b/crates/cubek-linalg/src/routines/base.rs new file mode 100644 index 000000000..ccdbec30c --- /dev/null +++ b/crates/cubek-linalg/src/routines/base.rs @@ -0,0 +1,41 @@ +use core::fmt::Debug; + +use cubecl::prelude::*; + +use crate::definition::{QRProblem, QRSetupError}; + +/// How the blueprint for a routine is obtained: either fully specified by the +/// caller (`Forced`) or derived from the problem and the hardware (`Inferred`). +#[derive(Debug, Clone)] +pub enum BlueprintStrategy { + /// Use this exact blueprint, only deriving the launch settings from it. + Forced(R::Blueprint), + /// Derive the blueprint from the problem, the hardware and the given + /// routine strategy. + Inferred(R::Strategy), +} + +/// A routine adapts a QR algorithm to the hardware it runs on. +/// +/// It does not make hard decisions about hardware specifics; instead it reads +/// the device properties and the problem description, and produces: +/// - a [`QRRoutine::Blueprint`]: the minimal comptime settings baked into the +/// generated kernels (a different blueprint triggers a new JIT compilation), +/// - a [`QRRoutine::LaunchSettings`]: the runtime parameters (cube dimensions, +/// cube counts, ...) used to launch those kernels. +pub trait QRRoutine: Debug + Clone + Sized { + /// Tunable knobs for this routine that do not depend on the hardware. + type Strategy: Debug + Clone + Default + Send + 'static; + /// Minimal comptime specialization settings for the kernels. + type Blueprint: Debug + Clone + Send + 'static; + /// Runtime launch parameters derived from the problem and the hardware. + type LaunchSettings: Clone; + + /// Adapt the algorithm to the problem and the hardware, producing the + /// blueprint for the compiler and the launch settings for the runtime. + fn prepare( + client: &ComputeClient, + problem: &QRProblem, + strategy: BlueprintStrategy, + ) -> Result<(Self::Blueprint, Self::LaunchSettings), QRSetupError>; +} diff --git a/crates/cubek-linalg/src/routines/mod.rs b/crates/cubek-linalg/src/routines/mod.rs new file mode 100644 index 000000000..3401b06c0 --- /dev/null +++ b/crates/cubek-linalg/src/routines/mod.rs @@ -0,0 +1,5 @@ +mod baht_tsqr; +mod base; + +pub use baht_tsqr::*; +pub use base::*; diff --git a/crates/cubek-linalg/tests/lib.rs b/crates/cubek-linalg/tests/lib.rs new file mode 100644 index 000000000..7a4d0e3b2 --- /dev/null +++ b/crates/cubek-linalg/tests/lib.rs @@ -0,0 +1 @@ +mod suite; diff --git a/crates/cubek-linalg/tests/suite/mod.rs b/crates/cubek-linalg/tests/suite/mod.rs new file mode 100644 index 000000000..c8253b1dc --- /dev/null +++ b/crates/cubek-linalg/tests/suite/mod.rs @@ -0,0 +1,2 @@ +pub mod qr; +pub mod utils; diff --git a/crates/cubek-linalg/tests/suite/qr/baht_tsqr.rs b/crates/cubek-linalg/tests/suite/qr/baht_tsqr.rs new file mode 100644 index 000000000..08c6f1e3d --- /dev/null +++ b/crates/cubek-linalg/tests/suite/qr/baht_tsqr.rs @@ -0,0 +1,194 @@ +use cubecl::{ + TestRuntime, + prelude::*, + std::tensor::{TensorHandle, into_contiguous}, +}; + +use crate::suite::utils::{ + assert_equals_approx, col_major_input, dtype_unsupported, row_major_input, +}; + +/// Run into_contiguous on a tensor and read the resulting tight row-major bytes. +fn read_contig( + client: &ComputeClient, + t: &TensorHandle, +) -> (Vec, Vec) { + let shape = t.shape().to_vec(); + let contig = into_contiguous::(client, t.clone().binding(), t.dtype); + let bytes = client.read_one(contig.handle.clone()).unwrap(); + (F::from_bytes(&bytes).to_vec(), shape) +} + +/// Reconstruct A = Q * R in f64 for maximum verification accuracy. +/// Q^T is row-major (after into_contiguous): Q[i,k] = q_t[k * rows + i]. +/// R is row-major (after into_contiguous): R[k,j] = r[k * cols + j]. +fn reconstruct_qr( + q_t_vals: &[F], + r_vals: &[F], + rows: usize, + cols: usize, + k_range: usize, + q_row_stride: usize, +) -> Vec { + let mut out = vec![0.0f64; rows * cols]; + for i in 0..rows { + for j in 0..cols { + let mut sum = 0.0f64; + for k in 0..k_range { + let q_ik = q_t_vals[k * q_row_stride + i].to_f64().unwrap(); + let r_kj = r_vals[k * cols + j].to_f64().unwrap(); + sum += q_ik * r_kj; + } + out[i * cols + j] = sum; + } + } + out.iter() + .map(|&v| ::from(v).unwrap()) + .collect() +} + +fn run_qr_square(dim: u32) { + let client = TestRuntime::client(&Default::default()); + if dtype_unsupported::(&client) { + return; + } + let dim_usize = dim as usize; + + let shape = vec![dim_usize, dim_usize]; + let num_elements = shape.iter().product(); + // Ones with 2s on the anti-diagonal (a symmetric matrix, logical + // row-major). + let mut data = vec![F::from_int(1); num_elements]; + let mut pos = dim_usize - 1; + for _i in 0..dim { + data[pos] = F::from_int(2); + pos += dim_usize - 1; + } + + let a = col_major_input(&client, shape.clone(), &data); + + let (q_t, r) = match cubek_linalg::qr::(&client, &a) { + Ok((q_t, r)) => (q_t, r), + Err(e) => panic!("QR launch failed: {e:?}"), + }; + + let (q_t_vals, _) = read_contig::(&client, &q_t); + let (r_vals_out, _) = read_contig::(&client, &r); + + // Reconstruct in f64 for accurate verification. + let out_data = reconstruct_qr( + &q_t_vals, + &r_vals_out, + dim_usize, + dim_usize, + dim_usize, + dim_usize, + ); + + assert_equals_approx::(&out_data, &data, shape, 2e-3); +} + +fn run_qr_rect(rows: u32, cols: u32, row_major: bool) { + let client = TestRuntime::client(&Default::default()); + if dtype_unsupported::(&client) { + return; + } + let rows_usize = rows as usize; + let cols_usize = cols as usize; + + let shape = vec![rows_usize, cols_usize]; + let num_elements = rows_usize * cols_usize; + + let mut row_major_data = vec![F::from_int(1); num_elements]; + for i in 0..rows_usize.min(cols_usize) { + row_major_data[i * cols_usize + i] = F::from_int(2); + } + + // `qr` normalizes any input layout to its internal col-major form, so + // both layouts must factorize the same logical matrix identically. + let a = if row_major { + row_major_input(&client, shape.clone(), &row_major_data) + } else { + col_major_input(&client, shape.clone(), &row_major_data) + }; + + let (q_t, r) = match cubek_linalg::qr::(&client, &a) { + Ok((q_t, r)) => (q_t, r), + Err(e) => panic!("QR launch failed: {e:?}"), + }; + + // Q^T row-major [rows × rows], R row-major [rows × cols]. + // q_t_vals[k * rows + i] = Q^T[k, i] = Q[i, k]. + let (q_t_vals, qt_shape) = read_contig::(&client, &q_t); + let (r_vals_out, _) = read_contig::(&client, &r); + + let out_data = reconstruct_qr( + &q_t_vals, + &r_vals_out, + rows_usize, + cols_usize, + qt_shape[0], + qt_shape[1], + ); + + assert_equals_approx::(&out_data, &row_major_data, shape, 2e-3); +} + +/// Same square reconstruction as [`run_qr_square`], but through +/// `qr_with_strategy` with `allow_tf32`. The trailing-update GEMMs then run on +/// tensor cores in tf32, so this asserts against the looser tolerance that +/// documents — it exists to pin the opt-in path down, not to prove accuracy. +fn run_qr_square_tf32(dim: u32) { + let client = TestRuntime::client(&Default::default()); + if dtype_unsupported::(&client) { + return; + } + let dim_usize = dim as usize; + + let shape = vec![dim_usize, dim_usize]; + let num_elements = shape.iter().product(); + let mut data = vec![F::from_int(1); num_elements]; + let mut pos = dim_usize - 1; + for _i in 0..dim { + data[pos] = F::from_int(2); + pos += dim_usize - 1; + } + + let a = col_major_input(&client, shape.clone(), &data); + + let strategy = cubek_linalg::routines::BahtTsqrStrategy { allow_tf32: true }; + let (q_t, r) = match cubek_linalg::qr_with_strategy::(&client, &a, strategy) { + Ok((q_t, r)) => (q_t, r), + Err(e) => panic!("QR launch failed: {e:?}"), + }; + + let (q_t_vals, _) = read_contig::(&client, &q_t); + let (r_vals_out, _) = read_contig::(&client, &r); + + let out_data = reconstruct_qr( + &q_t_vals, + &r_vals_out, + dim_usize, + dim_usize, + dim_usize, + dim_usize, + ); + + assert_equals_approx::(&out_data, &data, shape, 5e-2); +} + +pub fn test_qr(dim: u32) { + run_qr_square::(dim); +} + +pub fn test_qr_tf32(dim: u32) { + run_qr_square_tf32::(dim); +} + +pub fn test_qr_rect(rows: u32, cols: u32) { + run_qr_rect::(rows, cols, false); +} + +pub fn test_qr_rect_row_major(rows: u32, cols: u32) { + run_qr_rect::(rows, cols, true); +} diff --git a/crates/cubek-linalg/tests/suite/qr/errors.rs b/crates/cubek-linalg/tests/suite/qr/errors.rs new file mode 100644 index 000000000..6931577ec --- /dev/null +++ b/crates/cubek-linalg/tests/suite/qr/errors.rs @@ -0,0 +1,42 @@ +//! Validation tests for the launch-layer guards: shapes and element types +//! must be rejected with a `QRSetupError` before any kernel is dispatched. + +use cubecl::{TestRuntime, prelude::*}; +use cubek_linalg::QRSetupError; + +use crate::suite::utils::col_major_input; + +#[test] +fn qr_rejects_zero_columns() { + let client = TestRuntime::client(&Default::default()); + let a = col_major_input::(&client, vec![5, 0], &[]); + let result = cubek_linalg::qr::(&client, &a); + assert_eq!(result.err(), Some(QRSetupError::InvalidShape)); +} + +#[test] +fn qr_rejects_wide_matrix() { + let client = TestRuntime::client(&Default::default()); + let a = col_major_input::(&client, vec![2, 4], &[1.0f32; 8]); + let result = cubek_linalg::qr::(&client, &a); + assert_eq!(result.err(), Some(QRSetupError::InvalidShape)); +} + +#[test] +fn qr_rejects_dtype_mismatch() { + let client = TestRuntime::client(&Default::default()); + // f64 tensor launched through the f32 entry point: must error out + // before any buffer is sized or kernel dispatched. + let a = col_major_input::(&client, vec![4, 2], &[1.0f64; 8]); + let result = cubek_linalg::qr::(&client, &a); + assert!(matches!(result, Err(QRSetupError::TypeMismatch { .. }))); +} + +#[test] +fn solve_rejects_dtype_mismatch_b() { + let client = TestRuntime::client(&Default::default()); + let a = col_major_input::(&client, vec![2, 2], &[2.0, 0.0, 0.0, 2.0f32]); + let b = col_major_input::(&client, vec![2], &[1.0f64; 2]); + let result = cubek_linalg::solve::(&client, &a, &b); + assert!(matches!(result, Err(QRSetupError::TypeMismatch { .. }))); +} diff --git a/crates/cubek-linalg/tests/suite/qr/mod.rs b/crates/cubek-linalg/tests/suite/qr/mod.rs new file mode 100644 index 000000000..8cf64ec08 --- /dev/null +++ b/crates/cubek-linalg/tests/suite/qr/mod.rs @@ -0,0 +1,89 @@ +pub mod baht_tsqr; +pub mod errors; +pub mod solve; + +#[macro_export] +macro_rules! testgen_qr_baht_tsqr { + ($float:ident) => { + pub type FloatT = $float; + + #[test] + pub fn test_tiny() { + crate::suite::qr::baht_tsqr::test_qr::(3); + } + + #[test] + pub fn test_small() { + crate::suite::qr::baht_tsqr::test_qr::(47); + } + + #[test] + pub fn test_medium() { + crate::suite::qr::baht_tsqr::test_qr::(157); + } + + #[test] + pub fn test_big() { + crate::suite::qr::baht_tsqr::test_qr::(517); + } + + #[test] + pub fn test_rect() { + crate::suite::qr::baht_tsqr::test_qr_rect::(517, 157); + } + + #[test] + pub fn test_rect_row_major() { + crate::suite::qr::baht_tsqr::test_qr_rect_row_major::(157, 47); + } + + #[test] + pub fn test_tf32_opt_in() { + crate::suite::qr::baht_tsqr::test_qr_tf32::(157); + } + + }; + ([$($float:ident),*]) => { + mod baht_tsqr { + ::paste::paste! { + $(mod [<$float _ty>] { + $crate::testgen_qr_baht_tsqr!($float); + })* + } + } + }; +} + +mod test_baht_tsqr { + testgen_qr_baht_tsqr!([f32, f64]); +} + +#[macro_export] +macro_rules! testgen_solve { + ($float:ident) => { + pub type FloatT = $float; + + #[test] + pub fn test_solve_square() { + crate::suite::qr::solve::test_solve_square::(16); + } + + #[test] + pub fn test_solve_rect() { + crate::suite::qr::solve::test_solve_rect::(32, 16); + } + }; + ([$($float:ident),*]) => { + mod solve { + ::paste::paste! { + $(mod [<$float _ty>] { + $crate::testgen_solve!($float); + })* + } + } + }; +} + +mod test_solve { + testgen_solve!([f32, f64]); +} diff --git a/crates/cubek-linalg/tests/suite/qr/solve.rs b/crates/cubek-linalg/tests/suite/qr/solve.rs new file mode 100644 index 000000000..6cfb28fa9 --- /dev/null +++ b/crates/cubek-linalg/tests/suite/qr/solve.rs @@ -0,0 +1,59 @@ +use crate::suite::utils::{assert_equals_approx, col_major_input, dtype_unsupported}; +use cubecl::{TestRuntime, prelude::*}; + +/// Build a diagonally dominant `[rows, cols]` matrix (logical row-major), +/// a known solution `x_true`, and the matching right-hand side `b = A·x_true`. +fn make_system(rows: usize, cols: usize) -> (Vec, Vec, Vec) { + let mut a_data = vec![F::from_int(0); rows * cols]; + for i in 0..rows { + for j in 0..cols { + a_data[i * cols + j] = if i == j { + F::from_int(10) + } else { + F::from_int(1) + }; + } + } + + let x_true: Vec = (0..cols).map(|i| F::from_int((i + 1) as i64)).collect(); + + let mut b_data = vec![F::from_int(0); rows]; + for i in 0..rows { + let mut sum = F::from_int(0); + for j in 0..cols { + sum += a_data[i * cols + j] * x_true[j]; + } + b_data[i] = sum; + } + + (a_data, x_true, b_data) +} + +fn run_solve(rows: u32, cols: u32) { + let client = TestRuntime::client(&Default::default()); + if dtype_unsupported::(&client) { + return; + } + let rows_usize = rows as usize; + let cols_usize = cols as usize; + + let (a_data, x_true, b_data) = make_system::(rows_usize, cols_usize); + + let a = col_major_input(&client, vec![rows_usize, cols_usize], &a_data); + let b = col_major_input(&client, vec![rows_usize], &b_data); + + let x = cubek_linalg::solve::(&client, &a, &b).unwrap(); + + let x_bytes = client.read_one(x.handle.clone()).unwrap(); + let x_vals = F::from_bytes(&x_bytes).to_vec(); + + assert_equals_approx::(&x_vals, &x_true, vec![cols_usize], 5e-2); +} + +pub fn test_solve_square(dim: u32) { + run_solve::(dim, dim); +} + +pub fn test_solve_rect(rows: u32, cols: u32) { + run_solve::(rows, cols); +} diff --git a/crates/cubek-linalg/tests/suite/utils.rs b/crates/cubek-linalg/tests/suite/utils.rs new file mode 100644 index 000000000..ceb1c9d0d --- /dev/null +++ b/crates/cubek-linalg/tests/suite/utils.rs @@ -0,0 +1,112 @@ +use cubecl::features::TypeUsage; +use cubecl::std::tensor::TensorHandle; +use cubecl::zspace::{Shape, Strides}; +use cubecl::{TestRuntime, prelude::*}; +use cubek_test_utils::{HostData, HostDataVec, StridedLayout, TestInput, ValidationResult}; + +/// Returns `true` (and prints a skip notice) when the active backend can't do +/// reliable arithmetic in `F`, so the precision-sensitive QR tests would fail +/// for a reason unrelated to the algorithm. Tests call this first and return +/// early to avoid such false failures; CUDA (full f64) runs them all. +/// +/// Two signals are combined: +/// - The backend's advertised type support (`supported_uses`), which catches +/// backends that honestly report a missing type. +/// - A WGSL special case: the WGSL spec has **no 64-bit float**, yet the +/// `wgpu` backend still advertises full f64 support and then silently +/// produces garbage. We can't trust the feature flag there, so f64 on any +/// WGSL runtime is treated as unsupported by name. +pub(crate) fn dtype_unsupported( + client: &ComputeClient, +) -> bool { + let unsupported_by_features = !F::supported_uses(client).contains(TypeUsage::Arithmetic); + + // f64 is 8 bytes; f32/f16/bf16 are smaller. The QR test suite only uses + // f32 and f64, so size == 8 uniquely identifies f64. + let is_f64 = core::mem::size_of::() == 8; + let runtime = TestRuntime::name(client); + let wgsl_f64 = is_f64 && runtime.contains("wgsl"); + + if unsupported_by_features || wgsl_f64 { + println!( + "Skipping: backend `{runtime}` does not reliably support `{}` arithmetic.", + core::any::type_name::() + ); + true + } else { + false + } +} + +/// Build a device tensor from logical row-major values via the shared +/// `TestInput` builder (which owns the stride math). The builder's payload +/// type is f32; the QR test matrices hold small integers, so the round-trip +/// is lossless for both f32 and f64. +fn device_input( + client: &ComputeClient, + shape: Vec, + data_row_major: &[F], + layout: StridedLayout, +) -> TensorHandle { + TestInput::builder(client.clone(), Shape::from(shape)) + .dtype(F::as_type_native_unchecked().storage_type()) + .layout(layout) + .custom(data_row_major.iter().map(|v| v.to_f32().unwrap()).collect()) + .generate_without_host_data() +} + +pub(crate) fn col_major_input( + client: &ComputeClient, + shape: Vec, + data_row_major: &[F], +) -> TensorHandle { + // ColMajor requires rank >= 2; for vectors both layouts are the same. + let layout = if shape.len() >= 2 { + StridedLayout::ColMajor + } else { + StridedLayout::RowMajor + }; + device_input(client, shape, data_row_major, layout) +} + +pub(crate) fn row_major_input( + client: &ComputeClient, + shape: Vec, + data_row_major: &[F], +) -> TensorHandle { + device_input(client, shape, data_row_major, StridedLayout::RowMajor) +} + +/// Wrap logical row-major host values in a `HostData` blob (f64 payload for +/// f64 tensors, f32 otherwise) for the shared comparator. +fn host_data(shape: Vec, values: &[F]) -> HostData { + let storage = F::as_type_native_unchecked().storage_type(); + let mut strides = vec![1usize; shape.len()]; + for i in (0..shape.len().saturating_sub(1)).rev() { + strides[i] = strides[i + 1] * shape[i + 1]; + } + let values: Vec = values.iter().map(|v| v.to_f64().unwrap()).collect(); + HostData { + data: HostDataVec::from((values, storage)), + shape: Shape::from(shape), + strides: Strides::from(strides), + } +} + +/// Compare two logical row-major value slices through cubek-test-utils' +/// shared comparator (relative epsilon with an absolute floor), panicking +/// with its mismatch report on failure. +pub(crate) fn assert_equals_approx( + actual: &[F], + expected: &[F], + shape: Vec, + epsilon: f32, +) { + let actual = host_data::(shape.clone(), actual); + let expected = host_data::(shape, expected); + match cubek_test_utils::assert_equals_approx(&actual, &expected, epsilon) { + ValidationResult::Pass => {} + ValidationResult::Fail(msg) | ValidationResult::Error(msg) => panic!("{msg}"), + ValidationResult::Skipped(msg) => panic!("unexpected validation skip: {msg}"), + } +}