cubecl-runtime 0.11.0

use super::{AutotuneError, AutotuneKey, TuneFn, TuneInputs};

use alloc::boxed::Box;
use alloc::string::ToString;
use alloc::{string::String, sync::Arc, vec, vec::Vec};
use core::fmt::Display;
use core::sync::atomic::{AtomicU32, Ordering};
use cubecl_environment::collections::HashMap;

/// A single candidate for autotune: a named [`TuneFn`] plus the [groups](TuneGroup) it
/// belongs to. A tunable is autotuned whenever any of its groups is prioritized.
///
/// `Id` is what the tunables of a set are [identified](Self::identified) by, `()` for a set
/// whose tunables are [named](Self::new) and nothing more.
pub struct Tunable<K, F: TuneInputs, Output, Id = ()> {
    pub(crate) function: TuneFn<F, Output>,
    groups: Vec<(TuneGroup<K>, PriorityFunc<K>)>,
    identity: Id,
}

impl<K, F: TuneInputs, Output: 'static> Tunable<K, F, Output> {
    /// Create a tunable from a closure.
    ///
    /// The `for<'a> Fn(F::At<'a>) -> _` bound is spelled out in the `where`-clause rather
    /// than hidden behind a helper trait, so that closure inference sees it: otherwise
    /// `move |input| …` picks one concrete lifetime and fails with `implementation of
    /// FnOnce is not general enough` wherever `F::At<'a>` depends on `'a`.
    ///
    /// For multi-input kernels, destructure a tuple:
    /// `Tunable::new("name", |(lhs, rhs, out)| body)`.
    ///
    /// A tunable in no [group](Tunable::group) states no priority, so it trails the
    /// grouped candidates of its round and a short circuit among them leaves it
    /// unmeasured.
    pub fn new<Func, Err>(name: &str, func: Func) -> Self
    where
        Err: Into<String> + 'static,
        Func: for<'a> Fn(<F as TuneInputs>::At<'a>) -> Result<Output, Err> + Send + Sync + 'static,
    {
        Self::build(name, (), func)
    }
}

impl<K, F: TuneInputs, Output: 'static, Id> Tunable<K, F, Output, Id> {
    /// A tunable identified by `identity`, and named by it: the value
    /// [`LocalTuner::fastest_identity`](super::LocalTuner::fastest_identity) hands back once
    /// the tunable is the fastest for a key.
    ///
    /// For a caller that builds something else from the winner — a variant of the same kernel,
    /// say — and so needs to know *which* candidate won rather than only to run it again. The
    /// name a cached result is checksummed by is `identity`'s own, so the two cannot disagree.
    pub fn identified<Func, Err>(identity: Id, func: Func) -> Self
    where
        Id: Display,
        Err: Into<String> + 'static,
        Func: for<'a> Fn(<F as TuneInputs>::At<'a>) -> Result<Output, Err> + Send + Sync + 'static,
    {
        Self::build(&identity.to_string(), identity, func)
    }

    fn build<Func, Err>(name: &str, identity: Id, func: Func) -> Self
    where
        Err: Into<String> + 'static,
        Func: for<'a> Fn(<F as TuneInputs>::At<'a>) -> Result<Output, Err> + Send + Sync + 'static,
    {
        let name: String = name.into();
        let name_for_err = name.clone();
        Self {
            function: TuneFn::new(
                name,
                Box::new(move |inputs| {
                    func(inputs).map_err(|err| AutotuneError::Unknown {
                        name: name_for_err.to_string(),
                        err: err.into(),
                    })
                }),
            ),
            groups: Vec::new(),
            identity,
        }
    }

    /// What this tunable is [identified](Self::identified) by.
    pub(crate) fn identity(&self) -> &Id {
        &self.identity
    }

    /// Add this tunable to a [`TuneGroup`] with the given intra-group priority.
    ///
    /// Groups run in order of their own priority. Within a cutoff group the member
    /// priority picks one level as the round and leaves the levels under it as fallbacks.
    /// Within an [ordered](TuneGroup::ordered) group it orders one round that holds every
    /// member. A negative priority skips the tunable for this key.
    pub fn group(
        mut self,
        group: &TuneGroup<K>,
        priority: impl Fn(&K) -> i8 + Send + Sync + 'static,
    ) -> Self {
        self.groups.push((group.clone(), Arc::new(priority)));
        self
    }
}

