libmir-metal 0.1.0

Metal inference backend for libmir
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};

use super::{
    Result,
    compiled::CompiledGraphs,
    kernels::{Kernels, PageWriteOptions, PreparedPageWrite},
};

#[derive(Debug)]
pub struct Stream {
    native: mirtal::Stream,
    compiled: CompiledGraphs,
    kernels: Kernels,
    config: Arc<crate::MetalConfig>,
    graph_dumped: AtomicBool,
}

impl Stream {
    pub fn new_gpu() -> Result<Self> {
        Self::new(mirtal::Device::gpu(0), Arc::default())
    }

    pub(crate) fn new_gpu_with_config(config: Arc<crate::MetalConfig>) -> Result<Self> {
        Self::new(mirtal::Device::gpu(0), config)
    }

    pub fn new_cpu() -> Result<Self> {
        Self::new(mirtal::Device::cpu(0), Arc::default())
    }

    fn new(device: mirtal::Device, config: Arc<crate::MetalConfig>) -> Result<Self> {
        let native = device.new_stream()?;
        let compiled = CompiledGraphs::new(&native)?;
        let kernels = Kernels::new()?;
        Ok(Self {
            native,
            compiled,
            kernels,
            config,
            graph_dumped: AtomicBool::new(false),
        })
    }

    pub fn synchronize(&self) -> Result<()> {
        Ok(self.native.synchronize()?)
    }

    pub(super) const fn native(&self) -> &mirtal::Stream {
        &self.native
    }

    pub(crate) fn config(&self) -> &crate::MetalConfig {
        &self.config
    }

    pub(crate) fn take_graph_dump_path(&self) -> Option<&std::path::Path> {
        self.config
            .diagnostics
            .graph_dump
            .as_deref()
            .filter(|_| !self.graph_dumped.swap(true, Ordering::Relaxed))
    }

    pub(super) fn gelu_approx_mul(
        &self,
        gate: &mirtal::Array,
        input: &mirtal::Array,
    ) -> Result<mirtal::Array> {
        Ok(self.compiled.geglu(gate, input, &self.native)?)
    }

    pub(super) fn silu_mul(
        &self,
        gate: &mirtal::Array,
        input: &mirtal::Array,
    ) -> Result<mirtal::Array> {
        Ok(self.compiled.swiglu(gate, input, &self.native)?)
    }

    pub(super) fn precise_silu_mul(
        &self,
        reference: &mirtal::Array,
        gate: &mirtal::Array,
        input: &mirtal::Array,
    ) -> Result<mirtal::Array> {
        Ok(self.compiled.precise_swiglu(reference, gate, input, &self.native)?)
    }

    pub(super) fn logit_softcap(&self, input: &mirtal::Array, cap: f32) -> Result<mirtal::Array> {
        let cap = mirtal::Array::from_slice(&[cap], [])?;
        Ok(self.compiled.logit_softcap(input, &cap, &self.native)?)
    }

    pub(super) fn affine_router(&self, inputs: [&mirtal::Array; 6]) -> Result<[mirtal::Array; 2]> {
        Ok(self.compiled.router(inputs, &self.native)?)
    }

    pub(super) fn paged_attention(
        &self,
        inputs: [&mirtal::Array; 5],
        scratch: &super::attention::PagedAttentionScratch,
        page_size: usize,
        context_tokens: usize,
        scale: f32,
    ) -> Result<mirtal::Array> {
        self.kernels
            .paged_attention(&self.native, inputs, scratch, page_size, context_tokens, scale)
    }

    pub(super) fn gated_delta_gates(
        &self,
        inputs: [&mirtal::Array; 4],
    ) -> Result<[mirtal::Array; 2]> {
        self.kernels.gated_delta_gates(&self.native, inputs)
    }

    pub(super) fn gated_delta_recurrence(
        &self,
        inputs: [&mirtal::Array; 6],
    ) -> Result<[mirtal::Array; 2]> {
        self.kernels.gated_delta_recurrence(&self.native, inputs)
    }

    pub(super) fn gated_delta_decode(
        &self,
        inputs: [&mirtal::Array; 8],
        normalize: bool,
    ) -> Result<[mirtal::Array; 2]> {
        self.kernels.gated_delta_decode(&self.native, inputs, normalize)
    }

    pub(super) fn page_write(
        &self,
        inputs: [&mirtal::Array; 5],
        options: PageWriteOptions,
        prepared: &mut PreparedPageWrite,
    ) -> Result<[mirtal::Array; 2]> {
        self.kernels.page_write(&self.native, inputs, options, prepared)
    }
}