sim_fabric/
frame_gen.rs

1// Copyright (c) 2025 Graphcore Ltd. All rights reserved.
2
3use std::fmt;
4use std::rc::Rc;
5
6use gwr_models::data_frame::DataFrame;
7use gwr_models::fabric::FabricConfig;
8use gwr_track::entity::Entity;
9use rand::SeedableRng;
10use rand::seq::{IteratorRandom, SliceRandom};
11use rand_xoshiro::Xoshiro256PlusPlus;
12use serde::Serialize;
13
14#[derive(clap::ValueEnum, Clone, Copy, Default, Debug, Serialize, PartialEq)]
15#[serde(rename_all = "kebab-case")]
16pub enum TrafficPattern {
17    /// All sources will send to one dest chosen at random
18    #[default]
19    AllToOne,
20
21    /// All sources will send to random (valid) destinations
22    Random,
23
24    /// All sources will be allocated a fixed (random) destination
25    AllToAllFixed,
26
27    /// All sources will send to to all other destinations in sequence
28    AllToAllSeq,
29
30    /// All sources will send to to all other destinations in random (but
31    /// repeated) order
32    AllToAllRandom,
33}
34
35impl fmt::Display for TrafficPattern {
36    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
37        write!(f, "{self:?}")
38    }
39}
40
41/// A frame Generator that can be used by the `Source` to produce frames on
42/// the fly.
43///
44/// This allows each frame being created to be unique which aids debug of the
45/// system.
46pub struct FrameGen {
47    pub entity: Rc<Entity>,
48    config: Rc<FabricConfig>,
49    source_index: usize,
50    dest_index: usize,
51    traffic_pattern: TrafficPattern,
52    overhead_size_bytes: usize,
53    payload_size_bytes: usize,
54    num_send_frames: usize,
55    num_sent_frames: usize,
56    rng: Xoshiro256PlusPlus,
57    next_dests: Vec<usize>,
58}
59
60impl FrameGen {
61    #[expect(clippy::too_many_arguments)]
62    #[must_use]
63    pub fn new(
64        parent: &Rc<Entity>,
65        config: Rc<FabricConfig>,
66        source_index: usize,
67        initial_dest_index: usize,
68        traffic_pattern: TrafficPattern,
69        overhead_size_bytes: usize,
70        payload_size_bytes: usize,
71        num_send_frames: usize,
72        seed: u64,
73    ) -> Self {
74        // Create a local RNG which is different per source
75        let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed ^ (source_index as u64));
76        let num_ports = config.num_ports();
77
78        let dest_index = match traffic_pattern {
79            TrafficPattern::Random => (0..num_ports).choose(&mut rng).unwrap(),
80            TrafficPattern::AllToOne
81            | TrafficPattern::AllToAllFixed
82            | TrafficPattern::AllToAllSeq
83            | TrafficPattern::AllToAllRandom => initial_dest_index,
84        };
85
86        let mut next_dests: Vec<usize> = (0..num_ports).collect();
87        next_dests.shuffle(&mut rng);
88
89        Self {
90            entity: Rc::new(Entity::new(parent, format!("gen{source_index}").as_str())),
91            config,
92            source_index,
93            dest_index,
94            traffic_pattern,
95            overhead_size_bytes,
96            payload_size_bytes,
97            num_send_frames,
98            num_sent_frames: 0,
99            rng,
100            next_dests,
101        }
102    }
103
104    #[must_use]
105    fn next_dest(&mut self) -> usize {
106        let num_ports = self.config.max_num_ports();
107        match self.traffic_pattern {
108            TrafficPattern::Random => (0..num_ports).choose(&mut self.rng).unwrap(),
109            TrafficPattern::AllToOne => self.dest_index,
110            TrafficPattern::AllToAllFixed => self.dest_index,
111            TrafficPattern::AllToAllSeq => (self.dest_index + 1) % num_ports,
112            TrafficPattern::AllToAllRandom => {
113                let dest = self.next_dests.pop().unwrap();
114                if self.next_dests.is_empty() {
115                    self.next_dests = (0..num_ports).collect();
116                    self.next_dests.shuffle(&mut self.rng);
117                }
118                dest
119            }
120        }
121    }
122}
123
124impl Iterator for FrameGen {
125    type Item = DataFrame;
126    fn next(&mut self) -> Option<Self::Item> {
127        // If this can only send to self then there is nothing to do
128        let no_valid_packets = (self.dest_index == self.source_index)
129            && (self.traffic_pattern == TrafficPattern::AllToOne
130                || self.traffic_pattern == TrafficPattern::AllToAllFixed);
131
132        if no_valid_packets {
133            return None;
134        }
135
136        if self.num_sent_frames < self.num_send_frames {
137            while self.dest_index == self.source_index {
138                self.dest_index = self.next_dest();
139            }
140
141            let label = (self.num_sent_frames | (self.source_index << 32)) as u64;
142            self.num_sent_frames += 1;
143
144            // Send to the correct `dest`, but set `src` to a unique value to aid debug
145            // (frame count).
146            let dest = self.config.port_indices()[self.dest_index];
147            let frame = Some(
148                DataFrame::new(
149                    &self.entity,
150                    self.overhead_size_bytes,
151                    self.payload_size_bytes,
152                )
153                .set_dest(dest as u64)
154                .set_label(label),
155            );
156
157            self.dest_index = self.next_dest();
158            frame
159        } else {
160            None
161        }
162    }
163}