From f89620838aced1d69904e9a6edfc6a0275244b88 Mon Sep 17 00:00:00 2001 From: sj0618 Date: Sat, 18 Jul 2026 01:03:58 +0900 Subject: [PATCH] feat: add cached CUDA graph inference --- README.md | 9 +++- examples/benchmark.rs | 27 ++++++++--- src/engine.cpp | 107 ++++++++++++++++++++++++++++++++++++++++++ src/engine.h | 40 ++++++++++++++++ src/lib.rs | 71 ++++++++++++++++++++++++++++ 5 files changed, 247 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 920e62f..3f09e62 100644 --- a/README.md +++ b/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 diff --git a/examples/benchmark.rs b/examples/benchmark.rs index 0d9d920..67e127c 100644 --- a/examples/benchmark.rs +++ b/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; diff --git a/src/engine.cpp b/src/engine.cpp index 40c37c9..b0e9f3f 100644 --- a/src/engine.cpp +++ b/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(mMinBatchSize) || + batch_size > static_cast(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(); + 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(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(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 Engine::get_input_dims() const { rust::Vec result; for (const auto &meta : mTensorMetadata) { diff --git a/src/engine.h b/src/engine.h index 142a832..1dea789 100644 --- a/src/engine.h +++ b/src/engine.h @@ -3,11 +3,13 @@ #include "NvInfer.h" #include #include +#include #include #include #include #include #include +#include #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 get_input_dims() const; rust::Vec 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 input_ptrs; + std::vector 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 mTensorMetadata; std::vector mOutputLengths; @@ -114,6 +153,7 @@ class Engine { Logger mLogger; const std::string kEnginePath; + std::map> mCudaGraphs; }; std::unique_ptr load_engine(const Options &options); diff --git a/src/lib.rs b/src/lib.rs index b278115..34a250d 100644 --- a/src/lib.rs +++ b/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 {