/// A priority bucket for tunables, computed from the [autotune key](AutotuneKey).
///
/// Higher-priority groups are autotuned first; once any tunable in a group returns a
/// valid result, no later groups are tried.
pub struct TuneGroup<K> {
    id: GroupId,
    name: Arc<String>,
    pub(crate) priority: PriorityFunc<K>,
    order: GroupOrder,
}

/// Identifies a [`TuneGroup`]: two groups may share a name.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct GroupId(u32);

impl GroupId {
    #[cfg(test)]
    pub(crate) fn new(id: u32) -> Self {
        Self(id)
    }
}

/// How a [`TuneGroup`]'s member priorities shape its batch.
#[derive(Debug, Clone, Copy)]
enum GroupOrder {
    /// The highest level is the round, the levels under it are fallbacks.
    Cutoff,
    /// Every member is in one round, best first.
    Ordered(OrderedSettings),
}

/// The settings of an [ordered](TuneGroup::ordered) group.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct OrderedSettings {
    /// When to stop measuring the group's members. Without one, every member is measured.
    pub patience: Option<Patience>,
}

impl OrderedSettings {
    /// Stop measuring the group once its members stop beating its best one.
    pub fn with_patience(mut self, patience: Patience) -> Self {
        self.patience = Some(patience);
        self
    }
}

/// When to stop measuring an [ordered](TuneGroup::ordered) group: once at least
/// [`min_measured`](Self::min_measured) of its members were measured and the last
/// [`max_misses`](Self::max_misses) of them in a row missed, none faster than the group's
/// leader by more than 2%. The leader is the first member measured, then each member that
/// was faster than it by more than 2%.
///
/// The group is benchmarked best first, so a run of members that cannot beat its leader says
/// the rest of the order is unlikely to either, and they are never compiled. The members
/// already measured still compete as usual, the ones skipped are reported as such, and the
/// rest of the batch is measured as if the group had ended there. A member is skipped only
/// when every group that brought it into the batch ran out of patience. Should every member
/// measured fail later on, the skipped ones are measured after all.
///
/// A member is judged on its best first-pass time, which follows a single warmup: a lucky one
/// makes a leader the rest of the group struggles to beat. A member that fails is neither
/// measured nor a miss: it says nothing about how the order is faring.
///
/// Adaptive scheduler only: the fixed-count pass, the only one on wasm, measures every member.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Patience {
    /// Members measured before the group may stop. One member is always measured.
    pub min_measured: usize,
    /// Members in a row that miss before the group stops. Zero stops it at
    /// [`min_measured`](Self::min_measured), whatever they measured.
    pub max_misses: usize,
}

impl<K> core::fmt::Debug for TuneGroup<K> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("TuneGroup")
            .field("id", &self.id)
            .field("name", &self.name)
            .finish()
    }
}

impl<K> Clone for TuneGroup<K> {
    fn clone(&self) -> Self {
        Self {
            id: self.id,
            name: self.name.clone(),
            priority: self.priority.clone(),
            order: self.order,
        }
    }
}

impl<K> TuneGroup<K> {
    /// Create a new group based on a priority function.
    ///
    /// A member's own priority ([`Tunable::group`]) is a cutoff: the highest level is the
    /// round, and the levels under it are fallbacks reached as each round fails.
    pub fn new(name: &str, f: impl Fn(&K) -> i8 + Send + Sync + 'static) -> Self {
        Self::build(name, f, GroupOrder::Cutoff)
    }

    /// Create a group whose member priorities order one round rather than cut it off:
    /// every member with a non-negative priority is in the batch, best first.
    ///
    /// The batch is benchmarked in that order and the short circuit ends it at the first
    /// candidate under the [bounds](super::Bounds)' time limit, so the priority decides
    /// how much of the group is compiled rather than which of it can win. The batch costs
    /// its whole membership without a bounds generator, on wasm, or with the short circuit
    /// disabled.
    ///
    /// The batch sits at the priority of its best member, so a cutoff group at the same
    /// group priority interleaves with it by priority.
    ///
    /// [`OrderedSettings`] holds the group's settings, such as the [`Patience`] that stops
    /// measuring it once its members stop beating its best one.
    pub fn ordered(
        name: &str,
        f: impl Fn(&K) -> i8 + Send + Sync + 'static,
        settings: OrderedSettings,
    ) -> Self {
        Self::build(name, f, GroupOrder::Ordered(settings))
    }

    fn build(name: &str, f: impl Fn(&K) -> i8 + Send + Sync + 'static, order: GroupOrder) -> Self {
        let id = GroupId(GROUP_COUNTER.fetch_add(1, Ordering::Relaxed));

        Self {
            id,
            name: Arc::new(name.into()),
            priority: Arc::new(f),
            order,
        }
    }

