#[cutile::module]
mod kernels {
use cutile::core::*;
pub fn unary_store<const OP: i32, const B: i32>(
out: &mut Tensor<f64, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
let x: Tile<f64, { [B] }> = load_tile_like(input, out);
let zero = constant(0.0f64, out.shape());
let one = constant(1.0f64, out.shape());
let half = constant(0.5f64, out.shape());
let third = constant(0.3333333333333333f64, out.shape());
let three = constant(3.0f64, out.shape());
let six = constant(6.0f64, out.shape());
let pi = constant(std::f64::consts::PI, out.shape());
let one_eighty = constant(180.0f64, out.shape());
let selu_alpha = constant(1.6732632423543772f64, out.shape());
let selu_scale = constant(1.0507009873554805f64, out.shape());
let value = if OP == 0 {
zero - x
} else if OP == 1 {
absf(x)
} else if OP == 2 {
sqrt(x, rounding::NearestEven, ftz::Disabled)
} else if OP == 3 {
exp(x)
} else if OP == 4 {
log(x)
} else if OP == 5 {
sin(x)
} else if OP == 6 {
cos(x)
} else if OP == 7 {
tan(x)
} else if OP == 8 {
sinh(x)
} else if OP == 9 {
cosh(x)
} else if OP == 10 {
tanh(x)
} else if OP == 11 {
atan2(x, sqrt(one - x * x, rounding::NearestEven, ftz::Disabled))
} else if OP == 12 {
atan2(sqrt(one - x * x, rounding::NearestEven, ftz::Disabled), x)
} else if OP == 13 {
atan2(x, one)
} else if OP == 14 {
ceil(x)
} else if OP == 15 {
floor(x)
} else if OP == 16 {
let positive = cmpf(
x,
zero,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
);
select(positive, floor(x), ceil(x))
} else if OP == 17 {
one / x
} else if OP == 18 {
x * x
} else if OP == 19 {
rsqrt(x, ftz::Disabled)
} else if OP == 20 {
log(x + sqrt(x * x + one, rounding::NearestEven, ftz::Disabled))
} else if OP == 21 {
log(x + sqrt(x * x - one, rounding::NearestEven, ftz::Disabled))
} else if OP == 22 {
half * log((one + x) / (one - x))
} else if OP == 23 {
max_tile(x, zero)
} else if OP == 24 {
one / (one + exp(zero - x))
} else if OP == 25 {
x / (one + exp(zero - x))
} else if OP == 26 {
let sqrt_2_over_pi = constant(0.7978846f64, out.shape());
let cubic_coeff = constant(0.044715f64, out.shape());
half * x * (one + tanh(sqrt_2_over_pi * (x + cubic_coeff * x * x * x)))
} else if OP == 27 {
let negative = cmpf(x, zero, predicate::LessThan, cmp_ordering::Ordered);
let positive = cmpf(x, zero, predicate::GreaterThan, cmp_ordering::Ordered);
select(negative, zero - one, select(positive, one, zero))
} else if OP == 28 {
log(one + x)
} else if OP == 29 {
exp(x) - one
} else if OP == 30 {
exp2(x, ftz::Disabled)
} else if OP == 31 {
log2(x)
} else if OP == 32 {
let ln_10 = constant(std::f64::consts::LN_10, out.shape());
log(x) / ln_10
} else if OP == 33 {
let root = pow(absf(x), third);
let negative = cmpf(x, zero, predicate::LessThan, cmp_ordering::Ordered);
select(negative, zero - root, root)
} else if OP == 34 {
max_tile(x, zero) + log(one + exp(zero - absf(x)))
} else if OP == 35 {
min_tile(max_tile((x + three) / six, zero), one)
} else if OP == 36 {
x * min_tile(max_tile((x + three) / six, zero), one)
} else if OP == 37 {
x * tanh(max_tile(x, zero) + log(one + exp(zero - absf(x))))
} else if OP == 38 {
let positive = cmpf(
x,
zero,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
);
selu_scale * select(positive, x, selu_alpha * (exp(x) - one))
} else if OP == 39 {
x / (one + absf(x))
} else if OP == 40 {
x * pi / one_eighty
} else if OP == 41 {
x * one_eighty / pi
} else if OP == 42 {
let positive = cmpf(
x,
zero,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
);
let truncated = select(positive, floor(x), ceil(x));
x - truncated
} else if OP == 43 {
min_tile(max_tile(x, zero), six)
} else if OP == 44 {
x - tanh(x)
} else if OP == 45 {
zero - (max_tile(zero - x, zero) + log(one + exp(zero - absf(zero - x))))
} else {
x
};
out.store(value);
}
pub fn nan_to_num_store<const B: i32>(
out: &mut Tensor<f64, { [B] }>,
input: &Tensor<f64, { [-1] }>,
nan: f64,
posinf: f64,
neginf: f64,
) {
let x: Tile<f64, { [B] }> = load_tile_like(input, out);
let zero = constant(0.0f64, out.shape());
let infinity = constant(f64::INFINITY, out.shape());
let abs_x = absf(x);
let numeric = cmpf(x, x, predicate::Equal, cmp_ordering::Ordered);
let infinite = cmpf(abs_x, infinity, predicate::Equal, cmp_ordering::Ordered);
let positive = cmpf(x, zero, predicate::GreaterThan, cmp_ordering::Ordered);
let nan = broadcast_scalar(nan, out.shape());
let posinf = broadcast_scalar(posinf, out.shape());
let neginf = broadcast_scalar(neginf, out.shape());
let inf_value = select(positive, posinf, neginf);
out.store(select(numeric, select(infinite, inf_value, x), nan));
}
pub fn unary_bool_store<const OP: i32, const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
let x: Tile<f64, { [B] }> = load_tile_like(input, out);
let one = constant(1u8, out.shape());
let zero = constant(0u8, out.shape());
let infinity = constant(f64::INFINITY, out.shape());
let abs_x = absf(x);
let mask = if OP == 0 {
cmpf(x, x, predicate::NotEqual, cmp_ordering::Unordered)
} else if OP == 1 {
cmpf(abs_x, infinity, predicate::Equal, cmp_ordering::Ordered)
} else if OP == 2 {
cmpf(x, infinity, predicate::Equal, cmp_ordering::Ordered)
} else if OP == 3 {
cmpf(
x,
constant(f64::NEG_INFINITY, out.shape()),
predicate::Equal,
cmp_ordering::Ordered,
)
} else if OP == 4 {
let numeric = cmpf(x, x, predicate::Equal, cmp_ordering::Ordered);
let not_infinite = cmpf(abs_x, infinity, predicate::NotEqual, cmp_ordering::Ordered);
numeric & not_infinite
} else if OP == 5 {
let numeric = cmpf(x, x, predicate::Equal, cmp_ordering::Ordered);
let not_infinite = cmpf(abs_x, infinity, predicate::NotEqual, cmp_ordering::Ordered);
let min_normal = constant(f64::MIN_POSITIVE, out.shape());
let normal_magnitude = cmpf(
abs_x,
min_normal,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
);
numeric & not_infinite & normal_magnitude
} else if OP == 10 {
let numeric = cmpf(x, x, predicate::Equal, cmp_ordering::Ordered);
let nonzero = cmpf(
x,
constant(0.0f64, out.shape()),
predicate::NotEqual,
cmp_ordering::Ordered,
);
let min_normal = constant(f64::MIN_POSITIVE, out.shape());
let subnormal_magnitude = cmpf(
abs_x,
min_normal,
predicate::LessThan,
cmp_ordering::Ordered,
);
numeric & nonzero & subnormal_magnitude
} else if OP == 6 {
cmpf(
x,
constant(0.0f64, out.shape()),
predicate::Equal,
cmp_ordering::Ordered,
)
} else if OP == 7 {
cmpf(
x,
constant(0.0f64, out.shape()),
predicate::GreaterThan,
cmp_ordering::Ordered,
)
} else if OP == 9 {
cmpf(
x,
constant(0.0f64, out.shape()),
predicate::NotEqual,
cmp_ordering::Unordered,
)
} else {
cmpf(
x,
constant(0.0f64, out.shape()),
predicate::LessThan,
cmp_ordering::Ordered,
)
};
out.store(select(mask, one, zero));
}
#[cutile::entry()]
pub fn neg_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<0i32, B>(out, input);
}
#[cutile::entry()]
pub fn abs_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<1i32, B>(out, input);
}
#[cutile::entry()]
pub fn sqrt_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<2i32, B>(out, input);
}
#[cutile::entry()]
pub fn exp_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<3i32, B>(out, input);
}
#[cutile::entry()]
pub fn exp2_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<30i32, B>(out, input);
}
#[cutile::entry()]
pub fn log_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<4i32, B>(out, input);
}
#[cutile::entry()]
pub fn log2_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<31i32, B>(out, input);
}
#[cutile::entry()]
pub fn log10_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<32i32, B>(out, input);
}
#[cutile::entry()]
pub fn sin_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<5i32, B>(out, input);
}
#[cutile::entry()]
pub fn cos_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<6i32, B>(out, input);
}
#[cutile::entry()]
pub fn tan_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<7i32, B>(out, input);
}
#[cutile::entry()]
pub fn sinh_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<8i32, B>(out, input);
}
#[cutile::entry()]
pub fn cosh_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<9i32, B>(out, input);
}
#[cutile::entry()]
pub fn tanh_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<10i32, B>(out, input);
}
#[cutile::entry()]
pub fn asin_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<11i32, B>(out, input);
}
#[cutile::entry()]
pub fn acos_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<12i32, B>(out, input);
}
#[cutile::entry()]
pub fn atan_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<13i32, B>(out, input);
}
#[cutile::entry()]
pub fn ceil_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<14i32, B>(out, input);
}
#[cutile::entry()]
pub fn floor_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<15i32, B>(out, input);
}
#[cutile::entry()]
pub fn trunc_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<16i32, B>(out, input);
}
#[cutile::entry()]
pub fn frac_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<42i32, B>(out, input);
}
#[cutile::entry()]
pub fn reciprocal_f64<const B: i32>(
out: &mut Tensor<f64, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_store::<17i32, B>(out, input);
}
#[cutile::entry()]
pub fn square_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<18i32, B>(out, input);
}
#[cutile::entry()]
pub fn rsqrt_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<19i32, B>(out, input);
}
#[cutile::entry()]
pub fn cbrt_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<33i32, B>(out, input);
}
#[cutile::entry()]
pub fn asinh_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<20i32, B>(out, input);
}
#[cutile::entry()]
pub fn acosh_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<21i32, B>(out, input);
}
#[cutile::entry()]
pub fn atanh_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<22i32, B>(out, input);
}
#[cutile::entry()]
pub fn sign_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<27i32, B>(out, input);
}
#[cutile::entry()]
pub fn nan_to_num_f64<const B: i32>(
out: &mut Tensor<f64, { [B] }>,
input: &Tensor<f64, { [-1] }>,
nan: f64,
posinf: f64,
neginf: f64,
) {
nan_to_num_store(out, input, nan, posinf, neginf);
}
#[cutile::entry()]
pub fn log1p_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<28i32, B>(out, input);
}
#[cutile::entry()]
pub fn expm1_f64<const B: i32>(out: &mut Tensor<f64, { [B] }>, input: &Tensor<f64, { [-1] }>) {
unary_store::<29i32, B>(out, input);
}
#[cutile::entry()]
pub fn deg2rad_f64<const B: i32>(
out: &mut Tensor<f64, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_store::<40i32, B>(out, input);
}
#[cutile::entry()]
pub fn rad2deg_f64<const B: i32>(
out: &mut Tensor<f64, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_store::<41i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_nan_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<0i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_inf_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<1i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_pos_inf_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<2i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_neg_inf_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<3i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_finite_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<4i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_normal_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<5i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_subnormal_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<10i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_zero_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<6i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_nonzero_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<9i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_positive_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<7i32, B>(out, input);
}
#[cutile::entry()]
pub fn is_negative_bool_f64<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f64, { [-1] }>,
) {
unary_bool_store::<8i32, B>(out, input);
}
}
pub use kernels::*;