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

//! Unified compile options.
//!
//! Replaces the historical mix of `compile()`, `compile_with_precision()`,
//! `compile_with_options()` with a single `Backend::compile(graph, &options)`.
//! New compile-time knobs can be added to `CompileOptions` without
//! changing the trait — backends just read what they care about.
//!
//! Builder-pattern API for ergonomics:
//!
//! ```rust,ignore
//! let opts = CompileOptions::new()
//!     .precision(Precision::F16)
//!     .policy(PrecisionPolicy::AutoMixed)
//!     .with_dce(true)
//!     .with_constant_folding(true);
//! ```

use crate::Precision;
use rlx_ir::OpKind;
use rlx_ir::logical_kernel::{KernelDispatchConfig, KernelDispatchPolicy};
use rlx_opt::rlx_compile::scaled_quant_insert::ScaledQuantConfig;
use rlx_opt::{FusionOptions, FusionTarget, PrecisionPolicy};
use std::collections::HashMap;

/// All knobs the compile pipeline understands.
/// Add new fields here rather than introducing new compile entry points.
#[derive(Debug, Clone)]
pub struct CompileOptions {
    /// Target numeric precision for execution. Default: F32.
    pub precision: Precision,
    /// Optional per-op precision policy (mixed precision rewrite).
    pub policy: Option<PrecisionPolicy>,
    /// Opt-in native low-precision GEMM. When set, every 2-D `Op::MatMul` is
    /// rewritten to a `ScaledMatMul` whose operands are dynamically quantized to
    /// this element format + scale layout — any [`ScaledFormat`](rlx_ir::ScaledFormat), including a
    /// parameterized `Custom` (e.g. `f4e3m0`). Off by default; changes numerics.
    pub scaled_quant: Option<ScaledQuantConfig>,
    /// Run the graph as if every value were stored in this format (weights,
    /// inputs and every intermediate; see
    /// [`emulate_format`](rlx_opt::rlx_compile::emulate_format)). `None` and
    /// `Some(F32)` leave the graph untouched. Env: `RLX_NUMBER_FORMAT`.
    pub number_format: Option<rlx_ir::NumberFormat>,
    /// What the number-format rewrite leaves alone (pre-rounded weights, kept
    /// params / ops). Env: `RLX_NUMBER_FORMAT_PARAMS_PREROUNDED`.
    pub number_format_options: rlx_opt::rlx_compile::emulate_format::EmulateOptions,
    /// RNG policy for in-graph [`Op::RngNormal`](rlx_ir::Op::RngNormal) / [`Op::RngUniform`](rlx_ir::Op::RngUniform) nodes.
    pub rng: rlx_ir::RngOptions,
    /// Run dead-code elimination as part of compile. Default: true.
    pub dce: bool,
    /// Run constant folding. Default: true (cheap, only helps).
    pub constant_folding: bool,
    /// Verbose pass logging. Equivalent to `RLX_VERBOSE=1` or
    /// [`rlx_ir::env::set("RLX_VERBOSE", "1")`].
    pub verbose: bool,
    /// Override fusion pipeline target (default: inferred from device).
    pub fusion_target: Option<FusionTarget>,
    /// Per-target fusion toggles (Metal env overrides, skip fusion, …).
    pub fusion_opts: FusionOptions,
    /// Arena alignment for buffer planning. Default: 64.
    pub arena_alignment: usize,
    /// Panic at compile time if fusion diagnostics report missed patterns.
    pub assert_fusion_clean: bool,
    /// Backend op claim set for backend-aware fusion + post-fusion
    /// legalization. Set by `Backend::compile` implementations.
    pub supported_ops: Option<&'static [OpKind]>,
    /// When set, specialize symbolic dims before backend lowering.
    pub dim_binding: Option<rlx_ir::DimBinding>,
    /// Bake fixed param tensors into constants before DCE / constant folding.
    pub param_bindings: Option<HashMap<String, Vec<f32>>>,
    /// Packed U8 GGUF weights for `Op::Param` nodes in `DequantMatMul` (TPU HLO bake).
    pub quant_param_bindings: Option<HashMap<String, Vec<u8>>>,
    /// Native vs common IR lowering ([`KernelDispatchConfig`], `RLX_KERNEL_DISPATCH=common`).
    pub kernel_dispatch: KernelDispatchConfig,
    /// Prevent Metal from lowering the graph through MPSGraph.
    ///
    /// This keeps the regular Metal thunk schedule while retaining all other
    /// compile options. Off by default.
    pub disable_mpsgraph: bool,
    /// On backends that list `OpKind::Scan` (host-fallback), Scans with
    /// `length <= scan_unroll_max_length` are IR-unrolled into body replicas so
    /// they run as ordinary device ops. Longer Scans keep `ScanHost`.
    /// Default: **64**. `0` disables unroll on claiming backends; `u32::MAX`
    /// prefers unroll for every eligible Scan.
    pub scan_unroll_max_length: u32,
    /// Hoist the param-invariant subgraph (compute that depends only on weights,
    /// never on `Op::Input`) into a separate "prepare" graph run ONCE, feeding
    /// its outputs into the main graph across forwards. Complements
    /// `param_bindings` (which folds such compute away at compile when weight
    /// VALUES are known); this handles run-time-only weights. Opt-in; default
    /// off (also settable via `RLX_CACHE_PARAM_INVARIANT=1`).
    pub cache_param_invariant: bool,
    /// Static provable NaN/Inf lint during compile (`RLX_LINT_NUMERICS`).
    pub lint_numerics: bool,
    /// Print fusion before/after report (`RLX_FUSION_REPORT`).
    pub fusion_report: bool,
    /// Skip Conv+Bias+Act fusion (`RLX_DISABLE_CONV_BIAS_ACT_FUSION`).
    pub disable_conv_bias_act_fusion: bool,
    /// Disable IO-gated peaks-only fusion (`RLX_NO_IO_PEAKS_OUTPUT`).
    pub no_io_peaks_output: bool,
}