    fn is_ordered(&self) -> bool {
        matches!(self.order, GroupOrder::Ordered(_))
    }

    fn patience(&self) -> Option<Patience> {
        match self.order {
            GroupOrder::Ordered(settings) => settings.patience,
            GroupOrder::Cutoff => None,
        }
    }
}

/// One batch of the [`TunePlan`]: the [tunables](Tunable) to benchmark, in order, and the
/// [`Patience`] of each group that brought one in.
#[derive(Debug, Default, PartialEq, Eq)]
pub(crate) struct Batch {
    pub(crate) entries: Vec<BatchEntry>,
    pub(crate) patience: HashMap<GroupId, Patience>,
}

/// One [tunable](Tunable) of a [`Batch`], with every group that brought it into the batch.
/// An ungrouped tunable has none.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct BatchEntry {
    pub(crate) index: usize,
    pub(crate) groups: Vec<GroupId>,
}

impl Batch {
    /// A batch that measures every one of `indices`, no group's patience applying.
    pub(crate) fn unconditional(indices: impl IntoIterator<Item = usize>) -> Self {
        let mut batch = Self::default();
        for index in indices {
            batch.push_unconditional(index);
        }
        batch
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// The tunable indices, in the order they are benchmarked.
    pub(crate) fn indices(&self) -> Vec<usize> {
        self.entries.iter().map(|entry| entry.index).collect()
    }

    /// Add a tunable through one of its groups. A tunable several groups bring in is
    /// benchmarked once, at its first place, and answers to each of them.
    fn push(&mut self, planned: &Planned) {
        match self.entries.iter_mut().find(|e| e.index == planned.index) {
            Some(entry) if entry.groups.contains(&planned.group) => {}
            Some(entry) => entry.groups.push(planned.group),
            None => self.entries.push(BatchEntry {
                index: planned.index,
                groups: vec![planned.group],
            }),
        }
    }

    fn push_unconditional(&mut self, index: usize) {
        self.entries.push(BatchEntry {
            index,
            groups: Vec::new(),
        });
    }
}

#[derive(Debug, Clone)]
/// A group plan dictates which [tunables](Tunable) should be executed, and in what order.
pub(crate) struct TunePlan {
    priorities: Vec<i8>,
    no_groups: Vec<usize>,
    groups: HashMap<i8, GroupPlan>,
    returned: Vec<usize>,
    /// The [`Patience`] of every group that has one.
    patience: HashMap<GroupId, Patience>,
    /// Each tunable's place in each of its groups, by tunable index, as the
    /// key priced it.
    #[cfg(persistence)]
    placements: Vec<Vec<Placement>>,
}

/// A tunable's place in one of its [groups](TuneGroup), as one key priced it.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Placement {
    /// The group's name.
    pub group: String,
    /// Whether the group is [ordered](TuneGroup::ordered): its members'
    /// priorities order one round rather than cutting it into rounds.
    pub ordered: bool,
    /// The group's own priority: higher groups run first.
    pub group_priority: i8,
    /// The tunable's priority within the group; negative skips it.
    pub priority: i8,
}

#[derive(Default, Debug, Clone)]
struct GroupPlan {
    priorities: Vec<i8>,
    indices: HashMap<i8, Vec<Planned>>,
}

/// One tunable's place in a [`GroupPlan`]: which tunable, through which group, and the
/// priority that orders it within its batch.
///
/// The group is its id, not its name: a name is the caller's label and two groups may
/// share one, which the cross-level dedup in [`TunePlan::group_plan_next`] would then
/// read as a single group and strike a live candidate out of the plan.
#[derive(Debug, Clone)]
struct Planned {
    index: usize,
    group: GroupId,
    priority: i8,
}

#[derive(Debug)]
struct Cleanup {
    groups: Vec<i8>,
    tunables: Vec<(i8, i8)>,
    /// Within group priority is too low to even try.
    skipped: bool,
}

