nuvec_gpu/metal/

mod.rs

//! The Metal runtime on `objc2-metal` (macOS only).

mod allocator;
mod buffer;
mod command;
pub mod profile;
mod residency;

use std::collections::{HashMap, HashSet};
use std::ptr::NonNull;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Condvar, Mutex, OnceLock, PoisonError};

use allocator::Pool;
pub use buffer::Buffer;
pub use command::{Arg, CommandBuffer, Pending, Submitted, Timeline};
use nuvec_kernels::{MetallibMode, NaxMode};
use objc2::rc::{Retained, autoreleasepool};
use objc2::runtime::ProtocolObject;
use objc2_foundation::NSString;
use objc2_metal::{
    MTLCommandQueue, MTLComputePipelineState, MTLCreateSystemDefaultDevice, MTLDataType, MTLDevice,
    MTLFunctionConstantValues, MTLLibrary,
};

use crate::info::{NaxOverride, NaxReason, NaxStatus, nax_status, parse_architecture};
use crate::{DeviceInfo, GpuError, MemoryStats, Result};

type Library = Retained<ProtocolObject<dyn MTLLibrary>>;

/// nuvec's kernels: one precompiled library, or (runtime sources) one library per shader
/// file, with the library of each kernel name.
struct Libraries {
    libraries: Vec<Library>,
    /// Kernel names, sorted, with their library (the first that defines each). A sorted
    /// vector rather than a hash map: every allocation is referenced from its start, so
    /// `leaks` never mistakes this lifelong index for leaked memory (a hash table is
    /// referenced from inside its allocation).
    by_name: Vec<(Box<str>, usize)>,
}

impl Libraries {
    fn new(libraries: Vec<Library>) -> Self {
        let mut by_name: Vec<(Box<str>, usize)> = Vec::new();
        for (i, library) in libraries.iter().enumerate() {
            for name in library.functionNames().iter() {
                by_name.push((name.to_string().into_boxed_str(), i));
            }
        }
        // Stable: the first library's entry of a name comes first and is kept.
        by_name.sort_by(|a, b| a.0.cmp(&b.0));
        by_name.dedup_by(|later, first| later.0 == first.0);
        Libraries { libraries, by_name }
    }

    /// The library that defines kernel `name`.
    fn get(&self, name: &str) -> Result<&Library> {
        self.by_name
            .binary_search_by(|(n, _)| n.as_ref().cmp(name))
            .map(|at| &self.libraries[self.by_name[at].1])
            .map_err(|_| GpuError::MissingKernel(name.to_owned()))
    }
}

/// A pipeline cache entry: empty until compiled (a failed compilation leaves it empty).
type Slot = Arc<Mutex<Option<Pipeline>>>;

/// A function constant value for [`Device::pipeline_with_constants`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Constant {
    /// A `bool` constant.
    Bool(bool),
    /// An `int` constant.
    Int(i32),
}

/// A Metal GPU with its command queue, kernel library and pipeline cache. Cheap to clone.
#[derive(Clone)]
pub struct Device {
    inner: Arc<Inner>,
}

struct Inner {
    device: Retained<ProtocolObject<dyn MTLDevice>>,
    queue: Retained<ProtocolObject<dyn MTLCommandQueue>>,
    info: DeviceInfo,
    library: OnceLock<Result<Libraries>>,
    /// The Neural Accelerator (NAX) kernels' library, loaded on first use (macOS 26.2+).
    nax_library: OnceLock<Result<Libraries>>,
    /// The NAX decision, made once ([`Device::nax`]).
    nax: OnceLock<NaxStatus>,
    /// Set when a NAX pipeline failed to build: the existing kernels from then on.
    nax_failed: AtomicBool,
    /// Compiled pipelines by key: library kernels by name (and function constants), source
    /// kernels by `src:{name}`. The map's lock is held only to find a key's slot; a slot's
    /// lock is held while that one pipeline compiles, so a thread asking for the same key
    /// waits for that compilation (one compilation per key) and other keys are not blocked.
    pipelines: Mutex<HashMap<String, Slot>>,
    pool: Arc<Pool>,
    in_flight: InFlight,
}

