cubecl-runtime 0.12.0-pre.1

use super::{AutotuneKey, AutotuneOutput, TunableSet, TuneInputs, Tuner};
#[cfg(feature = "autotune-checks")]
use crate::tune::AutotuneLoggerExt;
use crate::{client::Client, device::HardwareIdentity, tune::TuneCacheResult};
#[cfg(feature = "autotune-checks")]
use alloc::string::ToString;
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::{
    any::{Any, TypeId},
    hash::Hash,
};
use cubecl_environment::collections::HashMap;
use cubecl_environment::sync::{Mutex, RwLock};

/// The tunable sets a [`LocalTuner`] has built, keyed by device as well as by
/// initializer: a set is built from the device it will run on — its client, its
/// hardware properties — so one device's set cannot answer for another's. See
/// [`LocalTuner::init`].
type Sets<ID> = RwLock<Option<HashMap<(TypeId, ID), Arc<dyn Any + Send + Sync>>>>;

/// A local tuner allows to create a tuner for a specific key that can be different from the server
/// key.
pub struct LocalTuner<AK: AutotuneKey, ID> {
    state: Mutex<Option<Tuners<AK, ID>>>,
    name: &'static str,
    sets: Sets<ID>,
}

/// The tuners a [`LocalTuner`] opened. An id keeps a tuner per candidate list it executed, shared
/// with the ids of its hardware, so ids of the same hardware running one list tune a key once for
/// all of them.
struct Tuners<AK: AutotuneKey, ID> {
    /// An id's tuners, by the checksum of the list each runs.
    by_id: HashMap<ID, HashMap<u64, Arc<Tuner<AK>>>>,
    opened: HashMap<TunerIdentity, Arc<Tuner<AK>>>,
}

/// Which tuner a device's executes reach: one per hardware and candidate list.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct TunerIdentity {
    hardware: HardwareIdentity,
    /// The low half of the candidate list's [`compute_checksum`](TunableSet::compute_checksum): each
    /// device builds its own list, and one of the same hardware may build another.
    checksum: u64,
}

impl<AK: AutotuneKey, ID: Hash + Eq + Clone> Tuners<AK, ID> {
    fn new() -> Self {
        Self {
            by_id: HashMap::new(),
            opened: HashMap::new(),
        }
    }

    /// The tuner `id` executes `operations` on, opened by the first id of its hardware to execute
    /// the same candidate list.
    fn get_or_open<I: TuneInputs, Out, Id>(
        &mut self,
        id: &ID,
        name: &str,
        client: &Client,
        operations: &TunableSet<AK, I, Out, Id>,
    ) -> Arc<Tuner<AK>> {
        let checksum = operations.checksum();
        if let Some(tuner) = self.by_id.get(id).and_then(|lists| lists.get(&checksum)) {
            return tuner.clone();
        }
        let identity = TunerIdentity {
            hardware: client.hardware_identity(),
            checksum,
        };
        let tuner = self
            .opened
            .entry(identity)
            .or_insert_with_key(|identity| {
                Arc::new(Tuner::new(&name.replace("::", "-"), &identity.hardware))
            })
            .clone();
        self.by_id
            .entry(id.clone())
            .or_default()
            .insert(checksum, tuner.clone());
        tuner
    }
}

/// Create a local tuner with the provided name.
#[macro_export]
macro_rules! local_tuner {
    ($name:expr) => {
        LocalTuner::new(concat!(module_path!(), "-", $name))
    };
    () => {
        LocalTuner::new(module_path!())
    };
}

pub use local_tuner;

