optirs_gpu/shaders/
msl.rs

1//! Metal Shading Language sources for the optimizer kernels.
2//!
3//! These mirror [`super::wgsl`] line for line in semantics. They exist because
4//! `scirs2-core` 0.6.5 registers its optimizer kernels with an *empty*
5//! `metal_source`, and because the WebGPU backend's runtime device detection
6//! does not yet enumerate wgpu adapters — on macOS the Metal backend is the
7//! path that actually reaches the GPU.
8//!
9//! Conventions (see [`super`]): buffers are bound to argument-table indices in
10//! the order `x, y, a, b, result, output`, every scalar travels in a buffer,
11//! and integer scalars are recovered with `as_type<uint>`.
12//!
13//! scirs2-core's Metal compiler extracts the entry point with
14//! `source.find("kernel void ")` up to the next `(`, so the name and the
15//! opening parenthesis must stay on one line.
16
17/// Adam with coupled L2 weight decay. See [`super::wgsl::ADAM`] for bindings.
18pub const ADAM: &str = r#"
19#include <metal_stdlib>
20using namespace metal;
21
22kernel void optirs_adam(
23    device float* x [[buffer(0)]],
24    device const float* y [[buffer(1)]],
25    device float* a [[buffer(2)]],
26    device float* b [[buffer(3)]],
27    device const float* result [[buffer(4)]],
28    uint idx [[thread_position_in_grid]])
29{
30    uint n = as_type<uint>(result[7]);
31    if (idx >= n) { return; }
32
33    float lr = result[0];
34    float beta1 = result[1];
35    float beta2 = result[2];
36    float eps = result[3];
37    float wd = result[4];
38    float bc1 = result[5];
39    float bc2 = result[6];
40
41    float p = x[idx];
42    float g = y[idx];
43    if (wd > 0.0f) { g = g + wd * p; }
44
45    float mi = beta1 * a[idx] + (1.0f - beta1) * g;
46    float vi = beta2 * b[idx] + (1.0f - beta2) * g * g;
47    a[idx] = mi;
48    b[idx] = vi;
49
50    float m_hat = mi / bc1;
51    float v_hat = vi / bc2;
52    x[idx] = p - lr * m_hat / (sqrt(v_hat) + eps);
53}
54"#;
55
56/// AdamW with decoupled weight decay. See [`super::wgsl::ADAMW`].
57pub const ADAMW: &str = r#"
58#include <metal_stdlib>
59using namespace metal;
60
61kernel void optirs_adamw(
62    device float* x [[buffer(0)]],
63    device const float* y [[buffer(1)]],
64    device float* a [[buffer(2)]],
65    device float* b [[buffer(3)]],
66    device const float* result [[buffer(4)]],
67    uint idx [[thread_position_in_grid]])
68{
69    uint n = as_type<uint>(result[7]);
70    if (idx >= n) { return; }
71
72    float lr = result[0];
73    float beta1 = result[1];
74    float beta2 = result[2];
75    float eps = result[3];
76    float wd = result[4];
77    float bc1 = result[5];
78    float bc2 = result[6];
79
80    float p = x[idx];
81    float g = y[idx];
82
83    float mi = beta1 * a[idx] + (1.0f - beta1) * g;
84    float vi = beta2 * b[idx] + (1.0f - beta2) * g * g;
85    a[idx] = mi;
86    b[idx] = vi;
87
88    float m_hat = mi / bc1;
89    float v_hat = vi / bc2;
90
91    float decayed = p - lr * wd * p;
92    x[idx] = decayed - lr * m_hat / (sqrt(v_hat) + eps);
93}
94"#;
95
96/// SGD with momentum / dampening / Nesterov. See [`super::wgsl::SGD`].
97pub const SGD: &str = r#"
98#include <metal_stdlib>
99using namespace metal;
100
101kernel void optirs_sgd(
102    device float* x [[buffer(0)]],
103    device const float* y [[buffer(1)]],
104    device float* a [[buffer(2)]],
105    device const float* b [[buffer(3)]],
106    uint idx [[thread_position_in_grid]])
107{
108    uint n = as_type<uint>(b[6]);
109    if (idx >= n) { return; }
110
111    float lr = b[0];
112    float momentum = b[1];
113    float dampening = b[2];
114    float wd = b[3];
115    float nesterov = b[4];
116    float first = b[5];
117
118    float p = x[idx];
119    float g = y[idx];
120    if (wd > 0.0f) { g = g + wd * p; }
121
122    if (momentum > 0.0f) {
123        float buf = g;
124        if (first < 0.5f) { buf = momentum * a[idx] + (1.0f - dampening) * g; }
125        a[idx] = buf;
126        if (nesterov > 0.5f) { g = g + momentum * buf; } else { g = buf; }
127    }
128
129    x[idx] = p - lr * g;
130}
131"#;
132
133/// RMSprop with centering and momentum. See [`super::wgsl::RMSPROP`].
134pub const RMSPROP: &str = r#"
135#include <metal_stdlib>
136using namespace metal;
137
138kernel void optirs_rmsprop(
139    device float* x [[buffer(0)]],
140    device const float* y [[buffer(1)]],
141    device float* a [[buffer(2)]],
142    device float* b [[buffer(3)]],
143    device float* result [[buffer(4)]],
144    device const float* output [[buffer(5)]],
145    uint idx [[thread_position_in_grid]])
146{
147    uint n = as_type<uint>(output[6]);
148    if (idx >= n) { return; }
149
150    float lr = output[0];
151    float alpha = output[1];
152    float eps = output[2];
153    float wd = output[3];
154    float momentum = output[4];
155    float centered = output[5];
156
157    float p = x[idx];
158    float g = y[idx];
159    if (wd > 0.0f) { g = g + wd * p; }
160
161    float sq = alpha * a[idx] + (1.0f - alpha) * g * g;
162    a[idx] = sq;
163
164    float avg = sq;
165    if (centered > 0.5f) {
166        float ga = alpha * b[idx] + (1.0f - alpha) * g;
167        b[idx] = ga;
168        avg = sq - ga * ga;
169    }
170
171    float denom = sqrt(max(avg, 0.0f)) + eps;
172
173    if (momentum > 0.0f) {
174        float buf = momentum * result[idx] + g / denom;
175        result[idx] = buf;
176        x[idx] = p - lr * buf;
177    } else {
178        x[idx] = p - lr * g / denom;
179    }
180}
181"#;
182
183/// Adagrad with learning-rate decay. See [`super::wgsl::ADAGRAD`].
184pub const ADAGRAD: &str = r#"
185#include <metal_stdlib>
186using namespace metal;
187
188kernel void optirs_adagrad(
189    device float* x [[buffer(0)]],
190    device const float* y [[buffer(1)]],
191    device float* a [[buffer(2)]],
192    device const float* b [[buffer(3)]],
193    uint idx [[thread_position_in_grid]])
194{
195    uint n = as_type<uint>(b[3]);
196    if (idx >= n) { return; }
197
198    float clr = b[0];
199    float eps = b[1];
200    float wd = b[2];
201
202    float p = x[idx];
203    float g = y[idx];
204    if (wd > 0.0f) { g = g + wd * p; }
205
206    float s = a[idx] + g * g;
207    a[idx] = s;
208    x[idx] = p - clr * g / (sqrt(s) + eps);
209}
210"#;
211
212/// Local reduction step for multi-GPU all-reduce-mean. See
213/// [`super::wgsl::ALL_REDUCE_MEAN`].
214pub const ALL_REDUCE_MEAN: &str = r#"
215#include <metal_stdlib>
216using namespace metal;
217
218kernel void optirs_all_reduce_mean(
219    device float* x [[buffer(0)]],
220    device const float* y [[buffer(1)]],
221    uint idx [[thread_position_in_grid]])
222{
223    uint n = as_type<uint>(y[0]);
224    if (idx >= n) { return; }
225
226    uint num_gpus = as_type<uint>(y[1]);
227    x[idx] = x[idx] / float(num_gpus);
228}
229"#;
230
231/// Two-phase LAMB with a threadgroup norm reduction. See [`super::wgsl::LAMB`].
232pub const LAMB: &str = r#"
233#include <metal_stdlib>
234using namespace metal;
235
236kernel void optirs_lamb(
237    device float* x [[buffer(0)]],
238    device const float* y [[buffer(1)]],
239    device float* a [[buffer(2)]],
240    device float* b [[buffer(3)]],
241    device float* result [[buffer(4)]],
242    device const float* output [[buffer(5)]],
243    uint idx [[thread_position_in_grid]],
244    uint lid [[thread_position_in_threadgroup]],
245    uint wgid [[threadgroup_position_in_grid]])
246{
247    threadgroup float scratch_p[256];
248    threadgroup float scratch_u[256];
249
250    uint n = as_type<uint>(output[7]);
251    uint phase = as_type<uint>(output[8]);
252
253    float lr = output[0];
254    float beta1 = output[1];
255    float beta2 = output[2];
256    float eps = output[3];
257    float wd = output[4];
258    float bc1 = output[5];
259    float bc2 = output[6];
260    float trust = output[9];
261
262    float sum_p = 0.0f;
263    float sum_u = 0.0f;
264
265    if (idx < n) {
266        if (phase == 0u) {
267            float p = x[idx];
268            float g = y[idx];
269            float mi = beta1 * a[idx] + (1.0f - beta1) * g;
270            float vi = beta2 * b[idx] + (1.0f - beta2) * g * g;
271            a[idx] = mi;
272            b[idx] = vi;
273            float m_hat = mi / bc1;
274            float v_hat = vi / bc2;
275            result[idx] = m_hat / (sqrt(v_hat) + eps) + wd * p;
276        } else {
277            x[idx] = x[idx] - lr * trust * result[idx];
278        }
279        float pv = x[idx];
280        float uv = result[idx];
281        sum_p = pv * pv;
282        sum_u = uv * uv;
283    }
284
285    scratch_p[lid] = sum_p;
286    scratch_u[lid] = sum_u;
287    threadgroup_barrier(mem_flags::mem_threadgroup);
288
289    for (uint stride = 128u; stride > 0u; stride >>= 1u) {
290        if (lid < stride) {
291            scratch_p[lid] = scratch_p[lid] + scratch_p[lid + stride];
292            scratch_u[lid] = scratch_u[lid] + scratch_u[lid + stride];
293        }
294        threadgroup_barrier(mem_flags::mem_threadgroup);
295    }
296
297    if (lid == 0u) {
298        result[n + wgid * 2u] = scratch_p[0];
299        result[n + wgid * 2u + 1u] = scratch_u[0];
300    }
301}
302"#;
optirs_gpu/shaders/msl.rs

optirs_gpu/shaders/
msl.rs
