Skip to main content

gwr_models/memory/memory_access_gen/
mod.rs

1// Copyright (c) 2023 Graphcore Ltd. All rights reserved.
2
3//! A data source.
4//!
5//! The data source produces data as defined by the [DataGenerator] that is
6//! provided.
7//!
8//! # Ports
9//!
10//! This component has one port:
11//!  - One [output port](gwr_engine::port::OutPort): `tx`
12
13use std::cell::RefCell;
14use std::collections::HashSet;
15use std::rc::Rc;
16
17use async_trait::async_trait;
18use gwr_components::types::DataGenerator;
19use gwr_components::{connect_tx, port_rx, take_option};
20use gwr_engine::engine::Engine;
21use gwr_engine::executor::Spawner;
22use gwr_engine::port::{InPort, OutPort, PortStateResult};
23use gwr_engine::sim_error;
24use gwr_engine::time::clock::Clock;
25use gwr_engine::traits::{Runnable, SimObject};
26use gwr_engine::types::{SimError, SimResult};
27use gwr_model_builder::{EntityDisplay, EntityGet};
28use gwr_track::Id;
29use gwr_track::entity::Entity;
30use gwr_track::tracker::aka::Aka;
31
32use crate::memory::traits::AccessMemory;
33
34pub mod random;
35pub mod strided;
36
37#[derive(EntityGet, EntityDisplay)]
38pub struct MemoryAccessGen<T>
39where
40    T: SimObject + AccessMemory,
41{
42    entity: Rc<Entity>,
43    spawner: Spawner,
44    data_generator: RefCell<Option<DataGenerator<T>>>,
45    rx: RefCell<Option<InPort<T>>>,
46    tx: RefCell<Option<OutPort<T>>>,
47    payload_bytes_received: Rc<RefCell<usize>>,
48}
49
50impl<T> MemoryAccessGen<T>
51where
52    T: SimObject + AccessMemory,
53{
54    pub fn new_and_register_with_renames(
55        engine: &Engine,
56        clock: &Clock,
57        parent: &Rc<Entity>,
58        name: &str,
59        aka: Option<&Aka>,
60        data_generator: DataGenerator<T>,
61    ) -> Result<Rc<Self>, SimError> {
62        let entity = Rc::new(Entity::new(parent, name));
63        let rx = InPort::new_with_renames(engine, clock, &entity, "rx", aka);
64        let tx = OutPort::new_with_renames(&entity, "tx", aka);
65        let rc_self = Rc::new(Self {
66            entity,
67            spawner: engine.spawner(),
68            data_generator: RefCell::new(Some(data_generator)),
69            rx: RefCell::new(Some(rx)),
70            tx: RefCell::new(Some(tx)),
71            payload_bytes_received: Rc::new(RefCell::new(0)),
72        });
73        engine.register(rc_self.clone());
74        Ok(rc_self)
75    }
76
77    pub fn new_and_register(
78        engine: &Engine,
79        clock: &Clock,
80        parent: &Rc<Entity>,
81        name: &str,
82        data_generator: DataGenerator<T>,
83    ) -> Result<Rc<Self>, SimError> {
84        Self::new_and_register_with_renames(engine, clock, parent, name, None, data_generator)
85    }
86
87    pub fn set_generator(&self, data_generator: Option<DataGenerator<T>>) {
88        *self.data_generator.borrow_mut() = data_generator;
89    }
90
91    pub fn connect_port_tx(&self, port_state: PortStateResult<T>) -> SimResult {
92        connect_tx!(self.tx, connect ; port_state)
93    }
94
95    pub fn port_rx(&self) -> PortStateResult<T> {
96        port_rx!(self.rx, state)
97    }
98
99    pub fn payload_bytes_received(&self) -> usize {
100        *self.payload_bytes_received.borrow()
101    }
102}
103
104#[async_trait(?Send)]
105impl<T> Runnable for MemoryAccessGen<T>
106where
107    T: SimObject + AccessMemory,
108{
109    async fn run(&self) -> SimResult {
110        let data_generator = match self.data_generator.borrow_mut().take() {
111            Some(data_generator) => data_generator,
112            None => return Ok(()),
113        };
114
115        // Use a HashSet so that memory accesses are permitted in any order
116        let expected = Rc::new(RefCell::new(HashSet::new()));
117        let rx = take_option!(self.rx);
118        let mut tx = take_option!(self.tx);
119
120        {
121            let expected = expected.clone();
122            let payload_bytes_received = self.payload_bytes_received.clone();
123            self.spawner.spawn(async move {
124                run_input(rx, expected, payload_bytes_received).await?;
125                Ok(())
126            });
127        }
128
129        for value in data_generator {
130            let id = value.id();
131            if !expected.borrow_mut().insert(id) {
132                return sim_error!("Generator produced duplicate ID {id}");
133            }
134            tx.put(value)?.await;
135        }
136
137        Ok(())
138    }
139}
140
141async fn run_input<T>(
142    mut rx: InPort<T>,
143    expected: Rc<RefCell<HashSet<Id>>>,
144    payload_bytes_received: Rc<RefCell<usize>>,
145) -> SimResult
146where
147    T: SimObject + AccessMemory,
148{
149    loop {
150        let received = rx.get()?.await;
151        let received_id = received.id();
152        if !expected.borrow_mut().remove(&received_id) {
153            return sim_error!("{received_id} received when not expected");
154        }
155        *payload_bytes_received.borrow_mut() += received.access_size_bytes();
156    }
157}