#[cutile::module]
mod kernels {
use cutile::core::*;
pub fn scalar_store<const OP: i32, const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
let x: Tile<f16, { [B] }> = load_tile_like(input, out);
let scalar = broadcast_scalar(scalar, out.shape());
let value = if OP == 0 {
x + scalar
} else if OP == 1 {
x - scalar
} else if OP == 2 {
scalar - x
} else if OP == 3 {
x * scalar
} else if OP == 4 {
x / scalar
} else if OP == 5 {
scalar / x
} else if OP == 10 {
x - floor(x / scalar) * scalar
} else if OP == 11 {
scalar - floor(scalar / x) * x
} else if OP == 12 {
atan2(x, scalar)
} else if OP == 13 {
atan2(scalar, x)
} else if OP == 6 {
pow(x, scalar)
} else if OP == 7 {
pow(scalar, x)
} else if OP == 8 {
min_tile(x, scalar)
} else if OP == 9 {
max_tile(x, scalar)
} else if OP == 14 {
sqrt(
x * x + scalar * scalar,
rounding::NearestEven,
ftz::Disabled,
)
} else if OP == 15 {
let diff = x - scalar;
diff * diff
} else if OP == 16 {
let max = maxf(x, scalar, nan::Enabled, ftz::Disabled);
let neg_inf = constant(f16::from_f32(f32::NEG_INFINITY), out.shape());
let both_neg_inf = cmpf(x, neg_inf, predicate::Equal, cmp_ordering::Ordered)
& cmpf(scalar, neg_inf, predicate::Equal, cmp_ordering::Ordered);
let value = max + log(exp(x - max) + exp(scalar - max));
select(both_neg_inf, neg_inf, value)
} else if OP == 17 {
let zero = constant(f16::from_f32(0.0), out.shape());
let zero_x = cmpf(x, zero, predicate::Equal, cmp_ordering::Ordered);
select(zero_x, zero, x * log(scalar))
} else if OP == 18 {
let zero = constant(f16::from_f32(0.0), out.shape());
let zero_scalar = cmpf(scalar, zero, predicate::Equal, cmp_ordering::Ordered);
select(zero_scalar, zero, scalar * log(x))
} else {
x
};
out.store(value);
}
pub fn scalar_cmp_bool_store<const OP: i32, const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
let x: Tile<f16, { [B] }> = load_tile_like(input, out);
let scalar = broadcast_scalar(scalar, out.shape());
let one = constant(1u8, out.shape());
let zero = constant(0u8, out.shape());
let mask = if OP == 0 {
cmpf(x, scalar, predicate::Equal, cmp_ordering::Ordered)
} else if OP == 1 {
cmpf(x, scalar, predicate::NotEqual, cmp_ordering::Unordered)
} else if OP == 2 {
cmpf(x, scalar, predicate::LessThan, cmp_ordering::Ordered)
} else if OP == 3 {
cmpf(x, scalar, predicate::LessThanOrEqual, cmp_ordering::Ordered)
} else if OP == 4 {
cmpf(x, scalar, predicate::GreaterThan, cmp_ordering::Ordered)
} else {
cmpf(
x,
scalar,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
)
};
out.store(select(mask, one, zero));
}
#[cutile::entry()]
pub fn add_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<0i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn sub_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<1i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn rsub_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<2i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn scale_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<3i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn div_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<4i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn rdiv_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<5i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn modulo_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<10i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn rmodulo_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<11i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn atan2_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<12i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn ratan2_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<13i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn pow_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<6i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn rpow_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<7i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn hypot_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<14i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn squared_difference_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<15i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn logaddexp_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<16i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn xlogy_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<17i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn rxlogy_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<18i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn min_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<8i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn max_scalar_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_store::<9i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn equal_scalar_bool_f16<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_cmp_bool_store::<0i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn not_equal_scalar_bool_f16<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_cmp_bool_store::<1i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn less_scalar_bool_f16<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_cmp_bool_store::<2i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn less_equal_scalar_bool_f16<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_cmp_bool_store::<3i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn greater_scalar_bool_f16<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_cmp_bool_store::<4i32, B>(out, input, scalar);
}
#[cutile::entry()]
pub fn greater_equal_scalar_bool_f16<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
input: &Tensor<f16, { [-1] }>,
scalar: f16,
) {
scalar_cmp_bool_store::<5i32, B>(out, input, scalar);
}
}
pub use kernels::*;