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