rlx_runtime/
browser_support.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//! Browser graph legality checks — reject transport/host-only paths before
6//! compile/run on wasm.
7
8use rlx_driver::Device;
9use rlx_ir::{Graph, Op, OpKind};
10
11/// `collective.*` custom op names blocked in the browser (no TCP transport).
12pub const COLLECTIVE_OPS: &[&str] = &[
13    "collective.all_reduce",
14    "collective.all_gather",
15    "collective.reduce_scatter",
16    "collective.copy_to_parallel",
17    "collective.reduce_from_parallel",
18    "collective.broadcast",
19    "collective.reduce",
20    "collective.all_to_all",
21    "collective.moe_dispatch",
22    "collective.moe_combine",
23    "collective.ppermute",
24    "collective.send",
25];
26
27/// Returns a human-readable reason when `graph` contains browser-forbidden ops.
28pub fn collective_block_reason(graph: &Graph) -> Option<String> {
29    for node in graph.nodes() {
30        if let Op::Custom { name, .. } = &node.op {
31            if COLLECTIVE_OPS.contains(&name.as_str()) {
32                return Some(format!(
33                    "collective op '{name}' unavailable in browser: no TCP transport on wasm32 \
34                     (rlx-driver ProcessGroup uses std::net sockets)"
35                ));
36            }
37        }
38    }
39    None
40}
41
42/// True when the graph has no browser-blocked transport/custom ops.
43pub fn passes_browser_preflight(graph: &Graph) -> bool {
44    collective_block_reason(graph).is_none()
45}
46
47/// Pick the highest-priority browser device that can legalize `graph`.
48pub fn select_browser_device_for_graph(graph: &Graph) -> Option<Device> {
49    for &device in crate::device_ext::BROWSER_DEVICE_PRIORITY {
50        if !crate::device_ext::is_available(device) {
51            continue;
52        }
53        if !passes_browser_preflight(graph) {
54            return None;
55        }
56        if crate::device_ext::supports_graph(device, graph) {
57            return Some(device);
58        }
59    }
60    None
61}
62
63/// First unsupported op kind for a browser device (for diagnostics).
64pub fn first_browser_unsupported_op(graph: &Graph, device: Device) -> Option<(usize, OpKind)> {
65    if collective_block_reason(graph).is_some() {
66        return graph.nodes().iter().find_map(|n| {
67            if let Op::Custom { name, .. } = &n.op {
68                if COLLECTIVE_OPS.contains(&name.as_str()) {
69                    return Some((n.id.0 as usize, OpKind::Custom));
70                }
71            }
72            None
73        });
74    }
75    crate::device_ext::first_unsupported_op(device, graph).map(|(i, op)| (i, op.kind()))
76}
rlx_runtime/browser_support.rs

rlx_runtime/
browser_support.rs
