rlx-compile 0.2.18

HIR → MIR → LIR compile pipeline for RLX
Documentation

rlx-compile

HIR → MIR → LIR compile pipeline for RLX: CompilePipeline, backend legalization, memory planning, precision / PTQ passes.

Depends on rlx-ir and rlx-fusion.

What's here

  • compiler — CompilePipeline, CompileResult, pipeline inspect hooks.
  • fusion_pipeline — fusion_passes, fusion_passes_for_supported, backend-aware pass lists.
  • legalize / legalize_broadcast — rewrite graphs for a target device.
  • memory — liveness analysis and arena slot assignment.
  • precision — auto-mixed precision policy (f32 ↔ f16/bf16 around matmul).
  • quant_insert / quant_propagate — PTQ Q/DQ insertion and propagation.
  • const_fold, dce, inline, promote_params, svg.
  • param_specialize — bake fixed Op::Param tensors into Op::Constant before DCE / constant folding (CompileOptions::param_bindings in runtime).
  • algebraic_simplify — same-shape x*0, x*1, x+0 rewrites (safe only when constant rank matches output; no scalar-broadcast folding).

Kernel dispatch transparency

prepare_graph_for_backend_with_report and friends produce a KernelDispatchReport (native vs common-IR vs rewritten vs unsupported). The runtime uses the same legalization path as compile; see the root README.md for env vars (RLX_DISPATCH_REPORT, RLX_KERNEL_DISPATCH).

Op sharding — decompose, price the IO, choose where

shard cuts a single op into pieces that have no data dependency on each other, so they can run on separate threads, devices or nodes:

use rlx_compile::shard::{choose, shardings, LinkCost, Residency, WorkerCost};

let options = shardings(&graph, node, 4);          // ways to cut it
let pick = choose(&options, &workers, link, Residency { shared_resident: true }, local_rate);

Three questions, and nothing else:

  1. Which axis, into what pieces. A conv can be cut by batch or by output channel; a matmul by rows or columns. Ranges are contiguous, disjoint and cover the axis exactly — a gap or an overlap would silently drop or double-count samples.
  2. What it costs in bytes. The axes differ in which operand gets replicated, and that is usually the whole story: a batch split replicates the weights, an out-channel split replicates the activations. Sliced operands cross a link once in total; shared ones cross once per worker, unless Residency { shared_resident: true } says they are already there — which is the data-parallel case, and it changes the answer.
  3. Where to run it. choose scores the slowest shard's compute against the transfer, and returns None for "run it whole" when no split beats that. A worker whose cost model is not rankable is excluded and reported, by the same CostCalibration::is_rankable rule device ranking already follows.

The module performs no graph surgery: ShardOption names the operand, axis and ranges, and the caller emits narrows and a concat, or ships ranges to workers. It also defines its own LinkCost/WorkerCost rather than importing rlx_collectives::planner::Link or rlx_runtime::BackendCostModel, which both sit above this crate — rlx_runtime::shard_placement is the adapter.

Exact, but not always bit-identical. Every option preserves the maths, yet the backend may block a narrower shape differently: batch and row splits came back bit-identical, while out-channel and column splits agreed to ~5e-7. A cohort that relies on every worker holding identical values must shard along batch/rows. Executed and checked in rlx-runtime/tests/shard_equivalence.rs — each option is run the way a worker would be fed it (given only its slice) and recombined, because a shard handed the wrong slice still returns numbers.

Install

[dependencies]
rlx-compile = "0.2"

Usually via rlx-opt or rlx.

Build / test

cargo test -p rlx-compile

License

MIT OR Apache-2.0.