style: format Rust workspace
This commit is contained in:
@@ -54,30 +54,21 @@ impl CudaGraph {
|
||||
// made by THIS thread invalidate the capture. With GLOBAL mode, TP rank
|
||||
// threads capturing concurrently would poison each other's captures.
|
||||
error::check(unsafe {
|
||||
ffi::cudaStreamBeginCapture(
|
||||
stream.as_raw(),
|
||||
ffi::CUDA_STREAM_CAPTURE_MODE_THREAD_LOCAL,
|
||||
)
|
||||
ffi::cudaStreamBeginCapture(stream.as_raw(), ffi::CUDA_STREAM_CAPTURE_MODE_THREAD_LOCAL)
|
||||
})
|
||||
}
|
||||
|
||||
/// End capture and instantiate the executable graph.
|
||||
pub fn end_capture(&mut self, stream: &CudaStream) -> Result<()> {
|
||||
error::check(unsafe {
|
||||
ffi::cudaStreamEndCapture(stream.as_raw(), &mut self.graph)
|
||||
})?;
|
||||
error::check(unsafe {
|
||||
ffi::cudaGraphInstantiate(&mut self.exec, self.graph, 0)
|
||||
})
|
||||
error::check(unsafe { ffi::cudaStreamEndCapture(stream.as_raw(), &mut self.graph) })?;
|
||||
error::check(unsafe { ffi::cudaGraphInstantiate(&mut self.exec, self.graph, 0) })
|
||||
}
|
||||
|
||||
/// Replay the captured graph on `stream`.
|
||||
/// Panics if no graph has been captured yet.
|
||||
pub fn launch(&self, stream: &CudaStream) -> Result<()> {
|
||||
assert!(self.is_ready(), "CudaGraph::launch called before capture");
|
||||
error::check(unsafe {
|
||||
ffi::cudaGraphLaunch(self.exec, stream.as_raw())
|
||||
})
|
||||
error::check(unsafe { ffi::cudaGraphLaunch(self.exec, stream.as_raw()) })
|
||||
}
|
||||
|
||||
fn destroy_inner(&mut self) {
|
||||
|
||||
Reference in New Issue
Block a user