rlx-runtime 0.2.18

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0
use super::*;
use rlx_ir::OpKind;
use rlx_wgpu::backend::WgpuExecutable;

pub struct WgpuBackend;

impl Backend for WgpuBackend {
    fn supported_ops(&self) -> &'static [OpKind] {
        rlx_wgpu::SUPPORTED_OPS
    }

    fn compile(&self, graph: Graph, options: &CompileOptions) -> Box<dyn ExecutableGraph> {
        use rlx_opt::pass::Pass as _;
        let graph = rlx_opt::LowerControlFlow.run(graph);
        let graph = rlx_opt::legalize_or_rewrite_for_backend(graph, rlx_wgpu::SUPPORTED_OPS)
            .unwrap_or_else(|errors| {
                panic!("{}", rlx_opt::format_legalize_error("wgpu", &errors));
            });
        let graph = apply_scan_device_preference(graph, options);
        let graph = crate::precompile::precompile_cleanup(graph, options);
        // Materialize mid-axis broadcasts before MarkElementwiseRegions:
        // wgpu Binary/region kernels only handle trailing/scalar broadcast
        // via modulus; EEG patch embed uses [1,C,1,D] + [1,C,P,D].
        let graph = rlx_opt::LegalizeBroadcast.run(graph);
        // ORDER MATTERS: targeted-pattern fusions run BEFORE the
        // catch-all `MarkElementwiseRegions`. Otherwise the region
        // pass swallows the Add / Activation nodes into chains and
        // FuseMatMulBiasAct / FuseResidualLN fail to match the
        // narrower patterns they look for. (Metal pipeline at line
        // ~377 already orders these correctly; wgpu was inverted
        // and silently shipped 13 unfused LayerNorms per BERT
        // forward where 12 should have been FusedResidualLN.)
        let compile_result = crate::stages::compile_graph_stages_for_backend(
            rlx_driver::Device::Gpu,
            graph,
            options,
            rlx_wgpu::SUPPORTED_OPS,
        );
        crate::stages::maybe_log_fusion(&compile_result.fusion);
        let graph = compile_result.lir.into_graph();
        let graph = match options.policy.clone() {
            Some(p) => rlx_opt::AutoMixedPrecision::new(p).run(graph),
            None => graph,
        };
        let (graph, io_manifest) = cpu_low_precision::prepare_f32_exec_graph(graph);
        Box::new(WgpuExecutableWrapper {
            inner: WgpuExecutable::compile_rng(graph, options.rng),
            io_manifest,
        })
    }

    fn compile_lir(&self, lir: LirModule, options: &CompileOptions) -> Box<dyn ExecutableGraph> {
        use rlx_opt::pass::Pass as _;
        // LIR may already contain fused ElementwiseRegions; legalize
        // broadcasts on the unfused graph shape before backend prep.
        let graph = rlx_opt::LegalizeBroadcast.run(lir.into_graph());
        let graph = prepare_fused_graph(graph, options, rlx_wgpu::SUPPORTED_OPS, "wgpu");
        let (graph, io_manifest) = cpu_low_precision::prepare_f32_exec_graph(graph);
        Box::new(WgpuExecutableWrapper {
            inner: WgpuExecutable::compile_rng(graph, options.rng),
            io_manifest,
        })
    }
}

struct WgpuExecutableWrapper {
    inner: WgpuExecutable,
    io_manifest: cpu_low_precision::IoDtypeManifest,
}

unsafe impl Send for WgpuExecutableWrapper {}

impl ExecutableGraph for WgpuExecutableWrapper {
    fn capabilities(&self) -> crate::ExecutableCapabilities {
        crate::ExecutableCapabilities {
            clone: true,
            gpu_handles: true,
            // NOT resident, despite having the feed. `register_kv_row_feed` +
            // `feed_kv_row` are here (shard-aware device-to-device arena copy),
            // but the recipe callers follow also needs `read_output_row` to keep
            // a host mirror for bucket re-seeding and rollback, and rlx-wgpu has
            // no row readback to delegate to — so this inherited the trait's
            // `None` and claimed the flag anyway. Callers bound fine and then
            // died on the first decode step (found via rlx-unlimited-ocr's
            // resident decode: "resident K row read, layer 0"). Flip this back
            // to `true` once rlx-wgpu grows a one-row staging readback and the
            // wrapper forwards `read_output_row`; `kv_resident_implies_a_row_
            // read_that_returns_the_row` is the test that will hold it honest.
            kv_resident: false,
            typed_io: true,
            active_extent: true,
            ..crate::ExecutableCapabilities::NONE
        }
    }