impl Drop for Inner {
    /// Releases the cached buffers in an autorelease pool (a Metal object's deallocation
    /// may autorelease others, and the last device handle may be dropped on any thread).
    fn drop(&mut self) {
        autoreleasepool(|_| self.pool.clear_cache());
    }
}

/// The number of committed command buffers that have not completed yet.
#[derive(Default)]
pub(crate) struct InFlight {
    count: Mutex<usize>,
    done: Condvar,
}

impl InFlight {
    pub(crate) fn begin(&self) {
        *self.count.lock().unwrap_or_else(PoisonError::into_inner) += 1;
    }

    pub(crate) fn end(&self) {
        let mut count = self.count.lock().unwrap_or_else(PoisonError::into_inner);
        *count = count.saturating_sub(1);
        self.done.notify_all();
    }

    fn get(&self) -> usize {
        *self.count.lock().unwrap_or_else(PoisonError::into_inner)
    }

    fn wait_below(&self, max: usize) {
        let mut count = self.count.lock().unwrap_or_else(PoisonError::into_inner);
        while *count >= max.max(1) {
            count = self
                .done
                .wait(count)
                .unwrap_or_else(PoisonError::into_inner);
        }
    }
}

impl Device {
    /// The system's default Metal device, or `None` if there is none (for example in some
    /// virtual machines).
    pub fn system_default() -> Option<Device> {
        autoreleasepool(|_| Self::create())
    }

    fn create() -> Option<Device> {
        let device = MTLCreateSystemDefaultDevice()?;
        let queue = device.newCommandQueue()?;
        let architecture = device.architecture().name().to_string();
        let (generation, class) = parse_architecture(&architecture);
        let info = DeviceInfo {
            name: device.name().to_string(),
            architecture,
            generation,
            class,
            recommended_working_set: device.recommendedMaxWorkingSetSize(),
            max_threads_per_threadgroup: device.maxThreadsPerThreadgroup().width,
            physical_memory: objc2_foundation::NSProcessInfo::processInfo().physicalMemory(),
        };
        let ram = info.physical_memory;
        let options = objc2_metal::MTLResourceOptions::StorageModeShared;
        let placeholder = device.newBufferWithLength_options(16, options);
        let residency = residency::Residency::new(residency::attach(&device, &queue));
        let pool = Arc::new(Pool::new(
            info.recommended_working_set,
            ram,
            placeholder,
            residency,
        ));
        Some(Device {
            inner: Arc::new(Inner {
                device,
                queue,
                info,
                library: OnceLock::new(),
                nax_library: OnceLock::new(),
                nax: OnceLock::new(),
                nax_failed: AtomicBool::new(false),
                pipelines: Mutex::new(HashMap::new()),
                pool,
                in_flight: InFlight::default(),
            }),
        })
    }

