rlx_runtime/backend/
wgpu_backend.rs

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

rlx_runtime/backend/
wgpu_backend.rs
