diff --git a/candle-core/Cargo.toml b/candle-core/Cargo.toml index cbf8f2007f..a05f966a3f 100644 --- a/candle-core/Cargo.toml +++ b/candle-core/Cargo.toml @@ -52,3 +52,7 @@ harness = false [[example]] name = "metal_basics" required-features = ["metal"] + +[[example]] +name = "cuda_basics" +required-features = ["cuda"] diff --git a/candle-core/examples/cuda_basics.rs b/candle-core/examples/cuda_basics.rs index 9af1b006e3..315fe0a282 100644 --- a/candle-core/examples/cuda_basics.rs +++ b/candle-core/examples/cuda_basics.rs @@ -7,8 +7,92 @@ extern crate intel_mkl_src; use anyhow::Result; use candle_core::{Device, Tensor}; +const USE_CUDA_GRAPH: bool = true; + +fn cuda_graph() -> Result<()> { + let device = Device::new_cuda_with_stream(0)?; + let cu_device = match &device { + Device::Cuda(dev) => dev, + _ => unreachable!(), + }; + let cu_stream = cu_device.cu_stream(); + { + // load_ptx cannot be called while capturing the stream so we need this to happen + // beforehand. + let u = Tensor::zeros((4096, 4096), candle_core::DType::F32, &device)? + .to_dtype(candle_core::DType::BF16)?; + let mut x = Tensor::zeros((4096, 4096), candle_core::DType::F32, &device)? + .to_dtype(candle_core::DType::BF16)?; + let v = Tensor::zeros(4096, candle_core::DType::F32, &device)? + .to_dtype(candle_core::DType::BF16)?; + let _x = x.mul(&u)?.broadcast_add(&v)?; + let _x = x.affine(1., 0.5)?; + x.slice_set(&u, 0, 0)?; + device.synchronize()?; + } + if USE_CUDA_GRAPH { + unsafe { + cudarc::driver::sys::lib() + .cuStreamBeginCapture_v2( + *cu_stream, + cudarc::driver::sys::CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, + ) + .result()? + } + } + { + let u = Tensor::zeros((4096, 4096), candle_core::DType::F32, &device)? + .to_dtype(candle_core::DType::BF16)?; + let mut x = Tensor::zeros((4096, 4096), candle_core::DType::F32, &device)? + .to_dtype(candle_core::DType::BF16)?; + let v = Tensor::zeros((4096, 1), candle_core::DType::F32, &device)? + .to_dtype(candle_core::DType::BF16)?; + for _i in 0..100 { + // x.slice_set(&u, 0, 0)?; + // x.broadcast_add(&v)?; + x = x.affine(1., 0.5)?; + // x = (&u + &x)?; + } + } + if USE_CUDA_GRAPH { + let cu_graph: cudarc::driver::sys::CUgraph = unsafe { + let mut cu_graph = std::mem::MaybeUninit::uninit(); + cudarc::driver::sys::lib() + .cuStreamEndCapture(*cu_stream, cu_graph.as_mut_ptr()) + .result()?; + cu_graph.assume_init() + }; + let cu_graph_e: cudarc::driver::sys::CUgraphExec = unsafe { + let mut cu_graph_e = std::mem::MaybeUninit::uninit(); + cudarc::driver::sys::lib() + .cuGraphInstantiateWithFlags(cu_graph_e.as_mut_ptr(), cu_graph, 0) + .result()?; + cu_graph_e.assume_init() + }; + println!("graph captured!"); + for i in 1..100 { + println!("graph exec {i}"); + unsafe { + cudarc::driver::sys::lib() + .cuGraphLaunch(cu_graph_e, *cu_stream) + .result()? + } + println!("sync"); + if let Err(err) = device.synchronize() { + println!("err: {err:?}") + } + println!("done syncing"); + } + } else { + device.synchronize()?; + } + Ok(()) +} + fn main() -> Result<()> { - let device = Device::new_cuda(0)?; + cuda_graph()?; + return Ok(()); + let device = Device::new_cuda_with_stream(0)?; let x = Tensor::randn(0f32, 1.0, (8 * 4096, 8 * 4096), &device)? .to_dtype(candle_core::DType::BF16)?; candle_core::cuda::set_gemm_reduced_precision_f32(false);