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(any(target_os = "macos", target_os = "ios"))]
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(any(target_os = "macos", target_os = "ios")))]
47        {
48            bail!("`--metal-gpu-trace` present but it is only available on MacOS and iOS")
49        }
50    } else if matches.get_flag("cuda-gpu-trace") {
51        #[cfg(all(any(target_os = "linux", target_os = "windows"), feature = "cuda"))]
52        {
53            ensure_cuda_runtime_dependencies(
54                "`--cuda-gpu-trace` present but no CUDA installation has been found",
55            )?;
56            let _prof = cudarc::driver::safe::Profiler::new()?;
57            func()
58        }
59        #[cfg(not(all(any(target_os = "linux", target_os = "windows"), feature = "cuda")))]
60        {
61            bail!(
62                "`--cuda-gpu-trace` present but tract was not built with CUDA support \
63                 (re-build on Linux/Windows with one of the cuda-XXXXX features)"
64            )
65        }
66    } else {
67        func()
68    }
69}
tract_libcli/lib.rs

tract_libcli/
lib.rs
