cubek-tile 0.3.0

CubeK: Tile
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

use cubecl::prelude::*;

use super::{Reads, Recipe, RecipeCoords, RecipeExpand};

/// Pointwise sum of two recipes: `(A + B)(coords) = A(coords) + B(coords)`.
#[derive(CubeType, Clone)]
pub struct Sum<A: CubeType, B: CubeType> {
    pub lhs: A,
    pub rhs: B,
}

/// Construct a [`Sum`].
#[cube]
pub fn sum_of<A: CubeType, B: CubeType>(lhs: A, rhs: B) -> Sum<A, B> {
    Sum::<A, B> { lhs, rhs }
}

#[cube]
impl<T: Numeric, A: Recipe<T>, B: Recipe<T>> Recipe<T> for Sum<A, B> {
    fn evaluate(&self, coordinates: &RecipeCoords) -> T {
        self.lhs.evaluate(coordinates) + self.rhs.evaluate(coordinates)
    }
}

impl<A: CubeType, B: CubeType> Reads for SumExpand<A, B>
where
    A::ExpandType: Reads,
    B::ExpandType: Reads,
{
    fn reads(&self, scope: &Scope, axis: crate::Axis) -> bool {
        self.lhs.reads(scope, axis) || self.rhs.reads(scope, axis)
    }
}