Skip to main content

gwr_models/memory/memory_access_gen/
random.rs

1// Copyright (c) 2025 Graphcore Ltd. All rights reserved.
2
3use std::rc::Rc;
4
5use gwr_engine::types::AccessType;
6use gwr_track::entity::Entity;
7use rand::rngs::StdRng;
8use rand::{Rng, SeedableRng};
9
10use crate::memory::memory_access::MemoryAccess;
11use crate::memory::memory_map::MemoryMap;
12
13/// A Random address access generator.
14///
15/// Will emit memory accesses in the range [base, end)
16pub struct Random {
17    entity: Rc<Entity>,
18
19    // For address -> device mapping
20    memory_map: Rc<MemoryMap>,
21
22    // Configuration
23    src_addr: u64,
24    base_addr: u64,
25    addr_range: u64,
26    alignment_mask: u64,
27    overhead_size_bytes: usize,
28    access_size_bytes: usize,
29    num_to_send: usize,
30
31    // State
32    num_sent: usize,
33    rng: StdRng,
34}
35
36impl Random {
37    #[expect(clippy::too_many_arguments)]
38    #[must_use]
39    pub fn new(
40        parent: &Rc<Entity>,
41        name: &str,
42        seed: u64,
43        memory_map: &Rc<MemoryMap>,
44        src_addr: u64,
45        base_addr: u64,
46        end_addr: u64,
47        alignment_mask: u64,
48        overhead_size_bytes: usize,
49        access_size_bytes: usize,
50        num_to_send: usize,
51    ) -> Self {
52        let rng = StdRng::seed_from_u64(seed);
53        Self {
54            entity: Rc::new(Entity::new(parent, name)),
55            memory_map: memory_map.clone(),
56            src_addr,
57            base_addr,
58            addr_range: end_addr - base_addr,
59            alignment_mask,
60            overhead_size_bytes,
61            access_size_bytes,
62            num_to_send,
63            num_sent: 0,
64            rng,
65        }
66    }
67}
68
69impl Iterator for Random {
70    type Item = MemoryAccess;
71
72    fn next(&mut self) -> Option<Self::Item> {
73        if self.num_sent < self.num_to_send {
74            self.num_sent += 1;
75
76            let dst_addr =
77                ((self.rng.next_u64() % self.addr_range) + self.base_addr) & self.alignment_mask;
78            let (dst_device, _) = self.memory_map.lookup(dst_addr)?;
79            let (src_device, _) = self.memory_map.lookup(self.src_addr)?;
80
81            let access = MemoryAccess::new(
82                &self.entity,
83                AccessType::ReadRequest,
84                self.access_size_bytes,
85                dst_addr,
86                self.src_addr,
87                dst_device,
88                src_device,
89                self.overhead_size_bytes,
90            );
91
92            Some(access)
93        } else {
94            None
95        }
96    }
97}