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