Skip to main content

cubecl_tile/tile/kind/procedural/
sum.rs

1use cubecl::prelude::*;
2
3use super::{Reads, Recipe, RecipeCoords, RecipeExpand};
4
5/// Pointwise sum of two recipes: `(A + B)(coords) = A(coords) + B(coords)`.
6#[derive(CubeType, Clone)]
7pub struct Sum<A: CubeType, B: CubeType> {
8    pub lhs: A,
9    pub rhs: B,
10}
11
12/// Construct a [`Sum`].
13#[cube]
14pub fn sum_of<A: CubeType, B: CubeType>(lhs: A, rhs: B) -> Sum<A, B> {
15    Sum::<A, B> { lhs, rhs }
16}
17
18#[cube]
19impl<T: Numeric, A: Recipe<T>, B: Recipe<T>> Recipe<T> for Sum<A, B> {
20    fn evaluate(&self, coordinates: &RecipeCoords) -> T {
21        self.lhs.evaluate(coordinates) + self.rhs.evaluate(coordinates)
22    }
23}
24
25impl<A: CubeType, B: CubeType> Reads for SumExpand<A, B>
26where
27    A::ExpandType: Reads,
28    B::ExpandType: Reads,
29{
30    fn reads(&self, scope: &Scope, axis: crate::Axis) -> bool {
31        self.lhs.reads(scope, axis) || self.rhs.reads(scope, axis)
32    }
33}