mirtal 0.1.0

Native, explicit-stream Rust API for Apple Metal
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

use crate::{DType, Shape};

#[derive(Debug, Clone, PartialEq, Eq)]
/// Shape and dtype requested for one Metal kernel output.
pub struct OutputSpec {
    /// Output array shape.
    pub shape: Shape,
    /// Output array dtype.
    pub dtype: DType,
}

impl OutputSpec {
    #[must_use]
    /// Creates an output specification.
    pub const fn new(shape: Shape, dtype: DType) -> Self {
        Self { shape, dtype }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
/// A named Metal function-template argument.
pub struct TemplateArg {
    /// Template parameter name.
    pub name: &'static str,
    /// Value supplied for the parameter.
    pub value: TemplateValue,
}

impl TemplateArg {
    #[must_use]
    /// Creates a signed-integer template argument.
    pub const fn int(name: &'static str, value: i32) -> Self {
        Self { name, value: TemplateValue::Int(value) }
    }

    #[must_use]
    /// Creates a Boolean template argument.
    pub const fn bool(name: &'static str, value: bool) -> Self {
        Self { name, value: TemplateValue::Bool(value) }
    }

    #[must_use]
    /// Creates a dtype template argument.
    pub const fn dtype(name: &'static str, value: DType) -> Self {
        Self { name, value: TemplateValue::DType(value) }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// A value accepted by a Metal function-template parameter.
pub enum TemplateValue {
    /// Signed integer template value.
    Int(i32),
    /// Boolean template value.
    Bool(bool),
    /// Dtype template value.
    DType(DType),
}

#[derive(Debug, Clone)]
/// Geometry and template values for one Metal kernel launch.
pub struct Dispatch {
    /// Number of threads in the three-dimensional launch grid.
    pub grid: [usize; 3],
    /// Threads per threadgroup in each dimension.
    pub threadgroup: [usize; 3],
    /// Named function-template arguments.
    pub templates: Vec<TemplateArg>,
    /// Optional initial value used to allocate output buffers.
    pub init_value: Option<f32>,
    /// Enables verbose MLX Metal diagnostics.
    pub verbose: bool,
}

impl Dispatch {
    #[must_use]
    /// Creates a launch with no template arguments or output initializer.
    pub const fn new(grid: [usize; 3], threadgroup: [usize; 3]) -> Self {
        Self {
            grid,
            threadgroup,
            templates: Vec::new(),
            init_value: None,
            verbose: false,
        }
    }

    #[must_use]
    /// Appends function-template arguments to the launch.
    pub fn templates(mut self, templates: impl IntoIterator<Item = TemplateArg>) -> Self {
        self.templates.extend(templates);
        self
    }

    #[must_use]
    /// Sets the value used to initialize output buffers.
    pub const fn init_value(mut self, value: f32) -> Self {
        self.init_value = Some(value);
        self
    }

    #[must_use]
    /// Enables or disables verbose MLX Metal diagnostics.
    pub const fn verbose(mut self, verbose: bool) -> Self {
        self.verbose = verbose;
        self
    }
}