Skip to main content

cubecl_tile/tile/kind/procedural/
product.rs

1use cubecl::prelude::*;
2
3use super::{Reads, Recipe, RecipeCoords, RecipeExpand};
4
5/// Pointwise product of two recipes: `(A * B)(coords) = A(coords) * B(coords)`.
6#[derive(CubeType, Clone)]
7pub struct Product<A: CubeType, B: CubeType> {
8    pub lhs: A,
9    pub rhs: B,
10}
11
12/// Construct a [`Product`].
13#[cube]
14pub fn product_of<A: CubeType, B: CubeType>(lhs: A, rhs: B) -> Product<A, B> {
15    Product::<A, B> { lhs, rhs }
16}
17
18#[cube]
19impl<T: Numeric, A: Recipe<T>, B: Recipe<T>> Recipe<T> for Product<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 ProductExpand<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}