rlx-metal 0.2.6

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
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

// 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/>.

// `objc` crate's `class!` / `msg_send!` macros expand to
// `cfg(feature = "cargo-clippy")` checks that aren't recognized by
// modern rustc. The warnings are third-party noise (~78 across this
// crate); they say nothing about our code. Silence at the crate root.
#![allow(unexpected_cfgs)]

//! RLX Metal backend — Apple Silicon GPU execution.
//!
//! Compiles RLX IR graphs to Metal compute pipelines + MPS matrix kernels.
//!
//! Architecture mirrors rlx-cpu:
//! - `device` — Metal device discovery and properties
//! - `arena`  — GPU buffer allocation from memory plan
//! - `blas`   — MPS matrix multiplication (analog of cblas_sgemm)
//! - `kernels`— custom MSL compute shaders (analog of NEON kernels)
//! - `thunk`  — pre-compiled command buffer with arena offsets
//! - `backend`— ExecutableGraph implementation
//!
//! Apple Silicon advantages:
//! - Unified memory: zero-copy CPU↔GPU
//! - 16-core GPU on M4 Pro: ~1.4 TFLOP/s peak
//! - 273 GB/s memory bandwidth (vs 120 on CPU)
//! - MPSMatrixMultiplication uses dedicated matmul hardware

#[cfg(target_os = "macos")]
pub mod device;

#[cfg(target_os = "macos")]
pub mod arena;

#[cfg(target_os = "macos")]
pub mod blas;

#[cfg(target_os = "macos")]
pub mod mps_blas;

#[cfg(target_os = "macos")]
pub mod mps_graph;

#[cfg(target_os = "macos")]
pub mod mps_graph_hybrid;
#[cfg(target_os = "macos")]
pub mod mps_graph_lower;

#[cfg(target_os = "macos")]
pub mod icb;

#[cfg(target_os = "macos")]
pub mod kernels;

#[cfg(target_os = "macos")]
pub mod fft_dispatch;

#[cfg(target_os = "macos")]
pub mod llada2_gate;

#[cfg(target_os = "macos")]
pub mod cost;

#[cfg(target_os = "macos")]
pub mod calibrate;

#[cfg(target_os = "macos")]
pub mod thunk;

#[cfg(target_os = "macos")]
pub mod backend;

#[cfg(target_os = "macos")]
pub mod attention_bwd_gpu;

#[cfg(target_os = "macos")]
pub mod thunk_profile;

#[cfg(target_os = "macos")]
pub mod mps_profile;

#[cfg(all(feature = "native-splat", target_os = "macos"))]
pub mod splat_adam;
#[cfg(all(feature = "native-splat", target_os = "macos"))]
pub mod splat_native;
#[cfg(all(feature = "native-splat", target_os = "macos"))]
pub mod splat_training;
#[cfg(all(feature = "native-splat", target_os = "macos"))]
pub mod splat_training_pipeline;

#[cfg(target_os = "macos")]
pub mod async_copy;

#[cfg(target_os = "macos")]
pub mod op_registry;

/// PLAN: Schedule splitting for the Metal MPSGraph path. Splits the
/// schedule at attention boundaries so the broken slice-of-computed
/// MPSGraph attention pattern is replaced by the parity-correct
/// thunk path; everything else still gets the MPSGraph dispatch-
/// overhead reduction. Scaffolding only today (data model +
/// segmenter + 3 unit tests); executor wiring + per-segment plan
/// compilation is the next chunk.
pub mod segmented;

/// Stub when not on macOS — Metal is only available on Apple platforms.
#[cfg(not(target_os = "macos"))]
pub fn is_available() -> bool {
    false
}

#[cfg(target_os = "macos")]
pub fn is_available() -> bool {
    device::has_metal_device()
}