tract_libcli/
lib.rs

1#![allow(clippy::collapsible_if)]
2#[macro_use]
3extern crate log;
4
5pub mod annotations;
6pub mod display_params;
7pub mod draw;
8pub mod export;
9pub mod model;
10pub mod profile;
11pub mod tensor;
12pub mod terminal;
13pub mod time;
14
15use tract_core::internal::*;
16#[allow(unused_imports)]
17use tract_cuda::utils::get_cuda_lib;
18
19pub fn capture_gpu_trace<F>(matches: &clap::ArgMatches, func: F) -> TractResult<()>
20where
21    F: FnOnce() -> TractResult<()>,
22{
23    if matches.is_present("metal-gpu-trace") {
24        #[cfg(any(target_os = "macos", target_os = "ios"))]
25        {
26            let gpu_trace_path =
27                std::path::Path::new(matches.value_of("metal-gpu-trace").unwrap()).to_path_buf();
28            ensure!(gpu_trace_path.is_absolute(), "Metal GPU trace file has to be absolute");
29            ensure!(
30                !gpu_trace_path.exists(),
31                format!("Given Metal GPU trace file {:?} already exists.", gpu_trace_path)
32            );
33
34            log::info!("Capturing Metal GPU trace at : {gpu_trace_path:?}");
35            tract_metal::METAL_STREAM.with_borrow(move |stream| {
36                stream.capture_trace(gpu_trace_path, move |_stream| func())
37            })
38        }
39        #[cfg(not(any(target_os = "macos", target_os = "ios")))]
40        {
41            bail!("`--metal-gpu-trace` present but it is only available on MacOS and iOS")
42        }
43    } else if matches.is_present("cuda-gpu-trace") {
44        if get_cuda_lib().is_none() {
45            bail!("`--cuda-gpu-trace` present but no CUDA insatllation has been found")
46        }
47
48        let _prof = cudarc::driver::safe::Profiler::new()?;
49        func()
50    } else {
51        func()
52    }
53}
tract_libcli/lib.rs

tract_libcli/
lib.rs