impl<AK, ID> LocalTuner<AK, ID>
where
    AK: AutotuneKey + 'static,
    ID: Hash + PartialEq + Eq + Clone,
{
    /// Create a new local tuner.
    pub const fn new(name: &'static str) -> Self {
        Self {
            state: Mutex::new(None),
            name,
            sets: RwLock::new(None),
        }
    }

    /// Get or initialize the [`TunableSet`] for `id`.
    ///
    /// Returns a cached `Arc<TunableSet>` keyed by the `TypeId` of `init_set`
    /// *and* by `id`, so the initializer runs once per device rather than once
    /// per process.
    ///
    /// The device is part of the key because a set is routinely built from the
    /// device it will run on: a closure captures that device's
    /// [`Client`] to ask what it supports, or reads its hardware
    /// properties to decide which tunables are worth offering at all. Keyed by
    /// the initializer alone, whichever device tuned first would answer those
    /// questions for every device that followed — promoting kernels onto
    /// hardware that cannot run them, or withholding kernels from hardware
    /// that can.
    pub fn init<I, Out, Id, F>(&self, id: &ID, init_set: F) -> Arc<TunableSet<AK, I, Out, Id>>
    where
        F: Fn() -> TunableSet<AK, I, Out, Id> + 'static + Send + Sync,
        Id: Send + Sync + 'static,
        I: TuneInputs,
        Out: AutotuneOutput,
    {
        let key = (TypeId::of::<F>(), id.clone());
        let sets = self.sets.read();

        static DOWNCAST_ERROR: &str = "Local tuner only support one set of tunable that must work on the same input and output declared with the init function.";

        if let Some(sets) = sets.as_ref()
            && let Some(set) = sets.get(&key)
        {
            return set.clone().downcast().expect(DOWNCAST_ERROR);
        };

        core::mem::drop(sets);

        let mut sets = self.sets.write();

        if let Some(sets) = sets.as_ref()
            && let Some(set) = sets.get(&key)
        {
            return set.clone().downcast().expect(DOWNCAST_ERROR);
        };

        let content = Arc::new(init_set());

        if let Some(sets) = sets.as_mut() {
            sets.insert(key, content.clone());
        } else {
            let mut map = HashMap::<(TypeId, ID), Arc<dyn Any + Send + Sync>>::new();
            map.insert(key, content.clone());
            *sets = Some(map);
        };

        content
    }

    /// What the fastest tunable in `operations` for `key` on `id` is
    /// [identified](super::Tunable::identified) by: `None` until a result for the key is
    /// settled, tuned by a round in this process or read back from disk by an
    /// [`execute`](Self::execute).
    ///
    /// Results are shared by every `id` of the same [hardware](crate::device::HardwareIdentity)
    /// executing the same candidates, from the first [`execute`](Self::execute) of each: an `id`
    /// that has not executed `operations` sees none.
    ///
    /// It reads results and never produces one: it never starts a round, never waits on one in
    /// flight, never reads the persistent cache, and never resets the tuner's cache after an
    /// environment switch. It reports nothing settled there until the next `execute`.
    pub fn fastest_identity<I, Out, Id>(
        &self,
        id: &ID,
        operations: &TunableSet<AK, I, Out, Id>,
        key: &AK,
    ) -> Option<Id>
    where
        I: TuneInputs,
        Id: Clone,
    {
        let tuner = self
            .state
            .lock()
            .as_ref()?
            .by_id
            .get(id)?
            .get(&operations.checksum())?
            .clone();
        let fastest_index = tuner.settled(key)?;
        Some(operations.identity(fastest_index).clone())
    }

    /// Clear the autotune state.
    pub fn clear(&self) {
        self.state.lock().take();
    }

    #[cfg(feature = "autotune-checks")]
    fn checks<'a, I: TuneInputs, Out: AutotuneOutput, Id>(
        &self,
        operations: &TunableSet<AK, I, Out, Id>,
        inputs: &<I as TuneInputs>::At<'a>,
    ) -> alloc::vec::Vec<crate::tune::log::CheckResult>
    where
        <I as TuneInputs>::At<'a>: Clone + Send,
    {
        let mut checks_outputs = Vec::new();
        for i in 0..operations.len() {
            let op = operations.fastest(i);
            let result = op.execute(inputs.clone());
            checks_outputs.push((op.name.to_string(), result));
        }
        super::check_autotune_outputs(checks_outputs)
    }

    /// Execute the fastest operation in a [`TunableSet`], triggering a tuning pass on
    /// the first call for a given key.
    pub fn execute<'a, I: TuneInputs, Out, Id>(
        &self,
        id: &ID,
        client: &Client,
        operations: Arc<TunableSet<AK, I, Out, Id>>,
        inputs: <I as TuneInputs>::At<'a>,
    ) -> Out
    where
        <I as TuneInputs>::At<'a>: Clone + Send,
        Out: AutotuneOutput,
        Id: Send + Sync,
    {
        let key = operations.generate_key(&inputs);

        let tuner = self
            .state
            .lock()
            .get_or_insert_with(Tuners::new)
            .get_or_open(id, self.name, client, &operations);

        // Fast path: a cached hit skips straight to the fastest operation.
        // `fastest` also resets the tuner cache if the environment switched, so
        // a miss here falls through to `check_tune`, which reads the store.
        if let TuneCacheResult::Hit { fastest_index } = tuner.fastest(&key) {
            // `autotune-checks` cross-checks the candidates of every launch, a hit's included.
            // Bound first, so the logger guard drops before the checks run every candidate.
            #[cfg(feature = "autotune-checks")]
            {
                let mut log_context =
                    crate::tune::AutotuneLogContext::new(&mut tuner.logger().lock());
                log_context.set_checks(|| self.checks::<I, Out, Id>(&operations, &inputs));
            }
            return operations
                .fastest(fastest_index)
                .execute(inputs)
                .expect("Should run when selected by autotune.");
        }

        // Built after the fast path: a hit logs nothing, so it takes no logger lock.
        #[allow(unused_mut)]
        let mut log_context = crate::tune::AutotuneLogContext::new(&mut tuner.logger().lock());

        #[cfg(feature = "autotune-checks")]
        log_context.set_checks(|| self.checks::<I, Out, Id>(&operations, &inputs));

        let fastest =
            tuner.check_tune::<I, Out, Id>(&key, &inputs, &operations, client, log_context);

        // Run the execution depending on the cache state.
        match fastest {
            TuneCacheResult::Hit { fastest_index } => operations
                .fastest(fastest_index)
                .execute(inputs)
                .expect("Should run when selected by autotune."),
            TuneCacheResult::Miss => {
                panic!(
                    "Somehow we STILL didn't check a tuning checksum or start tuning, something has gone wrong."
                )
            }
            // Still waiting (e.g. on wasm), stopped on a lost device, or its candidates' kernels
            // queued under a `CompileOnly` override: run the first operation that serves the
            // problem, which there only queues its kernels.
            TuneCacheResult::Pending | TuneCacheResult::Compiled => {
                let mut failures = Vec::new();
                for i in 0..operations.len() {
                    match operations.fastest(i).execute(inputs.clone()) {
                        Ok(output) => return output,
                        Err(err) => failures.push(err),
                    }
                }
                panic!("All autotune operations failed, no viable operation found.\n{failures:?}");
            }
        }
    }
}