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