cubek-interpolate 0.3.0

use super::{
    filter::{SeparableFilter, SeparableFilterFamily, SeparableWeights, TapDistance},
    space::*,
};
use crate::InputStage;
use cubecl::{ir::ElemType, prelude::*};
use cubek_tile::{
    Axis, Partitioning, Region, RegisterBlock, Semiring, StageStorage, Stages, Tile, TileArg,
    procedural::{Factors, Phase, Procedural, affine_along, sum_of},
};

/// The distance from a tap to the source coordinate the output position lands on.
///
/// The rational is comptime because it is the same one the gather's projection already carries
/// ([`Rational::tap_axis`](super::coordinate::Rational::tap_axis)), and the kernel is already
/// compiled per shape. Handed as launch scalars instead, its divisor reaches [`Phase`] as a
/// runtime value: nothing folds, and every output position pays a true integer division for the
/// residue and a float division to normalize it, per resampled axis.
#[cube]
fn tap_distance<E: Float>(
    #[comptime] tap: Axis,
    #[comptime] output: Axis,
    #[comptime] scale: u32,
    #[comptime] offset: i32,
    #[comptime] divisor: u32,
    #[comptime] radius: usize,
) -> TapDistance<E> {
    sum_of(
        affine_along(tap, E::new(-(radius as f32)), E::new(1.0_f32)),
        Phase::<E> {
            coefficient: E::new(-1.0_f32),
            numerator_scale: comptime!(scale),
            numerator_offset: comptime!(offset),
            divisor: comptime!(divisor),
            axis: output,
        },
    )
}

#[cube(launch)]
pub fn interpolate_tile_kernel<E: Float, V: Size, F: SeparableFilterFamily>(
    input: &TileArg<'_, E, V>,
    output: &TileArg<'_, E, V>,
    space: Partitioning,
    #[comptime] row_scale: u32,
    #[comptime] row_offset: i32,
    #[comptime] row_divisor: u32,
    #[comptime] col_scale: u32,
    #[comptime] col_offset: i32,
    #[comptime] col_divisor: u32,
    #[comptime] radius: usize,
    #[comptime] stage: InputStage,
    #[comptime] padded: Option<usize>,
    #[comptime] config: RegisterBlock,
    #[define(E)] _dtype: ElemType,
) {
    let input = input.tile(comptime!(space.clone()));

    let row = tap_distance(TAP_H, OUTPUT_H, row_scale, row_offset, row_divisor, radius);
    let col = tap_distance(TAP_W, OUTPUT_W, col_scale, col_offset, col_divisor, radius);
    let mut factors = Sequence::new();
    factors.push(F::Filter::<E>::along(row));
    factors.push(F::Filter::<E>::along(col));
    let weights = Procedural::<E>::separable::<SeparableWeights<E, F::Filter<E>>>(
        comptime!(
            space
                .space()
                .subspace(&[BATCH, OUTPUT_H, OUTPUT_W, TAP_H, TAP_W])
        ),
        Factors::new(factors),
    );
    let weights = match comptime!(F::NORMALIZATION) {
        Some((taps, guard)) => weights.normalized(comptime!(taps), comptime!(guard)),
        None => weights,
    }
    .tile();

    let output = output.tile(comptime!(space.clone()));

    // This cube's box of the output, walked one channel block at a time. Whether the input is
    // staged into shared memory for each block is the launch's call on the window's size, stated
    // as `stage`; the walk is the same either way.
    for cube in space {
        let blocks = cube.walk();
        match comptime!(stage) {
            InputStage::Smem => {
                let mut stages =
                    Stages::smem_single_at(&blocks, &input, StageStorage::Strided, padded, 1usize);
                stages.pipelined(blocks, |slot, block| {
                    let output_block = output.at(block);
                    let weights_block = weights.at(block);
                    slot.consume(|input_block| {
                        interpolate_block(
                            block,
                            &output_block,
                            &weights_block,
                            input_block,
                            config,
                        );
                    });
                });
            }
            InputStage::InPlace => {
                for block in blocks {
                    interpolate_block(
                        &block,
                        &output.at(&block),
                        &weights.at(&block),
                        &input.at(&block),
                        config,
                    );
                }
            }
        }
    }
}

/// One block of the cube's walk: this plane's rows, then this unit's columns and channel
/// lines, each output cell contracting its whole tap window at the leaf under `config`.
#[cube]
fn interpolate_block<E: Float>(
    block: &Region,
    output: &Tile<E>,
    weights: &Tile<E>,
    input: &Tile<E>,
    #[comptime] config: RegisterBlock,
) {
    for plane in block {
        let output_plane = output.at(&plane);
        let weights_plane = weights.at(&plane);
        let input_plane = input.at(&plane);
        for cell in plane {
            let mut output_cell = output_plane
                .at(&cell)
                .accumulating(config, Semiring::SUM_PROD);
            output_cell.mm(&weights_plane.at(&cell), &input_plane.at(&cell));
        }
    }
}