rlx_runtime/
browser_support.rs

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

rlx_runtime/
browser_support.rs
