gwr_components/arbiter/policy/
priority_round_robin.rs

1// Copyright (c) 2025 Graphcore Ltd. All rights reserved.
2
3//! Priority Round Robin arbitration policy
4
5use std::collections::BTreeMap;
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;
12
13use crate::arbiter::Arbitrate;
14
15#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
16pub enum Priority {
17    #[default]
18    Low = 0,
19    Medium,
20    High,
21}
22
23struct PriorityLevel {
24    current_candidate_index: usize,
25    candidates: Vec<usize>,
26}
27
28pub struct PriorityRoundRobin<P>
29where
30    P: Copy + Default + Ord,
31{
32    priority_vec: Vec<P>,
33    priority_map: Option<BTreeMap<P, PriorityLevel>>,
34}
35
36impl<P> PriorityRoundRobin<P>
37where
38    P: Copy + Default + Ord,
39{
40    #[must_use]
41    pub fn new(num_inputs: usize) -> Self {
42        let default_priority = P::default();
43        let priority_vec = vec![default_priority; num_inputs];
44
45        Self {
46            priority_vec,
47            priority_map: None,
48        }
49    }
50
51    pub fn from_priorities(priority_vec: Vec<P>, num_inputs: usize) -> Result<Self, SimError> {
52        if priority_vec.len() != num_inputs {
53            return sim_error!("The number of priorities must be equal to the number of inputs");
54        }
55
56        Ok(Self {
57            priority_vec,
58            priority_map: None,
59        })
60    }
61
62    pub fn set_priority(mut self, port_index: usize, new_priority: P) -> Self {
63        self.priority_vec[port_index] = new_priority;
64        self
65    }
66
67    fn create_map(&mut self) {
68        let mut priority_map = BTreeMap::new();
69        for (i, priority) in self.priority_vec.iter().enumerate() {
70            priority_map
71                .entry(*priority)
72                .and_modify(|e: &mut PriorityLevel| {
73                    e.candidates.push(i);
74                })
75                .or_insert(PriorityLevel {
76                    current_candidate_index: 0,
77                    candidates: Vec::from([i]),
78                });
79        }
80        self.priority_map = Some(priority_map);
81    }
82}
83
84impl<T, P> Arbitrate<T> for PriorityRoundRobin<P>
85where
86    T: SimObject,
87    P: Copy + Default + Ord,
88{
89    fn arbitrate(
90        &mut self,
91        _entity: &Rc<Entity>,
92        input_values: &mut [Option<T>],
93    ) -> Option<(usize, T)> {
94        if self.priority_map.is_none() {
95            self.create_map();
96        }
97        let priority_map = self.priority_map.as_mut().unwrap();
98        for (_, priority_level) in priority_map.iter_mut().rev() {
99            let priority_vec = &priority_level.candidates;
100            let num_inputs = priority_vec.len();
101            for i in 0..num_inputs {
102                let priority_index = (i + priority_level.current_candidate_index) % num_inputs;
103                let input_index = priority_vec[priority_index];
104                if let Some(value) = input_values[input_index].take() {
105                    priority_level.current_candidate_index = (priority_index + 1) % num_inputs;
106                    return Some((input_index, value));
107                }
108            }
109        }
110        None
111    }
112}