impl TunePlan {
    pub fn new<K: AutotuneKey, F: TuneInputs, Out, Id>(
        key: &K,
        tunables: &[Tunable<K, F, Out, Id>],
    ) -> Self {
        let mut priorities = Vec::<i8>::new();
        let mut no_groups = Vec::new();
        let mut groups = HashMap::<i8, GroupPlan>::new();

        // Each of the caller's priority functions is asked once and its answer carried:
        // an ordered group's level is its best member's priority, which is known only
        // once every member is priced.
        let mut priced = Vec::new();
        let mut ordered_levels = HashMap::<GroupId, i8>::new();
        let mut patience = HashMap::<GroupId, Patience>::new();
        #[cfg(persistence)]
        let mut placements = vec![Vec::new(); tunables.len()];

        for (index, tunable) in tunables.iter().enumerate() {
            if tunable.groups.is_empty() {
                no_groups.push(index);
                continue;
            }

            for (group, within_group_priority_fn) in tunable.groups.iter() {
                if let Some(group_patience) = group.patience() {
                    patience.insert(group.id, group_patience);
                }
                let group_priority = (group.priority)(key);
                let priority = within_group_priority_fn(key);
                #[cfg(persistence)]
                placements[index].push(Placement {
                    group: group.name.to_string(),
                    ordered: group.is_ordered(),
                    group_priority,
                    priority,
                });

                if group.is_ordered() && priority >= 0 {
                    let level = ordered_levels.entry(group.id).or_insert(priority);
                    *level = (*level).max(priority);
                }

                priced.push((index, group, group_priority, priority));
            }
        }

        for (index, group, group_priority, priority) in priced {
            if !priorities.contains(&group_priority) {
                priorities.push(group_priority);
            }

            let group_plan = match groups.get_mut(&group_priority) {
                Some(val) => val,
                None => {
                    groups.insert(group_priority, GroupPlan::default());
                    groups.get_mut(&group_priority).unwrap()
                }
            };

            // An ordered group is one batch at its best member's level, so a cutoff
            // group at the same group priority interleaves with it by priority.
            let level = match group.is_ordered() && priority >= 0 {
                true => ordered_levels[&group.id],
                false => priority,
            };
            let planned = Planned {
                index,
                group: group.id,
                priority,
            };

            if group_plan.priorities.contains(&level) {
                group_plan.indices.get_mut(&level).unwrap().push(planned);
            } else {
                group_plan.priorities.push(level);
                group_plan.indices.insert(level, vec![planned]);
            }
        }

        priorities.sort();

        for group in groups.iter_mut() {
            group.1.priorities.sort();
        }

        Self {
            priorities,
            no_groups,
            groups,
            returned: Vec::new(),
            patience,
            #[cfg(persistence)]
            placements,
        }
    }

    /// The whole plan as a recorded tune reports it: every tunable named by
    /// `names`, in the order the plan releases them, with the batch each
    /// comes in and its place in each of its groups. A tunable no batch
    /// holds, its every priority negative, comes last.
    #[cfg(persistence)]
    pub(crate) fn account(&self, names: &[&str]) -> Vec<super::record::PlannedCandidate> {
        let mut plan = self.clone();
        let batches = core::iter::from_fn(|| {
            Some(plan.next())
                .filter(|batch| !batch.is_empty())
                .map(|batch| batch.indices())
        });
        let mut account: Vec<_> = batches
            .enumerate()
            .flat_map(|(batch, indices)| indices.into_iter().map(move |index| (index, Some(batch))))
            .collect();
        let excluded =
            (0..names.len()).filter(|index| !account.iter().any(|(seen, _)| seen == index));
        account.extend(excluded.map(|index| (index, None)).collect::<Vec<_>>());
        account
            .into_iter()
            .map(|(index, batch)| super::record::PlannedCandidate {
                name: names[index].to_string(),
                index,
                batch,
                groups: self.placements[index].clone(),
            })
            .collect()
    }

    /// The next batch to benchmark. An empty one means the plan is exhausted.
    pub(crate) fn next(&mut self) -> Batch {
        // A tunable in no group states no priority, so it trails the batch: the grouped
        // candidates decide what is compiled and benchmarked first.
        let ungrouped = core::mem::take(&mut self.no_groups);
        let priority = self.priorities.last();

        let priority = match priority {
            Some(val) => *val,
            None => return Batch::unconditional(ungrouped),
        };
        let mut batch = Batch {
            entries: Vec::new(),
            patience: self.patience.clone(),
        };

        let (group_indices, cleanup) = self.group_plan_next(priority);
        // Some entries are skipped for this round of prioritizing.
        let skipped = cleanup.skipped || priority < 0;
        let mut all_skip = true;

        self.cleanup(cleanup);

        if priority >= 0 {
            for planned in group_indices {
                if self.returned.contains(&planned.index) {
                    continue;
                }
                all_skip = false;
                batch.push(&planned);
            }
        }

        for index in ungrouped {
            batch.push_unconditional(index);
        }

        // The batch is empty, but it doesn't mean we should stop
        // autotuning, since some entries were skipped.

        if batch.is_empty() && (skipped || all_skip) {
            self.next()
        } else {
            for entry in batch.entries.iter() {
                self.returned.push(entry.index);
            }
            batch
        }
    }