    fn set_param(&mut self, name: &str, data: &[f32]) {
        self.inner.set_param(name, data);
    }
    fn run(&mut self, inputs: &[(&str, &[f32])]) -> Vec<Vec<f32>> {
        self.inner.run(inputs)
    }
    fn run_read_outputs(
        &mut self,
        inputs: &[(&str, &[f32])],
        read_indices: Option<&[usize]>,
    ) -> Vec<Vec<f32>> {
        self.inner.run_read_outputs(inputs, read_indices)
    }
    fn bind_gpu_handle(&mut self, name: &str, data: &[f32]) -> bool {
        self.inner.bind_gpu_handle(name, data)
    }
    fn register_kv_row_feed(&mut self, handle_name: &str, output_index: usize) -> bool {
        self.inner.register_kv_row_feed(handle_name, output_index);
        true
    }
    fn feed_kv_row(&mut self, src_row: usize, dst_row: usize, row_elems: usize) -> bool {
        self.inner.feed_kv_row(src_row, dst_row, row_elems)
    }
    fn has_gpu_handle(&self, name: &str) -> bool {
        self.inner.has_gpu_handle(name)
    }
    fn set_gpu_handle_feed(&mut self, handle_name: &str, output_index: usize) -> bool {
        self.inner.set_gpu_handle_feed(handle_name, output_index);
        true
    }
    fn read_gpu_handle(&self, name: &str) -> Option<Vec<f32>> {
        self.inner.read_gpu_handle(name)
    }
    fn prepare_resident_gpu_handle(&mut self, name: &str) -> bool {
        self.inner.prepare_resident_gpu_handle(name)
    }
    fn set_active_extent(&mut self, extent: Option<(usize, usize)>) {
        self.inner.set_active_extent(extent);
    }

    fn set_rng(&mut self, rng: rlx_ir::RngOptions) {
        self.inner.set_rng(rng);
    }

    fn rng(&self) -> rlx_ir::RngOptions {
        self.inner.rng()
    }

    /// Typed param upload: the wgpu arena is f32-uniform, so widen F16/BF16 and
    /// small-integer/bool params (I64/I32/U32/Bool — e.g. `input_ids`, duration
    /// carry, alignment frame count) to F32 at the host boundary. U8/I8 packed
    /// quant weights keep their raw bytes.
    fn set_param_typed(&mut self, name: &str, data: &[u8], dtype: rlx_ir::DType) {
        match dtype {
            rlx_ir::DType::U8 | rlx_ir::DType::I8 => {
                self.inner.set_param_bytes(name, data);
            }
            rlx_ir::DType::F32 => {
                let f32_slice = rlx_ir::bytes::decode_le::<f32>(data);
                self.inner.set_param(name, f32_slice.as_ref());
            }
            rlx_ir::DType::BF16 => {
                // A BF16 matmul weight is kept PACKED (2 B/elem) in the wgpu
                // arena's bf16 side-buffer and unpacked in-shader — no host
                // widen, no ne*4 arena write.
                if !self.inner.set_param_bf16_packed(name, data) {
                    // Every OTHER BF16 param lands in the main arena, whose slot
                    // `plan_f32_uniform` sized NATIVE (2 B/elem — `F32Uniform`
                    // widens activations but keeps non-F32 *params* packed).
                    // Widening here wrote ne*4 bytes into an ne*2 slot: it
                    // overran the following param and disagreed with the host
                    // steps that read the slot back as bf16 (e.g.
                    // `Step::DequantGroupedMatmulMlxHost { scale_bf16 }`, whose
                    // flag comes from the IR dtype). MXFP4 MoE scales are the
                    // first non-matmul BF16 param to hit this. Write raw, like
                    // the U8 arm above.
                    self.inner.set_param_bytes(name, data);
                }
            }
            rlx_ir::DType::F16
            | rlx_ir::DType::F64
            | rlx_ir::DType::I64
            | rlx_ir::DType::I32
            | rlx_ir::DType::U32
            | rlx_ir::DType::Bool => {
                let f32 = super::widen_bytes_to_f32(data, dtype);
                self.inner.set_param(name, &f32);
            }
            other => panic!(
                "rlx-wgpu set_param_typed: dtype {other:?} unsupported \
                             (f32-uniform arena; widen small ints/floats to F32)"
            ),
        }
    }

