huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Tweaks to the graph experiment.

Browse Source
This commit is contained in:
laurent
2024-10-03 17:12:52 +02:00
parent b2956857ef
commit 1bb68854d3
1 changed files with 44 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
candle-core/examples/cuda_basics.rs
View File

@ -7,6 +7,8 @@ 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 {
@ -24,8 +26,11 @@ fn cuda_graph() -> Result<()> {
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(
@ -33,38 +38,53 @@ fn cuda_graph() -> Result<()> {
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, candle_core::DType::F32, &device)?
let v = Tensor::zeros((4096, 1), candle_core::DType::F32, &device)?
.to_dtype(candle_core::DType::BF16)?;
for _i in 0..1 {
x = x.mul(&u)?.broadcast_add(&v)?;
for _i in 0..100 {
// x.slice_set(&u, 0, 0)?;
// x.broadcast_add(&v)?;
x = x.affine(1., 0.5)?;
// x = (&u + &x)?;
}
}
let cu_graph = unsafe {
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 = unsafe {
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()
};
for _i in 0..100 {
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(())
}