feat: add cached CUDA graph inference

README.md

@@ -84,7 +84,14 @@ cargo test
 - The caller must ensure the CUDA context outlives the engine, particularly when cudarc's event tracking is disabled.
 - Only the batch dimension is dynamic. Non-batch dynamic shapes are yet not supported.
 - Engine files are not portable across TensorRT versions or GPU architectures. Rebuild from ONNX for each target.
-- CUDA graphs are not yet supported.
+- `infer_cuda_graph` and `infer_cuda_graph_async` capture and cache a graph for each
+  batch size and device-pointer combination. Keep those device buffers alive while
+  using the cached graph, and call `clear_cuda_graph_cache` before retiring them.
+To compare regular execution with CUDA graph replay:
+
+```
+cargo run --release --example benchmark -- --path /path/to/model.engine --cuda-graph
+```
 
 ## Credits
 

examples/benchmark.rs

@@ -25,6 +25,9 @@ struct Args {
 
     #[arg(short, long, value_name = "DEVICE", default_value_t = 0)]
     device: u32,
+    /// Capture and replay a CUDA graph for each batch/pointer combination.
+    #[arg(long)]
+    cuda_graph: bool,
 }
 
 fn main() {
@@ -88,9 +91,15 @@ fn main() {
 
     info!("Warming up...");
     for _ in 0..1024 {
-        engine
-            .infer(&input_ptrs, &output_ptrs, stream.cu_stream(), batch_size)
-            .unwrap();
+        if args.cuda_graph {
+            engine
+                .infer_cuda_graph(&input_ptrs, &output_ptrs, stream.cu_stream(), batch_size)
+                .unwrap();
+        } else {
+            engine
+                .infer(&input_ptrs, &output_ptrs, stream.cu_stream(), batch_size)
+                .unwrap();
+        }
     }
 
     info!("Running {} inference iterations...", args.iterations);
@@ -99,9 +108,15 @@ fn main() {
 
     for i in 0..args.iterations {
         let start = Instant::now();
-        engine
-            .infer(&input_ptrs, &output_ptrs, stream.cu_stream(), batch_size)
-            .unwrap();
+        if args.cuda_graph {
+            engine
+                .infer_cuda_graph(&input_ptrs, &output_ptrs, stream.cu_stream(), batch_size)
+                .unwrap();
+        } else {
+            engine
+                .infer(&input_ptrs, &output_ptrs, stream.cu_stream(), batch_size)
+                .unwrap();
+        }
         let elapsed = start.elapsed();
         latencies.push(elapsed);
         total_time += elapsed;

src/engine.cpp

@@ -269,6 +269,113 @@ void Engine::infer_async(const uint64_t *input_ptrs, size_t num_inputs,
           batch_size);
 }
 
+void Engine::enqueue_cuda_graph(const uint64_t *input_ptrs, size_t num_inputs,
+                                const uint64_t *output_ptrs, size_t num_outputs,
+                                cudaStream_t stream, uint32_t batch_size,
+                                bool synchronize) {
+  if (batch_size < static_cast<uint32_t>(mMinBatchSize) ||
+      batch_size > static_cast<uint32_t>(mMaxBatchSize)) {
+    throw std::runtime_error(
+        "Batch size " + std::to_string(batch_size) + " is outside [" +
+        std::to_string(mMinBatchSize) + "," + std::to_string(mMaxBatchSize) +
+        "]");
+  }
+  if (num_inputs != get_num_inputs()) {
+    throw std::runtime_error("Expected " + std::to_string(get_num_inputs()) +
+                             " input pointers, got " +
+                             std::to_string(num_inputs));
+  }
+  if (num_outputs != get_num_outputs()) {
+    throw std::runtime_error("Expected " + std::to_string(get_num_outputs()) +
+                             " output pointers, got " +
+                             std::to_string(num_outputs));
+  }
+
+  GraphKey key;
+  key.batch_size = batch_size;
+  key.input_ptrs.reserve(num_inputs);
+  key.output_ptrs.reserve(num_outputs);
+  for (size_t i = 0; i < num_inputs; ++i) {
+    key.input_ptrs.push_back(input_ptrs[i]);
+  }
+  for (size_t i = 0; i < num_outputs; ++i) {
+    key.output_ptrs.push_back(output_ptrs[i]);
+  }
+
+  auto graph = mCudaGraphs.find(key);
+  if (graph == mCudaGraphs.end()) {
+    cudaGraph_t captured_graph = nullptr;
+    cudaGraphExec_t executable = nullptr;
+    bool capturing = false;
+
+    try {
+      checkCudaErrorCode(
+          cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));
+      capturing = true;
+      enqueue(input_ptrs, num_inputs, output_ptrs, num_outputs, stream,
+              batch_size);
+      checkCudaErrorCode(cudaStreamEndCapture(stream, &captured_graph));
+      capturing = false;
+      checkCudaErrorCode(cudaGraphInstantiate(
+          &executable, captured_graph, nullptr, nullptr, 0));
+
+      auto instance = std::make_unique<GraphInstance>();
+      instance->graph = captured_graph;
+      instance->executable = executable;
+      captured_graph = nullptr;
+      executable = nullptr;
+
+      graph = mCudaGraphs.emplace(std::move(key), std::move(instance)).first;
+    } catch (...) {
+      if (capturing) {
+        cudaGraph_t discarded_graph = nullptr;
+        cudaStreamEndCapture(stream, &discarded_graph);
+        if (discarded_graph != nullptr) {
+          cudaGraphDestroy(discarded_graph);
+        }
+      }
+      if (executable != nullptr) {
+        cudaGraphExecDestroy(executable);
+      }
+      if (captured_graph != nullptr) {
+        cudaGraphDestroy(captured_graph);
+      }
+      throw;
+    }
+  }
+
+  checkCudaErrorCode(cudaGraphLaunch(graph->second->executable, stream));
+  if (synchronize) {
+    checkCudaErrorCode(cudaStreamSynchronize(stream));
+  }
+}
+
+void Engine::infer_cuda_graph(const uint64_t *input_ptrs, size_t num_inputs,
+                              const uint64_t *output_ptrs, size_t num_outputs,
+                              uint64_t stream, uint32_t batch_size) {
+  enqueue_cuda_graph(input_ptrs, num_inputs, output_ptrs, num_outputs,
+                     reinterpret_cast<cudaStream_t>(stream), batch_size, true);
+}
+
+void Engine::infer_cuda_graph_async(const uint64_t *input_ptrs,
+                                    size_t num_inputs,
+                                    const uint64_t *output_ptrs,
+                                    size_t num_outputs, uint64_t stream,
+                                    uint32_t batch_size) {
+  enqueue_cuda_graph(input_ptrs, num_inputs, output_ptrs, num_outputs,
+                     reinterpret_cast<cudaStream_t>(stream), batch_size, false);
+}
+
+void Engine::clear_cuda_graph_cache() { mCudaGraphs.clear(); }
+
+Engine::GraphInstance::~GraphInstance() {
+  if (executable != nullptr) {
+    cudaGraphExecDestroy(executable);
+  }
+  if (graph != nullptr) {
+    cudaGraphDestroy(graph);
+  }
+}
 rust::Vec<TensorInfo> Engine::get_input_dims() const {
   rust::Vec<TensorInfo> result;
   for (const auto &meta : mTensorMetadata) {

src/engine.h

@@ -3,11 +3,13 @@
 #include "NvInfer.h"
 #include <cstdlib>
 #include <cuda_runtime.h>
+#include <map>
 #include <memory>
 #include <spdlog/sinks/stdout_color_sinks.h>
 #include <spdlog/spdlog.h>
 #include <string>
 #include <vector>
+#include <tuple>
 
 #include "rust/cxx.h"
 
@@ -80,6 +82,21 @@ class Engine {
   void infer_async(const uint64_t *input_ptrs, size_t num_inputs,
                    const uint64_t *output_ptrs, size_t num_outputs,
                    uint64_t stream, uint32_t batch_size);
+  // Captures and replays a CUDA graph for the given batch and device
+  // pointers. The graph is cached by batch size and pointer addresses.
+  // Synchronizes the stream before returning.
+  void infer_cuda_graph(const uint64_t *input_ptrs, size_t num_inputs,
+                        const uint64_t *output_ptrs, size_t num_outputs,
+                        uint64_t stream, uint32_t batch_size);
+
+  // Same as infer_cuda_graph() but does not synchronize.
+  void infer_cuda_graph_async(const uint64_t *input_ptrs, size_t num_inputs,
+                              const uint64_t *output_ptrs, size_t num_outputs,
+                              uint64_t stream, uint32_t batch_size);
+
+  // Discards all captured graphs. Call this after retiring the device buffers
+  // used to capture them.
+  void clear_cuda_graph_cache();
 
   rust::Vec<TensorInfo> get_input_dims() const;
   rust::Vec<TensorInfo> get_output_dims() const;
@@ -101,6 +118,28 @@ class Engine {
   void enqueue(const uint64_t *input_ptrs, size_t num_inputs,
                const uint64_t *output_ptrs, size_t num_outputs,
                cudaStream_t stream, uint32_t batch_size);
+  struct GraphKey {
+    uint32_t batch_size;
+    std::vector<uint64_t> input_ptrs;
+    std::vector<uint64_t> output_ptrs;
+
+    bool operator<(const GraphKey &other) const {
+      return std::tie(batch_size, input_ptrs, output_ptrs) <
+             std::tie(other.batch_size, other.input_ptrs, other.output_ptrs);
+    }
+  };
+
+  struct GraphInstance {
+    cudaGraph_t graph = nullptr;
+    cudaGraphExec_t executable = nullptr;
+
+    ~GraphInstance();
+  };
+
+  void enqueue_cuda_graph(const uint64_t *input_ptrs, size_t num_inputs,
+                          const uint64_t *output_ptrs, size_t num_outputs,
+                          cudaStream_t stream, uint32_t batch_size,
+                          bool synchronize);
 
   std::vector<TensorMetadata> mTensorMetadata;
   std::vector<uint32_t> mOutputLengths;
@@ -114,6 +153,7 @@ class Engine {
   Logger mLogger;
 
   const std::string kEnginePath;
+  std::map<GraphKey, std::unique_ptr<GraphInstance>> mCudaGraphs;
 };
 
 std::unique_ptr<Engine> load_engine(const Options &options);

src/lib.rs

@@ -97,6 +97,30 @@ mod ffi {
             stream: u64,
             batch_size: u32,
         ) -> Result<()>;
+        /// Capture and replay a CUDA graph for this batch and device-pointer
+        /// set. The graph is cached until `clear_cuda_graph_cache` is called.
+        unsafe fn infer_cuda_graph(
+            self: Pin<&mut Engine>,
+            input_ptrs: *const u64,
+            num_inputs: usize,
+            output_ptrs: *const u64,
+            num_outputs: usize,
+            stream: u64,
+            batch_size: u32,
+        ) -> Result<()>;
+
+        /// Same as `infer_cuda_graph`, but does not synchronize the stream.
+        unsafe fn infer_cuda_graph_async(
+            self: Pin<&mut Engine>,
+            input_ptrs: *const u64,
+            num_inputs: usize,
+            output_ptrs: *const u64,
+            num_outputs: usize,
+            stream: u64,
+            batch_size: u32,
+        ) -> Result<()>;
+
+        fn clear_cuda_graph_cache(self: Pin<&mut Engine>);
     }
 }
 
@@ -193,6 +217,53 @@ impl Engine {
             )
         }
     }
+    /// Capture and replay a CUDA graph for this batch and device-pointer set.
+    ///
+    /// The graph is cached by batch size and device pointers. Keep those
+    /// device buffers alive until `clear_cuda_graph_cache` is called.
+    pub fn infer_cuda_graph(
+        &mut self,
+        inputs: &[CUdeviceptr],
+        outputs: &[CUdeviceptr],
+        stream: CUstream,
+        batch_size: u32,
+    ) -> Result<(), Exception> {
+        unsafe {
+            self.inner.pin_mut().infer_cuda_graph(
+                inputs.as_ptr(),
+                inputs.len(),
+                outputs.as_ptr(),
+                outputs.len(),
+                stream as u64,
+                batch_size,
+            )
+        }
+    }
+
+    /// Same as `infer_cuda_graph`, but does not synchronize the stream.
+    pub fn infer_cuda_graph_async(
+        &mut self,
+        inputs: &[CUdeviceptr],
+        outputs: &[CUdeviceptr],
+        stream: CUstream,
+        batch_size: u32,
+    ) -> Result<(), Exception> {
+        unsafe {
+            self.inner.pin_mut().infer_cuda_graph_async(
+                inputs.as_ptr(),
+                inputs.len(),
+                outputs.as_ptr(),
+                outputs.len(),
+                stream as u64,
+                batch_size,
+            )
+        }
+    }
+
+    /// Discard all cached CUDA graphs before retiring their device buffers.
+    pub fn clear_cuda_graph_cache(&mut self) {
+        self.inner.pin_mut().clear_cuda_graph_cache();
+    }
 
     /// Get input tensor metadata.
     pub fn get_input_dims(&self) -> Vec<TensorInfo> {