rlx_runtime/backend/
wgpu_backend.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4use super::*;
5use rlx_ir::OpKind;
6use rlx_wgpu::backend::WgpuExecutable;
7
8pub struct WgpuBackend;
9
10impl Backend for WgpuBackend {
11    fn supported_ops(&self) -> &'static [OpKind] {
12        rlx_wgpu::SUPPORTED_OPS
13    }
14
15    fn compile(&self, graph: Graph, options: &CompileOptions) -> Box<dyn ExecutableGraph> {
16        use rlx_opt::pass::Pass as _;
17        let graph = rlx_opt::LowerControlFlow.run(graph);
18        let graph = rlx_opt::legalize_or_rewrite_for_backend(graph, rlx_wgpu::SUPPORTED_OPS)
19            .unwrap_or_else(|errors| {
20                panic!("{}", rlx_opt::format_legalize_error("wgpu", &errors));
21            });
22        let graph = apply_scan_device_preference(graph, options);
23        let graph = crate::precompile::precompile_cleanup(graph, options);
24        // Materialize mid-axis broadcasts before MarkElementwiseRegions:
25        // wgpu Binary/region kernels only handle trailing/scalar broadcast
26        // via modulus; EEG patch embed uses [1,C,1,D] + [1,C,P,D].
27        let graph = rlx_opt::LegalizeBroadcast.run(graph);
28        // ORDER MATTERS: targeted-pattern fusions run BEFORE the
29        // catch-all `MarkElementwiseRegions`. Otherwise the region
30        // pass swallows the Add / Activation nodes into chains and
31        // FuseMatMulBiasAct / FuseResidualLN fail to match the
32        // narrower patterns they look for. (Metal pipeline at line
33        // ~377 already orders these correctly; wgpu was inverted
34        // and silently shipped 13 unfused LayerNorms per BERT
35        // forward where 12 should have been FusedResidualLN.)
36        let compile_result = crate::stages::compile_graph_stages_for_backend(
37            rlx_driver::Device::Gpu,
38            graph,
39            options,
40            rlx_wgpu::SUPPORTED_OPS,
41        );
42        crate::stages::maybe_log_fusion(&compile_result.fusion);
43        let graph = compile_result.lir.into_graph();
44        let graph = match options.policy.clone() {
45            Some(p) => rlx_opt::AutoMixedPrecision::new(p).run(graph),
46            None => graph,
47        };
48        let (graph, io_manifest) = cpu_low_precision::prepare_f32_exec_graph(graph);
49        Box::new(WgpuExecutableWrapper {
50            inner: WgpuExecutable::compile_rng(graph, options.rng),
51            io_manifest,
52        })
53    }
54
55    fn compile_lir(&self, lir: LirModule, options: &CompileOptions) -> Box<dyn ExecutableGraph> {
56        use rlx_opt::pass::Pass as _;
57        // LIR may already contain fused ElementwiseRegions; legalize
58        // broadcasts on the unfused graph shape before backend prep.
59        let graph = rlx_opt::LegalizeBroadcast.run(lir.into_graph());
60        let graph = prepare_fused_graph(graph, options, rlx_wgpu::SUPPORTED_OPS, "wgpu");
61        let (graph, io_manifest) = cpu_low_precision::prepare_f32_exec_graph(graph);
62        Box::new(WgpuExecutableWrapper {
63            inner: WgpuExecutable::compile_rng(graph, options.rng),
64            io_manifest,
65        })
66    }
67}
68
69struct WgpuExecutableWrapper {
70    inner: WgpuExecutable,
71    io_manifest: cpu_low_precision::IoDtypeManifest,
72}
73
74unsafe impl Send for WgpuExecutableWrapper {}
75
76impl ExecutableGraph for WgpuExecutableWrapper {
77    fn capabilities(&self) -> crate::ExecutableCapabilities {
78        crate::ExecutableCapabilities {
79            clone: true,
80            gpu_handles: true,
81            typed_io: true,
82            active_extent: true,
83            ..crate::ExecutableCapabilities::NONE
84        }
85    }
86
87    fn set_param(&mut self, name: &str, data: &[f32]) {
88        self.inner.set_param(name, data);
89    }
90    fn run(&mut self, inputs: &[(&str, &[f32])]) -> Vec<Vec<f32>> {
91        self.inner.run(inputs)
92    }
93    fn run_read_outputs(
94        &mut self,
95        inputs: &[(&str, &[f32])],
96        read_indices: Option<&[usize]>,
97    ) -> Vec<Vec<f32>> {
98        self.inner.run_read_outputs(inputs, read_indices)
99    }
100    fn bind_gpu_handle(&mut self, name: &str, data: &[f32]) -> bool {
101        self.inner.bind_gpu_handle(name, data)
102    }
103    fn has_gpu_handle(&self, name: &str) -> bool {
104        self.inner.has_gpu_handle(name)
105    }
106    fn set_gpu_handle_feed(&mut self, handle_name: &str, output_index: usize) -> bool {
107        self.inner.set_gpu_handle_feed(handle_name, output_index);
108        true
109    }
110    fn read_gpu_handle(&self, name: &str) -> Option<Vec<f32>> {
111        self.inner.read_gpu_handle(name)
112    }
113    fn prepare_resident_gpu_handle(&mut self, name: &str) -> bool {
114        self.inner.prepare_resident_gpu_handle(name)
115    }
116    fn set_active_extent(&mut self, extent: Option<(usize, usize)>) {
117        self.inner.set_active_extent(extent);
118    }
119
120    fn set_rng(&mut self, rng: rlx_ir::RngOptions) {
121        self.inner.set_rng(rng);
122    }
123
124    fn rng(&self) -> rlx_ir::RngOptions {
125        self.inner.rng()
126    }
127
128    /// Typed param upload: the wgpu arena is f32-uniform, so widen F16/BF16 and
129    /// small-integer/bool params (I64/I32/U32/Bool — e.g. `input_ids`, duration
130    /// carry, alignment frame count) to F32 at the host boundary. U8/I8 packed
131    /// quant weights keep their raw bytes.
132    fn set_param_typed(&mut self, name: &str, data: &[u8], dtype: rlx_ir::DType) {
133        match dtype {
134            rlx_ir::DType::U8 | rlx_ir::DType::I8 => {
135                self.inner.set_param_bytes(name, data);
136            }
137            rlx_ir::DType::F32 => {
138                let n = data.len() / 4;
139                let f32_slice =
140                    unsafe { std::slice::from_raw_parts(data.as_ptr() as *const f32, n) };
141                self.inner.set_param(name, f32_slice);
142            }
143            rlx_ir::DType::BF16 => {
144                // A BF16 matmul weight is kept PACKED (2 B/elem) in the wgpu
145                // arena's bf16 side-buffer and unpacked in-shader — no host
146                // widen, no ne*4 arena write.
147                if !self.inner.set_param_bf16_packed(name, data) {
148                    // Every OTHER BF16 param lands in the main arena, whose slot
149                    // `plan_f32_uniform` sized NATIVE (2 B/elem — `F32Uniform`
150                    // widens activations but keeps non-F32 *params* packed).
151                    // Widening here wrote ne*4 bytes into an ne*2 slot: it
152                    // overran the following param and disagreed with the host
153                    // steps that read the slot back as bf16 (e.g.
154                    // `Step::DequantGroupedMatmulMlxHost { scale_bf16 }`, whose
155                    // flag comes from the IR dtype). MXFP4 MoE scales are the
156                    // first non-matmul BF16 param to hit this. Write raw, like
157                    // the U8 arm above.
158                    self.inner.set_param_bytes(name, data);
159                }
160            }
161            rlx_ir::DType::F16
162            | rlx_ir::DType::F64
163            | rlx_ir::DType::I64
164            | rlx_ir::DType::I32
165            | rlx_ir::DType::U32
166            | rlx_ir::DType::Bool => {
167                let f32 = super::widen_bytes_to_f32(data, dtype);
168                self.inner.set_param(name, &f32);
169            }
170            other => panic!(
171                "rlx-wgpu set_param_typed: dtype {other:?} unsupported \
172                             (f32-uniform arena; widen small ints/floats to F32)"
173            ),
174        }
175    }
176
177    /// Typed run: widen each typed input to F32, run, then narrow each
178    /// output back to its declared dtype.
179    fn run_typed(
180        &mut self,
181        inputs: &[(&str, &[u8], rlx_ir::DType)],
182    ) -> Vec<(Vec<u8>, rlx_ir::DType)> {
183        let mut owned: Vec<(String, Vec<f32>)> = Vec::with_capacity(inputs.len());
184        for (name, data, dt) in inputs {
185            let v: Vec<f32> = match *dt {
186                rlx_ir::DType::F32 => {
187                    let n = data.len() / 4;
188                    unsafe { std::slice::from_raw_parts(data.as_ptr() as *const f32, n) }.to_vec()
189                }
190                rlx_ir::DType::F16 => {
191                    let n = data.len() / 2;
192                    let s =
193                        unsafe { std::slice::from_raw_parts(data.as_ptr() as *const half::f16, n) };
194                    s.iter().map(|h| h.to_f32()).collect()
195                }
196                rlx_ir::DType::BF16 => {
197                    let n = data.len() / 2;
198                    let s = unsafe {
199                        std::slice::from_raw_parts(data.as_ptr() as *const half::bf16, n)
200                    };
201                    s.iter().map(|h| h.to_f32()).collect()
202                }
203                // Integer/bool inputs (e.g. embedding indices `phone_ids`) are
204                // widened to f32, matching the f32-arena convention shared with
205                // the CPU backend (Gather etc. operate on f32-encoded indices).
206                rlx_ir::DType::I64 => {
207                    let n = data.len() / 8;
208                    let s = unsafe { std::slice::from_raw_parts(data.as_ptr() as *const i64, n) };
209                    s.iter().map(|&x| x as f32).collect()
210                }
211                rlx_ir::DType::I32 => {
212                    let n = data.len() / 4;
213                    let s = unsafe { std::slice::from_raw_parts(data.as_ptr() as *const i32, n) };
214                    s.iter().map(|&x| x as f32).collect()
215                }
216                rlx_ir::DType::U8 | rlx_ir::DType::I8 | rlx_ir::DType::Bool => {
217                    data.iter().map(|&b| b as f32).collect()
218                }
219                other => {
220                    panic!("rlx-wgpu run_typed: input '{name}' dtype {other:?} unsupported")
221                }
222            };
223            owned.push((name.to_string(), v));
224        }
225        let refs: Vec<(&str, &[f32])> = owned
226            .iter()
227            .map(|(n, d)| (n.as_str(), d.as_slice()))
228            .collect();
229        let dtypes = super::declared_output_dtypes(&self.io_manifest, self.inner.output_dtypes());
230        let outs = self.inner.run(&refs);
231        outs.into_iter()
232            .zip(
233                dtypes
234                    .into_iter()
235                    .chain(std::iter::repeat(rlx_ir::DType::F32)),
236            )
237            .map(|(v, dt)| (narrow_to_dtype(&v, dt), dt))
238            .collect()
239    }
240
241    fn clone_box(&self) -> Box<dyn ExecutableGraph> {
242        Box::new(WgpuExecutableWrapper {
243            inner: self.inner.clone_for_cache(),
244            io_manifest: self.io_manifest.clone(),
245        })
246    }
247
248    #[cfg(target_arch = "wasm32")]
249    fn wgpu_requires_host(&self) -> bool {
250        self.inner.requires_host_execution()
251    }
252
253    #[cfg(target_arch = "wasm32")]
254    fn wgpu_run_async<'a>(
255        &'a mut self,
256        inputs: &'a [(&'a str, &'a [f32])],
257    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<Vec<f32>>, String>> + 'a>>
258    {
259        if self.inner.requires_host_execution() {
260            return Box::pin(async move {
261                Err(
262                    "WebGPU wasm graph contains host-executed ops (GGUF dequant, LSTM, \
263                     collectives, FFT-host, …) — use WebGL fallback or CPU"
264                        .into(),
265                )
266            });
267        }
268        Box::pin(async move { Ok(self.inner.run_async(inputs).await) })
269    }
270}
271
272/// Cast every element of a wgpu f32 output buffer down to the
273/// declared output dtype, returning the corresponding byte stream.
274/// The arena keeps every value as f32; declared output dtypes
275/// (Bool, I8, I32, F16, ...) require an exit-time narrowing to be
276/// byte-identical with backends that store the native dtype.
277fn narrow_to_dtype(v: &[f32], dt: rlx_ir::DType) -> Vec<u8> {
278    use rlx_ir::DType;
279    match dt {
280        DType::F32 => {
281            let mut bytes = Vec::with_capacity(v.len() * 4);
282            for &x in v {
283                bytes.extend_from_slice(&x.to_le_bytes());
284            }
285            bytes
286        }
287        DType::F16 => {
288            let mut bytes = Vec::with_capacity(v.len() * 2);
289            for &x in v {
290                bytes.extend_from_slice(&half::f16::from_f32(x).to_le_bytes());
291            }
292            bytes
293        }
294        DType::BF16 => {
295            let mut bytes = Vec::with_capacity(v.len() * 2);
296            for &x in v {
297                bytes.extend_from_slice(&half::bf16::from_f32(x).to_le_bytes());
298            }
299            bytes
300        }
301        DType::F64 => {
302            let mut bytes = Vec::with_capacity(v.len() * 8);
303            for &x in v {
304                bytes.extend_from_slice(&(x as f64).to_le_bytes());
305            }
306            bytes
307        }
308        DType::I8 => v.iter().map(|&x| x as i8 as u8).collect(),
309        DType::U8 => v.iter().map(|&x| x as u8).collect(),
310        DType::I16 => {
311            let mut bytes = Vec::with_capacity(v.len() * 2);
312            for &x in v {
313                bytes.extend_from_slice(&(x as i16).to_le_bytes());
314            }
315            bytes
316        }
317        DType::I32 => {
318            let mut bytes = Vec::with_capacity(v.len() * 4);
319            for &x in v {
320                bytes.extend_from_slice(&(x as i32).to_le_bytes());
321            }
322            bytes
323        }
324        DType::U32 => {
325            let mut bytes = Vec::with_capacity(v.len() * 4);
326            for &x in v {
327                bytes.extend_from_slice(&(x as u32).to_le_bytes());
328            }
329            bytes
330        }
331        DType::I64 => {
332            let mut bytes = Vec::with_capacity(v.len() * 8);
333            for &x in v {
334                bytes.extend_from_slice(&(x as i64).to_le_bytes());
335            }
336            bytes
337        }
338        DType::Bool => v
339            .iter()
340            .map(|&x| if x != 0.0 { 1u8 } else { 0u8 })
341            .collect(),
342        // C64 (complex f32 pair) — the wgpu backend's f32 arena
343        // doesn't synthesize complex outputs today; this branch
344        // only fires if a graph somehow asks for a C64 output and
345        // the backend lowered it as 2N real floats. We pass the
346        // raw f32 stream straight through; downstream code that
347        // wants complex semantics is responsible for re-pairing.
348        DType::C64 => {
349            let mut bytes = Vec::with_capacity(v.len() * 4);
350            for &x in v {
351                bytes.extend_from_slice(&x.to_le_bytes());
352            }
353            bytes
354        }
355        // C128 (complex f64 pair) — like C64, the f32-uniform wgpu arena
356        // never synthesizes complex outputs (C128 casts/ops are rejected
357        // at lowering). Defensive only: widen the raw f32 lanes to f64 so
358        // the byte width matches the declared 16 B/elem C128 layout.
359        DType::C128 => {
360            let mut bytes = Vec::with_capacity(v.len() * 8);
361            for &x in v {
362                bytes.extend_from_slice(&(x as f64).to_le_bytes());
363            }
364            bytes
365        }
366    }
367}
rlx_runtime/backend/wgpu_backend.rs

rlx_runtime/backend/
wgpu_backend.rs
