#[cutile::module]
mod kernels {
use cutile::core::*;
#[cutile::entry()]
pub fn cast_f32_f16<const B: i32>(
out: &mut Tensor<f16, { [B] }>,
input: &Tensor<f32, { [-1] }>,
) {
let values: Tile<f32, { [B] }> = load_tile_like(input, out);
let values: Tile<f16, { [B] }> = convert_tile(values);
out.store(values);
}
#[cutile::entry()]
pub fn cast_f32_f64<const B: i32>(
out: &mut Tensor<f64, { [B] }>,
input: &Tensor<f32, { [-1] }>,
) {
let values: Tile<f32, { [B] }> = load_tile_like(input, out);
let values: Tile<f64, { [B] }> = convert_tile(values);
out.store(values);
}
#[cfg(feature = "dtype-bf16")]
#[cutile::entry()]
pub fn cast_f32_bf16<const B: i32>(
out: &mut Tensor<bf16, { [B] }>,
input: &Tensor<f32, { [-1] }>,
) {
let values: Tile<f32, { [B] }> = load_tile_like(input, out);
let values: Tile<bf16, { [B] }> = convert_tile(values);
out.store(values);
}
#[cutile::entry()]
pub fn cast_f32_u32<const B: i32>(
out: &mut Tensor<u32, { [B] }>,
input: &Tensor<f32, { [-1] }>,
) {
let values: Tile<f32, { [B] }> = load_tile_like(input, out);
out.store(ftoi(values, rounding::Zero));
}
#[cutile::entry()]
pub fn cast_f32_u8<const B: i32>(out: &mut Tensor<u8, { [B] }>, input: &Tensor<f32, { [-1] }>) {
let values: Tile<f32, { [B] }> = load_tile_like(input, out);
out.store(ftoi(values, rounding::Zero));
}
#[cutile::entry()]
pub fn cast_f32_i8<const B: i32>(out: &mut Tensor<i8, { [B] }>, input: &Tensor<f32, { [-1] }>) {
let values: Tile<f32, { [B] }> = load_tile_like(input, out);
out.store(ftoi(values, rounding::Zero));
}
#[cutile::entry()]
pub fn cast_f32_i32<const B: i32>(
out: &mut Tensor<i32, { [B] }>,
input: &Tensor<f32, { [-1] }>,
) {
let values: Tile<f32, { [B] }> = load_tile_like(input, out);
out.store(ftoi(values, rounding::Zero));
}
#[cutile::entry()]
pub fn cast_f32_u64<const B: i32>(
out: &mut Tensor<u64, { [B] }>,
input: &Tensor<f32, { [-1] }>,
) {
let values: Tile<f32, { [B] }> = load_tile_like(input, out);
out.store(ftoi(values, rounding::Zero));
}
#[cutile::entry()]
pub fn cast_f32_i64<const B: i32>(
out: &mut Tensor<i64, { [B] }>,
input: &Tensor<f32, { [-1] }>,
) {
let values: Tile<f32, { [B] }> = load_tile_like(input, out);
out.store(ftoi(values, rounding::Zero));
}
#[cutile::entry()]
pub fn cast_bool_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
input: &Tensor<u8, { [-1] }>,
) {
let values: Tile<u8, { [B] }> = load_tile_like(input, out);
let zero_u8: Tile<u8, { [B] }> = constant(0u8, out.shape());
let mask = cmpi(values, zero_u8, predicate::NotEqual);
let one: Tile<f32, { [B] }> = constant(1.0f32, out.shape());
let zero: Tile<f32, { [B] }> = constant(0.0f32, out.shape());
out.store(select(mask, one, zero));
}
#[cutile::entry()]
pub fn cast_u32_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
input: &Tensor<u32, { [-1] }>,
) {
let values: Tile<u32, { [B] }> = load_tile_like(input, out);
let values: Tile<f32, { [B] }> = convert_tile(values);
out.store(values);
}
#[cutile::entry()]
pub fn cast_i32_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
input: &Tensor<i32, { [-1] }>,
) {
let values: Tile<i32, { [B] }> = load_tile_like(input, out);
let values: Tile<f32, { [B] }> = convert_tile(values);
out.store(values);
}
#[cutile::entry()]
pub fn cast_u64_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
input: &Tensor<u64, { [-1] }>,
) {
let values: Tile<u64, { [B] }> = load_tile_like(input, out);
let values: Tile<f32, { [B] }> = convert_tile(values);
out.store(values);
}
#[cutile::entry()]
pub fn cast_i64_f32<const B: i32>(
out: &mut Tensor<f32, { [B] }>,
input: &Tensor<i64, { [-1] }>,
) {
let values: Tile<i64, { [B] }> = load_tile_like(input, out);
let values: Tile<f32, { [B] }> = convert_tile(values);
out.store(values);
}
#[cutile::entry()]
pub fn cast_u8_f32<const B: i32>(out: &mut Tensor<f32, { [B] }>, input: &Tensor<u8, { [-1] }>) {
let values: Tile<u8, { [B] }> = load_tile_like(input, out);
let values: Tile<f32, { [B] }> = convert_tile(values);
out.store(values);
}
#[cutile::entry()]
pub fn cast_i8_f32<const B: i32>(out: &mut Tensor<f32, { [B] }>, input: &Tensor<i8, { [-1] }>) {
let values: Tile<i8, { [B] }> = load_tile_like(input, out);
let values: Tile<f32, { [B] }> = convert_tile(values);
out.store(values);
}
#[cutile::entry()]
pub unsafe fn cast_f32_bool(out: *mut u8, input: *mut f32, len: i32) {
let pid: (i32, i32, i32) = get_tile_block_id();
let tile_shape = const_shape![128];
let offsets: Tile<i32, { [128] }> =
iota(tile_shape) + broadcast_scalar(pid.0 * 128i32, tile_shape);
let len_tile = broadcast_scalar(len, tile_shape);
let mask = cmpi(offsets, len_tile, predicate::LessThan);
let zero_offsets: Tile<i32, { [128] }> = constant(0i32, tile_shape);
let safe_offsets = select(mask, offsets, zero_offsets);
let input_base: PointerTile<*mut f32, { [] }> = pointer_to_tile(input);
let input_base: PointerTile<*mut f32, { [1] }> = input_base.reshape(const_shape![1]);
let input_ptrs: PointerTile<*mut f32, { [128] }> = input_base.broadcast(tile_shape);
let input_ptrs: PointerTile<*mut f32, { [128] }> = input_ptrs.offset_tile(safe_offsets);
let input_result: (Tile<f32, { [128] }>, Token) = load_ptr_tko(
input_ptrs,
ordering::Weak,
None::<scope::TileBlock>,
Some(mask),
Some(0.0f32),
None,
Latency::<0>,
);
let zero = constant(0.0f32, tile_shape);
let nonzero = cmpf(
input_result.0,
zero,
predicate::NotEqual,
cmp_ordering::Unordered,
);
let true_value: Tile<u8, { [128] }> = constant(1u8, tile_shape);
let false_value: Tile<u8, { [128] }> = constant(0u8, tile_shape);
let values = select(nonzero, true_value, false_value);
let out_base: PointerTile<*mut u8, { [] }> = pointer_to_tile(out);
let out_base: PointerTile<*mut u8, { [1] }> = out_base.reshape(const_shape![1]);
let out_ptrs: PointerTile<*mut u8, { [128] }> = out_base.broadcast(tile_shape);
let out_ptrs: PointerTile<*mut u8, { [128] }> = out_ptrs.offset_tile(offsets);
store_ptr_tko(
out_ptrs,
values,
ordering::Weak,
None::<scope::TileBlock>,
Some(mask),
None,
Latency::<0>,
);
}
}
pub use kernels::*;