ruDNN 0.1.101

use ruda_kernel::{dsl::{Runtime, RudaTuneId}, tensor::RudaTensor};
use super::{AttentionStrategy, attention};
use ruda_core::tensor::spatial::AttentionModuleOptions;
use ruda_kernel::dsl::tune::{LocalTuner, Tunable, TunableSet, TuneGroup};
use crate::attention::kernel_ir::{
    launch::AttentionAutotuneKey, routines::blackbox_accelerated::BlackboxAcceleratedStrategy,
};

/// Executes autotune on attention operations
pub fn attention_autotune<R: Runtime>(
    query: RudaTensor<R>,
    key: RudaTensor<R>,
    value: RudaTensor<R>,
    mask: Option<RudaTensor<R>>,
    attn_bias: Option<RudaTensor<R>>,
    options: AttentionModuleOptions,
) -> RudaTensor<R> {
    // A one-row bottom-right causal mask allows all nonempty key positions.
    // Normalize before both eligibility checks and launch/tuning key creation.
    let options = super::support::normalize_single_query_causal(
        options, query.meta.shape[2], key.meta.shape[2],
    );

    // Do not benchmark or reuse a cached FlashAttention candidate when it cannot
    // preserve the requested semantics. The fallback uses the same runtime,
    // device and tensor dtypes as the ordinary fallback strategy.
    if super::support::unsupported_flash_reason(
        &options,
        attn_bias.is_some(),
        query.meta.shape[2],
        key.meta.shape[2],
    ).is_some() {
        return crate::attention::fallback::attention_fallback::<super::ops::RudaAttentionOps<R>>(
            query, key, value, mask, attn_bias, options,
        );
    }

    let client = query.client.clone();

    static TUNER: LocalTuner<AttentionAutotuneKey, RudaTuneId> = LocalTuner::new("ruda_tensor_device::kernel::attention::tune");

    let tune_id = RudaTuneId::new(&client, &query.device);
    let tunables = TUNER.init_for_device(&tune_id, || {
        const PRIORITY_MAX: i8 = 3;
        const PRIORITY_MIN: i8 = 0;

        let flash_attention =
            TuneGroup::<AttentionAutotuneKey>::new("flash_attention", |_key| PRIORITY_MAX);

        let fallback = TuneGroup::<AttentionAutotuneKey>::new("fallback", |key| {
            if key.seq_q > 4096 {
                PRIORITY_MIN
            } else {
                PRIORITY_MAX
            }
        });

        let mut set = TunableSet::new(create_key::<R>, input_gen::<R>)
            .with_stack_tuning(0, "attention-whole-operator-v1", |(q, k, v, mask, bias, options)| {
                format!("q={};k={};v={};mask={:?};bias={:?};options={:?}",
                    q.autotune_signature(), k.autotune_signature(), v.autotune_signature(),
                    mask.as_ref().map(|t| t.autotune_signature()),
                    bias.as_ref().map(|t| t.autotune_signature()), options)
            });

        // First entry should always work, since it is considered the fallback.
        set = set.with(
            Tunable::new(
                "fallback",
                |(query, key, value, mask, attn_bias, options)| {
                    attention::<R>(
                        query,
                        key,
                        value,
                        mask,
                        attn_bias,
                        options,
                        AttentionStrategy::Fallback,
                    )
                    .map_err(|err| std::format!("{err:?}"))
                },
            )
            .group(&fallback, |_key| PRIORITY_MAX),
        );

        let seq_q = 1;
        let seq_kv = 1;
        for num_planes in [2, 4, 8] {
            let name = format!("blackbox_accelerated_{num_planes}_planes_p_{seq_q}-{seq_kv}");
            set = set.with(
                Tunable::new(
                    &name,
                    move |(query, key, value, mask, attn_bias, options)| {
                        attention::<R>(
                            query,
                            key,
                            value,
                            mask,
                            attn_bias,
                            options,
                            AttentionStrategy::FlashBlackboxAccelerated(
                                BlackboxAcceleratedStrategy {
                                    num_planes,
                                    seq_q,
                                    seq_kv,
                                },
                            ),
                        )
                        .map_err(|err| std::format!("{err:?}"))
                    },
                )
                .group(&flash_attention, |_key| PRIORITY_MAX),
            );
        }

        set = set.with(
            Tunable::new("unit", |(query, key, value, mask, attn_bias, options)| {
                attention::<R>(
                    query,
                    key,
                    value,
                    mask,
                    attn_bias,
                    options,
                    AttentionStrategy::FlashUnit,
                )
                .map_err(|err| std::format!("{err:?}"))
            })
            .group(&flash_attention, |_key| PRIORITY_MIN),
        );

        set
    });

    TUNER.execute(
        &tune_id,
        &client,
        tunables,
        (query, key, value, mask, attn_bias, options),
    )
}

#[allow(clippy::type_complexity)]
fn create_key<R: Runtime>(
    (query, key, value, mask, _attn_bias, _options): &(
        RudaTensor<R>,
        RudaTensor<R>,
        RudaTensor<R>,
        Option<RudaTensor<R>>,
        Option<RudaTensor<R>>,
        AttentionModuleOptions,
    ),
) -> AttentionAutotuneKey {
    let total_batches = query.meta.shape[0] * query.meta.shape[1];
    let seq_q = query.meta.shape[2];
    let head_dim = query.meta.shape[3];
    let seq_kv = value.meta.shape[2];
    let val_dim = value.meta.shape[3];

    AttentionAutotuneKey::generate(
        query.dtype.into(),
        key.dtype.into(),
        value.dtype.into(),
        query.dtype.into(),
        total_batches,
        seq_q,
        head_dim,
        seq_kv,
        val_dim,
        mask.is_some(),
    )
}

#[allow(clippy::type_complexity)]
fn input_gen<R: Runtime>(
    _key: &AttentionAutotuneKey,
    (query, key, value, mask, attn_bias, options): &(
        RudaTensor<R>,
        RudaTensor<R>,
        RudaTensor<R>,
        Option<RudaTensor<R>>,
        Option<RudaTensor<R>>,
        AttentionModuleOptions,
    ),
) -> (
    RudaTensor<R>,
    RudaTensor<R>,
    RudaTensor<R>,
    Option<RudaTensor<R>>,
    Option<RudaTensor<R>>,
    AttentionModuleOptions,
) {
    (
        query.clone(),
        key.clone(),
        value.clone(),
        mask.clone(),
        attn_bias.clone(),
        *options,
    )
}