sim_fabric/
frame_gen.rs

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