    /// The process-wide default device, created on first use: one allocator, kernel cache
    /// and queue shared by everything that uses the GPU.
    pub fn shared() -> Option<&'static Device> {
        static SHARED: OnceLock<Option<Device>> = OnceLock::new();
        SHARED.get_or_init(Device::system_default).as_ref()
    }

    /// Memory use and limits of this device's allocator.
    pub fn memory(&self) -> MemoryStats {
        self.inner.pool.stats()
    }

    /// Sets the memory limit and returns the previous one (as MLX's `set_memory_limit`).
    ///
    /// The limit bounds `active + mapped` ([`MemoryStats::active_with_mapped`]): an
    /// allocation or a file mapping that would pass it fails with
    /// [`GpuError::OutOfMemory`] before anything is wired, so a server's memory budget can
    /// refuse work instead of swapping. This deliberately differs from MLX 0.32.4, whose
    /// limit is "a guideline for the maximum amount of memory to use during graph
    /// evaluation" (allocations fail there only when no RAM or swap is left). A limit below
    /// what is already in use is accepted; the next allocation or mapping then fails. The
    /// default is `min(1.5 × recommended working set, 0.95 × RAM)`, as MLX.
    pub fn set_memory_limit(&self, bytes: usize) -> usize {
        autoreleasepool(|_| self.inner.pool.set_memory_limit(bytes))
    }

    /// Sets the cache limit, trimming the cache to it, and returns the previous one.
    pub fn set_cache_limit(&self, bytes: usize) -> usize {
        autoreleasepool(|_| self.inner.pool.set_cache_limit(bytes))
    }

    /// Sets the wired limit and returns the previous one (MLX's `set_wired_limit`; 0 by
    /// default). Buffers, mapped files included, are kept resident up to this many bytes
    /// while GPU work runs, instead of being made resident per command buffer, which can
    /// page memory in again under pressure. Lowering it unwires buffers until the wired set
    /// fits. A limit above the recommended working set is an error. On a GPU without
    /// residency sets (before Metal 3) the limit is recorded but nothing is wired.
    pub fn set_wired_limit(&self, bytes: usize) -> Result<usize> {
        if bytes as u64 > self.inner.info.recommended_working_set {
            return Err(GpuError::InvalidArgument(format!(
                "wired limit {bytes} exceeds the recommended working set ({} bytes)",
                self.inner.info.recommended_working_set
            )));
        }
        Ok(autoreleasepool(|_| self.inner.pool.set_wired_limit(bytes)))
    }

    /// Whether buffers can be wired on this system (residency sets: a Metal 3 GPU).
    pub fn can_wire(&self) -> bool {
        self.inner.pool.can_wire()
    }

    /// Releases every cached buffer.
    pub fn clear_cache(&self) {
        autoreleasepool(|_| self.inner.pool.clear_cache());
    }

    /// Command buffers committed on this device that have not completed yet.
    pub fn in_flight(&self) -> usize {
        self.inner.in_flight.get()
    }

    /// Blocks until fewer than `max` command buffers are in flight (back-pressure for
    /// callers that submit without waiting).
    pub fn wait_for_capacity(&self, max: usize) {
        self.inner.in_flight.wait_below(max);
    }

    pub(crate) fn flights(&self) -> &InFlight {
        &self.inner.in_flight
    }

    /// Resets the peaks to the current use: `peak` to `active`, `peak_with_mapped` to
    /// `active + mapped`.
    pub fn reset_peak(&self) {
        self.inner.pool.reset_peak();
    }

    pub(crate) fn pool(&self) -> &Arc<Pool> {
        &self.inner.pool
    }

    /// What this GPU is and can do.
    pub fn info(&self) -> &DeviceInfo {
        &self.inner.info
    }

    /// The names of the kernels in nuvec's kernel library, sorted (loading the library on
    /// first use).
    pub fn kernel_names(&self) -> Result<Vec<String>> {
        // Sorted already.
        Ok(self
            .library()?
            .by_name
            .iter()
            .map(|(n, _)| n.to_string())
            .collect())
    }

    /// The compute pipeline for the kernel `name` from nuvec's kernel library, compiled on
    /// first use and cached for the life of the device. The cache is keyed by kernel name
    /// and function constants, so it is bounded by the library's kernels times the constant
    /// combinations the operations use (a few thousand at most); pipelines from source are
    /// keyed by their name (one entry per distinct custom kernel instantiation).
    pub fn pipeline(&self, name: &str) -> Result<Pipeline> {
        self.pipeline_with(name, &[])
    }

    /// The pipeline for kernel `name` specialised with boolean function constants
    /// `(index, value)`, compiled on first use and cached per name and constants.
    pub fn pipeline_with(&self, name: &str, constants: &[(usize, bool)]) -> Result<Pipeline> {
        let constants: Vec<(usize, Constant)> = constants
            .iter()
            .map(|&(i, v)| (i, Constant::Bool(v)))
            .collect();
        self.pipeline_with_constants(name, &constants)
    }

    /// The pipeline for kernel `name` specialised with boolean or integer function constants
    /// `(index, value)`, compiled on first use and cached per name and constants.
    pub fn pipeline_with_constants(
        &self,
        name: &str,
        constants: &[(usize, Constant)],
    ) -> Result<Pipeline> {
        thread_local! {
            /// The cache key being looked up, reused across calls (a hit allocates nothing).
            static KEY: std::cell::RefCell<String> = const { std::cell::RefCell::new(String::new()) };
        }
        let hit = KEY.with(|k| {
            use std::fmt::Write;
            let mut k = k.borrow_mut();
            k.clear();
            k.push_str(name);
            for (i, v) in constants {
                let _ = match v {
                    Constant::Bool(b) => write!(k, "#{i}={}", u8::from(*b)),
                    Constant::Int(n) => write!(k, "#{i}=i{n}"),
                };
            }
            let slot = {
                let cache = self
                    .inner
                    .pipelines
                    .lock()
                    .unwrap_or_else(PoisonError::into_inner);
                cache.get(k.as_str()).cloned()
            };
            match slot.and_then(|slot| slot.lock().unwrap_or_else(PoisonError::into_inner).clone())
            {
                Some(pipeline) => Ok(pipeline),
                None => Err(k.clone()),
            }
        });
        let key = match hit {
            Ok(pipeline) => return Ok(pipeline),
            Err(key) => key,
        };
        // Loaded (or compiled, from runtime sources) before any cache lock is taken. Any
        // NAX failure while NAX is active (no library, a missing kernel, a pipeline error)
        // turns it off for the process, so later products take the existing kernels.
        let nax = name.starts_with("nuvec_nax_");
        let built = match nax {
            true => self.nax_library().and_then(|l| l.get(name)),
            false => self.library().and_then(|l| l.get(name)),
        }
        .and_then(|library| self.build_pipeline(name, &key, library, constants));
        if nax && built.is_err() && self.nax().is_active() {
            self.inner.nax_failed.store(true, Ordering::Relaxed);
        }
        built
    }

    fn build_pipeline(
        &self,
        name: &str,
        key: &str,
        library: &Library,
        constants: &[(usize, Constant)],
    ) -> Result<Pipeline> {
        let key = key.to_owned();
        self.cached(key.clone(), || {
            let ns_name = NSString::from_str(name);
            let function = if constants.is_empty() {
                library
                    .newFunctionWithName(&ns_name)
                    .ok_or_else(|| GpuError::MissingKernel(name.to_owned()))?
            } else {
                let values = MTLFunctionConstantValues::new();
                for &(index, value) in constants {
                    match value {
                        // SAFETY: `b` is a live local, which Metal copies during the call;
                        // `bool` matches `MTLDataType::Bool` (one byte).
                        Constant::Bool(b) => unsafe {
                            values.setConstantValue_type_atIndex(
                                NonNull::from(&b).cast(),
                                MTLDataType::Bool,
                                index,
                            );
                        },
                        // SAFETY: as above; `i32` matches `MTLDataType::Int`.
                        Constant::Int(n) => unsafe {
                            values.setConstantValue_type_atIndex(
                                NonNull::from(&n).cast(),
                                MTLDataType::Int,
                                index,
                            );
                        },
                    }
                }
                library
                    .newFunctionWithName_constantValues_error(&ns_name, &values)
                    .map_err(|e| GpuError::Pipeline {
                        name: key.clone(),
                        message: e.localizedDescription().to_string(),
                    })?
            };
            let state = self
                .inner
                .device
                .newComputePipelineStateWithFunction_error(&function)
                .map_err(|e| GpuError::Pipeline {
                    name: key.clone(),
                    message: e.localizedDescription().to_string(),
                })?;
            Ok(Pipeline {
                state,
                name: key.as_str().into(),
            })
        })
    }

    /// The pipeline for kernel `name` from Metal source that `source` generates, compiled
    /// at runtime on first use and cached by name (so a name must always mean the same
    /// source; `source` runs only on a cache miss). Source kernels have their own namespace:
    /// a source kernel named like one of nuvec's library kernels is a different pipeline.
    /// The source may include nuvec's headers (`#include "common.h"`); it is compiled as
    /// MLX compiles its runtime kernels: Metal 3.2, safe math, Metal's default (fast)
    /// math functions.
    ///
    /// The source is trusted code: the pipeline is run with the `unsafe`
    /// [`CommandBuffer::dispatch`] family, whose contract (in-bounds accesses only) the
    /// kernel must keep.
    pub fn pipeline_from_source(
        &self,
        name: &str,
        source: impl FnOnce() -> String,
    ) -> Result<Pipeline> {
        self.cached(format!("src:{name}"), || {
            let failed = |message: String| GpuError::Pipeline {
                name: name.to_owned(),
                message,
            };
            let source = inline_sources(&[("generated.metal", &source())]);
            // As MLX compiles its runtime kernels (`mx.compile`, custom kernels): safe math
            // with Metal's default, fast `metal::` functions (pow, exp, tanh, …), so their
            // results can differ in the low bits from the precompiled kernels' (D-104).
            let options = objc2_metal::MTLCompileOptions::new();
            options.setMathMode(objc2_metal::MTLMathMode::Safe);
            options.setLanguageVersion(objc2_metal::MTLLanguageVersion::Version3_2);
            let library = self
                .inner
                .device
                .newLibraryWithSource_options_error(&NSString::from_str(&source), Some(&options))
                .map_err(|e| failed(e.localizedDescription().to_string()))?;
            let function = library
                .newFunctionWithName(&NSString::from_str(name))
                .ok_or_else(|| GpuError::MissingKernel(name.to_owned()))?;
            let state = self
                .inner
                .device
                .newComputePipelineStateWithFunction_error(&function)
                .map_err(|e| failed(e.localizedDescription().to_string()))?;
            Ok(Pipeline {
                state,
                name: name.into(),
            })
        })
    }

    /// The cached pipeline for `key`, compiled by `compile` on a miss. The cache's lock is
    /// not held while compiling (only the key's own slot is); compilation runs in an
    /// autorelease pool, so the temporary Objective-C objects it creates are released on
    /// long-lived threads too.
    fn cached(&self, key: String, compile: impl FnOnce() -> Result<Pipeline>) -> Result<Pipeline> {
        let slot = {
            let mut cache = self
                .inner
                .pipelines
                .lock()
                .unwrap_or_else(PoisonError::into_inner);
            cache.entry(key).or_default().clone()
        };
        let mut slot = slot.lock().unwrap_or_else(PoisonError::into_inner);
        if let Some(p) = &*slot {
            return Ok(p.clone());
        }
        let pipeline = autoreleasepool(|_| compile())?;
        *slot = Some(pipeline.clone());
        Ok(pipeline)
    }

    /// nuvec's kernel libraries, loaded once.
    fn library(&self) -> Result<&Libraries> {
        self.inner
            .library
            .get_or_init(|| autoreleasepool(|_| load_library(&self.inner.device)))
            .as_ref()
            .map_err(Clone::clone)
    }

    fn nax_library(&self) -> Result<&Libraries> {
        self.inner
            .nax_library
            .get_or_init(|| autoreleasepool(|_| load_nax_library(&self.inner.device)))
            .as_ref()
            .map_err(Clone::clone)
    }

    /// Whether products run on the Neural Accelerator kernels (Metal 4 TensorOps), and why
    /// not: MLX's eligibility (an M5-class GPU on macOS 26.2 or later), the `NUVEC_NAX`
    /// override (`auto`, `on`, `off`), and whether this build has the kernels. `auto` keeps
    /// them off on eligible GPUs until an M5 measurement is recorded; `on` forces them on
    /// any GPU with TensorOps (macOS 26.2), for testing. After a NAX pipeline fails to build,
    /// [`NaxReason::PipelineError`].
    pub fn nax(&self) -> NaxStatus {
        if self.inner.nax_failed.load(Ordering::Relaxed) {
            return NaxStatus::Off(NaxReason::PipelineError);
        }
        *self.inner.nax.get_or_init(|| {
            let v = objc2_foundation::NSProcessInfo::processInfo().operatingSystemVersion();
            let os = (v.majorVersion as u32, v.minorVersion as u32);
            let info = &self.inner.info;
            nax_status(
                info.generation,
                info.class,
                os,
                NaxOverride::from_env(),
                nuvec_kernels::NAX_MODE != NaxMode::Absent,
            )
        })
    }

    pub(crate) fn raw(&self) -> &ProtocolObject<dyn MTLDevice> {
        &self.inner.device
    }

    pub(crate) fn queue(&self) -> &ProtocolObject<dyn MTLCommandQueue> {
        &self.inner.queue
    }
}

