gwr_models/processing_element/operators/
mod.rs

1// Copyright (c) 2026 Graphcore Ltd. All rights reserved.
2
3//! The Operators define what operations a Processing Element can perform
4
5use std::fmt::Display;
6use std::rc::Rc;
7
8use gwr_engine::sim_error;
9use gwr_engine::types::{SimError, SimResult};
10
11use crate::processing_element::operators::dtype::DataType;
12use crate::processing_element::{ComputeCapabilities, MachineOpCounts};
13
14pub mod dtype;
15
16#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub enum ExpansionDirection {
18    Backward,
19    Forward,
20}
21
22#[must_use]
23pub fn shape_string(dims: &[usize]) -> String {
24    dims.iter()
25        .map(|d| d.to_string())
26        .collect::<Vec<_>>()
27        .join("×")
28}
29
30pub trait HasShape {
31    /// Return the number of dimensions
32    #[must_use]
33    fn num_dims(&self) -> usize;
34
35    /// Return the number of elements
36    #[must_use]
37    fn num_elements(&self) -> usize;
38
39    /// Return the size of a given dimension within a larger space.
40    ///
41    /// Assumes the defined shape are the inner dimensions and will return 1
42    /// when the specified dimension is out of the defined shape.
43    ///
44    /// For example, a shape of the form [2, 4, 5] with calls:
45    ///  shape.get_dim(4, 0) will return 1 (dimension outside of defined shape)
46    ///  shape.get_dim(4, 1) will return 2
47    ///  shape.get_dim(4, 2) will return 4
48    ///  shape.get_dim(4, 3) will return 5
49    #[must_use]
50    fn get_dim(&self, total_dims: usize, i: usize) -> usize;
51
52    /// Get access to the underlying shape
53    #[must_use]
54    fn shape(&self) -> &Shape;
55}
56
57impl<T> HasShape for &T
58where
59    T: HasShape,
60{
61    fn num_dims(&self) -> usize {
62        (*self).num_dims()
63    }
64
65    fn num_elements(&self) -> usize {
66        (*self).num_elements()
67    }
68
69    fn get_dim(&self, total_dims: usize, i: usize) -> usize {
70        (*self).get_dim(total_dims, i)
71    }
72
73    fn shape(&self) -> &Shape {
74        (*self).shape()
75    }
76}
77
78#[derive(Clone, Debug, PartialEq)]
79pub struct Shape(Vec<usize>);
80
81impl Display for Shape {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        write!(f, "{}", shape_string(&self.0))
84    }
85}
86
87impl Shape {
88    #[must_use]
89    pub fn new(dims: &[usize]) -> Self {
90        Self(dims.to_vec())
91    }
92
93    #[must_use]
94    pub fn get_dims(&self) -> &Vec<usize> {
95        &self.0
96    }
97}
98
99impl HasShape for Shape {
100    fn num_dims(&self) -> usize {
101        self.0.len()
102    }
103
104    fn num_elements(&self) -> usize {
105        self.0.iter().product()
106    }
107
108    fn get_dim(&self, total_dims: usize, i: usize) -> usize {
109        let dim_index = total_dims - i;
110        let rank = self.num_dims();
111        if dim_index <= rank {
112            self.0[rank - dim_index]
113        } else {
114            1
115        }
116    }
117
118    fn shape(&self) -> &Shape {
119        self
120    }
121}
122
123#[derive(Clone, Debug, PartialEq)]
124pub struct Offsets(Vec<usize>);
125
126impl Offsets {
127    #[must_use]
128    pub fn get_dims(&self) -> &Vec<usize> {
129        &self.0
130    }
131}
132
133#[derive(Clone, Debug)]
134pub struct Tensor {
135    id: Option<String>,
136    dtype: DataType,
137    shape: Shape,
138    addr: u64,
139}
140
141impl Tensor {
142    /// Create a tensor
143    #[must_use]
144    pub fn new(dims: &[usize], dtype: &DataType, addr: u64) -> Self {
145        Self {
146            id: None,
147            shape: Shape(dims.to_vec()),
148            dtype: *dtype,
149            addr,
150        }
151    }
152
153    #[must_use]
154    pub fn with_id(mut self, id: impl Into<String>) -> Self {
155        self.id = Some(id.into());
156        self
157    }
158
159    pub fn set_id(&mut self, id: impl Into<String>) {
160        self.id = Some(id.into());
161    }
162
163    #[must_use]
164    pub fn id(&self) -> Option<&str> {
165        self.id.as_deref()
166    }
167
168    /// Return the number of bytes this entire Tensor will consume in memory.
169    ///
170    /// This currently assumes it is optimally packed into memory bytes
171    #[must_use]
172    pub fn num_bytes(&self) -> usize {
173        (self.num_elements() * self.dtype.num_bits()).div_ceil(8)
174    }
175
176    #[must_use]
177    pub fn dtype(&self) -> &DataType {
178        &self.dtype
179    }
180
181    #[must_use]
182    pub fn addr(&self) -> u64 {
183        self.addr
184    }
185
186    pub fn set_addr(&mut self, addr: u64) {
187        self.addr = addr;
188    }
189}
190
191impl HasShape for Tensor {
192    fn num_dims(&self) -> usize {
193        self.shape.num_dims()
194    }
195
196    fn num_elements(&self) -> usize {
197        self.shape.num_elements()
198    }
199
200    fn get_dim(&self, total_dims: usize, i: usize) -> usize {
201        self.shape.get_dim(total_dims, i)
202    }
203
204    fn shape(&self) -> &Shape {
205        &self.shape
206    }
207}
208
209/// A view into a tensor
210#[derive(Clone, Debug)]
211pub struct TensorView {
212    tensor: Tensor,
213    shape: Shape,
214    offsets: Offsets,
215}
216
217impl TensorView {
218    /// Create a view into the given tensor
219    #[must_use]
220    pub fn new(tensor: Tensor, shape: &[usize], offsets: &[usize]) -> Self {
221        Self {
222            tensor,
223            shape: Shape(shape.to_vec()),
224            offsets: Offsets(offsets.to_vec()),
225        }
226    }
227
228    /// Create a view which is the full tensor
229    #[must_use]
230    pub fn new_full(tensor: Tensor) -> Self {
231        let shape = Shape(tensor.shape().get_dims().to_vec());
232        let offsets = Offsets(vec![0; tensor.num_dims()]);
233        Self {
234            tensor,
235            shape,
236            offsets,
237        }
238    }
239
240    #[must_use]
241    pub fn tensor(&self) -> &Tensor {
242        &self.tensor
243    }
244
245    #[must_use]
246    pub fn offsets(&self) -> &Offsets {
247        &self.offsets
248    }
249
250    #[must_use]
251    pub fn is_full_view(&self) -> bool {
252        self.shape == *self.tensor.shape() && self.offsets.0.iter().all(|offset| *offset == 0)
253    }
254
255    #[must_use]
256    pub fn from_output_partition(
257        tensor: Tensor,
258        output_rank: usize,
259        partition_dim: usize,
260        partition_offset: usize,
261        partition_len: usize,
262    ) -> Self {
263        Self::from_output_partitions(
264            tensor,
265            output_rank,
266            &[DimPartition {
267                dim: partition_dim,
268                offset: partition_offset,
269                len: partition_len,
270            }],
271        )
272    }
273
274    #[must_use]
275    pub fn from_output_partitions(
276        tensor: Tensor,
277        output_rank: usize,
278        partitions: &[DimPartition],
279    ) -> Self {
280        let base_view = Self::new_full(tensor);
281        Self::from_output_partitions_on_view(&base_view, output_rank, partitions)
282    }
283
284    /// Create a view by applying output partitions to an existing base view.
285    ///
286    /// This preserves the base view's offsets and shape constraints.
287    #[must_use]
288    pub fn from_output_partitions_on_view(
289        base_view: &TensorView,
290        output_rank: usize,
291        partitions: &[DimPartition],
292    ) -> Self {
293        let view_rank = base_view.num_dims();
294        let rank_pad = output_rank.saturating_sub(view_rank);
295        let mut shape = base_view.shape().get_dims().clone();
296        let mut offsets = base_view.offsets().get_dims().clone();
297
298        for partition in partitions {
299            if partition.dim < rank_pad {
300                continue;
301            }
302
303            let view_dim = partition.dim - rank_pad;
304            if view_dim < view_rank && shape[view_dim] > 1 {
305                offsets[view_dim] += partition.offset;
306                shape[view_dim] = partition.len;
307            }
308        }
309
310        Self::new(base_view.tensor().clone(), &shape, &offsets)
311    }
312
313    #[must_use]
314    pub fn num_bytes(&self) -> usize {
315        let dtype = self.tensor.dtype();
316        let num_bits = dtype.num_bits();
317        let num_elements = self.num_elements();
318        (num_bits * num_elements).div_ceil(8)
319    }
320
321    /// Return the offset of the first element (in number of elements)
322    pub fn element_offset(&self) -> Result<usize, SimError> {
323        let shape = &self.tensor.shape.0;
324        let offsets = &self.offsets.0;
325        if shape.len() != offsets.len() {
326            return sim_error!(
327                "shape rank {} does not match offset rank {}",
328                shape.len(),
329                offsets.len()
330            );
331        }
332
333        let mut stride = 1;
334        let mut total = 0;
335        for (dim, offset) in shape.iter().rev().zip(offsets.iter().rev()) {
336            if offset >= dim {
337                return sim_error!("offset {offset} is out of range for dimension of size {dim}");
338            }
339            total += offset * stride;
340            stride *= *dim;
341        }
342        Ok(total)
343    }
344}
345
346impl HasShape for TensorView {
347    fn num_dims(&self) -> usize {
348        self.shape.num_dims()
349    }
350
351    fn num_elements(&self) -> usize {
352        self.shape.num_elements()
353    }
354
355    fn get_dim(&self, total_dims: usize, i: usize) -> usize {
356        self.shape.get_dim(total_dims, i)
357    }
358
359    fn shape(&self) -> &Shape {
360        &self.shape
361    }
362}
363
364#[derive(Clone, Debug)]
365pub struct TensorPartition {
366    pub inputs: Vec<Option<TensorView>>,
367    pub outputs: Vec<Option<TensorView>>,
368}
369
370pub trait Operator {
371    /// Validate that the input and output tensors are valid shapes and
372    /// datatypes
373    fn validate_tensors(&self, inputs: &[Option<Tensor>], outputs: &[Option<Tensor>]) -> SimResult;
374
375    /// Returns the number of clock ticks needed to perform the
376    /// specified computation give the machine capabilities
377    fn compute_delay_ticks(
378        &self,
379        compute_capabilities: &Rc<ComputeCapabilities>,
380        inputs: &[Option<TensorView>],
381        outputs: &[Option<TensorView>],
382    ) -> Result<usize, SimError>;
383
384    /// Returns the total number of FLOPs performed by the specified
385    /// computation.
386    fn compute_flops(
387        &self,
388        inputs: &[Option<TensorView>],
389        outputs: &[Option<TensorView>],
390    ) -> Result<usize, SimError> {
391        Ok(self.compute_machine_ops(inputs, outputs)?.total())
392    }
393
394    /// Returns the number of machine operations performed by the specified
395    /// computation, broken down by operation type.
396    fn compute_machine_ops(
397        &self,
398        inputs: &[Option<TensorView>],
399        outputs: &[Option<TensorView>],
400    ) -> Result<MachineOpCounts, SimError>;
401
402    /// Partition the operation into one or more views that can be executed in
403    /// parallel. Implementations may return fewer than `num_partitions` if the
404    /// operator cannot be split that finely.
405    fn partition_views(
406        &self,
407        input_views: &[Option<TensorView>],
408        output_views: &[Option<TensorView>],
409        num_partitions: usize,
410    ) -> Result<Vec<TensorPartition>, SimError>;
411}
412
413/// Create partitions from full Tensors
414///
415/// This is a wrapper function to create TensorViews and then call
416/// `create_partitions` using those views.
417pub fn partition_tensors<T: Operator>(
418    operator: &T,
419    input_tensors: &[Option<Tensor>],
420    output_tensors: &[Option<Tensor>],
421    num_partitions: usize,
422) -> Result<Vec<TensorPartition>, SimError> {
423    let input_views = input_tensors
424        .iter()
425        .map(|maybe_tensor| {
426            maybe_tensor
427                .as_ref()
428                .map(|tensor| TensorView::new_full(tensor.clone()))
429        })
430        .collect::<Vec<_>>();
431    let output_views = output_tensors
432        .iter()
433        .map(|maybe_tensor| {
434            maybe_tensor
435                .as_ref()
436                .map(|tensor| TensorView::new_full(tensor.clone()))
437        })
438        .collect::<Vec<_>>();
439    operator.partition_views(&input_views, &output_views, num_partitions)
440}
441
442fn partition_into_ranges(total: usize, requested: usize) -> Vec<(usize, usize)> {
443    // Determine a valid number of partitions such that: total >= partitions >=1
444    let partitions = requested.clamp(1, total.max(1));
445
446    // All ranges get this number of entries
447    let base_range_size = total / partitions;
448
449    // This number of ranges get an extra 1
450    let remainder = total % partitions;
451
452    let mut start = 0;
453    let mut ranges = Vec::with_capacity(partitions);
454
455    for i in 0..partitions {
456        let len = base_range_size + usize::from(i < remainder);
457        if len == 0 {
458            continue;
459        }
460        ranges.push((start, len));
461        start += len;
462    }
463
464    if ranges.is_empty() {
465        ranges.push((0, total.max(1)));
466    }
467
468    ranges
469}
470
471#[derive(Clone, Debug, PartialEq, Eq)]
472pub struct DimPartition {
473    pub dim: usize,
474    pub offset: usize,
475    pub len: usize,
476}
477
478#[must_use]
479pub fn partition_across_dimensions(
480    dims: &[usize],
481    candidate_dims: &[usize],
482    requested: usize,
483) -> Vec<Vec<DimPartition>> {
484    let requested = requested.max(1);
485    let mut split_dims = Vec::new();
486    let mut achieved_partitions = 1usize;
487
488    for &dim in candidate_dims {
489        let dim_extent = dims[dim];
490        if dim_extent <= 1 {
491            continue;
492        }
493
494        let needed = requested.div_ceil(achieved_partitions).max(1);
495        let splits = dim_extent.min(needed);
496        if splits <= 1 {
497            continue;
498        }
499
500        split_dims.push((dim, partition_into_ranges(dim_extent, splits)));
501        achieved_partitions *= splits;
502        if achieved_partitions >= requested {
503            break;
504        }
505    }
506
507    if split_dims.is_empty() {
508        // In the case we are just requesting a single partition we just preserve the
509        // shape
510        let preserve_shape: Vec<DimPartition> = dims
511            .iter()
512            .enumerate()
513            .map(|(idx, dim)| DimPartition {
514                dim: idx,
515                offset: 0,
516                len: *dim,
517            })
518            .collect();
519        return vec![preserve_shape];
520    }
521
522    let mut partitions = vec![Vec::new()];
523    for (dim, ranges) in split_dims {
524        let mut next = Vec::with_capacity(partitions.len() * ranges.len());
525        for base in &partitions {
526            for (offset, len) in &ranges {
527                let mut partition = base.clone();
528                partition.push(DimPartition {
529                    dim,
530                    offset: *offset,
531                    len: *len,
532                });
533                next.push(partition);
534            }
535        }
536        partitions = next;
537    }
538
539    partitions
540}
541
542#[must_use]
543pub fn apply_dim_partitions(
544    dims: &[usize],
545    partitions: &[DimPartition],
546) -> (Vec<usize>, Vec<usize>) {
547    let mut shape = dims.to_vec();
548    let mut offsets = vec![0; dims.len()];
549
550    for partition in partitions {
551        shape[partition.dim] = partition.len;
552        offsets[partition.dim] = partition.offset;
553    }
554
555    (shape, offsets)
556}
557
558pub mod add;
559pub mod custom;
560pub mod gemm;
561pub mod maxpool;
gwr_models/processing_element/operators/mod.rs

gwr_models/processing_element/operators/
mod.rs
