1use std::fmt::Debug;
2use std::sync::Mutex;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::mpsc::{Receiver, Sender, channel};
5use std::thread;
6use std::time::Duration;
7
8use crate::internal::*;
9
10#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct LaneTable {
20 taken: Vec<bool>,
21}
22
23impl LaneTable {
24 pub fn new(max_lanes: usize) -> TractResult<LaneTable> {
25 ensure!(max_lanes > 0, "A laned state needs at least one lane");
26 Ok(LaneTable { taken: vec![false; max_lanes] })
27 }
28
29 pub fn max_lanes(&self) -> usize {
32 self.taken.len()
33 }
34
35 pub fn taken(&self) -> usize {
36 self.taken.iter().filter(|t| **t).count()
37 }
38
39 pub fn take(&mut self) -> Option<LaneId> {
43 let lane = self.taken.iter().position(|t| !t)?;
44 self.taken[lane] = true;
45 Some(LaneId(lane))
46 }
47
48 pub fn give_back(&mut self, lane: LaneId) -> TractResult<()> {
50 ensure!(self.is_taken(lane), "Lane {} is not taken, so it can not be given back", lane.0);
51 self.taken[lane.0] = false;
52 Ok(())
53 }
54
55 pub fn is_taken(&self, lane: LaneId) -> bool {
56 self.taken.get(lane.0).copied().unwrap_or(false)
57 }
58
59 pub fn seat(&self, lanes: impl IntoIterator<Item = LaneId>) -> TractResult<Seating> {
63 let lanes: Vec<LaneId> = lanes.into_iter().collect();
64 for lane in &lanes {
65 ensure!(self.is_taken(*lane), "Seating lane {}, which no stream took", lane.0);
66 }
67 Seating::new(self.max_lanes(), lanes)
68 }
69}
70
71#[cfg(test)]
72mod test {
73 use super::*;
74
75 #[test]
76 fn takes_the_lowest_free_lane() -> TractResult<()> {
77 let mut table = LaneTable::new(3)?;
78 assert_eq!(table.take(), Some(LaneId(0)));
79 assert_eq!(table.take(), Some(LaneId(1)));
80 table.give_back(LaneId(0))?;
81 assert_eq!(table.take(), Some(LaneId(0)));
82 assert_eq!(table.taken(), 2);
83 Ok(())
84 }
85
86 #[test]
87 fn runs_out_of_lanes() -> TractResult<()> {
88 let mut table = LaneTable::new(1)?;
89 assert_eq!(table.take(), Some(LaneId(0)));
90 assert_eq!(table.take(), None);
91 Ok(())
92 }
93
94 #[test]
95 fn gives_back_a_taken_lane_only() -> TractResult<()> {
96 let mut table = LaneTable::new(2)?;
97 assert!(table.give_back(LaneId(0)).is_err());
98 table.take();
99 table.give_back(LaneId(0))?;
100 assert!(table.give_back(LaneId(0)).is_err());
101 assert!(table.give_back(LaneId(7)).is_err());
102 Ok(())
103 }
104
105 #[test]
106 fn seats_taken_lanes_in_order() -> TractResult<()> {
107 let mut table = LaneTable::new(4)?;
108 table.take();
109 table.take();
110 table.take();
111 table.give_back(LaneId(1))?;
112 let seating = table.seat([LaneId(2), LaneId(0)])?;
113 assert_eq!(seating.max_lanes(), 4);
114 assert_eq!(seating.occupancy(), 2);
115 assert_eq!(seating.address(0), (Some(0), Some(2)));
116 assert_eq!(seating.address(1), (Some(1), Some(0)));
117 assert!(table.seat([LaneId(0), LaneId(1)]).is_err());
118 assert!(table.seat([LaneId(0), LaneId(0)]).is_err());
119 Ok(())
120 }
121}
122
123crate::declare_knob!(
124 TRACT_MAX_SEATS,
125 usize,
126 256,
127 "Most streams a laned runtime serves in one turn, clamped to the state's lanes."
128);
129
130crate::declare_knob!(
131 TRACT_TURN_LINGER_US,
132 usize,
133 0,
134 "How long a laned runtime waits for more streams once one is ready to run."
135);
136
137#[derive(Clone)]
152pub struct LanedRunnable {
153 shared: Arc<Shared>,
154}
155
156struct Shared {
157 requests: Mutex<Sender<Request>>,
160 inner: Arc<dyn Runnable>,
161 model: Option<Arc<TypedModel>>,
162 plan: Option<Arc<TypedSimplePlan>>,
163 batch: Symbol,
164 max_lanes: usize,
165 counts: Arc<Counts>,
166}
167
168#[derive(Debug, Default)]
171struct Counts {
172 turns: AtomicU64,
173 seats: AtomicU64,
174}
175
176impl LanedRunnable {
177 pub fn wrap(inner: Arc<dyn Runnable>, max_lanes: usize) -> TractResult<LanedRunnable> {
181 let model = inner.typed_model().cloned();
182 let plan = inner.typed_plan().cloned();
183 let mut symbols: Vec<Symbol> = vec![];
184 let mut batch_in: Vec<bool> = vec![];
185 for ix in 0..inner.input_count() {
186 let symbol = batch_symbol(inner.input_fact(ix)?);
187 batch_in.push(symbol.is_some());
188 symbols.extend(symbol);
189 }
190 let mut batch_out: Vec<bool> = vec![];
191 for ix in 0..inner.output_count() {
192 let symbol = batch_symbol(inner.output_fact(ix)?);
193 batch_out.push(symbol.is_some());
194 symbols.extend(symbol);
195 }
196 symbols.sort();
197 symbols.dedup();
198 ensure!(
199 symbols.len() == 1,
200 "A laned model carries one batch symbol on axis 0, this one carries {symbols:?}"
201 );
202 ensure!(batch_out.iter().any(|b| *b), "A laned model must batch one output at least");
203 let batch = symbols.remove(0);
204 let counts = Arc::new(Counts::default());
205 let max_seats = TRACT_MAX_SEATS.get().min(max_lanes);
206 let linger = Duration::from_micros(TRACT_TURN_LINGER_US.get() as u64);
207 let (requests, queue) = channel::<Request>();
208 let (spawned, ready) = channel::<TractResult<()>>();
209 let worker_counts = counts.clone();
210 let worker_inner = inner.clone();
211 thread::Builder::new().name("tract-lanes".into()).spawn(move || {
212 let mut state = match worker_inner.spawn().and_then(|mut state| {
213 let lanes: Vec<LaneId> = (0..max_lanes).map(LaneId).collect();
214 state.reset_lanes(&lanes).context("Preparing a laned model")?;
215 Ok(state)
216 }) {
217 Ok(state) => {
218 let _ = spawned.send(Ok(()));
219 state
220 }
221 Err(e) => {
222 let _ = spawned.send(Err(e));
223 return;
224 }
225 };
226 worker(
227 &mut *state,
228 queue,
229 Table { batch_in, batch_out, max_seats, linger, max_lanes, counts: worker_counts },
230 );
231 })?;
232 ready.recv().map_err(|_| format_err!("The laned worker died spawning the state"))??;
233 Ok(LanedRunnable {
234 shared: Arc::new(Shared {
235 requests: Mutex::new(requests),
236 inner,
237 model,
238 plan,
239 batch,
240 max_lanes,
241 counts,
242 }),
243 })
244 }
245
246 pub fn max_lanes(&self) -> usize {
247 self.shared.max_lanes
248 }
249
250 pub fn inner(&self) -> &Arc<dyn Runnable> {
253 &self.shared.inner
254 }
255
256 pub fn batch_symbol(&self) -> &Symbol {
260 &self.shared.batch
261 }
262
263 pub fn turns_and_seats(&self) -> (u64, u64) {
266 (
267 self.shared.counts.turns.load(Ordering::Relaxed),
268 self.shared.counts.seats.load(Ordering::Relaxed),
269 )
270 }
271
272 fn request(&self) -> TractResult<Sender<Request>> {
273 Ok(self.shared.requests.lock().map_err(|_| format_err!("Poisoned laned sender"))?.clone())
274 }
275}
276
277fn batch_symbol(fact: &TypedFact) -> Option<Symbol> {
281 match fact.shape.dims().first() {
282 Some(TDim::Sym(sym)) => Some(sym.clone()),
283 _ => None,
284 }
285}
286
287impl Debug for LanedRunnable {
288 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
289 write!(f, "LanedRunnable({} lanes)", self.shared.max_lanes)
290 }
291}
292
293impl Runnable for LanedRunnable {
294 fn spawn(&self) -> TractResult<Box<dyn State>> {
295 let requests = self.request()?;
296 let (taken, lane) = channel();
297 requests.send(Request::Take(taken)).map_err(|_| format_err!("The laned worker is gone"))?;
298 let lane = lane.recv().map_err(|_| format_err!("The laned worker dropped a lane"))??;
299 Ok(Box::new(SessionHandle {
300 lease: Arc::new(Lease { lane, requests }),
301 runnable: self.clone(),
302 }))
303 }
304
305 fn typed_plan(&self) -> Option<&Arc<TypedSimplePlan>> {
306 self.shared.plan.as_ref()
307 }
308
309 fn typed_model(&self) -> Option<&Arc<TypedModel>> {
310 self.shared.model.as_ref()
311 }
312}
313
314#[derive(Clone, Debug)]
318pub struct SessionHandle {
319 lease: Arc<Lease>,
320 runnable: LanedRunnable,
321}
322
323#[derive(Debug)]
324struct Lease {
325 lane: LaneId,
326 requests: Sender<Request>,
327}
328
329impl Drop for Lease {
330 fn drop(&mut self) {
331 let _ = self.requests.send(Request::GiveBack(self.lane));
332 }
333}
334
335impl State for SessionHandle {
336 fn run(&mut self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
337 let (done, outputs) = channel();
338 self.lease
339 .requests
340 .send(Request::Turn(Turn { lane: self.lease.lane, inputs, done }))
341 .map_err(|_| format_err!("The laned worker is gone"))?;
342 outputs.recv().map_err(|_| format_err!("The laned worker dropped a turn"))?
343 }
344
345 fn runnable(&self) -> &dyn Runnable {
346 &self.runnable
347 }
348}
349
350enum Request {
351 Take(Sender<TractResult<LaneId>>),
352 GiveBack(LaneId),
353 Turn(Turn),
354}
355
356struct Turn {
357 lane: LaneId,
358 inputs: TVec<TValue>,
359 done: Sender<TractResult<TVec<TValue>>>,
360}
361
362struct Table {
365 batch_in: Vec<bool>,
366 batch_out: Vec<bool>,
367 max_seats: usize,
368 linger: Duration,
369 max_lanes: usize,
370 counts: Arc<Counts>,
371}
372
373fn worker(state: &mut dyn State, queue: Receiver<Request>, table: Table) {
374 let mut lanes = match LaneTable::new(table.max_lanes) {
375 Ok(lanes) => lanes,
376 Err(_) => return,
377 };
378 let mut queued: Vec<Turn> = vec![];
379 loop {
380 if queued.is_empty() {
381 match queue.recv() {
382 Ok(request) => serve(state, &mut lanes, &mut queued, request),
383 Err(_) => return,
384 }
385 if !table.linger.is_zero() {
386 thread::sleep(table.linger);
387 }
388 }
389 while let Ok(request) = queue.try_recv() {
390 serve(state, &mut lanes, &mut queued, request);
391 }
392 let mut seated: Vec<Turn> = vec![];
393 let mut waiting: Vec<Turn> = vec![];
394 for turn in queued.drain(..) {
395 if seated.len() < table.max_seats && !seated.iter().any(|s| s.lane == turn.lane) {
396 seated.push(turn);
397 } else {
398 waiting.push(turn);
399 }
400 }
401 queued = waiting;
402 if seated.is_empty() {
403 continue;
404 }
405 table.counts.turns.fetch_add(1, Ordering::Relaxed);
406 table.counts.seats.fetch_add(seated.len() as u64, Ordering::Relaxed);
407 match run_turn(state, &lanes, &seated, &table) {
408 Ok(per_seat) => {
409 for (turn, outputs) in seated.into_iter().zip(per_seat) {
410 let _ = turn.done.send(Ok(outputs));
411 }
412 }
413 Err(e) => {
414 let e = format!("{e:#}");
415 for turn in seated {
416 let _ = turn.done.send(Err(format_err!("Laned turn failed: {e}")));
417 }
418 }
419 }
420 }
421}
422
423fn serve(state: &mut dyn State, lanes: &mut LaneTable, queued: &mut Vec<Turn>, request: Request) {
427 match request {
428 Request::Take(taken) => {
429 let lane = lanes.take().ok_or_else(|| {
430 format_err!("Every one of the {} lanes is taken", lanes.max_lanes())
431 });
432 let lane = lane.and_then(|lane| {
433 state.reset_lanes(&[lane]).map(|_| lane).inspect_err(|_| {
434 let _ = lanes.give_back(lane);
435 })
436 });
437 let _ = taken.send(lane);
438 }
439 Request::GiveBack(lane) => {
440 let _ = lanes.give_back(lane);
441 }
442 Request::Turn(turn) => queued.push(turn),
443 }
444}
445
446fn run_turn(
447 state: &mut dyn State,
448 lanes: &LaneTable,
449 seated: &[Turn],
450 table: &Table,
451) -> TractResult<Vec<TVec<TValue>>> {
452 let seating = lanes.seat(seated.iter().map(|turn| turn.lane))?;
453 let mut batched: TVec<TValue> = tvec!();
454 for turn in seated {
455 ensure!(
456 turn.inputs.len() == table.batch_in.len(),
457 "A turn feeds {} inputs, the model takes {}",
458 turn.inputs.len(),
459 table.batch_in.len()
460 );
461 }
462 for (ix, is_batched) in table.batch_in.iter().enumerate() {
463 if *is_batched {
464 let rows: TVec<&Tensor> = seated.iter().map(|turn| &*turn.inputs[ix]).collect();
465 for row in &rows {
466 ensure!(
467 row.rank() > 0 && row.shape()[0] == 1,
468 "A stream feeds one row per turn, input {ix} carries {:?}",
469 row.shape()
470 );
471 }
472 batched.push(Tensor::stack_tensors(0, &rows)?.into_tvalue());
473 } else {
474 let shared = &seated[0].inputs[ix];
475 for (seat, turn) in seated.iter().enumerate().skip(1) {
476 ensure!(
477 turn.inputs[ix] == *shared,
478 "Input {ix} carries no batch axis, so one value of it serves the whole \
479 turn, and seats 0 and {seat} feed it different ones"
480 );
481 }
482 batched.push(shared.clone());
483 }
484 }
485 state.seat(seating)?;
486 let outputs = state.run(batched)?;
487 let mut per_seat: Vec<TVec<TValue>> = seated.iter().map(|_| tvec!()).collect();
488 for (ix, output) in outputs.into_iter().enumerate() {
489 if table.batch_out.get(ix).copied().unwrap_or(false) {
490 ensure!(
491 output.shape()[0] == seated.len(),
492 "The turn seats {} streams, output {ix} carries {:?}",
493 seated.len(),
494 output.shape()
495 );
496 for (seat, outputs) in per_seat.iter_mut().enumerate() {
497 outputs.push(output.slice(0, seat, seat + 1)?.into_tvalue());
498 }
499 } else {
500 for outputs in per_seat.iter_mut() {
501 outputs.push(output.clone());
502 }
503 }
504 }
505 Ok(per_seat)
506}
507
508#[cfg(all(test, not(target_family = "wasm")))]
510mod laned_test {
511 use super::*;
512 use crate::ops::math::{add, mul};
513
514 fn doubler(max_lanes: usize) -> TractResult<LanedRunnable> {
517 let mut model = TypedModel::default();
518 let batch = model.symbols.sym("B");
519 let input = model.add_source("input", f32::fact(dims!(batch, 3)))?;
520 let two = model.add_const("two", tensor2(&[[2f32]]))?;
521 let doubled = model.wire_node("doubled", mul(), &[input, two])?;
522 model.select_output_outlets(&doubled)?;
523 let inner = DefaultRuntime.prepare(model)?;
524 LanedRunnable::wrap(inner.into(), max_lanes)
525 }
526
527 fn turn(handle: &mut Box<dyn State>, stream: usize, turn: usize) -> TractResult<()> {
528 let input = tensor2(&[[stream as f32, turn as f32, 1.]]);
529 let output = handle.run(tvec!(input.into_tvalue()))?;
530 assert_eq!(&*output[0], &tensor2(&[[2. * stream as f32, 2. * turn as f32, 2.]]));
531 Ok(())
532 }
533
534 static LINGER: Mutex<()> = Mutex::new(());
537
538 fn spawn_once_free(runnable: &LanedRunnable) -> TractResult<Box<dyn State>> {
541 for _ in 0..100 {
542 if let Ok(handle) = runnable.spawn() {
543 return Ok(handle);
544 }
545 std::thread::sleep(Duration::from_millis(10));
546 }
547 runnable.spawn()
548 }
549
550 #[test]
551 fn one_stream_at_a_time() -> TractResult<()> {
552 let runnable = doubler(2)?;
553 let mut handle = runnable.spawn()?;
554 for t in 0..4 {
555 turn(&mut handle, 0, t)?;
556 }
557 Ok(())
558 }
559
560 #[test]
561 fn every_stream_gets_its_own_row() -> TractResult<()> {
562 let runnable = doubler(8)?;
563 let streams: Vec<_> = (0..8)
564 .map(|stream| {
565 let runnable = runnable.clone();
566 std::thread::spawn(move || -> TractResult<()> {
567 let mut handle = runnable.spawn()?;
568 for t in 0..32 {
569 turn(&mut handle, stream, t)?;
570 }
571 Ok(())
572 })
573 })
574 .collect();
575 for stream in streams {
576 stream.join().unwrap()?;
577 }
578 Ok(())
579 }
580
581 #[test]
582 fn a_turn_seats_the_streams_that_are_ready() -> TractResult<()> {
583 let _linger = LINGER.lock().unwrap_or_else(|e| e.into_inner());
584 TRACT_TURN_LINGER_US.set(20_000);
585 let runnable = doubler(8);
586 TRACT_TURN_LINGER_US.clear();
587 let runnable = runnable?;
588 let streams: Vec<_> = (0..8)
589 .map(|stream| {
590 let runnable = runnable.clone();
591 std::thread::spawn(move || -> TractResult<()> {
592 let mut handle = runnable.spawn()?;
593 for t in 0..4 {
594 turn(&mut handle, stream, t)?;
595 }
596 Ok(())
597 })
598 })
599 .collect();
600 for stream in streams {
601 stream.join().unwrap()?;
602 }
603 let (turns, seats) = runnable.turns_and_seats();
604 assert!(seats > turns, "{seats} seats over {turns} turns, none of them shared");
605 Ok(())
606 }
607
608 fn biased(max_lanes: usize) -> TractResult<LanedRunnable> {
611 let mut model = TypedModel::default();
612 let batch = model.symbols.sym("B");
613 let input = model.add_source("input", f32::fact(dims!(batch, 3)))?;
614 let bias = model.add_source("bias", f32::fact(dims!(1, 1)))?;
615 let two = model.add_const("two", tensor2(&[[2f32]]))?;
616 let doubled = model.wire_node("doubled", mul(), &[input, two])?;
617 let biased = model.wire_node("biased", add(), &[doubled[0], bias])?;
618 model.select_output_outlets(&biased)?;
619 let inner = DefaultRuntime.prepare(model)?;
620 LanedRunnable::wrap(inner.into(), max_lanes)
621 }
622
623 fn biased_turns(runnable: &LanedRunnable, biases: &[f32]) -> TractResult<Vec<TractResult<()>>> {
627 let handles: Vec<Box<dyn State>> =
628 biases.iter().map(|_| runnable.spawn()).collect::<TractResult<_>>()?;
629 let streams: Vec<_> = handles
630 .into_iter()
631 .zip(biases.iter().copied())
632 .map(|(mut handle, bias)| {
633 std::thread::spawn(move || -> TractResult<()> {
634 handle.run(tvec!(
635 tensor2(&[[1f32, 2., 3.]]).into_tvalue(),
636 tensor2(&[[bias]]).into_tvalue()
637 ))?;
638 Ok(())
639 })
640 })
641 .collect();
642 Ok(streams.into_iter().map(|stream| stream.join().unwrap()).collect())
643 }
644
645 #[test]
646 fn seats_agreeing_on_a_shared_input_share_a_turn() -> TractResult<()> {
647 let _linger = LINGER.lock().unwrap_or_else(|e| e.into_inner());
648 TRACT_TURN_LINGER_US.set(100_000);
649 let runnable = biased(2);
650 TRACT_TURN_LINGER_US.clear();
651 let runnable = runnable?;
652 let served = biased_turns(&runnable, &[7., 7.])?;
653 assert!(served.iter().all(|s| s.is_ok()), "{served:?}");
654 assert_eq!(runnable.turns_and_seats(), (1, 2));
655 Ok(())
656 }
657
658 #[test]
659 fn seats_disagreeing_on_a_shared_input_fail_the_turn() -> TractResult<()> {
660 let _linger = LINGER.lock().unwrap_or_else(|e| e.into_inner());
661 TRACT_TURN_LINGER_US.set(100_000);
662 let runnable = biased(2);
663 TRACT_TURN_LINGER_US.clear();
664 let runnable = runnable?;
665 let served = biased_turns(&runnable, &[7., 8.])?;
666 assert_eq!(runnable.turns_and_seats(), (1, 2));
667 for stream in &served {
668 let error = format!("{:#}", stream.as_ref().unwrap_err());
669 assert!(error.contains("seats 0 and 1 feed it different ones"), "{error}");
670 }
671 Ok(())
672 }
673
674 #[test]
675 fn a_dropped_stream_gives_its_lane_back() -> TractResult<()> {
676 let runnable = doubler(1)?;
677 let mut handle = runnable.spawn()?;
678 turn(&mut handle, 0, 0)?;
679 assert!(runnable.spawn().is_err());
680 let clone = dyn_clone::clone_box(&*handle);
681 drop(handle);
682 assert!(runnable.spawn().is_err());
683 drop(clone);
684 let mut handle = spawn_once_free(&runnable)?;
685 turn(&mut handle, 1, 0)?;
686 Ok(())
687 }
688}