#[cutile::module]
mod kernels {
use cutile::core::*;
pub fn reduce_impl<const OP: i32, const BN: i32>(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
let pid: (i32, i32, i32) = get_tile_block_id();
let tile_shape = const_shape![1i32, BN];
let input_tile_raw: Tile<f32, { [1, BN] }> =
input.partition(tile_shape).load([pid.0, 0i32]);
let col_iota: Tile<i32, { [BN] }> = iota(const_shape![BN]);
let col_iota = col_iota.reshape(tile_shape);
let input_shape: [i32; 2] = get_tensor_shape(input);
let cols_tile = broadcast_scalar(input_shape[1], tile_shape);
let valid_cols = cmpi(col_iota, cols_tile, predicate::LessThan);
let value: Tile<f32, { [1] }> = if OP == 0 || OP == 1 {
let zero = constant(0.0f32, tile_shape);
let input_tile = select(valid_cols, input_tile_raw, zero);
let sum: Tile<f32, { [1] }> = reduce_sum(input_tile, 1i32);
if OP == 1 {
let divisor = broadcast_scalar(divisor, const_shape![1]);
sum / divisor
} else {
sum
}
} else if OP == 5 || OP == 6 {
let zero = constant(0.0f32, tile_shape);
let input_tile = select(valid_cols, input_tile_raw, zero);
let divisor = broadcast_scalar(divisor, const_shape![1]);
let sum: Tile<f32, { [1] }> = reduce_sum(input_tile, 1i32);
let mean: Tile<f32, { [1] }> = sum / divisor;
let mean_broadcast = mean.reshape(const_shape![1i32, 1i32]).broadcast(tile_shape);
let diff = select(valid_cols, input_tile_raw - mean_broadcast, zero);
let variance_sum: Tile<f32, { [1] }> = reduce_sum(diff * diff, 1i32);
let variance: Tile<f32, { [1] }> = variance_sum / divisor;
if OP == 6 {
sqrt(variance, rounding::NearestEven, ftz::Disabled)
} else {
variance
}
} else if OP == 2 {
let fill = constant(f32::NEG_INFINITY, tile_shape);
reduce_max(select(valid_cols, input_tile_raw, fill), 1i32)
} else if OP == 3 {
let fill = constant(f32::INFINITY, tile_shape);
reduce_min(select(valid_cols, input_tile_raw, fill), 1i32)
} else {
let one = constant(1.0f32, tile_shape);
reduce_prod(select(valid_cols, input_tile_raw, one), 1i32)
};
out.store(value.reshape(const_shape![1i32, 1i32]));
}
pub fn arg_reduce_impl<const OP: i32, const BN: i32>(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
let pid: (i32, i32, i32) = get_tile_block_id();
let tile_shape = const_shape![1i32, BN];
let input_tile_raw: Tile<f32, { [1, BN] }> =
input.partition(tile_shape).load([pid.0, 0i32]);
let col_iota: Tile<i32, { [BN] }> = iota(const_shape![BN]);
let col_iota = col_iota.reshape(tile_shape);
let input_shape: [i32; 2] = get_tensor_shape(input);
let cols_tile = broadcast_scalar(input_shape[1], tile_shape);
let valid_cols = cmpi(col_iota, cols_tile, predicate::LessThan);
let numeric = cmpf(
input_tile_raw,
input_tile_raw,
predicate::Equal,
cmp_ordering::Ordered,
);
let one_i32: Tile<i32, { [1, BN] }> = constant(1i32, tile_shape);
let zero_i32: Tile<i32, { [1, BN] }> = constant(0i32, tile_shape);
let numeric_i32 = select(valid_cols & numeric, one_i32, zero_i32);
let numeric_count: Tile<i32, { [1] }> = reduce_sum(numeric_i32, 1i32);
let zero_row: Tile<i32, { [1, 1] }> = constant(0i32, const_shape![1i32, 1i32]);
let has_numeric = cmpi(
numeric_count.reshape(const_shape![1i32, 1i32]),
zero_row,
predicate::GreaterThan,
)
.broadcast(tile_shape);
let no_numeric = cmpi(
numeric_count.reshape(const_shape![1i32, 1i32]),
zero_row,
predicate::Equal,
)
.broadcast(tile_shape);
let candidates = if OP == 0 {
let fill: Tile<f32, { [1, BN] }> = constant(f32::NEG_INFINITY, tile_shape);
select(valid_cols & numeric, input_tile_raw, fill)
} else {
let fill: Tile<f32, { [1, BN] }> = constant(f32::INFINITY, tile_shape);
select(valid_cols & numeric, input_tile_raw, fill)
};
let best_value: Tile<f32, { [1] }> = if OP == 0 {
reduce_max(candidates, 1i32)
} else {
reduce_min(candidates, 1i32)
};
let best_value = best_value
.reshape(const_shape![1i32, 1i32])
.broadcast(tile_shape);
let best_numeric = cmpf(
input_tile_raw,
best_value,
predicate::Equal,
cmp_ordering::Ordered,
);
let best_nan = no_numeric & valid_cols;
let best = (has_numeric & valid_cols & numeric & best_numeric) | best_nan;
let sentinel: Tile<i32, { [1, BN] }> = broadcast_scalar(input_shape[1], tile_shape);
let best_indices = select(best, col_iota, sentinel);
let best_index: Tile<i32, { [1] }> = reduce_min(best_indices, 1i32);
let lane_iota: Tile<i32, { [2] }> = iota(const_shape![2]);
let lane_iota = lane_iota.reshape(const_shape![1i32, 2i32]);
let zero_pair: Tile<i32, { [1, 2] }> = constant(0i32, const_shape![1i32, 2i32]);
let first_lane = cmpi(lane_iota, zero_pair, predicate::Equal);
let best_index = best_index
.reshape(const_shape![1i32, 1i32])
.broadcast(const_shape![1i32, 2i32]);
out.store(select(first_lane, best_index, zero_pair));
}
#[cutile::entry()]
pub fn reduce_sum_f32_1(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<0i32, 1i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_sum_f32_8(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<0i32, 8i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_sum_f32_32(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<0i32, 32i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_sum_f32_64(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<0i32, 64i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_sum_f32_128(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<0i32, 128i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_sum_f32_256(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<0i32, 256i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_sum_f32_512(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<0i32, 512i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_sum_f32_1024(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<0i32, 1024i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_mean_f32_1(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
reduce_impl::<1i32, 1i32>(out, input, divisor);
}
#[cutile::entry()]
pub fn reduce_mean_f32_8(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
reduce_impl::<1i32, 8i32>(out, input, divisor);
}
#[cutile::entry()]
pub fn reduce_mean_f32_32(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
reduce_impl::<1i32, 32i32>(out, input, divisor);
}
#[cutile::entry()]
pub fn reduce_mean_f32_64(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
reduce_impl::<1i32, 64i32>(out, input, divisor);
}
#[cutile::entry()]
pub fn reduce_mean_f32_128(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
reduce_impl::<1i32, 128i32>(out, input, divisor);
}
#[cutile::entry()]
pub fn reduce_mean_f32_256(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
reduce_impl::<1i32, 256i32>(out, input, divisor);
}
#[cutile::entry()]
pub fn reduce_mean_f32_512(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
reduce_impl::<1i32, 512i32>(out, input, divisor);
}
#[cutile::entry()]
pub fn reduce_mean_f32_1024(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
reduce_impl::<1i32, 1024i32>(out, input, divisor);
}
#[cutile::entry()]
pub fn reduce_max_f32_1(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<2i32, 1i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_max_f32_8(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<2i32, 8i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_max_f32_32(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<2i32, 32i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_max_f32_64(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<2i32, 64i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_max_f32_128(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<2i32, 128i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_max_f32_256(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<2i32, 256i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_max_f32_512(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<2i32, 512i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_max_f32_1024(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<2i32, 1024i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_min_f32_1(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<3i32, 1i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_min_f32_8(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<3i32, 8i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_min_f32_32(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<3i32, 32i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_min_f32_64(out: &mut Tensor<f32, { [1, 1] }>, input: &Tensor<f32, { [-1, -1] }>) {
reduce_impl::<3i32, 64i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_min_f32_128(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<3i32, 128i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_min_f32_256(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<3i32, 256i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_min_f32_512(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<3i32, 512i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_min_f32_1024(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<3i32, 1024i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_product_f32_1(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<4i32, 1i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_product_f32_8(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<4i32, 8i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_product_f32_32(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<4i32, 32i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_product_f32_64(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<4i32, 64i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_product_f32_128(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<4i32, 128i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_product_f32_256(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<4i32, 256i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_product_f32_512(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<4i32, 512i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_product_f32_1024(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
reduce_impl::<4i32, 1024i32>(out, input, 1.0f32);
}
#[cutile::entry()]
pub fn reduce_variance_f32<const BN: i32>(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
reduce_impl::<5i32, BN>(out, input, divisor);
}
#[cutile::entry()]
pub fn reduce_std_f32<const BN: i32>(
out: &mut Tensor<f32, { [1, 1] }>,
input: &Tensor<f32, { [-1, -1] }>,
divisor: f32,
) {
reduce_impl::<6i32, BN>(out, input, divisor);
}
#[cutile::entry()]
pub fn reduce_argmax_f32_1(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<0i32, 1i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmax_f32_8(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<0i32, 8i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmax_f32_128(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<0i32, 128i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmax_f32_256(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<0i32, 256i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmax_f32_512(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<0i32, 512i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmax_f32_1024(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<0i32, 1024i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmin_f32_1(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<1i32, 1i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmin_f32_8(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<1i32, 8i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmin_f32_128(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<1i32, 128i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmin_f32_256(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<1i32, 256i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmin_f32_512(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<1i32, 512i32>(out, input);
}
#[cutile::entry()]
pub fn reduce_argmin_f32_1024(
out: &mut Tensor<i32, { [1, 2] }>,
input: &Tensor<f32, { [-1, -1] }>,
) {
arg_reduce_impl::<1i32, 1024i32>(out, input);
}
}
pub use kernels::*;