#[cutile::module]
mod kernels {
use cutile::core::*;
pub fn binary_store<const OP: i32, const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
let lhs_tile: Tile<f32, { [B] }> = load_tile_like(lhs, out);
let rhs_tile: Tile<f32, { [B] }> = load_tile_like(rhs, out);
let value = if OP == 0 {
lhs_tile + rhs_tile
} else if OP == 1 {
lhs_tile - rhs_tile
} else if OP == 2 {
lhs_tile * rhs_tile
} else if OP == 3 {
lhs_tile / rhs_tile
} else if OP == 4 {
maxf(lhs_tile, rhs_tile, nan::Enabled, ftz::Disabled)
} else if OP == 5 {
minf(lhs_tile, rhs_tile, nan::Enabled, ftz::Disabled)
} else if OP == 6 {
atan2(lhs_tile, rhs_tile)
} else if OP == 7 {
lhs_tile - floor(lhs_tile / rhs_tile) * rhs_tile
} else if OP == 8 {
pow(lhs_tile, rhs_tile)
} else if OP == 9 {
sqrt(
lhs_tile * lhs_tile + rhs_tile * rhs_tile,
rounding::NearestEven,
ftz::Disabled,
)
} else if OP == 10 {
let diff = lhs_tile - rhs_tile;
diff * diff
} else if OP == 11 {
let max = maxf(lhs_tile, rhs_tile, nan::Enabled, ftz::Disabled);
let neg_inf = constant(f32::NEG_INFINITY, out.shape());
let both_neg_inf = cmpf(lhs_tile, neg_inf, predicate::Equal, cmp_ordering::Ordered)
& cmpf(rhs_tile, neg_inf, predicate::Equal, cmp_ordering::Ordered);
let value = max + log(exp(lhs_tile - max) + exp(rhs_tile - max));
select(both_neg_inf, neg_inf, value)
} else if OP == 12 {
let zero = constant(0.0f32, out.shape());
let zero_lhs = cmpf(lhs_tile, zero, predicate::Equal, cmp_ordering::Ordered);
select(zero_lhs, zero, lhs_tile * log(rhs_tile))
} else {
lhs_tile
};
out.store(value);
}
pub fn cmp_store<const OP: i32, const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
let lhs_tile: Tile<f32, { [B] }> = load_tile_like(lhs, out);
let rhs_tile: Tile<f32, { [B] }> = load_tile_like(rhs, out);
let one = constant(1.0f32, out.shape());
let zero = constant(0.0f32, out.shape());
let mask = if OP == 0 {
cmpf(lhs_tile, rhs_tile, predicate::Equal, cmp_ordering::Ordered)
} else if OP == 1 {
cmpf(
lhs_tile,
rhs_tile,
predicate::NotEqual,
cmp_ordering::Unordered,
)
} else if OP == 2 {
cmpf(
lhs_tile,
rhs_tile,
predicate::LessThan,
cmp_ordering::Ordered,
)
} else if OP == 3 {
cmpf(
lhs_tile,
rhs_tile,
predicate::LessThanOrEqual,
cmp_ordering::Ordered,
)
} else if OP == 4 {
cmpf(
lhs_tile,
rhs_tile,
predicate::GreaterThan,
cmp_ordering::Ordered,
)
} else {
cmpf(
lhs_tile,
rhs_tile,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
)
};
out.store(select(mask, one, zero));
}
pub fn cmp_bool_store<const OP: i32, const B: i32>(
out: &mut Tensor<u8, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
let lhs_tile: Tile<f32, { [B] }> = load_tile_like(lhs, out);
let rhs_tile: Tile<f32, { [B] }> = load_tile_like(rhs, out);
let one = constant(1u8, out.shape());
let zero = constant(0u8, out.shape());
let mask = if OP == 0 {
cmpf(lhs_tile, rhs_tile, predicate::Equal, cmp_ordering::Ordered)
} else if OP == 1 {
cmpf(
lhs_tile,
rhs_tile,
predicate::NotEqual,
cmp_ordering::Unordered,
)
} else if OP == 2 {
cmpf(
lhs_tile,
rhs_tile,
predicate::LessThan,
cmp_ordering::Ordered,
)
} else if OP == 3 {
cmpf(
lhs_tile,
rhs_tile,
predicate::LessThanOrEqual,
cmp_ordering::Ordered,
)
} else if OP == 4 {
cmpf(
lhs_tile,
rhs_tile,
predicate::GreaterThan,
cmp_ordering::Ordered,
)
} else {
cmpf(
lhs_tile,
rhs_tile,
predicate::GreaterThanOrEqual,
cmp_ordering::Ordered,
)
};
out.store(select(mask, one, zero));
}
#[cutile::entry()]
pub fn add_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<0i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn sub_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<1i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn mul_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<2i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn div_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<3i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn pow_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<8i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn fma_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
acc: &Tensor<f32, { [-1] }>,
) {
let lhs_tile: Tile<f32, { [B] }> = load_tile_like(lhs, out);
let rhs_tile: Tile<f32, { [B] }> = load_tile_like(rhs, out);
let acc_tile: Tile<f32, { [B] }> = load_tile_like(acc, out);
out.store(fma(
lhs_tile,
rhs_tile,
acc_tile,
rounding::NearestEven,
ftz::Disabled,
));
}
#[cutile::entry()]
pub fn maximum_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<4i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn minimum_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<5i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn atan2_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<6i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn modulo_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<7i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn hypot_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<9i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn squared_difference_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<10i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn logaddexp_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<11i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn xlogy_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
binary_store::<12i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn equal_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_store::<0i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn not_equal_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_store::<1i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn less_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_store::<2i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn less_equal_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_store::<3i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn greater_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_store::<4i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn greater_equal_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_store::<5i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn equal_bool_f32<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_bool_store::<0i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn not_equal_bool_f32<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_bool_store::<1i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn less_bool_f32<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_bool_store::<2i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn less_equal_bool_f32<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_bool_store::<3i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn greater_bool_f32<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_bool_store::<4i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn greater_equal_bool_f32<const B: i32>(
out: &mut Tensor<u8, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
) {
cmp_bool_store::<5i32, B>(out, lhs, rhs);
}
#[cutile::entry()]
pub fn clamp_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
input: &Tensor<f32, { [-1] }>,
min_value: f32,
max_value: f32,
) {
let x: Tile<f32, { [B] }> = load_tile_like(input, out);
let min_value = broadcast_scalar(min_value, out.shape());
let max_value = broadcast_scalar(max_value, out.shape());
out.store(min_tile(max_tile(x, min_value), max_value));
}
#[cutile::entry()]
pub fn lerp_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
lhs: &Tensor<f32, { [-1] }>,
rhs: &Tensor<f32, { [-1] }>,
scalar: f32,
) {
let lhs_tile: Tile<f32, { [B] }> = load_tile_like(lhs, out);
let rhs_tile: Tile<f32, { [B] }> = load_tile_like(rhs, out);
let scalar = broadcast_scalar(scalar, out.shape());
out.store(lhs_tile + scalar * (rhs_tile - lhs_tile));
}
#[cutile::entry()]
pub fn where_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
condition: &Tensor<f32, { [-1] }>,
x: &Tensor<f32, { [-1] }>,
y: &Tensor<f32, { [-1] }>,
) {
let condition_tile: Tile<f32, { [B] }> = load_tile_like(condition, out);
let x_tile: Tile<f32, { [B] }> = load_tile_like(x, out);
let y_tile: Tile<f32, { [B] }> = load_tile_like(y, out);
let zero = constant(0.0f32, out.shape());
let mask = cmpf(
condition_tile,
zero,
predicate::NotEqual,
cmp_ordering::Ordered,
);
out.store(select(mask, x_tile, y_tile));
}
#[cutile::entry()]
pub fn where_bool_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
condition: &Tensor<u8, { [-1] }>,
x: &Tensor<f32, { [-1] }>,
y: &Tensor<f32, { [-1] }>,
) {
let condition_tile: Tile<u8, { [B] }> = load_tile_like(condition, out);
let x_tile: Tile<f32, { [B] }> = load_tile_like(x, out);
let y_tile: Tile<f32, { [B] }> = load_tile_like(y, out);
let zero = constant(0u8, out.shape());
let mask = cmpi(condition_tile, zero, predicate::NotEqual);
out.store(select(mask, x_tile, y_tile));
}
#[cutile::entry()]
pub fn masked_fill_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
input: &Tensor<f32, { [-1] }>,
mask: &Tensor<u8, { [-1] }>,
value: f32,
) {
let input_tile: Tile<f32, { [B] }> = load_tile_like(input, out);
let mask_tile: Tile<u8, { [B] }> = load_tile_like(mask, out);
let zero = constant(0u8, out.shape());
let mask = cmpi(mask_tile, zero, predicate::NotEqual);
let value = broadcast_scalar(value, out.shape());
out.store(select(mask, value, input_tile));
}
}
pub use kernels::*;