    /// Typed run: widen each typed input to F32, run, then narrow each
    /// output back to its declared dtype.
    fn run_typed(
        &mut self,
        inputs: &[(&str, &[u8], rlx_ir::DType)],
    ) -> Vec<(Vec<u8>, rlx_ir::DType)> {
        let mut owned: Vec<(String, Vec<f32>)> = Vec::with_capacity(inputs.len());
        for (name, data, dt) in inputs {
            let v: Vec<f32> = match *dt {
                rlx_ir::DType::F32 => rlx_ir::bytes::decode_le_vec::<f32>(data),
                rlx_ir::DType::F16 => rlx_ir::bytes::decode_le::<u16>(data)
                    .iter()
                    .map(|&b| half::f16::from_bits(b).to_f32())
                    .collect(),
                rlx_ir::DType::BF16 => rlx_ir::bytes::decode_le::<u16>(data)
                    .iter()
                    .map(|&b| half::bf16::from_bits(b).to_f32())
                    .collect(),
                // Integer/bool inputs (e.g. embedding indices `phone_ids`) are
                // widened to f32, matching the f32-arena convention shared with
                // the CPU backend (Gather etc. operate on f32-encoded indices).
                rlx_ir::DType::I64 => rlx_ir::bytes::decode_le::<i64>(data)
                    .iter()
                    .map(|&x| x as f32)
                    .collect(),
                rlx_ir::DType::I32 => rlx_ir::bytes::decode_le::<i32>(data)
                    .iter()
                    .map(|&x| x as f32)
                    .collect(),
                rlx_ir::DType::U8 | rlx_ir::DType::I8 | rlx_ir::DType::Bool => {
                    data.iter().map(|&b| b as f32).collect()
                }
                other => {
                    panic!("rlx-wgpu run_typed: input '{name}' dtype {other:?} unsupported")
                }
            };
            owned.push((name.to_string(), v));
        }
        let refs: Vec<(&str, &[f32])> = owned
            .iter()
            .map(|(n, d)| (n.as_str(), d.as_slice()))
            .collect();
        let dtypes = super::declared_output_dtypes(&self.io_manifest, self.inner.output_dtypes());
        let outs = self.inner.run(&refs);
        outs.into_iter()
            .zip(
                dtypes
                    .into_iter()
                    .chain(std::iter::repeat(rlx_ir::DType::F32)),
            )
            .map(|(v, dt)| (narrow_to_dtype(&v, dt), dt))
            .collect()
    }

    fn clone_box(&self) -> Box<dyn ExecutableGraph> {
        Box::new(WgpuExecutableWrapper {
            inner: self.inner.clone_for_cache(),
            io_manifest: self.io_manifest.clone(),
        })
    }

    #[cfg(target_arch = "wasm32")]
    fn wgpu_requires_host(&self) -> bool {
        self.inner.requires_host_execution()
    }

    fn wgpu_host_summary(&self) -> Option<String> {
        self.inner.host_summary()
    }