impl std::fmt::Debug for Device {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Device")
            .field("info", &self.inner.info)
            .finish_non_exhaustive()
    }
}

/// Loads the precompiled library, or compiles the embedded sources as the build script
/// would have: each shader file on its own (headers inlined, as the runtime compiler has no
/// include path for them), with the same language version, IEEE-conformant math and
/// floating-point contraction setting, so that both modes give the same results. The files
/// compile in parallel.
fn load_library(device: &ProtocolObject<dyn MTLDevice>) -> Result<Libraries> {
    let error = |e: Retained<objc2_foundation::NSError>| {
        GpuError::LibraryLoad(e.localizedDescription().to_string())
    };
    match nuvec_kernels::MODE {
        MetallibMode::Precompiled => {
            let bytes = nuvec_kernels::metallib().ok_or(GpuError::NoKernels)?;
            let library = device
                .newLibraryWithData_error(&dispatch2::DispatchData::from_static_bytes(bytes))
                .map_err(error)?;
            Ok(Libraries::new(vec![library]))
        }
        MetallibMode::RuntimeSource => {
            let files: Vec<(&str, &str)> = nuvec_kernels::sources()
                .iter()
                .copied()
                .filter(|(p, _)| p.ends_with(".metal"))
                .collect();
            let compile = |&(path, text): &(&str, &str)| -> Result<Library> {
                let mut source = String::new();
                if !nuvec_kernels::contracted(path) {
                    // The build script's `-ffp-contract=off`.
                    source.push_str("#pragma clang fp contract(off)\n");
                }
                source.push_str(&inline_sources(&[(path, text)]));
                let options = objc2_metal::MTLCompileOptions::new();
                options.setMathMode(objc2_metal::MTLMathMode::Safe);
                // The build's `-fno-fast-math`: `metal::` functions (pow, exp, …) are the precise
                // ones (the default is the fast ones).
                options.setMathFloatingPointFunctions(
                    objc2_metal::MTLMathFloatingPointFunctions::Precise,
                );
                options.setLanguageVersion(objc2_metal::MTLLanguageVersion::Version3_2);
                device
                    .newLibraryWithSource_options_error(
                        &NSString::from_str(&source),
                        Some(&options),
                    )
                    .map_err(|e| {
                        GpuError::LibraryLoad(format!("{path}: {}", e.localizedDescription()))
                    })
            };
            let libraries = std::thread::scope(|scope| {
                let handles: Vec<_> = files
                    .iter()
                    .map(|file| scope.spawn(move || compile(file)))
                    .collect();
                handles
                    .into_iter()
                    .map(|h| h.join().expect("shader compilation does not panic"))
                    .collect::<Result<Vec<_>>>()
            })?;
            Ok(Libraries::new(libraries))
        }
        _ => Err(GpuError::NoKernels),
    }
}

