mlx-native 0.8.1

Pure-Rust Metal GPU compute library for MLX-compatible inference on Apple Silicon
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
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
144

//! Per-row sum reduction along the last dimension of a 2-D tensor +
//! its broadcast-along-cols backward.
//!
//! Used by reverse-mode autograd in downstream crates (hf2q ADR-020
//! Track 1: KL-divergence loss composition needs Σ_j p · (log_p − log_q)
//! per row).

use metal::MTLSize;

use crate::buffer::MlxBuffer;
use crate::dtypes::DType;
use crate::encoder::CommandEncoder;
use crate::error::{MlxError, Result};
use crate::kernel_registry::KernelRegistry;

pub static ROW_SUM_SHADER_SOURCE: &str = include_str!("../shaders/row_sum.metal");

pub fn register(registry: &mut KernelRegistry) {
    registry.register_source("row_sum_f32", ROW_SUM_SHADER_SOURCE);
    registry.register_source("row_sum_backward_f32", ROW_SUM_SHADER_SOURCE);
}

/// Encode `output[b] = Σ_j input[b, j]` for a 2-D `[rows, cols]` f32 input.
#[allow(clippy::too_many_arguments)]
pub fn dispatch_row_sum_f32(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    input: &MlxBuffer,
    output: &MlxBuffer,
    params_buf: &MlxBuffer,
    rows: u32,
    cols: u32,
) -> Result<()> {
    if rows == 0 || cols == 0 {
        return Err(MlxError::InvalidArgument(
            "row_sum_f32: rows and cols must be > 0".into(),
        ));
    }
    let in_expected = (rows as usize) * (cols as usize);
    if input.element_count() != in_expected {
        return Err(MlxError::InvalidArgument(format!(
            "row_sum_f32: input element count {} != rows({}) * cols({})",
            input.element_count(),
            rows,
            cols
        )));
    }
    if output.element_count() != rows as usize {
        return Err(MlxError::InvalidArgument(format!(
            "row_sum_f32: output element count {} != rows({})",
            output.element_count(),
            rows
        )));
    }
    if input.dtype() != DType::F32 || output.dtype() != DType::F32 {
        return Err(MlxError::InvalidArgument(format!(
            "row_sum_f32: only f32 supported; got input={} output={}",
            input.dtype(),
            output.dtype()
        )));
    }
    if params_buf.byte_len() < 8 {
        return Err(MlxError::InvalidArgument(format!(
            "row_sum_f32: params_buf too small (need 8 bytes, got {})",
            params_buf.byte_len()
        )));
    }

    let pipeline = registry.get_pipeline("row_sum_f32", device)?;
    let tg_size = std::cmp::min(256, cols.next_power_of_two()) as u64;
    let shared_mem_bytes = tg_size * 4;

    encoder.encode_threadgroups_with_shared(
        pipeline,
        &[(0, input), (1, output), (2, params_buf)],
        &[(0, shared_mem_bytes)],
        MTLSize::new(rows as u64, 1, 1),
        MTLSize::new(tg_size, 1, 1),
    );

    Ok(())
}

/// Encode `dx[b, i] = d_out[b]` (broadcast along the cols dim).  This
/// is the backward of [`dispatch_row_sum_f32`].
#[allow(clippy::too_many_arguments)]
pub fn dispatch_row_sum_backward_f32(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &metal::DeviceRef,
    d_out: &MlxBuffer,
    dx: &MlxBuffer,
    params_buf: &MlxBuffer,
    rows: u32,
    cols: u32,
) -> Result<()> {
    if rows == 0 || cols == 0 {
        return Err(MlxError::InvalidArgument(
            "row_sum_backward_f32: rows and cols must be > 0".into(),
        ));
    }
    if d_out.element_count() != rows as usize {
        return Err(MlxError::InvalidArgument(format!(
            "row_sum_backward_f32: d_out element count {} != rows({})",
            d_out.element_count(),
            rows
        )));
    }
    let dx_expected = (rows as usize) * (cols as usize);
    if dx.element_count() != dx_expected {
        return Err(MlxError::InvalidArgument(format!(
            "row_sum_backward_f32: dx element count {} != rows({}) * cols({})",
            dx.element_count(),
            rows,
            cols
        )));
    }
    if d_out.dtype() != DType::F32 || dx.dtype() != DType::F32 {
        return Err(MlxError::InvalidArgument(format!(
            "row_sum_backward_f32: only f32; d_out={} dx={}",
            d_out.dtype(),
            dx.dtype()
        )));
    }
    if params_buf.byte_len() < 8 {
        return Err(MlxError::InvalidArgument(format!(
            "row_sum_backward_f32: params_buf too small (need 8 bytes, got {})",
            params_buf.byte_len()
        )));
    }

    let pipeline = registry.get_pipeline("row_sum_backward_f32", device)?;
    let tg_size = std::cmp::min(256, cols.next_power_of_two()) as u64;

    encoder.encode_threadgroups(
        pipeline,
        &[(0, d_out), (1, dx), (2, params_buf)],
        MTLSize::new(rows as u64, 1, 1),
        MTLSize::new(tg_size, 1, 1),
    );

    Ok(())
}