rlx-metal 0.2.2

Metal backend for RLX — Apple Silicon GPU via Metal Performance Shaders + custom MSL 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

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.
// RLX — LLaDA2 group-limited gate on Metal (host delegate, unified memory).

use crate::op_registry::{MetalKernel, register_metal_kernel};
use rlx_ir::Shape;
use std::sync::Arc;

pub const OP_NAME: &str = "llada2.group_limited_gate";

#[derive(Debug)]
struct Llada2GateMetal;

impl MetalKernel for Llada2GateMetal {
    fn name(&self) -> &str {
        OP_NAME
    }

    fn execute(
        &self,
        inputs: &[(&[u8], &Shape)],
        output: (&mut [u8], &Shape),
        attrs: &[u8],
    ) -> Result<(), String> {
        let sig_bytes = inputs[0].0;
        let route_bytes = inputs[1].0;
        let out_bytes = output.0;
        if !sig_bytes.len().is_multiple_of(4)
            || !route_bytes.len().is_multiple_of(4)
            || !out_bytes.len().is_multiple_of(4)
        {
            return Err("gate: non-f32-aligned buffers".into());
        }
        let sig = bytemuck::cast_slice::<u8, f32>(sig_bytes);
        let route = bytemuck::cast_slice::<u8, f32>(route_bytes);
        let out = bytemuck::cast_slice_mut::<u8, f32>(out_bytes);
        rlx_cpu::llada2_gate::execute_gate_f32(sig, route, out, attrs)
    }
}

pub fn register() {
    register_metal_kernel(Arc::new(Llada2GateMetal));
}