diff --git a/Cargo.toml b/Cargo.toml index 27f0b25a4..88508bc41 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -153,6 +153,8 @@ cudarc = { version = "0.19.6", features = [ "std", "driver", "nvrtc", + "cublas", + "cublaslt", "nccl", "fallback-dynamic-loading", "cuda-version-from-build-system", diff --git a/crates/cubecl-cuda/examples/cublas_gemm.rs b/crates/cubecl-cuda/examples/cublas_gemm.rs new file mode 100644 index 000000000..54ba64c2e --- /dev/null +++ b/crates/cubecl-cuda/examples/cublas_gemm.rs @@ -0,0 +1,290 @@ +use cubecl_core::{ + Runtime, + ir::{ElemType, FloatKind}, + server::{GemmDescriptor, GemmMatrix}, +}; +use cubecl_cuda::{CudaDevice, CudaRuntime}; +use cubecl_environment::future; +use half::bf16; +use std::time::Instant; + +#[derive(Clone, Copy)] +struct Shape { + m: usize, + n: usize, + k: usize, + batch: usize, + lhs_t: bool, + rhs_t: bool, + lhs_broadcast: bool, + rhs_broadcast: bool, + repeats: usize, +} + +fn main() { + let args = std::env::args().skip(1).collect::>(); + if args.first().is_some_and(|arg| arg == "check") { + correctness_matrix(); + return; + } + if args.first().is_some_and(|arg| arg == "enqueue") { + enqueue_overhead(args.get(1).map_or(10_000, |value| parse(value))); + return; + } + if args.len() != 9 { + eprintln!( + "usage: cublas_gemm M N K BATCH LHS_T RHS_T LHS_BROADCAST RHS_BROADCAST REPEATS\n cublas_gemm check\n cublas_gemm enqueue [REPEATS]" + ); + std::process::exit(2); + } + let shape = Shape { + m: parse(&args[0]), + n: parse(&args[1]), + k: parse(&args[2]), + batch: parse(&args[3]), + lhs_t: parse_flag(&args[4]), + rhs_t: parse_flag(&args[5]), + lhs_broadcast: parse_flag(&args[6]), + rhs_broadcast: parse_flag(&args[7]), + repeats: parse(&args[8]), + }; + let ms = benchmark(shape); + let flops = 2.0 * shape.m as f64 * shape.n as f64 * shape.k as f64 * shape.batch as f64; + println!( + "m={} n={} k={} batch={} lhs_t={} rhs_t={} lhs_broadcast={} rhs_broadcast={} ms={ms:.4} tflops={:.2}", + shape.m, + shape.n, + shape.k, + shape.batch, + shape.lhs_t, + shape.rhs_t, + shape.lhs_broadcast, + shape.rhs_broadcast, + flops / (ms * 1.0e9), + ); +} + +fn enqueue_overhead(repeats: usize) { + let client = CudaRuntime::client(&CudaDevice::default()); + let one = bf16::ONE.to_bits(); + let lhs = client.create_from_slice(bytemuck::cast_slice(&[one])); + let rhs = client.create_from_slice(bytemuck::cast_slice(&[one])); + let out = client.empty(2); + let descriptor = descriptor( + Shape { + m: 1, + n: 1, + k: 1, + batch: 1, + lhs_t: false, + rhs_t: false, + lhs_broadcast: false, + rhs_broadcast: false, + repeats, + }, + lhs.binding(), + rhs.binding(), + out.binding(), + ); + client.gemm(descriptor.clone()); + future::block_on(client.sync()).unwrap(); + + let start = Instant::now(); + for _ in 0..repeats { + client.gemm(descriptor.clone()); + } + let enqueue = start.elapsed(); + let drain_start = Instant::now(); + future::block_on(client.sync()).unwrap(); + let drain = drain_start.elapsed(); + println!( + "repeats={repeats} enqueue_us_per_call={:.3} final_gpu_drain_ms={:.3}", + enqueue.as_secs_f64() * 1.0e6 / repeats as f64, + drain.as_secs_f64() * 1.0e3, + ); +} + +fn benchmark(shape: Shape) -> f64 { + let client = CudaRuntime::client(&CudaDevice::default()); + let lhs_batches = if shape.lhs_broadcast { 1 } else { shape.batch }; + let rhs_batches = if shape.rhs_broadcast { 1 } else { shape.batch }; + let lhs_elems = lhs_batches * shape.m * shape.k; + let rhs_elems = rhs_batches * shape.k * shape.n; + let out_elems = shape.batch * shape.m * shape.n; + let lhs = client.create_from_slice(bytemuck::cast_slice(&values(lhs_elems, 3))); + let rhs = client.create_from_slice(bytemuck::cast_slice(&values(rhs_elems, 7))); + let out = client.empty(out_elems * 2); + let descriptor = descriptor(shape, lhs.binding(), rhs.binding(), out.binding()); + + for _ in 0..5 { + client.gemm(descriptor.clone()); + } + future::block_on(client.sync()).unwrap(); + + let start = Instant::now(); + for _ in 0..shape.repeats { + client.gemm(descriptor.clone()); + } + future::block_on(client.sync()).unwrap(); + start.elapsed().as_secs_f64() * 1_000.0 / shape.repeats as f64 +} + +fn correctness_matrix() { + for lhs_t in [false, true] { + for rhs_t in [false, true] { + for batch in [1, 3] { + for lhs_broadcast in [false, true] { + for rhs_broadcast in [false, true] { + if batch == 1 && (lhs_broadcast || rhs_broadcast) { + continue; + } + check(Shape { + m: 5, + n: 7, + k: 3, + batch, + lhs_t, + rhs_t, + lhs_broadcast, + rhs_broadcast, + repeats: 1, + }); + } + } + } + } + } + println!("cuBLAS BF16 layout correctness matrix passed"); +} + +fn check(shape: Shape) { + let client = CudaRuntime::client(&CudaDevice::default()); + let lhs_batches = if shape.lhs_broadcast { 1 } else { shape.batch }; + let rhs_batches = if shape.rhs_broadcast { 1 } else { shape.batch }; + let lhs_bits = values(lhs_batches * shape.m * shape.k, 3); + let rhs_bits = values(rhs_batches * shape.k * shape.n, 7); + let lhs = client.create_from_slice(bytemuck::cast_slice(&lhs_bits)); + let rhs = client.create_from_slice(bytemuck::cast_slice(&rhs_bits)); + let out = client.empty(shape.batch * shape.m * shape.n * 2); + client.gemm(descriptor( + shape, + lhs.binding(), + rhs.binding(), + out.clone().binding(), + )); + let bytes = client.read_one_unchecked(out); + let actual = bytemuck::cast_slice::(&bytes); + + for batch in 0..shape.batch { + for row in 0..shape.m { + for col in 0..shape.n { + let expected = (0..shape.k) + .map(|inner| { + get( + &lhs_bits, + shape.m, + shape.k, + shape.lhs_t, + if shape.lhs_broadcast { 0 } else { batch }, + row, + inner, + ) * get( + &rhs_bits, + shape.k, + shape.n, + shape.rhs_t, + if shape.rhs_broadcast { 0 } else { batch }, + inner, + col, + ) + }) + .sum::(); + let index = (batch * shape.m + row) * shape.n + col; + let actual = bf16::from_bits(actual[index]).to_f32(); + assert!( + (actual - expected).abs() <= 0.06, + "layout failed: lhs_t={} rhs_t={} lhs_broadcast={} rhs_broadcast={} batch={batch} row={row} col={col}: {actual} != {expected}", + shape.lhs_t, + shape.rhs_t, + shape.lhs_broadcast, + shape.rhs_broadcast, + ); + } + } + } +} + +fn descriptor( + shape: Shape, + lhs: cubecl_core::server::Binding, + rhs: cubecl_core::server::Binding, + out: cubecl_core::server::Binding, +) -> GemmDescriptor { + GemmDescriptor::new( + GemmMatrix::new( + lhs, + if shape.lhs_t { shape.m } else { shape.k } as u32, + if shape.lhs_broadcast { + 0 + } else { + (shape.m * shape.k) as u64 + }, + shape.lhs_t, + ), + GemmMatrix::new( + rhs, + if shape.rhs_t { shape.k } else { shape.n } as u32, + if shape.rhs_broadcast { + 0 + } else { + (shape.k * shape.n) as u64 + }, + shape.rhs_t, + ), + GemmMatrix::new(out, shape.n as u32, (shape.m * shape.n) as u64, false), + shape.m as u32, + shape.n as u32, + shape.k as u32, + shape.batch as u32, + ElemType::Float(FloatKind::BF16), + ) +} + +fn values(len: usize, offset: usize) -> Vec { + (0..len) + .map(|index| { + let value = ((index + offset) % 11) as f32 / 32.0 - 0.15; + bf16::from_f32(value).to_bits() + }) + .collect() +} + +fn get( + bits: &[u16], + rows: usize, + cols: usize, + transposed: bool, + batch: usize, + row: usize, + col: usize, +) -> f32 { + let base = batch * rows * cols; + let index = if transposed { + base + col * rows + row + } else { + base + row * cols + col + }; + bf16::from_bits(bits[index]).to_f32() +} + +fn parse(value: &str) -> usize { + value.parse().unwrap() +} + +fn parse_flag(value: &str) -> bool { + match value { + "0" => false, + "1" => true, + _ => panic!("flag must be 0 or 1"), + } +} diff --git a/crates/cubecl-cuda/examples/cublas_gemm_check.rs b/crates/cubecl-cuda/examples/cublas_gemm_check.rs new file mode 100644 index 000000000..7ae6e5287 --- /dev/null +++ b/crates/cubecl-cuda/examples/cublas_gemm_check.rs @@ -0,0 +1,510 @@ +use cubecl_core as cubecl; +use cubecl_core::{ + Runtime, + ir::{ElemType, FloatKind}, + prelude::*, + server::{Binding, GemmDescriptor, GemmMatrix, GroupedGemmDescriptor, Handle}, +}; +use cubecl_cuda::{CudaDevice, CudaRuntime}; +use cubecl_environment::{future, stream::StreamId}; +use half::bf16; + +#[cube(launch)] +fn copy_bf16(input: &[bf16], output: &mut [bf16]) { + if ABSOLUTE_POS < output.len() { + output[ABSOLUTE_POS] = input[ABSOLUTE_POS]; + } +} + +#[derive(Clone, Copy, Debug)] +struct Problem { + m: usize, + n: usize, + k: usize, + batch: usize, + lhs_t: bool, + rhs_t: bool, + lhs_broadcast: bool, + rhs_broadcast: bool, + padding: usize, + batch_gap: usize, + offset: usize, +} + +struct Matrix { + bits: Vec, + rows: usize, + cols: usize, + batches: usize, + transposed: bool, + ld: usize, + batch_stride: usize, + offset: usize, +} + +fn main() { + correctness_matrix(); + grouped_correctness(); + cross_stream_ordering(); + overlapping_output_is_rejected(); + aliased_output_is_rejected(); + foreign_output_is_rejected(); + zero_k_is_rejected(); + prior_error_is_not_bypassed(); + println!("cuBLAS BF16 padded/offset/batched/grouped/multistream checks passed"); +} + +fn correctness_matrix() { + for lhs_t in [false, true] { + for rhs_t in [false, true] { + for lhs_broadcast in [false, true] { + for rhs_broadcast in [false, true] { + check(Problem { + m: 5, + n: 7, + k: 3, + batch: 3, + lhs_t, + rhs_t, + lhs_broadcast, + rhs_broadcast, + padding: 5, + batch_gap: 11, + offset: 13, + }); + } + } + } + } +} + +fn check(problem: Problem) { + let client = CudaRuntime::client(&CudaDevice::default()); + let lhs = Matrix::new( + problem.m, + problem.k, + if problem.lhs_broadcast { + 1 + } else { + problem.batch + }, + problem.lhs_t, + problem.padding, + problem.batch_gap, + problem.offset, + 3, + ); + let rhs = Matrix::new( + problem.k, + problem.n, + if problem.rhs_broadcast { + 1 + } else { + problem.batch + }, + problem.rhs_t, + problem.padding + 2, + problem.batch_gap + 3, + problem.offset + 5, + 7, + ); + let out = Matrix::zeros( + problem.m, + problem.n, + problem.batch, + false, + problem.padding + 4, + problem.batch_gap + 7, + problem.offset + 9, + ); + + let lhs_base = client.create_from_slice(bytemuck::cast_slice(&lhs.bits)); + let rhs_base = client.create_from_slice(bytemuck::cast_slice(&rhs.bits)); + let out_base = client.empty(out.bits.len() * 2); + let lhs_view = view(&lhs_base, lhs.offset, lhs.bits.len()); + let rhs_view = view(&rhs_base, rhs.offset, rhs.bits.len()); + let out_view = view(&out_base, out.offset, out.bits.len()); + let descriptor = GemmDescriptor::new( + matrix_arg(&lhs, lhs_view, problem.lhs_broadcast), + matrix_arg(&rhs, rhs_view, problem.rhs_broadcast), + matrix_arg(&out, out_view, false), + problem.m as u32, + problem.n as u32, + problem.k as u32, + problem.batch as u32, + ElemType::Float(FloatKind::BF16), + ); + client.gemm(descriptor); + let bytes = client.read_one_unchecked(out_base); + let actual = bytemuck::cast_slice::(&bytes); + + for batch in 0..problem.batch { + for row in 0..problem.m { + for col in 0..problem.n { + let expected = (0..problem.k) + .map(|inner| { + lhs.get(if problem.lhs_broadcast { 0 } else { batch }, row, inner) + * rhs.get(if problem.rhs_broadcast { 0 } else { batch }, inner, col) + }) + .sum::(); + let index = out.index(batch, row, col); + let actual = bf16::from_bits(actual[index]).to_f32(); + assert!( + (actual - expected).abs() <= 0.06, + "{problem:?}, b={batch}, row={row}, col={col}: {actual} != {expected}" + ); + } + } + } +} + +fn grouped_correctness() { + let client = CudaRuntime::client(&CudaDevice::default()); + let elem = ElemType::Float(FloatKind::BF16); + if !client + .features() + .matmul + .accelerated_grouped_gemm + .contains(&elem) + { + return; + } + let problems = [(3, 5, 4), (7, 2, 3), (4, 6, 5)]; + let mut matrices = Vec::with_capacity(problems.len()); + let mut groups = Vec::with_capacity(problems.len()); + + for (index, (m, n, k)) in problems.into_iter().enumerate() { + let lhs = Matrix::new(m, k, 1, index % 2 == 0, 3 + index, 0, 7, index + 2); + let rhs = Matrix::new(k, n, 1, index % 2 != 0, 5 + index, 0, 11, index + 5); + let out = Matrix::zeros(m, n, 1, false, 2 + index, 0, 13); + let lhs_base = client.create_from_slice(bytemuck::cast_slice(&lhs.bits)); + let rhs_base = client.create_from_slice(bytemuck::cast_slice(&rhs.bits)); + let out_base = client.empty(out.bits.len() * 2); + groups.push(GemmDescriptor::new( + matrix_arg(&lhs, view(&lhs_base, lhs.offset, lhs.bits.len()), false), + matrix_arg(&rhs, view(&rhs_base, rhs.offset, rhs.bits.len()), false), + matrix_arg(&out, view(&out_base, out.offset, out.bits.len()), false), + m as u32, + n as u32, + k as u32, + 1, + elem, + )); + matrices.push((lhs, rhs, out, lhs_base, rhs_base, out_base)); + } + + let descriptor = GroupedGemmDescriptor::new(groups); + for _ in 0..16 { + client.grouped_gemm(descriptor.clone()); + } + for (lhs, rhs, out, _lhs_base, _rhs_base, out_base) in matrices { + let bytes = client.read_one_unchecked(out_base); + let actual = bytemuck::cast_slice::(&bytes); + for row in 0..out.rows { + for col in 0..out.cols { + let expected = (0..lhs.cols) + .map(|inner| lhs.get(0, row, inner) * rhs.get(0, inner, col)) + .sum::(); + let actual = bf16::from_bits(actual[out.index(0, row, col)]).to_f32(); + assert!( + (actual - expected).abs() <= 0.06, + "grouped m={}, n={}, k={}, row={row}, col={col}: {actual} != {expected}", + out.rows, + out.cols, + lhs.cols + ); + } + } + } +} + +fn cross_stream_ordering() { + let mut producer = CudaRuntime::client(&CudaDevice::default()); + let mut gemm = producer.clone(); + let mut consumer = producer.clone(); + unsafe { + producer.set_stream(StreamId { value: 100 }); + gemm.set_stream(StreamId { value: 101 }); + consumer.set_stream(StreamId { value: 102 }); + } + + let problem = Problem { + m: 8, + n: 6, + k: 4, + batch: 2, + lhs_t: false, + rhs_t: true, + lhs_broadcast: false, + rhs_broadcast: true, + padding: 0, + batch_gap: 0, + offset: 0, + }; + let lhs = Matrix::new(8, 4, 2, false, 0, 0, 0, 3); + let rhs = Matrix::new(4, 6, 1, true, 0, 0, 0, 7); + let lhs_handle = producer.create_from_slice(bytemuck::cast_slice(&lhs.bits)); + let rhs_handle = producer.create_from_slice(bytemuck::cast_slice(&rhs.bits)); + let gemm_out = gemm.empty(problem.batch * problem.m * problem.n * 2); + let descriptor = GemmDescriptor::new( + matrix_arg(&lhs, lhs_handle.binding(), false), + matrix_arg(&rhs, rhs_handle.binding(), true), + GemmMatrix::new( + gemm_out.clone().binding(), + problem.n as u32, + (problem.m * problem.n) as u64, + false, + ), + problem.m as u32, + problem.n as u32, + problem.k as u32, + problem.batch as u32, + ElemType::Float(FloatKind::BF16), + ); + gemm.gemm(descriptor); + + let consumed = consumer.empty(problem.batch * problem.m * problem.n * 2); + copy_bf16::launch::( + &consumer, + CubeCount::Static(1, 1, 1), + CubeDim::new(&consumer, 128), + unsafe { BufferArg::from_raw_parts(gemm_out, problem.batch * problem.m * problem.n) }, + unsafe { + BufferArg::from_raw_parts(consumed.clone(), problem.batch * problem.m * problem.n) + }, + ); + let bytes = consumer.read_one_unchecked(consumed); + let actual = bytemuck::cast_slice::(&bytes); + for batch in 0..problem.batch { + for row in 0..problem.m { + for col in 0..problem.n { + let expected = (0..problem.k) + .map(|inner| lhs.get(batch, row, inner) * rhs.get(0, inner, col)) + .sum::(); + let index = (batch * problem.m + row) * problem.n + col; + let actual = bf16::from_bits(actual[index]).to_f32(); + assert!((actual - expected).abs() <= 0.06); + } + } + } +} + +fn overlapping_output_is_rejected() { + let mut client = CudaRuntime::client(&CudaDevice::default()); + unsafe { client.set_stream(StreamId { value: 103 }) }; + + let m = 2; + let n = 3; + let k = 4; + let batches = 2; + let lhs = client.create_from_slice(bytemuck::cast_slice(&vec![ + bf16::ONE.to_bits(); + batches * m * k + ])); + let rhs = client.create_from_slice(bytemuck::cast_slice(&vec![ + bf16::ONE.to_bits(); + batches * k * n + ])); + let out = client.empty(batches * m * n * 2); + let descriptor = GemmDescriptor::new( + GemmMatrix::new(lhs.binding(), k as u32, (m * k) as u64, false), + GemmMatrix::new(rhs.binding(), n as u32, (k * n) as u64, false), + GemmMatrix::new(out.binding(), n as u32, (m * n - 1) as u64, false), + m as u32, + n as u32, + k as u32, + batches as u32, + ElemType::Float(FloatKind::BF16), + ); + client.gemm(descriptor); + assert!(future::block_on(client.sync()).is_err()); +} + +fn aliased_output_is_rejected() { + let mut client = CudaRuntime::client(&CudaDevice::default()); + unsafe { client.set_stream(StreamId { value: 104 }) }; + + let values = vec![bf16::ONE.to_bits(); 4]; + let lhs_and_out = client.create_from_slice(bytemuck::cast_slice(&values)); + let rhs = client.create_from_slice(bytemuck::cast_slice(&values)); + client.gemm(GemmDescriptor::new( + GemmMatrix::new(lhs_and_out.clone().binding(), 2, 0, false), + GemmMatrix::new(rhs.binding(), 2, 0, false), + GemmMatrix::new(lhs_and_out.binding(), 2, 0, false), + 2, + 2, + 2, + 1, + ElemType::Float(FloatKind::BF16), + )); + assert!(future::block_on(client.sync()).is_err()); +} + +fn foreign_output_is_rejected() { + let mut origin = CudaRuntime::client(&CudaDevice::default()); + let mut execution = origin.clone(); + unsafe { + origin.set_stream(StreamId { value: 105 }); + execution.set_stream(StreamId { value: 106 }); + } + + let values = vec![bf16::ONE.to_bits(); 4]; + let lhs = execution.create_from_slice(bytemuck::cast_slice(&values)); + let rhs = execution.create_from_slice(bytemuck::cast_slice(&values)); + let foreign_out = origin.empty(8); + execution.gemm(GemmDescriptor::new( + GemmMatrix::new(lhs.binding(), 2, 0, false), + GemmMatrix::new(rhs.binding(), 2, 0, false), + GemmMatrix::new(foreign_out.binding(), 2, 0, false), + 2, + 2, + 2, + 1, + ElemType::Float(FloatKind::BF16), + )); + assert!(future::block_on(execution.sync()).is_err()); +} + +fn zero_k_is_rejected() { + let mut client = CudaRuntime::client(&CudaDevice::default()); + unsafe { client.set_stream(StreamId { value: 107 }) }; + + let placeholder = client.empty(2); + let out = client.empty(12); + client.gemm(GemmDescriptor::new( + GemmMatrix::new(placeholder.clone().binding(), 1, 0, false), + GemmMatrix::new(placeholder.binding(), 3, 0, false), + GemmMatrix::new(out.binding(), 3, 0, false), + 2, + 3, + 0, + 1, + ElemType::Float(FloatKind::BF16), + )); + assert!(future::block_on(client.sync()).is_err()); +} + +fn prior_error_is_not_bypassed() { + let client = CudaRuntime::client(&CudaDevice::default()); + let input = client.create_from_slice(bytemuck::cast_slice(&[bf16::ONE.to_bits()])); + let output = client.empty(2); + + // The CUDA launch limit is at most 1024 units per cube. The fire-and-forget + // launch records a server error; the following GEMM must observe that + // unhealthy stream instead of enqueueing through it. + copy_bf16::launch::( + &client, + CubeCount::Static(1, 1, 1), + CubeDim::new_1d(2048), + unsafe { BufferArg::from_raw_parts(input.clone(), 1) }, + unsafe { BufferArg::from_raw_parts(output.clone(), 1) }, + ); + let descriptor = GemmDescriptor::new( + GemmMatrix::new(input.clone().binding(), 1, 0, false), + GemmMatrix::new(input.binding(), 1, 0, false), + GemmMatrix::new(output.clone().binding(), 1, 0, false), + 1, + 1, + 1, + 1, + ElemType::Float(FloatKind::BF16), + ); + client.gemm(descriptor.clone()); + // Surfaces the original launch error and restores stream health. + assert!(future::block_on(client.sync()).is_err()); + + client.gemm(descriptor); + let bytes = client.read_one_unchecked(output); + assert_eq!( + bytemuck::cast_slice::(&bytes), + &[bf16::ONE.to_bits()] + ); +} + +impl Matrix { + #[allow(clippy::too_many_arguments)] + fn new( + rows: usize, + cols: usize, + batches: usize, + transposed: bool, + padding: usize, + batch_gap: usize, + offset: usize, + seed: usize, + ) -> Self { + let mut matrix = Self::zeros(rows, cols, batches, transposed, padding, batch_gap, offset); + for batch in 0..batches { + for row in 0..rows { + for col in 0..cols { + let logical = (batch * rows + row) * cols + col; + let value = ((logical + seed) % 11) as f32 / 32.0 - 0.15; + let index = matrix.index(batch, row, col); + matrix.bits[index] = bf16::from_f32(value).to_bits(); + } + } + } + matrix + } + + #[allow(clippy::too_many_arguments)] + fn zeros( + rows: usize, + cols: usize, + batches: usize, + transposed: bool, + padding: usize, + batch_gap: usize, + offset: usize, + ) -> Self { + let ld = (if transposed { rows } else { cols }) + padding; + let span = (if transposed { cols } else { rows }) * ld; + let batch_stride = span + batch_gap; + let suffix = 17; + Self { + bits: vec![0; offset + batches * batch_stride + suffix], + rows, + cols, + batches, + transposed, + ld, + batch_stride, + offset, + } + } + + fn index(&self, batch: usize, row: usize, col: usize) -> usize { + assert!(batch < self.batches && row < self.rows && col < self.cols); + let matrix = self.offset + batch * self.batch_stride; + if self.transposed { + matrix + col * self.ld + row + } else { + matrix + row * self.ld + col + } + } + + fn get(&self, batch: usize, row: usize, col: usize) -> f32 { + bf16::from_bits(self.bits[self.index(batch, row, col)]).to_f32() + } +} + +fn matrix_arg(matrix: &Matrix, binding: Binding, broadcast: bool) -> GemmMatrix { + GemmMatrix::new( + binding, + matrix.ld as u32, + if broadcast { + 0 + } else { + matrix.batch_stride as u64 + }, + matrix.transposed, + ) +} + +fn view(base: &Handle, offset: usize, total: usize) -> Binding { + let suffix = total - offset; + base.clone() + .offset_start((offset * 2) as u64) + .offset_end((suffix.min(17) * 2) as u64) + .binding() +} diff --git a/crates/cubecl-cuda/src/compute/cublas.rs b/crates/cubecl-cuda/src/compute/cublas.rs new file mode 100644 index 000000000..8ef2df633 --- /dev/null +++ b/crates/cubecl-cuda/src/compute/cublas.rs @@ -0,0 +1,866 @@ +use std::collections::HashMap; +use std::ffi::c_void; + +#[cfg(cuda_12050)] +use cubecl_core::server::GroupedGemmDescriptor; +use cubecl_core::{ + ir::{ElemType, FloatKind}, + server::{GemmDescriptor, GemmMatrix, ServerError}, +}; +use cubecl_environment::backtrace::BackTrace; +use cudarc::cublas::sys::cublasOperation_t; +#[cfg(cuda_12050)] +use cudarc::cublas::{result as blas, sys as blas_sys}; +use cudarc::cublaslt::{result as lt, sys}; + +use super::storage::gpu::GpuResource; + +/// Workspace given to every cublasLt matmul. 32 MiB matches the size the +/// cuBLAS documentation recommends for Ampere+ so the heuristic can pick +/// split-K and other workspace-hungry algorithms. +const WORKSPACE_BYTES: usize = 32 * 1024 * 1024; + +/// A cached execution plan: the descriptor/layout objects plus the algorithm +/// the cublasLt heuristic selected for one GEMM shape. +struct MatmulPlan { + desc: sys::cublasLtMatmulDesc_t, + a_layout: sys::cublasLtMatrixLayout_t, + b_layout: sys::cublasLtMatrixLayout_t, + d_layout: sys::cublasLtMatrixLayout_t, + algo: sys::cublasLtMatmulAlgo_t, +} + +/// Pinned-host and device storage for one in-flight grouped-GEMM pointer list. +/// +/// cuBLAS consumes the matrix-pointer arrays on the device. The completion +/// event prevents either side of this staging pair from being overwritten +/// while an earlier grouped launch is still reading it. +#[cfg(cuda_12050)] +struct GroupedPointerStaging { + host: *mut std::ffi::c_void, + device: cudarc::driver::sys::CUdeviceptr, + event: cudarc::driver::sys::CUevent, + capacity: usize, + in_flight: bool, + /// Captured graph nodes retain both staging addresses for replay, so a + /// captured slot must remain immutable until the server is destroyed. + captured: bool, +} + +#[cfg(cuda_12050)] +impl GroupedPointerStaging { + fn new(capacity: usize) -> Result { + let bytes = capacity + .checked_mul(core::mem::size_of::()) + .ok_or_else(|| validation_error("grouped GEMM pointer staging size overflow"))?; + // SAFETY: both allocations are owned by the returned staging slot and + // released exactly once in `destroy`. + let host = unsafe { + cudarc::driver::result::malloc_host( + bytes, + cudarc::driver::sys::CU_MEMHOSTALLOC_WRITECOMBINED, + ) + } + .map_err(cuda_error)?; + let device = match unsafe { cudarc::driver::result::malloc_sync(bytes) } { + Ok(device) => device, + Err(err) => { + // SAFETY: `host` was allocated immediately above and has not + // been exposed or freed. + unsafe { + let _ = cudarc::driver::result::free_host(host); + } + return Err(cuda_error(err)); + } + }; + let event = match cudarc::driver::result::event::create( + cudarc::driver::sys::CUevent_flags_enum::CU_EVENT_DISABLE_TIMING, + ) { + Ok(event) => event, + Err(err) => { + // SAFETY: both allocations are uniquely owned here. + unsafe { + let _ = cudarc::driver::result::free_sync(device); + let _ = cudarc::driver::result::free_host(host); + } + return Err(cuda_error(err)); + } + }; + Ok(Self { + host, + device, + event, + capacity, + in_flight: false, + captured: false, + }) + } + + fn available(&self) -> bool { + !self.captured + && (!self.in_flight + // SAFETY: `event` remains live for this slot's entire lifetime. + || unsafe { cudarc::driver::result::event::query(self.event) }.is_ok()) + } + + fn destroy(self) { + // The server is shutting down. Waiting here protects the pinned host + // buffer if a final pointer upload is still in flight. + if self.in_flight { + // SAFETY: `event` is live and was recorded after the last grouped + // launch using this slot. + let _ = unsafe { cudarc::driver::result::event::synchronize(self.event) }; + } + // SAFETY: all three resources are uniquely owned by this slot. + unsafe { + let _ = cudarc::driver::result::event::destroy(self.event); + let _ = cudarc::driver::result::free_sync(self.device); + let _ = cudarc::driver::result::free_host(self.host); + } + } +} + +#[derive(PartialEq, Eq, Hash)] +struct MatrixKey { + leading_dimension: u32, + batch_stride: u64, + transposed: bool, +} + +#[derive(PartialEq, Eq, Hash)] +struct PlanKey { + m: u32, + n: u32, + k: u32, + batch_count: u32, + lhs: MatrixKey, + rhs: MatrixKey, + out: MatrixKey, +} + +impl PlanKey { + fn new(descriptor: &GemmDescriptor) -> Self { + let matrix = |m: &GemmMatrix| MatrixKey { + leading_dimension: m.leading_dimension, + batch_stride: m.batch_stride, + transposed: m.transposed, + }; + Self { + m: descriptor.m, + n: descriptor.n, + k: descriptor.k, + batch_count: descriptor.batch_count, + lhs: matrix(&descriptor.lhs), + rhs: matrix(&descriptor.rhs), + out: matrix(&descriptor.out), + } + } +} + +#[derive(Default)] +pub(crate) struct CublasState { + handle: Option, + #[cfg(cuda_12050)] + grouped_handle: Option, + /// One workspace per CUDA stream: concurrent matmuls on different + /// streams must not share scratch memory. + workspaces: HashMap, + plans: HashMap, + #[cfg(cuda_12050)] + grouped_staging: Vec, +} + +impl core::fmt::Debug for CublasState { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut debug = f.debug_struct("CublasState"); + debug.field("initialized", &self.handle.is_some()); + #[cfg(cuda_12050)] + debug.field("grouped_initialized", &self.grouped_handle.is_some()); + debug.field("plans", &self.plans.len()); + #[cfg(cuda_12050)] + debug.field("grouped_staging", &self.grouped_staging.len()); + debug.finish() + } +} + +impl CublasState { + pub(crate) fn launch( + &mut self, + descriptor: &GemmDescriptor, + lhs: &GpuResource, + rhs: &GpuResource, + out: &GpuResource, + stream: cudarc::driver::sys::CUstream, + ) -> Result<(), ServerError> { + if descriptor.elem != ElemType::Float(FloatKind::BF16) { + return Err(validation_error("cuBLAS GEMM currently supports only BF16")); + } + if descriptor.m == 0 || descriptor.n == 0 || descriptor.batch_count == 0 { + return Ok(()); + } + if descriptor.k == 0 { + return Err(validation_error( + "cuBLAS GEMM does not initialize a nonempty zero-K output", + )); + } + validate(descriptor, lhs, rhs, out)?; + + let handle = match self.handle { + Some(handle) => handle, + None => { + let handle = lt::create_handle().map_err(cublas_error)?; + self.handle = Some(handle); + handle + } + }; + let workspace = match self.workspaces.get(&(stream as usize)) { + Some(workspace) => *workspace, + None => { + // SAFETY: the server made its CUDA context current before + // resolving the streams for this launch. + let workspace = unsafe { cudarc::driver::result::malloc_sync(WORKSPACE_BYTES) } + .map_err(|err| ServerError::Generic { + reason: format!("cublasLt workspace allocation failed: {err}"), + backtrace: BackTrace::capture(), + })?; + self.workspaces.insert(stream as usize, workspace); + workspace + } + }; + + let key = PlanKey::new(descriptor); + if !self.plans.contains_key(&key) { + let plan = build_plan(handle, descriptor)?; + self.plans.insert(PlanKey::new(descriptor), plan); + } + let plan = self + .plans + .get(&key) + .expect("cublasLt plan was just inserted"); + + let alpha = 1.0f32; + let beta = 0.0f32; + + // cuBLAS is column-major. Swapping the row-major operands computes + // D^T = rhs^T @ lhs^T without copies, so A carries `rhs` and B + // carries `lhs`. The call is asynchronous on the given stream. + unsafe { + lt::matmul( + handle, + plan.desc, + (&alpha as *const f32).cast::(), + (&beta as *const f32).cast::(), + rhs.ptr as *const c_void, + plan.a_layout, + lhs.ptr as *const c_void, + plan.b_layout, + out.ptr as *const c_void, + plan.d_layout, + out.ptr as *mut c_void, + plan.d_layout, + &plan.algo, + workspace as *mut c_void, + WORKSPACE_BYTES, + stream.cast(), + ) + } + .map_err(cublas_error)?; + + Ok(()) + } + + #[cfg(cuda_12050)] + pub(crate) fn launch_grouped( + &mut self, + descriptor: &GroupedGemmDescriptor, + lhs: &[GpuResource], + rhs: &[GpuResource], + out: &[GpuResource], + stream: cudarc::driver::sys::CUstream, + ) -> Result<(), ServerError> { + let Some(elem) = descriptor.groups.first().map(|group| group.elem) else { + return Ok(()); + }; + if elem != ElemType::Float(FloatKind::BF16) { + return Err(validation_error( + "cuBLAS grouped GEMM currently supports only BF16", + )); + } + if descriptor.groups.len() != lhs.len() + || descriptor.groups.len() != rhs.len() + || descriptor.groups.len() != out.len() + { + return Err(validation_error( + "grouped GEMM descriptor/resource count mismatch", + )); + } + + let mut active = Vec::with_capacity(descriptor.groups.len()); + let mut pointer_count = 0usize; + for (index, group) in descriptor.groups.iter().enumerate() { + if group.elem != elem { + return Err(validation_error( + "every grouped GEMM entry must use the descriptor element type", + )); + } + if group.m == 0 || group.n == 0 || group.batch_count == 0 { + continue; + } + if group.k == 0 { + return Err(validation_error( + "cuBLAS grouped GEMM does not initialize a nonempty zero-K output", + )); + } + validate(group, &lhs[index], &rhs[index], &out[index])?; + pointer_count = pointer_count + .checked_add(group.batch_count as usize) + .ok_or_else(|| validation_error("grouped GEMM batch count overflow"))?; + active.push(index); + } + if active.is_empty() { + return Ok(()); + } + for (position, &index) in active.iter().enumerate() { + for &other in &active[position + 1..] { + if resources_overlap(&out[index], &out[other]) + || resources_overlap(&out[index], &lhs[other]) + || resources_overlap(&out[index], &rhs[other]) + || resources_overlap(&out[other], &lhs[index]) + || resources_overlap(&out[other], &rhs[index]) + { + return Err(validation_error( + "grouped GEMM outputs may not overlap another group", + )); + } + } + } + let group_count: i32 = active + .len() + .try_into() + .map_err(|_| validation_error("grouped GEMM has more than i32::MAX groups"))?; + let staging_values = pointer_count + .checked_mul(3) + .ok_or_else(|| validation_error("grouped GEMM pointer count overflow"))?; + + let slot_index = match self + .grouped_staging + .iter() + .position(|slot| slot.capacity >= staging_values && slot.available()) + { + Some(index) => index, + None => { + let capacity = staging_values + .checked_next_power_of_two() + .ok_or_else(|| validation_error("grouped GEMM staging capacity overflow"))?; + self.grouped_staging + .push(GroupedPointerStaging::new(capacity)?); + self.grouped_staging.len() - 1 + } + }; + let slot = &mut self.grouped_staging[slot_index]; + // CUDA graphs replay the captured host-to-device copy, so both the + // pinned source and device destination addresses must remain stable. + // Retaining a tiny slot per captured grouped launch provides that + // lifetime without imposing synchronization on ordinary launches. + // SAFETY: `stream` is the live execution stream resolved by the server. + slot.captured = !matches!( + unsafe { cudarc::driver::result::stream::is_capturing(stream) }.map_err(cuda_error)?, + cudarc::driver::sys::CUstreamCaptureStatus::CU_STREAM_CAPTURE_STATUS_NONE + ); + // SAFETY: the pinned allocation contains `capacity` u64 values and the + // selected slot is no longer in use by a previous launch. + let host = + unsafe { core::slice::from_raw_parts_mut(slot.host.cast::(), slot.capacity) }; + + let mut trans_a = Vec::with_capacity(active.len()); + let mut trans_b = Vec::with_capacity(active.len()); + let mut m = Vec::with_capacity(active.len()); + let mut n = Vec::with_capacity(active.len()); + let mut k = Vec::with_capacity(active.len()); + let mut lda = Vec::with_capacity(active.len()); + let mut ldb = Vec::with_capacity(active.len()); + let mut ldc = Vec::with_capacity(active.len()); + let mut group_sizes = Vec::with_capacity(active.len()); + let alpha = vec![1.0_f32; active.len()]; + let beta = vec![0.0_f32; active.len()]; + let elem_size = core::mem::size_of::() as u64; + let (a_values, rest) = host[..staging_values].split_at_mut(pointer_count); + let (b_values, c_values) = rest.split_at_mut(pointer_count); + let mut pointer = 0; + + for &index in &active { + let group = &descriptor.groups[index]; + // cuBLAS is column-major. Swapping the row-major operands computes + // D^T = rhs^T @ lhs^T without materialization. + trans_a.push(operation(&group.rhs)); + trans_b.push(operation(&group.lhs)); + m.push(group.n as i32); + n.push(group.m as i32); + k.push(group.k as i32); + lda.push(group.rhs.leading_dimension as i32); + ldb.push(group.lhs.leading_dimension as i32); + ldc.push(group.out.leading_dimension as i32); + group_sizes.push(group.batch_count as i32); + for batch in 0..group.batch_count as u64 { + a_values[pointer] = batch_pointer(&rhs[index], &group.rhs, batch, elem_size)?; + b_values[pointer] = batch_pointer(&lhs[index], &group.lhs, batch, elem_size)?; + c_values[pointer] = batch_pointer(&out[index], &group.out, batch, elem_size)?; + pointer += 1; + } + } + debug_assert_eq!(pointer, pointer_count); + + // SAFETY: the host slice is pinned and remains untouched until the + // completion event recorded below. The device allocation is large + // enough for exactly `staging_values` pointers. + unsafe { + cudarc::driver::result::memcpy_htod_async(slot.device, &host[..staging_values], stream) + } + .map_err(cuda_error)?; + + let handle = match self.grouped_handle { + Some(handle) => handle, + None => { + let handle = blas::create_handle().map_err(grouped_cublas_error)?; + self.grouped_handle = Some(handle); + handle + } + }; + // SAFETY: the handle and CubeCL stream are live for the server. + unsafe { blas::set_stream(handle, stream.cast()) }.map_err(grouped_cublas_error)?; + let a_device = slot.device as *const *const c_void; + let b_device = (slot.device + (pointer_count * core::mem::size_of::()) as u64) + as *const *const c_void; + let c_device = (slot.device + (2 * pointer_count * core::mem::size_of::()) as u64) + as *const *mut c_void; + // SAFETY: dimensions and leading dimensions were validated; all + // pointer arrays reside in the staging device allocation and refer to + // live CubeCL resources ordered on `stream`. + let launch = unsafe { + blas_sys::cublasGemmGroupedBatchedEx( + handle, + trans_a.as_ptr(), + trans_b.as_ptr(), + m.as_ptr(), + n.as_ptr(), + k.as_ptr(), + alpha.as_ptr().cast::(), + a_device, + blas_sys::cudaDataType_t::CUDA_R_16BF, + lda.as_ptr(), + b_device, + blas_sys::cudaDataType_t::CUDA_R_16BF, + ldb.as_ptr(), + beta.as_ptr().cast::(), + c_device, + blas_sys::cudaDataType_t::CUDA_R_16BF, + ldc.as_ptr(), + group_count, + group_sizes.as_ptr(), + blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, + ) + .result() + }; + // Record even when cuBLAS rejects the launch: the pointer upload is + // already queued and must complete before the slot is reused. + // SAFETY: `event` and `stream` are live. + unsafe { cudarc::driver::result::event::record(slot.event, stream) }.map_err(cuda_error)?; + slot.in_flight = true; + launch.map_err(grouped_cublas_error)?; + Ok(()) + } + + pub(crate) fn destroy(&mut self) { + for (_, plan) in self.plans.drain() { + // SAFETY: each plan uniquely owns its descriptor objects and is + // destroyed exactly once. + unsafe { + let _ = lt::destroy_matmul_desc(plan.desc); + let _ = lt::destroy_matrix_layout(plan.a_layout); + let _ = lt::destroy_matrix_layout(plan.b_layout); + let _ = lt::destroy_matrix_layout(plan.d_layout); + } + } + for (_, workspace) in self.workspaces.drain() { + // SAFETY: the workspace was allocated by this state on the + // server's context and freed exactly once. + if let Err(err) = unsafe { cudarc::driver::result::free_sync(workspace) } { + log::warn!("Unable to free cublasLt workspace: {err}"); + } + } + #[cfg(cuda_12050)] + for staging in self.grouped_staging.drain(..) { + staging.destroy(); + } + if let Some(handle) = self.handle.take() { + // SAFETY: this state uniquely owns the handle and destroys it once. + if let Err(err) = unsafe { lt::destroy_handle(handle) } { + log::warn!("Unable to destroy cublasLt handle: {err}"); + } + } + #[cfg(cuda_12050)] + if let Some(handle) = self.grouped_handle.take() { + // SAFETY: this state uniquely owns the handle and destroys it once. + if let Err(err) = unsafe { blas::destroy_handle(handle) } { + log::warn!("Unable to destroy cuBLAS grouped-GEMM handle: {err}"); + } + } + } +} + +/// Build the descriptor, layouts, and heuristic-selected algorithm for one +/// GEMM shape. All dimensions below are in cuBLAS column-major terms, i.e. +/// the row-major operands are swapped: `m_lt = n`, `n_lt = m`, `A = rhs`, +/// `B = lhs`. +fn build_plan( + handle: sys::cublasLtHandle_t, + descriptor: &GemmDescriptor, +) -> Result { + let m_lt = descriptor.n as u64; + let n_lt = descriptor.m as u64; + let k_lt = descriptor.k as u64; + let op_a = operation(&descriptor.rhs); + let op_b = operation(&descriptor.lhs); + + let desc = lt::create_matmul_desc( + sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, + sys::cudaDataType::CUDA_R_32F, + ) + .map_err(cublas_error)?; + let destroy_desc = || { + // SAFETY: created above, not yet owned by a plan. + unsafe { + let _ = lt::destroy_matmul_desc(desc); + } + }; + for (attr, op) in [ + ( + sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA, + op_a, + ), + ( + sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB, + op_b, + ), + ] { + let value = op as i32; + // SAFETY: `desc` is live and the attribute is an i32 by contract. + if let Err(err) = unsafe { + lt::set_matmul_desc_attribute( + desc, + attr, + (&value as *const i32).cast::(), + core::mem::size_of::(), + ) + } { + destroy_desc(); + return Err(cublas_error(err)); + } + } + + // Stored (pre-transpose) dimensions of each column-major operand. + let a_dims = if matches!(op_a, cublasOperation_t::CUBLAS_OP_N) { + (m_lt, k_lt) + } else { + (k_lt, m_lt) + }; + let b_dims = if matches!(op_b, cublasOperation_t::CUBLAS_OP_N) { + (k_lt, n_lt) + } else { + (n_lt, k_lt) + }; + let layout = |rows: u64, cols: u64, matrix: &GemmMatrix| -> Result<_, ServerError> { + let layout = lt::create_matrix_layout( + sys::cudaDataType::CUDA_R_16BF, + rows, + cols, + matrix.leading_dimension as i64, + ) + .map_err(cublas_error)?; + let batch_count = descriptor.batch_count as i32; + let batch_stride = matrix.batch_stride as i64; + // SAFETY: `layout` is live; both attributes take the given widths. + let result = unsafe { + lt::set_matrix_layout_attribute( + layout, + sys::cublasLtMatrixLayoutAttribute_t::CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, + (&batch_count as *const i32).cast::(), + core::mem::size_of::(), + ) + .and_then(|_| { + lt::set_matrix_layout_attribute( + layout, + sys::cublasLtMatrixLayoutAttribute_t:: + CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, + (&batch_stride as *const i64).cast::(), + core::mem::size_of::(), + ) + }) + }; + if let Err(err) = result { + // SAFETY: created above, not yet owned by a plan. + unsafe { + let _ = lt::destroy_matrix_layout(layout); + } + return Err(cublas_error(err)); + } + Ok(layout) + }; + + let a_layout = match layout(a_dims.0, a_dims.1, &descriptor.rhs) { + Ok(layout) => layout, + Err(err) => { + destroy_desc(); + return Err(err); + } + }; + let b_layout = match layout(b_dims.0, b_dims.1, &descriptor.lhs) { + Ok(layout) => layout, + Err(err) => { + destroy_desc(); + // SAFETY: created above, not yet owned by a plan. + unsafe { + let _ = lt::destroy_matrix_layout(a_layout); + } + return Err(err); + } + }; + let d_layout = match layout(m_lt, n_lt, &descriptor.out) { + Ok(layout) => layout, + Err(err) => { + destroy_desc(); + // SAFETY: created above, not yet owned by a plan. + unsafe { + let _ = lt::destroy_matrix_layout(a_layout); + let _ = lt::destroy_matrix_layout(b_layout); + } + return Err(err); + } + }; + let cleanup = || { + destroy_desc(); + // SAFETY: created above, not yet owned by a plan. + unsafe { + let _ = lt::destroy_matrix_layout(a_layout); + let _ = lt::destroy_matrix_layout(b_layout); + let _ = lt::destroy_matrix_layout(d_layout); + } + }; + + let pref = match lt::create_matmul_pref() { + Ok(pref) => pref, + Err(err) => { + cleanup(); + return Err(cublas_error(err)); + } + }; + let workspace_bytes = WORKSPACE_BYTES as u64; + // SAFETY: `pref` is live and the attribute is a u64 by contract. + if let Err(err) = unsafe { + lt::set_matmul_pref_attribute( + pref, + sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + (&workspace_bytes as *const u64).cast::(), + core::mem::size_of::(), + ) + } { + // SAFETY: created above. + unsafe { + let _ = lt::destroy_matmul_pref(pref); + } + cleanup(); + return Err(cublas_error(err)); + } + + // SAFETY: every descriptor is live; the C layout equals the D layout + // because beta is always zero. + let heuristic = unsafe { + lt::get_matmul_algo_heuristic(handle, desc, a_layout, b_layout, d_layout, d_layout, pref) + }; + // SAFETY: created above. + unsafe { + let _ = lt::destroy_matmul_pref(pref); + } + let heuristic = match heuristic { + Ok(heuristic) => heuristic, + Err(_) => { + cleanup(); + return Err(validation_error( + "cublasLt has no algorithm for this GEMM shape", + )); + } + }; + + Ok(MatmulPlan { + desc, + a_layout, + b_layout, + d_layout, + algo: heuristic.algo, + }) +} + +fn operation(matrix: &GemmMatrix) -> cublasOperation_t { + if matrix.transposed { + cublasOperation_t::CUBLAS_OP_T + } else { + cublasOperation_t::CUBLAS_OP_N + } +} + +fn validate( + descriptor: &GemmDescriptor, + lhs: &GpuResource, + rhs: &GpuResource, + out: &GpuResource, +) -> Result<(), ServerError> { + let max_i32 = i32::MAX as u32; + if descriptor.m > max_i32 + || descriptor.n > max_i32 + || descriptor.k > max_i32 + || descriptor.batch_count > max_i32 + { + return Err(validation_error( + "GEMM dimensions exceed the cuBLAS i32 API", + )); + } + if descriptor.out.transposed { + return Err(validation_error("GEMM output must be row-major")); + } + if resources_overlap(out, lhs) || resources_overlap(out, rhs) { + return Err(validation_error("GEMM output may not overlap either input")); + } + + validate_matrix( + "lhs", + &descriptor.lhs, + descriptor.m, + descriptor.k, + descriptor.batch_count, + lhs, + true, + )?; + validate_matrix( + "rhs", + &descriptor.rhs, + descriptor.k, + descriptor.n, + descriptor.batch_count, + rhs, + true, + )?; + validate_matrix( + "out", + &descriptor.out, + descriptor.m, + descriptor.n, + descriptor.batch_count, + out, + false, + ) +} + +fn resources_overlap(lhs: &GpuResource, rhs: &GpuResource) -> bool { + if lhs.size == 0 || rhs.size == 0 { + return false; + } + let lhs_end = lhs.ptr.saturating_add(lhs.size); + let rhs_end = rhs.ptr.saturating_add(rhs.size); + lhs.ptr < rhs_end && rhs.ptr < lhs_end +} + +#[cfg(cuda_12050)] +fn batch_pointer( + resource: &GpuResource, + matrix: &GemmMatrix, + batch: u64, + elem_size: u64, +) -> Result { + batch + .checked_mul(matrix.batch_stride) + .and_then(|offset| offset.checked_mul(elem_size)) + .and_then(|offset| resource.ptr.checked_add(offset)) + .ok_or_else(|| validation_error("grouped GEMM batch pointer overflow")) +} + +#[allow(clippy::too_many_arguments)] +fn validate_matrix( + name: &str, + matrix: &GemmMatrix, + rows: u32, + cols: u32, + batches: u32, + resource: &GpuResource, + allow_broadcast: bool, +) -> Result<(), ServerError> { + let required_ld = if matrix.transposed { rows } else { cols }; + if matrix.leading_dimension < required_ld || matrix.leading_dimension > i32::MAX as u32 { + return Err(validation_error(&format!( + "{name} leading dimension {} is invalid for logical shape [{rows}, {cols}]", + matrix.leading_dimension + ))); + } + if matrix.batch_stride > i64::MAX as u64 { + return Err(validation_error(&format!( + "{name} batch stride exceeds the cuBLAS i64 API" + ))); + } + let matrix_elems = if rows == 0 || cols == 0 { + 0 + } else if matrix.transposed { + (cols as u64 - 1) * matrix.leading_dimension as u64 + rows as u64 + } else { + (rows as u64 - 1) * matrix.leading_dimension as u64 + cols as u64 + }; + if !allow_broadcast && batches > 1 && matrix.batch_stride < matrix_elems { + return Err(validation_error("GEMM output batches may not overlap")); + } + let batch_offset = if batches <= 1 || matrix.batch_stride == 0 { + 0 + } else { + (batches as u64 - 1) + .checked_mul(matrix.batch_stride) + .ok_or_else(|| validation_error("GEMM batch stride overflow"))? + }; + let required_bytes = batch_offset + .checked_add(matrix_elems) + .and_then(|elements| elements.checked_mul(2)) + .ok_or_else(|| validation_error("GEMM buffer size overflow"))?; + if required_bytes > resource.size { + return Err(validation_error(&format!( + "{name} requires {required_bytes} bytes but its binding contains {}", + resource.size + ))); + } + Ok(()) +} + +fn validation_error(message: &str) -> ServerError { + ServerError::Validation { + message: message.into(), + backtrace: BackTrace::capture(), + } +} + +fn cublas_error(error: lt::CublasError) -> ServerError { + ServerError::Generic { + reason: format!("cuBLAS error: {error:?}"), + backtrace: BackTrace::capture(), + } +} + +#[cfg(cuda_12050)] +fn grouped_cublas_error(error: blas::CublasError) -> ServerError { + ServerError::Generic { + reason: format!("cuBLAS grouped GEMM error: {error:?}"), + backtrace: BackTrace::capture(), + } +} + +#[cfg(cuda_12050)] +fn cuda_error(error: cudarc::driver::DriverError) -> ServerError { + ServerError::Generic { + reason: format!("CUDA grouped GEMM staging error: {error:?}"), + backtrace: BackTrace::capture(), + } +} diff --git a/crates/cubecl-cuda/src/compute/mod.rs b/crates/cubecl-cuda/src/compute/mod.rs index c595738e3..a3ed1a094 100644 --- a/crates/cubecl-cuda/src/compute/mod.rs +++ b/crates/cubecl-cuda/src/compute/mod.rs @@ -1,6 +1,7 @@ pub(crate) mod command; pub(crate) mod communication; pub(crate) mod context; +pub(crate) mod cublas; pub(crate) mod io; pub(crate) mod storage; pub(crate) mod stream; diff --git a/crates/cubecl-cuda/src/compute/server.rs b/crates/cubecl-cuda/src/compute/server.rs index d8f3107cc..c1717e610 100644 --- a/crates/cubecl-cuda/src/compute/server.rs +++ b/crates/cubecl-cuda/src/compute/server.rs @@ -5,20 +5,23 @@ use crate::{ command::Command, communication::{get_nccl_comm_id, get_nccl_dtype_count, to_nccl_op}, context::CudaContext, + cublas::CublasState, stream::CudaStreamBackend, sync::Fence, }, }; use cubecl_common::{bytes::Bytes, profile::ProfileDuration}; +#[cfg(cuda_12050)] +use cubecl_core::server::GroupedGemmDescriptor; use cubecl_core::{ MemoryConfiguration, device::DeviceId, ir::{ElemType, FloatKind, IntKind, MemoryDeviceProperties, StorageType, UIntKind}, prelude::*, server::{ - Binding, CommunicationId, CopyDescriptor, Handle, KernelArguments, LaunchError, - ProfileError, ProfilingToken, ReduceOperation, ServerCommunication, ServerError, - ServerUtilities, StreamErrorMode, TensorMapBinding, TensorMapMeta, + Binding, CommunicationId, CopyDescriptor, GemmDescriptor, Handle, KernelArguments, + LaunchError, ProfileError, ProfilingToken, ReduceOperation, ServerCommunication, + ServerError, ServerUtilities, StreamErrorMode, TensorMapBinding, TensorMapMeta, }, }; use cubecl_environment::backtrace::BackTrace; @@ -56,6 +59,7 @@ pub struct CudaServer { utilities: Arc>, comm_stream: *mut CUstream_st, communicators: HashMap, + cublas: CublasState, } // SAFETY: `CudaServer` is only accessed from one thread at a time via the `DeviceHandle`, @@ -166,6 +170,27 @@ impl ComputeServer for CudaServer { } } + fn gemm(&mut self, descriptor: GemmDescriptor, stream_id: StreamId) { + if let Err(err) = self.gemm_checked(descriptor, stream_id) { + let mut stream = match self.streams.resolve(stream_id, [].into_iter(), false) { + Ok(stream) => stream, + Err(err) => unreachable!("{err}"), + }; + stream.current().errors.push(err); + } + } + + #[cfg(cuda_12050)] + fn grouped_gemm(&mut self, descriptor: GroupedGemmDescriptor, stream_id: StreamId) { + if let Err(err) = self.grouped_gemm_checked(descriptor, stream_id) { + let mut stream = match self.streams.resolve(stream_id, [].into_iter(), false) { + Ok(stream) => stream, + Err(err) => unreachable!("{err}"), + }; + stream.current().errors.push(err); + } + } + fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError> { let mut command = self.command_no_inputs( stream_id, @@ -597,6 +622,7 @@ impl CudaServer { utilities: Arc::new(utilities), comm_stream, communicators: HashMap::default(), + cublas: CublasState::default(), } } @@ -657,6 +683,80 @@ impl CudaServer { errors } + fn gemm_checked( + &mut self, + descriptor: GemmDescriptor, + stream_id: StreamId, + ) -> Result<(), ServerError> { + // A binding records its allocation stream, which is also the stream + // later consumers use for dependency tracking. Executing a write on a + // different stream would leave that metadata stale and let a consumer + // on the allocation stream race the GEMM. + if descriptor.out.binding.stream != stream_id { + return Err(ServerError::Validation { + message: "GEMM output must be allocated on the execution stream".into(), + backtrace: BackTrace::capture(), + }); + } + let bindings = [ + &descriptor.lhs.binding, + &descriptor.rhs.binding, + &descriptor.out.binding, + ]; + self.unsafe_set_current(); + let streams = self + .streams + .resolve(stream_id, bindings.into_iter(), true)?; + let cublas = &mut self.cublas; + let mut command = Command::new(&mut self.ctx, streams); + let lhs = command.resource(descriptor.lhs.binding.clone())?; + let rhs = command.resource(descriptor.rhs.binding.clone())?; + let out = command.resource(descriptor.out.binding.clone())?; + let stream = command.streams.current().sys; + + cublas.launch(&descriptor, &lhs, &rhs, &out, stream) + } + + #[cfg(cuda_12050)] + fn grouped_gemm_checked( + &mut self, + descriptor: GroupedGemmDescriptor, + stream_id: StreamId, + ) -> Result<(), ServerError> { + if descriptor + .groups + .iter() + .any(|group| group.out.binding.stream != stream_id) + { + return Err(ServerError::Validation { + message: "grouped GEMM outputs must be allocated on the execution stream".into(), + backtrace: BackTrace::capture(), + }); + } + let bindings = descriptor + .groups + .iter() + .flat_map(|group| [&group.lhs.binding, &group.rhs.binding, &group.out.binding]) + .collect::>(); + self.unsafe_set_current(); + let streams = self + .streams + .resolve(stream_id, bindings.into_iter(), true)?; + let cublas = &mut self.cublas; + let mut command = Command::new(&mut self.ctx, streams); + let mut lhs = Vec::with_capacity(descriptor.groups.len()); + let mut rhs = Vec::with_capacity(descriptor.groups.len()); + let mut out = Vec::with_capacity(descriptor.groups.len()); + for group in &descriptor.groups { + lhs.push(command.resource(group.lhs.binding.clone())?); + rhs.push(command.resource(group.rhs.binding.clone())?); + out.push(command.resource(group.out.binding.clone())?); + } + let stream = command.streams.current().sys; + + cublas.launch_grouped(&descriptor, &lhs, &rhs, &out, stream) + } + fn launch_checked( &mut self, kernel: Box>, @@ -919,6 +1019,13 @@ impl CudaServer { } } +impl Drop for CudaServer { + fn drop(&mut self) { + self.unsafe_set_current(); + self.cublas.destroy(); + } +} + fn elem_to_tensor_map_type(ty: StorageType) -> CUtensorMapDataType { use cudarc::driver::sys::CUtensorMapDataType::*; match ty { diff --git a/crates/cubecl-cuda/src/runtime.rs b/crates/cubecl-cuda/src/runtime.rs index 2a32feb78..b6136e1d5 100644 --- a/crates/cubecl-cuda/src/runtime.rs +++ b/crates/cubecl-cuda/src/runtime.rs @@ -212,6 +212,19 @@ impl DeviceService for CudaServer { if arch_version >= 80 { device_props.features.copy_async = true; + device_props + .features + .matmul + .accelerated_gemm + .insert(ElemType::Float(FloatKind::BF16)); + #[cfg(cuda_12050)] + { + device_props + .features + .matmul + .accelerated_grouped_gemm + .insert(ElemType::Float(FloatKind::BF16)); + } } // NOTE: I commented that since I observed synchronisation issues with atomic add for bf16. diff --git a/crates/cubecl-ir/src/features.rs b/crates/cubecl-ir/src/features.rs index 54fb027d8..d42de0bd7 100644 --- a/crates/cubecl-ir/src/features.rs +++ b/crates/cubecl-ir/src/features.rs @@ -1,4 +1,4 @@ -use crate::{AddressType, OpaqueType, SemanticType, StorageType, Type}; +use crate::{AddressType, ElemType, OpaqueType, SemanticType, StorageType, Type}; use alloc::collections::{BTreeMap, BTreeSet}; use enumset::EnumSetType; @@ -53,6 +53,10 @@ pub struct Types { /// Matrix multiplication-related features #[derive(Debug, Clone, PartialEq, Eq, Default, Hash)] pub struct MatmulFeatures { + /// Element types supported by a backend-native GEMM implementation. + pub accelerated_gemm: BTreeSet, + /// Element types supported by a backend-native grouped GEMM implementation. + pub accelerated_grouped_gemm: BTreeSet, /// The cmma feature enables cooperative matrix-multiply and accumulate operations. pub cmma: BTreeSet, /// Cube MMA is like cmma but at the cube level, rather than the plane level. diff --git a/crates/cubecl-runtime/src/client.rs b/crates/cubecl-runtime/src/client.rs index b213f1f10..fe68d170f 100644 --- a/crates/cubecl-runtime/src/client.rs +++ b/crates/cubecl-runtime/src/client.rs @@ -7,10 +7,10 @@ use crate::{ memory_management::{MemoryAllocationMode, MemoryConfiguration, MemoryUsage}, runtime::Runtime, server::{ - CommunicationId, ComputeServer, CopyDescriptor, CubeCount, ExecutionMode, Handle, IoError, - KernelArguments, MemoryLayout, MemoryLayoutDescriptor, MemoryLayoutPolicy, - MemoryLayoutStrategy, ProfileError, ReduceOperation, ServerCommunication, ServerError, - ServerUtilities, + CommunicationId, ComputeServer, CopyDescriptor, CubeCount, ExecutionMode, GemmDescriptor, + GroupedGemmDescriptor, Handle, IoError, KernelArguments, MemoryLayout, + MemoryLayoutDescriptor, MemoryLayoutPolicy, MemoryLayoutStrategy, ProfileError, + ReduceOperation, ServerCommunication, ServerError, ServerUtilities, }, storage::{ComputeStorage, ManagedResource}, throughput::{ @@ -366,6 +366,28 @@ impl ComputeClient { .unwrap_or_resume() } + /// Enqueue a backend-accelerated GEMM. + /// + /// The operation is ordered on this client's `CubeCL` stream and returns + /// immediately, like a kernel launch. Callers must first check the device's + /// `features.matmul.accelerated_gemm` capability. + pub fn gemm(&self, descriptor: GemmDescriptor) { + let stream_id = self.stream_id(); + self.device + .submit(move |server| server.gemm(descriptor, stream_id)); + } + + /// Enqueue a backend-accelerated grouped GEMM. + /// + /// The operation is ordered on this client's `CubeCL` stream and returns + /// immediately, like a kernel launch. Callers must first check the device's + /// `features.matmul.accelerated_grouped_gemm` capability. + pub fn grouped_gemm(&self, descriptor: GroupedGemmDescriptor) { + let stream_id = self.stream_id(); + self.device + .submit(move |server| server.grouped_gemm(descriptor, stream_id)); + } + fn do_create_from_slices( &self, descriptors: Vec, diff --git a/crates/cubecl-runtime/src/memory_management/memory_manage.rs b/crates/cubecl-runtime/src/memory_management/memory_manage.rs index 4f6b396ac..ef2207550 100644 --- a/crates/cubecl-runtime/src/memory_management/memory_manage.rs +++ b/crates/cubecl-runtime/src/memory_management/memory_manage.rs @@ -938,20 +938,35 @@ impl MemoryManagement { ); // Find first pool that fits this allocation - let pool = self - .pools - .iter_mut() - .find(|p| p.accept(size)) - .ok_or(IoError::BufferTooBig { - size, - backtrace: BackTrace::capture(), - })?; + let pool_index = + self.pools + .iter() + .position(|p| p.accept(size)) + .ok_or(IoError::BufferTooBig { + size, + backtrace: BackTrace::capture(), + })?; - if let Some(slice) = pool.try_reserve(size) { + if let Some(slice) = self.pools[pool_index].try_reserve(size) { return Ok(slice); } - let allocated = pool.alloc(&mut self.storage, size); + let allocated = self.pools[pool_index].alloc(&mut self.storage, size); + + // Autotuning can leave one-off pages cached in several size buckets. + // If the device is full, release every unused page and retry once + // before surfacing the allocation error. + let allocated = match allocated { + Err(IoError::BufferTooBig { .. }) => { + self.cleanup(true); + if let Some(slice) = self.pools[pool_index].try_reserve(size) { + Ok(slice) + } else { + self.pools[pool_index].alloc(&mut self.storage, size) + } + } + result => result, + }; self.logger.log_memory( |level| matches!(level, MemoryLogLevel::Full), @@ -1070,8 +1085,12 @@ impl core::fmt::Debug for MemoryManagement { #[cfg(test)] mod tests { use super::*; - use crate::{memory_management::MemoryManagement, storage::BytesStorage}; + use crate::{ + memory_management::MemoryManagement, + storage::{BytesResource, BytesStorage, ComputeStorage, StorageHandle, StorageId}, + }; use alloc::vec; + use hashbrown::HashMap; const DUMMY_MEM_PROPS: MemoryDeviceProperties = MemoryDeviceProperties { max_page_size: 128 * 1024 * 1024, @@ -1085,6 +1104,86 @@ mod tests { } } + #[derive(Debug)] + struct LimitedStorage { + inner: BytesStorage, + capacity: u64, + allocated: u64, + sizes: HashMap, + } + + impl LimitedStorage { + fn new(capacity: u64) -> Self { + Self { + inner: BytesStorage::default(), + capacity, + allocated: 0, + sizes: HashMap::new(), + } + } + } + + impl ComputeStorage for LimitedStorage { + type Resource = BytesResource; + + fn alignment(&self) -> usize { + self.inner.alignment() + } + + fn get(&mut self, handle: &StorageHandle) -> Self::Resource { + self.inner.get(handle) + } + + fn alloc(&mut self, size: u64) -> Result { + if self.allocated + size > self.capacity { + return Err(IoError::BufferTooBig { + size, + backtrace: BackTrace::capture(), + }); + } + let handle = self.inner.alloc(size)?; + self.allocated += size; + self.sizes.insert(handle.id, size); + Ok(handle) + } + + fn dealloc(&mut self, id: StorageId) { + if let Some(size) = self.sizes.remove(&id) { + self.allocated -= size; + } + self.inner.dealloc(id); + } + + fn flush(&mut self) { + self.inner.flush(); + } + } + + #[test_log::test] + fn allocation_oom_reclaims_cached_pages_and_retries() { + let mut memory_management = MemoryManagement::from_configuration( + LimitedStorage::new(1024), + &DUMMY_MEM_PROPS, + MemoryConfiguration::Custom { + pool_options: vec![MemoryPoolOptions { + pool_type: PoolType::ExclusivePages { + max_alloc_size: 1024, + }, + dealloc_period: None, + }], + }, + Arc::new(ServerLogger::default()), + options(), + ); + + let cached = memory_management.reserve(768).unwrap(); + drop(cached); + + let _live = memory_management.reserve(1024).unwrap(); + assert_eq!(memory_management.memory_usage().bytes_in_use, 1024); + assert_eq!(memory_management.memory_usage().bytes_reserved, 1024); + } + // Test pools with slices. #[test_log::test] #[cfg(not(exclusive_memory_only))] diff --git a/crates/cubecl-runtime/src/server/base.rs b/crates/cubecl-runtime/src/server/base.rs index cc940dbec..04195d9e2 100644 --- a/crates/cubecl-runtime/src/server/base.rs +++ b/crates/cubecl-runtime/src/server/base.rs @@ -352,6 +352,63 @@ pub struct StreamErrorMode { pub flush: bool, } +/// Describes one matrix operand of an accelerated GEMM. +/// +/// Strides are expressed in elements. A matrix is either row-major +/// (`transposed == false`) or a transposed view of a row-major matrix. The +/// latter has a unit row stride and `leading_dimension` as its column stride. +#[derive(new, Clone, Debug)] +pub struct GemmMatrix { + /// Buffer containing the matrix data. + pub binding: Binding, + /// Leading dimension of the underlying row-major allocation. + pub leading_dimension: u32, + /// Distance between consecutive batch matrices, in elements. Zero + /// broadcasts this operand across batches. + pub batch_stride: u64, + /// Whether the logical matrix is a transposed row-major view. + pub transposed: bool, +} + +/// Backend-optional accelerated GEMM descriptor. +/// +/// Computes `out = lhs @ rhs` for `batch_count` matrices. All matrices use +/// `elem`; accumulation precision is selected by the backend. The output must +/// be row-major, allocated on the execution stream, and disjoint from both +/// inputs. Backends may reject a nonempty zero-K problem rather than provide a +/// separate output-initialization path. +#[allow(clippy::too_many_arguments)] +#[derive(new, Clone, Debug)] +pub struct GemmDescriptor { + /// Left-hand matrix, logically shaped `[m, k]`. + pub lhs: GemmMatrix, + /// Right-hand matrix, logically shaped `[k, n]`. + pub rhs: GemmMatrix, + /// Row-major output matrix, logically shaped `[m, n]`. + pub out: GemmMatrix, + /// Number of rows in the output. + pub m: u32, + /// Number of columns in the output. + pub n: u32, + /// Contracting dimension. + pub k: u32, + /// Number of independent matrix products. + pub batch_count: u32, + /// Scalar element type shared by all operands. + pub elem: ElemType, +} + +/// Backend-optional accelerated grouped GEMM descriptor. +/// +/// Every entry computes one strided batch of `out = lhs @ rhs`. Entries may +/// have different matrix dimensions and batch counts, but must share an +/// element type. A backend may execute the entries concurrently in one launch. +#[derive(new, Clone, Debug)] +pub struct GroupedGemmDescriptor { + /// Independent GEMM groups. An empty collection is a no-op. + pub groups: Vec, +} + /// The compute server is responsible for handling resources and computations over resources. /// /// Everything in the server is mutable, therefore it should be solely accessed through the @@ -435,6 +492,28 @@ where launch_mode: LaunchMode, ); + /// Enqueue an accelerated GEMM when the backend provides one. + /// + /// Callers must first check + /// [`Features::matmul.accelerated_gemm`](cubecl_ir::features::MatmulFeatures). + /// The enqueue returns immediately and records failures on the stream so + /// they surface at its next synchronization point. + fn gemm(&mut self, descriptor: GemmDescriptor, stream_id: StreamId) { + let _ = (descriptor, stream_id); + panic!("Compute server advertised an accelerated GEMM without implementing it") + } + + /// Enqueue an accelerated grouped GEMM when the backend provides one. + /// + /// Callers must first check + /// [`Features::matmul.accelerated_grouped_gemm`](cubecl_ir::features::MatmulFeatures). + /// The enqueue returns immediately and records failures on the stream so + /// they surface at its next synchronization point. + fn grouped_gemm(&mut self, descriptor: GroupedGemmDescriptor, stream_id: StreamId) { + let _ = (descriptor, stream_id); + panic!("Compute server advertised an accelerated grouped GEMM without implementing it") + } + /// Flush all outstanding tasks in the server. fn flush(&mut self, stream_id: StreamId) -> Result<(), ServerError>;