/// Loads the precompiled NAX library, or compiles its embedded sources with Metal 4.0 (only
/// called when [`Device::nax`] is active, so on macOS 26.2 or later).
fn load_nax_library(device: &ProtocolObject<dyn MTLDevice>) -> Result<Libraries> {
    match nuvec_kernels::NAX_MODE {
        NaxMode::Precompiled => {
            let bytes = nuvec_kernels::nax_metallib().ok_or(GpuError::NoKernels)?;
            let library = device
                .newLibraryWithData_error(&dispatch2::DispatchData::from_static_bytes(bytes))
                .map_err(|e| GpuError::LibraryLoad(e.localizedDescription().to_string()))?;
            Ok(Libraries::new(vec![library]))
        }
        NaxMode::RuntimeSource => {
            let libraries = nuvec_kernels::nax_sources()
                .iter()
                .filter(|(p, _)| p.ends_with(".metal"))
                .map(|(path, text)| {
                    let options = objc2_metal::MTLCompileOptions::new();
                    options.setMathMode(objc2_metal::MTLMathMode::Safe);
                    options.setMathFloatingPointFunctions(
                        objc2_metal::MTLMathFloatingPointFunctions::Precise,
                    );
                    options.setLanguageVersion(objc2_metal::MTLLanguageVersion::Version4_0);
                    device
                        .newLibraryWithSource_options_error(
                            &NSString::from_str(text),
                            Some(&options),
                        )
                        .map_err(|e| {
                            GpuError::LibraryLoad(format!("{path}: {}", e.localizedDescription()))
                        })
                })
                .collect::<Result<Vec<_>>>()?;
            Ok(Libraries::new(libraries))
        }
        _ => Err(GpuError::NoKernels),
    }
}