impl Default for CompileOptions {
    fn default() -> Self {
        Self::from_env()
    }
}

impl CompileOptions {
    /// Defaults with **no** process-env overlay (tests / explicit builders).
    pub fn baseline() -> Self {
        Self {
            precision: Precision::F32,
            policy: None,
            scaled_quant: None,
            number_format: None,
            number_format_options: Default::default(),
            rng: rlx_ir::RngOptions::default(),
            dce: true,
            constant_folding: true,
            verbose: false,
            fusion_target: None,
            fusion_opts: FusionOptions::default(),
            arena_alignment: 64,
            assert_fusion_clean: false,
            supported_ops: None,
            dim_binding: None,
            param_bindings: None,
            quant_param_bindings: None,
            kernel_dispatch: KernelDispatchConfig::default(),
            disable_mpsgraph: false,
            scan_unroll_max_length: 64,
            cache_param_invariant: false,
            lint_numerics: false,
            fusion_report: false,
            disable_conv_bias_act_fusion: false,
            no_io_peaks_output: false,
        }
    }

    /// `baseline` plus Public/compile-layer `RLX_*` overlays from the registry.
    pub fn from_env() -> Self {
        let mut opts = Self::baseline();
        opts.apply_env();
        opts
    }

    /// Overlay compile-layer env keys (idempotent). Uses [`rlx_ir::env_registry`]
    /// so aliases resolve.
    pub fn apply_env(&mut self) {
        use rlx_ir::env_registry;
        if env_registry::flag("RLX_VERBOSE") {
            self.verbose = true;
        }
        if env_registry::flag("RLX_CACHE_PARAM_INVARIANT") {
            self.cache_param_invariant = true;
        }
        if env_registry::flag("RLX_DISABLE_MPSGRAPH") {
            self.disable_mpsgraph = true;
        }
        self.kernel_dispatch = KernelDispatchConfig::from_env();
        if env_registry::flag("RLX_LINT_NUMERICS") {
            self.lint_numerics = true;
        }
        if env_registry::flag("RLX_FUSION_REPORT") {
            self.fusion_report = true;
        }
        if env_registry::flag("RLX_DISABLE_CONV_BIAS_ACT_FUSION") {
            self.disable_conv_bias_act_fusion = true;
        }
        if env_registry::flag("RLX_NO_IO_PEAKS_OUTPUT") {
            self.no_io_peaks_output = true;
        }
        if let Some(spec) = env_registry::var("RLX_NUMBER_FORMAT") {
            match spec.parse::<rlx_ir::NumberFormat>() {
                Ok(f) => self.number_format = Some(f),
                // A typo must not silently run the graph in f32.
                Err(e) => panic!("RLX_NUMBER_FORMAT: {e}"),
            }
        }
        if env_registry::flag("RLX_NUMBER_FORMAT_PARAMS_PREROUNDED") {
            self.number_format_options.params_prerounded = true;
        }
        if let Some(list) = env_registry::var("RLX_NUMBER_FORMAT_KEEP") {
            self.number_format_options.keep_ops = parse_keep_ops(&list);
        }
    }

    pub fn new() -> Self {
        Self::from_env()
    }

