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;
11// None of it works without threads.
12#[cfg(not(target_family = "wasm"))]
13pub mod profile_autobatch;
14pub mod tensor;
15pub mod terminal;
16pub mod time;
17
18use tract_core::internal::*;
19#[allow(unused_imports)]
20#[cfg(all(any(target_os = "linux", target_os = "windows"), feature = "cuda"))]
21use tract_cuda::utils::ensure_cuda_runtime_dependencies;
22
23pub fn capture_gpu_trace<F>(matches: &clap::ArgMatches, func: F) -> TractResult<()>
24where
25    F: FnOnce() -> TractResult<()>,
26{
27    if matches.contains_id("metal-gpu-trace")
28        && matches.get_one::<String>("metal-gpu-trace").is_some()
29    {
30        #[cfg(all(any(target_os = "macos", target_os = "ios"), feature = "metal"))]
31        {
32            let gpu_trace_path =
33                std::path::Path::new(matches.get_one::<String>("metal-gpu-trace").unwrap())
34                    .to_path_buf();
35            ensure!(gpu_trace_path.is_absolute(), "Metal GPU trace file has to be absolute");
36            ensure!(
37                !gpu_trace_path.exists(),
38                format!("Given Metal GPU trace file {:?} already exists.", gpu_trace_path)
39            );
40
41            log::info!("Capturing Metal GPU trace at : {gpu_trace_path:?}");
42            tract_metal::with_metal_stream(move |stream| {
43                stream.capture_trace(gpu_trace_path, move |_stream| func())
44            })
45        }
46        #[cfg(not(all(any(target_os = "macos", target_os = "ios"), feature = "metal")))]
47        {
48            bail!(
49                "`--metal-gpu-trace` present but tract was not built with Metal support \
50                 (re-build on MacOS/iOS with the metal feature)"
51            )
52        }
53    } else if matches.get_flag("cuda-gpu-trace") {
54        #[cfg(all(any(target_os = "linux", target_os = "windows"), feature = "cuda"))]
55        {
56            ensure_cuda_runtime_dependencies(
57                "`--cuda-gpu-trace` present but no CUDA installation has been found",
58            )?;
59            let _prof = cudarc::driver::safe::Profiler::new()?;
60            func()
61        }
62        #[cfg(not(all(any(target_os = "linux", target_os = "windows"), feature = "cuda")))]
63        {
64            bail!(
65                "`--cuda-gpu-trace` present but tract was not built with CUDA support \
66                 (re-build on Linux/Windows with one of the cuda-XXXXX features)"
67            )
68        }
69    } else {
70        func()
71    }
72}
tract_libcli/lib.rs

tract_libcli/
lib.rs
