rlx_ir/
opcodes.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Canonical integer opcodes for the standalone elementwise kernels — the
6//! **single source of truth** for the small `u32` "op id" every backend hands to
7//! its unary / binary / compare / reduce kernel switch.
8//!
9//! Historically each backend hand-rolled its own `binary_op_id` / `act_id` /
10//! `reduce_id` match. Being copies, they drifted: the forward-activation
11//! numbering on CUDA/wgpu (`Relu=0 … Gelu=9`) diverged from Vulkan/oneAPI
12//! (`Gelu=0 … Relu=3`), and Vulkan even carried two disagreeing activation
13//! tables. Divergent-but-same-shaped opcode tables are exactly the silent
14//! wrong-output failure mode we want to design out.
15//!
16//! These inherent methods pin every id in one place. Each backend wrapper now
17//! delegates here (a value-preserving change — the numbers are unchanged), and
18//! the `#[cfg(test)]` bijection guards below fail the moment a new enum variant
19//! is added without an id, or two variants collide on one id. The exhaustive
20//! `match` means adding an [`Activation`]/[`BinaryOp`]/… variant won't even
21//! compile until it is given an opcode.
22//!
23//! ## Two activation schemes, both pinned here
24//!
25//! There are two real, self-consistent forward-activation numberings, each baked
26//! into a matching set of GPU shader switches:
27//!
28//! * [`Activation::opcode_relu_first`] — CUDA / wgpu / ROCm **forward** unary
29//!   kernels, and **every** backend's activation-*backward* kernel
30//!   (Vulkan/oneAPI/Metal).
31//! * [`Activation::opcode_gelu_first`] — Vulkan and oneAPI **forward** unary
32//!   kernels (`act_id`).
33//!
34//! Collapsing the two into one is a separate, rig-validated change: it requires
35//! editing the paired `.comp` / `.cl` shader switches in lockstep and re-running
36//! the cross-backend parity sweep on real Vulkan/oneAPI hardware. Until then,
37//! defining each scheme exactly once here stops any *further* drift — and the
38//! parity harness (see `rlx-runtime/tests/elementwise_backend_parity.rs`)
39//! detects at runtime if a shader ever disagrees with its id table.
40
41use crate::op::{Activation, BinaryOp, CmpOp, ReduceOp};
42
43impl BinaryOp {
44    /// Every [`BinaryOp`] variant, for exhaustive test/parity fan-out.
45    pub const ALL: [BinaryOp; 14] = [
46        BinaryOp::Add,
47        BinaryOp::Sub,
48        BinaryOp::Mul,
49        BinaryOp::Div,
50        BinaryOp::Max,
51        BinaryOp::Min,
52        BinaryOp::Pow,
53        BinaryOp::Mod,
54        BinaryOp::BitAnd,
55        BinaryOp::BitOr,
56        BinaryOp::BitXor,
57        BinaryOp::Shl,
58        BinaryOp::Shr,
59        BinaryOp::Atan2,
60    ];
61
62    /// Opcode passed to the `binary` kernel on every backend (identical across
63    /// CPU/CUDA/ROCm/Metal/wgpu/Vulkan/oneAPI).
64    pub const fn opcode(self) -> u32 {
65        match self {
66            BinaryOp::Add => 0,
67            BinaryOp::Sub => 1,
68            BinaryOp::Mul => 2,
69            BinaryOp::Div => 3,
70            BinaryOp::Max => 4,
71            BinaryOp::Min => 5,
72            BinaryOp::Pow => 6,
73            BinaryOp::Mod => 7,
74            BinaryOp::BitAnd => 8,
75            BinaryOp::BitOr => 9,
76            BinaryOp::BitXor => 10,
77            BinaryOp::Shl => 11,
78            BinaryOp::Shr => 12,
79            BinaryOp::Atan2 => 13,
80        }
81    }
82}
83
84impl CmpOp {
85    /// Every [`CmpOp`] variant, for exhaustive test/parity fan-out.
86    pub const ALL: [CmpOp; 6] = [
87        CmpOp::Eq,
88        CmpOp::Ne,
89        CmpOp::Lt,
90        CmpOp::Le,
91        CmpOp::Gt,
92        CmpOp::Ge,
93    ];
94
95    /// Opcode passed to the `compare` kernel on every backend.
96    pub const fn opcode(self) -> u32 {
97        match self {
98            CmpOp::Eq => 0,
99            CmpOp::Ne => 1,
100            CmpOp::Lt => 2,
101            CmpOp::Le => 3,
102            CmpOp::Gt => 4,
103            CmpOp::Ge => 5,
104        }
105    }
106}
107
108impl ReduceOp {
109    /// Every [`ReduceOp`] variant, for exhaustive test/parity fan-out.
110    pub const ALL: [ReduceOp; 5] = [
111        ReduceOp::Sum,
112        ReduceOp::Mean,
113        ReduceOp::Max,
114        ReduceOp::Min,
115        ReduceOp::Prod,
116    ];
117
118    /// Opcode passed to the `reduce` kernel (`sum=0, mean=1, max=2, min=3,
119    /// prod=4`).
120    pub const fn opcode(self) -> u32 {
121        match self {
122            ReduceOp::Sum => 0,
123            ReduceOp::Mean => 1,
124            ReduceOp::Max => 2,
125            ReduceOp::Min => 3,
126            ReduceOp::Prod => 4,
127        }
128    }
129
130    /// Opcode passed to the `pool{1,2,3}d` kernels, whose legend differs from
131    /// [`ReduceOp::opcode`]: `max=0, mean=1, sum=2, min=3, prod=4` (Max and Sum
132    /// are swapped). Using the reduce legend here made max-pooling compute the
133    /// window sum — see the CUDA `pool_op_id` note this replaced.
134    pub const fn pool_opcode(self) -> u32 {
135        match self {
136            ReduceOp::Max => 0,
137            ReduceOp::Mean => 1,
138            ReduceOp::Sum => 2,
139            ReduceOp::Min => 3,
140            ReduceOp::Prod => 4,
141        }
142    }
143}
144
145impl Activation {
146    /// Every [`Activation`] variant, for exhaustive test/parity fan-out.
147    pub const ALL: [Activation; 29] = [
148        Activation::Gelu,
149        Activation::GeluApprox,
150        Activation::Silu,
151        Activation::Relu,
152        Activation::Sigmoid,
153        Activation::Tanh,
154        Activation::Exp,
155        Activation::Log,
156        Activation::Sqrt,
157        Activation::Rsqrt,
158        Activation::Neg,
159        Activation::Abs,
160        Activation::Sin,
161        Activation::Cos,
162        Activation::Tan,
163        Activation::Atan,
164        Activation::Recip,
165        Activation::Round,
166        Activation::Floor,
167        Activation::Ceil,
168        Activation::Sign,
169        Activation::Softplus,
170        Activation::Elu,
171        Activation::Erf,
172        Activation::HardSwish,
173        Activation::HardSigmoid,
174        Activation::Mish,
175        Activation::Softsign,
176        Activation::LogSigmoid,
177    ];
178
179    /// "Relu-first" activation opcode: CUDA / wgpu / ROCm **forward** unary
180    /// kernels, and **every** backend's activation-*backward* kernel
181    /// (`activation_backward.{cu,comp,cl}` + Metal). `Relu=0 … Gelu=9 …
182    /// Recip=17 … LogSigmoid=28`.
183    pub const fn opcode_relu_first(self) -> u32 {
184        match self {
185            Activation::Relu => 0,
186            Activation::Sigmoid => 1,
187            Activation::Tanh => 2,
188            Activation::Exp => 3,
189            Activation::Log => 4,
190            Activation::Sqrt => 5,
191            Activation::Rsqrt => 6,
192            Activation::Neg => 7,
193            Activation::Abs => 8,
194            Activation::Gelu => 9,
195            Activation::Silu => 10,
196            Activation::GeluApprox => 11,
197            Activation::Round => 12,
198            Activation::Sin => 13,
199            Activation::Cos => 14,
200            Activation::Tan => 15,
201            Activation::Atan => 16,
202            Activation::Recip => 17,
203            Activation::Floor => 18,
204            Activation::Ceil => 19,
205            Activation::Sign => 20,
206            Activation::Softplus => 21,
207            Activation::Elu => 22,
208            Activation::Erf => 23,
209            Activation::HardSwish => 24,
210            Activation::HardSigmoid => 25,
211            Activation::Mish => 26,
212            Activation::Softsign => 27,
213            Activation::LogSigmoid => 28,
214        }
215    }
216
217    /// "Gelu-first" activation opcode: Vulkan and oneAPI **forward** unary
218    /// kernels (`act_id`). `Gelu=0, GeluApprox=1, Silu=2, Relu=3 …
219    /// LogSigmoid=28`. Note `Round=16, Recip=17` (swapped vs enum order).
220    pub const fn opcode_gelu_first(self) -> u32 {
221        match self {
222            Activation::Gelu => 0,
223            Activation::GeluApprox => 1,
224            Activation::Silu => 2,
225            Activation::Relu => 3,
226            Activation::Sigmoid => 4,
227            Activation::Tanh => 5,
228            Activation::Exp => 6,
229            Activation::Log => 7,
230            Activation::Sqrt => 8,
231            Activation::Rsqrt => 9,
232            Activation::Neg => 10,
233            Activation::Abs => 11,
234            Activation::Sin => 12,
235            Activation::Cos => 13,
236            Activation::Tan => 14,
237            Activation::Atan => 15,
238            Activation::Round => 16,
239            Activation::Recip => 17,
240            Activation::Floor => 18,
241            Activation::Ceil => 19,
242            Activation::Sign => 20,
243            Activation::Softplus => 21,
244            Activation::Elu => 22,
245            Activation::Erf => 23,
246            Activation::HardSwish => 24,
247            Activation::HardSigmoid => 25,
248            Activation::Mish => 26,
249            Activation::Softsign => 27,
250            Activation::LogSigmoid => 28,
251        }
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258
259    /// A scheme is valid iff its ids are a bijection onto `0..N` — no gaps, no
260    /// collisions. This is what fails if someone adds a variant and reuses an id
261    /// or leaves a hole.
262    fn assert_bijection(name: &str, ids: &[u32]) {
263        let n = ids.len();
264        let mut sorted = ids.to_vec();
265        sorted.sort_unstable();
266        for (i, &id) in sorted.iter().enumerate() {
267            assert_eq!(
268                id as usize, i,
269                "{name}: opcodes must be a contiguous bijection onto 0..{n}; \
270                 got {sorted:?} (duplicate or gap at index {i})"
271            );
272        }
273    }
274
275    #[test]
276    fn binary_opcodes_are_a_bijection() {
277        let ids: Vec<u32> = BinaryOp::ALL.iter().map(|&o| o.opcode()).collect();
278        assert_bijection("BinaryOp::opcode", &ids);
279    }
280
281    #[test]
282    fn compare_opcodes_are_a_bijection() {
283        let ids: Vec<u32> = CmpOp::ALL.iter().map(|&o| o.opcode()).collect();
284        assert_bijection("CmpOp::opcode", &ids);
285    }
286
287    #[test]
288    fn reduce_opcodes_are_a_bijection() {
289        let ids: Vec<u32> = ReduceOp::ALL.iter().map(|&o| o.opcode()).collect();
290        assert_bijection("ReduceOp::opcode", &ids);
291        let pool: Vec<u32> = ReduceOp::ALL.iter().map(|&o| o.pool_opcode()).collect();
292        assert_bijection("ReduceOp::pool_opcode", &pool);
293    }
294
295    #[test]
296    fn activation_schemes_are_bijections() {
297        let relu: Vec<u32> = Activation::ALL
298            .iter()
299            .map(|&a| a.opcode_relu_first())
300            .collect();
301        assert_bijection("Activation::opcode_relu_first", &relu);
302        let gelu: Vec<u32> = Activation::ALL
303            .iter()
304            .map(|&a| a.opcode_gelu_first())
305            .collect();
306        assert_bijection("Activation::opcode_gelu_first", &gelu);
307    }
308
309    /// `ALL` must list every variant exactly once (guards against forgetting to
310    /// extend `ALL` when a variant is added — the bijection would still pass on
311    /// a short list, but the count check here won't).
312    #[test]
313    fn all_arrays_have_expected_lengths() {
314        assert_eq!(BinaryOp::ALL.len(), 14);
315        assert_eq!(CmpOp::ALL.len(), 6);
316        assert_eq!(ReduceOp::ALL.len(), 5);
317        assert_eq!(Activation::ALL.len(), 29);
318    }
319
320    /// Pin the *exact* opcode of every variant to the value its GPU shader
321    /// switch expects. A bijection check alone would still pass if two ids were
322    /// swapped — but a swap silently mismatches the paired `.cu`/`.wgsl`/`.comp`
323    /// switch. This is the lock: if the canonical table above is edited, this
324    /// fails, forcing whoever changes it to also update the shaders (and the
325    /// cross-backend parity sweep to re-confirm on hardware).
326    #[test]
327    fn canonical_ids_are_pinned() {
328        // Binary — identical on every backend.
329        for (op, want) in [
330            (BinaryOp::Add, 0),
331            (BinaryOp::Div, 3),
332            (BinaryOp::Pow, 6),
333            (BinaryOp::Mod, 7),
334            (BinaryOp::Shr, 12),
335            (BinaryOp::Atan2, 13),
336        ] {
337            assert_eq!(op.opcode(), want, "BinaryOp::{op:?}");
338        }
339        // Compare / reduce.
340        assert_eq!(CmpOp::Eq.opcode(), 0);
341        assert_eq!(CmpOp::Ge.opcode(), 5);
342        assert_eq!(ReduceOp::Sum.opcode(), 0);
343        assert_eq!(ReduceOp::Prod.opcode(), 4);
344        // Pool legend swaps Max/Sum vs the reduce legend.
345        assert_eq!(ReduceOp::Max.opcode(), 2);
346        assert_eq!(ReduceOp::Max.pool_opcode(), 0);
347        assert_eq!(ReduceOp::Sum.pool_opcode(), 2);
348        // Relu-first (CUDA/wgpu/ROCm forward + all backward kernels).
349        assert_eq!(Activation::Relu.opcode_relu_first(), 0);
350        assert_eq!(Activation::Gelu.opcode_relu_first(), 9);
351        assert_eq!(Activation::Recip.opcode_relu_first(), 17);
352        assert_eq!(Activation::LogSigmoid.opcode_relu_first(), 28);
353        // Gelu-first (Vulkan/oneAPI forward) — note Round=16, Recip=17.
354        assert_eq!(Activation::Gelu.opcode_gelu_first(), 0);
355        assert_eq!(Activation::Relu.opcode_gelu_first(), 3);
356        assert_eq!(Activation::Atan.opcode_gelu_first(), 15);
357        assert_eq!(Activation::Round.opcode_gelu_first(), 16);
358        assert_eq!(Activation::Recip.opcode_gelu_first(), 17);
359        assert_eq!(Activation::LogSigmoid.opcode_gelu_first(), 28);
360    }
361}
rlx_ir/opcodes.rs

rlx_ir/
opcodes.rs
