1use tract_core::internal::*;
2use tract_core::num_traits::Zero;
3use tract_core::ops::scan::State;
4use tract_core::ops::submodel::TypedModelOpState;
5
6use crate::annotations::*;
7use crate::model::Model;
8use crate::tensor::make_inputs_for_model;
9use std::any::TypeId;
10use std::time::Duration;
11
12pub struct BenchLimits {
13 pub warmup_loops: usize,
14 pub warmup_time: std::time::Duration,
15 pub max_loops: usize,
16 pub max_time: std::time::Duration,
17}
18
19impl Default for BenchLimits {
20 fn default() -> Self {
21 BenchLimits {
22 warmup_loops: 0,
23 warmup_time: Duration::default(),
24 max_loops: 100_000,
25 max_time: std::time::Duration::from_secs(5),
26 }
27 }
28}
29
30impl BenchLimits {
31 pub fn warmup(&self, model: &TypedModel, inputs: &TVec<TValue>) -> TractResult<()> {
32 if self.warmup_time.is_zero() && self.warmup_loops.is_zero() {
33 return Ok(());
34 }
35 let plan = TypedSimplePlan::new(model.clone())?;
36 let mut state = TypedSimpleState::new(Arc::new(plan))?;
37 let mut iters = 0;
38 let max_loops = if self.warmup_loops.is_zero() { usize::MAX } else { self.warmup_loops };
39 let max_time = if self.warmup_time.is_zero() { Duration::MAX } else { self.warmup_time };
40
41 let start_warmup = crate::time::now();
42 debug!("Warming up before profiling...");
43 while iters < max_loops && start_warmup.elapsed() < max_time {
44 state.run(inputs.clone())?;
45 iters += 1;
46 }
47 debug!("Done warming up.");
48 Ok(())
49 }
50}
51
52pub fn profile(
53 model: &TypedModel,
54 bench_limits: &BenchLimits,
55 dg: &mut Annotations,
56 plan_options: &PlanOptions,
57 inputs: &TVec<TValue>,
58 custom_profiler: Option<HashMap<TypeId, Profiler>>,
59 folded: bool,
60) -> TractResult<()> {
61 info!("Running entire network");
62 let mut iters = 0usize;
63 let prefix = tvec!();
64
65 bench_limits.warmup(model, inputs)?;
66
67 let plan = TypedSimplePlan::new_with_options(model.clone(), plan_options)?;
68 let mut state = TypedSimpleState::new(Arc::new(plan))?;
69
70 let start = crate::time::now();
71 let mut time_accounted_by_inner_nodes = Duration::default();
72 while iters < bench_limits.max_loops && start.elapsed() < bench_limits.max_time {
73 rec_profiler(
74 &mut state,
75 dg,
76 inputs,
77 custom_profiler.as_ref(),
78 &prefix,
79 None,
80 &mut time_accounted_by_inner_nodes,
81 folded,
82 )?;
83
84 iters += 1;
85 }
86
87 let entire = start.elapsed() - time_accounted_by_inner_nodes;
88
89 info!("Running {} iterations max. for each node.", bench_limits.max_loops);
90 info!("Running for {} ms max. for each node.", bench_limits.max_time.as_millis());
91
92 let denum = (iters as f32).recip();
93 let entire = entire.mul_f32(denum);
94 for d in dg.tags.values_mut() {
95 if let Some(d) = d.profile.as_mut() {
96 *d = d.mul_f32(denum);
97 }
98
99 if let Some(d) = d.accelerator_profile.as_mut() {
100 *d = d.mul_f32(denum);
101 }
102 }
103 let max = dg.tags.values().filter_map(|t| t.profile).max().unwrap();
104 let sum = dg.tags.values().filter_map(|t| t.profile).sum::<Duration>();
105 let accel_sum = dg.tags.values().filter_map(|t| t.accelerator_profile).sum::<Duration>();
106 dg.profile_summary = Some(ProfileSummary { max, sum, accel_sum, entire, iters });
107 Ok(())
108}
109
110#[cfg(any(target_os = "macos", target_os = "ios"))]
111pub fn profile_metal(
112 model: &TypedModel,
113 bench_limits: &BenchLimits,
114 dg: &mut Annotations,
115 plan_options: &PlanOptions,
116 inputs: &TVec<TValue>,
117) -> TractResult<()> {
118 info!("Running entire network");
119 let mut iters = 0usize;
120 let prefix = tvec!();
121
122 bench_limits.warmup(model, inputs)?;
123
124 let mut plan = TypedSimplePlan::new_with_options(model.clone(), plan_options)?;
125 let state = TypedSimpleState::new_from_inputs(&plan, inputs.clone())?;
126
127 let session_handler =
128 tract_gpu::session_handler::DeviceSessionHandler::from_plan(&plan, &state.session_state.resolved_symbols)?;
129
130 plan = plan.with_session_handler(session_handler);
131
132 let mut state = TypedSimpleState::new(Arc::new(plan))?;
133
134 let mut entire = Duration::default();
135 while iters < bench_limits.max_loops && entire < bench_limits.max_time {
136 entire += rec_profiler_metal(&mut state, dg, inputs, &prefix)?.1;
137
138 iters += 1;
139 }
140
141 info!("Running {} iterations max. for each node.", bench_limits.max_loops);
142 info!("Running for {} ms max. for each node.", bench_limits.max_time.as_millis());
143
144 let denum = (iters as f32).recip();
145 let entire = entire.mul_f32(denum);
146 for d in dg.tags.values_mut() {
147 if let Some(d) = d.profile.as_mut() {
148 *d = d.mul_f32(denum);
149 }
150
151 if let Some(d) = d.accelerator_profile.as_mut() {
152 *d = d.mul_f32(denum);
153 }
154 }
155 let max = dg.tags.values().filter_map(|t| t.profile).max().unwrap();
156 let sum = dg.tags.values().filter_map(|t| t.profile).sum::<Duration>();
157 let accel_sum = dg.tags.values().filter_map(|t| t.accelerator_profile).sum::<Duration>();
158 dg.profile_summary = Some(ProfileSummary { max, sum, accel_sum, entire, iters });
159 Ok(())
160}
161
162#[cfg(any(target_os = "macos", target_os = "ios"))]
163pub fn rec_profiler_metal(
164 state: &mut TypedSimpleState<TypedModel, Arc<TypedSimplePlan<TypedModel>>>,
165 dg: &mut Annotations,
166 inputs: &TVec<TValue>,
167 prefix: &[(usize, String)],
168) -> TractResult<(TVec<TValue>, Duration)> {
169 let profile_start = crate::time::now();
170 let r = state.run_plan_with_eval(
171 inputs.clone(),
172 |session_state, mut node_state, node, input| {
173 let start = crate::time::now();
175 let res = tract_core::plan::eval(
176 session_state,
177 node_state.as_deref_mut(),
178 node,
179 input.clone(),
180 );
181 let elapsed = start.elapsed();
182 let node_id = NodeQId(prefix.into(), node.id);
183 *dg.node_mut(node_id).profile.get_or_insert(Duration::default()) += elapsed;
184
185 res
186 },
187 )?;
188
189 Ok((r, profile_start.elapsed()))
190}
191
192#[allow(clippy::too_many_arguments)]
193pub fn rec_profiler(
194 state: &mut TypedSimpleState<TypedModel, Arc<TypedSimplePlan<TypedModel>>>,
195 dg: &mut Annotations,
196 inputs: &TVec<TValue>,
197 profilers: Option<&HashMap<TypeId, Profiler>>,
198 prefix: &[(usize, String)],
199 multiplier: Option<usize>,
200 time_accounted_by_inner_nodes: &mut Duration,
201 folded: bool,
202) -> TractResult<TVec<TValue>> {
203 let r = state.run_plan_with_eval(
204 inputs.clone(),
205 |session_state, mut node_state, node, input| {
206 let start = crate::time::now();
208 let res = tract_core::plan::eval(
209 session_state,
210 node_state.as_deref_mut(),
211 node,
212 input.clone(),
213 );
214 let elapsed = start.elapsed().mul_f32(multiplier.unwrap_or(1) as _);
215 let node_id = NodeQId(prefix.into(), node.id);
216 *dg.node_mut(node_id).profile.get_or_insert(Duration::default()) += elapsed;
217
218 if !folded {
219 let start = crate::time::now();
220 profile_submodel(
221 node,
222 node_state,
223 input,
224 dg,
225 profilers,
226 prefix,
227 time_accounted_by_inner_nodes,
228 )?;
229 *time_accounted_by_inner_nodes += start.elapsed();
230 }
231
232 let prefix_vec = prefix.to_vec();
234 if !prefix_vec.is_empty() {
235 (1..prefix_vec.len() + 1).map(|idx| prefix_vec[..idx].to_vec()).for_each(
236 |parent_path| {
237 let parent_node = parent_path.last().map(|it| it.0).unwrap();
238 let parent = dg
239 .node_mut(NodeQId(
240 parent_path[..parent_path.len() - 1].into(),
241 parent_node,
242 ))
243 .profile
244 .get_or_insert(Duration::default());
245 *parent -= elapsed.min(*parent);
246 },
247 );
248 }
249 res
250 },
251 )?;
252 Ok(r)
253}
254
255fn profile_submodel(
256 node: &TypedNode,
257 mut node_state: Option<&mut dyn OpState>,
258 input: TVec<TValue>,
259 dg: &mut Annotations,
260 profilers: Option<&HashMap<TypeId, Profiler>>,
261 prefix: &[(usize, String)],
262 time_accounted_by_inner_nodes: &mut Duration,
263) -> TractResult<()> {
264 if let Some(ref mut op_state) = node_state {
265 if let Some(profiler) = profilers.and_then(|it| it.get(&op_state.type_id())) {
266 let mut new_prefix: TVec<_> = prefix.into();
267 new_prefix.push((node.id, "submodel".to_string()));
268
269 let (_, _) =
270 (profiler.func)(*op_state, input, dg, &new_prefix, time_accounted_by_inner_nodes)?;
271 } else if let Some(scan_state) = op_state.downcast_mut::<State>() {
272 let mut new_prefix: TVec<_> = prefix.into();
273 new_prefix.push((node.id, "loop".to_string()));
274
275 let scan_inputs = make_inputs_for_model(scan_state.model_state.model())?;
276 let multi = scan_state.iteration_count(&input);
277
278 rec_profiler(
279 &mut scan_state.model_state,
280 dg,
281 &scan_inputs,
282 None,
283 &new_prefix,
284 Some(multi),
285 time_accounted_by_inner_nodes,
286 false,
287 )?;
288 } else if let Some(typed_model_state) = op_state.downcast_mut::<TypedModelOpState>() {
289 let mut new_prefix: TVec<_> = prefix.into();
290 new_prefix.push((node.id, "submodel".to_string()));
291
292 rec_profiler(
293 typed_model_state,
294 dg,
295 &input,
296 None,
297 &new_prefix,
298 None,
299 time_accounted_by_inner_nodes,
300 false,
301 )?;
302 }
303 }
304
305 Ok(())
306}
307
308type ProfilerFn = fn(
309 &mut dyn OpState,
310 TVec<TValue>,
311 &mut Annotations,
312 &[(usize, String)],
313 &mut Duration,
314) -> TractResult<(TractResult<TVec<TValue>>, Duration)>;
315
316#[derive(Clone)]
317pub struct Profiler {
318 pub func: ProfilerFn,
319 pub name: &'static str,
320}
321
322impl Hash for Profiler {
323 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
324 self.name.hash(state)
325 }
326}
327
328pub fn extract_costs(
329 annotations: &mut Annotations,
330 model: &dyn Model,
331 extra_symbols: &SymbolValues,
332) -> TractResult<()> {
333 fn extract_costs_rec(
334 annotations: &mut Annotations,
335 model: &dyn Model,
336 prefix: &[(usize, String)],
337 multiplier: TDim,
338 extra_symbols: &SymbolValues,
339 ) -> TractResult<()> {
340 if let Some(model) = model.downcast_ref::<TypedModel>() {
341 for node_id in 0..model.nodes().len() {
342 let inputs = model.node_input_facts(node_id)?;
343 let cost = model
344 .node(node_id)
345 .op
346 .cost(&inputs)
347 .with_context(|| format!("costing node {}", model.node(node_id)))?;
348 annotations.node_mut(NodeQId(prefix.into(), node_id)).cost = cost
349 .into_iter()
350 .map(|(k, v)| {
351 let cost = if k.is_compute() { v * &multiplier } else { v };
352 (k, cost.eval(extra_symbols))
353 })
354 .collect();
355
356 let nested_subs = model.nested_models(node_id);
357 let nested_multis = (model as &dyn Model).nested_models_iters(node_id, &inputs);
358 for (name, sub) in nested_subs {
359 let mut prefix: TVec<_> = prefix.into();
360 prefix.push((node_id, name.to_string()));
361 extract_costs_rec(
362 annotations,
363 sub,
364 &prefix,
365 nested_multis.clone().unwrap_or_else(|| 1.into()) * &multiplier,
366 extra_symbols,
367 )?;
368 }
369 }
370 }
371 Ok(())
372 }
373 extract_costs_rec(annotations, model, &[], 1.into(), extra_symbols)
374}