gwr_models/processing_element/operators/
add.rs

1// Copyright (c) 2026 Graphcore Ltd. All rights reserved.
2
3//! The Add operator
4//!
5//! See <https://onnx.ai/onnx/operators/onnx__Add.html#l-onnx-doc-add>
6
7use std::rc::Rc;
8
9use gwr_engine::sim_error;
10use gwr_engine::types::{SimError, SimResult};
11use rand::RngExt;
12
13use super::{Operator, Shape, Tensor, TensorPartition};
14use crate::processing_element::operators::{
15    HasShape, TensorView, apply_dim_partitions, partition_across_dimensions,
16};
17use crate::processing_element::{ComputeCapabilities, MachineOp, MachineOpCounts};
18
19const NAME: &str = "Add";
20
21fn choose_partition_dims<T: HasShape>(output: &T) -> Vec<usize> {
22    let dims = output.shape().get_dims();
23    let mut candidate_dims = dims
24        .iter()
25        .enumerate()
26        .filter_map(|(dim, size)| (*size > 1).then_some(dim))
27        .collect::<Vec<_>>();
28
29    if candidate_dims.is_empty() {
30        candidate_dims.push(output.num_dims().saturating_sub(1));
31    }
32
33    candidate_dims
34}
35
36pub struct OperatorAdd {}
37
38fn broadcast_shapes(a: &Shape, b: &Shape) -> Result<Shape, SimError> {
39    let rank_a = a.num_dims();
40    let rank_b = b.num_dims();
41    let rank = rank_a.max(rank_b);
42    let mut result = vec![1; rank];
43
44    for (i, result_i) in result.iter_mut().enumerate() {
45        let a_dim = a.get_dim(rank, i);
46        let b_dim = b.get_dim(rank, i);
47
48        *result_i = if a_dim == b_dim {
49            a_dim
50        } else if a_dim == 1 {
51            b_dim
52        } else if b_dim == 1 {
53            a_dim
54        } else {
55            return sim_error!("{NAME}: cannot broadcast shapes {:?} and {:?}", a, b);
56        };
57    }
58
59    Ok(Shape(result))
60}
61
62fn choose_input_shape(output: &Tensor, rng: &mut impl RngExt, expand_ratio: f64) -> Shape {
63    let keep_prob = expand_ratio.clamp(0.0, 1.0);
64    let mut dims = output.shape.0.clone();
65
66    for dim in &mut dims {
67        if *dim > 1 && rng.random_bool(1.0 - keep_prob) {
68            *dim = 1;
69        }
70    }
71
72    while dims.len() > 1 && dims[0] == 1 && rng.random_bool(1.0 - keep_prob) {
73        dims.remove(0);
74    }
75
76    Shape(dims)
77}
78
79fn validate_inputs<T: HasShape>(inputs: &[Option<T>]) -> Result<(&T, &T), SimError> {
80    if inputs.len() != 2 {
81        return sim_error!("{NAME}: {} inputs found - expected 2", inputs.len());
82    }
83    let input_a = inputs[0]
84        .as_ref()
85        .ok_or(SimError(format!("{NAME}: missing input 0")))?;
86    let input_b = inputs[1]
87        .as_ref()
88        .ok_or(SimError(format!("{NAME}: missing input 1")))?;
89    Ok((input_a, input_b))
90}
91
92fn validate_outputs<T: HasShape>(outputs: &[Option<T>]) -> Result<&T, SimError> {
93    if outputs.len() != 1 {
94        return sim_error!("{NAME}: {} outputs found - expected 1", outputs.len());
95    }
96    outputs[0]
97        .as_ref()
98        .ok_or(SimError(format!("{NAME}: missing output")))
99}
100
101fn validate_input_outputs<'a, 'b, T: HasShape>(
102    inputs: &'a [Option<T>],
103    outputs: &'b [Option<T>],
104) -> Result<(&'a T, &'a T, &'b T), SimError> {
105    let (input_a, input_b) = validate_inputs(inputs)?;
106    let output = validate_outputs(outputs)?;
107
108    let expected_shape = broadcast_shapes(input_a.shape(), input_b.shape())?;
109    if expected_shape != *output.shape() {
110        return sim_error!(
111            "{NAME}: Invalid output shape - expected {:?}, found {:?}",
112            expected_shape,
113            output.shape()
114        );
115    }
116    Ok((input_a, input_b, output))
117}
118
119fn num_add_flops<T: HasShape>(
120    inputs: &[Option<T>],
121    outputs: &[Option<T>],
122) -> Result<usize, SimError> {
123    let (_, _, output) = validate_input_outputs(inputs, outputs)?;
124    Ok(output.num_elements())
125}
126
127impl OperatorAdd {
128    pub fn create_outputs(
129        &self,
130        inputs: &[Option<Tensor>],
131        _expand_ratio: f64,
132        _rng: &mut impl RngExt,
133    ) -> Result<Vec<Option<Tensor>>, gwr_engine::types::SimError> {
134        let (input_a, input_b) = validate_inputs(inputs)?;
135
136        let output_shape = broadcast_shapes(&input_a.shape, &input_b.shape)?;
137        let output_dtype = if input_a.dtype > input_b.dtype {
138            input_a.dtype
139        } else {
140            input_b.dtype
141        };
142
143        Ok(vec![Some(Tensor {
144            id: None,
145            shape: output_shape,
146            dtype: output_dtype,
147            addr: 0,
148        })])
149    }
150
151    pub fn create_inputs(
152        &self,
153        outputs: &[Option<Tensor>],
154        expand_ratio: f64,
155        rng: &mut impl RngExt,
156    ) -> Result<Vec<Option<Tensor>>, gwr_engine::types::SimError> {
157        let output = validate_outputs(outputs)?;
158
159        // We cannot shrink both inputs because we need to preserve one that will
160        // cause the output to be of the right shape. So we choose one and update
161        // that.
162        let (input_a_shape, input_b_shape) = if rng.random_bool(0.5) {
163            (
164                choose_input_shape(output, rng, expand_ratio),
165                output.shape.clone(),
166            )
167        } else {
168            (
169                output.shape.clone(),
170                choose_input_shape(output, rng, expand_ratio),
171            )
172        };
173
174        Ok(vec![
175            Some(Tensor {
176                id: None,
177                shape: input_a_shape,
178                dtype: output.dtype,
179                addr: 0,
180            }),
181            Some(Tensor {
182                id: None,
183                shape: input_b_shape,
184                dtype: output.dtype,
185                addr: 0,
186            }),
187        ])
188    }
189}
190
191impl Operator for OperatorAdd {
192    fn validate_tensors(&self, inputs: &[Option<Tensor>], outputs: &[Option<Tensor>]) -> SimResult {
193        validate_input_outputs(inputs, outputs)?;
194        Ok(())
195    }
196
197    fn compute_delay_ticks(
198        &self,
199        compute_capabilities: &Rc<ComputeCapabilities>,
200        inputs: &[Option<TensorView>],
201        outputs: &[Option<TensorView>],
202    ) -> Result<usize, SimError> {
203        let num_adds = num_add_flops(inputs, outputs)?;
204        compute_capabilities.cycles_for_ops(num_adds, MachineOp::Add)
205    }
206
207    fn compute_machine_ops(
208        &self,
209        inputs: &[Option<TensorView>],
210        outputs: &[Option<TensorView>],
211    ) -> Result<MachineOpCounts, SimError> {
212        Ok(MachineOpCounts {
213            adds: num_add_flops(inputs, outputs)?,
214            ..MachineOpCounts::default()
215        })
216    }
217
218    fn partition_views(
219        &self,
220        input_views: &[Option<TensorView>],
221        output_views: &[Option<TensorView>],
222        num_partitions: usize,
223    ) -> Result<Vec<TensorPartition>, SimError> {
224        let (_, _, output_view) = validate_input_outputs(input_views, output_views)?;
225
226        let partition_dims = choose_partition_dims(&output_view);
227        let output_view_dims = output_view.shape().get_dims();
228        let partition_specs =
229            partition_across_dimensions(output_view_dims, &partition_dims, num_partitions);
230        let output_rank = output_view.num_dims();
231
232        let mut partitions = Vec::with_capacity(partition_specs.len());
233        for spec in partition_specs {
234            let (output_shape, partition_offsets) = apply_dim_partitions(output_view_dims, &spec);
235            let output_offsets = output_view
236                .offsets()
237                .get_dims()
238                .iter()
239                .zip(partition_offsets.iter())
240                .map(|(base, offset)| base + offset)
241                .collect::<Vec<_>>();
242            let output_view =
243                TensorView::new(output_view.tensor().clone(), &output_shape, &output_offsets);
244
245            let input_views = input_views
246                .iter()
247                .map(|maybe_input_view| {
248                    maybe_input_view.as_ref().map(|input_view| {
249                        TensorView::from_output_partitions_on_view(input_view, output_rank, &spec)
250                    })
251                })
252                .collect::<Vec<_>>();
253
254            partitions.push(TensorPartition {
255                inputs: input_views,
256                outputs: vec![Some(output_view)],
257            });
258        }
259
260        Ok(partitions)
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use std::convert::Infallible;
267
268    use rand::TryRng;
269
270    use super::*;
271    use crate::processing_element::operators::dtype::DataType;
272    use crate::processing_element::operators::partition_tensors;
273
274    fn tensor(dims: &[usize]) -> Option<Tensor> {
275        Some(Tensor::new(dims, &DataType::Bf16, 0))
276    }
277
278    fn tensor_view(dims: &[usize]) -> Option<TensorView> {
279        let tensor = Tensor::new(dims, &DataType::Bf16, 0);
280        Some(TensorView::new_full(tensor))
281    }
282
283    struct FixedBoolRng {
284        values: Vec<bool>,
285    }
286
287    impl FixedBoolRng {
288        fn with_bool_values(values: impl IntoIterator<Item = bool>) -> Self {
289            Self {
290                values: values.into_iter().collect(),
291            }
292        }
293
294        fn next_bool(&mut self) -> bool {
295            if self.values.is_empty() {
296                true
297            } else {
298                self.values.remove(0)
299            }
300        }
301    }
302
303    impl TryRng for FixedBoolRng {
304        type Error = Infallible;
305
306        fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
307            Ok(if self.next_bool() { 0 } else { u32::MAX })
308        }
309
310        fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
311            Ok(if self.next_bool() { 0 } else { u64::MAX })
312        }
313
314        fn try_fill_bytes(&mut self, dst: &mut [u8]) -> Result<(), Self::Error> {
315            for chunk in dst.chunks_mut(size_of::<u64>()) {
316                let bytes = self.try_next_u64()?.to_le_bytes();
317                chunk.copy_from_slice(&bytes[..chunk.len()]);
318            }
319            Ok(())
320        }
321    }
322
323    #[test]
324    fn create_outputs_broadcasts_same_rank_inputs() {
325        let op = OperatorAdd {};
326        let inputs = vec![tensor(&[2, 3, 4]), tensor(&[1, 3, 1])];
327        let mut rng = rand::rng();
328
329        let outputs = op.create_outputs(&inputs, 1.0, &mut rng).unwrap();
330
331        assert_eq!(outputs.len(), 1);
332        assert_eq!(outputs[0].as_ref().unwrap().shape, Shape(vec![2, 3, 4]));
333    }
334
335    #[test]
336    fn create_outputs_broadcasts_different_rank_inputs() {
337        let op = OperatorAdd {};
338        let inputs = vec![tensor(&[3, 4]), tensor(&[2, 1, 4])];
339        let mut rng = rand::rng();
340
341        let outputs = op.create_outputs(&inputs, 1.0, &mut rng).unwrap();
342
343        assert_eq!(outputs.len(), 1);
344        assert_eq!(outputs[0].as_ref().unwrap().shape, Shape(vec![2, 3, 4]));
345    }
346
347    #[test]
348    fn create_outputs_rejects_non_broadcastable_inputs() {
349        let op = OperatorAdd {};
350        let inputs = vec![tensor(&[2, 3]), tensor(&[4, 3])];
351        let mut rng = rand::rng();
352
353        let err = op.create_outputs(&inputs, 1.0, &mut rng).unwrap_err();
354
355        assert!(format!("{err}").contains("cannot broadcast shapes"));
356    }
357
358    #[test]
359    fn validate_tensors_accepts_broadcasted_output_shape() {
360        let op = OperatorAdd {};
361        let inputs = vec![tensor(&[3, 4]), tensor(&[2, 1, 4])];
362        let outputs = vec![tensor(&[2, 3, 4])];
363
364        op.validate_tensors(&inputs, &outputs).unwrap();
365    }
366
367    #[test]
368    fn validate_tensors_rejects_wrong_output_shape() {
369        let op = OperatorAdd {};
370        let inputs = vec![tensor(&[3, 4]), tensor(&[2, 1, 4])];
371        let outputs = vec![tensor(&[3, 4])];
372
373        let err = op.validate_tensors(&inputs, &outputs).unwrap_err();
374
375        assert!(format!("{err}").contains("Invalid output shape"));
376    }
377
378    #[test]
379    fn create_inputs_with_expand_ratio_one_preserves_output_shape() {
380        let op = OperatorAdd {};
381        let outputs = vec![tensor(&[2, 3, 4])];
382
383        for shrink_input_a in [false, true] {
384            let mut rng = FixedBoolRng::with_bool_values([shrink_input_a]);
385
386            let inputs = op.create_inputs(&outputs, 1.0, &mut rng).unwrap();
387
388            assert_eq!(inputs.len(), 2);
389            assert_eq!(inputs[0].as_ref().unwrap().shape, Shape(vec![2, 3, 4]));
390            assert_eq!(inputs[1].as_ref().unwrap().shape, Shape(vec![2, 3, 4]));
391            op.validate_tensors(&inputs, &outputs).unwrap();
392        }
393    }
394
395    #[test]
396    fn create_inputs_with_expand_ratio_zero_creates_broadcastable_inputs() {
397        let op = OperatorAdd {};
398        let outputs = vec![tensor(&[2, 3, 4])];
399
400        for shrink_input_a in [false, true] {
401            let mut rng = FixedBoolRng::with_bool_values([shrink_input_a]);
402
403            let inputs = op.create_inputs(&outputs, 0.0, &mut rng).unwrap();
404
405            assert_eq!(inputs.len(), 2);
406            op.validate_tensors(&inputs, &outputs).unwrap();
407
408            let inputs: Vec<Tensor> = inputs.into_iter().map(|input| input.unwrap()).collect();
409            let outputs: Vec<Tensor> = outputs
410                .iter()
411                .map(|output| output.clone().unwrap())
412                .collect();
413            assert_eq!(inputs[usize::from(shrink_input_a)].shape, outputs[0].shape);
414            assert!(inputs[0].num_dims() <= outputs[0].num_dims());
415            assert!(inputs[1].num_dims() <= outputs[0].num_dims());
416            assert!(
417                inputs[0]
418                    .shape
419                    .0
420                    .iter()
421                    .all(|dim| *dim == 1 || outputs[0].shape.0.contains(dim))
422            );
423            assert!(
424                inputs[1]
425                    .shape
426                    .0
427                    .iter()
428                    .all(|dim| *dim == 1 || outputs[0].shape.0.contains(dim))
429            );
430        }
431    }
432
433    #[test]
434    fn delay_ticks() {
435        let compute_capabilities = Rc::new(ComputeCapabilities {
436            adds_per_tick: 1.0,
437            muls_per_tick: 100.0,
438            compares_per_tick: 200.0,
439            sram_bytes: 1024,
440        });
441        let operator = OperatorAdd {};
442        let delay_ticks = operator
443            .compute_delay_ticks(
444                &compute_capabilities,
445                &[tensor_view(&[4, 5]), tensor_view(&[4, 5])],
446                &[tensor_view(&[4, 5])],
447            )
448            .unwrap();
449        assert_eq!(delay_ticks, 20);
450
451        let compute_capabilities = Rc::new(ComputeCapabilities {
452            adds_per_tick: 2.0,
453            muls_per_tick: 100.0,
454            compares_per_tick: 100.0,
455            sram_bytes: 1024,
456        });
457        let delay_ticks = operator
458            .compute_delay_ticks(
459                &compute_capabilities,
460                &[tensor_view(&[4, 5]), tensor_view(&[4, 5])],
461                &[tensor_view(&[4, 5])],
462            )
463            .unwrap();
464        assert_eq!(delay_ticks, 10);
465
466        let delay_ticks = operator
467            .compute_delay_ticks(
468                &compute_capabilities,
469                &[tensor_view(&[10, 4, 5]), tensor_view(&[10, 4, 5])],
470                &[tensor_view(&[10, 4, 5])],
471            )
472            .unwrap();
473        assert_eq!(delay_ticks, 100);
474    }
475
476    #[test]
477    fn flop_count_matches_output_elements() {
478        let operator = OperatorAdd {};
479        assert_eq!(
480            operator
481                .compute_flops(
482                    &[tensor_view(&[2, 3, 4]), tensor_view(&[1, 3, 1])],
483                    &[tensor_view(&[2, 3, 4])],
484                )
485                .unwrap(),
486            24
487        );
488    }
489
490    type OffsetsShapes = (&'static [usize], &'static [usize]);
491
492    // Ensure that both inputs and the output are all of the expected shape as they
493    // should all be the same.
494    fn check_partitions(partitions: &[TensorPartition], expected: &[OffsetsShapes]) {
495        assert_eq!(partitions.len(), expected.len());
496
497        for (partition, (expected_offsets, expected_shape)) in
498            partitions.iter().zip(expected.iter())
499        {
500            let views = partition
501                .inputs
502                .iter()
503                .chain(partition.outputs.iter())
504                .map(|view| view.as_ref().unwrap());
505
506            for view in views {
507                assert_eq!(view.offsets().get_dims().as_slice(), *expected_offsets);
508                assert_eq!(view.shape().get_dims().as_slice(), *expected_shape);
509            }
510        }
511    }
512
513    #[test]
514    fn can_partition_across_one_dimension() {
515        let op = OperatorAdd {};
516        let inputs = vec![tensor(&[1, 5, 3, 4]), tensor(&[1, 5, 3, 4])];
517        let outputs = vec![tensor(&[1, 5, 3, 4])];
518
519        let partitions = partition_tensors(&op, &inputs, &outputs, 5).unwrap();
520        assert_eq!(partitions.len(), 5);
521
522        let expected: &[OffsetsShapes] = &[
523            (&[0, 0, 0, 0], &[1, 1, 3, 4]),
524            (&[0, 1, 0, 0], &[1, 1, 3, 4]),
525            (&[0, 2, 0, 0], &[1, 1, 3, 4]),
526            (&[0, 3, 0, 0], &[1, 1, 3, 4]),
527            (&[0, 4, 0, 0], &[1, 1, 3, 4]),
528        ];
529        check_partitions(&partitions, expected);
530    }
531
532    #[test]
533    fn can_partition_across_multiple_dimensions() {
534        let op = OperatorAdd {};
535        let inputs = vec![tensor(&[2, 3, 4]), tensor(&[2, 3, 4])];
536        let outputs = vec![tensor(&[2, 3, 4])];
537
538        let partitions = partition_tensors(&op, &inputs, &outputs, 5).unwrap();
539        assert_eq!(partitions.len(), 6);
540
541        let expected: &[OffsetsShapes] = &[
542            (&[0, 0, 0], &[1, 1, 4]),
543            (&[0, 1, 0], &[1, 1, 4]),
544            (&[0, 2, 0], &[1, 1, 4]),
545            (&[1, 0, 0], &[1, 1, 4]),
546            (&[1, 1, 0], &[1, 1, 4]),
547            (&[1, 2, 0], &[1, 1, 4]),
548        ];
549        check_partitions(&partitions, expected);
550    }
551
552    #[test]
553    fn partitions_preserve_subset_view_offsets() {
554        let op = OperatorAdd {};
555        let input_a = Tensor::new(&[4, 5, 4], &DataType::Bf16, 0);
556        let input_b = Tensor::new(&[4, 5, 4], &DataType::Bf16, 0);
557        let output = Tensor::new(&[4, 5, 4], &DataType::Bf16, 0);
558
559        let input_views = vec![
560            Some(TensorView::new(input_a, &[2, 2, 4], &[1, 2, 0])),
561            Some(TensorView::new(input_b, &[2, 2, 4], &[1, 2, 0])),
562        ];
563        let output_views = vec![Some(TensorView::new(output, &[2, 2, 4], &[1, 2, 0]))];
564
565        let partitions = op.partition_views(&input_views, &output_views, 2).unwrap();
566        assert_eq!(partitions.len(), 2);
567
568        let expected: &[OffsetsShapes] = &[(&[1, 2, 0], &[1, 2, 4]), (&[2, 2, 0], &[1, 2, 4])];
569        check_partitions(&partitions, expected);
570    }
571}
gwr_models/processing_element/operators/add.rs

gwr_models/processing_element/operators/
add.rs