    fn cleanup(&mut self, cleanup: Cleanup) {
        for group_p in cleanup.groups {
            let index = self
                .priorities
                .iter()
                .enumerate()
                .find(|p| *p.1 == group_p)
                .unwrap();

            self.priorities.remove(index.0);
            self.groups.remove(&group_p);
        }

        for (group_p, tunable_p) in cleanup.tunables {
            if let Some(group) = self.groups.get_mut(&group_p) {
                let index = group
                    .priorities
                    .iter()
                    .enumerate()
                    .find(|p| *p.1 == tunable_p)
                    .unwrap();
                group.priorities.remove(index.0);
                group.indices.remove(&tunable_p);
            }
        }
    }

    fn group_plan_next(&mut self, priority: i8) -> (Vec<Planned>, Cleanup) {
        let group_plan = self.groups.get_mut(&priority).expect("To be filled");
        let within_group_prio = group_plan.priorities.pop().unwrap();
        let mut next_indices = group_plan.indices.remove(&within_group_prio).unwrap();
        // Highest priority first, registration order among equals.
        next_indices.sort_by_key(|planned| (core::cmp::Reverse(planned.priority), planned.index));

        let mut cleanup_groups = Vec::new();
        let mut cleanup_tunables = Vec::new();

        for (pg, group) in self.groups.iter_mut() {
            let mut num_empty_tunables = 0;
            let num_tunables = group.priorities.len();

            for (pt, indices) in group.indices.iter_mut() {
                for n in &next_indices {
                    let entry = indices
                        .iter()
                        .position(|p| p.index == n.index && p.group == n.group);
                    if let Some(entry) = entry {
                        indices.remove(entry);
                    }
                }

                if indices.is_empty() {
                    num_empty_tunables += 1;
                    cleanup_tunables.push((*pg, *pt));
                }
            }

            if num_empty_tunables == num_tunables {
                cleanup_groups.push(*pg);
            }
        }

        if within_group_prio < 0 {
            // Discard algorithms with negative priority
            next_indices.clear();
        }

        (
            next_indices,
            Cleanup {
                groups: cleanup_groups,
                tunables: cleanup_tunables,
                skipped: within_group_prio < 0,
            },
        )
    }
}

type PriorityFunc<K> = Arc<dyn Fn(&K) -> i8 + Send + Sync>;

static GROUP_COUNTER: AtomicU32 = AtomicU32::new(0);

#[cfg(test)]
mod tests {
    use core::fmt::Display;

    use serde::{Deserialize, Serialize};

    use super::*;

    #[derive(Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize, Debug)]
    struct FakeAutotuneKey;

