showcase: rust execution engine - #5732
Conversation
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
b6c76d8 to
107f313
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #5732 +/- ##
==========================================
- Coverage 98.11% 97.90% -0.22%
==========================================
Files 340 346 +6
Lines 32743 33965 +1222
==========================================
+ Hits 32127 33252 +1125
- Misses 616 713 +97 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
awesome @BradyPlanden :) I'll have a more detailed look at the code, and I'll look into the diffsol slowdowns in comparison with idaklu, some possible sources are:
|
Three defects, none of which depend on the API shape: The cache scratch vector lived on the solver object, so two threads evaluating one Function shared it. Moved onto BrentMemory, which CasADi hands out per concurrent evaluation. The bracket test read `fa * fb <= 0`, and two residuals near the underflow limit multiply to zero, which reads as a sign change. Now tested by sign, with NaN handled explicitly, so a bracket holding no root is reported as one. The NumPy path returned NaN where the plugin raises: a blanket `except (ValueError, RuntimeError)` around brentq swallowed both an empty bracket and a failure to converge along with the tree probes it was meant to absorb. Those two now raise SolverError, and only a probe returns NaN. `Brent` and `BrentUnknown` also lose their public names, becoming `_Brent` and `_BrentUnknown` alongside `_BaseAverage`, and the docs page goes. The node has no in-tree consumer yet, and lifting the residual into a closed sub-expression -- which is where this is headed, and what the Rust tape backend in #5732 wants too -- will change the constructor. Cheaper to withdraw the commitment now than to break it later. The changelog bullet goes with it: nothing user-facing is left to announce. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
martinjrobins
left a comment
martinjrobins
left a comment
There was a problem hiding this comment.
Hi @BradyPlanden, I've had a first pass through, I skimmed a lot of it cause it is kinda big :) so I may have misunderstood some things, but these are my initial comments. I'll keep on having a look when I can
| /// | ||
| /// `Arc` marks the state shared with the `PreparedSolver` across solves, which | ||
| /// is `Send + Sync`; everything built per solve is owned outright. | ||
| pub struct Equations { |
There was a problem hiding this comment.
you already have a ModelEvaluator struct, so this struct seems to be duplicating this
functionality
| pub(crate) compiled: Arc<CompiledModel>, | ||
| /// Shared with the caller, which evaluates output tapes against the same | ||
| /// scratch between steps on the output-variable path. | ||
| pub(crate) ws: Rc<RefCell<Workspace>>, |
There was a problem hiding this comment.
why the Rc? should only need a RefCell, then either share a reference to your
ops, or better yet use a ModelEvaluator struct containing a RefCell and share a referfence
to that.
| //! `Equations::reset()` always returns `None` and `ResetOp` is never | ||
| //! constructed. It exists solely to satisfy the `Reset` associated-type bounds | ||
| //! of `OdeEquationsImplicitSens`. | ||
|
|
There was a problem hiding this comment.
don't need this, use diffsol's UnitCallable
| //! `Equations::reset()` always returns `None` and `ResetOp` is never | ||
| //! constructed. It exists solely to satisfy the `Reset` associated-type bounds | ||
| //! of `OdeEquationsImplicitSens`. | ||
|
|
There was a problem hiding this comment.
if you use pybamm's unified experiment mode, this will give you the reset and you can then
get proper gradients through the reset.
| /// diffsol mints one of these per callback invocation, so it holds nothing of | ||
| /// its own: every field is a reference into the equations, making a mint a | ||
| /// handful of pointer copies. | ||
| pub struct RhsOp<'a> { |
There was a problem hiding this comment.
second paragraph of this docstring is redundent
| pub name: Option<String>, | ||
| } | ||
|
|
||
| impl FunctionSignature { |
There was a problem hiding this comment.
filename should be function_signature.rs
| pub name: Option<String>, | ||
| } | ||
|
|
||
| impl FunctionSignature { |
There was a problem hiding this comment.
I spent a long time trying to find out where this was used in the ffi, but I don't think it
is. perhaps it should be called FunctionMetadata.
| ) | ||
|
|
||
|
|
||
| class ObservationBackend(ABC): |
There was a problem hiding this comment.
Observation sounds like a data measurement, perhaps SolutionBackend?
|
|
||
|
|
||
| class ObservationBackend(ABC): | ||
| """How a Solution lowers a variable name into evaluable leaves. |
There was a problem hiding this comment.
unclear what a leaf is?
I read this a couple of times but can't figure it out. what is the key
in backend[key]? is it a variable name? a segment index? a slice of segments?
Looking further down, it looks like the backends use interior mutability,
so you might want to say they provide an immutable api.
| def __getitem__(self, key: slice) -> ObservationBackend: | ||
| """This backend restricted to a slice of the Solution's segments.""" | ||
|
|
||
| @abstractmethod |
There was a problem hiding this comment.
missing return type on signature
|
Thanks for the initial review @martinjrobins! I'm now on hols now for a few weeks so once I'm back I'll take a deeper look. From an initial glance, I think you've correctly found some of the duplications / sloppy abstracts that have crept into the implementation over the few months of building. I'm working towards tightening these at the moment and will update this branch with the corrections. That said, I'm not very keen on supporting this branch long-term, so I will start splitting this into smaller PRs pretty quickly :) |
2 participants
This PR adds a Rust execution engine as discussed in #5667. This engine compiles a discretised model's expression trees: RHS, algebraic residuals, Jacobians, events, output variables, into a flat instruction tape evaluated entirely in Rust. It updates the IDAKLU solver with an FFI to integrate with this engine without crossing the Python layer, and adds a pure-Rust diffsol solver implementation.
The engine is opt-in through
BaseModel.convert_to_format:This engine lands as an optional backend for PyBaMM, but it offers the structure to replace CasADi as the default once it stabilises. This branch moves the default to
'rust', showcasing full test suite coverage.Mental model
convert_to_formatalready selected the expression backend ("casadi","jax","python"); this adds"rust". Four components sit behind it:pybamm-rust/pybamm-corepybamm-rust/pybamm-pythonRustFunctionssubstituted forCasadiFunctions.pybammsolvers/.../Expressions/Rustpybamm.DiffsolSolver: pure-Rust BDF (diffsol + faer). No SUNDIALS, no CasADi in the path.solvers/diffsol_solver.pyPerformance has two axes
ProcessedVariabledrives CasADi with a per-timepoint Python loop; the Rust path runs one native tangent sweep per parameter across the whole trajectory, GIL released. Scales with output points.Performance
The baseline is
origin/main. Every CasADi number below was measured in a second checkout oforigin/main(abfdff33f) with its own replicated environment. The Rust numbers come from this branch.Top-level
rust_idaklubeatsmain's CasADi default in all 72 cells, from 1.22× to 12.36×.rust_diffsolwins 69 of 72, see Gradients.1C discharge
Triangle-wave drive cycle
Output variables
Passing
output_variables=["Voltage [V]"]moves observation inside the solve and skips storing the full state.Repeated evaluation
One solve per parameter draw, read through the interpolating call interface, representing the inner loop of a fitting or inference run. With
output_variables=set therust_idakluwin rises to 1.35–6.98×.Gradients
Forward sensitivities. The largest ratios in the set, and the ones needing the most care to read.
mainrust_idaklu, 1C dischargerust_idaklu, drive cyclecasadi_idaklu_aot, 1C dischargeAOT-compiling the CasADi kernels is worth ~1.5× when integration dominates and decays to nothing once observation does. The Rust path moves the opposite way and grows to 8.08×.
Two caveats.
main's IDAKLU never callsIDASetSensParams, so it runs at IDAS'spbar = 1default; this PR passespbar_i = |p_i|, which moves step selection for both IDAKLU backends.rust_diffsolerror-controls the scaled sensitivities while neither IDAKLU backend does (IDASetSensErrConis never called), so it integrates strictly more state under error control than the rows it is timed against, which is why it is slower thanmain's CasADi.Setting
output_variables=shrinks the gradient win to 1.13–3.88×, because the observation saving is already taken by the restriction.The experiment path
pybamm.Experimentwith four steps, 10 s period againstmain's CasADi, asrust_idaklu/rust_diffsol. Runnable by using--protocols experimentin the benchmark suite.All 9 cells win. This protocol sits outside the sweep above, as its step period fixes the output grid.
output_variables=pluscalculate_sensitivitiesis not supported across a step boundary, on any backend or either branch: the handoff needsd(state)/dpat the boundary and an outputs-only solve does not store it. This PR makes that fail uniformly with a namedSolverError. Onmainthe same combination silently returns wrong gradients (off by ~110–120× tolerance), which is how it was found.Where the time goes
DFN sensitivities, 1C discharge, at 10,000 output points (ms):
main)main)Rust's integration lands level with AOT-compiled CasADi (109 ms against 106 ms) with no compile step in front of it, and the end-to-end 7.87× comes mostly from the other axis: 1,560 ms of observation becomes 93 ms.
Setup cost
The fair "fast CasADi" comparator is
options={"compile": True}, which shells out togcc -O3 -march=native. Rust's "compile" is an in-memory IR build:(first-solve setup, 1C discharge at 1,000 points; CasADi measured on
main, cache isolated per case, each row verified as a fresh compile followed by a fresh-process disk reload. The non-AOTcasadi_idakluprepare is 14.6 / 36.2 / 40.3 ms.)Numerical parity and accuracy
Correctness of the Rust backend is established in the test suite, where the kernels (RHS, Jacobians, events, outputs, tangents) are checked against CasADi expression-for-expression at near machine precision (the shared parity fixture asserts rtol
1e-10, atol1e-14).rust_idaklufeeds those kernels to the same SUNDIALS integrator, so its trajectories are the same computation. The tables below confirm that parity holds end-to-end through the FFI asrust_idaklu's worst errors matchcasadi_idaklu's to the digit in every cell.We can use the benchmark to measure the solution accuracy of
rust_diffsolin reference to IDAKLU/CasADi. This is done below, where every row is scored against a converged reference ofatol = rtol = 1e-10(falling back to1e-9where that will not converge) with no solution interpolation. Rows outside tolerance, by lane and output-point count:At 1,000 and 10,000 output points 286 of 288 rows pass. The two misses are a narrow state-accuracy gap for a DFN drive cycle at 1.16–1.17× tolerance.
The inference-lane block at 10 and 100 points is a property of the output grid, where that lane reads strictly between stored nodes, and a 10-point grid cannot resolve a discharge between its nodes however the trajectory was computed.
mainmisses in the same proportion. Where the interpolants differ the Rust diffsol path is the most accurate of the three, because it supplies derivatives for Hermite interpolation where both IDAKLU backends fall back to a linear chord:casadi_idaklu(main)rust_idaklurust_diffsolThe identical
casadi_idaklu/rust_idaklurows are the end-to-end form of the parity the shared-integrator design was meant to buy.Binary size
This is mostly because CasADi's wheel carries IPOPT, HiGHS, Cbc/Clp, qpOASES, MUMPS and a bundled gfortran that we never use and can't decline. The Rust improvement is shown in the dependency footprint below where diffsol, faer and the whole CAS/AD stack are statically linked into one object.
Size Overview
+73,938 / −923 against
origin/main, with the majority coming from tests:src/IDASetSensParamssrc/uv.lockand the RustCargo.lockReproducing
From the root directory, run the below command. Additional flags are documented in the benchmark suite.
uv run python benchmarks/run_rust_observability.py \ --lane solver --models SPM SPMe DFN \ --protocols cc_discharge drive_cycle \ --warmup 5 --aot all --output-points 1000 --json solver_1000.jsonEvery JSON payload embeds its own provenance (git revision, dirty flag, platform, package versions).