rlx_runtime/
pkg.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//! Load `.rlxp` packages into a [`Session`](crate::Session).
6//!
7//! ```ignore
8//! let compiled = rlx_runtime::pkg::compile_rlxp(&session, "model.rlxp")?;
9//! ```
10//!
11//! Also: [`open_rlxp`], [`load_rlxp_graph`], placement helpers, and
12//! `rlxp://path#tensor` URI resolution in the weight cache.
13
14use crate::CompileOptions;
15use crate::compiled::CompiledGraph;
16use crate::session::Session;
17use anyhow::{Context, Result};
18use rlx_ir::Graph;
19use rlx_pkg::{Package, Placement};
20use std::path::Path;
21
22/// Open a package and return its MIR graph (hot-tier materialize).
23///
24/// Fails if the pack is weight-only (`encoding = "none"`).
25pub fn load_rlxp_graph(path: impl AsRef<Path>) -> Result<Graph> {
26    let pack = Package::open(path.as_ref())
27        .with_context(|| format!("open package {}", path.as_ref().display()))?;
28    pack.graph()
29}
30
31/// Open a flat `.rlxp`, ZIP, or package directory.
32pub fn open_rlxp(path: impl AsRef<Path>) -> Result<Package> {
33    Package::open(path.as_ref())
34        .with_context(|| format!("open package {}", path.as_ref().display()))
35}
36
37/// Compile the embedded MIR graph of an `.rlxp` package.
38///
39/// Requires `include_graph` / `encoding = bincode_graph_v1` at write time.
40pub fn compile_rlxp(session: &Session, path: impl AsRef<Path>) -> Result<CompiledGraph> {
41    let graph = load_rlxp_graph(path)?;
42    Ok(session.compile(graph))
43}
44
45/// Compile an `.rlxp` and bind every catalog tensor onto matching Params
46/// (and leave Constants to [`Package::graph`] materialize).
47///
48/// Used for MLX `--keep-packed` packs where DequantMatMul weights are Params
49/// named `{base}.weight/.scales/.biases`.
50pub fn compile_rlxp_bind_params(
51    session: &Session,
52    path: impl AsRef<Path>,
53) -> Result<CompiledGraph> {
54    let pack = open_rlxp(path.as_ref())?;
55    let graph = pack
56        .graph_with(rlx_pkg::MaterializeMode::HotOnly)
57        .with_context(|| format!("graph {}", path.as_ref().display()))?;
58    let mut compiled = session.compile(graph);
59    for (name, bytes, dtype) in pack.typed_weight_bindings()? {
60        compiled.set_param_typed(&name, &bytes, dtype);
61    }
62    Ok(compiled)
63}
64
65/// Compile with explicit [`CompileOptions`].
66pub fn compile_rlxp_with(
67    session: &Session,
68    path: impl AsRef<Path>,
69    options: &CompileOptions,
70) -> Result<CompiledGraph> {
71    let graph = load_rlxp_graph(path)?;
72    Ok(session.compile_with(graph, options))
73}
74
75/// Placement metadata from a package, if present.
76pub fn load_rlxp_placement(path: impl AsRef<Path>) -> Result<Option<Placement>> {
77    let pack = open_rlxp(path)?;
78    Ok(pack.placement().cloned())
79}
80
81/// Tensor names this rank should load for a placement (metadata-only packs).
82///
83/// When `placement` has no entry for a tensor, it is treated as replicated
84/// (all ranks load it). When an entry exists, only listed ranks load it.
85pub fn tensors_for_rank(placement: &Placement, rank: u32, all_names: &[&str]) -> Vec<String> {
86    all_names
87        .iter()
88        .filter_map(|name| match placement.ranks_for(name) {
89            Some(ranks) if ranks.contains(&rank) => Some((*name).to_string()),
90            Some(_) => None,
91            None => Some((*name).to_string()),
92        })
93        .collect()
94}
95
96/// Prefer a physical `dist/rank-{rank}/` tree when present; otherwise filter
97/// the global weight index by [`tensors_for_rank`].
98pub fn weight_names_for_rank(pack: &Package, rank: u32) -> Vec<String> {
99    if pack.rank_root(rank).is_some() {
100        return pack
101            .tensors_for_rank(rank)
102            .into_iter()
103            .filter(|e| e.rank == Some(rank))
104            .map(|e| e.name.clone())
105            .collect();
106    }
107    let Some(idx) = pack.weights_index() else {
108        return Vec::new();
109    };
110    let names: Vec<&str> = idx.names().collect();
111    match pack.placement() {
112        Some(pl) => tensors_for_rank(pl, rank, &names),
113        None => names.into_iter().map(str::to_string).collect(),
114    }
115}
rlx_runtime/pkg.rs

rlx_runtime/
pkg.rs
