gwr_models/processing_element/operators/
gemm.rs

1// Copyright (c) 2026 Graphcore Ltd. All rights reserved.
2
3//! The Gemm operator
4//!
5//! See <https://onnx.ai/onnx/operators/onnx__Gemm.html#l-onnx-doc-gemm>
6
7use std::rc::Rc;
8
9use gwr_engine::sim_error;
10use gwr_engine::types::{SimError, SimResult};
11use rand::RngExt;
12
13use super::{Operator, Tensor, TensorPartition};
14use crate::processing_element::operators::{
15    HasShape, Shape, TensorView, apply_dim_partitions, partition_across_dimensions,
16};
17use crate::processing_element::{ComputeCapabilities, MachineOp, MachineOpCounts};
18
19const NAME: &str = "Gemm";
20
21/// Return all dimensions that can be partitioned.
22///
23/// Prefer outer dims first. Add M and N in case required.
24fn choose_partition_dims<T: HasShape>(output: &T) -> Vec<usize> {
25    let dims = output.shape().get_dims();
26    let mut candidate_dims = dims
27        .iter()
28        .enumerate()
29        .filter_map(|(dim, size)| (*size > 1).then_some(dim))
30        .collect::<Vec<_>>();
31
32    if candidate_dims.is_empty() {
33        candidate_dims.push(output.num_dims().saturating_sub(1));
34    }
35
36    candidate_dims
37}
38
39// Define offsets from the inner-most dimension for named dimensions
40const INPUT_A_OFFSET_M: usize = 2;
41const INPUT_A_OFFSET_K: usize = 1;
42const INPUT_B_OFFSET_K: usize = 2;
43const INPUT_B_OFFSET_N: usize = 1;
44const OUTPUT_OFFSET_M: usize = 2;
45const OUTPUT_OFFSET_N: usize = 1;
46
47/// Return the value of a dimension starting from the inner most
48fn get_inner_dim<T: HasShape>(has_shape: &T, i: usize) -> usize {
49    has_shape.shape().get_dims()[has_shape.num_dims() - i].max(1)
50}
51
52/// Construct an input B shape for a Gemm consuming `input_a` as input A.
53pub fn gemm_rhs_shape<T: HasShape>(input_a: &T) -> Result<Shape, SimError> {
54    let rank = input_a.num_dims();
55    if rank < 2 {
56        return sim_error!("{NAME}: input A must be at least 2D");
57    }
58
59    let mut rhs_dims = input_a.shape().get_dims().clone();
60    rhs_dims[rank - INPUT_B_OFFSET_K] = get_inner_dim(input_a, INPUT_A_OFFSET_K);
61    rhs_dims[rank - INPUT_B_OFFSET_N] = get_inner_dim(input_a, INPUT_A_OFFSET_M);
62    Ok(Shape::new(&rhs_dims))
63}
64
65/// Choose a value for the K in a (M,K)x(K,N) -> (M,N) Gemm
66///
67/// Starts with the maximum of M and N and then grows or shrinks depending
68/// on the `expand_ratio` specified
69fn choose_gemm_k(output: &Tensor, rng: &mut impl RngExt, expand_ratio: f64) -> usize {
70    let m = get_inner_dim(output, OUTPUT_OFFSET_M);
71    let n = get_inner_dim(output, OUTPUT_OFFSET_N);
72    let reference = m.max(n);
73    let scaled = ((reference as f64) * expand_ratio).round().max(1.0) as usize;
74    let lower = ((scaled as f64) * 0.75).round().max(1.0) as usize;
75    let upper = ((scaled as f64) * 1.25).round().max(lower as f64) as usize;
76
77    if lower == upper {
78        lower
79    } else {
80        rng.random_range(lower..=upper)
81    }
82}
83
84fn should_add_input_c(rng: &mut impl RngExt, expand_ratio: f64) -> bool {
85    if !expand_ratio.is_finite() || expand_ratio <= 0.0 {
86        false
87    } else if expand_ratio >= 1.0 {
88        true
89    } else {
90        rng.random_bool(expand_ratio)
91    }
92}
93
94fn output_tensor_from_inputs(inputs: &[Option<Tensor>]) -> Result<Tensor, SimError> {
95    let (input_a, input_b) = validate_inputs(inputs)?;
96
97    let output_shape = broadcast_shapes(&input_a.shape, &input_b.shape)?;
98    let output_dtype = if input_a.dtype > input_b.dtype {
99        input_a.dtype
100    } else {
101        input_b.dtype
102    };
103
104    Ok(Tensor {
105        id: None,
106        shape: output_shape,
107        dtype: output_dtype,
108        addr: 0,
109    })
110}
111
112pub fn maybe_add_input_c(
113    inputs: &mut Vec<Option<Tensor>>,
114    expand_ratio: f64,
115    rng: &mut impl RngExt,
116) -> Result<bool, SimError> {
117    if inputs.len() >= 3 || !should_add_input_c(rng, expand_ratio) {
118        return Ok(false);
119    }
120
121    inputs.push(Some(output_tensor_from_inputs(inputs)?));
122    Ok(true)
123}
124
125fn broadcast_shapes(a: &Shape, b: &Shape) -> Result<Shape, SimError> {
126    let rank_a = a.num_dims();
127    let rank_b = b.num_dims();
128    let rank_result = rank_a.max(rank_b);
129
130    if rank_result < 2 {
131        return sim_error!("{NAME}: inputs must be at least 2D ({:?} and {:?})", a, b);
132    }
133
134    let mut result = vec![1; rank_result];
135
136    for (i, result_i) in result.iter_mut().enumerate() {
137        let a_dim = a.get_dim(rank_result, i);
138        let b_dim = b.get_dim(rank_result, i);
139
140        *result_i = if i == (rank_result.saturating_sub(OUTPUT_OFFSET_M)) {
141            // M
142            a_dim
143        } else if i == (rank_result.saturating_sub(OUTPUT_OFFSET_N)) {
144            // N
145            b_dim
146        } else if a_dim == b_dim {
147            a_dim
148        } else if a_dim == 1 {
149            b_dim
150        } else if b_dim == 1 {
151            a_dim
152        } else {
153            return sim_error!("{NAME}: cannot broadcast shapes {:?} and {:?}", a, b);
154        };
155    }
156
157    Ok(Shape(result))
158}
159
160fn validate_inputs<T: HasShape>(inputs: &[Option<T>]) -> Result<(&T, &T), SimError> {
161    let input_a = inputs[0]
162        .as_ref()
163        .ok_or(SimError(format!("{NAME}: missing input 0")))?;
164    let input_b = inputs[1]
165        .as_ref()
166        .ok_or(SimError(format!("{NAME}: missing input 1")))?;
167    let shape_a = input_a.shape();
168    let shape_b = input_b.shape();
169    let output_shape = broadcast_shapes(shape_a, shape_b)?;
170
171    let k_a = get_inner_dim(shape_a, INPUT_A_OFFSET_K);
172    let k_b = get_inner_dim(shape_b, INPUT_B_OFFSET_K);
173    if k_a != k_b {
174        return sim_error!("{NAME}: incompatible K in {:?} x {:?}", shape_a, shape_b);
175    }
176
177    if inputs.len() == 2 {
178        // Input C is a scalar that is broadcast - so no issues
179    } else if inputs.len() == 3
180        && let Some(input_c) = &inputs[2]
181    {
182        // Validate the Tensor C inputs
183        let out_m = get_inner_dim(&output_shape, OUTPUT_OFFSET_M);
184        let out_n = get_inner_dim(&output_shape, OUTPUT_OFFSET_N);
185
186        let shape_c = input_c.shape();
187        let c_m = get_inner_dim(&shape_c, OUTPUT_OFFSET_M);
188        let c_n = get_inner_dim(&shape_c, OUTPUT_OFFSET_N);
189        if (c_m != 1 && c_m != out_m) || (c_n != 1 && c_n != out_n) {
190            return sim_error!(
191                "{NAME}: input C incompatible ({:?} x {:?}) + {:?}",
192                shape_a,
193                shape_b,
194                shape_c
195            );
196        }
197    } else {
198        return sim_error!(
199            "{NAME}: {} input tensors found - expected 2 or 3",
200            inputs.len()
201        );
202    }
203
204    Ok((input_a, input_b))
205}
206
207fn validate_outputs<T: HasShape>(outputs: &[Option<T>]) -> Result<&T, SimError> {
208    if outputs.len() != 1 {
209        return sim_error!("{NAME}: {} outputs found - expected 1", outputs.len());
210    }
211    outputs[0]
212        .as_ref()
213        .ok_or(SimError(format!("{NAME}: missing output")))
214}
215
216fn validate_input_outputs<'a, 'b, T: HasShape>(
217    inputs: &'a [Option<T>],
218    outputs: &'b [Option<T>],
219) -> Result<(&'a T, &'a T, &'b T), SimError> {
220    let (input_a, input_b) = validate_inputs(inputs)?;
221    let output = validate_outputs(outputs)?;
222
223    let shape_a = input_a.shape();
224    let shape_b = input_b.shape();
225    let rank_inputs = input_a.num_dims().max(input_b.num_dims());
226    let shape_output = output.shape();
227    let rank_output = output.num_dims();
228
229    if rank_inputs != rank_output {
230        return sim_error!(
231            "{NAME}: incompatible ranks ({:?} x {:?} => {:?}",
232            shape_a,
233            shape_b,
234            shape_output
235        );
236    }
237
238    // We just need to check input vs output as the inputs have already been
239    // validated against each other
240    for i in 0..(rank_inputs - 2) {
241        let dim_a = input_a.get_dim(rank_inputs, i);
242        let dim_b = input_b.get_dim(rank_inputs, i);
243        let dim_in = if dim_a == 1 { dim_b } else { dim_a };
244        let dim_out = output.get_dim(rank_output, i);
245        if dim_in != 1 && dim_out != dim_in {
246            return sim_error!(
247                "{NAME}: Invalid output dimension {i} in {:?} x {:?} => {:?}",
248                shape_a,
249                shape_b,
250                shape_output
251            );
252        }
253    }
254
255    let input_m = get_inner_dim(input_a, INPUT_A_OFFSET_M);
256    let input_n = get_inner_dim(input_b, INPUT_B_OFFSET_N);
257    let output_m = get_inner_dim(output, OUTPUT_OFFSET_M);
258    let output_n = get_inner_dim(output, OUTPUT_OFFSET_N);
259
260    if (input_m != output_m) || (input_n != output_n) {
261        return sim_error!(
262            "{NAME}: Invalid M or N {:?} x {:?} => {:?}",
263            shape_a,
264            shape_b,
265            shape_output
266        );
267    }
268
269    Ok((input_a, input_b, output))
270}
271
272fn gemm_op_counts<T: HasShape>(
273    inputs: &[Option<T>],
274    outputs: &[Option<T>],
275) -> Result<(usize, usize), SimError> {
276    let (input_a_view, input_b_view, output_view) = validate_input_outputs(inputs, outputs)?;
277    let m = get_inner_dim(input_a_view, INPUT_A_OFFSET_M);
278    let k = get_inner_dim(input_a_view, INPUT_A_OFFSET_K);
279    let n = get_inner_dim(input_b_view, INPUT_B_OFFSET_N);
280
281    let num_matmuls = output_view
282        .shape()
283        .get_dims()
284        .iter()
285        .take(output_view.num_dims().saturating_sub(2))
286        .product::<usize>();
287
288    let num_muls = m * n * k * num_matmuls;
289    let num_matmul_adds = m * n * (k - 1) * num_matmuls;
290    // When there is a C input tensor each output element has one extra add
291    let num_c_adds = usize::from(inputs.len() == 3) * output_view.shape().num_elements();
292    let num_adds = num_matmul_adds + num_c_adds;
293    Ok((num_muls, num_adds))
294}
295
296pub struct OperatorGemm {}
297
298impl OperatorGemm {
299    pub fn create_outputs(
300        &self,
301        inputs: &[Option<Tensor>],
302        _expand_ratio: f64,
303        _rng: &mut impl RngExt,
304    ) -> Result<Vec<Option<Tensor>>, gwr_engine::types::SimError> {
305        Ok(vec![Some(output_tensor_from_inputs(inputs)?)])
306    }
307
308    pub fn create_inputs(
309        &self,
310        outputs: &[Option<Tensor>],
311        expand_ratio: f64,
312        rng: &mut impl RngExt,
313    ) -> Result<Vec<Option<Tensor>>, gwr_engine::types::SimError> {
314        let output = validate_outputs(outputs)?;
315
316        let mut input_a_shape = output.shape.clone();
317        let mut input_b_shape = output.shape.clone();
318
319        let k = choose_gemm_k(output, rng, expand_ratio);
320
321        let rank = output.num_dims();
322        input_a_shape.0[rank.saturating_sub(INPUT_A_OFFSET_K)] = k;
323        input_b_shape.0[rank.saturating_sub(INPUT_B_OFFSET_K)] = k;
324
325        let mut inputs = vec![
326            Some(Tensor {
327                id: None,
328                shape: input_a_shape,
329                dtype: output.dtype,
330                addr: 0,
331            }),
332            Some(Tensor {
333                id: None,
334                shape: input_b_shape,
335                dtype: output.dtype,
336                addr: 0,
337            }),
338        ];
339
340        maybe_add_input_c(&mut inputs, expand_ratio, rng)?;
341
342        Ok(inputs)
343    }
344}
345
346impl Operator for OperatorGemm {
347    fn validate_tensors(&self, inputs: &[Option<Tensor>], outputs: &[Option<Tensor>]) -> SimResult {
348        validate_input_outputs(inputs, outputs)?;
349        Ok(())
350    }
351
352    fn compute_delay_ticks(
353        &self,
354        compute_capabilities: &Rc<ComputeCapabilities>,
355        inputs: &[Option<TensorView>],
356        outputs: &[Option<TensorView>],
357    ) -> Result<usize, SimError> {
358        let (num_muls, num_adds) = gemm_op_counts(inputs, outputs)?;
359        Ok(
360            compute_capabilities.cycles_for_ops(num_muls, MachineOp::Mul)?
361                + compute_capabilities.cycles_for_ops(num_adds, MachineOp::Add)?,
362        )
363    }
364
365    fn compute_machine_ops(
366        &self,
367        inputs: &[Option<TensorView>],
368        outputs: &[Option<TensorView>],
369    ) -> Result<MachineOpCounts, SimError> {
370        let (num_muls, num_adds) = gemm_op_counts(inputs, outputs)?;
371        Ok(MachineOpCounts {
372            adds: num_adds,
373            muls: num_muls,
374            ..MachineOpCounts::default()
375        })
376    }
377
378    fn partition_views(
379        &self,
380        input_views: &[Option<TensorView>],
381        output_views: &[Option<TensorView>],
382        num_partitions: usize,
383    ) -> Result<Vec<TensorPartition>, SimError> {
384        let (input_a_view, input_b_view, output_view) =
385            validate_input_outputs(input_views, output_views)?;
386
387        let input_c_view = if input_views.len() > 2 {
388            input_views[2].clone()
389        } else {
390            None
391        };
392
393        let rank = output_view.num_dims();
394        let partition_dims = choose_partition_dims(&output_view);
395        let output_view_dims = output_view.shape().get_dims();
396        let partition_specs =
397            partition_across_dimensions(output_view_dims, &partition_dims, num_partitions);
398        let m_dim = rank.saturating_sub(OUTPUT_OFFSET_M);
399        let n_dim = rank.saturating_sub(OUTPUT_OFFSET_N);
400
401        let mut partitions = Vec::with_capacity(partition_specs.len());
402        for spec in partition_specs {
403            let (output_shape, partition_offsets) = apply_dim_partitions(output_view_dims, &spec);
404            let output_offsets = output_view
405                .offsets()
406                .get_dims()
407                .iter()
408                .zip(partition_offsets.iter())
409                .map(|(base, offset)| base + offset)
410                .collect::<Vec<_>>();
411            let output_view =
412                TensorView::new(output_view.tensor().clone(), &output_shape, &output_offsets);
413
414            let a_spec = spec
415                .iter()
416                .filter(|partition| partition.dim != n_dim)
417                .cloned()
418                .collect::<Vec<_>>();
419            let b_spec = spec
420                .iter()
421                .filter(|partition| partition.dim != m_dim)
422                .cloned()
423                .collect::<Vec<_>>();
424
425            let split_m = spec.iter().any(|partition| partition.dim == m_dim);
426            let split_n = spec.iter().any(|partition| partition.dim == n_dim);
427            let split_outer = spec
428                .iter()
429                .any(|partition| partition.dim != m_dim && partition.dim != n_dim);
430
431            let input_a_view = if split_outer || split_m {
432                TensorView::from_output_partitions_on_view(input_a_view, rank, &a_spec)
433            } else {
434                input_a_view.clone()
435            };
436
437            let input_b_view = if split_outer || split_n {
438                TensorView::from_output_partitions_on_view(input_b_view, rank, &b_spec)
439            } else {
440                input_b_view.clone()
441            };
442
443            let input_c_view = input_c_view
444                .as_ref()
445                .map(|view| TensorView::from_output_partitions_on_view(view, rank, &spec));
446
447            let mut partition_inputs = vec![Some(input_a_view), Some(input_b_view)];
448            if let Some(view) = input_c_view {
449                partition_inputs.push(Some(view));
450            }
451
452            partitions.push(TensorPartition {
453                inputs: partition_inputs,
454                outputs: vec![Some(output_view)],
455            });
456        }
457
458        Ok(partitions)
459    }
460}
461
462#[cfg(test)]
463mod tests {
464    use super::*;
465    use crate::processing_element::operators::dtype::DataType;
466    use crate::processing_element::operators::{Operator, Shape, Tensor, partition_tensors};
467
468    fn tensor(shape: &[usize]) -> Option<Tensor> {
469        Some(Tensor::new(shape, &DataType::Bf16, 0))
470    }
471
472    fn tensor_view(shape: &[usize]) -> Option<TensorView> {
473        let tensor = Tensor::new(shape, &DataType::Bf16, 0);
474        Some(TensorView::new_full(tensor))
475    }
476
477    #[test]
478    fn test_broadcast_shapes() {
479        let a = Shape(vec![1, 1, 4, 5]);
480        let b = Shape(vec![1, 1, 5, 10]);
481        let c = broadcast_shapes(&a, &b).unwrap();
482        assert_eq!(c, Shape(vec![1, 1, 4, 10]));
483
484        let a = Shape(vec![3, 1, 4, 5]);
485        let b = Shape(vec![1, 5, 5, 10]);
486        let c = broadcast_shapes(&a, &b).unwrap();
487        assert_eq!(c, Shape(vec![3, 5, 4, 10]));
488    }
489
490    #[test]
491    fn create_outputs_uses_gemm_m_and_n_before_broadcasting() {
492        let operator = OperatorGemm {};
493        let mut rng = rand::rng();
494
495        let outputs = operator
496            .create_outputs(
497                &[tensor(&[1, 48, 1, 25]), tensor(&[1, 48, 25, 1])],
498                1.0,
499                &mut rng,
500            )
501            .unwrap();
502
503        assert_eq!(
504            outputs[0].as_ref().unwrap().shape(),
505            &Shape::new(&[1, 48, 1, 1])
506        );
507    }
508
509    #[test]
510    fn validate_gemm() {
511        let operator = OperatorGemm {};
512
513        operator
514            .validate_tensors(
515                &[tensor(&[1, 4, 5]), tensor(&[1, 5, 8])],
516                &[tensor(&[1, 4, 8])],
517            )
518            .unwrap();
519
520        operator
521            .validate_tensors(
522                &[tensor(&[3, 2, 10, 5]), tensor(&[5, 12])],
523                &[tensor(&[3, 2, 10, 12])],
524            )
525            .unwrap();
526    }
527
528    #[test]
529    fn invalid_broadcast_1() {
530        let operator = OperatorGemm {};
531
532        let err = operator
533            .validate_tensors(
534                &[tensor(&[2, 4, 5]), tensor(&[3, 5, 8])],
535                &[tensor(&[1, 4, 8])],
536            )
537            .unwrap_err();
538        assert!(format!("{err}").contains("cannot broadcast shapes"));
539    }
540
541    #[test]
542    fn invalid_broadcast_2() {
543        let operator = OperatorGemm {};
544
545        let err = operator
546            .validate_tensors(
547                &[tensor(&[4, 1, 4, 5]), tensor(&[3, 1, 5, 8])],
548                &[tensor(&[1, 1, 4, 8])],
549            )
550            .unwrap_err();
551        assert!(format!("{err}").contains("cannot broadcast shapes"));
552    }
553
554    #[test]
555    fn invalid_k() {
556        let operator = OperatorGemm {};
557
558        let err = operator
559            .validate_tensors(
560                &[tensor(&[9, 2, 4, 5]), tensor(&[9, 2, 4, 8])],
561                &[tensor(&[9, 2, 4, 8])],
562            )
563            .unwrap_err();
564        assert!(format!("{err}").contains("incompatible K"));
565    }
566
567    #[test]
568    fn get_inner_dim_uses_one_based_inner_offsets() {
569        let shape = Shape::new(&[7, 11, 13, 17]);
570
571        assert_eq!(get_inner_dim(&shape, 1), 17);
572        assert_eq!(get_inner_dim(&shape, 2), 13);
573        assert_eq!(get_inner_dim(&shape, 3), 11);
574        assert_eq!(get_inner_dim(&shape, 4), 7);
575    }
576
577    #[test]
578    fn gemm_named_offsets_map_to_expected_dimensions() {
579        let input_a = Shape::new(&[19, 23, 29, 31]);
580        let input_b = Shape::new(&[19, 23, 31, 37]);
581        let output = Shape::new(&[19, 23, 29, 37]);
582
583        assert_eq!(get_inner_dim(&input_a, INPUT_A_OFFSET_M), 29);
584        assert_eq!(get_inner_dim(&input_a, INPUT_A_OFFSET_K), 31);
585        assert_eq!(get_inner_dim(&input_b, INPUT_B_OFFSET_K), 31);
586        assert_eq!(get_inner_dim(&input_b, INPUT_B_OFFSET_N), 37);
587        assert_eq!(get_inner_dim(&output, OUTPUT_OFFSET_M), 29);
588        assert_eq!(get_inner_dim(&output, OUTPUT_OFFSET_N), 37);
589    }
590
591    #[test]
592    fn gemm_rhs_shape_uses_gemm_named_offsets() {
593        let input_a = Shape::new(&[19, 23, 29, 31]);
594        let input_b = gemm_rhs_shape(&input_a).unwrap();
595
596        assert_eq!(input_b, Shape::new(&[19, 23, 31, 29]));
597
598        OperatorGemm {}
599            .validate_tensors(
600                &[tensor(input_a.get_dims()), tensor(input_b.get_dims())],
601                &[tensor(&[19, 23, 29, 29])],
602            )
603            .unwrap();
604    }
605
606    #[test]
607    fn gemm_rhs_shape_rejects_rank_below_two() {
608        let err = gemm_rhs_shape(&Shape::new(&[31])).unwrap_err();
609
610        assert!(format!("{err}").contains("input A must be at least 2D"));
611    }
612
613    #[test]
614    fn delay_ticks_uses_m_k_and_n_from_the_innermost_dimensions() {
615        let operator = OperatorGemm {};
616        let compute_capabilities = Rc::new(ComputeCapabilities {
617            adds_per_tick: 1.0,
618            muls_per_tick: 1.0,
619            compares_per_tick: 100.0,
620            sram_bytes: 1024,
621        });
622        let delay_ticks = operator
623            .compute_delay_ticks(
624                &compute_capabilities,
625                &[tensor_view(&[2, 3, 4, 5]), tensor_view(&[2, 3, 5, 7])],
626                &[tensor_view(&[2, 3, 4, 7])],
627            )
628            .unwrap();
629
630        // Expect outer dimension GEMMS of M * K * N muls + M * (K - 1) * N adds
631        assert_eq!(delay_ticks, (2 * 3) * ((4 * 5 * 7) + (4 * 4 * 7)));
632    }
633
634    #[test]
635    fn delay_ticks() {
636        let operator = OperatorGemm {};
637        let compute_capabilities = Rc::new(ComputeCapabilities {
638            adds_per_tick: 1.0,
639            muls_per_tick: 1.0,
640            compares_per_tick: 100.0,
641            sram_bytes: 1024,
642        });
643        let delay_ticks = operator
644            .compute_delay_ticks(
645                &compute_capabilities,
646                &[tensor_view(&[4, 5]), tensor_view(&[5, 8])],
647                &[tensor_view(&[4, 8])],
648            )
649            .unwrap();
650        assert_eq!(delay_ticks, 160 + 128);
651
652        let delay_ticks = operator
653            .compute_delay_ticks(
654                &compute_capabilities,
655                &[tensor_view(&[10, 11, 4, 5]), tensor_view(&[5, 8])],
656                &[tensor_view(&[10, 11, 4, 8])],
657            )
658            .unwrap();
659        assert_eq!(delay_ticks, 17600 + 14080);
660    }
661
662    #[test]
663    fn flop_count_adds_multiplies_and_accumulates() {
664        let operator = OperatorGemm {};
665        assert_eq!(
666            operator
667                .compute_flops(
668                    &[tensor_view(&[4, 5]), tensor_view(&[5, 8])],
669                    &[tensor_view(&[4, 8])],
670                )
671                .unwrap(),
672            (4 * 5 * 8) + (4 * (5 - 1) * 8)
673        );
674    }
675
676    #[test]
677    fn flop_count_includes_optional_c_elementwise_add() {
678        let operator = OperatorGemm {};
679        assert_eq!(
680            operator
681                .compute_flops(
682                    &[
683                        tensor_view(&[4, 5]),
684                        tensor_view(&[5, 8]),
685                        tensor_view(&[4, 8]),
686                    ],
687                    &[tensor_view(&[4, 8])],
688                )
689                .unwrap(),
690            (4 * 5 * 8) + (4 * (5 - 1) * 8) + (4 * 8)
691        );
692    }
693
694    #[test]
695    fn create_inputs_with_expand_ratio_zero_omits_input_c() {
696        let operator = OperatorGemm {};
697        let mut rng = rand::rng();
698
699        let inputs = operator
700            .create_inputs(&[tensor(&[4, 8])], 0.0, &mut rng)
701            .unwrap();
702
703        assert_eq!(inputs.len(), 2);
704        operator
705            .validate_tensors(&inputs, &[tensor(&[4, 8])])
706            .unwrap();
707    }
708
709    #[test]
710    fn create_inputs_with_expand_ratio_one_adds_input_c() {
711        let operator = OperatorGemm {};
712        let mut rng = rand::rng();
713
714        let inputs = operator
715            .create_inputs(&[tensor(&[4, 8])], 1.0, &mut rng)
716            .unwrap();
717
718        assert_eq!(inputs.len(), 3);
719        assert_eq!(inputs[2].as_ref().unwrap().shape(), &Shape::new(&[4, 8]));
720        operator
721            .validate_tensors(&inputs, &[tensor(&[4, 8])])
722            .unwrap();
723    }
724
725    type OffsetsShapes = (&'static [usize], &'static [usize]);
726    type PartitionOffsetsShapes = (OffsetsShapes, OffsetsShapes, OffsetsShapes);
727    fn check_partitions(partitions: &[TensorPartition], expected: &[PartitionOffsetsShapes]) {
728        assert_eq!(partitions.len(), expected.len());
729        for (partition, expected_partition) in partitions.iter().zip(expected.iter()) {
730            let in_a = partition.inputs[0].as_ref().unwrap();
731            let in_b = partition.inputs[1].as_ref().unwrap();
732            let out = partition.outputs[0].as_ref().unwrap();
733
734            assert_eq!(
735                in_a.offsets().get_dims().as_slice(),
736                (expected_partition.0).0
737            );
738            assert_eq!(in_a.shape().get_dims().as_slice(), (expected_partition.0).1);
739            assert_eq!(
740                in_b.offsets().get_dims().as_slice(),
741                (expected_partition.1).0
742            );
743            assert_eq!(in_b.shape().get_dims().as_slice(), (expected_partition.1).1);
744            assert_eq!(
745                out.offsets().get_dims().as_slice(),
746                (expected_partition.2).0
747            );
748            assert_eq!(out.shape().get_dims().as_slice(), (expected_partition.2).1);
749        }
750    }
751
752    #[test]
753    fn partitions_prefer_outer_dims_before_m_or_n() {
754        let operator = OperatorGemm {};
755        let input_tensors = vec![tensor(&[3, 20, 10, 5]), tensor(&[3, 20, 5, 12])];
756        let output_tensors = vec![tensor(&[3, 20, 10, 12])];
757
758        let partitions = partition_tensors(&operator, &input_tensors, &output_tensors, 4).unwrap();
759        assert_eq!(partitions.len(), 6);
760
761        let expected: &[PartitionOffsetsShapes] = &[
762            (
763                (&[0, 0, 0, 0], &[1, 10, 10, 5]),
764                (&[0, 0, 0, 0], &[1, 10, 5, 12]),
765                (&[0, 0, 0, 0], &[1, 10, 10, 12]),
766            ),
767            (
768                (&[0, 10, 0, 0], &[1, 10, 10, 5]),
769                (&[0, 10, 0, 0], &[1, 10, 5, 12]),
770                (&[0, 10, 0, 0], &[1, 10, 10, 12]),
771            ),
772            (
773                (&[1, 0, 0, 0], &[1, 10, 10, 5]),
774                (&[1, 0, 0, 0], &[1, 10, 5, 12]),
775                (&[1, 0, 0, 0], &[1, 10, 10, 12]),
776            ),
777            (
778                (&[1, 10, 0, 0], &[1, 10, 10, 5]),
779                (&[1, 10, 0, 0], &[1, 10, 5, 12]),
780                (&[1, 10, 0, 0], &[1, 10, 10, 12]),
781            ),
782            (
783                (&[2, 0, 0, 0], &[1, 10, 10, 5]),
784                (&[2, 0, 0, 0], &[1, 10, 5, 12]),
785                (&[2, 0, 0, 0], &[1, 10, 10, 12]),
786            ),
787            (
788                (&[2, 10, 0, 0], &[1, 10, 10, 5]),
789                (&[2, 10, 0, 0], &[1, 10, 5, 12]),
790                (&[2, 10, 0, 0], &[1, 10, 10, 12]),
791            ),
792        ];
793        check_partitions(&partitions, expected);
794    }
795
796    #[test]
797    fn partitions_fall_back_to_m_when_no_outer_dims_are_available() {
798        let operator = OperatorGemm {};
799        let inputs = vec![tensor(&[10, 5]), tensor(&[5, 12])];
800        let outputs = vec![tensor(&[10, 12])];
801
802        let partitions = partition_tensors(&operator, &inputs, &outputs, 4).unwrap();
803        assert_eq!(partitions.len(), 4);
804
805        let expected_m_offsets = [0, 3, 6, 8];
806        let expected_m_lengths = [3, 3, 2, 2];
807
808        for (partition_idx, partition) in partitions.iter().enumerate() {
809            let a_view = partition.inputs[0].as_ref().unwrap();
810            let b_view = partition.inputs[1].as_ref().unwrap();
811            let out_view = partition.outputs[0].as_ref().unwrap();
812
813            assert_eq!(
814                a_view.shape().get_dims().as_slice(),
815                &[expected_m_lengths[partition_idx], 5]
816            );
817            assert_eq!(
818                a_view.offsets().get_dims().as_slice(),
819                &[expected_m_offsets[partition_idx], 0]
820            );
821
822            assert_eq!(b_view.shape().get_dims().as_slice(), &[5, 12]);
823            assert_eq!(b_view.offsets().get_dims().as_slice(), &[0, 0]);
824
825            assert_eq!(
826                out_view.shape().get_dims().as_slice(),
827                &[expected_m_lengths[partition_idx], 12]
828            );
829            assert_eq!(
830                out_view.offsets().get_dims().as_slice(),
831                &[expected_m_offsets[partition_idx], 0]
832            );
833        }
834    }
835
836    #[test]
837    fn can_partition_m_and_n_when_needed() {
838        let operator = OperatorGemm {};
839        let inputs = vec![tensor(&[4, 5]), tensor(&[5, 6])];
840        let outputs = vec![tensor(&[4, 6])];
841
842        let partitions = partition_tensors(&operator, &inputs, &outputs, 8).unwrap();
843        assert_eq!(partitions.len(), 8);
844
845        let expected: &[PartitionOffsetsShapes] = &[
846            ((&[0, 0], &[1, 5]), (&[0, 0], &[5, 3]), (&[0, 0], &[1, 3])),
847            ((&[0, 0], &[1, 5]), (&[0, 3], &[5, 3]), (&[0, 3], &[1, 3])),
848            ((&[1, 0], &[1, 5]), (&[0, 0], &[5, 3]), (&[1, 0], &[1, 3])),
849            ((&[1, 0], &[1, 5]), (&[0, 3], &[5, 3]), (&[1, 3], &[1, 3])),
850            ((&[2, 0], &[1, 5]), (&[0, 0], &[5, 3]), (&[2, 0], &[1, 3])),
851            ((&[2, 0], &[1, 5]), (&[0, 3], &[5, 3]), (&[2, 3], &[1, 3])),
852            ((&[3, 0], &[1, 5]), (&[0, 0], &[5, 3]), (&[3, 0], &[1, 3])),
853            ((&[3, 0], &[1, 5]), (&[0, 3], &[5, 3]), (&[3, 3], &[1, 3])),
854        ];
855        check_partitions(&partitions, expected);
856    }
857
858    #[test]
859    fn partitions_preserve_subset_view_offsets() {
860        let operator = OperatorGemm {};
861        let input_a = Tensor::new(&[8, 5], &DataType::Bf16, 0);
862        let input_b = Tensor::new(&[5, 9], &DataType::Bf16, 0);
863        let output = Tensor::new(&[8, 9], &DataType::Bf16, 0);
864
865        let input_views = vec![
866            Some(TensorView::new(input_a, &[4, 5], &[2, 0])),
867            Some(TensorView::new(input_b, &[5, 6], &[0, 3])),
868        ];
869        let output_views = vec![Some(TensorView::new(output, &[4, 6], &[2, 3]))];
870
871        let partitions = operator
872            .partition_views(&input_views, &output_views, 8)
873            .unwrap();
874        assert_eq!(partitions.len(), 8);
875
876        let expected: &[PartitionOffsetsShapes] = &[
877            ((&[2, 0], &[1, 5]), (&[0, 3], &[5, 3]), (&[2, 3], &[1, 3])),
878            ((&[2, 0], &[1, 5]), (&[0, 6], &[5, 3]), (&[2, 6], &[1, 3])),
879            ((&[3, 0], &[1, 5]), (&[0, 3], &[5, 3]), (&[3, 3], &[1, 3])),
880            ((&[3, 0], &[1, 5]), (&[0, 6], &[5, 3]), (&[3, 6], &[1, 3])),
881            ((&[4, 0], &[1, 5]), (&[0, 3], &[5, 3]), (&[4, 3], &[1, 3])),
882            ((&[4, 0], &[1, 5]), (&[0, 6], &[5, 3]), (&[4, 6], &[1, 3])),
883            ((&[5, 0], &[1, 5]), (&[0, 3], &[5, 3]), (&[5, 3], &[1, 3])),
884            ((&[5, 0], &[1, 5]), (&[0, 6], &[5, 3]), (&[5, 6], &[1, 3])),
885        ];
886        check_partitions(&partitions, expected);
887    }
888}
gwr_models/processing_element/operators/gemm.rs

gwr_models/processing_element/operators/
gemm.rs