/// The `.metal` sources as one translation unit, each `#include "header.h"` of one of
/// nuvec's headers (in a source or in another header) replaced by the header's text on
/// first use.
fn inline_sources(sources: &[(&str, &str)]) -> String {
    let headers: HashMap<&str, &str> = nuvec_kernels::sources()
        .iter()
        .chain(sources)
        .copied()
        .chain([("common.h", nuvec_kernels::COMMON_H)])
        .filter(|(p, _)| p.ends_with(".h"))
        .collect();
    let mut included = HashSet::new();
    let mut out = String::new();
    for (_, source) in sources.iter().filter(|(p, _)| p.ends_with(".metal")) {
        inline_into(source, &headers, &mut included, &mut out);
    }
    out
}

fn inline_into<'a>(
    source: &str,
    headers: &HashMap<&'a str, &'a str>,
    included: &mut HashSet<&'a str>,
    out: &mut String,
) {
    for line in source.lines() {
        let header = line
            .trim()
            .strip_prefix("#include \"")
            .and_then(|r| r.strip_suffix('"'));
        match header.and_then(|h| headers.get_key_value(h)) {
            Some((&h, &text)) => {
                if included.insert(h) {
                    inline_into(text, headers, included, out);
                    out.push('\n');
                }
            }
            None => {
                out.push_str(line);
                out.push('\n');
            }
        }
    }
}

