cubecl-core 0.11.0-pre.3

CubeCL core create
Documentation 
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

use crate::{prelude::*, unexpanded};

/// Fused multiply-add `A*B+C`.
#[allow(unused_variables)]
pub fn fma<C: CubePrimitive>(a: C, b: C, c: C) -> C {
    unexpanded!()
}

/// Expand method of [`fma()`].
pub mod fma {
    use super::*;
    use cubecl_ir::{Scope, dialect::math::FmaOp};

    pub fn expand<C: CubePrimitive>(
        scope: &Scope,
        a: NativeExpand<C>,
        b: NativeExpand<C>,
        c: NativeExpand<C>,
    ) -> NativeExpand<C> {
        let a = a.read_value(scope);
        let b = b.read_value(scope);
        let c = c.read_value(scope);

        let [a, b, c] = normalize_same_vectorization(scope, [a, b, c]);

        let op = FmaOp::new(scope.ctx_mut(), a, b, c);
        scope.register_with_result(&op).into()
    }
}