gwr_models/processing_element/
mod.rs

1// Copyright (c) 2026 Graphcore Ltd. All rights reserved.
2
3//! A Processing Element (PE) for a simulation.
4//!
5//! The PE performs computations defined by a timetable.
6//!
7//! The PE comprises:
8//!  - Load/Store
9//!  - Internal Buffers
10//!  - Compute
11//!
12//! Identifies all compute nodes that can execute
13//! because their dependencies are satisfied (or they have no dependencies).
14
15//! # Ports
16//!
17//! Each PE has:
18//!  - One [input port](gwr_engine::port::InPort): `rx`
19//!  - One [output port](gwr_engine::port::OutPort): `tx`
20//!
21//! that are managed by the `LoadStoreUnit`
22
23use std::cell::RefCell;
24use std::fmt::{self, Display};
25use std::rc::Rc;
26
27use async_trait::async_trait;
28use gwr_engine::engine::Engine;
29use gwr_engine::executor::Spawner;
30use gwr_engine::port::PortStateResult;
31use gwr_engine::time::clock::{Clock, phase};
32use gwr_engine::traits::Runnable;
33use gwr_engine::types::{AccessType, SimError, SimResult};
34use gwr_model_builder::{EntityDisplay, EntityGet};
35use gwr_track::debug;
36use gwr_track::entity::{Entity, EntityGroup, EntityLane};
37use gwr_track::tracker::aka::Aka;
38use serde::{Deserialize, Serialize};
39
40use crate::log_stats;
41use crate::memory::memory_access::MemoryAccess;
42use crate::memory::memory_map::{DeviceId, MemoryMap};
43use crate::processing_element::dispatch::Dispatch;
44use crate::processing_element::flop_monitor::FlopMonitor;
45use crate::processing_element::load_store_unit::LoadStoreUnit;
46use crate::processing_element::operators::TensorView;
47use crate::processing_element::task::{ComputeTaskConfig, Task};
48
49pub mod dispatch;
50mod flop_monitor;
51mod load_store_unit;
52pub mod operators;
53pub mod task;
54
55#[derive(Clone, Copy, Eq, Hash, PartialEq)]
56pub enum MachineOp {
57    Add,
58    Compare,
59    Mul,
60}
61
62#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Deserialize, Serialize)]
63#[serde(deny_unknown_fields)]
64pub struct MachineOpCounts {
65    #[serde(default)]
66    pub adds: usize,
67    #[serde(default)]
68    pub compares: usize,
69    #[serde(default)]
70    pub muls: usize,
71}
72
73impl MachineOpCounts {
74    #[must_use]
75    pub fn total(&self) -> usize {
76        self.adds + self.compares + self.muls
77    }
78
79    pub fn add_assign(&mut self, other: Self) {
80        self.adds += other.adds;
81        self.compares += other.compares;
82        self.muls += other.muls;
83    }
84}
85
86pub struct ProcessingElementStatsDisplay {
87    prefix: String,
88    time_now_ns: f64,
89    machine_ops: MachineOpCounts,
90}
91
92impl ProcessingElementStatsDisplay {
93    #[must_use]
94    pub fn new(prefix: impl Into<String>, time_now_ns: f64, machine_ops: MachineOpCounts) -> Self {
95        Self {
96            prefix: prefix.into(),
97            time_now_ns,
98            machine_ops,
99        }
100    }
101}
102
103impl Display for ProcessingElementStatsDisplay {
104    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
105        let total_flops = self.machine_ops.total();
106        let time_now_s = self.time_now_ns / 1e9;
107        let total_gflops = total_flops as f64 / 1e9;
108        let average_gflops_per_second = if time_now_s == 0.0 {
109            0.0
110        } else {
111            total_gflops / time_now_s
112        };
113
114        writeln!(f, "{}:", self.prefix)?;
115        writeln!(
116            f,
117            "  FLOPs: {total_flops}, {total_gflops:.2} GFLOPs, {average_gflops_per_second:.2} GFLOP/s"
118        )?;
119        write!(
120            f,
121            "  Machine ops: {} total, {} add, {} mul, {} compare",
122            self.machine_ops.total(),
123            self.machine_ops.adds,
124            self.machine_ops.muls,
125            self.machine_ops.compares
126        )
127    }
128}
129
130type Dispatcher = Rc<dyn Dispatch>;
131
132pub struct ProcessingElementConfig {
133    /// The number of outstanding requests can the LSU handle at once
134    pub num_active_requests: usize,
135
136    /// The maximum number of bytes in each memory access
137    pub lsu_access_bytes: usize,
138
139    /// The number of bytes of protocol overhead for each memory transaction
140    pub overhead_size_bytes: usize,
141
142    /// The total number of local SRAM storage bytes available to the PE
143    pub sram_bytes: usize,
144
145    /// Number of add operations per tick
146    pub adds_per_tick: f64,
147
148    /// Number of multiply operations per tick
149    pub muls_per_tick: f64,
150
151    /// Number of compare operations per tick
152    pub compares_per_tick: f64,
153}
154
155pub struct ComputeCapabilities {
156    adds_per_tick: f64,
157    muls_per_tick: f64,
158    compares_per_tick: f64,
159    sram_bytes: usize,
160}
161
162impl ComputeCapabilities {
163    #[must_use]
164    pub fn ops_per_tick(&self, op: MachineOp) -> f64 {
165        match op {
166            MachineOp::Add => self.adds_per_tick,
167            MachineOp::Compare => self.compares_per_tick,
168            MachineOp::Mul => self.muls_per_tick,
169        }
170    }
171
172    pub fn cycles_for_ops(&self, num_ops: usize, op: MachineOp) -> Result<usize, SimError> {
173        if num_ops == 0 {
174            return Ok(0);
175        }
176
177        let ops_per_tick = self.ops_per_tick(op);
178        if !ops_per_tick.is_finite() || ops_per_tick <= 0.0 {
179            return Err(SimError(format!(
180                "invalid compute throughput {ops_per_tick} ops/tick"
181            )));
182        }
183
184        Ok(((num_ops as f64) / ops_per_tick).ceil() as usize)
185    }
186}
187
188#[derive(Default)]
189struct ProcessingElementStats {
190    machine_ops: MachineOpCounts,
191}
192
193struct Lane {
194    lane: EntityLane,
195    active: bool,
196}
197
198pub(crate) struct ActivityLanes {
199    entity: Rc<Entity>,
200    track_name: String,
201    lanes: Vec<Lane>,
202}
203
204impl ActivityLanes {
205    fn new(entity: Rc<Entity>, track_name: &str) -> Self {
206        Self {
207            entity,
208            track_name: track_name.to_string(),
209            lanes: Vec::new(),
210        }
211    }
212
213    fn begin_in_group(
214        lanes: &Rc<RefCell<Self>>,
215        name: &str,
216        group: &Rc<EntityGroup>,
217    ) -> ActivityLaneGuard {
218        let mut lanes_ref = lanes.borrow_mut();
219        let lane_idx = match lanes_ref.lanes.iter().position(|lane| !lane.active) {
220            Some(lane_idx) => lane_idx,
221            None => lanes_ref.add_new_lane(),
222        };
223
224        let lane = &mut lanes_ref.lanes[lane_idx];
225        lane.lane.begin_in_group(name, group);
226        lane.active = true;
227
228        ActivityLaneGuard {
229            lanes: lanes.clone(),
230            lane_idx,
231            active: true,
232        }
233    }
234
235    fn add_new_lane(&mut self) -> usize {
236        let lane_idx = self.lanes.len();
237        let lane = EntityLane::new(&self.entity, &format!("{}::{lane_idx}", self.track_name));
238        self.lanes.push(Lane {
239            lane,
240            active: false,
241        });
242        lane_idx
243    }
244
245    fn end(&mut self, lane_idx: usize) {
246        let lane = &mut self.lanes[lane_idx];
247        lane.lane.end();
248        lane.active = false;
249    }
250}
251
252struct ActivityLaneGuard {
253    lanes: Rc<RefCell<ActivityLanes>>,
254    lane_idx: usize,
255    active: bool,
256}
257
258impl Drop for ActivityLaneGuard {
259    fn drop(&mut self) {
260        if self.active {
261            self.lanes.borrow_mut().end(self.lane_idx);
262            self.active = false;
263        }
264    }
265}
266
267struct ProcessingElementActivityLanes {
268    entity: Rc<Entity>,
269    compute: Rc<RefCell<ActivityLanes>>,
270    lsu_read: Rc<RefCell<ActivityLanes>>,
271    lsu_write: Rc<RefCell<ActivityLanes>>,
272}
273
274impl ProcessingElementActivityLanes {
275    fn new(entity: Rc<Entity>) -> Self {
276        Self {
277            entity: entity.clone(),
278            compute: Rc::new(RefCell::new(ActivityLanes::new(
279                entity.clone(),
280                "lane::compute",
281            ))),
282            lsu_read: Rc::new(RefCell::new(ActivityLanes::new(
283                entity.clone(),
284                "lane::lsu_read",
285            ))),
286            lsu_write: Rc::new(RefCell::new(ActivityLanes::new(entity, "lane::lsu_write"))),
287        }
288    }
289
290    fn create_group(&self, name: &str) -> Rc<EntityGroup> {
291        Rc::new(EntityGroup::new(&self.entity, name))
292    }
293}
294
295#[derive(EntityGet, EntityDisplay)]
296pub struct ProcessingElement {
297    entity: Rc<Entity>,
298    lsu: Rc<LoadStoreUnit>,
299    clock: Clock,
300    spawner: Spawner,
301
302    compute_capabilities: Rc<ComputeCapabilities>,
303    stats: Rc<RefCell<ProcessingElementStats>>,
304    activity_lanes: Rc<ProcessingElementActivityLanes>,
305    dispatcher: RefCell<Option<Dispatcher>>,
306    flop_monitor: Option<Rc<FlopMonitor>>,
307}
308
309impl ProcessingElement {
310    #[expect(clippy::too_many_arguments)]
311    pub fn new_and_register_with_renames(
312        engine: &Engine,
313        clock: &Clock,
314        parent: &Rc<Entity>,
315        name: &str,
316        aka: Option<&Aka>,
317        memory_map: &Rc<MemoryMap>,
318        pe_config: &ProcessingElementConfig,
319        device_id: DeviceId,
320    ) -> Result<Rc<Self>, SimError> {
321        let entity = Rc::new(Entity::new(parent, name));
322
323        let lsu = LoadStoreUnit::new_and_register(
324            engine, clock, &entity, aka, pe_config, memory_map, device_id,
325        )?;
326        let monitor_window_size = entity.tracker.monitoring_window_size_for(entity.id);
327        let flop_monitor = monitor_window_size.map(|window_size_ticks| {
328            FlopMonitor::new_and_register(engine, &entity, clock, window_size_ticks)
329        });
330
331        let rc_self = Rc::new(Self {
332            entity: entity.clone(),
333            lsu,
334            clock: clock.clone(),
335            spawner: engine.spawner(),
336
337            compute_capabilities: Rc::new(ComputeCapabilities {
338                adds_per_tick: pe_config.adds_per_tick,
339                muls_per_tick: pe_config.muls_per_tick,
340                compares_per_tick: pe_config.compares_per_tick,
341                sram_bytes: pe_config.sram_bytes,
342            }),
343            stats: Rc::new(RefCell::new(ProcessingElementStats::default())),
344            activity_lanes: Rc::new(ProcessingElementActivityLanes::new(entity.clone())),
345
346            dispatcher: RefCell::new(None),
347            flop_monitor,
348        });
349        engine.register(rc_self.clone());
350        Ok(rc_self)
351    }
352
353    pub fn new_and_register(
354        engine: &Engine,
355        clock: &Clock,
356        parent: &Rc<Entity>,
357        name: &str,
358        memory_map: &Rc<MemoryMap>,
359        pe_config: &ProcessingElementConfig,
360        device_id: DeviceId,
361    ) -> Result<Rc<Self>, SimError> {
362        Self::new_and_register_with_renames(
363            engine, clock, parent, name, None, memory_map, pe_config, device_id,
364        )
365    }
366
367    pub fn set_dispatcher(&self, dispatcher: &Dispatcher) {
368        *self.dispatcher.borrow_mut() = Some(dispatcher.clone());
369    }
370
371    pub fn connect_port_tx(&self, port_state: PortStateResult<MemoryAccess>) -> SimResult {
372        self.lsu.connect_port_tx(port_state)
373    }
374
375    pub fn port_rx(&self) -> PortStateResult<MemoryAccess> {
376        self.lsu.port_rx()
377    }
378
379    #[must_use]
380    pub fn total_graph_nodes(&self) -> usize {
381        match self.dispatcher.borrow().as_ref() {
382            None => 0,
383            Some(dispatcher) => dispatcher.total_tasks_for_pe(self.entity.name.as_str()),
384        }
385    }
386
387    #[must_use]
388    pub fn total_flops(&self) -> usize {
389        self.stats.borrow().machine_ops.total()
390    }
391
392    #[must_use]
393    pub fn machine_ops(&self) -> MachineOpCounts {
394        self.stats.borrow().machine_ops
395    }
396
397    pub fn dump_stats(&self, time_now_ns: f64) {
398        let stats = self.stats.borrow();
399        log_stats(
400            &self.entity,
401            ProcessingElementStatsDisplay::new(
402                format!("ProcessingElement {}", self.entity.full_name()),
403                time_now_ns,
404                stats.machine_ops,
405            ),
406        );
407    }
408}
409
410#[async_trait(?Send)]
411impl Runnable for ProcessingElement {
412    async fn run(&self) -> SimResult {
413        let dispatcher = self
414            .dispatcher
415            .borrow()
416            .as_ref()
417            .ok_or_else(|| SimError("Started without dispatcher".to_string()))?
418            .clone();
419
420        let pe_name = self.entity.name.as_str();
421        let (mut complete, mut ready_node_indices) = dispatcher.ready_task_indices(pe_name)?;
422
423        loop {
424            if complete {
425                break;
426            }
427            if ready_node_indices.is_empty() {
428                // Wait for something to change
429                dispatcher.wait_for_change().await;
430            } else {
431                // Spawn all so they can run in parallel
432                for task_idx in ready_node_indices.drain(..) {
433                    dispatcher.set_task_active(task_idx)?;
434
435                    let clock = self.clock.clone();
436                    let dispatcher = dispatcher.clone();
437                    let lsu = self.lsu.clone();
438                    let compute_capabilities = self.compute_capabilities.clone();
439                    let stats = self.stats.clone();
440                    let entity = self.entity.clone();
441                    let activity_lanes = self.activity_lanes.clone();
442                    let flop_monitor = self.flop_monitor.clone();
443                    self.spawner.spawn(async move {
444                        handle_task(
445                            entity,
446                            clock,
447                            dispatcher,
448                            lsu,
449                            compute_capabilities,
450                            stats,
451                            activity_lanes,
452                            flop_monitor,
453                            task_idx,
454                        )
455                        .await
456                    });
457                }
458            }
459
460            (complete, ready_node_indices) = dispatcher.ready_task_indices(pe_name)?;
461        }
462        debug!(self.entity ; "PE {pe_name} DONE");
463        Ok(())
464    }
465}
466
467#[expect(clippy::too_many_arguments)]
468async fn handle_task(
469    entity: Rc<Entity>,
470    clock: Clock,
471    dispatcher: Dispatcher,
472    lsu: Rc<LoadStoreUnit>,
473    compute_capabilities: Rc<ComputeCapabilities>,
474    stats: Rc<RefCell<ProcessingElementStats>>,
475    activity_lanes: Rc<ProcessingElementActivityLanes>,
476    flop_monitor: Option<Rc<FlopMonitor>>,
477    task_idx: usize,
478) -> SimResult {
479    let task = dispatcher.task_by_id(task_idx)?;
480    match task {
481        Task::ComputeTask { config } => handle_compute_task(
482            clock,
483            dispatcher,
484            lsu,
485            task_idx,
486            compute_capabilities,
487            stats,
488            activity_lanes,
489            flop_monitor,
490            &config,
491        )
492        .await
493        .map_err(|err| SimError(format!("{entity} had error on task {}:\n{err}", config.id))),
494        Task::SyncTask { .. } => {
495            todo!();
496        }
497    }
498}
499
500fn tensor_view_num_bytes(view: &TensorView) -> usize {
501    view.num_bytes()
502}
503
504fn tensor_view_base_addr(view: &TensorView) -> Result<u64, SimError> {
505    let base_addr = view.tensor().addr();
506    let element_offset = view.element_offset()?;
507    let dtype = view.tensor().dtype();
508    let byte_offset = (dtype.num_bits() * element_offset).div_ceil(8) as u64;
509    Ok(base_addr + byte_offset)
510}
511
512#[expect(clippy::too_many_arguments)]
513async fn handle_compute_task(
514    clock: Clock,
515    dispatcher: Dispatcher,
516    lsu: Rc<LoadStoreUnit>,
517    task_idx: usize,
518    compute_capabilities: Rc<ComputeCapabilities>,
519    stats: Rc<RefCell<ProcessingElementStats>>,
520    activity_lanes: Rc<ProcessingElementActivityLanes>,
521    flop_monitor: Option<Rc<FlopMonitor>>,
522    config: &ComputeTaskConfig,
523) -> SimResult {
524    let total_num_bytes: usize = config
525        .inputs
526        .iter()
527        .chain(config.outputs.iter())
528        .filter_map(|view| view.as_ref())
529        .map(tensor_view_num_bytes)
530        .sum();
531
532    let num_partitions = total_num_bytes
533        .div_ceil(compute_capabilities.sram_bytes.max(1))
534        .max(1);
535
536    let partitions =
537        config
538            .op
539            .create_partitions(&config.inputs, &config.outputs, num_partitions)?;
540    let activity_name = config.activity_name();
541    let group = activity_lanes.create_group(&format!("{activity_name} operation"));
542
543    for partition in partitions {
544        for (idx, view) in partition.inputs.iter().enumerate() {
545            let Some(view) = view else {
546                continue;
547            };
548            lsu.do_access(
549                AccessType::ReadRequest,
550                tensor_view_num_bytes(view),
551                tensor_view_base_addr(view)?,
552                &activity_lanes.lsu_read,
553                &format!("{activity_name} tensor {idx} read"),
554                &group,
555            )
556            .await?;
557        }
558
559        let compute_ticks = config.op.compute_delay_ticks(
560            &compute_capabilities,
561            &partition.inputs,
562            &partition.outputs,
563        )?;
564        let machine_ops = config
565            .op
566            .compute_machine_ops(&partition.inputs, &partition.outputs)?;
567        let compute_flops = machine_ops.total();
568        if let Some(flop_monitor) = &flop_monitor {
569            flop_monitor.record_interval(compute_ticks as u64, compute_flops as f64);
570        }
571        {
572            // Lanes cannot support overlapping activity. If a lane will be released
573            // in the current clock cycle then we want to re-use it rather than allocate
574            // a new lane. Hence we wait here for the end of the current clock cycle
575            // to ensure all lanes that will be released in this cycle have been.
576            clock.wait_phase(phase::END).await;
577
578            let _activity = ActivityLanes::begin_in_group(
579                &activity_lanes.compute,
580                &format!("{activity_name} compute"),
581                &group,
582            );
583            clock.wait_ticks(compute_ticks as u64).await;
584        }
585        stats.borrow_mut().machine_ops.add_assign(machine_ops);
586
587        for (idx, view) in partition.outputs.iter().enumerate() {
588            let Some(view) = view else {
589                continue;
590            };
591            lsu.do_access(
592                AccessType::WriteNonPostedRequest,
593                tensor_view_num_bytes(view),
594                tensor_view_base_addr(view)?,
595                &activity_lanes.lsu_write,
596                &format!("{activity_name} tensor {idx} write"),
597                &group,
598            )
599            .await?;
600        }
601    }
602
603    dispatcher.set_task_completed(task_idx)?;
604    Ok(())
605}
606
607#[cfg(test)]
608mod tests {
609    use super::*;
610
611    #[test]
612    fn cycles_for_ops_uses_ceil_for_fractional_throughput() {
613        let compute_capabilities = ComputeCapabilities {
614            adds_per_tick: 0.5,
615            muls_per_tick: 2.5,
616            compares_per_tick: 4.0,
617            sram_bytes: 1024,
618        };
619
620        assert_eq!(
621            compute_capabilities
622                .cycles_for_ops(3, MachineOp::Add)
623                .unwrap(),
624            6
625        );
626        assert_eq!(
627            compute_capabilities
628                .cycles_for_ops(6, MachineOp::Mul)
629                .unwrap(),
630            3
631        );
632        assert_eq!(
633            compute_capabilities
634                .cycles_for_ops(0, MachineOp::Compare)
635                .unwrap(),
636            0
637        );
638    }
639
640    #[test]
641    fn cycles_for_ops_rejects_invalid_throughput() {
642        let compute_capabilities = ComputeCapabilities {
643            adds_per_tick: 0.0,
644            muls_per_tick: -1.0,
645            compares_per_tick: f64::INFINITY,
646            sram_bytes: 1024,
647        };
648
649        assert!(
650            compute_capabilities
651                .cycles_for_ops(1, MachineOp::Add)
652                .is_err()
653        );
654        assert!(
655            compute_capabilities
656                .cycles_for_ops(1, MachineOp::Mul)
657                .is_err()
658        );
659        assert!(
660            compute_capabilities
661                .cycles_for_ops(1, MachineOp::Compare)
662                .is_err()
663        );
664
665        let compute_capabilities = ComputeCapabilities {
666            adds_per_tick: f64::NAN,
667            muls_per_tick: 1.0,
668            compares_per_tick: 1.0,
669            sram_bytes: 1024,
670        };
671
672        assert!(
673            compute_capabilities
674                .cycles_for_ops(1, MachineOp::Add)
675                .is_err()
676        );
677    }
678}
gwr_models/processing_element/mod.rs

gwr_models/processing_element/
mod.rs
