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::{Arc, Condvar, Mutex, OnceLock, PoisonError};

use allocator::Pool;
pub use buffer::Buffer;
pub use command::{Arg, CommandBuffer, Pending, Submitted, Timeline};
use nuvec_kernels::MetallibMode;
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::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>,
    by_name: HashMap<String, usize>,
}

impl Libraries {
    fn new(libraries: Vec<Library>) -> Self {
        let mut by_name = HashMap::new();
        for (i, library) in libraries.iter().enumerate() {
            for name in library.functionNames().iter() {
                by_name.entry(name.to_string()).or_insert(i);
            }
        }
        Libraries { libraries, by_name }
    }

    /// The library that defines kernel `name`.
    fn get(&self, name: &str) -> Result<&Library> {
        self.by_name
            .get(name)
            .map(|&i| &self.libraries[i])
            .ok_or_else(|| 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>>,
    /// 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,
        };
        let ram = objc2_foundation::NSProcessInfo::processInfo().physicalMemory();
        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(),
                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 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> {
        let key = constants.iter().fold(name.to_owned(), |mut k, (i, v)| {
            match v {
                Constant::Bool(b) => k.push_str(&format!("#{i}={}", u8::from(*b))),
                Constant::Int(n) => k.push_str(&format!("#{i}=i{n}")),
            }
            k
        });
        // Loaded (or compiled, from runtime sources) before any cache lock is taken.
        let library = self.library()?.get(name)?;
        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
    /// nuvec's own kernels are, with Metal 3.2 and IEEE-conformant math.
    ///
    /// 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())]);
            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)
    }

    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);
                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),
    }
}

/// The `.metal` sources as one translation unit, each `#include "header.h"` of one of
/// nuvec's headers 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")) {
        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(h).map(|text| (h, text))) {
                Some((h, text)) => {
                    if included.insert(h) {
                        out.push_str(text);
                        out.push('\n');
                    }
                }
                None => {
                    out.push_str(line);
                    out.push('\n');
                }
            }
        }
    }
    out
}

/// 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\""));
    }
}