ruDNN 0.1.103

use {ruda_core::tensor::spatial::ConvOptions};
use {ruda_core::tensor::Shape};
use {ruda_kernel::dsl::ir::StorageType, ruda_kernel::dsl::tune::LocalTuner, ruda_kernel::dsl::tune::Tunable, ruda_kernel::dsl::tune::TunableSet, ruda_kernel::dsl::tune::anchor};
use {crate::convolution::AcceleratedTileKind};

use {crate::convolution::tensor::tune_key::ConvTuneKey, ruda_kernel::dsl::Runtime, ruda_kernel::dsl::RudaTuneId, crate::convolution::tensor::ConvAutotuneKey, crate::convolution::tensor::backward_data::fallback::conv_data_backward_fallback, crate::convolution::tensor::backward_data::implicit_gemm::*, ruda_kernel::tensor::RudaTensor};

/// Executes autotune on conv2d operations
pub fn dgrad_autotune<R: Runtime, const N: usize>(
    out_grad: RudaTensor<R>,
    weights: RudaTensor<R>,
    input_shape: Shape,
    options: ConvOptions<N>,
) -> RudaTensor<R> {
    let client = out_grad.client.clone();

    static TUNER: LocalTuner<ConvTuneKey, RudaTuneId> = LocalTuner::new("ruda_tensor_device::kernel::conv::backward_data::tune::strict_f32_v1");

    // Note: TMA isn't currently implemented properly, and will always error.
    // It's kept here so it gets automatically enabled as soon as the convolution kernels update.
    // No CMMA for TMA because swizzling will be mandatory for good performance on dgrad.
    let tune_id = RudaTuneId::new(&out_grad.client, &out_grad.device);
    let tunables = TUNER.init_for_device(&tune_id, || {
        TunableSet::new(create_key::<R, N>, create_wgrad_input::<R, N>)
            .with(Tunable::new(
                "wgrad_fallback",
                |(out_grad, weights, input_shape, options)| {
                    conv_data_backward_fallback::<R, N>(out_grad, weights, input_shape, options)
                },
            ))
            .with(Tunable::new(
                "simple_sync_cmma",
                |(input, grad, shape, options)| {
                    dgrad_gemm_simple_sync(input, grad, shape, options, AcceleratedTileKind::Cmma)
                },
            ))
            .with(Tunable::new(
                "simple_sync_mma",
                |(input, grad, shape, options)| {
                    dgrad_gemm_simple_sync(input, grad, shape, options, AcceleratedTileKind::Mma)
                },
            ))
            .with(Tunable::new(
                "simple_async_cmma",
                |(input, grad, shape, options)| {
                    dgrad_gemm_simple_async(input, grad, shape, options, AcceleratedTileKind::Cmma)
                },
            ))
            .with(Tunable::new(
                "simple_async_mma",
                |(input, grad, shape, options)| {
                    dgrad_gemm_simple_async(input, grad, shape, options, AcceleratedTileKind::Mma)
                },
            ))
            .with(Tunable::new(
                "simple_tma_mma",
                |(input, grad, shape, options)| {
                    dgrad_gemm_simple_tma(input, grad, shape, options, AcceleratedTileKind::Mma)
                },
            ))
    });

    TUNER.execute(
        &tune_id,
        &client,
        tunables,
        (out_grad, weights, input_shape, options),
    )
}

pub fn create_wgrad_input<R: Runtime, const N: usize>(
    _key: &ConvTuneKey,
    (out_grad, weights, input_shape, options): &(
        RudaTensor<R>,
        RudaTensor<R>,
        Shape,
        ConvOptions<N>,
    ),
) -> (RudaTensor<R>, RudaTensor<R>, Shape, ConvOptions<N>) {
    (
        out_grad.clone(),
        weights.clone(),
        input_shape.clone(),
        options.clone(),
    )
}

fn create_key<R: Runtime, const N: usize>(
    (out_grad, weights, input_shape, options): &(
        RudaTensor<R>,
        RudaTensor<R>,
        Shape,
        ConvOptions<N>,
    ),
) -> ConvTuneKey {
    let dtype = out_grad.dtype;
    let rank = out_grad.meta.num_dims();
    let dim_c = rank - 1;

    let batch_size = out_grad.meta.shape()[0];
    let in_channels = input_shape[dim_c];
    let out_channels = out_grad.meta.shape()[dim_c];

    let kernel_size = weights.meta.shape()[1..dim_c].to_vec();
    let in_shape = input_shape[1..dim_c]
        .iter()
        .map(|shape| anchor(*shape, None, None, None))
        .collect();

    let ConvOptions {
        stride,
        padding,
        dilation,
        groups,
    } = options.clone();

    let lhs_stride_align = if out_grad.meta.strides()[dim_c] == 1 {
        stride_align(out_grad.meta.strides(), out_grad.dtype.into())
    } else {
        0
    };
    let lhs_shape_align = pow2_factor(out_channels).min(lhs_stride_align);
    let rhs_stride_align = if weights.meta.strides()[dim_c] == 1 {
        stride_align(weights.meta.strides(), weights.dtype.into())
    } else {
        0
    };
    let rhs_shape_align = pow2_factor(in_channels).min(rhs_stride_align);

    ConvTuneKey::Conv(ConvAutotuneKey::new(
        kernel_size,
        stride.to_vec(),
        padding.to_vec(),
        dilation.to_vec(),
        groups,
        in_channels,
        out_channels,
        in_shape,
        batch_size,
        false,
        dtype,
        lhs_shape_align,
        lhs_stride_align,
        rhs_shape_align,
        rhs_stride_align,
    ))
}

/// Maximum factor relevant for strides. Currently set to 2^10 because that's 128-byte swizzle's
/// repeat number, so it's the largest align that can have performance impacts.
const MAX_STRIDE_FACTOR: u32 = 10;

/// Defines the non-contiguous stride alignment in terms of powers of two
fn stride_align(strides: &[usize], elem: StorageType) -> u8 {
    let max = MAX_STRIDE_FACTOR;
    let dim_c = strides.len() - 1;
    let factor = strides[..dim_c]
        .iter()
        .map(|it| (*it * elem.size_bits()) / 8)
        .map(|it| it.trailing_zeros())
        .min()
        .unwrap_or(max);
    factor.min(max) as u8
}

/// Defines the potential vectorization.
fn pow2_factor(axis: usize) -> u8 {
    axis.trailing_zeros().min(4) as u8
}