From 1bb68854d3a37c05a452312b633174b7b9c6d633 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Thu, 3 Oct 2024 17:12:52 +0200
Subject: [PATCH] Tweaks to the graph experiment.

---
 candle-core/examples/cuda_basics.rs | 68 +++++++++++++++++++----------
 1 file changed, 44 insertions(+), 24 deletions(-)

diff --git a/candle-core/examples/cuda_basics.rs b/candle-core/examples/cuda_basics.rs
index 0cd933b59..315fe0a28 100644
--- a/candle-core/examples/cuda_basics.rs
+++ b/candle-core/examples/cuda_basics.rs
@@ -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,47 +26,65 @@ 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()?;
     }
-    unsafe {
-        cudarc::driver::sys::lib()
+    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, 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 {
-        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 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 {
-        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: cudarc::driver::sys::CUgraphExec = unsafe {
+            let mut cu_graph_e = std::mem::MaybeUninit::uninit();
             cudarc::driver::sys::lib()
-                .cuGraphLaunch(cu_graph_e, *cu_stream)
-                .result()?
+                .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(())
 }