    impl Display for FakeAutotuneKey {
        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
            f.write_str("FakeAutotuneKey")
        }
    }

    impl AutotuneKey for FakeAutotuneKey {}

    #[test_log::test]
    fn test_plan_order() {
        let group0 = TuneGroup::<FakeAutotuneKey>::new("group0", |_| 2);
        let group1 = TuneGroup::<FakeAutotuneKey>::new("group1", |_| 1);

        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 1);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 2);
        let tunable3 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group1, |_| 2);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2, tunable3]);

        // tunable0 is in no group, so it trails the batch rather than leading it.
        assert_eq!(plan.next().indices(), vec![2, 0]);
        assert_eq!(plan.next().indices(), vec![1]);
        assert_eq!(plan.next().indices(), vec![3]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_order_multi_groups_same_priority() {
        let group0 = TuneGroup::<FakeAutotuneKey>::new("group0", |_| 2);
        let group1 = TuneGroup::<FakeAutotuneKey>::new("group1", |_| 1);
        let group2 = TuneGroup::<FakeAutotuneKey>::new("group2", |_| 1);

        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 1);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 2);
        let tunable3 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group1, |_| 2);
        let tunable4 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group2, |_| 2);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2, tunable3, tunable4]);

        assert_eq!(plan.next().indices(), vec![2, 0]);
        assert_eq!(plan.next().indices(), vec![1]);
        assert_eq!(plan.next().indices(), vec![3, 4]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_order_tunable_multiple_groups() {
        let group0 = TuneGroup::<FakeAutotuneKey>::new("group0", |_| 1);
        let group1 = TuneGroup::<FakeAutotuneKey>::new("group1", |_| 2);

        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel);
        let tunable1 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel)
            .group(&group0, |_| 1)
            .group(&group1, |_| 2);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 2);
        let tunable3 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group1, |_| 3);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2, tunable3]);

        assert_eq!(plan.next().indices(), vec![3, 0]);
        assert_eq!(plan.next().indices(), vec![1]);
        assert_eq!(plan.next().indices(), vec![2]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_negative_priority() {
        let group0 = TuneGroup::<FakeAutotuneKey>::new("group0", |_| 2);
        let group1 = TuneGroup::<FakeAutotuneKey>::new("group1", |_| 1);

        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| -1);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 2);
        let tunable3 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group1, |_| 2);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2, tunable3]);

        assert_eq!(plan.next().indices(), vec![2, 0]);
        assert_eq!(plan.next().indices(), vec![3]);
        assert!(plan.next().is_empty());
    }

    #[cfg(persistence)]
    #[test_log::test]
    fn test_plan_account_lists_every_tunable_in_release_order() {
        let group0 = TuneGroup::<FakeAutotuneKey>::new("group0", |_| 2);
        let group1 = TuneGroup::<FakeAutotuneKey>::new("group1", |_| 1);

        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| -1);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 2);
        let tunable3 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group1, |_| 2);

        let key = FakeAutotuneKey;
        let plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2, tunable3]);
        let account = plan.account(&["t0", "t1", "t2", "t3"]);

        let order: Vec<_> = account
            .iter()
            .map(|planned| (planned.name.as_str(), planned.batch))
            .collect();
        assert_eq!(
            order,
            vec![
                ("t2", Some(0)),
                ("t0", Some(0)),
                ("t3", Some(1)),
                ("t1", None)
            ]
        );
        assert_eq!(
            account[0].groups,
            vec![Placement {
                group: "group0".into(),
                ordered: false,
                group_priority: 2,
                priority: 2,
            }]
        );
        assert!(account[1].groups.is_empty());
        // Accounting for the plan leaves it to run from its first batch.
        let mut plan = plan;
        assert_eq!(plan.next().indices(), vec![2, 0]);
    }

    #[test_log::test]
    fn test_plan_no_group() {
        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel);
        let tunable1 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1]);

        assert_eq!(plan.next().indices(), vec![0, 1]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_falls_through_when_all_group_tunables_fail() {
        // Every tunable lives in exactly one group; the caller treats every batch as a failure
        // by continuing to call next(). The plan must still surface every tunable, in priority
        // order, before going empty.
        let group0 = TuneGroup::<FakeAutotuneKey>::new("group0", |_| 2);
        let group1 = TuneGroup::<FakeAutotuneKey>::new("group1", |_| 1);

        let tunable0 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 1);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 2);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group1, |_| 1);
        let tunable3 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group1, |_| 2);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2, tunable3]);

        let mut all_returned: Vec<usize> = Vec::new();
        loop {
            let batch = plan.next();
            if batch.is_empty() {
                break;
            }
            all_returned.extend(batch.indices());
        }

        // Highest group (prio 2) drains first from highest intra-priority down, then next group.
        assert_eq!(all_returned, vec![1, 0, 3, 2]);
    }

    #[test_log::test]
    fn test_plan_single_group_exhausts_all_intra_priorities() {
        // A single group with multiple intra-priorities should yield each batch separately,
        // allowing the caller to continue on failures until the group is exhausted.
        let group0 = TuneGroup::<FakeAutotuneKey>::new("group0", |_| 0);

        let tunable0 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 1);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 2);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 3);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2]);

        assert_eq!(plan.next().indices(), vec![2]);
        assert_eq!(plan.next().indices(), vec![1]);
        assert_eq!(plan.next().indices(), vec![0]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_all_negative_group_advances_to_next_group() {
        // A group whose every tunable has a negative intra-priority should be skipped entirely
        // without stopping autotuning — the next group must still be reached.
        let group0 = TuneGroup::<FakeAutotuneKey>::new("group0", |_| 2);
        let group1 = TuneGroup::<FakeAutotuneKey>::new("group1", |_| 1);

        let tunable0 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| -1);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| -2);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group1, |_| 1);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2]);

        assert_eq!(plan.next().indices(), vec![2]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_no_group_tunables_only_emitted_once_even_on_failures() {
        // The ungrouped tunables are emitted together with the first group batch. If the caller
        // keeps calling next() (treating the first batch as failing), they must not be
        // re-emitted, and the plan must still advance to later groups.
        let group0 = TuneGroup::<FakeAutotuneKey>::new("group0", |_| 2);
        let group1 = TuneGroup::<FakeAutotuneKey>::new("group1", |_| 1);

        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 1);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group1, |_| 1);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2]);

        assert_eq!(plan.next().indices(), vec![1, 0]);
        assert_eq!(plan.next().indices(), vec![2]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_multi_group_tunable_not_duplicated_across_failed_groups() {
        // tunable1 belongs to both group0 and group1. It must be returned exactly once (via its
        // higher-priority group), even if the caller continues iterating after failures.
        let group0 = TuneGroup::<FakeAutotuneKey>::new("group0", |_| 1);
        let group1 = TuneGroup::<FakeAutotuneKey>::new("group1", |_| 2);

        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel)
            .group(&group0, |_| 1)
            .group(&group1, |_| 1);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group0, |_| 2);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1]);

        let mut all_returned: Vec<usize> = Vec::new();
        loop {
            let batch = plan.next();
            if batch.is_empty() {
                break;
            }
            all_returned.extend(batch.indices());
        }

        // tunable0 comes from group1 (higher priority). tunable1 is the sole member of group0
        // after cross-group dedup. No duplicates.
        assert_eq!(all_returned, vec![0, 1]);
    }

    #[test_log::test]
    fn test_plan_recurses_when_batch_is_fully_already_returned() {
        // Regression test: a tunable that lives in multiple groups was already emitted via its
        // higher-priority group, so when its lower-priority group's batch fires the only index
        // is one already present in `returned`. The plan must NOT return an empty batch here
        // (that signals "no more work" to the caller and aborts with NoValidKernelFound); it
        // must recurse to the next intra-priority and surface the remaining tunable.
        //
        // Cross-group dedup in group_plan_next compares (index, Arc<String> group_name), so a
        // tunable appearing in both group_hi and group_lo isn't auto-removed from group_lo
        // when popped from group_hi — the `returned` + `all_skip` path is the only guard.
        let group_hi = TuneGroup::<FakeAutotuneKey>::new("hi", |_| 2);
        let group_lo = TuneGroup::<FakeAutotuneKey>::new("lo", |_| 1);

        // tunable0 is in both groups. tunable1 is only in group_lo at a lower intra-priority.
        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel)
            .group(&group_hi, |_| 1)
            .group(&group_lo, |_| 2);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group_lo, |_| 1);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1]);

        // First call: group_hi yields tunable0.
        assert_eq!(plan.next().indices(), vec![0]);
        // Second call: group_lo's higher intra-priority batch is just tunable0 (already
        // returned). Without the fix this returns [] and the autotuner aborts. With the fix
        // the plan recurses and yields tunable1.
        assert_eq!(plan.next().indices(), vec![1]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_ordered_group_is_one_batch_best_first() {
        // Every member is in the round, best first, and a negative priority still skips.
        let group =
            TuneGroup::<FakeAutotuneKey>::ordered("ordered", |_| 1, OrderedSettings::default());

        let tunable0 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group, |_| 1);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group, |_| 3);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group, |_| -1);
        let tunable3 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group, |_| 3);
        let tunable4 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group, |_| 2);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2, tunable3, tunable4]);

        assert_eq!(plan.next().indices(), vec![1, 3, 4, 0]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_batch_carries_every_group_of_an_entry() {
        // The batch carries patience only for the group that has one, and a tunable two
        // groups bring in is one entry naming both.
        let patience = Patience {
            min_measured: 10,
            max_misses: 5,
        };
        let ordered = TuneGroup::<FakeAutotuneKey>::ordered(
            "ordered",
            |_| 1,
            OrderedSettings::default().with_patience(patience),
        );
        let cutoff = TuneGroup::<FakeAutotuneKey>::new("cutoff", |_| 1);

        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&ordered, |_| 2);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&cutoff, |_| 2);
        let tunable3 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&ordered, |_| 1);
        let tunable4 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel)
            .group(&cutoff, |_| 2)
            .group(&ordered, |_| 1);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2, tunable3, tunable4]);
        let batch = plan.next();
        let entry = |index, groups: &[GroupId]| BatchEntry {
            index,
            groups: groups.to_vec(),
        };

        assert_eq!(
            batch.entries,
            vec![
                entry(1, &[ordered.id]),
                entry(2, &[cutoff.id]),
                entry(4, &[cutoff.id, ordered.id]),
                entry(3, &[ordered.id]),
                entry(0, &[]),
            ]
        );
        assert_eq!(batch.patience, HashMap::from_iter([(ordered.id, patience)]));
    }

    #[test_log::test]
    fn test_plan_entry_names_a_group_listed_twice_once() {
        // A group listed twice on one tunable is still one group: naming it twice would
        // count each of the tunable's times twice against its patience.
        let ordered =
            TuneGroup::<FakeAutotuneKey>::ordered("ordered", |_| 1, OrderedSettings::default());

        let tunable = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel)
            .group(&ordered, |_| 2)
            .group(&ordered, |_| 1);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable]);

        assert_eq!(
            plan.next().entries,
            vec![BatchEntry {
                index: 0,
                groups: vec![ordered.id],
            }]
        );
    }

    #[test_log::test]
    fn test_plan_ordered_group_keeps_the_group_cutoff() {
        // The order is within the group. A lower-priority group is still a fallback,
        // reached only once the ordered batch fails.
        let first =
            TuneGroup::<FakeAutotuneKey>::ordered("first", |_| 2, OrderedSettings::default());
        let fallback = TuneGroup::<FakeAutotuneKey>::new("fallback", |_| 1);

        let tunable0 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&first, |_| 1);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&first, |_| 2);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&fallback, |_| 1);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2]);

        assert_eq!(plan.next().indices(), vec![1, 0]);
        assert_eq!(plan.next().indices(), vec![2]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_ordered_batch_leads_the_ungrouped_tunables() {
        // A tunable in no group must not preempt the ordered group's best candidate. The
        // batch stops at the first candidate under the bound, so whatever leads it is
        // what gets compiled.
        let group =
            TuneGroup::<FakeAutotuneKey>::ordered("ordered", |_| 1, OrderedSettings::default());

        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group, |_| 1);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&group, |_| 5);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2]);

        assert_eq!(plan.next().indices(), vec![2, 1, 0]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_ordered_batch_is_not_jumped_by_a_cutoff_group_beside_it() {
        // Both groups sit at group priority 1. The ordered batch is planned at its best
        // member's priority, so the cutoff member under it is a fallback behind the batch
        // and cannot take a round of its own in front of it, where a short circuit would
        // skip the ordered group whole.
        let ordered =
            TuneGroup::<FakeAutotuneKey>::ordered("ordered", |_| 1, OrderedSettings::default());
        let cutoff = TuneGroup::<FakeAutotuneKey>::new("cutoff", |_| 1);

        let tunable0 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&ordered, |_| 3);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&ordered, |_| 1);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&cutoff, |_| 2);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2]);

        assert_eq!(plan.next().indices(), vec![0, 1]);
        assert_eq!(plan.next().indices(), vec![2]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_cutoff_member_above_the_ordered_batch_still_leads() {
        // The interleaving cuts both ways. A cutoff member priced above the ordered
        // group's best keeps its round in front, and one priced level with it joins the
        // batch, ordered among the members by priority.
        let ordered =
            TuneGroup::<FakeAutotuneKey>::ordered("ordered", |_| 1, OrderedSettings::default());
        let cutoff = TuneGroup::<FakeAutotuneKey>::new("cutoff", |_| 1);

        let tunable0 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&ordered, |_| 2);
        let tunable1 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&ordered, |_| 1);
        let tunable2 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&cutoff, |_| 3);
        let tunable3 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&cutoff, |_| 2);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2, tunable3]);

        assert_eq!(plan.next().indices(), vec![2]);
        assert_eq!(plan.next().indices(), vec![0, 3, 1]);
        assert!(plan.next().is_empty());
    }

    #[test_log::test]
    fn test_plan_same_named_groups_do_not_strike_each_other_out() {
        // Two groups may share a name, and the cross-level dedup must still tell them
        // apart. Keyed on the name, popping `hi`'s discarded negative level takes
        // tunable1 out of `lo`'s plan and the only viable candidate is never benchmarked.
        let hi = TuneGroup::<FakeAutotuneKey>::new("shared", |_| 2);
        let lo = TuneGroup::<FakeAutotuneKey>::new("shared", |_| 1);

        let tunable0 =
            Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel).group(&hi, |_| 1);
        let tunable1 = Tunable::<FakeAutotuneKey, (), ()>::new("fake", fake_kernel)
            .group(&hi, |_| -1)
            .group(&lo, |_| 1);

        let key = FakeAutotuneKey;
        let mut plan = TunePlan::new(&key, &[tunable0, tunable1]);

        assert_eq!(plan.next().indices(), vec![0]);
        assert_eq!(plan.next().indices(), vec![1]);
        assert!(plan.next().is_empty());
    }

    fn fake_kernel(_: ()) -> Result<(), String> {
        Ok(())
    }
}