cubecl_cpp/metal/
unary.rs

1use cubecl_core::{
2    self as cubecl,
3    frontend::polyfills::expm1,
4    ir::{
5        CanMaterialize, Pure, cube_op,
6        dialect::{
7            bitwise::*,
8            general::ReinterpretCastOp,
9            math::{Expm1Op, FAbsOp, TanhOp},
10        },
11        interfaces::TypedExt,
12        prelude::op_traits,
13    },
14    prelude::*,
15};
16use pliron::value::Value;
17
18use crate::{
19    metal::metal_op_with_out,
20    shared::{ty::TypeExtCPP, unary::lower_target_unop, unroll::unrolling},
21    target::Metal,
22};
23
24metal_op_with_out!(FAbsOp, |op, ctx| {
25    format!("abs({})", op.input(ctx).name(ctx))
26});
27
28metal_op_with_out!(CountOnesOp, |op, ctx| {
29    format!("popcount({})", op.input(ctx).name(ctx))
30});
31
32metal_op_with_out!(ReverseBitsOp, |op, ctx| {
33    format!("reverse_bits({})", op.input(ctx).name(ctx))
34});
35
36metal_op_with_out!(LeadingZerosBitsOp, |op, ctx| {
37    format!("clz({})", op.input(ctx).name(ctx))
38});
39
40unrolling!(TrailingZerosBitsOp);
41metal_op_with_out!(TrailingZerosBitsOp, |op, ctx| {
42    format!("ctz({})", op.input(ctx).name(ctx))
43});
44
45metal_op_with_out!(ReinterpretCastOp, |op, ctx| {
46    let input = op.input(ctx);
47    let ty = op.get_result(ctx).get_type(ctx).to_cpp(ctx);
48    if input.is_ptr(ctx) {
49        format!("reinterpret_cast<{ty}>({})", input.name(ctx))
50    } else {
51        format!("reinterpret_cast<const thread {ty}&>({})", input.name(ctx))
52    }
53});
54
55lower_target_unop!(Expm1Op, expm1, Metal);
56lower_target_unop!(TanhOp, safe_tanh, Metal);
57
58#[cube_op(name = "msl.tanh")]
59#[result_ty(same_as = input)]
60#[op_traits(Pure, CanMaterialize)]
61pub struct MslTanhOp {
62    pub input: Value,
63}
64
65unrolling!(MslTanhOp);
66metal_op_with_out!(MslTanhOp, |op, ctx| {
67    format!("tanh({})", op.input(ctx).name(ctx))
68});
69
70/// use the simple version because otherwise we'd get an infinite lowering loop
71#[cube]
72fn simple_tanh<T: Float, N: Size>(input: Vector<T, N>) -> Vector<T, N> {
73    intrinsic!(|scope| {
74        let input = input.read_value(scope);
75        let tanh = MslTanhOp::new(scope.ctx_mut(), input);
76        scope.register_with_result(&tanh).into()
77    })
78}
79
80#[cube]
81fn safe_tanh<T: Float, N: Size>(x: Vector<T, N>) -> Vector<T, N> {
82    let threshold = Vector::new(T::new(43.0_f32));
83    select(x > threshold, Vector::one(), simple_tanh(x))
84}
cubecl_cpp/metal/unary.rs

cubecl_cpp/metal/
unary.rs