/// A compiled compute kernel. Cheap to clone.
#[derive(Clone)]
pub struct Pipeline {
    state: Retained<ProtocolObject<dyn MTLComputePipelineState>>,
    name: Arc<str>,
}

impl Pipeline {
    /// The kernel's name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// The most threads one threadgroup of this kernel may have.
    pub fn max_threads_per_threadgroup(&self) -> usize {
        self.state.maxTotalThreadsPerThreadgroup()
    }

    /// The SIMD-group width (threads that execute in lockstep).
    pub fn execution_width(&self) -> usize {
        self.state.threadExecutionWidth()
    }

    /// Whether two handles refer to the same compiled pipeline.
    pub fn same_as(&self, other: &Pipeline) -> bool {
        Retained::as_ptr(&self.state) == Retained::as_ptr(&other.state)
    }

    pub(crate) fn raw(&self) -> &ProtocolObject<dyn MTLComputePipelineState> {
        &self.state
    }
}

impl std::fmt::Debug for Pipeline {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Pipeline")
            .field("name", &self.name)
            .finish_non_exhaustive()
    }
}

#[cfg(test)]
mod tests {
    use super::inline_sources;

    #[test]
    fn inlines_each_header_once() {
        let sources = [
            ("a.metal", "#include \"h.h\"\nkernel A;"),
            ("b.metal", "  #include \"h.h\"\nkernel B;"),
            ("h.h", "#pragma once\ninline int f();"),
        ];
        let out = inline_sources(&sources);
        assert_eq!(out.matches("inline int f();").count(), 1);
        assert!(out.contains("kernel A;") && out.contains("kernel B;"));
        assert!(!out.contains("#include \"h.h\""));
    }

    #[test]
    fn inlines_headers_included_by_headers() {
        let sources = [
            (
                "a.metal",
                "#include \"outer.h\"\n#include \"inner.h\"\nkernel A;",
            ),
            ("outer.h", "#include \"inner.h\"\ninline int g();"),
            ("inner.h", "inline int f();"),
        ];
        let out = inline_sources(&sources);
        assert_eq!(out.matches("inline int f();").count(), 1);
        assert!(out.find("inline int f();") < out.find("inline int g();"));
        assert!(!out.contains("#include"));
    }
}