libmir-metal 0.3.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

use super::super::{Array, Result, Stream};

#[derive(Debug)]
pub(in crate::engine) struct GptqLinear {
    weight: Array,
    zero_points: Array,
    scales: Array,
    group_indices: Array,
    input: usize,
    output: usize,
    group_size: usize,
    legacy: bool,
}

impl GptqLinear {
    pub(in crate::engine) fn new(
        arrays: [Array; 4],
        input: usize,
        output: usize,
        group_size: usize,
        legacy: bool,
    ) -> Self {
        let [weight, zero_points, scales, group_indices] = arrays;
        Self {
            weight,
            zero_points,
            scales,
            group_indices,
            input,
            output,
            group_size,
            legacy,
        }
    }

    pub(in crate::engine) fn forward(&self, input: &Array, stream: &Stream) -> Result<Array> {
        stream.kernels().gptq_linear(
            stream,
            [input, &self.weight, &self.zero_points, &self.scales, &self.group_indices],
            self.input,
            self.output,
            self.group_size,
            self.legacy,
        )
    }

    pub(in crate::engine) const fn group_size(&self) -> usize {
        self.group_size
    }
}