    #[cfg(target_arch = "wasm32")]
    fn wgpu_run_async<'a>(
        &'a mut self,
        inputs: &'a [(&'a str, &'a [f32])],
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<Vec<f32>>, String>> + 'a>>
    {
        if self.inner.requires_host_execution() {
            // Name what is in the way rather than listing what it might have
            // been: a conv graph's answer is one `Conv2dHost` per convolution,
            // which tells a caller whether to change backend or change model.
            let detail = self
                .inner
                .host_summary()
                .unwrap_or_else(|| "host-executed steps".to_string());
            return Box::pin(async move {
                Err(format!(
                    "WebGPU wasm cannot run this graph — {detail}. Use WebGL (if the \
                     graph trains there) or CPU."
                ))
            });
        }
        Box::pin(async move { Ok(self.inner.run_async(inputs).await) })
    }
}

/// Cast every element of a wgpu f32 output buffer down to the
/// declared output dtype, returning the corresponding byte stream.
/// The arena keeps every value as f32; declared output dtypes
/// (Bool, I8, I32, F16, ...) require an exit-time narrowing to be
/// byte-identical with backends that store the native dtype.
fn narrow_to_dtype(v: &[f32], dt: rlx_ir::DType) -> Vec<u8> {
    use rlx_ir::DType;
    match dt {
        DType::F32 => {
            let mut bytes = Vec::with_capacity(v.len() * 4);
            for &x in v {
                bytes.extend_from_slice(&x.to_le_bytes());
            }
            bytes
        }
        DType::F16 => {
            let mut bytes = Vec::with_capacity(v.len() * 2);
            for &x in v {
                bytes.extend_from_slice(&half::f16::from_f32(x).to_le_bytes());
            }
            bytes
        }
        DType::BF16 => {
            let mut bytes = Vec::with_capacity(v.len() * 2);
            for &x in v {
                bytes.extend_from_slice(&half::bf16::from_f32(x).to_le_bytes());
            }
            bytes
        }
        DType::F64 => {
            let mut bytes = Vec::with_capacity(v.len() * 8);
            for &x in v {
                bytes.extend_from_slice(&(x as f64).to_le_bytes());
            }
            bytes
        }
        DType::I8 => v.iter().map(|&x| x as i8 as u8).collect(),
        DType::U8 => v.iter().map(|&x| x as u8).collect(),
        DType::I16 => {
            let mut bytes = Vec::with_capacity(v.len() * 2);
            for &x in v {
                bytes.extend_from_slice(&(x as i16).to_le_bytes());
            }
            bytes
        }
        DType::I32 => {
            let mut bytes = Vec::with_capacity(v.len() * 4);
            for &x in v {
                bytes.extend_from_slice(&(x as i32).to_le_bytes());
            }
            bytes
        }
        DType::U32 => {
            let mut bytes = Vec::with_capacity(v.len() * 4);
            for &x in v {
                bytes.extend_from_slice(&(x as u32).to_le_bytes());
            }
            bytes
        }
        DType::I64 => {
            let mut bytes = Vec::with_capacity(v.len() * 8);
            for &x in v {
                bytes.extend_from_slice(&(x as i64).to_le_bytes());
            }
            bytes
        }
        DType::Bool => v
            .iter()
            .map(|&x| if x != 0.0 { 1u8 } else { 0u8 })
            .collect(),
        // C64 (complex f32 pair) — the wgpu backend's f32 arena
        // doesn't synthesize complex outputs today; this branch
        // only fires if a graph somehow asks for a C64 output and
        // the backend lowered it as 2N real floats. We pass the
        // raw f32 stream straight through; downstream code that
        // wants complex semantics is responsible for re-pairing.
        DType::C64 => {
            let mut bytes = Vec::with_capacity(v.len() * 4);
            for &x in v {
                bytes.extend_from_slice(&x.to_le_bytes());
            }
            bytes
        }
        // C128 (complex f64 pair) — like C64, the f32-uniform wgpu arena
        // never synthesizes complex outputs (C128 casts/ops are rejected
        // at lowering). Defensive only: widen the raw f32 lanes to f64 so
        // the byte width matches the declared 16 B/elem C128 layout.
        DType::C128 => {
            let mut bytes = Vec::with_capacity(v.len() * 8);
            for &x in v {
                bytes.extend_from_slice(&(x as f64).to_le_bytes());
            }
            bytes
        }
    }
}