    pub fn precision(mut self, p: Precision) -> Self {
        self.precision = p;
        self
    }
    pub fn policy(mut self, p: PrecisionPolicy) -> Self {
        self.policy = Some(p);
        self
    }
    /// Enable native low-precision GEMM for every 2-D matmul in the graph, in
    /// the given element format + scale layout. Accepts any [`ScaledFormat`](rlx_ir::ScaledFormat) via
    /// [`ScaledQuantConfig`] — e.g. `ScaledQuantConfig::fp8_e4m3()`, or a
    /// parameterized `ScaledFormat::custom(3, 0)` (`f4e3m0`).
    pub fn scaled_quant(mut self, cfg: ScaledQuantConfig) -> Self {
        self.scaled_quant = Some(cfg);
        self
    }
    /// Round weights to `format` as they are uploaded, inserting nothing into
    /// the graph. For a training graph whose forward was already rewritten
    /// with [`emulate_format`](crate::emulate_format): the backward must not
    /// round gradients, but each step's weight upload still should.
    pub fn round_weights_on_upload(mut self, format: rlx_ir::NumberFormat) -> Self {
        self.number_format = Some(format);
        self.number_format_options.keep_ops = OpKind::ALL.to_vec();
        self
    }
    /// Run the graph in `format` (see [`Self::number_format`]).
    pub fn number_format(mut self, format: rlx_ir::NumberFormat) -> Self {
        self.number_format = Some(format);
        self
    }
    pub fn rng(mut self, rng: rlx_ir::RngOptions) -> Self {
        self.rng = rng;
        self
    }
    pub fn rng_backend(mut self, backend: rlx_ir::RngBackend) -> Self {
        self.rng.backend = backend;
        self
    }
    pub fn rng_seed(mut self, seed: u64) -> Self {
        self.rng.seed = seed;
        self
    }
    pub fn no_policy(mut self) -> Self {
        self.policy = None;
        self
    }
    pub fn with_dce(mut self, on: bool) -> Self {
        self.dce = on;
        self
    }
    pub fn with_constant_folding(mut self, on: bool) -> Self {
        self.constant_folding = on;
        self
    }
    pub fn with_verbose(mut self, on: bool) -> Self {
        self.verbose = on;
        self
    }
    pub fn fusion_target(mut self, target: FusionTarget) -> Self {
        self.fusion_target = Some(target);
        self
    }
    pub fn fusion_opts(mut self, opts: FusionOptions) -> Self {
        self.fusion_opts = opts;
        self
    }
    pub fn arena_alignment(mut self, bytes: usize) -> Self {
        self.arena_alignment = bytes;
        self
    }
    pub fn supported_ops(mut self, ops: &'static [OpKind]) -> Self {
        self.supported_ops = Some(ops);
        self
    }
    pub fn assert_fusion_clean(mut self, on: bool) -> Self {
        self.assert_fusion_clean = on;
        self
    }
    pub fn dim_binding(mut self, binding: rlx_ir::DimBinding) -> Self {
        self.dim_binding = Some(binding);
        self
    }
    pub fn param_bindings(mut self, bindings: HashMap<String, Vec<f32>>) -> Self {
        self.param_bindings = Some(bindings);
        self
    }
    /// Packed U8 weights for TPU GGUF `DequantMatMul` param bake at compile time.
    pub fn quant_param_bindings(mut self, bindings: HashMap<String, Vec<u8>>) -> Self {
        self.quant_param_bindings = Some(bindings);
        self
    }
    pub fn kernel_dispatch(mut self, policy: KernelDispatchPolicy) -> Self {
        self.kernel_dispatch.policy = policy;
        self
    }

    pub fn kernel_dispatch_config(mut self, config: KernelDispatchConfig) -> Self {
        self.kernel_dispatch = config;
        self
    }

    /// Prefer IR-unrolling Scans with `length <= max` on Scan-claiming backends.
    pub fn scan_unroll_max_length(mut self, max: u32) -> Self {
        self.scan_unroll_max_length = max;
        self
    }

    /// Force listed logical kernels to use common IR even when native is in `supported_ops`.
    pub fn force_common_kinds(mut self, kinds: &'static [OpKind]) -> Self {
        self.kernel_dispatch.force_common_kinds = kinds;
        self
    }

    /// Keep listed logical kernels native even under `ForceCommon` / missing from `supported_ops`.
    pub fn force_native_kinds(mut self, kinds: &'static [OpKind]) -> Self {
        self.kernel_dispatch.force_native_kinds = kinds;
        self
    }

    pub fn with_lint_numerics(mut self, on: bool) -> Self {
        self.lint_numerics = on;
        self
    }

    pub fn with_fusion_report(mut self, on: bool) -> Self {
        self.fusion_report = on;
        self
    }

    pub fn with_disable_mpsgraph(mut self, on: bool) -> Self {
        self.disable_mpsgraph = on;
        self
    }

    pub fn with_cache_param_invariant(mut self, on: bool) -> Self {
        self.cache_param_invariant = on;
        self
    }

    pub fn with_no_io_peaks_output(mut self, on: bool) -> Self {
        self.no_io_peaks_output = on;
        self
    }

    pub fn with_disable_conv_bias_act_fusion(mut self, on: bool) -> Self {
        self.disable_conv_bias_act_fusion = on;
        self
    }
}

/// `Softmax,RmsNorm` → those kinds; `!MatMul` → every kind but `MatMul`.
/// Names are `OpKind`'s, case-insensitive. A typo panics rather than silently
/// keeping nothing.
fn parse_keep_ops(list: &str) -> Vec<OpKind> {
    let kind = |name: &str| {
        OpKind::ALL
            .iter()
            .copied()
            .find(|k| format!("{k:?}").eq_ignore_ascii_case(name.trim()))
            .unwrap_or_else(|| panic!("RLX_NUMBER_FORMAT_KEEP: unknown op kind {name:?}"))
    };
    if let Some(only) = list.trim().strip_prefix('!') {
        let only = kind(only);
        return OpKind::ALL.iter().copied().filter(|k| *k != only).collect();
    }
    list.split(',')
        .filter(|s| !s.trim().is_empty())
        .map(kind)
        .collect()
}