gwr_components/
test_helpers.rs

1// Copyright (c) 2023 Graphcore Ltd. All rights reserved.
2
3use std::cmp::min;
4use std::collections::HashMap;
5use std::rc::Rc;
6
7use gwr_engine::engine::Engine;
8use gwr_engine::port::InPort;
9use gwr_track::entity::{Entity, GetEntity};
10
11use crate::arbiter::Arbiter;
12use crate::arbiter::policy::{Priority, PriorityRoundRobin};
13use crate::flow_controls::limiter::Limiter;
14use crate::source::Source;
15use crate::store::Store;
16use crate::{connect_port, option_box_repeat, rc_limiter};
17
18#[derive(Clone)]
19pub struct ArbiterInputData {
20    pub val: usize,
21    pub count: usize,
22    pub weight: usize,
23    pub priority: Priority,
24}
25
26pub fn check_round_robin(inputs: &[ArbiterInputData], data: &[usize]) {
27    let total_count: usize = inputs.iter().map(|i| i.count).sum();
28    assert_eq!(data.len(), total_count);
29
30    let mut inputs = inputs.to_vec();
31    let mut offset = 0;
32    loop {
33        // Determine the count for each input value in the next window. Note that this
34        // copes with inputs producing the same value and inputs not producing
35        // their full weight in the window.
36        let mut expected_window_counts: HashMap<usize, usize> = HashMap::new();
37        let mut window_length = 0;
38        let max_priority = inputs
39            .iter()
40            .map(|i| {
41                if i.count > 0 {
42                    i.priority
43                } else {
44                    Priority::default()
45                }
46            })
47            .max()
48            .unwrap();
49        for input in &mut inputs {
50            let value_count = min(input.count, input.weight);
51            if input.priority == max_priority && value_count > 0 {
52                expected_window_counts
53                    .entry(input.val)
54                    .and_modify(|e| *e += value_count)
55                    .or_insert(value_count);
56
57                window_length += value_count;
58                input.count -= value_count;
59            }
60        }
61        if window_length == 0 {
62            return;
63        }
64
65        let mut window_counts = HashMap::new();
66        for value in data.iter().skip(offset).take(window_length) {
67            window_counts
68                .entry(*value)
69                .and_modify(|e| *e += 1)
70                .or_insert(1);
71        }
72        assert_eq!(window_counts, expected_window_counts);
73
74        offset += window_length;
75    }
76}
77
78pub fn priority_policy_test_core(engine: &mut Engine, inputs: &[ArbiterInputData]) {
79    let clock = engine.default_clock();
80    let num_inputs = inputs.len();
81    let total_count = inputs.iter().map(|e| e.count).sum();
82    let mut policy = PriorityRoundRobin::new(num_inputs);
83    for (i, input) in inputs.iter().enumerate() {
84        policy = policy.set_priority(i, input.priority);
85    }
86
87    let arbiter = Arbiter::new_and_register(
88        engine,
89        &clock,
90        engine.top(),
91        "arb",
92        num_inputs,
93        Box::new(policy),
94    )
95    .unwrap();
96    let mut sources = Vec::new();
97    for (i, input) in inputs.iter().enumerate() {
98        sources.push(
99            Source::new_and_register(
100                engine,
101                engine.top(),
102                &("source_".to_owned() + &i.to_string()),
103                option_box_repeat!(input.val; input.count),
104            )
105            .unwrap(),
106        );
107    }
108
109    let write_limiter = rc_limiter!(&clock, 1);
110    let store_limiter =
111        Limiter::new_and_register(engine, &clock, engine.top(), "limit_wr", write_limiter).unwrap();
112    let store =
113        Store::new_and_register(engine, &clock, engine.top(), "store", total_count).unwrap();
114    connect_port!(store_limiter, tx => store, rx).unwrap();
115
116    for (i, source) in sources.iter_mut().enumerate() {
117        connect_port!(source, tx => arbiter, rx, i).unwrap();
118    }
119    connect_port!(arbiter, tx => store_limiter, rx).unwrap();
120
121    let port = InPort::new(
122        engine,
123        &clock,
124        &Rc::new(Entity::new(engine.top(), "port")),
125        "test_rx",
126    );
127    store.connect_port_tx(port.state()).unwrap();
128
129    let check_inputs = inputs.to_owned();
130    engine.spawn(async move {
131        let mut store_get = vec![0; total_count];
132        for i in &mut store_get {
133            *i = port.get()?.await;
134        }
135
136        check_round_robin(&check_inputs, &store_get);
137        Ok(())
138    });
139}
gwr_components/test_helpers.rs

gwr_components/
test_helpers.rs
