gwr_components/arbiter/policy/
weighted_round_robin.rs

1// Copyright (c) 2025 Graphcore Ltd. All rights reserved.
2
3//! Weighted Round Robin policy
4
5use std::fmt::Write;
6use std::rc::Rc;
7
8use gwr_engine::sim_error;
9use gwr_engine::traits::SimObject;
10use gwr_engine::types::SimError;
11use gwr_track::entity::Entity;
12use gwr_track::trace;
13
14use crate::arbiter::Arbitrate;
15
16pub struct WeightedRoundRobin {
17    candidate: usize,
18    grants: Vec<usize>,
19    weights: Vec<usize>,
20}
21
22impl WeightedRoundRobin {
23    pub fn new(weights: Vec<usize>, num_inputs: usize) -> Result<Self, SimError> {
24        if weights.len() != num_inputs {
25            return sim_error!("The number of weights must be equal to the number of inputs");
26        }
27
28        Ok(Self {
29            candidate: 0,
30            grants: vec![0; weights.len()],
31            weights,
32        })
33    }
34}
35
36impl WeightedRoundRobin {
37    pub fn state_str<T>(&self, input_values: &[Option<T>]) -> String
38    where
39        T: SimObject,
40    {
41        let mut s = String::new();
42        let _ = write!(s, "{}: ", self.candidate);
43        for (i, grant) in self.grants.iter().enumerate() {
44            let req = if input_values[i].is_some() { "r" } else { "-" };
45            let _ = write!(s, "{req}/{grant}/{}, ", self.weights[i]);
46        }
47        s
48    }
49}
50
51impl<T> Arbitrate<T> for WeightedRoundRobin
52where
53    T: SimObject,
54{
55    fn arbitrate(
56        &mut self,
57        entity: &Rc<Entity>,
58        input_values: &mut [Option<T>],
59    ) -> Option<(usize, T)> {
60        trace!(entity ; "wrr: arbitrate {}", self.state_str(input_values));
61
62        let num_inputs = input_values.len();
63        let mut selected_candidate = None;
64        for i in 0..num_inputs {
65            let index = (i + self.candidate) % num_inputs;
66            if input_values[index].is_none() {
67                continue;
68            }
69            if self.weights[index] > self.grants[index] {
70                selected_candidate = Some(index);
71                break;
72            } else if selected_candidate.is_none() {
73                selected_candidate = Some(index);
74            }
75        }
76        if let Some(index) = selected_candidate {
77            if self.weights[index] == self.grants[index] {
78                self.grants[index] = 0;
79            }
80            self.grants[index] += 1;
81
82            let value = input_values[index].take().unwrap();
83            self.candidate = (index + 1) % num_inputs;
84            Some((index, value))
85        } else {
86            None
87        }
88    }
89}