nuvec 0.6.0

Pure-Rust MLX for Apple Silicon: lazy arrays on unified memory, composable transforms and Metal kernels
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

//! The Metal device (MLX's `mx.metal`): availability, device facts, the Neural
//! Accelerator kernels' status, GPU captures and the kernel profile. Every function works
//! without a Metal device.

use std::path::Path;

use nuvec_gpu::Device;

pub use nuvec_gpu::{DeviceInfo, NaxReason, NaxStatus};

use crate::{Error, Result};

/// Whether a Metal GPU is available (MLX's `mx.metal.is_available`).
pub fn is_available() -> bool {
    Device::shared().is_some()
}

/// What the GPU is and can do (MLX's `mx.metal.device_info`), or `None` without one.
pub fn device_info() -> Option<&'static DeviceInfo> {
    Device::shared().map(Device::info)
}

/// Whether products run on the Neural Accelerator (NAX) kernels, and why not.
pub fn nax() -> NaxStatus {
    Device::shared().map_or(NaxStatus::Off(NaxReason::NoDevice), Device::nax)
}

/// Starts capturing GPU work into `path`, a `.gputrace` document that does not exist yet
/// (MLX's `mx.metal.start_capture`). Metal allows it only with `MTL_CAPTURE_ENABLED=1` in
/// the environment, or under Xcode.
pub fn start_capture(path: impl AsRef<Path>) -> Result<()> {
    let d = Device::shared()
        .ok_or_else(|| Error::Backend("GPU capture: there is no Metal device".into()))?;
    d.start_capture(path)
        .map_err(|e| Error::Backend(e.to_string()))
}

/// Stops the running capture and writes its document (MLX's `mx.metal.stop_capture`);
/// nothing if none runs.
pub fn stop_capture() {
    if let Some(d) = Device::shared() {
        d.stop_capture();
    }
}

/// Per-kernel GPU times, collected when `NUVEC_PROFILE` is set (to anything but `0`): every
/// dispatch runs alone and is timed, so use the shares to find hotspots and the benchmarks
/// for speed.
pub mod profile {
    /// Whether profiling is on.
    pub fn enabled() -> bool {
        nuvec_gpu::profile::enabled()
    }

    /// The totals so far, `(kernel, dispatches, GPU seconds)`, the most expensive first.
    pub fn report() -> Vec<(String, u64, f64)> {
        nuvec_gpu::profile::report()
    }

    /// Clears the totals.
    pub fn reset() {
        nuvec_gpu::profile::reset();
    }
}