candding 0.1.0

//! Verification status types every catalog shares: whether a model's output has been checked on a backend or in a
//! dtype, the per-device dtype axis derived from those two, and the serde shape a `DType` field uses.

mod cells;

use candle_core::DType;
use candle_core::Device;
pub use cells::Backend;
pub(crate) use cells::CellException;
#[cfg(test)]
pub(crate) use cells::rules;
use serde::Deserialize;
use serde::Serialize;

/// Whether a model's output has been checked on a backend or in a dtype with the golden or parity suites.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum Support {
  Verified,
  Pending,
  Untested,
  /// Checked and found not to work: the run missed the tolerance or the combination cannot execute.
  Unsupported,
  /// No such axis: the entry computes with no device and no dtype, as `Qdrant/bm25` does.
  NotApplicable,
}

/// One status per backend: the CPU, Metal and CUDA.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct BackendSupport {
  pub cpu: Support,
  pub metal: Support,
  pub cuda: Support,
}

impl BackendSupport {
  /// Untested on every backend, what a descriptor outside the registry starts from; assign a field to change one.
  pub fn untested() -> Self {
    Self {
      cpu: Support::Untested,
      metal: Support::Untested,
      cuda: Support::Untested,
    }
  }

  /// Whether this row's dtype runs only on an accelerator: refused on the CPU and not refused on Metal.
  pub fn needs_accelerator(&self) -> bool {
    self.cpu == Support::Unsupported && self.metal != Support::Unsupported
  }

  /// No device axis at all, as for `Qdrant/bm25`: unlike `untested`, there is nothing to check.
  pub const fn not_applicable() -> Self {
    Self {
      cpu: Support::NotApplicable,
      metal: Support::NotApplicable,
      cuda: Support::NotApplicable,
    }
  }

  /// The status for the backend a `Device` belongs to.
  pub fn for_device(&self, device: &Device) -> Support {
    self.get(Backend::of(device))
  }
}

/// Per-dtype verification status shared by every device, the fact each entry states;
/// `DtypeBackendSupport::derive` turns it into per-device rows.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct DtypeSupport {
  pub f32: Support,
  pub f16: Support,
  pub bf16: Support,
}

impl DtypeSupport {
  /// A new entry before its first golden run; a registry test fails while any entry is left this way.
  pub const fn pending() -> Self {
    Self {
      f32: Support::Pending,
      f16: Support::Pending,
      bf16: Support::Pending,
    }
  }

  /// No numeric-format axis at all, the dtype counterpart of `BackendSupport::not_applicable`.
  pub const fn not_applicable() -> Self {
    Self {
      f32: Support::NotApplicable,
      f16: Support::NotApplicable,
      bf16: Support::NotApplicable,
    }
  }

  /// Every dtype verified against the reference.
  pub const fn verified() -> Self {
    Self {
      f32: Support::Verified,
      f16: Support::Verified,
      bf16: Support::Verified,
    }
  }

  pub fn status(&self, dtype: DType) -> Support {
    match dtype {
      DType::F16 => self.f16,
      DType::BF16 => self.bf16,
      _ => self.f32,
    }
  }
}

/// One `BackendSupport` row per dtype: the catalog's cells, derived from each entry's statuses, and the row every
/// builder refuses from.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct DtypeBackendSupport {
  pub f32: BackendSupport,
  pub f16: BackendSupport,
  pub bf16: BackendSupport,
}

impl DtypeBackendSupport {
  /// `BackendSupport::untested` for every dtype, what a descriptor outside the registry starts from.
  pub fn untested() -> Self {
    Self {
      f32: BackendSupport::untested(),
      f16: BackendSupport::untested(),
      bf16: BackendSupport::untested(),
    }
  }

  /// The cpu/metal/cuda row for one dtype.
  pub fn row(&self, dtype: DType) -> BackendSupport {
    match dtype {
      DType::F16 => self.f16,
      DType::BF16 => self.bf16,
      _ => self.f32,
    }
  }

  /// One dtype on one device.
  pub fn status(&self, dtype: DType, device: &Device) -> Support {
    self.row(dtype).for_device(device)
  }

  /// Derives the per-device rows: a flatly `Unsupported` dtype everywhere, BF16 and `accel_only` dtypes refused on
  /// the CPU, an `Untested` or `Pending` device for every dtype, and the flat status otherwise.
  pub(crate) fn derive(
    backends: BackendSupport,
    dtypes: DtypeSupport,
    accel_only: &[DType],
  ) -> Self {
    let row = |dtype: DType| -> BackendSupport {
      let flat = dtypes.status(dtype);
      let cell = |device_status: Support, refused_here: bool| -> Support {
        if flat == Support::Unsupported || refused_here {
          Support::Unsupported
        } else if matches!(device_status, Support::Untested | Support::Pending) {
          device_status
        } else {
          flat
        }
      };
      let no_cpu_bf16 = dtype == DType::BF16 && backends.cpu != Support::NotApplicable;
      BackendSupport {
        cpu: cell(backends.cpu, no_cpu_bf16 || accel_only.contains(&dtype)),
        metal: cell(backends.metal, false),
        cuda: cell(backends.cuda, false),
      }
    };
    Self {
      f32: row(DType::F32),
      f16: row(DType::F16),
      bf16: row(DType::BF16),
    }
  }
}

/// `DType` as the lowercase name the CLI, JSON and docs use; only the three float dtypes are descriptor values.
pub(crate) mod dtype_serde {
  use candle_core::DType;
  use serde::Deserialize;
  use serde::Deserializer;
  use serde::Serializer;

  pub fn serialize<S: Serializer>(dtype: &DType, serializer: S) -> Result<S::Ok, S::Error> {
    serializer.serialize_str(dtype.as_str())
  }

  pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<DType, D::Error> {
    let name = String::deserialize(deserializer)?;
    match name.as_str() {
      "f32" => Ok(DType::F32),
      "f16" => Ok(DType::F16),
      "bf16" => Ok(DType::BF16),
      other => Err(serde::de::Error::custom(format!(
        "unsupported model dtype `{other}`; expected f32, f16 or bf16"
      ))),
    }
  }
}

#[cfg(test)]
mod tests;