rustytorch/examples/pinn_mre_helmholtz/src/unsafe_graph.rs

//! Raw FFI CUDA Graph Capture
//!
//! This module bypasses cudarc's safe abstractions and uses raw NVIDIA Driver API
//! calls for CUDA graph capture. This is necessary because cudarc's internal safety
//! mechanisms (Mutex guards, stream tracking) interfere with graph capture semantics.
//!
//! # Safety
//!
//! The `UnsafeGraph` struct uses raw CUDA pointers and requires that:
//! 1. All GPU pointers used during capture remain valid for the graph's lifetime
//! 2. The closure passed to `capture()` only issues GPU commands (no CPU allocations)
//! 3. The same stream is used for capture and launch

use std::ptr;
use std::sync::Arc;
use cudarc::driver::safe::CudaStream;
use cudarc::driver::sys::{
    cuStreamBeginCapture_v2, cuStreamEndCapture,
    cuGraphInstantiateWithFlags, cuGraphLaunch,
    cuGraphDestroy, cuGraphExecDestroy,
    CUgraph, CUgraphExec, CUstreamCaptureMode, CUresult,
};

/// A CUDA graph captured using raw FFI calls.
///
/// This bypasses cudarc's safe wrappers to avoid any hidden synchronization
/// or mutex locks that might interfere with graph capture.
pub struct UnsafeGraph {
    /// The executable graph (compiled version)
    graph_exec: CUgraphExec,
    /// The graph definition (kept for cleanup)
    graph: CUgraph,
    /// Reference to the stream used for capture (for launch)
    stream: Arc<CudaStream>,
    /// Number of times the graph has been launched
    launch_count: u64,
}

// CUgraph and CUgraphExec are raw pointers, which are Send but not Sync
// We require explicit synchronization for thread safety
unsafe impl Send for UnsafeGraph {}

impl Drop for UnsafeGraph {
    fn drop(&mut self) {
        unsafe {
            // Clean up in reverse order of creation
            if !self.graph_exec.is_null() {
                cuGraphExecDestroy(self.graph_exec);
            }
            if !self.graph.is_null() {
                cuGraphDestroy(self.graph);
            }
        }
    }
}

impl UnsafeGraph {
    /// Capture a closure into a CUDA graph using raw FFI.
    ///
    /// # Arguments
    /// * `stream` - The CUDA stream to use for capture
    /// * `func` - Closure that issues GPU commands to be captured
    ///
    /// # Safety
    /// The closure must ONLY issue GPU commands. Any CPU-side allocations,
    /// mutex locks, or other blocking operations may cause capture to fail.
    ///
    /// # Returns
    /// * `Ok(UnsafeGraph)` - Successfully captured and compiled graph
    /// * `Err(CUresult)` - CUDA driver error (e.g., STREAM_CAPTURE_ISOLATION)
    pub fn capture<F, E>(stream: Arc<CudaStream>, func: F) -> Result<Self, CUresult>
    where
        F: FnOnce() -> Result<(), E>,
        E: std::fmt::Debug,
    {
        let raw_stream = stream.cu_stream();

        unsafe {
            // 1. Begin Capture with GLOBAL mode
            // GLOBAL mode captures all operations from any stream, which is most permissive
            let res = cuStreamBeginCapture_v2(
                raw_stream,
                CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_GLOBAL
            );
            if res != CUresult::CUDA_SUCCESS {
                #[cfg(debug_assertions)]
                eprintln!("[UnsafeGraph] cuStreamBeginCapture_v2 failed: {:?}", res);
                return Err(res);
            }

            // 2. Execute User Logic (PURE GPU COMMANDS ONLY)
            // If the closure fails, we must end capture to avoid leaving the stream in a bad state
            if let Err(e) = func() {
                #[cfg(debug_assertions)]
                eprintln!("[UnsafeGraph] Capture closure failed: {:?}", e);

                // Abort capture - end it but discard the graph
                let mut tmp_graph = ptr::null_mut();
                let _ = cuStreamEndCapture(raw_stream, &mut tmp_graph);
                if !tmp_graph.is_null() {
                    cuGraphDestroy(tmp_graph);
                }
                return Err(CUresult::CUDA_ERROR_UNKNOWN);
            }

            // 3. End Capture
            let mut graph = ptr::null_mut();
            let res = cuStreamEndCapture(raw_stream, &mut graph);
            if res != CUresult::CUDA_SUCCESS {
                #[cfg(debug_assertions)]
                eprintln!("[UnsafeGraph] cuStreamEndCapture failed: {:?}", res);
                return Err(res);
            }

            if graph.is_null() {
                #[cfg(debug_assertions)]
                eprintln!("[UnsafeGraph] cuStreamEndCapture returned null graph");
                return Err(CUresult::CUDA_ERROR_UNKNOWN);
            }

            // 4. Instantiate (Compile) the graph
            let mut graph_exec = ptr::null_mut();
            // flags = 0 for default behavior
            let res = cuGraphInstantiateWithFlags(
                &mut graph_exec,
                graph,
                0  // flags
            );

            if res != CUresult::CUDA_SUCCESS {
                #[cfg(debug_assertions)]
                eprintln!("[UnsafeGraph] cuGraphInstantiateWithFlags failed: {:?}", res);
                cuGraphDestroy(graph);
                return Err(res);
            }

            #[cfg(debug_assertions)]
            eprintln!("[UnsafeGraph] Graph captured and instantiated successfully");

            Ok(Self {
                graph_exec,
                graph,
                stream,
                launch_count: 0,
            })
        }
    }

    /// Launch the captured graph on its stream.
    ///
    /// This replays all captured GPU operations with minimal CPU overhead.
    /// Typical launch overhead is ~4µs vs ~40µs for 11 individual kernel launches.
    pub fn launch(&mut self) -> Result<(), CUresult> {
        unsafe {
            let res = cuGraphLaunch(self.graph_exec, self.stream.cu_stream());
            if res != CUresult::CUDA_SUCCESS {
                #[cfg(debug_assertions)]
                eprintln!("[UnsafeGraph] cuGraphLaunch failed: {:?}", res);
                return Err(res);
            }
            self.launch_count += 1;
            Ok(())
        }
    }

    /// Get the number of times this graph has been launched.
    pub fn launch_count(&self) -> u64 {
        self.launch_count
    }

    /// Check if the graph is valid (has been successfully captured and instantiated).
    pub fn is_valid(&self) -> bool {
        !self.graph_exec.is_null() && !self.graph.is_null()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_null_pointer_cleanup() {
        // Test that cuGraphDestroy/cuGraphExecDestroy handle null pointers gracefully
        // We test the raw FFI calls directly since we can't safely construct a dummy CudaStream
        unsafe {
            // These should be no-ops for null pointers (CUDA spec allows this)
            let null_graph: CUgraph = ptr::null_mut();
            let null_exec: CUgraphExec = ptr::null_mut();

            // According to CUDA docs, destroying null handles should be safe
            // We just verify no panic/crash occurs
            if !null_exec.is_null() {
                cuGraphExecDestroy(null_exec);
            }
            if !null_graph.is_null() {
                cuGraphDestroy(null_graph);
            }
        }
        // If we get here without panic, the null check logic is correct
    }
}