diff --git a/.github/scripts/run_cargo.py b/.github/scripts/run_cargo.py new file mode 100644 index 0000000000..e8a381f280 --- /dev/null +++ b/.github/scripts/run_cargo.py @@ -0,0 +1,26 @@ +"""Run `cargo` with the given arguments, or skip where no Rust toolchain is installed. + +The cargo pre-commit hooks go through this so that `pre-commit run --all-files` works +without Rust; CI installs Rust, so the `style` job still runs them. +""" + +from __future__ import annotations + +import shutil +import subprocess +import sys + + +def main() -> int: + cargo = shutil.which("cargo") + if cargo is None: + print( + "cargo is not on PATH, so this Rust hook is skipped; install Rust " + "(https://rustup.rs) to run it." + ) + return 0 + return subprocess.call([cargo, *sys.argv[1:]]) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index f173b473f4..40519194db 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -25,7 +25,7 @@ jobs: strategy: fail-fast: false matrix: - language: [python, c-cpp] + language: [python, c-cpp, rust] steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: diff --git a/.github/workflows/test_on_push.yml b/.github/workflows/test_on_push.yml index 66eab22119..5b1a3851af 100644 --- a/.github/workflows/test_on_push.yml +++ b/.github/workflows/test_on_push.yml @@ -59,6 +59,7 @@ jobs: - '.github/workflows/_build_solver_wheels.yml' pybamm: - 'packages/pybamm/**' + - 'packages/pybamm-rust/**' - 'pyproject.toml' - 'uv.lock' - 'noxfile.py' @@ -103,6 +104,50 @@ jobs: echo "macos_runners=$macos" } >> "$GITHUB_OUTPUT" + rust_msrv: + needs: changes + if: ${{ github.event_name == 'push' || needs.changes.outputs.pybamm == 'true' }} + runs-on: ubuntu-latest + permissions: + contents: read + name: Rust MSRV (1.89) + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - name: Install the MSRV toolchain + run: rustup toolchain install 1.89 --profile minimal --no-self-update + - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + with: + workspaces: packages/pybamm-rust + key: msrv-1.89 + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Check against the declared MSRV + run: > + cargo +1.89 check --manifest-path packages/pybamm-rust/Cargo.toml + --locked --all-targets --all-features + + rust_tests: + needs: changes + if: ${{ github.event_name == 'push' || needs.changes.outputs.pybamm == 'true' }} + runs-on: ubuntu-latest + permissions: + contents: read + name: Rust tests + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + with: + workspaces: packages/pybamm-rust + key: tests + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Run the Rust test suite + run: > + cargo test --manifest-path packages/pybamm-rust/Cargo.toml + --locked --workspace --all-features + style: runs-on: ubuntu-latest permissions: @@ -125,6 +170,15 @@ jobs: with: python-version: 3.12 + # Pinned, as each Rust release adds pedantic and nursery lints that `-D warnings` fails on. + - name: Install the pinned Rust toolchain + run: | + rustup toolchain install 1.99.1 --profile minimal --component clippy,rustfmt --no-self-update + rustup default 1.99.1 + - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + with: + workspaces: packages/pybamm-rust + save-if: ${{ github.ref == 'refs/heads/main' }} - uses: astral-sh/setup-uv@1c37ad07a6a961277cf70c0d37d6f313000f5884 # v10.3.0 with: version: "latest" @@ -355,6 +409,8 @@ jobs: needs: - changes - style + - rust_msrv + - rust_tests - build_solver - run_unit_tests - run_coverage diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f3c9658788..8f461a8ebd 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,8 +2,8 @@ ci: autoupdate_commit_msg: "chore: update pre-commit hooks" autofix_commit_msg: "style: pre-commit fixes" autoupdate_schedule: "weekly" - # No base ref here, so comment-slop would scan nothing; `style` gates it. - skip: [comment-slop] + # No base ref or Rust toolchain here; the `style` CI job gates these. + skip: [comment-slop, cargo-fmt, cargo-clippy] repos: - repo: https://github.com/astral-sh/ruff-pre-commit @@ -83,3 +83,19 @@ repos: additional_dependencies: [pyyaml] files: ^\.github/(workflows/test_on_push\.yml|scripts/check_ci_gate\.py)$ pass_filenames: false + # Skipped, saying so, where cargo is not installed. + - id: cargo-fmt + name: cargo fmt + entry: python .github/scripts/run_cargo.py fmt --manifest-path packages/pybamm-rust/Cargo.toml --all + language: python + types: [rust] + pass_filenames: false + verbose: true + - id: cargo-clippy + name: cargo clippy + entry: python .github/scripts/run_cargo.py clippy --quiet --manifest-path packages/pybamm-rust/Cargo.toml --all-targets --all-features + language: python + types: [rust] + pass_filenames: false + verbose: true + args: ["--", "-D", "warnings"] diff --git a/CHANGELOG.md b/CHANGELOG.md index c99486239e..2fec95f394 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ ## Features - Replaced Codacy with open checks that also run locally. Ruff's `flake8-bandit` (`S`) rules replace Bandit, and CodeQL scans the Python and C++ code. `nox -s type-check` runs Pyrefly over `pybamm`, which now type-checks cleanly, and a new type error fails CI. clang-tidy reports solver findings to code scanning as an advisory check. ([#5974](https://github.com/pybamm-team/PyBaMM/pull/5974)) +- Added `pybamm-core` (`packages/pybamm-rust`), a Rust expression compiler: an arena DAG with simplification, CSE and zero propagation, a typed IR, scalar and lane-batched interpreters, and forward-mode AD. CI runs `cargo test`, an MSRV check, `cargo fmt`, `cargo clippy` and CodeQL. ([#5956](https://github.com/pybamm-team/PyBaMM/pull/5956)) +- `pybamm-core` evaluates transcendentals through the platform C library, as CasADi does: bitwise identical across evaluation paths, last-bit differences across platforms. `pow` is exact for exponents `0`, `1`, `2` and `-1`; `erf`, `asinh` and IEEE 754-2019 `minimum`/`maximum` use `libm`. ([#5956](https://github.com/pybamm-team/PyBaMM/pull/5956)) ## Bug fixes diff --git a/packages/pybamm-rust/.gitignore b/packages/pybamm-rust/.gitignore new file mode 100644 index 0000000000..fcf99788dd --- /dev/null +++ b/packages/pybamm-rust/.gitignore @@ -0,0 +1,7 @@ +/target/ +*/target/ +*.so +*.dylib +*.pyd +.venv/ +uv.lock diff --git a/packages/pybamm-rust/Cargo.lock b/packages/pybamm-rust/Cargo.lock new file mode 100644 index 0000000000..18d5018073 --- /dev/null +++ b/packages/pybamm-rust/Cargo.lock @@ -0,0 +1,497 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + +[[package]] +name = "bit-set" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "bytemuck" +version = "1.25.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "bytemuck", + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "proptest" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" +dependencies = [ + "bit-set", + "bit-vec", + "bitflags", + "num-traits", + "rand", + "rand_chacha", + "rand_xorshift", + "regex-syntax", + "rusty-fork", + "tempfile", + "unarray", +] + +[[package]] +name = "pulp" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "046aa45b989642ec2e4717c8e72d677b13edd831a4d3b6cf37d9a3e54912496a" +dependencies = [ + "bytemuck", + "cfg-if", + "libm", + "num-complex", + "paste", + "pulp-wasm-simd-flag", + "raw-cpuid", + "reborrow", + "version_check", +] + +[[package]] +name = "pulp-wasm-simd-flag" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d8f70e07b9c3962945a74e59ca1c511bba65b6419468acc217c457d93f3c740" + +[[package]] +name = "pybamm-core" +version = "0.1.0" +dependencies = [ + "bincode", + "libm", + "proptest", + "pulp", + "rustc-hash", + "serde", + "thiserror", +] + +[[package]] +name = "quick-error" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rand_xorshift" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" +dependencies = [ + "rand_core", +] + +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags", +] + +[[package]] +name = "reborrow" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430" + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "rusty-fork" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc6bf79ff24e648f6da1f8d1f011e9cac26491b619e6b9280f2b47f1774e6ee2" +dependencies = [ + "fnv", + "quick-error", + "tempfile", + "wait-timeout", +] + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys", +] + +[[package]] +name = "thiserror" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "unarray" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wait-timeout" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11" +dependencies = [ + "libc", +] + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zerocopy" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] diff --git a/packages/pybamm-rust/Cargo.toml b/packages/pybamm-rust/Cargo.toml new file mode 100644 index 0000000000..0270b3ddda --- /dev/null +++ b/packages/pybamm-rust/Cargo.toml @@ -0,0 +1,102 @@ +[workspace] +resolver = "3" +members = [ + "pybamm-core", +] + +[workspace.package] +rust-version = "1.89" + +[workspace.lints.rust] +unsafe_code = "warn" +missing_debug_implementations = "warn" +rust_2018_idioms = { level = "warn", priority = -1 } +trivial_casts = "warn" +trivial_numeric_casts = "warn" +unused_lifetimes = "warn" +unused_qualifications = "warn" + +[workspace.lints.clippy] +correctness = { level = "deny", priority = -1 } +pedantic = { level = "warn", priority = -1 } +nursery = { level = "warn", priority = -1 } +perf = { level = "warn", priority = -1 } +complexity = { level = "warn", priority = -1 } +style = { level = "warn", priority = -1 } + +cast_possible_truncation = "warn" +cast_possible_wrap = "warn" +cast_sign_loss = "warn" +checked_conversions = "warn" +clone_on_ref_ptr = "warn" +debug_assert_with_mut_call = "warn" +empty_enums = "warn" +enum_glob_use = "warn" +expl_impl_clone_on_copy = "warn" +explicit_deref_methods = "warn" +explicit_into_iter_loop = "warn" +explicit_iter_loop = "warn" +filter_map_next = "warn" +flat_map_option = "warn" +float_cmp = "warn" +fn_params_excessive_bools = "warn" +implicit_clone = "warn" +inefficient_to_string = "warn" +invalid_upcast_comparisons = "warn" +large_digit_groups = "warn" +large_stack_arrays = "warn" +large_types_passed_by_value = "warn" +manual_ok_or = "warn" +map_flatten = "warn" +map_unwrap_or = "warn" +match_bool = "warn" +match_same_arms = "warn" +match_wild_err_arm = "warn" +match_wildcard_for_single_variants = "warn" +mut_mut = "warn" +needless_borrow = "warn" +needless_continue = "warn" +needless_for_each = "warn" +needless_pass_by_value = "warn" +option_option = "warn" +pub_underscore_fields = "warn" +range_minus_one = "warn" +range_plus_one = "warn" +redundant_closure_for_method_calls = "warn" +redundant_else = "warn" +ref_option_ref = "warn" +same_functions_in_if_condition = "warn" +semicolon_if_nothing_returned = "warn" +string_add_assign = "warn" +struct_excessive_bools = "warn" +trait_duplication_in_bounds = "warn" +trivially_copy_pass_by_ref = "warn" +uninlined_format_args = "warn" +unnested_or_patterns = "warn" +unreadable_literal = "warn" +unused_self = "warn" +used_underscore_binding = "warn" +zero_sized_map_values = "warn" + +missing_errors_doc = "allow" +missing_panics_doc = "allow" +module_name_repetitions = "allow" +must_use_candidate = "allow" +too_many_lines = "allow" +similar_names = "allow" +indexing_slicing = "allow" +cast_precision_loss = "allow" +many_single_char_names = "allow" +# Fallible APIs keep a Result signature even where no path errors yet +unnecessary_wraps = "allow" +# Baseline x86-64 wheels lower mul_add to a libm call (no FMA), and its single +# rounding breaks CasADi parity. +suboptimal_flops = "allow" + +[profile.release] +lto = "thin" +codegen-units = 1 + +[profile.bench] +lto = "thin" +codegen-units = 1 diff --git a/packages/pybamm-rust/pybamm-core/Cargo.toml b/packages/pybamm-rust/pybamm-core/Cargo.toml new file mode 100644 index 0000000000..1d936b8f66 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "pybamm-core" +version = "0.1.0" +edition = "2024" +rust-version.workspace = true + +[dependencies] +rustc-hash = "2" +thiserror = "2" +# `erf`, which `f64` lacks; `asinh`, as `f64`'s overflows past `f64::MAX / 2`; and IEEE +# 754-2019 `minimumNumber`/`maximumNumber`, as `f64::min`/`max` leave signed zeros unordered. +libm = "0.2" +serde = { version = "1", features = ["derive"], optional = true } +bincode = { version = "1", optional = true } + +[target.'cfg(target_arch = "x86_64")'.dependencies] +# Compiles matrix products for the detected instruction set, so `accumulate` needs no `unsafe`. +pulp = "0.22" + +[dev-dependencies] +proptest = "1" + +[features] +serialize = ["serde", "bincode"] + +[lints] +workspace = true diff --git a/packages/pybamm-rust/pybamm-core/src/accumulate.rs b/packages/pybamm-rust/pybamm-core/src/accumulate.rs new file mode 100644 index 0000000000..f91425800e --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/accumulate.rs @@ -0,0 +1,122 @@ +//! Fused multiply-add accumulation for matrix products. +//! +//! Stencil rows cancel large terms, so an unfused product's rounding lands at their scale. +//! A CPU without an FMA instruction accumulates unfused, as a software `mul_add` halves the speed. + +// Products must inline into `pulp`'s feature-enabled wrapper to compile as FMA instructions. +#![allow(clippy::inline_always)] + +/// How a matrix product accumulates `acc + a * b`. +pub trait Accumulate { + fn mul_acc(acc: f64, a: f64, b: f64) -> f64; +} + +/// `acc + a * b` with a single rounding. +pub enum Fused {} + +/// `acc + a * b` with the product and the sum each rounded. +pub enum Unfused {} + +impl Accumulate for Fused { + #[inline(always)] + fn mul_acc(acc: f64, a: f64, b: f64) -> f64 { + a.mul_add(b, acc) + } +} + +impl Accumulate for Unfused { + #[inline(always)] + fn mul_acc(acc: f64, a: f64, b: f64) -> f64 { + acc + a * b + } +} + +/// A matrix-product loop, generic over its accumulation. +pub trait Product { + type Output; + fn run(self) -> Self::Output; +} + +/// The widest instruction set this CPU supports, detected once. +#[cfg(target_arch = "x86_64")] +static ARCH: std::sync::LazyLock = std::sync::LazyLock::new(pulp::Arch::new); + +/// Whether the target guarantees an FMA instruction. +#[cfg(not(target_arch = "x86_64"))] +const TARGET_FMA: bool = cfg!(any(target_arch = "aarch64", target_feature = "fma")); + +/// Hands a [`Product`] to `pulp`, which compiles it for the detected instruction set. +#[cfg(target_arch = "x86_64")] +struct Fma

(P); + +#[cfg(target_arch = "x86_64")] +impl pulp::WithSimd for Fma

{ + type Output = P::Output; + + #[inline(always)] + fn with_simd(self, _simd: S) -> P::Output { + self.0.run::() + } +} + +/// Run `product`, fused where the CPU has an FMA instruction. +#[inline] +pub fn dispatch(product: P) -> P::Output { + // Baseline x86-64 lacks FMA, so `pulp` compiles the fused loop for the detected set. + #[cfg(target_arch = "x86_64")] + { + match *ARCH { + pulp::Arch::Scalar => product.run::(), + arch => arch.dispatch(Fma(product)), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + if TARGET_FMA { + product.run::() + } else { + product.run::() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// `-(1 + 2^-29) + (1 + 2^-30)^2`, whose exact value `2^-60` an unfused product rounds away. + struct CancellingRow; + + impl Product for CancellingRow { + type Output = f64; + fn run(self) -> f64 { + let (a, b) = (1.0 + 2f64.powi(-29), 1.0 + 2f64.powi(-30)); + A::mul_acc(A::mul_acc(0.0, -1.0, a), b, b) + } + } + + #[test] + fn a_cancelling_row_rounds_once_per_term_only_when_fused() { + assert_eq!( + CancellingRow.run::().to_bits(), + 2f64.powi(-60).to_bits() + ); + assert_eq!(CancellingRow.run::().to_bits(), 0f64.to_bits()); + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn aarch64_always_fuses() { + assert_eq!(dispatch(CancellingRow).to_bits(), 2f64.powi(-60).to_bits()); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn an_x86_cpu_fuses_exactly_when_it_has_avx2_and_fma() { + // Guards the detection, so a CPU that could fuse never silently runs unfused. + let has_fma = std::arch::is_x86_feature_detected!("avx2") + && std::arch::is_x86_feature_detected!("fma"); + let expected = if has_fma { 2f64.powi(-60) } else { 0.0 }; + assert_eq!(dispatch(CancellingRow).to_bits(), expected.to_bits()); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/arena.rs b/packages/pybamm-rust/pybamm-core/src/arena.rs new file mode 100644 index 0000000000..ca4fddb9f0 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/arena.rs @@ -0,0 +1,599 @@ +//! Flat storage for the expression DAG: every [`Node`] lives in one [`Arena`], +//! addressed by a `u32` [`NodeId`] that is only valid against that arena. + +use std::ops::Index; + +use crate::error::CoreError; +use crate::node::Node; + +/// Handle to a node in one [`Arena`], valid only against that arena. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub struct NodeId(u32); + +impl NodeId { + /// Position in the arena, for indexing side tables. + #[inline] + pub const fn index(self) -> usize { + self.0 as usize + } + + /// The underlying `u32`. + #[inline] + pub const fn raw(self) -> u32 { + self.0 + } +} + +impl From for NodeId { + #[inline] + fn from(val: u32) -> Self { + Self(val) + } +} + +/// Owner of an expression DAG: nodes in allocation order, addressed by [`NodeId`]. +/// +/// Append-only, so ids stay valid. Every child id is below its parent's, so ids are in +/// range and the graph is acyclic; [`alloc`](Self::alloc) and deserialising enforce it. +#[derive(Clone, Debug, Default)] +#[cfg_attr( + feature = "serialize", + derive(serde::Serialize, serde::Deserialize), + serde(try_from = "PersistedArena") +)] +pub struct Arena { + nodes: Vec, +} + +/// The first child of `node` not allocated before position `at`. +fn first_unallocated_child(node: &Node, at: usize) -> Option { + let mut unallocated = None; + node.for_each_child(|child| { + if child.index() >= at && unallocated.is_none() { + unallocated = Some(child); + } + }); + unallocated +} + +impl Arena { + /// An empty arena. + pub const fn new() -> Self { + Self { nodes: Vec::new() } + } + + /// Append `node` and return its id. + /// + /// Panics if a child is not yet allocated here, or past `u32::MAX` nodes. + #[must_use] + pub fn alloc(&mut self, node: Node) -> NodeId { + let at = self.nodes.len(); + if let Some(child) = first_unallocated_child(&node, at) { + panic!( + "node {at} names child {}, which is not allocated before it", + child.raw() + ); + } + let id = NodeId(u32::try_from(at).expect("arena exceeded u32::MAX nodes")); + self.nodes.push(node); + id + } + + /// Borrow a node. + #[inline] + pub fn get(&self, id: NodeId) -> &Node { + &self.nodes[id.index()] + } + + /// Nodes allocated so far, reachable or not. + pub const fn len(&self) -> usize { + self.nodes.len() + } + + /// Every allocated node in allocation order, reachable or not. + #[inline] + pub(crate) fn nodes(&self) -> &[Node] { + &self.nodes + } + + /// Whether nothing has been allocated yet. + pub const fn is_empty(&self) -> bool { + self.nodes.is_empty() + } + + /// Nodes reachable from `root`, children before parents and each visited once. + pub fn topological_order(&self, root: NodeId) -> Vec { + let n = self.len(); + let mut visited = vec![false; n]; + let mut order = Vec::with_capacity(n); + let mut stack: Vec = Vec::with_capacity(n); + stack.push(TopoFrame::Pending(root)); + while let Some(frame) = stack.pop() { + match frame { + // Pushed once per node: any other `Pending(id)` sits below this frame, since one + // pushed above it would come from a descendant of `id`, i.e. a cycle. + TopoFrame::Visit(id) => { + visited[id.index()] = true; + order.push(id); + }, + TopoFrame::Pending(id) => { + if visited[id.index()] { + continue; + } + stack.push(TopoFrame::Visit(id)); + self.get(id).for_each_child(|c| { + if !visited[c.index()] { + stack.push(TopoFrame::Pending(c)); + } + }); + }, + } + } + order + } +} + +/// The on-disk form of an [`Arena`], refused unless every child precedes its parent. +#[cfg(feature = "serialize")] +#[derive(serde::Deserialize)] +struct PersistedArena { + nodes: Vec, +} + +#[cfg(feature = "serialize")] +impl TryFrom for Arena { + type Error = CoreError; + + fn try_from(persisted: PersistedArena) -> Result { + for (at, node) in persisted.nodes.iter().enumerate() { + if let Some(child) = first_unallocated_child(node, at) { + return Err(CoreError::ChildNotBeforeParent { + node: at, + child: child.index(), + }); + } + } + Ok(Self { + nodes: persisted.nodes, + }) + } +} + +/// An expression DAG together with its root node. +/// +/// Fields are private so a rebuilt arena cannot be paired with a stale root. +#[derive(Clone, Debug)] +#[cfg_attr( + feature = "serialize", + derive(serde::Serialize, serde::Deserialize), + serde(try_from = "PersistedDag") +)] +pub struct Dag { + arena: Arena, + root: NodeId, +} + +/// The on-disk form of a [`Dag`], refused unless the root is a node of the arena. +#[cfg(feature = "serialize")] +#[derive(serde::Deserialize)] +struct PersistedDag { + arena: Arena, + root: NodeId, +} + +#[cfg(feature = "serialize")] +impl TryFrom for Dag { + type Error = CoreError; + + fn try_from(persisted: PersistedDag) -> Result { + let (root, len) = (persisted.root.index(), persisted.arena.len()); + if root >= len { + return Err(CoreError::RootOutOfRange { root, len }); + } + Ok(Self::new(persisted.arena, persisted.root)) + } +} + +impl Dag { + /// Pair an arena with a root it issued; panics if `root` is out of range. + pub const fn new(arena: Arena, root: NodeId) -> Self { + assert!( + root.index() < arena.len(), + "a DAG's root must be a node of its arena" + ); + Self { arena, root } + } + + /// The DAG's nodes. + pub const fn arena(&self) -> &Arena { + &self.arena + } + + /// The node this DAG evaluates to. + pub const fn root(&self) -> NodeId { + self.root + } + + /// Append to the DAG and move the root onto the appended node. + /// + /// Panics if `rewrite` returns a root outside the arena it leaves behind. + pub fn rewrite_in_place(&mut self, rewrite: impl FnOnce(&mut Arena, NodeId) -> NodeId) { + let root = rewrite(&mut self.arena, self.root); + assert!( + root.index() < self.arena.len(), + "a DAG's root must be a node of its arena" + ); + self.root = root; + } + + /// Borrow the pair. + pub const fn as_ref(&self) -> DagRef<'_> { + DagRef { + arena: &self.arena, + root: self.root, + } + } +} + +/// A borrowed [`Dag`]. +#[derive(Clone, Copy, Debug)] +pub struct DagRef<'a> { + arena: &'a Arena, + root: NodeId, +} + +impl<'a> DagRef<'a> { + /// Pair a borrowed arena with a root it issued; panics if `root` is out of range. + pub const fn new(arena: &'a Arena, root: NodeId) -> Self { + assert!( + root.index() < arena.len(), + "a DAG's root must be a node of its arena" + ); + Self { arena, root } + } + + /// The DAG's nodes. + pub const fn arena(self) -> &'a Arena { + self.arena + } + + /// The node this DAG evaluates to. + pub const fn root(self) -> NodeId { + self.root + } +} + +/// How far into `y`, `y_dot` and the packed inputs the sub-DAG at a root reads. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct ReadExtents { + /// One past the highest element of `y` read. + pub y: usize, + /// One past the highest element of `y_dot` read, or `None` if no `StateVectorDot` is reached. + pub y_dot: Option, + /// One past the highest packed input element read. + pub inputs: usize, +} + +impl ReadExtents { + /// Scan the nodes `root` reaches; tangent leaves read no input and are not counted. + pub fn scan(arena: &Arena, root: NodeId) -> Self { + // Children precede parents, so one descending sweep marks everything `root` reaches. + let mut reached = vec![false; root.index() + 1]; + reached[root.index()] = true; + let mut extents = Self::default(); + for (at, node) in arena.nodes()[..=root.index()].iter().enumerate().rev() { + if !reached[at] { + continue; + } + match node { + Node::StateVector { end, .. } => extents.y = extents.y.max(*end), + Node::StateVectorDot { end, .. } => { + extents.y_dot = Some(extents.y_dot.unwrap_or(0).max(*end)); + }, + Node::InputParameter { offset, width, .. } => { + extents.inputs = extents.inputs.max(offset.saturating_add(*width)); + }, + _ => {}, + } + node.for_each_child(|child| reached[child.index()] = true); + } + extents + } + + /// The state count both `y` and `y_dot` must cover. + pub fn n_states(self) -> usize { + self.y.max(self.y_dot.unwrap_or(0)) + } +} + +/// Reconcile a declared state count with the extent an expression reads. +/// +/// Declaring fewer than `extent` is an error (it would panic inside evaluation), as is +/// declaring more than the tape's `u32` state indices reach; `None` adopts the extent. +/// Per-state tables are sized by the count, so one near `u32::MAX` can still exhaust memory. +pub const fn resolve_state_count( + declared: Option, + extent: usize, +) -> Result { + match declared { + Some(n) if n < extent => Err(CoreError::StateCountTooSmall { + declared: n, + extent, + }), + Some(n) if n > u32::MAX as usize => Err(CoreError::StateCountTooLarge { + declared: n, + max: u32::MAX, + }), + Some(n) => Ok(n), + None => Ok(extent), + } +} + +enum TopoFrame { + Pending(NodeId), + /// Children are all emitted; emit this node now. + Visit(NodeId), +} + +impl Index for Arena { + type Output = Node; + + #[inline] + fn index(&self, id: NodeId) -> &Node { + self.get(id) + } +} + +/// Dense per-node side table, keyed by [`NodeId`]. +#[derive(Clone, Debug)] +pub(crate) struct NodeMap { + slots: Vec>, +} + +impl NodeMap { + /// A table sized for an arena of `arena_len` nodes, all entries empty. + #[inline] + pub(crate) fn new(arena_len: usize) -> Self { + let mut slots = Vec::with_capacity(arena_len); + slots.resize_with(arena_len, || None); + Self { slots } + } + + /// The value for `id`, or `None` if unset or out of range. + #[inline] + pub(crate) fn get(&self, id: NodeId) -> Option<&T> { + self.slots.get(id.index()).and_then(|opt| opt.as_ref()) + } + + /// Set the value for `id`, growing the table if needed; returns the old value. + #[inline] + pub(crate) fn insert(&mut self, id: NodeId, value: T) -> Option { + if id.index() >= self.slots.len() { + self.slots.resize_with(id.index() + 1, || None); + } + self.slots[id.index()].replace(value) + } + + /// Whether `id` has a value set. + #[inline] + pub(crate) fn contains_key(&self, id: NodeId) -> bool { + self.slots.get(id.index()).is_some_and(Option::is_some) + } +} + +impl NodeMap { + /// `id`'s counterpart in a pass's output, given children are rebuilt before parents. + #[inline] + pub(crate) fn rebuilt(&self, id: NodeId) -> NodeId { + *self + .get(id) + .expect("children are rebuilt before their parents") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::node::Node; + + #[test] + fn test_alloc_and_get() { + let mut arena = Arena::new(); + let id = arena.alloc(Node::Scalar(42.5)); + match arena.get(id) { + Node::Scalar(v) => assert!((*v - 42.5).abs() < f64::EPSILON), + _ => panic!("expected Scalar"), + } + } + + #[test] + fn test_sequential_ids() { + let mut arena = Arena::new(); + let id0 = arena.alloc(Node::Scalar(1.0)); + let id1 = arena.alloc(Node::Scalar(2.0)); + assert_eq!(id0.index(), 0); + assert_eq!(id1.index(), 1); + } + + #[test] + fn test_topological_order_visits_children_first() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let b = arena.alloc(Node::Scalar(2.0)); + let add = arena.alloc(Node::Add(a, b)); + + let order = arena.topological_order(add); + + let pos_a = order.iter().position(|&n| n == a).unwrap(); + let pos_b = order.iter().position(|&n| n == b).unwrap(); + let pos_add = order.iter().position(|&n| n == add).unwrap(); + assert!(pos_a < pos_add); + assert!(pos_b < pos_add); + assert_eq!(order.len(), 3); + } + + #[test] + fn test_topological_order_deduplicates_shared_subexpression() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let neg_a = arena.alloc(Node::Neg(a)); + let combined = arena.alloc(Node::Add(neg_a, neg_a)); + + let order = arena.topological_order(combined); + assert_eq!(order.len(), 3); + assert_eq!(order.iter().filter(|&&n| n == a).count(), 1); + assert_eq!(order.iter().filter(|&&n| n == neg_a).count(), 1); + } + + #[test] + fn test_topological_order_skips_unreachable_nodes() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let _unreachable = arena.alloc(Node::Scalar(99.0)); + let neg_a = arena.alloc(Node::Neg(a)); + + let order = arena.topological_order(neg_a); + assert_eq!(order.len(), 2); + assert!(order.contains(&a)); + assert!(order.contains(&neg_a)); + } + + #[test] + fn read_extents_cover_only_what_the_root_reaches() { + let mut arena = Arena::new(); + let wide = arena.alloc(Node::StateVector { start: 0, end: 9 }); + let sv = arena.alloc(Node::StateVector { start: 2, end: 5 }); + let svd = arena.alloc(Node::StateVectorDot { start: 0, end: 1 }); + let p = arena.alloc(Node::InputParameter { + name: "p".to_string(), + index: 0, + offset: 1, + width: 2, + }); + let product = arena.alloc(Node::Mul(svd, p)); + let sum = arena.alloc(Node::Add(sv, product)); + let _unreached = arena.alloc(Node::Neg(wide)); + let extents = ReadExtents::scan(&arena, sum); + assert_eq!( + extents, + ReadExtents { + y: 5, + y_dot: Some(1), + inputs: 3 + } + ); + assert_eq!(extents.n_states(), 5); + + let empty_dot = arena.alloc(Node::StateVectorDot { start: 0, end: 0 }); + let extents = ReadExtents::scan(&arena, empty_dot); + assert_eq!( + extents.y_dot, + Some(0), + "an empty y_dot read still reads y_dot" + ); + assert_eq!(extents.n_states(), 0); + + let lone = arena.alloc(Node::Scalar(1.0)); + assert_eq!(ReadExtents::scan(&arena, lone), ReadExtents::default()); + } + + #[test] + fn a_state_count_past_the_tapes_u32_indices_is_refused() { + let max = u32::MAX as usize; + assert_eq!(resolve_state_count(Some(max), 1).ok(), Some(max)); + // Unrefused, a Jacobian sized by this count aborts the process on allocation. + let too_large = resolve_state_count(Some(1 << 45), 1); + assert!(matches!( + too_large, + Err(CoreError::StateCountTooLarge { declared, max: u32::MAX }) if declared == 1 << 45 + )); + assert!(resolve_state_count(Some(max + 1), 1).is_err()); + } + + #[test] + fn test_node_map_basic_get_insert() { + let mut map: NodeMap = NodeMap::new(4); + let id0: NodeId = 0u32.into(); + let id3: NodeId = 3u32.into(); + + assert!(!map.contains_key(id0)); + assert_eq!(map.get(id0), None); + + assert_eq!(map.insert(id0, 42), None); + assert_eq!(map.get(id0), Some(&42)); + assert!(map.contains_key(id0)); + + assert_eq!(map.insert(id0, 99), Some(42)); + assert_eq!(map.get(id0), Some(&99)); + + assert!(!map.contains_key(id3)); + map.insert(id3, 7); + assert_eq!(map.get(id3), Some(&7)); + } + + #[test] + #[should_panic(expected = "node 1 names child 1, which is not allocated before it")] + fn alloc_refuses_a_child_not_yet_allocated() { + let mut arena = Arena::new(); + let _ = arena.alloc(Node::Scalar(1.0)); + let _ = arena.alloc(Node::Neg(NodeId::from(1u32))); + } + + #[test] + #[should_panic(expected = "a DAG's root must be a node of its arena")] + fn a_rewrite_cannot_leave_the_root_outside_its_arena() { + let mut arena = Arena::new(); + let root = arena.alloc(Node::Time); + let mut dag = Dag::new(arena, root); + dag.rewrite_in_place(|arena, root| { + *arena = Arena::new(); + root + }); + } + + #[test] + #[should_panic(expected = "a DAG's root must be a node of its arena")] + fn a_dag_refuses_a_root_outside_its_arena() { + let mut arena = Arena::new(); + let _ = arena.alloc(Node::Scalar(1.0)); + let _ = DagRef::new(&arena, NodeId::from(1u32)); + } +} + +#[cfg(all(test, feature = "serialize"))] +mod persisted_tests { + use super::*; + + #[test] + fn an_arena_persists_as_its_bare_node_list() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::Time); + let root = arena.alloc(Node::Neg(x)); + let bytes = bincode::serialize(&Dag::new(arena.clone(), root)).expect("serialise"); + let bare = bincode::serialize(&(arena.nodes(), root.raw())).expect("serialise"); + assert_eq!(bytes, bare); + + let back: Dag = bincode::deserialize(&bytes).expect("deserialise"); + assert_eq!(back.arena().nodes(), arena.nodes()); + assert_eq!(back.root(), root); + } + + #[test] + fn a_persisted_child_at_or_after_its_parent_is_refused() { + let ahead = vec![Node::Neg(NodeId::from(1u32)), Node::Time]; + let looped = vec![Node::Time, Node::Neg(NodeId::from(1u32))]; + for nodes in [ahead, looped] { + let bytes = bincode::serialize(&nodes).expect("serialise"); + let err = bincode::deserialize::(&bytes).expect_err("must be refused"); + assert!(err.to_string().contains("which is not allocated before it")); + } + } + + #[test] + fn a_persisted_root_outside_its_arena_is_refused() { + let bytes = bincode::serialize(&(vec![Node::Time], 1u32)).expect("serialise"); + let err = bincode::deserialize::(&bytes).expect_err("must be refused"); + assert!(err.to_string().contains("root 1 is out of range")); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/branch_regions.rs b/packages/pybamm-rust/pybamm-core/src/branch_regions.rs new file mode 100644 index 0000000000..01a8b516ad --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/branch_regions.rs @@ -0,0 +1,945 @@ +//! Conditional branch region analysis: which nodes each `Conditional` branch owns exclusively. +//! +//! Ambiguous cases degrade to [`Ownership::Common`], costing speed but never correctness. + +use crate::arena::{Arena, Dag, DagRef, NodeId}; +use crate::node::Node; + +/// One branch of one conditional: `branches[index]` of `cond`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BranchLabel { + pub cond: NodeId, + pub index: usize, +} + +/// Which conditional branches exclusively need a node's value. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Ownership { + /// No consumer seen (unreachable from the root). + Unreached, + /// Needed only by the listed branches of one conditional (sorted, deduplicated). + Branches { cond: NodeId, indices: Vec }, + /// Needed regardless of which branch runs. + Common, +} + +impl Ownership { + /// Join in the lattice `Unreached < Branches < Common`; different conditionals give `Common`. + fn join(&mut self, other: &Self) { + match (&mut *self, other) { + (_, Self::Unreached) | (Self::Common, _) => {}, + (Self::Unreached, _) => *self = other.clone(), + (_, Self::Common) => *self = Self::Common, + ( + Self::Branches { cond, indices }, + Self::Branches { + cond: other_cond, + indices: other_indices, + }, + ) => { + if cond == other_cond { + for &i in other_indices { + if let Err(pos) = indices.binary_search(&i) { + indices.insert(pos, i); + } + } + } else { + *self = Self::Common; + } + }, + } + } +} + +/// The single branch that exclusively owns this node, if exactly one does. +#[must_use] +fn sole_owner(ownership: &Ownership) -> Option { + match ownership { + Ownership::Branches { cond, indices } if indices.len() == 1 => Some(BranchLabel { + cond: *cond, + index: indices[0], + }), + _ => None, + } +} + +/// Cheap over-approximate check for any conditional that could become blocks. +#[must_use] +fn has_multi_branch_conditional(arena: &Arena) -> bool { + arena + .nodes() + .iter() + .any(|node| matches!(node, Node::Conditional { branches, .. } if branches.len() >= 2)) +} + +/// Whether `cond`'s branches may become instruction blocks (it is `Common`, i.e. top-level). +/// +/// A deeply nested conditional degraded to `Common` also qualifies, but its runs are empty. +#[must_use] +fn is_conditional_blockable(arena: &Arena, owners: &[Ownership], cond: NodeId) -> bool { + matches!(arena.get(cond), Node::Conditional { branches, .. } if branches.len() >= 2) + && owners[cond.index()] == Ownership::Common +} + +/// Label every node in `eval_order` with the branches that exclusively need it. +/// +/// Indexed by `NodeId::index()`; nodes outside `eval_order` stay [`Ownership::Unreached`]. +#[must_use] +pub fn owner_sets(arena: &Arena, eval_order: &[NodeId]) -> Vec { + let mut owners = vec![Ownership::Unreached; arena.len()]; + if let Some(&root) = eval_order.last() { + owners[root.index()] = Ownership::Common; + } + + for &id in eval_order.iter().rev() { + let label = owners[id.index()].clone(); + match arena.get(id) { + Node::Conditional { selector, branches } if branches.len() >= 2 => { + owners[selector.index()].join(&Ownership::Common); + for (i, &b) in branches.iter().enumerate() { + let edge = Ownership::Branches { + cond: id, + indices: vec![i], + }; + owners[b.index()].join(&edge); + } + }, + node => node.for_each_child(|c| owners[c.index()].join(&label)), + } + } + + close_full_branch_coverage(arena, &mut owners); + force_nested_conditionals_common(arena, eval_order, &mut owners); + owners +} + +/// Degrade a node to `Common` once its `Branches` label spans every branch. +fn close_full_branch_coverage(arena: &Arena, owners: &mut [Ownership]) { + for owner in owners.iter_mut() { + if let Ownership::Branches { cond, indices } = owner + && let Node::Conditional { branches, .. } = arena.get(*cond) + && indices.len() >= branches.len() + { + *owner = Ownership::Common; + } + } +} + +/// Degrade the cone of every non-top-level conditional to `Common`. +/// +/// Keeps blocks flat and disjoint, which the evaluator's block skipping relies on. +fn force_nested_conditionals_common( + arena: &Arena, + eval_order: &[NodeId], + owners: &mut [Ownership], +) { + let mut nested = vec![false; arena.len()]; + let mut any_nested = false; + for &id in eval_order { + if matches!(arena.get(id), Node::Conditional { branches, .. } if branches.len() >= 2) + && owners[id.index()] != Ownership::Common + { + nested[id.index()] = true; + any_nested = true; + } + } + if !any_nested { + return; + } + for &id in eval_order { + if let Ownership::Branches { cond, .. } = &owners[id.index()] + && nested[cond.index()] + { + owners[id.index()] = Ownership::Common; + } + } +} + +/// Ceiling on the privatised arena's node count, as a multiple of the input's. +/// +/// Chain-nested cones grow the arena quadratically in branch count; this caps them. +const CLONE_BUDGET_MULTIPLE: usize = 4; + +/// Clone subgraphs shared by a strict subset of one conditional's branches, one copy per branch. +/// +/// Returns `None` when nothing needs cloning or the clone budget is exceeded. +/// Runs inside lowering; no hash-consing pass may follow it, or `cse` would merge the +/// per-branch clones back. +#[must_use] +pub fn privatise_conditionals(dag: DagRef<'_>) -> Option { + let (arena, root) = (dag.arena(), dag.root()); + if !has_multi_branch_conditional(arena) { + return None; + } + let eval_order = arena.topological_order(root); + let owners = owner_sets(arena, &eval_order); + + // Privatisable iff owned by 2..n of one blockable conditional's branches; a label covering + // all n was already degraded to `Common`. + let privatisable = |id: NodeId| match &owners[id.index()] { + Ownership::Branches { cond, indices } + if indices.len() >= 2 && is_conditional_blockable(arena, &owners, *cond) => + { + Some(indices.as_slice()) + }, + _ => None, + }; + + let mut projected_nodes = 0_usize; + let mut any_privatisable = false; + for &id in &eval_order { + match privatisable(id) { + Some(indices) => { + any_privatisable = true; + projected_nodes += indices.len(); + }, + None => projected_nodes += 1, + } + } + if !any_privatisable || projected_nodes > CLONE_BUDGET_MULTIPLE * arena.len() { + return None; + } + + let mut new_arena = Arena::new(); + let mut common: Vec> = vec![None; arena.len()]; + let mut per_branch: Vec> = vec![Vec::new(); arena.len()]; + + for &id in &eval_order { + let node = arena.get(id); + match privatisable(id) { + None => { + let remapped = + remap_children(node, owner_branch_of(&owners, id), &common, &per_branch); + common[id.index()] = Some(new_arena.alloc(remapped)); + }, + Some(indices) => { + for &i in indices { + let remapped = remap_children(node, Some(i), &common, &per_branch); + let clone = new_arena.alloc(remapped); + per_branch[id.index()].push((i, clone)); + } + }, + } + } + + let new_root = common[root.index()].expect("the root is never privatised"); + Some(Dag::new(new_arena, new_root)) +} + +/// Rebuild `node` against the privatised copies, reading each child from `branch`'s copy. +/// +/// A multi-branch `Conditional` resolves per slot, since `cse` can alias two branch slots +/// onto one node; a single-branch one is owned like any node, as in [`owner_sets`]. +fn remap_children( + node: &Node, + branch: Option, + common: &[Option], + per_branch: &[Vec<(usize, NodeId)>], +) -> Node { + match node { + Node::Conditional { selector, branches } if branches.len() >= 2 => Node::Conditional { + selector: resolve_child(common, per_branch, *selector, branch), + branches: branches + .iter() + .enumerate() + .map(|(i, &b)| resolve_child(common, per_branch, b, Some(i))) + .collect(), + }, + other => other.map_children(|c| resolve_child(common, per_branch, c, branch)), + } +} + +fn owner_branch_of(owners: &[Ownership], id: NodeId) -> Option { + sole_owner(&owners[id.index()]).map(|label| label.index) +} + +/// One conditional's branch nodes, as a contiguous range of the emission order. +#[derive(Clone, Debug)] +pub struct RegionGroup { + pub cond: NodeId, + /// Index in [`RegionSchedule::order`] where this group's block nodes start. + pub anchor: usize, + /// Node count per branch; branches sit back to back from `anchor`. + pub branch_lens: Vec, +} + +/// A block-contiguous emission order plus the groups to wrap in a `Dispatch`. +#[derive(Clone, Debug)] +pub struct RegionSchedule { + /// Emission order for every reachable node. + pub order: Vec, + pub groups: Vec, +} + +/// [`owner_sets`] over `eval_order`, or `None` when no conditional can become blocks. +/// +/// Computed once per lowering and shared by every schedule of that arena. +#[must_use] +pub fn region_owners(arena: &Arena, eval_order: &[NodeId]) -> Option> { + has_multi_branch_conditional(arena).then(|| owner_sets(arena, eval_order)) +} + +/// Reorder topological `base_order` so each blockable conditional's branch-owned nodes are +/// contiguous, grouped by branch, and immediately precede the `Conditional`. +#[must_use] +pub fn schedule_regions(arena: &Arena, base_order: &[NodeId]) -> RegionSchedule { + let owners = region_owners(arena, base_order); + schedule_regions_partitioned(arena, base_order, base_order, owners.as_deref()) +} + +/// [`schedule_regions`] for `partition`, an ordered subset of `base_order`, whose +/// [`region_owners`] are `owners`. +/// +/// A group whose `Conditional` is outside `partition` anchors at the end, and only +/// forms when `partition` holds the selector. +#[must_use] +pub fn schedule_regions_partitioned( + arena: &Arena, + partition: &[NodeId], + base_order: &[NodeId], + owners: Option<&[Ownership]>, +) -> RegionSchedule { + let unscheduled = || RegionSchedule { + order: partition.to_vec(), + groups: Vec::new(), + }; + let Some(owners) = owners else { + return unscheduled(); + }; + let mut in_partition = vec![false; arena.len()]; + for &id in partition { + in_partition[id.index()] = true; + } + + // One run per blockable conditional that can group here, keyed by its node. + let mut run_of: Vec> = vec![None; arena.len()]; + let mut runs: Vec<(NodeId, Vec>)> = Vec::new(); + for &id in base_order { + if let Node::Conditional { selector, branches } = arena.get(id) + && is_conditional_blockable(arena, owners, id) + && (in_partition[id.index()] || in_partition[selector.index()]) + { + run_of[id.index()] = Some(runs.len()); + runs.push((id, vec![Vec::new(); branches.len()])); + } + } + if runs.is_empty() { + return unscheduled(); + } + let owning_run = |id: NodeId| { + let label = sole_owner(&owners[id.index()])?; + run_of[label.cond.index()].map(|run| (run, label.index)) + }; + for &id in partition { + if let Some((run, branch)) = owning_run(id) { + runs[run].1[branch].push(id); + } + } + + let mut order = Vec::with_capacity(partition.len()); + let mut groups = Vec::with_capacity(runs.len()); + let mut pending = vec![true; runs.len()]; + + for &id in partition { + if owning_run(id).is_some() { + continue; + } + if let Some(pos) = run_of[id.index()] { + let branch_runs = std::mem::take(&mut runs[pos].1); + flush_group(&mut order, &mut groups, id, &branch_runs); + pending[pos] = false; + } + order.push(id); + } + for (pos, run) in runs.iter_mut().enumerate() { + if pending[pos] { + let branch_runs = std::mem::take(&mut run.1); + flush_group(&mut order, &mut groups, run.0, &branch_runs); + } + } + + // Release-mode: a dropped node lowers to silently wrong values, not a crash. + assert_eq!( + order.len(), + partition.len(), + "region scheduling dropped or duplicated a node" + ); + RegionSchedule { order, groups } +} + +/// Append one group's branch runs to `order`; an empty group is dropped to keep anchors distinct. +fn flush_group( + order: &mut Vec, + groups: &mut Vec, + cond: NodeId, + branch_runs: &[Vec], +) { + if branch_runs.iter().all(Vec::is_empty) { + return; + } + let anchor = order.len(); + let branch_lens = branch_runs.iter().map(Vec::len).collect(); + for run in branch_runs { + order.extend_from_slice(run); + } + groups.push(RegionGroup { + cond, + anchor, + branch_lens, + }); +} + +/// The active branch for `selector` (1-based, round-to-nearest), or `None` when nothing matches. +/// +/// Shared by every evaluator so their branch choices cannot drift apart. +#[inline] +#[must_use] +pub fn active_branch(selector: f64, n_branches: usize) -> Option { + for i in 0..n_branches { + let idx = (i + 1) as f64; + if selector > idx - 0.5 && selector < idx + 0.5 { + return Some(i); + } + } + None +} + +/// Resolve `child` to its `branch` copy when one was made, otherwise the common copy. +fn resolve_child( + common: &[Option], + per_branch: &[Vec<(usize, NodeId)>], + child: NodeId, + branch: Option, +) -> NodeId { + if let Some(b) = branch + && let Some(&(_, id)) = per_branch[child.index()].iter().find(|&&(i, _)| i == b) + { + return id; + } + common[child.index()].unwrap_or_else(|| { + panic!("child {child:?} has no copy reachable from its consumer's branch") + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::eval::CompiledExpr; + use crate::ir::{TypedIr, UnaryOp}; + use crate::node::Node; + + /// Assert `before` and `after` agree bitwise at `y = 0.3` for every selector value. + #[track_caller] + fn assert_same_for_selectors(before: &CompiledExpr, after: &CompiledExpr, selectors: &[f64]) { + for &sel_val in selectors { + let mut s1 = vec![0.0; before.scratch_len()]; + let mut s2 = vec![0.0; after.scratch_len()]; + let a = before.eval(&mut s1, 0.0, &[0.3], &[], &[sel_val]); + let b = after.eval(&mut s2, 0.0, &[0.3], &[], &[sel_val]); + assert_eq!(a[0].to_bits(), b[0].to_bits(), "selector {sel_val}"); + } + } + + #[test] + fn exclusive_branch_cones_are_labelled_per_branch() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Cos(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + + let order = arena.topological_order(cond); + let owners = owner_sets(&arena, &order); + + assert_eq!( + sole_owner(&owners[b1.index()]), + Some(BranchLabel { cond, index: 0 }) + ); + assert_eq!( + sole_owner(&owners[b2.index()]), + Some(BranchLabel { cond, index: 1 }) + ); + assert_eq!(owners[y.index()], Ownership::Common); + assert_eq!(owners[sel.index()], Ownership::Common); + assert_eq!(owners[cond.index()], Ownership::Common); + } + + #[test] + fn strict_subset_sharing_records_every_sharing_branch() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::Scalar(1.0)); + let shared = arena.alloc(Node::Exp(y)); + let b1 = arena.alloc(Node::Sin(shared)); + let b2 = arena.alloc(Node::Cos(shared)); + let b3 = arena.alloc(Node::Neg(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2, b3], + }); + + let order = arena.topological_order(cond); + let owners = owner_sets(&arena, &order); + + assert_eq!( + owners[shared.index()], + Ownership::Branches { + cond, + indices: vec![0, 1] + } + ); + assert_eq!(sole_owner(&owners[shared.index()]), None); + } + + #[test] + fn nodes_used_outside_the_conditional_stay_common() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::Scalar(1.0)); + let shared = arena.alloc(Node::Exp(y)); + let b1 = arena.alloc(Node::Sin(shared)); + let b2 = arena.alloc(Node::Cos(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + let root = arena.alloc(Node::Add(cond, shared)); + + let order = arena.topological_order(root); + let owners = owner_sets(&arena, &order); + + assert_eq!(owners[shared.index()], Ownership::Common); + assert_eq!( + sole_owner(&owners[b1.index()]), + Some(BranchLabel { cond, index: 0 }) + ); + } + + #[test] + fn selector_cone_is_common_even_when_otherwise_exclusive() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::Floor(y)); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Cos(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + + let order = arena.topological_order(cond); + let owners = owner_sets(&arena, &order); + + assert_eq!(owners[sel.index()], Ownership::Common); + } + + #[test] + fn nested_conditional_cone_is_forced_common() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let outer_sel = arena.alloc(Node::Scalar(1.0)); + let inner_sel = arena.alloc(Node::Scalar(2.0)); + let inner_b1 = arena.alloc(Node::Sin(y)); + let inner_b2 = arena.alloc(Node::Cos(y)); + let inner = arena.alloc(Node::Conditional { + selector: inner_sel, + branches: vec![inner_b1, inner_b2], + }); + let outer_b2 = arena.alloc(Node::Neg(y)); + let outer = arena.alloc(Node::Conditional { + selector: outer_sel, + branches: vec![inner, outer_b2], + }); + + let order = arena.topological_order(outer); + let owners = owner_sets(&arena, &order); + + assert!(is_conditional_blockable(&arena, &owners, outer)); + assert!(!is_conditional_blockable(&arena, &owners, inner)); + assert_eq!(owners[inner_b1.index()], Ownership::Common); + assert_eq!(owners[inner_b2.index()], Ownership::Common); + assert_eq!( + sole_owner(&owners[inner.index()]), + Some(BranchLabel { + cond: outer, + index: 0 + }) + ); + } + + #[test] + fn sharing_across_two_conditionals_degrades_to_common() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::Scalar(1.0)); + let shared = arena.alloc(Node::Exp(y)); + let a1 = arena.alloc(Node::Sin(shared)); + let a2 = arena.alloc(Node::Neg(y)); + let c1 = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![a1, a2], + }); + let b1 = arena.alloc(Node::Cos(shared)); + let b2 = arena.alloc(Node::Abs(y)); + let c2 = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + let root = arena.alloc(Node::Add(c1, c2)); + + let order = arena.topological_order(root); + let owners = owner_sets(&arena, &order); + + assert_eq!(owners[shared.index()], Ownership::Common); + } + + #[test] + fn plain_expression_is_all_common() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let root = arena.alloc(Node::Sin(y)); + let order = arena.topological_order(root); + let owners = owner_sets(&arena, &order); + assert!(order.iter().all(|n| owners[n.index()] == Ownership::Common)); + } + + #[test] + fn privatise_is_a_no_op_without_strict_subset_sharing() { + let mut plain = Arena::new(); + let y = plain.alloc(Node::StateVector { start: 0, end: 1 }); + let root = plain.alloc(Node::Sin(y)); + assert!(privatise_conditionals(DagRef::new(&plain, root)).is_none()); + + let mut all = Arena::new(); + let y = all.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = all.alloc(Node::Scalar(1.0)); + let shared = all.alloc(Node::Exp(y)); + let b1 = all.alloc(Node::Sin(shared)); + let b2 = all.alloc(Node::Cos(shared)); + let cond = all.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + assert!(privatise_conditionals(DagRef::new(&all, cond)).is_none()); + } + + #[test] + fn privatise_clones_a_strict_subset_cone_per_branch() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::Scalar(1.0)); + let inner = arena.alloc(Node::Exp(y)); + let shared = arena.alloc(Node::Sqrt(inner)); + let b1 = arena.alloc(Node::Sin(shared)); + let b2 = arena.alloc(Node::Cos(shared)); + let b3 = arena.alloc(Node::Neg(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2, b3], + }); + + let private = privatise_conditionals(DagRef::new(&arena, cond)) + .expect("strict subset must privatise"); + assert!(private.arena().len() <= CLONE_BUDGET_MULTIPLE * arena.len()); + + let order = private.arena().topological_order(private.root()); + let n_sqrt = order + .iter() + .filter(|&&id| matches!(private.arena().get(id), Node::Sqrt(_))) + .count(); + let n_exp = order + .iter() + .filter(|&&id| matches!(private.arena().get(id), Node::Exp(_))) + .count(); + assert_eq!((n_sqrt, n_exp), (2, 2)); + + let owners = owner_sets(private.arena(), &order); + for &id in &order { + if matches!(private.arena().get(id), Node::Sqrt(_) | Node::Exp(_)) { + assert!( + sole_owner(&owners[id.index()]).is_some(), + "cone node {id:?} is still shared after privatisation" + ); + } + } + } + + #[test] + fn privatise_preserves_values_for_every_selector() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let shared = arena.alloc(Node::Exp(y)); + let b1 = arena.alloc(Node::Sin(shared)); + let b2 = arena.alloc(Node::Cos(shared)); + let b3 = arena.alloc(Node::Neg(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2, b3], + }); + + let private = privatise_conditionals(DagRef::new(&arena, cond)).expect("privatises"); + let before = CompiledExpr::from_ir(TypedIr::from_dag_raw(DagRef::new(&arena, cond))); + let after = CompiledExpr::from_ir(TypedIr::from_dag_raw(private.as_ref())); + let selectors = [0.0_f64, 1.0, 2.0, 3.0, 4.0, 0.5, 1.5, f64::NAN]; + assert_same_for_selectors(&before, &after, &selectors); + } + + /// Regression: `cse` aliasing two branch slots onto one node used to panic. + #[test] + fn aliased_branch_slot_privatises_per_position() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let shared = arena.alloc(Node::Scalar(1.0)); + let other = arena.alloc(Node::Sin(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![shared, other, shared], + }); + + let order = arena.topological_order(cond); + let owners = owner_sets(&arena, &order); + assert_eq!( + owners[shared.index()], + Ownership::Branches { + cond, + indices: vec![0, 2] + }, + "an aliased branches slot is owned by exactly the positions that name it" + ); + + let before = CompiledExpr::from_ir(TypedIr::from_dag_raw(DagRef::new(&arena, cond))); + let private = + privatise_conditionals(DagRef::new(&arena, cond)).expect("the aliased slot privatises"); + let Node::Conditional { branches, .. } = private.arena().get(private.root()) else { + panic!("the root is still the conditional") + }; + assert_ne!( + branches[0], branches[2], + "each aliased position must resolve to its own clone" + ); + let after = CompiledExpr::from_ir(TypedIr::from_dag_raw(private.as_ref())); + let selectors = [0.0_f64, 1.0, 2.0, 3.0, 4.0, 0.5, 1.5, f64::NAN]; + assert_same_for_selectors(&before, &after, &selectors); + } + + /// Regression: a single-branch conditional is no conditional to `owner_sets`, so its + /// branch must resolve through the enclosing branch's copy, not slot position 0. + #[test] + fn a_single_branch_conditional_in_a_shared_cone_privatises() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let always = arena.alloc(Node::Scalar(1.0)); + let sin = arena.alloc(Node::Sin(y)); + let single = arena.alloc(Node::Conditional { + selector: always, + branches: vec![sin], + }); + let b0 = arena.alloc(Node::Neg(y)); + let b1 = arena.alloc(Node::Exp(single)); + let b2 = arena.alloc(Node::Cos(single)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b0, b1, b2], + }); + + let dag = DagRef::new(&arena, cond); + assert!( + privatise_conditionals(dag).is_some(), + "the single-branch cone is shared by branches 1 and 2" + ); + let raw = CompiledExpr::from_ir(TypedIr::from_dag_raw(dag)); + let private = CompiledExpr::new(dag); + assert_same_for_selectors(&raw, &private, &[1.0, 2.0, 3.0, 0.5, f64::NAN]); + } + + /// `cond(sel, [tanh(y), tanh²(y), ..., tanhⁿ(y)])`: each cone contains the previous. + fn chain_nested_conditional(n: usize) -> (Arena, NodeId) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let mut node = y; + let mut branches = Vec::with_capacity(n); + for _ in 0..n { + node = arena.alloc(Node::Tanh(node)); + branches.push(node); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches, + }); + (arena, cond) + } + + fn expected_chain(y: f64, n: usize, selector: f64) -> f64 { + active_branch(selector, n).map_or(0.0, |i| { + let mut v = y; + for _ in 0..=i { + v = UnaryOp::Tanh.apply(v); + } + v + }) + } + + #[test] + fn a_chain_nested_conditional_bails_out_of_the_clone_budget() { + const N: usize = 32; + let (arena, cond) = chain_nested_conditional(N); + assert!( + privatise_conditionals(DagRef::new(&arena, cond)).is_none(), + "the projected clone count must exceed the budget" + ); + + let ir = TypedIr::from_dag(DagRef::new(&arena, cond)); + let budget = CLONE_BUDGET_MULTIPLE * arena.len(); + assert!( + ir.instructions().len() <= budget, + "raw tape of {} instructions exceeds the {budget}-instruction budget", + ir.instructions().len() + ); + + let expr = CompiledExpr::from_ir(ir); + let mut scratch = vec![0.0; expr.scratch_len()]; + let y0 = 0.4_f64; + for sel_val in [ + 0.0_f64, + 1.0, + 2.0, + 17.0, + 32.0, + 33.0, + 0.5, + 1.5, + -1.0, + f64::NAN, + f64::INFINITY, + ] { + let got = expr.eval(&mut scratch, 0.0, &[y0], &[], &[sel_val])[0]; + let want = expected_chain(y0, N, sel_val); + assert_eq!(got.to_bits(), want.to_bits(), "selector {sel_val}"); + } + } + + #[test] + fn a_short_chain_still_privatises_within_the_budget() { + const N: usize = 4; + let (arena, cond) = chain_nested_conditional(N); + let private = privatise_conditionals(DagRef::new(&arena, cond)) + .expect("a 4-branch chain fits the budget"); + assert!( + private.arena().len() > arena.len(), + "cones must have been cloned" + ); + + let expr = CompiledExpr::new(DagRef::new(&arena, cond)); + let mut scratch = vec![0.0; expr.scratch_len()]; + let y0 = 0.4_f64; + for sel_val in [0.0_f64, 1.0, 2.0, 3.0, 4.0, 5.0, 0.5, 1.5, f64::NAN] { + let got = expr.eval(&mut scratch, 0.0, &[y0], &[], &[sel_val])[0]; + let want = expected_chain(y0, N, sel_val); + assert_eq!(got.to_bits(), want.to_bits(), "selector {sel_val}"); + } + } + + #[test] + fn aliased_and_disjoint_shared_cones_privatise_independently() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let x = arena.alloc(Node::Scalar(1.0)); + let inner = arena.alloc(Node::Exp(y)); + let q = arena.alloc(Node::Sqrt(inner)); + let b1 = arena.alloc(Node::Sin(q)); + let b3 = arena.alloc(Node::Cos(q)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![x, b1, x, b3], + }); + + let order = arena.topological_order(cond); + let owners = owner_sets(&arena, &order); + assert_eq!( + owners[x.index()], + Ownership::Branches { + cond, + indices: vec![0, 2] + }, + "the aliased slot is owned by exactly the positions that name it" + ); + assert_eq!( + owners[q.index()], + Ownership::Branches { + cond, + indices: vec![1, 3] + }, + "a disjoint shared cone privatises on its own index-set" + ); + + let private = privatise_conditionals(DagRef::new(&arena, cond)) + .expect("the disjoint shared cone must privatise"); + + let new_order = private.arena().topological_order(private.root()); + let n_sqrt = new_order + .iter() + .filter(|&&id| matches!(private.arena().get(id), Node::Sqrt(_))) + .count(); + let n_exp = new_order + .iter() + .filter(|&&id| matches!(private.arena().get(id), Node::Exp(_))) + .count(); + assert_eq!( + (n_sqrt, n_exp), + (2, 2), + "q and inner must be cloned once per sharing branch" + ); + let n_aliased_scalar = new_order + .iter() + .filter(|&&id| matches!(private.arena().get(id), Node::Scalar(v) if v.to_bits() == 1.0_f64.to_bits())) + .count(); + assert_eq!( + n_aliased_scalar, 2, + "the aliased slot is cloned once per position that names it" + ); + + let before = CompiledExpr::from_ir(TypedIr::from_dag_raw(DagRef::new(&arena, cond))); + let after = CompiledExpr::from_ir(TypedIr::from_dag_raw(private.as_ref())); + let selectors = [0.0_f64, 1.0, 2.0, 3.0, 4.0, 5.0, 0.5, 1.5, f64::NAN]; + assert_same_for_selectors(&before, &after, &selectors); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/error.rs b/packages/pybamm-rust/pybamm-core/src/error.rs new file mode 100644 index 0000000000..c2304811ed --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/error.rs @@ -0,0 +1,112 @@ +//! Crate-wide error type for failures driven by caller-supplied data. +//! +//! Panics remain reserved for broken internal invariants. + +/// Errors returned by fallible `pybamm-core` entry points. +#[derive(Debug, thiserror::Error)] +pub enum CoreError { + /// A half-open `[start, end)` node extent ran backwards. + #[error("{node}: start ({start}) must not exceed end ({end})")] + InvertedExtent { + node: &'static str, + start: usize, + end: usize, + }, + + /// A deserialised node named a child at or after itself, which could cycle or overrun. + #[error("node {node} names child {child}, which is not allocated before it")] + ChildNotBeforeParent { node: usize, child: usize }, + + #[error("root {root} is out of range for an arena of {len} nodes")] + RootOutOfRange { root: usize, len: usize }, + + #[error("Conditional requires at least one branch")] + EmptyConditional, + + #[error("InterpolantNd: expected {axes} children (one per axis), got {children}")] + InterpolantNdArity { axes: usize, children: usize }, + + /// The declared state count was below the extent the expression reads. + #[error("n_states={declared} is below the expression's state extent {extent}")] + StateCountTooSmall { declared: usize, extent: usize }, + + /// The declared state count was past the `u32` range tape state indices are stored in. + #[error("n_states={declared} exceeds the largest state count a tape can index, {max}")] + StateCountTooLarge { declared: usize, max: u32 }, + + /// The declared packed input width was below what the graph's inputs occupy. + /// + /// The FFI borrows the declared width from the caller's pointer, so raising it would over-read. + #[error( + "n_inputs={declared} is below the packed width {extent} of the inputs the graph registers" + )] + InputWidthTooSmall { declared: usize, extent: usize }, + + /// A name was registered twice at different widths, which would shift every later offset. + #[error( + "input parameter '{name}' re-registered with width {width}, previously registered with width {registered}" + )] + ParamWidthConflict { + name: String, + width: usize, + registered: usize, + }, + + /// Two `InputParameter` nodes describe one registration index differently. + #[error( + "input parameter index {index} carries two slots: '{name}' at offset {offset} width {width}, and '{other_name}' at offset {other_offset} width {other_width}" + )] + ParamSlotConflict { + index: usize, + name: String, + offset: usize, + width: usize, + other_name: String, + other_offset: usize, + other_width: usize, + }, + + /// One name registered at two indices. + #[error("input parameter '{name}' is registered at both index {index} and index {other_index}")] + ParamNameConflict { + name: String, + index: usize, + other_index: usize, + }, + + /// Two parameters claim the same packed element. + #[error( + "input parameters '{name}' (offset {offset}, width {width}) and '{other_name}' (offset {other_offset}, width {other_width}) overlap in the packed inputs" + )] + ParamSlotOverlap { + name: String, + offset: usize, + width: usize, + other_name: String, + other_offset: usize, + other_width: usize, + }, + + /// A parameter's packed slot ran past the `u32` range packed offsets are stored in. + #[error( + "input parameter '{name}' (offset {offset}, width {width}) runs past the packed-input index range" + )] + ParamSlotRange { + name: String, + offset: usize, + width: usize, + }, + + /// A parameter node's registration index reached past its arena, which no registration yields. + #[error("input parameter index {index} is out of range for an arena of {n_nodes} nodes")] + ParamIndexRange { index: usize, n_nodes: usize }, + + #[error("invalid CSR matrix: {0}")] + Csr(String), + + #[error("invalid array: {0}")] + Array(String), + + #[error("invalid interpolant: {0}")] + Interpolant(String), +} diff --git a/packages/pybamm-rust/pybamm-core/src/eval.rs b/packages/pybamm-rust/pybamm-core/src/eval.rs new file mode 100644 index 0000000000..7b715430e5 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/eval.rs @@ -0,0 +1,1717 @@ +//! Scalar interpreter for a `TypedIr` tape, including tangent inputs for forward-mode AD. + +// The op helpers must inline as one chain into the interpreter loop. +#![allow(clippy::inline_always)] + +use crate::accumulate::{Accumulate, Product, dispatch}; +use crate::arena::DagRef; +use crate::branch_regions::active_branch; +use crate::ir::ids::{NdInterpolantIdx, SlotTableIdx}; +use crate::ir::{ + BinaryOp, BroadcastKind, ConstPool, Instruction, Slot, SplitEvalInfo, TypedIr, UnaryOp, +}; +use crate::node::{CsrData, InterpolantData, NdInterpolantData}; + +/// `match $op`, expanding `$body` in each arm with `$f` bound to that arm's function. +/// +/// Each arm binds its own closure type, so `$body` monomorphises per op. +macro_rules! dispatch_op { + ($op:expr, $f:ident => $body:expr, { $($variant:path => $func:expr,)+ }) => { + match $op { + $($variant => { + let $f = $func; + $body + },)+ + } + }; +} +pub(crate) use dispatch_op; + +/// Expand `$body` with `$f` bound to `$op`'s scalar `Fn(f64, f64) -> f64`. +/// +/// The one definition of each [`BinaryOp`]'s value, shared by every interpreter so +/// their primal results cannot drift apart. Transcendentals take the platform C library's +/// values through `f64`, as `CasADi` does, so platforms may differ in the last bits. `min`/`max` +/// take `libm`'s, as `f64`'s leave the order of signed zeros to the target. +macro_rules! with_binary_op { + ($op:expr, $f:ident => $body:expr) => { + $crate::eval::dispatch_op!($op, $f => $body, { + $crate::ir::BinaryOp::Add => |x: f64, y: f64| x + y, + $crate::ir::BinaryOp::Sub => |x: f64, y: f64| x - y, + $crate::ir::BinaryOp::Mul => |x: f64, y: f64| x * y, + $crate::ir::BinaryOp::Div => |x: f64, y: f64| x / y, + $crate::ir::BinaryOp::Pow => $crate::eval::pow, + $crate::ir::BinaryOp::Minimum => libm::fminimum_num, + $crate::ir::BinaryOp::Maximum => libm::fmaximum_num, + $crate::ir::BinaryOp::Modulo => |x: f64, y: f64| x % y, + $crate::ir::BinaryOp::Hypot => f64::hypot, + $crate::ir::BinaryOp::EqualHeaviside => { + |x: f64, y: f64| if x <= y { 1.0 } else { 0.0 } + }, + $crate::ir::BinaryOp::NotEqualHeaviside => { + |x: f64, y: f64| if x < y { 1.0 } else { 0.0 } + }, + $crate::ir::BinaryOp::Equality => |x: f64, y: f64| { + if (x - y).abs() < $crate::eval::EQUALITY_EPS { 1.0 } else { 0.0 } + }, + }) + }; +} +pub(crate) use with_binary_op; + +/// Expand `$body` with `$f` bound to `$op`'s scalar `Fn(f64) -> f64`. +/// +/// The [`UnaryOp`] counterpart of [`with_binary_op`]. +macro_rules! with_unary_op { + ($op:expr, $f:ident => $body:expr) => { + $crate::eval::dispatch_op!($op, $f => $body, { + $crate::ir::UnaryOp::Neg => |x: f64| -x, + $crate::ir::UnaryOp::Abs => f64::abs, + $crate::ir::UnaryOp::Sqrt => f64::sqrt, + $crate::ir::UnaryOp::Exp => f64::exp, + $crate::ir::UnaryOp::Log => f64::ln, + $crate::ir::UnaryOp::Sin => f64::sin, + $crate::ir::UnaryOp::Cos => f64::cos, + $crate::ir::UnaryOp::Tanh => f64::tanh, + $crate::ir::UnaryOp::Sinh => f64::sinh, + $crate::ir::UnaryOp::Cosh => f64::cosh, + // `f64::asinh` is Rust's, not the C library's, and overflows past `f64::MAX / 2`. + $crate::ir::UnaryOp::Arcsinh => libm::asinh, + $crate::ir::UnaryOp::Arctan => f64::atan, + $crate::ir::UnaryOp::Erf => $crate::eval::erf, + $crate::ir::UnaryOp::Sign => $crate::eval::sign, + $crate::ir::UnaryOp::Floor => f64::floor, + $crate::ir::UnaryOp::Ceiling => f64::ceil, + }) + }; +} +pub(crate) use with_unary_op; + +impl BinaryOp { + /// This op's value at one pair of elements, as every interpreter computes it. + #[inline] + pub(crate) fn apply(self, x: f64, y: f64) -> f64 { + with_binary_op!(self, f => f(x, y)) + } +} + +impl UnaryOp { + /// This op's value at one element, as every interpreter computes it. + #[inline] + pub(crate) fn apply(self, x: f64) -> f64 { + with_unary_op!(self, f => f(x)) + } +} + +/// The segment `[lo, lo + 1]` of `knots` (at least two) holding `x`: right-continuous at +/// an interior knot, so never a zero-width segment between repeated ones, and clamped to the +/// first or last segment outside the domain. +/// +/// The one knot search behind every 1-D and N-D interpolant. +#[inline(always)] +fn locate_segment(knots: &[f64], x: f64) -> usize { + debug_assert!(knots.len() >= 2); + let n = knots.len(); + if x <= knots[0] { + return 0; + } + if x >= knots[n - 1] { + return n - 2; + } + let mut lo = 0; + let mut hi = n - 1; + while hi - lo > 1 { + let mid = usize::midpoint(lo, hi); + if knots[mid] <= x { + lo = mid; + } else { + hi = mid; + } + } + lo +} + +/// Linear interpolation, extending the boundary segment outside the domain (no flat clamp). +/// +/// Unfused `y0 + t * (y1 - y0)`, bit for bit what `CasADi`'s `interpn_linear` computes. +#[inline(always)] +pub(crate) fn interp_linear_1d(table: &InterpolantData, x: f64) -> f64 { + let (x_data, y_data) = (&table.x_data, &table.y_data); + if x_data.len() == 1 { + return y_data[0]; + } + let lo = locate_segment(x_data, x); + let t = (x - x_data[lo]) / (x_data[lo + 1] - x_data[lo]); + y_data[lo] + t * (y_data[lo + 1] - y_data[lo]) +} + +/// Slope of the linear interpolant at `x`, from the table's pre-computed slopes. +#[inline(always)] +pub(crate) fn interp_linear_1d_deriv(table: &InterpolantData, x: f64) -> f64 { + let (x_data, slopes) = (&table.x_data, &table.slopes); + if slopes.is_empty() { + return 0.0; + } + slopes[locate_segment(x_data, x)] +} + +/// Piecewise-cubic interpolation (cubic spline / pchip), power-basis Horner. +#[inline(always)] +pub(crate) fn interp_cubic_1d(breakpoints: &[f64], coeffs: &[[f64; 4]], x: f64) -> f64 { + let i = locate_segment(breakpoints, x); + let dx = x - breakpoints[i]; + let [c0, c1, c2, c3] = coeffs[i]; + c0 + (c1 + (c2 + c3 * dx) * dx) * dx +} + +/// Derivative of the piecewise-cubic interpolant: `p'(dx) = c1 + 2*c2*dx + 3*c3*dx^2`. +#[inline(always)] +pub(crate) fn interp_cubic_1d_deriv(breakpoints: &[f64], coeffs: &[[f64; 4]], x: f64) -> f64 { + let i = locate_segment(breakpoints, x); + let dx = x - breakpoints[i]; + let [_c0, c1, c2, c3] = coeffs[i]; + c1 + (2.0 * c2 + 3.0 * c3 * dx) * dx +} + +/// An N-D interpolant's cell geometry, computed once per instruction rather than per element. +#[derive(Clone, Copy)] +pub(crate) struct NdCell { + /// Per-axis polynomial order. + order: usize, + ndim: usize, + /// Coefficient stride of each axis within a cell: `order^(ndim - 1 - axis)`. + strides: [usize; 3], +} + +impl NdCell { + pub(crate) fn of(interp: &NdInterpolantData) -> Self { + let (order, ndim) = (interp.order as usize, interp.breakpoints.len()); + debug_assert!((2..=3).contains(&ndim) && order >= 2); + let mut strides = [1; 3]; + for axis in (0..ndim - 1).rev() { + strides[axis] = strides[axis + 1] * order; + } + Self { + order, + ndim, + strides, + } + } + + pub(crate) const fn ndim(self) -> usize { + self.ndim + } + + /// The coefficients of the cell holding `coords`, with per-axis offsets written to `dxs`; + /// clamped per axis like scipy `RegularGridInterpolator(fill_value=None)`. + pub(crate) fn locate<'a>( + self, + interp: &'a NdInterpolantData, + coords: &[f64], + dxs: &mut [f64; 3], + ) -> &'a [f64] { + let mut cell = 0; + for (a, knots) in interp.breakpoints.iter().enumerate() { + let i = locate_segment(knots, coords[a]); + dxs[a] = coords[a] - knots[i]; + cell = cell * (knots.len() - 1) + i; + } + let size = self.strides[0] * self.order; + &interp.coeffs[cell * size..][..size] + } + + /// The polynomial at offsets `dxs` within the cell whose coefficients are `coeffs`. + pub(crate) fn value(self, coeffs: &[f64], dxs: &[f64; 3]) -> f64 { + tensor_horner(coeffs, &dxs[..self.ndim], &self.strides, self.order) + } + + /// Its partial derivative along `axis`. + pub(crate) fn partial(self, coeffs: &[f64], dxs: &[f64; 3], axis: usize) -> f64 { + debug_assert!(axis < self.ndim); + tensor_horner_partial(coeffs, &dxs[..self.ndim], &self.strides, self.order, axis) + } +} + +/// Buffer offset of output element `i`'s coordinate in N-D interpolant input `source`; +/// a length-1 source broadcasts over the output. +#[inline(always)] +pub(crate) const fn nd_source_offset(source: Slot, i: usize) -> usize { + source.offset_usize() + if source.is_scalar() { 0 } else { i } +} + +/// Coordinates of output element `i` of an N-D interpolant, `read` from each source. +#[inline(always)] +pub(crate) fn nd_coords(sources: &[Slot], i: usize, read: impl Fn(usize) -> f64) -> [f64; 3] { + let mut coords = [0.0; 3]; + for (coord, &source) in coords.iter_mut().zip(sources) { + *coord = read(nd_source_offset(source, i)); + } + coords +} + +/// `InterpNd` (`axis` `None`) or `InterpNdPartial` into `out`, over `k` lanes of a +/// lane-minor buffer; the scalar interpreter is `k = 1`. +pub(crate) fn interp_nd_lanes( + buf: &mut [f64], + k: usize, + consts: &ConstPool, + interp_idx: NdInterpolantIdx, + sources_idx: SlotTableIdx, + axis: Option, + out: Slot, +) { + let interp = &consts.nd_interpolants[interp_idx.index()]; + let cell = NdCell::of(interp); + let sources = &consts.interp_nd_sources[sources_idx.index()..][..cell.ndim()]; + let mut dxs = [0.0_f64; 3]; + for e in 0..out.len_usize() { + for l in 0..k { + let coords = nd_coords(sources, e, |offset| buf[offset * k + l]); + let coeffs = cell.locate(interp, &coords, &mut dxs); + buf[(out.offset_usize() + e) * k + l] = axis.map_or_else( + || cell.value(coeffs, &dxs), + |axis| cell.partial(coeffs, &dxs, axis), + ); + } + } +} + +/// Evaluate an N-D tensor-product polynomial in nested Horner form. +/// +/// `coeffs` has `order^(dxs.len())` entries, axis-0 power slowest, ascending; `strides[0]` is +/// axis 0's coefficient stride, `order^(dxs.len() - 1)`. +fn tensor_horner(coeffs: &[f64], dxs: &[f64], strides: &[usize], order: usize) -> f64 { + if dxs.len() == 1 { + let mut acc = coeffs[order - 1]; + for a in (0..order - 1).rev() { + acc = acc * dxs[0] + coeffs[a]; + } + return acc; + } + let (stride, inner_dxs, inner_strides) = (strides[0], &dxs[1..], &strides[1..]); + let slab = |a: usize| &coeffs[a * stride..(a + 1) * stride]; + let mut acc = tensor_horner(slab(order - 1), inner_dxs, inner_strides, order); + for a in (0..order - 1).rev() { + acc = acc * dxs[0] + tensor_horner(slab(a), inner_dxs, inner_strides, order); + } + acc +} + +/// Partial derivative of the N-D tensor-product polynomial along `axis`. +fn tensor_horner_partial( + coeffs: &[f64], + dxs: &[f64], + strides: &[usize], + order: usize, + axis: usize, +) -> f64 { + if axis == 0 && dxs.len() == 1 { + let mut acc = (order - 1) as f64 * coeffs[order - 1]; + for a in (1..order - 1).rev() { + acc = acc * dxs[0] + a as f64 * coeffs[a]; + } + return acc; + } + let (stride, inner_dxs, inner_strides) = (strides[0], &dxs[1..], &strides[1..]); + let slab = |a: usize| &coeffs[a * stride..(a + 1) * stride]; + if axis == 0 { + let top = tensor_horner(slab(order - 1), inner_dxs, inner_strides, order); + let mut acc = (order - 1) as f64 * top; + for a in (1..order - 1).rev() { + acc = acc * dxs[0] + a as f64 * tensor_horner(slab(a), inner_dxs, inner_strides, order); + } + return acc; + } + let partial = + |a: usize| tensor_horner_partial(slab(a), inner_dxs, inner_strides, order, axis - 1); + let mut acc = partial(order - 1); + for a in (0..order - 1).rev() { + acc = acc * dxs[0] + partial(a); + } + acc +} + +/// Tolerance `BinaryOp::Equality` compares within. +pub(crate) const EQUALITY_EPS: f64 = 1e-14; + +/// `x^y` by the platform's `pow`, except exponents `0`, `1`, `2` and `-1`, which take `1`, +/// `x`, `x*x` and `1/x`. +/// +/// Those forms round once from exact operands, so are correctly rounded where a platform +/// `pow` need not be, and equal what `simplify` lowers the same constant exponents to. The +/// model residual compiles unsimplified, so its squares reach this as `pow(x, 2)`. +#[allow(clippy::float_cmp)] +#[inline(always)] +pub(crate) fn pow(x: f64, y: f64) -> f64 { + if y == 2.0 { + x * x + } else if y == 1.0 { + x + } else if y == -1.0 { + 1.0 / x + } else if y == 0.0 { + 1.0 + } else { + x.powf(y) + } +} + +/// The error function, to within an ulp, from `libm` as `f64` has none. +/// +/// Must be accurate, as `Erf` is differentiated exactly as `2/sqrt(pi) * exp(-a^2)`. +#[inline(always)] +pub(crate) fn erf(x: f64) -> f64 { + libm::erf(x) +} + +/// Sign function with `sign(0) = 0` (unlike `f64::signum`). +#[inline(always)] +pub(crate) fn sign(x: f64) -> f64 { + if x == 0.0 { 0.0 } else { x.signum() } +} + +/// `MaxReduce`: the largest of `values` by strict `>` from `-inf`, so a NaN never wins. +#[inline(always)] +pub(crate) fn max_reduce(values: impl IntoIterator) -> f64 { + values + .into_iter() + .fold(f64::NEG_INFINITY, |best, v| if v > best { v } else { best }) +} + +/// `MinReduce`: the smallest of `values` by strict `<` from `+inf`, so a NaN never wins. +#[inline(always)] +pub(crate) fn min_reduce(values: impl IntoIterator) -> f64 { + values + .into_iter() + .fold(f64::INFINITY, |best, v| if v < best { v } else { best }) +} + +/// Position of the element [`max_reduce`] (`is_max`) or [`min_reduce`] returns, `0` if none: +/// the one `ReduceArgSelect` picks and reverse mode credits. Ties keep the earliest. +#[inline(always)] +pub(crate) fn arg_reduce(values: impl IntoIterator, is_max: bool) -> usize { + // The first non-NaN element is the incumbent, so a NaN never wins, even against + // an all `-inf` (`+inf`) remainder the reduction returns. + let mut best: Option<(usize, f64)> = None; + for (i, v) in values.into_iter().enumerate() { + let wins = match best { + None => !v.is_nan(), + Some((_, b)) if is_max => v > b, + Some((_, b)) => v < b, + }; + if wins { + best = Some((i, v)); + } + } + best.map_or(0, |(i, _)| i) +} + +/// Split `buf` into a read window `[src, src+src_len)` and a disjoint write window +/// `[dst, dst+dst_len)`. +#[inline(always)] +pub(crate) fn split_read_write( + buf: &mut [f64], + src: usize, + src_len: usize, + dst: usize, + dst_len: usize, +) -> (&[f64], &mut [f64]) { + debug_assert!( + src + src_len <= dst || dst + dst_len <= src, + "read and write windows overlap (src={src}+{src_len}, dst={dst}+{dst_len})" + ); + if dst < src { + let (left, right) = buf.split_at_mut(src); + (&right[..src_len], &mut left[dst..dst + dst_len]) + } else { + let (left, right) = buf.split_at_mut(dst); + (&left[src..src + src_len], &mut right[..dst_len]) + } +} + +/// Split `buf` into read windows `a`, `b` and a write window `dst`, each `(start, len)`. +/// +/// `dst` must be disjoint from `a` and `b`, which may overlap each other. +#[inline(always)] +pub(crate) fn split_two_read_write( + buf: &mut [f64], + a: (usize, usize), + b: (usize, usize), + (dst, dst_len): (usize, usize), +) -> (&[f64], &[f64], &mut [f64]) { + let (left, rest) = buf.split_at_mut(dst); + let (dst_s, right) = rest.split_at_mut(dst_len); + let (left, right): (&[f64], &[f64]) = (left, right); + let read = |(start, len): (usize, usize)| { + debug_assert!( + start + len <= dst || dst + dst_len <= start, + "read window {start}+{len} overlaps write window {dst}+{dst_len}" + ); + if start + len <= dst { + &left[start..start + len] + } else { + &right[start - (dst + dst_len)..][..len] + } + }; + (read(a), read(b), dst_s) +} + +/// `out = A @ x` for a CSR `A`, each row summed from `0.0` in stored order. +pub(crate) struct CsrProduct<'a> { + pub(crate) csr: &'a CsrData, + pub(crate) x: &'a [f64], + pub(crate) out: &'a mut [f64], +} + +impl Product for CsrProduct<'_> { + type Output = (); + + #[inline(always)] + fn run(self) { + for (o, (cols, vals)) in self.out.iter_mut().zip(csr_rows(self.csr)) { + let mut sum = 0.0; + for (&col, &val) in cols.iter().zip(vals) { + sum = A::mul_acc(sum, val, self.x[col]); + } + *o = sum; + } + } +} + +/// `buf[dst..] = M @ buf[vec..]` for the row-major `rows x cols` matrix at `buf[mat..]`. +struct DenseProduct<'a> { + buf: &'a mut [f64], + mat: usize, + vec: usize, + dst: usize, + rows: usize, + cols: usize, +} + +impl Product for DenseProduct<'_> { + type Output = (); + + #[inline(always)] + fn run(self) { + let Self { + buf, + mat, + vec, + dst, + rows, + cols, + } = self; + for row in 0..rows { + let mat_row = &buf[mat + row * cols..][..cols]; + let mut sum = 0.0; + for (&m, &v) in mat_row.iter().zip(&buf[vec..vec + cols]) { + sum = A::mul_acc(sum, m, v); + } + buf[dst + row] = sum; + } + } +} + +/// Each row of `csr` as its stored `(columns, values)`, in row order. +/// +/// The one row walk behind every sparse product; `CsrData::try_new` guarantees the splits fit. +#[inline(always)] +pub(crate) fn csr_rows(csr: &CsrData) -> impl Iterator { + let (mut cols, mut vals) = (&csr.indices[..], &csr.data[..]); + csr.indptr.windows(2).map(move |w| { + let (row_cols, rest_cols) = cols.split_at(w[1] - w[0]); + let (row_vals, rest_vals) = vals.split_at(w[1] - w[0]); + (cols, vals) = (rest_cols, rest_vals); + (row_cols, row_vals) + }) +} + +/// Apply a binary operation with broadcasting. +#[inline(always)] +fn broadcast_apply f64>( + buf: &mut [f64], + f: F, + a: usize, + b: usize, + dst: usize, + len: usize, + kind: BroadcastKind, +) { + match kind { + BroadcastKind::ScalarScalar => { + buf[dst] = f(buf[a], buf[b]); + }, + BroadcastKind::ScalarVector => { + let scalar = buf[a]; + let (b_s, d_s) = split_read_write(buf, b, len, dst, len); + for (o, &y) in d_s.iter_mut().zip(b_s) { + *o = f(scalar, y); + } + }, + BroadcastKind::VectorScalar => { + let scalar = buf[b]; + let (a_s, d_s) = split_read_write(buf, a, len, dst, len); + for (o, &x) in d_s.iter_mut().zip(a_s) { + *o = f(x, scalar); + } + }, + BroadcastKind::VectorVector => zip_apply(buf, f, a, b, dst, len), + } +} + +/// Write `f(a[i], b[i])` to `dst[i]` for each `i < len`, with `dst` disjoint from `a` and `b`. +#[inline(always)] +pub(crate) fn zip_apply f64>( + buf: &mut [f64], + f: F, + a: usize, + b: usize, + dst: usize, + len: usize, +) { + let (a_s, b_s, d_s) = split_two_read_write(buf, (a, len), (b, len), (dst, len)); + for ((o, &x), &y) in d_s.iter_mut().zip(a_s).zip(b_s) { + *o = f(x, y); + } +} + +#[inline(always)] +fn eval_binary_op( + buf: &mut [f64], + op: BinaryOp, + a: usize, + b: usize, + dst: usize, + len: usize, + kind: BroadcastKind, +) { + with_binary_op!(op, f => broadcast_apply(buf, f, a, b, dst, len, kind)); +} + +/// Write `f` of each element of `[src, src+len)` to the disjoint `[dst, dst+len)`. +#[inline(always)] +pub(crate) fn unary_apply f64>( + buf: &mut [f64], + f: F, + src: usize, + dst: usize, + len: usize, +) { + let (src_s, dst_s) = split_read_write(buf, src, len, dst, len); + for (o, &x) in dst_s.iter_mut().zip(src_s) { + *o = f(x); + } +} + +#[inline(always)] +fn eval_unary_op(buf: &mut [f64], op: UnaryOp, src: usize, dst: usize, len: usize) { + with_unary_op!(op, f => unary_apply(buf, f, src, dst, len)); +} + +/// Seed vectors for tangent evaluation (forward-mode AD); `None` reads as zero. +#[derive(Debug, Default)] +pub struct TangentInputs<'a> { + /// Tangent of the state vector, indexed like `y`. + pub dy: Option<&'a [f64]>, + /// Tangent of the parameters, indexed by `InputParameter::index`, not input offset. + pub dp: Option<&'a [f64]>, +} + +#[derive(Debug, Clone, Copy)] +struct PrimalEvalInputs<'a> { + t: f64, + y: &'a [f64], + y_dot: &'a [f64], + inputs: &'a [f64], +} + +#[derive(Debug, Clone, Copy)] +struct EvalContext<'a> { + primal: Option>, + tangent: Option<&'a TangentInputs<'a>>, +} + +impl<'a> EvalContext<'a> { + const fn with_primal( + t: f64, + y: &'a [f64], + y_dot: &'a [f64], + inputs: &'a [f64], + tangent: Option<&'a TangentInputs<'a>>, + ) -> Self { + Self { + primal: Some(PrimalEvalInputs { + t, + y, + y_dot, + inputs, + }), + tangent, + } + } + + const fn tangent_only(tangent: &'a TangentInputs<'a>) -> Self { + Self { + primal: None, + tangent: Some(tangent), + } + } + + const fn primal(self) -> PrimalEvalInputs<'a> { + self.primal + .expect("primal instructions require full evaluation inputs") + } + + fn tangent_state(self) -> Option<&'a [f64]> { + self.tangent.and_then(|tangent| tangent.dy) + } + + fn tangent_parameter(self, index: usize) -> f64 { + self.tangent + .and_then(|tangent| tangent.dp) + .map_or(0.0, |dp| dp[index]) + } +} + +/// Execute a slice of instructions against a shared buffer, returning how many it +/// skipped in inactive `Dispatch` blocks. +/// +/// Counting skips rather than runs keeps bookkeeping out of the per-instruction path. +#[inline] +fn eval_instructions( + buf: &mut [f64], + instructions: &[Instruction], + consts: &ConstPool, + ctx: EvalContext<'_>, +) -> usize { + let mut skipped = 0_usize; + let mut pc = 0_usize; + while pc < instructions.len() { + match instructions[pc] { + Instruction::Dispatch { + selector, + blocks_idx, + blocks_len, + span_len, + } => { + let span_end = pc + span_len.index(); + let active = active_branch(buf[selector.index()], blocks_len.index()); + let ran = active.map_or(0, |active| { + let block = consts.branch_blocks[blocks_idx.index() + active].bounds(pc); + block.len() - eval_instructions(buf, &instructions[block], consts, ctx) + }); + skipped += span_end - (pc + 1) - ran; + pc = span_end; + continue; + }, + Instruction::LoadScalar { value, dst } => { + buf[dst.index()] = value; + }, + Instruction::LoadTime { dst } => { + buf[dst.index()] = ctx.primal().t; + }, + Instruction::LoadArray { data_idx, len, dst } => { + let src = consts.get_array(data_idx, len); + buf[dst.index()..dst.index() + len.index()].copy_from_slice(src); + }, + Instruction::FillZero { dst, len } => { + buf[dst.index()..dst.index() + len.index()].fill(0.0); + }, + Instruction::LoadStateVector { start, width, dst } => { + let primal = ctx.primal(); + let (start, width) = (start.index(), width.index()); + buf[dst.index()..dst.index() + width] + .copy_from_slice(&primal.y[start..start + width]); + }, + Instruction::LoadStateVectorDot { start, width, dst } => { + let primal = ctx.primal(); + let (start, width) = (start.index(), width.index()); + buf[dst.index()..dst.index() + width] + .copy_from_slice(&primal.y_dot[start..start + width]); + }, + Instruction::LoadInputParameter { offset, width, dst } => { + let primal = ctx.primal(); + let offset = offset.index(); + let width = width.index(); + buf[dst.index()..dst.index() + width] + .copy_from_slice(&primal.inputs[offset..offset + width]); + }, + Instruction::LoadTangentState { start, width, dst } => { + let (start, width) = (start.index(), width.index()); + let slot = &mut buf[dst.index()..dst.index() + width]; + if let Some(dy) = ctx.tangent_state() { + slot.copy_from_slice(&dy[start..start + width]); + } else { + slot.fill(0.0); + } + }, + Instruction::LoadTangentParameter { index, dst } => { + buf[dst.index()] = ctx.tangent_parameter(index.index()); + }, + + Instruction::Binary { + op, + a, + b, + dst, + len, + kind, + } => { + eval_binary_op( + buf, + op, + a.index(), + b.index(), + dst.index(), + len.index(), + kind, + ); + }, + + Instruction::Unary { op, src, dst, len } => { + eval_unary_op(buf, op, src.index(), dst.index(), len.index()); + }, + + Instruction::MaxReduce { src, src_len, dst } => { + let src = src.index(); + buf[dst.index()] = max_reduce(buf[src..src + src_len.index()].iter().copied()); + }, + + Instruction::MinReduce { src, src_len, dst } => { + let src = src.index(); + buf[dst.index()] = min_reduce(buf[src..src + src_len.index()].iter().copied()); + }, + + Instruction::ReduceArgSelect { + basis_src, + picker_src, + len, + is_max, + dst, + } => { + let picker = picker_src.index(); + let best = arg_reduce(buf[picker..picker + len.index()].iter().copied(), is_max); + buf[dst.index()] = buf[basis_src.index() + best]; + }, + + Instruction::Index { + src, + start, + dst, + len, + } => { + let src = src.offset_by(start); + buf.copy_within(src..src + len.index(), dst.index()); + }, + + Instruction::Concat { + sources_idx, + sources_len, + dst, + } => { + let mut write_pos = dst.index(); + for i in 0..sources_len.index() { + let slot = consts.concat_sources[sources_idx.index() + i]; + buf.copy_within(slot.bounds(), write_pos); + write_pos += slot.len_usize(); + } + }, + + Instruction::MatMul { + csr_idx, + vec_src, + dst, + } => { + let csr = &consts.csr_data[csr_idx.index()]; + let (x, out) = split_read_write( + buf, + vec_src.index(), + csr.shape.cols, + dst.index(), + csr.shape.rows, + ); + dispatch(CsrProduct { csr, x, out }); + }, + + Instruction::DenseMatMul { + mat_src, + rows, + cols, + vec_src, + dst, + } => { + dispatch(DenseProduct { + buf, + mat: mat_src.index(), + vec: vec_src.index(), + dst: dst.index(), + rows: rows.index(), + cols: cols.index(), + }); + }, + + Instruction::Interp1DLinear { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.interpolants[interp_idx.index()]; + let f = |x| interp_linear_1d(interp, x); + unary_apply(buf, f, src.index(), dst.index(), len.index()); + }, + + Instruction::Interp1DLinearDeriv { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.interpolants[interp_idx.index()]; + let f = |x| interp_linear_1d_deriv(interp, x); + unary_apply(buf, f, src.index(), dst.index(), len.index()); + }, + + Instruction::Interp1DCubic { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.cubic_interpolants[interp_idx.index()]; + let f = |x| interp_cubic_1d(&interp.breakpoints, &interp.coeffs, x); + unary_apply(buf, f, src.index(), dst.index(), len.index()); + }, + Instruction::Interp1DCubicDeriv { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.cubic_interpolants[interp_idx.index()]; + let f = |x| interp_cubic_1d_deriv(&interp.breakpoints, &interp.coeffs, x); + unary_apply(buf, f, src.index(), dst.index(), len.index()); + }, + + Instruction::InterpNd { + interp_idx, + sources_idx, + dst, + len, + } => { + let out = Slot::new(dst, len); + interp_nd_lanes(buf, 1, consts, interp_idx, sources_idx, None, out); + }, + Instruction::InterpNdPartial { + interp_idx, + sources_idx, + axis, + dst, + len, + } => { + let (axis, out) = (Some(axis.index()), Slot::new(dst, len)); + interp_nd_lanes(buf, 1, consts, interp_idx, sources_idx, axis, out); + }, + + Instruction::Conditional { + selector, + branches_idx, + branches_len, + dst, + out_len, + } => { + let dst = dst.index(); + let out_len = out_len.index(); + match active_branch(buf[selector.index()], branches_len.index()) { + Some(i) => { + let slot = consts.branch_offsets[branches_idx.index() + i]; + buf.copy_within(slot.offset_usize()..slot.offset_usize() + out_len, dst); + }, + None => buf[dst..dst + out_len].fill(0.0), + } + }, + } + pc += 1; + } + skipped +} + +/// Zero-allocation expression evaluator. +/// +/// Callers supply the scratch buffer, so one instance can be shared across concurrent solves. +#[derive(Debug, Clone)] +pub struct CompiledExpr { + ir: TypedIr, +} + +impl CompiledExpr { + /// Compile an expression DAG into a new `CompiledExpr`. + pub fn new(dag: DagRef<'_>) -> Self { + Self::from_ir(TypedIr::from_dag(dag)) + } + + /// Wrap a pre-built `TypedIr`. + pub const fn from_ir(ir: TypedIr) -> Self { + Self { ir } + } + + /// Reference to the underlying `TypedIr`. + #[inline] + pub const fn ir(&self) -> &TypedIr { + &self.ir + } + + /// Length of the scratch buffer this expression needs for evaluation. + #[inline] + pub const fn scratch_len(&self) -> usize { + self.ir.buffer_size() + } + + /// Output length of the root expression. + #[inline] + pub const fn output_len(&self) -> usize { + self.ir.output_len() + } + + /// Evaluate the expression into `scratch` (length `>= scratch_len()`), + /// returning the slice holding the root result. + pub fn eval<'s>( + &self, + scratch: &'s mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + ) -> &'s [f64] { + self.eval_internal(scratch, t, y, y_dot, inputs, None).1 + } + + /// Evaluate with tangent inputs for forward-mode AD (JVP). + pub fn eval_with_tangent<'s>( + &self, + scratch: &'s mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + tangent: &TangentInputs<'_>, + ) -> &'s [f64] { + self.eval_internal(scratch, t, y, y_dot, inputs, Some(tangent)) + .1 + } + + /// [`eval_with_tangent`](Self::eval_with_tangent), also returning how many instructions ran. + pub fn eval_counted_with_tangent<'s>( + &self, + scratch: &'s mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + tangent: &TangentInputs<'_>, + ) -> (usize, &'s [f64]) { + self.eval_internal(scratch, t, y, y_dot, inputs, Some(tangent)) + } + + fn eval_internal<'s>( + &self, + scratch: &'s mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + tangent: Option<&TangentInputs<'_>>, + ) -> (usize, &'s [f64]) { + let instructions = self.ir.instructions(); + let skipped = eval_instructions( + scratch, + instructions, + self.ir.consts(), + EvalContext::with_primal(t, y, y_dot, inputs, tangent), + ); + ( + instructions.len() - skipped, + &scratch[self.ir.root_slot().bounds()], + ) + } + + /// Evaluate and copy the result into `out`, using `scratch` for working space. + #[inline] + pub fn eval_into( + &self, + scratch: &mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + out: &mut [f64], + ) { + let result = self.eval(scratch, t, y, y_dot, inputs); + out[..result.len()].copy_from_slice(result); + } +} + +/// A [`CompiledExpr`] split into a primal and a tangent section. +/// +/// No batch path; rewrapped via [`CompiledExpr::from_ir`], a tangent load errs in `eval_batch`: +/// +/// ```compile_fail +/// # use pybamm_core::{Arena, DagRef, Node, SplitCompiledExpr}; +/// # let mut arena = Arena::new(); +/// # let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); +/// let split = SplitCompiledExpr::compile(DagRef::new(&arena, y)); +/// let _ = split.eval_batch(&mut [0.0; 8], 1, &[0.0], &[1.0], 1, &[]); +/// ``` +#[derive(Debug, Clone)] +pub struct SplitCompiledExpr { + expr: CompiledExpr, + split: SplitEvalInfo, +} + +impl SplitCompiledExpr { + /// Lower and compile `dag` with primal and tangent instructions partitioned. + pub fn compile(dag: DagRef<'_>) -> Self { + let (ir, split) = TypedIr::from_dag_split_eval(dag); + Self { + expr: CompiledExpr::from_ir(ir), + split, + } + } + + /// The whole tape, primal section first. + #[inline] + pub const fn ir(&self) -> &TypedIr { + self.expr.ir() + } + + /// Where the primal section ends, in instructions and in buffer slots. + #[inline] + pub const fn split_info(&self) -> SplitEvalInfo { + self.split + } + + /// Length of the scratch buffer both sections share. + #[inline] + pub const fn scratch_len(&self) -> usize { + self.expr.scratch_len() + } + + /// Both sections in one walk, for a single tangent. + pub fn eval_with_tangent<'s>( + &self, + scratch: &'s mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + tangent: &TangentInputs<'_>, + ) -> &'s [f64] { + self.expr + .eval_with_tangent(scratch, t, y, y_dot, inputs, tangent) + } + + /// Evaluate the primal section into `scratch`, returning the [`PrimalCache`] that + /// runs the tangent sweep on the same buffer. + pub fn eval_primal<'a>( + &'a self, + scratch: &'a mut [f64], + t: f64, + y: &[f64], + y_dot: &[f64], + inputs: &[f64], + ) -> PrimalCache<'a> { + eval_instructions( + scratch, + &self.ir().instructions()[..self.split.primal_end], + self.ir().consts(), + EvalContext::with_primal(t, y, y_dot, inputs, None), + ); + PrimalCache { + expr: self, + scratch, + } + } +} + +/// Borrow-based typestate for split evaluation. +/// +/// The only route to the tangent sweep, so tangent-before-primal is a compile error. +pub struct PrimalCache<'a> { + expr: &'a SplitCompiledExpr, + scratch: &'a mut [f64], +} + +impl std::fmt::Debug for PrimalCache<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PrimalCache").finish_non_exhaustive() + } +} + +impl PrimalCache<'_> { + /// Evaluate the tangent section on the cached primal region, returning the root slice. + pub fn eval_tangent(&mut self, tangent: &TangentInputs<'_>) -> &[f64] { + let ir = self.expr.ir(); + eval_instructions( + self.scratch, + &ir.instructions()[self.expr.split.primal_end..], + ir.consts(), + EvalContext::tangent_only(tangent), + ); + &self.scratch[ir.root_slot().bounds()] + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arena::{Arena, NodeId}; + use crate::ir::Slot; + use crate::ir::ids::{BufSlot, Width}; + use crate::node::{ + ArrayData, CsrData, CubicInterpolantData, InterpolantData, NdInterpolantData, Node, Shape, + }; + use crate::params::ParamIdx; + + /// Compile `root` and evaluate it once at `(t, y)`, with no `y_dot` or inputs. + fn eval_at(arena: &Arena, root: NodeId, t: f64, y: &[f64]) -> Vec { + let compiled = CompiledExpr::new(DagRef::new(arena, root)); + let mut scratch = vec![0.0; compiled.scratch_len()]; + compiled.eval(&mut scratch, t, y, &[], &[]).to_vec() + } + + fn scalar(arena: &mut Arena, value: f64) -> NodeId { + arena.alloc(Node::Scalar(value)) + } + + fn array(arena: &mut Arena, data: &[f64]) -> NodeId { + arena.alloc(Node::Array(Box::new(ArrayData { + data: data.to_vec(), + shape: Shape::vector(data.len()), + }))) + } + + #[test] + fn erf_matches_the_reference_to_an_ulp() { + // scipy.special.erf, printed at full precision. + let reference = [ + (0.0, 0.0), + (1e-8, 1.128_379_167_095_512_5e-8), + (0.1, 0.112_462_916_018_284_9), + (0.5, 0.520_499_877_813_046_5), + (0.843_75, 0.767_225_661_232_341_6), + (1.0, 0.842_700_792_949_714_8), + (1.25, 0.922_900_128_256_458_3), + (2.0, 0.995_322_265_018_952_7), + (3.0, 0.999_977_909_503_001_4), + (4.755_430_192_319_827, 0.999_999_999_982_464_1), + (6.0, 1.0), + ]; + for (x, want) in reference { + for signum in [1.0, -1.0] { + let got = erf(signum * x); + let want = signum * want; + assert!( + (got - want).abs() <= 2.0 * f64::EPSILON * want.abs(), + "erf({}) = {got:.17e}, reference {want:.17e}", + signum * x + ); + } + } + // scipy.special.erfc; `1 - erf(x)` is exact, hence the absolute floor. + let tail = [ + (3.0, 2.209_049_699_858_544_5e-5), + (4.0, 1.541_725_790_028_002e-8), + (4.755_430_192_319_827, 1.753_587_985_582_014e-11), + ]; + for (x, want) in tail { + let got = 1.0 - erf(x); + assert!( + (got - want).abs() <= 1e-5 * want + 2.0 * f64::EPSILON, + "1 - erf({x}) = {got:.17e}, erfc reference {want:.17e}" + ); + } + + assert_eq!( + erf(0.0).to_bits(), + 0.0_f64.to_bits(), + "erf(0) must be exactly +0" + ); + assert_eq!( + erf(-0.0).to_bits(), + (-0.0_f64).to_bits(), + "erf(-0) must be exactly -0" + ); + } + + #[test] + fn the_erf_derivative_rule_matches_the_primal_it_differentiates() { + let rule = |x: f64| std::f64::consts::FRAC_2_SQRT_PI * (-x * x).exp(); + for i in 0..=60 { + let x = 0.05 + f64::from(i) * 0.05; + let step = 1e-4 * x; + let quotient = (erf(x + step) - erf(x - step)) / (2.0 * step); + // 1e-6 relative covers the central difference's O(step^2) truncation. + assert!( + (quotient - rule(x)).abs() <= 1e-6 * rule(x) + 1e-12, + "at x = {x}: quotient {quotient:.12e} vs rule {:.12e}", + rule(x) + ); + } + } + + #[test] + fn size_check() { + use std::mem::size_of; + // Guard hot per-node types against accidental bloat. + assert_eq!(size_of::(), 4); + assert_eq!(size_of::(), 4); + assert_eq!(size_of::(), 4); + assert_eq!(size_of::(), 8); + assert_eq!(size_of::(), 4); + assert_eq!(size_of::(), 1); + assert_eq!(size_of::(), 1); + assert_eq!(size_of::(), 1); + } + + #[test] + fn single_nodes_evaluate_to_their_closed_form() { + type Build = fn(&mut Arena) -> NodeId; + let cases: [(&str, Build, &[f64]); 17] = [ + ("scalar", |a| scalar(a, 42.0), &[42.0]), + ("time", |a| a.alloc(Node::Time), &[3.5]), + ( + "state", + |a| a.alloc(Node::StateVector { start: 1, end: 3 }), + &[20.0, 30.0], + ), + ("array", |a| array(a, &[1.0, 2.0, 3.0]), &[1.0, 2.0, 3.0]), + ( + "add scalars", + |a| { + let (x, y) = (scalar(a, 2.0), scalar(a, 3.0)); + a.alloc(Node::Add(x, y)) + }, + &[5.0], + ), + ( + "add vectors", + |a| { + let x = array(a, &[1.0, 2.0, 3.0]); + let y = array(a, &[10.0, 20.0, 30.0]); + a.alloc(Node::Add(x, y)) + }, + &[11.0, 22.0, 33.0], + ), + ( + "scalar times vector", + |a| { + let (s, v) = (scalar(a, 2.0), array(a, &[1.0, 2.0, 3.0])); + a.alloc(Node::Mul(s, v)) + }, + &[2.0, 4.0, 6.0], + ), + ( + "(2 + 3) * 4 - 1", + |a| { + let (two, three) = (scalar(a, 2.0), scalar(a, 3.0)); + let (four, one) = (scalar(a, 4.0), scalar(a, 1.0)); + let sum = a.alloc(Node::Add(two, three)); + let product = a.alloc(Node::Mul(sum, four)); + a.alloc(Node::Sub(product, one)) + }, + &[19.0], + ), + ( + "neg", + |a| { + let x = scalar(a, 5.0); + a.alloc(Node::Neg(x)) + }, + &[-5.0], + ), + ( + "sqrt", + |a| { + let x = array(a, &[1.0, 4.0, 9.0]); + a.alloc(Node::Sqrt(x)) + }, + &[1.0, 2.0, 3.0], + ), + ( + "sin(0) + cos(0)", + |a| { + let zero = scalar(a, 0.0); + let (sin, cos) = (a.alloc(Node::Sin(zero)), a.alloc(Node::Cos(zero))); + a.alloc(Node::Add(sin, cos)) + }, + &[1.0], + ), + ( + "max reduce", + |a| { + let x = array(a, &[1.0, 5.0, 3.0]); + a.alloc(Node::MaxReduce(x)) + }, + &[5.0], + ), + ( + "index", + |a| { + let child = array(a, &[10.0, 20.0, 30.0, 40.0, 50.0]); + a.alloc(Node::Index { + child, + start: 1, + end: 4, + }) + }, + &[20.0, 30.0, 40.0], + ), + ( + "concat", + |a| { + let (x, y, z) = (scalar(a, 1.0), array(a, &[2.0, 3.0]), scalar(a, 4.0)); + a.alloc(Node::Concat(vec![x, y, z])) + }, + &[1.0, 2.0, 3.0, 4.0], + ), + ( + "sparse matmul", + |a| { + let m = a.alloc(Node::SparseMatrix(Box::new(CsrData { + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![2.0, 3.0], + shape: Shape::matrix(2, 3), + }))); + let v = array(a, &[1.0, 2.0, 3.0]); + a.alloc(Node::MatMul(m, v)) + }, + &[2.0, 6.0], + ), + ( + "dense matmul", + |a| { + let m = a.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], // row-major 2x3 + shape: Shape::matrix(2, 3), + }))); + let v = array(a, &[1.0, 2.0, 3.0]); + a.alloc(Node::MatMul(m, v)) + }, + &[14.0, 32.0], + ), + ( + "linear interpolant", + |a| { + let child = scalar(a, 1.5); + let table = + InterpolantData::try_new(vec![0.0, 1.0, 2.0], vec![0.0, 10.0, 20.0]) + .expect("valid interpolant"); + a.alloc(Node::Interpolant1DLinear { + data: Box::new(table), + child, + }) + }, + &[15.0], + ), + ]; + for (label, build, expected) in cases { + let mut arena = Arena::new(); + let root = build(&mut arena); + assert_eq!( + eval_at(&arena, root, 3.5, &[10.0, 20.0, 30.0, 40.0]), + expected, + "{label}" + ); + } + } + + #[test] + fn reduce_arg_select_picks_the_basis_at_the_first_extremum() { + for (picker_values, is_max, expected) in [ + ([1.0, 5.0, 3.0], true, 20.0), + ([1.0, 5.0, 3.0], false, 10.0), + ([5.0, 5.0, 3.0], true, 10.0), + ([5.0, 3.0, 3.0], false, 20.0), + ] { + let mut arena = Arena::new(); + let picker = array(&mut arena, &picker_values); + let basis = array(&mut arena, &[10.0, 20.0, 30.0]); + let node = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max, + }); + assert_eq!( + eval_at(&arena, node, 0.0, &[]), + [expected], + "picker {picker_values:?}, is_max={is_max}" + ); + } + } + + /// The subgradient must credit the element the reduction returned, NaNs included. + #[test] + fn arg_select_skips_a_nan_as_its_reduction_does() { + // A remainder at the reduction's starting value still beats a leading NaN. + assert_eq!(arg_reduce([f64::NAN, f64::NEG_INFINITY], true), 1); + assert_eq!(arg_reduce([f64::NAN, f64::INFINITY], false), 1); + assert_eq!(arg_reduce([f64::NAN, f64::NAN], true), 0); + + let picker = [f64::NAN, 1.0, 3.0, f64::NAN, -2.0]; + assert_eq!(max_reduce(picker).to_bits(), 3.0_f64.to_bits()); + assert_eq!(min_reduce(picker).to_bits(), (-2.0_f64).to_bits()); + // Positions 2 and 4, where the reductions found 3 and -2. + for (is_max, expected) in [(true, 30.0_f64), (false, 50.0)] { + let mut arena = Arena::new(); + let picker = array(&mut arena, &picker); + let basis = array(&mut arena, &[10.0, 20.0, 30.0, 40.0, 50.0]); + let node = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max, + }); + let picked = eval_at(&arena, node, 0.0, &[])[0]; + assert_eq!(picked.to_bits(), expected.to_bits(), "is_max={is_max}"); + } + } + + #[test] + fn a_conditional_takes_the_branch_whose_open_window_holds_the_selector() { + // Branch i (1-based) owns (i - 0.5, i + 0.5); any other selector, NaN included, is zero. + for (selector, branches, expected) in [ + (2.0, &[100.0, 200.0, 300.0][..], 200.0), + (0.6, &[100.0, 200.0], 100.0), + (1.4, &[100.0, 200.0], 100.0), + (0.5, &[100.0, 200.0], 0.0), + (0.0, &[100.0, 200.0], 0.0), + (10.0, &[100.0, 200.0], 0.0), + (-1.0, &[100.0, 200.0], 0.0), + (f64::NAN, &[999.0], 0.0), + ] { + let mut arena = Arena::new(); + let selector_id = scalar(&mut arena, selector); + let branches = branches.iter().map(|&b| scalar(&mut arena, b)).collect(); + let cond = arena.alloc(Node::Conditional { + selector: selector_id, + branches, + }); + assert_eq!( + eval_at(&arena, cond, 0.0, &[]), + [expected], + "selector {selector}" + ); + } + } + + #[test] + fn test_eval_interpolant_breakpoint_exactness() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + InterpolantData::try_new(vec![0.0, 1.0, 3.0, 7.0], vec![10.0, 20.0, 60.0, -5.0]) + .expect("valid interpolant"), + ), + child: x, + }); + for (xv, expected) in [(0.0, 10.0), (1.0, 20.0), (3.0, 60.0), (7.0, -5.0)] { + let result = eval_at(&arena, interp, 0.0, &[xv]); + assert!( + (result[0] - expected).abs() < 1e-14, + "at x={xv}: expected {expected}, got {}", + result[0] + ); + } + } + + #[test] + fn linear_interpolant_matches_casadi_bit_for_bit() { + // `casadi.MX.interpn_linear` values, as PyBaMM builds it; a fused multiply-add + // misses every one of these by an ulp. + let table = InterpolantData::try_new(vec![0.0, 0.1, 0.3, 0.7], vec![0.2, -1.3, 2.9, 0.05]) + .expect("valid interpolant"); + let casadi: [(f64, f64); 6] = [ + (0.079_896_907_216_494_84, -0.998_453_608_247_422_7), + (0.141_752_577_319_587_62, -0.423_195_876_288_659_84), + (0.152_061_855_670_103_1, -0.206_701_030_927_834_76), + (0.162_371_134_020_618_53, 0.009_793_814_432_989_212), + (-0.219_072_164_948_453_6, 3.486_082_474_226_804), + (0.708_762_886_597_938_2, -0.012_435_567_010_309_79), + ]; + for (x, want) in casadi { + let got = interp_linear_1d(&table, x); + assert_eq!( + got.to_bits(), + want.to_bits(), + "at x={x}: {got:e} vs CasADi {want:e}" + ); + } + } + + #[test] + #[allow(clippy::float_cmp)] + fn a_repeated_knot_is_a_right_continuous_step() { + // As `casadi.MX.interpn_linear` takes it: the knot belongs to the segment after it. + let table = InterpolantData::try_new(vec![0.0, 1.0, 1.0, 2.0], vec![0.0, 1.0, 3.0, 5.0]) + .expect("a repeated interior knot is valid"); + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DLinear { + data: Box::new(table.clone()), + child: x, + }); + let cases = [ + (-1.0, -1.0, 1.0), + (0.5, 0.5, 1.0), + (0.75, 0.75, 1.0), + (1.0, 3.0, 2.0), + (1.5, 4.0, 2.0), + (3.0, 7.0, 2.0), + ]; + for (xv, value, slope) in cases { + assert_eq!( + eval_at(&arena, interp, 0.0, &[xv]), + [value], + "value at {xv}" + ); + assert_eq!(interp_linear_1d_deriv(&table, xv), slope, "slope at {xv}"); + } + } + + #[test] + fn test_eval_interpolant_linear_extrapolation() { + // Matches scipy `interp1d(fill_value="extrapolate")`. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + InterpolantData::try_new(vec![1.0, 2.0, 3.0], vec![100.0, 200.0, 300.0]) + .expect("valid interpolant"), + ), + child: x, + }); + assert_eq!(eval_at(&arena, interp, 0.0, &[-5.0]), [-500.0]); + assert_eq!(eval_at(&arena, interp, 0.0, &[5.0]), [500.0]); + } + + #[test] + fn test_eval_interpolant_cubic() { + // p(dx) = 1 + 2*dx + 3*dx^2 + 4*dx^3 on [0, 2]. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DCubic { + data: Box::new(CubicInterpolantData { + breakpoints: vec![0.0, 2.0], + coeffs: vec![[1.0, 2.0, 3.0, 4.0]], + }), + child: x, + }); + for (xv, expected) in [(1.0, 10.0), (0.0, 1.0), (3.0, 142.0)] { + assert_eq!(eval_at(&arena, interp, 0.0, &[xv]), [expected], "x={xv}"); + } + } + + #[test] + fn test_eval_interpolant_cubic_breakpoints() { + // Right-continuous at breakpoints, like scipy `PPoly`. + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let interp = arena.alloc(Node::Interpolant1DCubic { + data: Box::new(CubicInterpolantData { + breakpoints: vec![0.0, 1.0, 3.0], + coeffs: vec![[5.0, 0.0, 0.0, 0.0], [7.0, 0.0, 0.0, 0.0]], + }), + child: x, + }); + for (xv, expected) in [(0.5, 5.0), (2.0, 7.0), (0.0, 5.0), (1.0, 7.0), (3.0, 7.0)] { + assert_eq!(eval_at(&arena, interp, 0.0, &[xv]), [expected], "x={xv}"); + } + } + + #[test] + fn test_eval_interpolant_nd_bilinear() { + // p = 1 + 3*dx1 + 2*dx0 + 4*dx0*dx1, coeffs indexed `a0*order + a1`. + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let interp = arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 2.0], vec![0.0, 2.0]], + coeffs: vec![1.0, 3.0, 2.0, 4.0], + order: 2, + }), + children: vec![x0, x1], + }); + for (y, expected) in [([1.0, 1.0], 10.0), ([0.0, 0.0], 1.0), ([3.0, 4.0], 67.0)] { + assert_eq!(eval_at(&arena, interp, 0.0, &y), [expected], "at {y:?}"); + } + } + + #[test] + fn test_eval_interpolant_nd_cell_selection() { + // Right-continuous like scipy, clamped at both edges. + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let interp = arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 1.0, 3.0], vec![0.0, 1.0]], + coeffs: vec![5.0, 0.0, 0.0, 0.0, 7.0, 0.0, 0.0, 0.0], + order: 2, + }), + children: vec![x0, x1], + }); + for (x, expected) in [(0.5, 5.0), (2.0, 7.0), (1.0, 7.0), (0.0, 5.0), (3.0, 7.0)] { + assert_eq!( + eval_at(&arena, interp, 0.0, &[x, 0.5]), + [expected], + "x0={x}" + ); + } + } + + #[test] + fn test_eval_interpolant_nd_tricubic() { + // p = 3 + dx0^3 + 2*dx1^2 + 5*dx2, coeffs indexed `(a0*4 + a1)*4 + a2`. + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let x2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let mut coeffs = vec![0.0; 64]; + coeffs[0] = 3.0; // (0,0,0) + coeffs[(3 * 4) * 4] = 1.0; // dx0^3 + coeffs[2 * 4] = 2.0; // dx1^2 + coeffs[1] = 5.0; // dx2^1 + let interp = arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 2.0], vec![0.0, 2.0], vec![0.0, 2.0]], + coeffs, + order: 4, + }), + children: vec![x0, x1, x2], + }); + assert_eq!(eval_at(&arena, interp, 0.0, &[1.0, 1.0, 1.0]), [11.0]); + assert_eq!(eval_at(&arena, interp, 0.0, &[3.0, 0.0, 0.0]), [30.0]); + } + + #[test] + fn test_eval_interpolant_nd_vector_children_broadcast() { + // p = dx0 + 10*dx1; the length-1 axis-1 child broadcasts. + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let x1 = arena.alloc(Node::StateVector { start: 3, end: 4 }); + let interp = arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 10.0], vec![0.0, 10.0]], + coeffs: vec![0.0, 10.0, 1.0, 0.0], + order: 2, + }), + children: vec![x0, x1], + }); + assert_eq!( + eval_at(&arena, interp, 0.0, &[1.0, 2.0, 3.0, 0.5]), + [6.0, 7.0, 8.0] + ); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/eval_batch.rs b/packages/pybamm-rust/pybamm-core/src/eval_batch.rs new file mode 100644 index 0000000000..be625a1dd2 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/eval_batch.rs @@ -0,0 +1,473 @@ +//! Lane-batched (K-wide) primal evaluator: slot `i` maps to lanes `[i*k, (i+1)*k)`, +//! with results bitwise identical to `k` independent `eval` calls. + +// The op helpers must inline as one chain into the interpreter loop. +#![allow(clippy::inline_always)] + +use crate::accumulate::{Accumulate, Product, dispatch}; +use crate::branch_regions::active_branch; +use crate::eval::{ + CompiledExpr, arg_reduce, csr_rows, interp_cubic_1d, interp_cubic_1d_deriv, interp_linear_1d, + interp_linear_1d_deriv, interp_nd_lanes, max_reduce, min_reduce, split_read_write, + split_two_read_write, unary_apply, with_binary_op, with_unary_op, zip_apply, +}; +use crate::ir::{BroadcastKind, ConstPool, Instruction, Slot}; +use crate::node::CsrData; + +/// Error returned when a tape cannot be evaluated by the primal batch path. +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum BatchEvalError { + #[error("eval_batch is primal-only but the tape contains a tangent load")] + NonPrimalInstruction, + #[error("eval_batch does not supply y_dot but the tape references state derivatives")] + StateDotUnsupported, +} + +/// Apply a binary broadcast op over `k` lanes with a monomorphic closure. +#[allow(clippy::too_many_arguments)] +#[inline(always)] +fn batch_broadcast_apply f64>( + buf: &mut [f64], + k: usize, + f: F, + a: usize, + b: usize, + dst: usize, + len: usize, + kind: BroadcastKind, +) { + match kind { + // A scalar-scalar op has `len == 1`. + BroadcastKind::ScalarScalar | BroadcastKind::VectorVector => { + zip_apply(buf, f, a * k, b * k, dst * k, len * k); + }, + BroadcastKind::ScalarVector => { + let (scalar, b_s, d_s) = + split_two_read_write(buf, (a * k, k), (b * k, len * k), (dst * k, len * k)); + for (d_chunk, b_chunk) in d_s.chunks_exact_mut(k).zip(b_s.chunks_exact(k)) { + for ((o, &s), &y) in d_chunk.iter_mut().zip(scalar).zip(b_chunk) { + *o = f(s, y); + } + } + }, + BroadcastKind::VectorScalar => { + let (a_s, scalar, d_s) = + split_two_read_write(buf, (a * k, len * k), (b * k, k), (dst * k, len * k)); + for (d_chunk, a_chunk) in d_s.chunks_exact_mut(k).zip(a_s.chunks_exact(k)) { + for ((o, &x), &s) in d_chunk.iter_mut().zip(a_chunk).zip(scalar) { + *o = f(x, s); + } + } + }, + } +} + +/// Fill the front of `dst` with `values` repeated across `lanes` lanes, lane-minor. +#[inline] +pub fn broadcast_lanes(dst: &mut [f64], lanes: usize, values: &[f64]) { + for (chunk, &value) in dst[..values.len() * lanes] + .chunks_exact_mut(lanes) + .zip(values) + { + chunk.fill(value); + } +} + +/// Lane `l` of a lane-minor window `k` wide: its elements `l`, `l + k`, `l + 2k`, ... +#[inline(always)] +fn lane(window: &[f64], k: usize, l: usize) -> impl Iterator { + window.iter().skip(l).step_by(k).copied() +} + +/// The evaluation point of every lane, shared by a whole batched walk. +#[derive(Clone, Copy)] +struct BatchInputs<'a> { + /// Lane count. + k: usize, + ts: &'a [f64], + /// The `(n_states, k)` column-major state matrix. + y_cols: &'a [f64], + n_states: usize, + inputs: &'a [f64], +} + +/// Execute a primal instruction slice against a lane-batched buffer. +/// +/// `Dispatch` runs every block any lane selects; `Conditional` still selects per lane. +fn eval_batch_instructions( + buf: &mut [f64], + instructions: &[Instruction], + consts: &ConstPool, + lanes: BatchInputs<'_>, +) -> Result<(), BatchEvalError> { + let BatchInputs { + k, + ts, + y_cols, + n_states, + inputs, + } = lanes; + let mut pc = 0_usize; + while pc < instructions.len() { + match instructions[pc] { + Instruction::Dispatch { + selector, + blocks_idx, + blocks_len, + span_len, + } => { + let n = blocks_len.index(); + let selectors = selector.index() * k..(selector.index() + 1) * k; + for b in 0..n { + // No block writes the selector: the `Conditional` after them still reads it. + if buf[selectors.clone()] + .iter() + .any(|&s| active_branch(s, n) == Some(b)) + { + let block = consts.branch_blocks[blocks_idx.index() + b].bounds(pc); + eval_batch_instructions(buf, &instructions[block], consts, lanes)?; + } + } + pc += span_len.index(); + continue; + }, + Instruction::LoadScalar { value, dst } => { + let base = dst.index() * k; + buf[base..base + k].fill(value); + }, + Instruction::LoadTime { dst } => { + let base = dst.index() * k; + buf[base..base + k].copy_from_slice(&ts[..k]); + }, + Instruction::LoadArray { data_idx, len, dst } => { + broadcast_lanes( + &mut buf[dst.index() * k..], + k, + consts.get_array(data_idx, len), + ); + }, + Instruction::FillZero { dst, len } => { + let base = dst.index() * k; + buf[base..base + len.index() * k].fill(0.0); + }, + Instruction::LoadStateVector { start, width, dst } => { + let base = dst.index() * k; + // `y_cols` is the `(n_states, k)` column-major state matrix. + for e in 0..width.index() { + let o = base + e * k; + let row = start.index() + e; + for l in 0..k { + buf[o + l] = y_cols[l * n_states + row]; + } + } + }, + Instruction::LoadStateVectorDot { .. } => { + return Err(BatchEvalError::StateDotUnsupported); + }, + Instruction::LoadInputParameter { offset, width, dst } => { + let values = &inputs[offset.index()..offset.index() + width.index()]; + broadcast_lanes(&mut buf[dst.index() * k..], k, values); + }, + Instruction::LoadTangentState { .. } | Instruction::LoadTangentParameter { .. } => { + return Err(BatchEvalError::NonPrimalInstruction); + }, + + Instruction::Binary { + op, + a, + b, + dst, + len, + kind, + } => { + let (a, b, dst, len) = (a.index(), b.index(), dst.index(), len.index()); + with_binary_op!(op, f => batch_broadcast_apply(buf, k, f, a, b, dst, len, kind)); + }, + + Instruction::Unary { op, src, dst, len } => { + let (src, dst, len) = (src.index(), dst.index(), len.index()); + with_unary_op!(op, f => unary_apply(buf, f, src * k, dst * k, len * k)); + }, + + Instruction::MaxReduce { src, src_len, dst } => { + let (src, src_len, dst) = (src.index(), src_len.index(), dst.index()); + for l in 0..k { + buf[dst * k + l] = max_reduce(lane(&buf[src * k..(src + src_len) * k], k, l)); + } + }, + + Instruction::MinReduce { src, src_len, dst } => { + let (src, src_len, dst) = (src.index(), src_len.index(), dst.index()); + for l in 0..k { + buf[dst * k + l] = min_reduce(lane(&buf[src * k..(src + src_len) * k], k, l)); + } + }, + + Instruction::ReduceArgSelect { + basis_src, + picker_src, + len, + is_max, + dst, + } => { + let (basis, picker, len) = (basis_src.index(), picker_src.index(), len.index()); + let dbase = dst.index() * k; + for l in 0..k { + let best = arg_reduce(lane(&buf[picker * k..(picker + len) * k], k, l), is_max); + buf[dbase + l] = buf[(basis + best) * k + l]; + } + }, + + Instruction::Index { + src, + start, + dst, + len, + } => { + let s = src.offset_by(start) * k; + let n = len.index() * k; + buf.copy_within(s..s + n, dst.index() * k); + }, + + Instruction::Concat { + sources_idx, + sources_len, + dst, + } => { + let mut write_pos = dst.index() * k; + for i in 0..sources_len.index() { + let slot = consts.concat_sources[sources_idx.index() + i]; + let n = slot.len_usize() * k; + buf.copy_within( + slot.offset_usize() * k..slot.offset_usize() * k + n, + write_pos, + ); + write_pos += n; + } + }, + + Instruction::MatMul { + csr_idx, + vec_src, + dst, + } => { + let csr = &consts.csr_data[csr_idx.index()]; + let (x, out) = split_read_write( + buf, + vec_src.index() * k, + csr.shape.cols * k, + dst.index() * k, + csr.shape.rows * k, + ); + dispatch(CsrLanes { csr, x, out, k }); + }, + + Instruction::DenseMatMul { + mat_src, + rows, + cols, + vec_src, + dst, + } => { + let (rows, cols) = (rows.index(), cols.index()); + let (m, v, out) = split_two_read_write( + buf, + (mat_src.index() * k, rows * cols * k), + (vec_src.index() * k, cols * k), + (dst.index() * k, rows * k), + ); + dispatch(DenseLanes { m, v, out, cols, k }); + }, + + Instruction::Interp1DLinear { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.interpolants[interp_idx.index()]; + let f = |x| interp_linear_1d(interp, x); + unary_apply(buf, f, src.index() * k, dst.index() * k, len.index() * k); + }, + + Instruction::Interp1DLinearDeriv { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.interpolants[interp_idx.index()]; + let f = |x| interp_linear_1d_deriv(interp, x); + unary_apply(buf, f, src.index() * k, dst.index() * k, len.index() * k); + }, + + Instruction::Interp1DCubic { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.cubic_interpolants[interp_idx.index()]; + let f = |x| interp_cubic_1d(&interp.breakpoints, &interp.coeffs, x); + unary_apply(buf, f, src.index() * k, dst.index() * k, len.index() * k); + }, + + Instruction::Interp1DCubicDeriv { + interp_idx, + src, + dst, + len, + } => { + let interp = &consts.cubic_interpolants[interp_idx.index()]; + let f = |x| interp_cubic_1d_deriv(&interp.breakpoints, &interp.coeffs, x); + unary_apply(buf, f, src.index() * k, dst.index() * k, len.index() * k); + }, + + Instruction::InterpNd { + interp_idx, + sources_idx, + dst, + len, + } => { + let out = Slot::new(dst, len); + interp_nd_lanes(buf, k, consts, interp_idx, sources_idx, None, out); + }, + + Instruction::InterpNdPartial { + interp_idx, + sources_idx, + axis, + dst, + len, + } => { + let (axis, out) = (Some(axis.index()), Slot::new(dst, len)); + interp_nd_lanes(buf, k, consts, interp_idx, sources_idx, axis, out); + }, + + Instruction::Conditional { + selector, + branches_idx, + branches_len, + dst, + out_len, + } => { + let sel_base = selector.index() * k; + let dst = dst.index(); + let out_len = out_len.index(); + for l in 0..k { + match active_branch(buf[sel_base + l], branches_len.index()) { + Some(i) => { + let bbase = + consts.branch_offsets[branches_idx.index() + i].offset_usize(); + for e in 0..out_len { + buf[(dst + e) * k + l] = buf[(bbase + e) * k + l]; + } + }, + None => { + for e in 0..out_len { + buf[(dst + e) * k + l] = 0.0; + } + }, + } + } + }, + } + pc += 1; + } + Ok(()) +} + +impl CompiledExpr { + /// Evaluate the tape for `k` lanes (time points) at once. + /// + /// `scratch` needs `scratch_len() * k` elements; `y_cols` is the `(n_states, k)` + /// column-major state matrix. Returns the root slot lane-minor (element `e`, lane + /// `l` at `e*k + l`), bitwise identical to `k` [`eval`](Self::eval) calls. + /// + /// # Panics + /// If `k` is zero or `y_cols` does not hold exactly `n_states * k` values. + pub fn eval_batch<'s>( + &self, + scratch: &'s mut [f64], + k: usize, + ts: &[f64], + y_cols: &[f64], + n_states: usize, + inputs: &[f64], + ) -> Result<&'s [f64], BatchEvalError> { + let ir = self.ir(); + assert!(k >= 1, "eval_batch needs at least one lane"); + assert_eq!( + y_cols.len(), + n_states * k, + "eval_batch needs n_states = {n_states} values for each of {k} lanes" + ); + debug_assert!(ts.len() >= k, "ts must have at least k time values"); + debug_assert!( + scratch.len() >= ir.buffer_size() * k, + "scratch too small for k lanes" + ); + let lanes = BatchInputs { + k, + ts, + y_cols, + n_states, + inputs, + }; + eval_batch_instructions(scratch, ir.instructions(), ir.consts(), lanes)?; + let root = ir.root_slot(); + Ok(&scratch[root.offset_usize() * k..(root.offset_usize() + root.len_usize()) * k]) + } +} + +/// `out = A @ x` across `k` lanes, each lane accumulated as the scalar `CsrProduct` does. +struct CsrLanes<'a> { + csr: &'a CsrData, + x: &'a [f64], + out: &'a mut [f64], + k: usize, +} + +impl Product for CsrLanes<'_> { + type Output = (); + + #[inline(always)] + fn run(self) { + let Self { csr, x, out, k } = self; + for (acc, (cols, vals)) in out.chunks_exact_mut(k).zip(csr_rows(csr)) { + acc.fill(0.0); + for (&col, &val) in cols.iter().zip(vals) { + for (a, &xi) in acc.iter_mut().zip(&x[col * k..][..k]) { + *a = A::mul_acc(*a, val, xi); + } + } + } + } +} + +/// `out = M @ v` across `k` lanes for a row-major `M` with `cols` columns. +struct DenseLanes<'a> { + m: &'a [f64], + v: &'a [f64], + out: &'a mut [f64], + cols: usize, + k: usize, +} + +impl Product for DenseLanes<'_> { + type Output = (); + + #[inline(always)] + fn run(self) { + let Self { m, v, out, cols, k } = self; + for (row, acc) in out.chunks_exact_mut(k).enumerate() { + acc.fill(0.0); + let m_row = &m[row * cols * k..][..cols * k]; + for (m_lanes, v_lanes) in m_row.chunks_exact(k).zip(v.chunks_exact(k)) { + for ((a, &mi), &vi) in acc.iter_mut().zip(m_lanes).zip(v_lanes) { + *a = A::mul_acc(*a, mi, vi); + } + } + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/ir/ids.rs b/packages/pybamm-rust/pybamm-core/src/ir/ids.rs new file mode 100644 index 0000000000..41be303938 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/ir/ids.rs @@ -0,0 +1,196 @@ +//! Typed `u32` indices for the instruction tape and its side tables, so buffer +//! offsets, element counts and pool indices cannot be swapped by accident. + +/// Define a `u32` index over the instruction tape, with `usize` conversions. +macro_rules! tape_index { + ($(#[$attr:meta])* $name:ident) => { + $crate::newtype::index_newtype! { $(#[$attr])* $name(u32) } + + impl $name { + /// Widened for indexing. + pub const fn index(self) -> usize { + self.get() as usize + } + + /// Narrow a `usize` count or offset; panics past `u32::MAX` rather than truncating. + pub fn from_usize(value: usize) -> Self { + Self::new( + u32::try_from(value) + .expect(concat!(stringify!($name), " exceeded u32::MAX")), + ) + } + } + }; +} + +pub(crate) use tape_index; + +/// Define a tape index that also indexes a constant-pool table. +macro_rules! pool_index { + ($(#[$attr:meta])* $name:ident) => { + tape_index! { $(#[$attr])* $name } + + impl $name { + /// The index the next `push` onto `table` will occupy. + pub(crate) fn of_next(table: &[T]) -> Self { + Self::from_usize(table.len()) + } + } + }; +} + +tape_index! { + /// Element offset of a [`Slot`](super::Slot)'s start in the evaluation buffer. + BufSlot +} + +tape_index! { + /// A count of `f64` elements: a slot's length, or an input's packed width. + Width +} + +tape_index! { + /// Element offset relative to the start of a slot, not to the buffer. + SlotOffset +} + +impl BufSlot { + /// The buffer index `within` elements into this slot. + #[must_use] + pub(crate) const fn offset_by(self, within: SlotOffset) -> usize { + self.index() + within.index() + } +} + +tape_index! { + /// A count of consecutive entries in one constant-pool side table. + EntryCount +} + +tape_index! { + /// A count of consecutive instructions on the tape. + InstructionCount +} + +tape_index! { + /// Row count of a dense matrix operand. + Rows +} + +tape_index! { + /// Column count of a dense matrix operand. + Cols +} + +impl Rows { + /// Element count of a column with this many rows. + pub(crate) const fn to_width(self) -> Width { + Width::new(self.get()) + } + + /// Element count of a dense `rows x cols` matrix. + pub(crate) fn times(self, cols: Cols) -> Width { + Width::from_usize(self.index() * cols.index()) + } +} + +impl Cols { + /// Element count of a row with this many columns. + pub(crate) const fn to_width(self) -> Width { + Width::new(self.get()) + } +} + +tape_index! { + /// Axis of an N-D interpolant, in `0..ndim`. + Axis +} + +tape_index! { + /// Index into the state vector `y` (or its derivative `y'`). + StateIdx +} + +tape_index! { + /// Element offset into the packed input vector, not a + /// [`ParamIdx`](crate::params::ParamIdx) registration index. + InputOffset +} + +pool_index! { + /// Index into [`ConstPool::array_offsets`](super::ConstPool::array_offsets). + ArrayIdx +} + +pool_index! { + /// Index into [`ConstPool::csr_data`](super::ConstPool::csr_data). + CsrIdx +} + +pool_index! { + /// Index into one of the [`Slot`](super::Slot) side tables; the instruction fixes which. + SlotTableIdx +} + +pool_index! { + /// Index into + /// [`ConstPool::branch_blocks`](super::ConstPool::branch_blocks). + BlockSpanIdx +} + +pool_index! { + /// Index into [`ConstPool::interpolants`](super::ConstPool::interpolants). + InterpolantIdx +} + +pool_index! { + /// Index into + /// [`ConstPool::cubic_interpolants`](super::ConstPool::cubic_interpolants). + CubicInterpolantIdx +} + +pool_index! { + /// Index into + /// [`ConstPool::nd_interpolants`](super::ConstPool::nd_interpolants). + NdInterpolantIdx +} + +/// A contiguous run of instructions one `Dispatch` branch owns. +/// +/// Offsets are relative to the owning `Dispatch`, since split tapes execute as sub-slices. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub struct BlockSpan { + rel_start: u32, + len: u32, +} + +impl BlockSpan { + /// A run of `len` instructions starting `rel_start` past its `Dispatch`. + pub(crate) const fn new(rel_start: u32, len: u32) -> Self { + Self { rel_start, len } + } + + /// A span from the `usize` instruction positions lowering counts in. + pub(crate) fn from_usize(rel_start: usize, len: usize) -> Self { + Self { + rel_start: u32::try_from(rel_start).expect("branch block offset exceeded u32::MAX"), + len: u32::try_from(len).expect("branch block length exceeded u32::MAX"), + } + } + + /// Instruction count; `0` when this branch owns none here. + pub const fn len(self) -> usize { + self.len as usize + } + + /// Whether this branch owns no instructions here. + pub const fn is_empty(self) -> bool { + self.len == 0 + } + + /// Absolute instruction bounds, given the `Dispatch`'s index in the executing tape. + pub const fn bounds(self, dispatch_pc: usize) -> std::ops::Range { + let start = dispatch_pc + self.rel_start as usize; + start..start + self.len as usize + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/ir/mod.rs b/packages/pybamm-rust/pybamm-core/src/ir/mod.rs new file mode 100644 index 0000000000..03b615bdae --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/ir/mod.rs @@ -0,0 +1,2277 @@ +//! Lowering from the expression DAG to a flat [`TypedIr`] instruction tape, with +//! operands as [`Slot`]s into one scratch buffer and constants in a [`ConstPool`]. + +pub mod ids; + +use crate::arena::{Arena, DagRef, NodeId}; +use crate::branch_regions::{ + RegionGroup, privatise_conditionals, region_owners, schedule_regions, + schedule_regions_partitioned, +}; +use crate::ir::ids::{ + ArrayIdx, Axis, BlockSpan, BlockSpanIdx, BufSlot, Cols, CsrIdx, CubicInterpolantIdx, + EntryCount, InputOffset, InstructionCount, InterpolantIdx, NdInterpolantIdx, Rows, SlotOffset, + SlotTableIdx, StateIdx, Width, +}; +use crate::node::{CsrData, CubicInterpolantData, InterpolantData, NdInterpolantData, Node}; +use crate::params::ParamIdx; + +/// A contiguous window of the evaluation buffer: where one value lives. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Slot { + pub offset: BufSlot, + pub len: Width, +} + +impl Slot { + #[inline] + pub const fn new(offset: BufSlot, len: Width) -> Self { + Self { offset, len } + } + + #[inline] + pub(crate) fn from_usize(offset: usize, len: usize) -> Self { + Self { + offset: BufSlot::from_usize(offset), + len: Width::from_usize(len), + } + } + + #[inline] + pub(crate) const fn offset_usize(self) -> usize { + self.offset.index() + } + + #[inline] + pub(crate) const fn len_usize(self) -> usize { + self.len.index() + } + + #[inline] + pub const fn bounds(self) -> std::ops::Range { + self.offset_usize()..self.offset_usize() + self.len_usize() + } + + /// Whether this slot holds a single element (and so broadcasts). + #[inline] + pub(crate) const fn is_scalar(self) -> bool { + self.len.get() == 1 + } +} + +/// Compile-time broadcast pattern for binary operations. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(u8)] +pub enum BroadcastKind { + ScalarScalar = 0, + /// Left operand is a single element, broadcast across the right. + ScalarVector = 1, + /// Right operand is a single element, broadcast across the left. + VectorScalar = 2, + VectorVector = 3, +} + +impl BroadcastKind { + /// Classify a pair of operand widths; panics on a mismatch [`validate`] rejects. + #[inline] + pub(crate) fn from_lens(a_len: usize, b_len: usize) -> Self { + assert!( + broadcast_widths_compatible(a_len, b_len), + "binary operand widths are incompatible: left={a_len}, right={b_len}" + ); + match (a_len, b_len) { + (1, 1) => Self::ScalarScalar, + (1, _) => Self::ScalarVector, + (_, 1) => Self::VectorScalar, + _ => Self::VectorVector, + } + } +} + +#[inline] +const fn broadcast_widths_compatible(a_len: usize, b_len: usize) -> bool { + a_len == b_len || (a_len == 1 && b_len > 1) || (b_len == 1 && a_len > 1) +} + +/// Element-wise binary operation, without domain guards; comparisons return 1.0 or 0.0. +/// +/// IEEE 754's exactly specified operations and the transcendentals take `f64`'s values, so a +/// transcendental may differ across platforms in its last bits. `Minimum` and `Maximum` take +/// `libm`'s, as `f64`'s leave the order of signed zeros open. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(u8)] +pub enum BinaryOp { + Add, + Sub, + Mul, + Div, + /// The platform's `pow`, but exact at the exponents `0`, `1`, `2` and `-1`. + Pow, + /// IEEE 754-2019 `minimumNumber`: the non-NaN operand if exactly one is NaN, and + /// `-0 < +0`. + Minimum, + /// IEEE 754-2019 `maximumNumber`: the non-NaN operand if exactly one is NaN, and + /// `-0 < +0`. + Maximum, + /// Truncated remainder (`%`): the sign follows the dividend, unlike Python's `%`. + Modulo, + Hypot, + /// `a <= b`, the branch-inclusive Heaviside: 1.0 when the operands are equal. + EqualHeaviside, + /// `a < b`, the strict Heaviside: 0.0 when the operands are equal. + NotEqualHeaviside, + /// `|a - b| < 1e-14`, a tolerance test rather than an exact comparison. + Equality, +} + +/// Element-wise unary `f64` function, without domain guards. +/// +/// IEEE 754's exactly specified operations and the transcendentals take `f64`'s values, so a +/// transcendental may differ across platforms in its last bits. `Erf` takes `libm`'s, as `f64` +/// has none, and so does `Arcsinh`, as `f64`'s overflows past `f64::MAX / 2`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(u8)] +pub enum UnaryOp { + Neg, + Abs, + Sqrt, + Exp, + Log, + Sin, + Cos, + Tanh, + Sinh, + Cosh, + Arcsinh, + Arctan, + Erf, + /// `sign(0) = 0`, unlike `f64::signum`. + Sign, + Floor, + Ceiling, +} + +/// Single evaluation step. +/// +/// `dst`/`src`/`a`/`b` are buffer offsets, `len` counts `f64` elements, and `*_idx` +/// addresses the [`ConstPool`]. Write windows never overlap read windows. +#[derive(Clone, Copy, Debug)] +pub enum Instruction { + LoadScalar { + value: f64, + dst: BufSlot, + }, + LoadTime { + dst: BufSlot, + }, + LoadArray { + data_idx: ArrayIdx, + len: Width, + dst: BufSlot, + }, + FillZero { + dst: BufSlot, + len: Width, + }, + LoadStateVector { + start: StateIdx, + width: Width, + dst: BufSlot, + }, + LoadStateVectorDot { + start: StateIdx, + width: Width, + dst: BufSlot, + }, + LoadInputParameter { + offset: InputOffset, + width: Width, + dst: BufSlot, + }, + LoadTangentState { + start: StateIdx, + width: Width, + dst: BufSlot, + }, + /// Load the parameter seed `dp[index]` into `dst`. + LoadTangentParameter { + index: ParamIdx, + dst: BufSlot, + }, + /// Element-wise binary operation: `dst = a b` with broadcast. + Binary { + op: BinaryOp, + a: BufSlot, + b: BufSlot, + dst: BufSlot, + len: Width, + kind: BroadcastKind, + }, + /// Element-wise unary operation: `dst = op(src)`. + Unary { + op: UnaryOp, + src: BufSlot, + dst: BufSlot, + len: Width, + }, + /// Reduce to a scalar: `dst = max(src[..src_len])`. + MaxReduce { + src: BufSlot, + src_len: Width, + dst: BufSlot, + }, + /// Reduce to a scalar: `dst = min(src[..src_len])`. + MinReduce { + src: BufSlot, + src_len: Width, + dst: BufSlot, + }, + /// Reduce subgradient: `dst = basis[k]`, `k` = argmax/argmin of + /// `picker[..len]` (first occurrence, strict comparison). + ReduceArgSelect { + basis_src: BufSlot, + picker_src: BufSlot, + len: Width, + is_max: bool, + dst: BufSlot, + }, + /// Slice: `dst = src[start..start+len]`. + Index { + src: BufSlot, + start: SlotOffset, + dst: BufSlot, + len: Width, + }, + /// Concatenate `sources_len` slot ranges into `dst`. + Concat { + sources_idx: SlotTableIdx, + sources_len: EntryCount, + dst: BufSlot, + }, + /// Sparse matrix-vector multiply: `dst = CSR(csr_idx) @ vec(vec_src)`. + MatMul { + csr_idx: CsrIdx, + vec_src: BufSlot, + dst: BufSlot, + }, + /// Dense matrix-vector multiply: `dst = A @ v`, A row-major at slot `mat_src` (rows × cols). + DenseMatMul { + mat_src: BufSlot, + rows: Rows, + cols: Cols, + vec_src: BufSlot, + dst: BufSlot, + }, + /// 1-D linear interpolation: `dst[i] = interp(src[i])`. + Interp1DLinear { + interp_idx: InterpolantIdx, + src: BufSlot, + dst: BufSlot, + len: Width, + }, + /// 1-D linear interpolation derivative, read from the table's slopes. + Interp1DLinearDeriv { + interp_idx: InterpolantIdx, + src: BufSlot, + dst: BufSlot, + len: Width, + }, + /// 1-D cubic (pchip) interpolation. + Interp1DCubic { + interp_idx: CubicInterpolantIdx, + src: BufSlot, + dst: BufSlot, + len: Width, + }, + /// 1-D cubic (pchip) interpolation derivative. + Interp1DCubicDeriv { + interp_idx: CubicInterpolantIdx, + src: BufSlot, + dst: BufSlot, + len: Width, + }, + /// N-D tensor-product interpolation. + InterpNd { + interp_idx: NdInterpolantIdx, + sources_idx: SlotTableIdx, + dst: BufSlot, + len: Width, + }, + /// N-D tensor-product interpolation partial derivative along `axis`. + InterpNdPartial { + interp_idx: NdInterpolantIdx, + sources_idx: SlotTableIdx, + axis: Axis, + dst: BufSlot, + len: Width, + }, + /// Piecewise conditional: select branches by `selector` value. + Conditional { + selector: BufSlot, + branches_idx: SlotTableIdx, + branches_len: EntryCount, + dst: BufSlot, + out_len: Width, + }, + /// Skip all but the active branch's instruction block. + /// + /// Precedes the blocks it guards; their [`BlockSpan`]s are relative to this instruction. + Dispatch { + selector: BufSlot, + blocks_idx: BlockSpanIdx, + blocks_len: EntryCount, + /// Instructions this `Dispatch` and its blocks occupy, so evaluation resumes past them. + span_len: InstructionCount, + }, +} + +// The interpreter streams the tape; keep each instruction to 24 bytes. +const _: () = assert!(size_of::() == 24); + +/// Side tables an instruction addresses by index instead of carrying inline. +#[derive(Debug, Default, Clone)] +pub struct ConstPool { + /// Flattened array data; `array_offsets[i]` gives the start of array `i`. + pub array_data: Vec, + pub array_offsets: Vec, + + pub csr_data: Vec, + + pub interpolants: Vec, + + pub cubic_interpolants: Vec, + + pub nd_interpolants: Vec, + + /// Concat source windows, one per operand. + pub concat_sources: Vec, + + /// Conditional branch windows, one per branch. + pub branch_offsets: Vec, + + /// `Dispatch` branch blocks, in instruction space. + pub branch_blocks: Vec, + + /// N-D interpolant child windows, one per input axis. + pub interp_nd_sources: Vec, +} + +impl ConstPool { + pub(crate) fn new() -> Self { + Self::default() + } + + pub(crate) fn add_array(&mut self, data: &[f64]) -> ArrayIdx { + let idx = ArrayIdx::of_next(&self.array_offsets); + self.array_offsets.push(self.array_data.len()); + self.array_data.extend_from_slice(data); + idx + } + + pub(crate) fn get_array(&self, idx: ArrayIdx, len: Width) -> &[f64] { + let start = self.array_offsets[idx.index()]; + &self.array_data[start..start + len.index()] + } + + pub(crate) fn add_csr(&mut self, csr: &CsrData) -> CsrIdx { + let idx = CsrIdx::of_next(&self.csr_data); + self.csr_data.push(csr.clone()); + idx + } + + pub(crate) fn add_interpolant(&mut self, data: &InterpolantData) -> InterpolantIdx { + let idx = InterpolantIdx::of_next(&self.interpolants); + self.interpolants.push(data.clone()); + idx + } + + pub(crate) fn add_cubic_interpolant( + &mut self, + data: &CubicInterpolantData, + ) -> CubicInterpolantIdx { + let idx = CubicInterpolantIdx::of_next(&self.cubic_interpolants); + self.cubic_interpolants.push(data.clone()); + idx + } + + pub(crate) fn add_nd_interpolant(&mut self, data: &NdInterpolantData) -> NdInterpolantIdx { + let idx = NdInterpolantIdx::of_next(&self.nd_interpolants); + self.nd_interpolants.push(data.clone()); + idx + } +} + +/// Partition of a split-eval tape: primal instructions and slots precede the tangent ones. +#[derive(Debug, Clone, Copy)] +pub struct SplitEvalInfo { + /// Index into the instruction stream where tangent instructions begin. + pub primal_end: usize, + /// Buffer region `[0, primal_buffer_size)` holds primal values. + pub primal_buffer_size: usize, +} + +/// Typed intermediate representation: a flat instruction tape with resolved shapes. +#[derive(Debug, Clone)] +pub struct TypedIr { + instructions: Vec, + + root_slot: Slot, + + buffer_size: usize, + + consts: ConstPool, +} + +impl TypedIr { + /// Compile an expression DAG into the IR, privatising conditional branch cones first. + pub fn from_dag(dag: DagRef<'_>) -> Self { + Self::privatise_then(dag, IRBuilder::build) + } + + /// [`from_dag`](Self::from_dag) without branch privatisation. + #[cfg(test)] + pub(crate) fn from_dag_raw(dag: DagRef<'_>) -> Self { + IRBuilder::build(dag.arena(), dag.root()) + } + + /// Run [`privatise_conditionals`] on `dag`, then lower the result with `build`. + fn privatise_then(dag: DagRef<'_>, build: impl FnOnce(&Arena, NodeId) -> T) -> T { + match privatise_conditionals(dag) { + Some(privatised) => build(privatised.arena(), privatised.root()), + None => build(dag.arena(), dag.root()), + } + } + + /// Output length of the root expression. + #[inline] + pub const fn output_len(&self) -> usize { + self.root_slot.len.index() + } + + /// Instruction stream in evaluation order. + #[inline] + pub fn instructions(&self) -> &[Instruction] { + &self.instructions + } + + /// Per-branch block lengths in tape order, one per branch of every `Dispatch`. + #[must_use] + pub fn branch_block_lens(&self) -> Vec { + self.block_lens().collect() + } + + /// Lowering emits one `branch_blocks` entry per branch of each `Dispatch`, in tape order. + fn block_lens(&self) -> impl Iterator + '_ { + self.consts + .branch_blocks + .iter() + .copied() + .map(BlockSpan::len) + } + + /// Number of `Dispatch` instructions, one per short-circuited conditional. + #[must_use] + pub fn dispatch_count(&self) -> usize { + self.instructions + .iter() + .filter(|instr| matches!(instr, Instruction::Dispatch { .. })) + .count() + } + + /// Instruction count excluding conditional branch blocks (the `Dispatch`es included). + #[must_use] + pub fn common_instruction_count(&self) -> usize { + self.instructions.len() - self.block_lens().sum::() + } + + /// Output slot of the root expression. + #[inline] + pub(crate) const fn root_slot(&self) -> Slot { + self.root_slot + } + + /// Total evaluation buffer size in elements. + #[inline] + pub const fn buffer_size(&self) -> usize { + self.buffer_size + } + + #[inline] + pub const fn consts(&self) -> &ConstPool { + &self.consts + } + + /// Compile with a primal/tangent split, so only the tangent section reruns per colour. + pub(crate) fn from_dag_split_eval(dag: DagRef<'_>) -> (Self, SplitEvalInfo) { + Self::privatise_then(dag, IRBuilder::build_split_eval) + } + + /// Compile with a no-reuse (SSA) slot layout, so the scratch is the reverse-AD value tape. + pub fn from_dag_pinned(dag: DagRef<'_>) -> Self { + Self::privatise_then(dag, IRBuilder::build_pinned) + } +} + +/// Builder for constructing `TypedIr` from an expression DAG. +struct IRBuilder { + instructions: Vec, + consts: ConstPool, + /// The buffer window of each `NodeId` in the arena. + slots: Vec, +} + +impl IRBuilder { + /// A builder over the `(offset, len)` windows slot assignment produced. + fn new(slots: Vec<(usize, usize)>, n_nodes: usize) -> Self { + Self { + instructions: Vec::with_capacity(n_nodes), + consts: ConstPool::new(), + slots: slots + .into_iter() + .map(|(offset, len)| Slot::from_usize(offset, len)) + .collect(), + } + } + + fn build(arena: &Arena, root: NodeId) -> TypedIr { + Self::build_scheduled(arena, root, |sizes, order| { + assign_slots(arena, sizes, order, &[root]) + }) + } + + /// Like [`Self::build`] but never recycles a slot (see [`assign_slots_pinned`]). + fn build_pinned(arena: &Arena, root: NodeId) -> TypedIr { + Self::build_scheduled(arena, root, assign_slots_pinned) + } + + /// Lower in one region schedule, with slots from `assign(sizes, order)`. + fn build_scheduled( + arena: &Arena, + root: NodeId, + assign: impl FnOnce(&[usize], &[NodeId]) -> (Vec<(usize, usize)>, usize), + ) -> TypedIr { + let base_order = arena.topological_order(root); + let schedule = schedule_regions(arena, &base_order); + let sizes = infer_sizes(arena, &schedule.order); + let (slots, buffer_size) = assign(&sizes, &schedule.order); + + let mut builder = Self::new(slots, schedule.order.len()); + builder.emit_scheduled(arena, &schedule.order, &schedule.groups); + builder.assert_block_slots_private(buffer_size); + builder.finish(root, buffer_size) + } + + fn build_split_eval(arena: &Arena, root: NodeId) -> (TypedIr, SplitEvalInfo) { + let eval_order = arena.topological_order(root); + let sizes = infer_sizes(arena, &eval_order); + let is_tangent = classify_tangent_nodes(arena, &eval_order); + + let mut primal_order = Vec::new(); + let mut tangent_order = Vec::new(); + for &nid in &eval_order { + if is_tangent[nid.index()] { + tangent_order.push(nid); + } else { + primal_order.push(nid); + } + } + // Each half schedules alone, so a `Dispatch` and its blocks share a half. + let owners = region_owners(arena, &eval_order); + let schedule = |half: &[NodeId]| { + schedule_regions_partitioned(arena, half, &eval_order, owners.as_deref()) + }; + let (primal_schedule, tangent_schedule) = + (schedule(&primal_order), schedule(&tangent_order)); + let partitioned_order: Vec = primal_schedule + .order + .iter() + .chain(tangent_schedule.order.iter()) + .copied() + .collect(); + // Whole-tape positions pin primal values the tangent half reads; tangent ones are local. + let primal_last_use = last_use_positions(arena, &partitioned_order, &[root]); + let tangent_last_use = last_use_positions(arena, &tangent_schedule.order, &[root]); + + // Primal slots occupy `[0, primal_hw)`, tangent slots `[primal_hw, total_hw)`. + let (primal_slots, primal_hw) = + assign_slots_configured(&sizes, &primal_schedule.order, &primal_last_use, 0); + + let (tangent_slots, total_hw) = assign_slots_configured( + &sizes, + &tangent_schedule.order, + &tangent_last_use, + primal_hw, + ); + + let mut slots = vec![(0_usize, 0_usize); sizes.len()]; + for &nid in &primal_schedule.order { + slots[nid.index()] = primal_slots[nid.index()]; + } + for &nid in &tangent_schedule.order { + slots[nid.index()] = tangent_slots[nid.index()]; + } + + let mut builder = Self::new(slots, partitioned_order.len()); + builder.emit_scheduled(arena, &primal_schedule.order, &primal_schedule.groups); + let primal_end = builder.instructions.len(); + builder.emit_scheduled(arena, &tangent_schedule.order, &tangent_schedule.groups); + builder.assert_no_block_straddles(primal_end); + builder.assert_block_slots_private(total_hw); + let split = SplitEvalInfo { + primal_end, + primal_buffer_size: primal_hw, + }; + (builder.finish(root, total_hw), split) + } + + /// The finished tape, whose result is `root`'s slot. + fn finish(self, root: NodeId, buffer_size: usize) -> TypedIr { + TypedIr { + root_slot: self.slot_for(root), + instructions: self.instructions, + buffer_size, + consts: self.consts, + } + } + + /// The buffer window allocated for a node's value. + fn slot_for(&self, id: NodeId) -> Slot { + self.slots[id.index()] + } + + /// Emit `order`, wrapping each annotated group range in a `Dispatch`. + fn emit_scheduled(&mut self, arena: &Arena, order: &[NodeId], groups: &[RegionGroup]) { + let mut group_at: Vec> = vec![None; order.len() + 1]; + for g in groups { + assert!( + group_at[g.anchor].is_none(), + "two region groups share anchor {}", + g.anchor + ); + group_at[g.anchor] = Some(g); + } + + let mut pos = 0; + while pos < order.len() { + if let Some(group) = group_at[pos] { + let total: usize = group.branch_lens.iter().sum(); + assert!(total > 0, "an empty region group must not be recorded"); + self.emit_dispatch(arena, group, &order[pos..pos + total]); + pos += total; + continue; + } + self.emit_node(arena, order[pos]); + pos += 1; + } + } + + /// Emit a `Dispatch` followed by one contiguous block per branch. + /// + /// Block lengths count instructions, not nodes, since `Node::SparseMatrix` emits none. + fn emit_dispatch(&mut self, arena: &Arena, group: &RegionGroup, nodes: &[NodeId]) { + let Node::Conditional { selector, .. } = arena.get(group.cond) else { + unreachable!("a region group is always anchored on a Conditional") + }; + let selector_slot = self.slot_for(*selector).offset; + let dispatch_at = self.instructions.len(); + let blocks_idx = BlockSpanIdx::of_next(&self.consts.branch_blocks); + // Reserve the table so the blocks can be emitted before it is filled. + self.consts.branch_blocks.resize( + blocks_idx.index() + group.branch_lens.len(), + BlockSpan::new(0, 0), + ); + self.instructions.push(Instruction::Dispatch { + selector: selector_slot, + blocks_idx, + blocks_len: EntryCount::from_usize(group.branch_lens.len()), + span_len: InstructionCount::new(1), + }); + + let mut taken = 0; + for (i, &n_nodes) in group.branch_lens.iter().enumerate() { + let start = self.instructions.len(); + for &node_id in &nodes[taken..taken + n_nodes] { + self.emit_node(arena, node_id); + } + taken += n_nodes; + let len = self.instructions.len() - start; + self.consts.branch_blocks[blocks_idx.index() + i] = + BlockSpan::from_usize(start - dispatch_at, len); + } + assert_eq!(taken, nodes.len(), "group range and branch_lens disagree"); + let end = self.instructions.len(); + let Instruction::Dispatch { span_len, .. } = &mut self.instructions[dispatch_at] else { + unreachable!("the Dispatch was pushed above") + }; + *span_len = InstructionCount::from_usize(end - dispatch_at); + } + + /// Assert every `Dispatch` block lies wholly on one side of `split`. + /// + /// A straddling block would run unconditionally in the tangent section. + fn assert_no_block_straddles(&self, split: usize) { + for (pc, instr) in self.instructions.iter().enumerate() { + if let Instruction::Dispatch { + blocks_idx, + blocks_len, + .. + } = *instr + { + for b in 0..blocks_len.index() { + let block = self.consts.branch_blocks[blocks_idx.index() + b].bounds(pc); + assert!( + split <= block.start || split >= block.end, + "split point {split} straddles block [{}, {})", + block.start, + block.end + ); + } + } + } + } + + /// Which instructions sit inside a branch block. + /// + /// Panics on a nested `Dispatch`, as evaluators skip only one level of block. + fn block_instruction_mask(&self) -> Vec { + let mut in_block = vec![false; self.instructions.len()]; + for (pc, instr) in self.instructions.iter().enumerate() { + if let Instruction::Dispatch { + blocks_idx, + blocks_len, + .. + } = *instr + { + for b in 0..blocks_len.index() { + let block = self.consts.branch_blocks[blocks_idx.index() + b].bounds(pc); + for at in block { + assert!( + !matches!(self.instructions[at], Instruction::Dispatch { .. }), + "instruction {at} is a Dispatch inside the block of the Dispatch at \ + {pc}: branch blocks must be flat" + ); + in_block[at] = true; + } + } + } + } + in_block + } + + /// Assert no instruction outside a branch block reads a value a block defined. + /// + /// That would silently read a stale recycled slot, so this runs in release too. + fn assert_block_slots_private(&self, buffer_size: usize) { + if self.consts.branch_blocks.is_empty() { + return; + } + let in_block = self.block_instruction_mask(); + let mut last_writer: Vec> = vec![None; buffer_size]; + + for (pc, instr) in self.instructions.iter().enumerate() { + if !in_block[pc] { + if let Instruction::Dispatch { selector, .. } = *instr { + assert!( + last_writer[selector.index()].is_some(), + "Dispatch at {pc} selects on buffer element {}, which no earlier \ + instruction defines", + selector.index() + ); + } + for_each_src_extent(instr, &self.consts, |src| { + // A `Conditional`'s branch slots are the one legitimate outside read. + if src.is_branch_slot { + return; + } + let start = src.slot.offset_usize(); + let end = start + src.slot.len_usize(); + for (index, writer) in last_writer[start..end].iter().enumerate() { + if let Some(writer) = *writer + && in_block[writer] + { + panic!( + "instruction {pc} ({instr:?}) outside a branch block reads buffer \ + element {}, defined by block-private instruction {writer} ({:?})", + start + index, + self.instructions[writer] + ); + } + } + }); + } + if let Some(dst) = instruction_dst_extent(instr, &self.consts) { + for writer in &mut last_writer[dst.bounds()] { + *writer = Some(pc); + } + } + } + } + + fn emit_node(&mut self, arena: &Arena, node_id: NodeId) { + let Slot { + offset: dst, + len: out_len, + } = self.slot_for(node_id); + + let instr = match &arena[node_id] { + Node::Scalar(v) => Instruction::LoadScalar { value: *v, dst }, + Node::Time => Instruction::LoadTime { dst }, + Node::Array(arr) => { + let data_idx = self.consts.add_array(&arr.data); + Instruction::LoadArray { + data_idx, + len: Width::from_usize(arr.data.len()), + dst, + } + }, + Node::ZeroVector { len } => Instruction::FillZero { + dst, + len: Width::from_usize(*len), + }, + Node::StateVector { start, end } => Instruction::LoadStateVector { + start: StateIdx::from_usize(*start), + width: state_width(*start, *end), + dst, + }, + Node::StateVectorDot { start, end } => Instruction::LoadStateVectorDot { + start: StateIdx::from_usize(*start), + width: state_width(*start, *end), + dst, + }, + Node::InputParameter { offset, width, .. } => Instruction::LoadInputParameter { + offset: InputOffset::from_usize(*offset), + width: Width::from_usize(*width), + dst, + }, + Node::TangentStateVector { start, end } => Instruction::LoadTangentState { + start: StateIdx::from_usize(*start), + width: state_width(*start, *end), + dst, + }, + Node::TangentParameter { index } => Instruction::LoadTangentParameter { + index: ParamIdx::from_usize(*index), + dst, + }, + Node::SparseMatrix(_) => return, // Not emitted as instruction + + Node::Add(a, b) => self.emit_binary(BinaryOp::Add, *a, *b, dst, out_len), + Node::Sub(a, b) => self.emit_binary(BinaryOp::Sub, *a, *b, dst, out_len), + Node::Mul(a, b) => self.emit_binary(BinaryOp::Mul, *a, *b, dst, out_len), + Node::Div(a, b) => self.emit_binary(BinaryOp::Div, *a, *b, dst, out_len), + Node::Pow(a, b) => self.emit_binary(BinaryOp::Pow, *a, *b, dst, out_len), + Node::Minimum(a, b) => self.emit_binary(BinaryOp::Minimum, *a, *b, dst, out_len), + Node::Maximum(a, b) => self.emit_binary(BinaryOp::Maximum, *a, *b, dst, out_len), + Node::Modulo(a, b) => self.emit_binary(BinaryOp::Modulo, *a, *b, dst, out_len), + Node::Hypot(a, b) => self.emit_binary(BinaryOp::Hypot, *a, *b, dst, out_len), + Node::EqualHeaviside(a, b) => { + self.emit_binary(BinaryOp::EqualHeaviside, *a, *b, dst, out_len) + }, + Node::NotEqualHeaviside(a, b) => { + self.emit_binary(BinaryOp::NotEqualHeaviside, *a, *b, dst, out_len) + }, + Node::Equality(a, b) => self.emit_binary(BinaryOp::Equality, *a, *b, dst, out_len), + + Node::Neg(a) => self.emit_unary(UnaryOp::Neg, *a, dst, out_len), + Node::Abs(a) => self.emit_unary(UnaryOp::Abs, *a, dst, out_len), + Node::Sqrt(a) => self.emit_unary(UnaryOp::Sqrt, *a, dst, out_len), + Node::Exp(a) => self.emit_unary(UnaryOp::Exp, *a, dst, out_len), + Node::Log(a) => self.emit_unary(UnaryOp::Log, *a, dst, out_len), + Node::Sin(a) => self.emit_unary(UnaryOp::Sin, *a, dst, out_len), + Node::Cos(a) => self.emit_unary(UnaryOp::Cos, *a, dst, out_len), + Node::Tanh(a) => self.emit_unary(UnaryOp::Tanh, *a, dst, out_len), + Node::Sinh(a) => self.emit_unary(UnaryOp::Sinh, *a, dst, out_len), + Node::Cosh(a) => self.emit_unary(UnaryOp::Cosh, *a, dst, out_len), + Node::Arcsinh(a) => self.emit_unary(UnaryOp::Arcsinh, *a, dst, out_len), + Node::Arctan(a) => self.emit_unary(UnaryOp::Arctan, *a, dst, out_len), + Node::Erf(a) => self.emit_unary(UnaryOp::Erf, *a, dst, out_len), + Node::Sign(a) => self.emit_unary(UnaryOp::Sign, *a, dst, out_len), + Node::Floor(a) => self.emit_unary(UnaryOp::Floor, *a, dst, out_len), + Node::Ceiling(a) => self.emit_unary(UnaryOp::Ceiling, *a, dst, out_len), + + Node::MaxReduce(a) => { + let src = self.slot_for(*a); + Instruction::MaxReduce { + src: src.offset, + src_len: src.len, + dst, + } + }, + Node::MinReduce(a) => { + let src = self.slot_for(*a); + Instruction::MinReduce { + src: src.offset, + src_len: src.len, + dst, + } + }, + Node::ReduceArgSelect { + basis, + picker, + is_max, + } => { + // `infer_sizes` checked the scan's `picker[0..len]` matches the basis width. + let basis_slot = self.slot_for(*basis); + Instruction::ReduceArgSelect { + basis_src: basis_slot.offset, + picker_src: self.slot_for(*picker).offset, + len: basis_slot.len, + is_max: *is_max, + dst, + } + }, + + Node::Index { child, start, .. } => Instruction::Index { + src: self.slot_for(*child).offset, + start: SlotOffset::from_usize(*start), + dst, + len: out_len, + }, + Node::Concat(children) => { + let sources_idx = SlotTableIdx::of_next(&self.consts.concat_sources); + for c in children { + let slot = self.slot_for(*c); + self.consts.concat_sources.push(slot); + } + Instruction::Concat { + sources_idx, + sources_len: EntryCount::from_usize(children.len()), + dst, + } + }, + + Node::MatMul(a, b) => match &arena[*a] { + Node::SparseMatrix(csr) => { + let csr_idx = self.consts.add_csr(csr); + Instruction::MatMul { + csr_idx, + vec_src: self.slot_for(*b).offset, + dst, + } + }, + Node::Array(arr) => Instruction::DenseMatMul { + mat_src: self.slot_for(*a).offset, + rows: Rows::from_usize(arr.shape.rows), + cols: Cols::from_usize(arr.shape.cols), + vec_src: self.slot_for(*b).offset, + dst, + }, + _ => panic!("MatMul requires a constant matrix on the left"), + }, + + Node::Interpolant1DLinear { data, child } => { + let interp_idx = self.consts.add_interpolant(data); + Instruction::Interp1DLinear { + interp_idx, + src: self.slot_for(*child).offset, + dst, + len: out_len, + } + }, + Node::Interpolant1DLinearDeriv { data, child } => { + let interp_idx = self.consts.add_interpolant(data); + Instruction::Interp1DLinearDeriv { + interp_idx, + src: self.slot_for(*child).offset, + dst, + len: out_len, + } + }, + Node::Interpolant1DCubic { data, child } => { + let interp_idx = self.consts.add_cubic_interpolant(data); + Instruction::Interp1DCubic { + interp_idx, + src: self.slot_for(*child).offset, + dst, + len: out_len, + } + }, + Node::Interpolant1DCubicDeriv { data, child } => { + let interp_idx = self.consts.add_cubic_interpolant(data); + Instruction::Interp1DCubicDeriv { + interp_idx, + src: self.slot_for(*child).offset, + dst, + len: out_len, + } + }, + Node::InterpolantNd { data, children } => { + let interp_idx = self.consts.add_nd_interpolant(data); + let sources_idx = SlotTableIdx::of_next(&self.consts.interp_nd_sources); + for c in children { + let slot = self.slot_for(*c); + self.consts.interp_nd_sources.push(slot); + } + Instruction::InterpNd { + interp_idx, + sources_idx, + dst, + len: out_len, + } + }, + Node::InterpolantNdPartial { + data, + children, + axis, + } => { + let interp_idx = self.consts.add_nd_interpolant(data); + let sources_idx = SlotTableIdx::of_next(&self.consts.interp_nd_sources); + for c in children { + let slot = self.slot_for(*c); + self.consts.interp_nd_sources.push(slot); + } + Instruction::InterpNdPartial { + interp_idx, + sources_idx, + axis: Axis::new(*axis), + dst, + len: out_len, + } + }, + + Node::Conditional { selector, branches } => { + let branches_idx = SlotTableIdx::of_next(&self.consts.branch_offsets); + for b in branches { + let slot = self.slot_for(*b); + self.consts.branch_offsets.push(slot); + } + Instruction::Conditional { + selector: self.slot_for(*selector).offset, + branches_idx, + branches_len: EntryCount::from_usize(branches.len()), + dst, + out_len, + } + }, + }; + + self.instructions.push(instr); + } + + fn emit_binary( + &self, + op: BinaryOp, + a: NodeId, + b: NodeId, + dst: BufSlot, + len: Width, + ) -> Instruction { + let a_slot = self.slot_for(a); + let b_slot = self.slot_for(b); + Instruction::Binary { + op, + a: a_slot.offset, + b: b_slot.offset, + dst, + len, + kind: BroadcastKind::from_lens(a_slot.len_usize(), b_slot.len_usize()), + } + } + + fn emit_unary(&self, op: UnaryOp, a: NodeId, dst: BufSlot, len: Width) -> Instruction { + let src = self.slot_for(a).offset; + Instruction::Unary { op, src, dst, len } + } +} + +#[cfg(test)] +impl IRBuilder { + /// A tape where an instruction outside the `Dispatch` block reads that block's slot. + fn test_tape_reading_across_block_boundary() -> Self { + let mut consts = ConstPool::new(); + consts.branch_blocks.push(BlockSpan::new(1, 1)); + + let instructions = vec![ + Instruction::LoadScalar { + value: 1.0, + dst: BufSlot::new(0), + }, + Instruction::Dispatch { + selector: BufSlot::new(0), + blocks_idx: BlockSpanIdx::new(0), + blocks_len: EntryCount::new(1), + span_len: InstructionCount::new(2), + }, + Instruction::LoadScalar { + value: 1.0, + dst: BufSlot::new(5), + }, + Instruction::Unary { + op: UnaryOp::Neg, + src: BufSlot::new(5), // outside the block, illegally reads the block's slot + dst: BufSlot::new(6), + len: Width::new(1), + }, + ]; + + Self { + instructions, + consts, + ..Self::new(Vec::new(), 0) + } + } + + /// A tape whose `Dispatch` selects on a slot no instruction writes. + fn test_tape_dispatching_on_an_undefined_selector() -> Self { + let mut consts = ConstPool::new(); + consts.branch_blocks.push(BlockSpan::new(1, 1)); + + let instructions = vec![ + Instruction::Dispatch { + selector: BufSlot::new(0), // never written + blocks_idx: BlockSpanIdx::new(0), + blocks_len: EntryCount::new(1), + span_len: InstructionCount::new(2), + }, + Instruction::LoadScalar { + value: 1.0, + dst: BufSlot::new(5), + }, + ]; + + Self { + instructions, + consts, + ..Self::new(Vec::new(), 0) + } + } +} + +/// Width of the state slice `[start, end)`; an inverted extent here is a lowering bug. +fn state_width(start: usize, end: usize) -> Width { + Width::from_usize( + end.checked_sub(start) + .expect("inverted state extent reached lowering"), + ) +} + +/// Why [`validate`] rejects a graph. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum GraphError { + /// A construct lowering cannot handle: a non-constant `MatMul` left operand. + #[error("{0}")] + Unsupported(String), + /// A width relationship the evaluator relies on is broken. + #[error("{0}")] + Invalid(String), +} + +/// Check a caller-built graph at `root` can be lowered, in one walk of its topological order. +/// +/// The earliest unsupported node wins over any invalid one; otherwise the earliest invalid. +/// A node kind only the rewrite passes build is invalid here: the reads a forged tangent +/// load makes are not part of the extents callers check their buffers against. +pub fn validate(arena: &Arena, root: NodeId) -> Result<(), GraphError> { + let mut sizes = vec![0_usize; arena.len()]; + let mut invalid = None; + for node_id in arena.topological_order(root) { + let node = &arena[node_id]; + if let Node::MatMul(a, _) = node + && !matches!(&arena[*a], Node::SparseMatrix(_) | Node::Array(_)) + { + return Err(GraphError::Unsupported(format!( + "MatMul left operand must be a constant matrix, got {:?}", + arena[*a] + ))); + } + // Widths are meaningless past an invalid node, but a later unsupported one still wins. + if invalid.is_none() { + let width = caller_buildable(node_id, node) + .and_then(|()| checked_width(arena, node_id, node, &sizes)); + match width { + Ok(width) => sizes[node_id.index()] = width, + Err(message) => invalid = Some(message), + } + } + } + invalid.map_or(Ok(()), |message| Err(GraphError::Invalid(message))) +} + +/// Refuse a node kind only the rewrite passes build. +fn caller_buildable(node_id: NodeId, node: &Node) -> Result<(), String> { + let kind = match node { + Node::ZeroVector { .. } => "ZeroVector", + Node::ReduceArgSelect { .. } => "ReduceArgSelect", + Node::TangentStateVector { .. } => "TangentStateVector", + Node::TangentParameter { .. } => "TangentParameter", + Node::Interpolant1DLinearDeriv { .. } => "Interpolant1DLinearDeriv", + Node::Interpolant1DCubicDeriv { .. } => "Interpolant1DCubicDeriv", + Node::InterpolantNdPartial { .. } => "InterpolantNdPartial", + _ => return Ok(()), + }; + Err(format!( + "node {node_id:?} is an internal {kind} node, which no caller builds" + )) +} + +fn infer_sizes_checked(arena: &Arena, eval_order: &[NodeId]) -> Result, String> { + let mut sizes = vec![0usize; arena.len()]; + for &node_id in eval_order { + sizes[node_id.index()] = checked_width(arena, node_id, &arena[node_id], &sizes)?; + } + Ok(sizes) +} + +/// `node`'s width, once its relationships hold and it fits the tape's `u32` indices. +fn checked_width( + arena: &Arena, + node_id: NodeId, + node: &Node, + sizes: &[usize], +) -> Result { + check_widths(arena, node_id, node, sizes)?; + let width = node_width(arena, node, |child| sizes[child.index()]); + if u32::try_from(width).is_err() { + return Err(format!( + "node {node_id:?} has width {width}, past the tape's u32 index range" + )); + } + Ok(width) +} + +/// The first width relationship `node` breaks, given its children's `sizes`. +fn check_widths( + arena: &Arena, + node_id: NodeId, + node: &Node, + sizes: &[usize], +) -> Result<(), String> { + let width = |id: &NodeId| sizes[id.index()]; + match node { + Node::StateVector { start, end } + | Node::StateVectorDot { start, end } + | Node::TangentStateVector { start, end } + | Node::Index { start, end, .. } + if start > end => + { + Err(format!( + "node {node_id:?} has an inverted extent: start={start}, end={end}" + )) + }, + Node::StateVector { end, .. } + | Node::StateVectorDot { end, .. } + | Node::TangentStateVector { end, .. } + if u32::try_from(*end).is_err() => + { + Err(format!( + "node {node_id:?} reads states up to {end}, past the tape's u32 index range" + )) + }, + Node::InputParameter { offset, width, .. } + if offset + .checked_add(*width) + .is_none_or(|end| u32::try_from(end).is_err()) => + { + Err(format!( + "node {node_id:?} reads inputs from offset {offset} with width {width}, past \ + the tape's u32 index range" + )) + }, + Node::Index { child, end, .. } if *end > width(child) => Err(format!( + "Index node {node_id:?} ends at {end}, beyond child width {}", + width(child) + )), + Node::Add(a, b) + | Node::Sub(a, b) + | Node::Mul(a, b) + | Node::Div(a, b) + | Node::Pow(a, b) + | Node::Minimum(a, b) + | Node::Maximum(a, b) + | Node::Modulo(a, b) + | Node::Hypot(a, b) + | Node::EqualHeaviside(a, b) + | Node::NotEqualHeaviside(a, b) + | Node::Equality(a, b) + if !broadcast_widths_compatible(width(a), width(b)) => + { + Err(format!( + "binary node {node_id:?} has incompatible operand widths {} and {}; widths \ + must match or one operand must be scalar", + width(a), + width(b) + )) + }, + Node::MatMul(a, b) => { + let cols = match &arena[*a] { + Node::SparseMatrix(csr) => csr.shape.cols, + Node::Array(arr) => arr.shape.cols, + _ => return Err("MatMul requires a constant matrix on the left".to_string()), + }; + if width(b) == cols { + Ok(()) + } else { + Err(format!( + "MatMul node {node_id:?} has {cols} columns but its vector operand has \ + width {}", + width(b) + )) + } + }, + Node::InterpolantNd { data, children } + | Node::InterpolantNdPartial { data, children, .. } => { + let axes = data.breakpoints.len(); + if children.len() != axes { + return Err(format!( + "N-D interpolant node {node_id:?} has {} children for {axes} axes", + children.len() + )); + } + let out_len = children.iter().map(width).max().unwrap_or(1); + if let Some(child) = children + .iter() + .find(|c| width(c) != 1 && width(c) != out_len) + { + return Err(format!( + "N-D interpolant node {node_id:?} has child widths that cannot broadcast: \ + {} and {out_len}", + width(child) + )); + } + match node { + Node::InterpolantNdPartial { axis, .. } if *axis as usize >= axes => Err(format!( + "N-D interpolant partial node {node_id:?} has axis {axis} for {axes} axes" + )), + _ => Ok(()), + } + }, + Node::Conditional { selector, branches } => { + if width(selector) != 1 { + return Err(format!( + "Conditional node {node_id:?} requires a scalar selector, got width {}", + width(selector) + )); + } + let out_len = branches.first().map_or(1, width); + branches + .iter() + .find(|branch| width(branch) != out_len) + .map_or(Ok(()), |branch| { + Err(format!( + "Conditional node {node_id:?} has branch widths {out_len} and {}", + width(branch) + )) + }) + }, + Node::ReduceArgSelect { basis, picker, .. } => { + if width(basis) == 0 { + Err(format!( + "ReduceArgSelect node {node_id:?} requires non-empty operands" + )) + } else if width(basis) != width(picker) { + Err(format!( + "ReduceArgSelect node {node_id:?} has basis width {} and picker width {}", + width(basis), + width(picker) + )) + } else { + Ok(()) + } + }, + _ => Ok(()), + } +} + +/// `node`'s output width from its children's: the one rule lowering and every rewrite share. +/// +/// Assumes the relationships [`validate`] checks, e.g. a non-inverted extent. +pub(crate) fn node_width( + arena: &Arena, + node: &Node, + mut width_of: impl FnMut(NodeId) -> usize, +) -> usize { + match node { + Node::Scalar(_) + | Node::Time + | Node::TangentParameter { .. } + | Node::MaxReduce(_) + | Node::MinReduce(_) + | Node::ReduceArgSelect { .. } => 1, + Node::InputParameter { width, .. } => *width, + Node::Array(arr) => arr.data.len(), + Node::ZeroVector { len } => *len, + Node::SparseMatrix(_) => 0, + Node::StateVector { start, end } + | Node::StateVectorDot { start, end } + | Node::TangentStateVector { start, end } + | Node::Index { start, end, .. } => end - start, + Node::Add(a, b) + | Node::Sub(a, b) + | Node::Mul(a, b) + | Node::Div(a, b) + | Node::Pow(a, b) + | Node::Minimum(a, b) + | Node::Maximum(a, b) + | Node::Modulo(a, b) + | Node::Hypot(a, b) + | Node::EqualHeaviside(a, b) + | Node::NotEqualHeaviside(a, b) + | Node::Equality(a, b) => width_of(*a).max(width_of(*b)), + Node::MatMul(a, b) => match arena.get(*a) { + Node::SparseMatrix(csr) => csr.shape.rows, + Node::Array(arr) => arr.shape.rows, + _ => width_of(*b), + }, + Node::Concat(children) => children.iter().map(|&c| width_of(c)).sum(), + Node::Neg(a) + | Node::Abs(a) + | Node::Sqrt(a) + | Node::Exp(a) + | Node::Log(a) + | Node::Sin(a) + | Node::Cos(a) + | Node::Tanh(a) + | Node::Sinh(a) + | Node::Cosh(a) + | Node::Arcsinh(a) + | Node::Arctan(a) + | Node::Erf(a) + | Node::Sign(a) + | Node::Floor(a) + | Node::Ceiling(a) + | Node::Interpolant1DLinear { child: a, .. } + | Node::Interpolant1DLinearDeriv { child: a, .. } + | Node::Interpolant1DCubic { child: a, .. } + | Node::Interpolant1DCubicDeriv { child: a, .. } => width_of(*a), + Node::InterpolantNd { children, .. } + | Node::InterpolantNdPartial { children, .. } + | Node::Conditional { + branches: children, .. + } => children.iter().map(|&c| width_of(c)).max().unwrap_or(1), + } +} + +/// Infer the output size of each node in topological evaluation order. +pub fn infer_sizes(arena: &Arena, eval_order: &[NodeId]) -> Vec { + infer_sizes_checked(arena, eval_order) + .unwrap_or_else(|message| panic!("invalid expression graph: {message}")) +} + +/// One slot range an instruction reads. +#[derive(Clone, Copy, Debug)] +struct SrcExtent { + slot: Slot, + /// A `Conditional`'s branch slot, which may legitimately read block-owned values. + is_branch_slot: bool, +} + +impl SrcExtent { + const fn read(offset: BufSlot, len: Width) -> Self { + Self { + slot: Slot::new(offset, len), + is_branch_slot: false, + } + } + + const fn of(slot: Slot) -> Self { + Self { + slot, + is_branch_slot: false, + } + } + + const fn branch_slot(offset: BufSlot, len: Width) -> Self { + Self { + slot: Slot::new(offset, len), + is_branch_slot: true, + } + } +} + +/// The exact buffer window an instruction writes, or `None` for `Dispatch`. +fn instruction_dst_extent(instr: &Instruction, consts: &ConstPool) -> Option { + let extent = match *instr { + Instruction::LoadScalar { dst, .. } + | Instruction::LoadTime { dst } + | Instruction::LoadTangentParameter { dst, .. } + | Instruction::MaxReduce { dst, .. } + | Instruction::MinReduce { dst, .. } + | Instruction::ReduceArgSelect { dst, .. } => Slot::new(dst, Width::new(1)), + Instruction::LoadArray { dst, len, .. } + | Instruction::FillZero { dst, len } + | Instruction::Binary { dst, len, .. } + | Instruction::Unary { dst, len, .. } + | Instruction::Index { dst, len, .. } + | Instruction::Interp1DLinear { dst, len, .. } + | Instruction::Interp1DLinearDeriv { dst, len, .. } + | Instruction::Interp1DCubic { dst, len, .. } + | Instruction::Interp1DCubicDeriv { dst, len, .. } + | Instruction::InterpNd { dst, len, .. } + | Instruction::InterpNdPartial { dst, len, .. } => Slot::new(dst, len), + Instruction::LoadStateVector { width, dst, .. } + | Instruction::LoadStateVectorDot { width, dst, .. } + | Instruction::LoadTangentState { width, dst, .. } + | Instruction::LoadInputParameter { width, dst, .. } => Slot::new(dst, width), + Instruction::Concat { + sources_idx, + sources_len, + dst, + } => { + // Sum in `usize`: a `u32` sum of `u32` widths wraps in release. + let width: usize = (0..sources_len.index()) + .map(|i| consts.concat_sources[sources_idx.index() + i].len.index()) + .sum(); + Slot::new(dst, Width::from_usize(width)) + }, + Instruction::MatMul { csr_idx, dst, .. } => Slot::new( + dst, + Width::from_usize(consts.csr_data[csr_idx.index()].shape.rows), + ), + Instruction::DenseMatMul { rows, dst, .. } => Slot::new(dst, rows.to_width()), + Instruction::Conditional { dst, out_len, .. } => Slot::new(dst, out_len), + Instruction::Dispatch { .. } => return None, + }; + Some(extent) +} + +/// Hand each slot range `instr` reads to `visit`, with exact widths. +fn for_each_src_extent(instr: &Instruction, consts: &ConstPool, mut visit: impl FnMut(SrcExtent)) { + match *instr { + Instruction::LoadScalar { .. } + | Instruction::LoadTime { .. } + | Instruction::LoadArray { .. } + | Instruction::FillZero { .. } + | Instruction::LoadStateVector { .. } + | Instruction::LoadStateVectorDot { .. } + | Instruction::LoadInputParameter { .. } + | Instruction::LoadTangentState { .. } + | Instruction::LoadTangentParameter { .. } => {}, + Instruction::Binary { + a, b, len, kind, .. + } => { + let one = Width::new(1); + let (a_len, b_len) = match kind { + BroadcastKind::ScalarScalar => (one, one), + BroadcastKind::ScalarVector => (one, len), + BroadcastKind::VectorScalar => (len, one), + BroadcastKind::VectorVector => (len, len), + }; + visit(SrcExtent::read(a, a_len)); + visit(SrcExtent::read(b, b_len)); + }, + Instruction::Unary { src, len, .. } + | Instruction::Interp1DLinear { src, len, .. } + | Instruction::Interp1DLinearDeriv { src, len, .. } + | Instruction::Interp1DCubic { src, len, .. } + | Instruction::Interp1DCubicDeriv { src, len, .. } => visit(SrcExtent::read(src, len)), + Instruction::MaxReduce { src, src_len, .. } + | Instruction::MinReduce { src, src_len, .. } => visit(SrcExtent::read(src, src_len)), + Instruction::ReduceArgSelect { + basis_src, + picker_src, + len, + .. + } => { + visit(SrcExtent::read(basis_src, len)); + visit(SrcExtent::read(picker_src, len)); + }, + Instruction::Index { + src, start, len, .. + } => visit(SrcExtent::read( + BufSlot::from_usize(src.offset_by(start)), + len, + )), + Instruction::Concat { + sources_idx, + sources_len, + .. + } => { + for &slot in &consts.concat_sources[sources_idx.index()..][..sources_len.index()] { + visit(SrcExtent::of(slot)); + } + }, + Instruction::MatMul { + csr_idx, vec_src, .. + } => visit(SrcExtent::read( + vec_src, + Width::from_usize(consts.csr_data[csr_idx.index()].shape.cols), + )), + Instruction::DenseMatMul { + mat_src, + rows, + cols, + vec_src, + .. + } => { + visit(SrcExtent::read(mat_src, rows.times(cols))); + visit(SrcExtent::read(vec_src, cols.to_width())); + }, + Instruction::InterpNd { + interp_idx, + sources_idx, + .. + } + | Instruction::InterpNdPartial { + interp_idx, + sources_idx, + .. + } => { + let ndim = consts.nd_interpolants[interp_idx.index()].breakpoints.len(); + for &slot in &consts.interp_nd_sources[sources_idx.index()..][..ndim] { + visit(SrcExtent::of(slot)); + } + }, + Instruction::Conditional { + selector, + branches_idx, + branches_len, + out_len, + .. + } => { + visit(SrcExtent::read(selector, Width::new(1))); + for slot in &consts.branch_offsets[branches_idx.index()..][..branches_len.index()] { + visit(SrcExtent::branch_slot(slot.offset, out_len)); + } + }, + Instruction::Dispatch { selector, .. } => visit(SrcExtent::read(selector, Width::new(1))), + } +} + +/// Classify nodes as primal or tangent via forward taint propagation. +fn classify_tangent_nodes(arena: &Arena, eval_order: &[NodeId]) -> Vec { + let mut is_tangent = vec![false; arena.len()]; + + for &nid in eval_order { + let node = arena.get(nid); + let tainted = match node { + Node::TangentStateVector { .. } | Node::TangentParameter { .. } => true, + _ => { + let mut any_child_tangent = false; + node.for_each_child(|c| { + if is_tangent[c.index()] { + any_child_tangent = true; + } + }); + any_child_tangent + }, + }; + is_tangent[nid.index()] = tainted; + } + + is_tangent +} + +/// The last position in `order` reading each node, with every node in `roots` pinned. +fn last_use_positions(arena: &Arena, order: &[NodeId], roots: &[NodeId]) -> Vec { + let mut last_use = vec![0_usize; arena.len()]; + for (pos, &node_id) in order.iter().enumerate() { + arena + .get(node_id) + .for_each_child(|c| last_use[c.index()] = pos); + } + for &r in roots { + last_use[r.index()] = usize::MAX; + } + last_use +} + +/// Assign buffer slots, computing last-use positions from the evaluation order. +fn assign_slots( + arena: &Arena, + sizes: &[usize], + eval_order: &[NodeId], + roots: &[NodeId], +) -> (Vec<(usize, usize)>, usize) { + let last_use = last_use_positions(arena, eval_order, roots); + assign_slots_configured(sizes, eval_order, &last_use, 0) +} + +/// Slot allocation with no reuse, so the primal scratch preserves every intermediate. +fn assign_slots_pinned(sizes: &[usize], eval_order: &[NodeId]) -> (Vec<(usize, usize)>, usize) { + let last_use = vec![usize::MAX; sizes.len()]; + assign_slots_configured(sizes, eval_order, &last_use, 0) +} + +/// Sweep-line buffer slot allocator, reusing freed regions. +/// +/// `last_use[id]` is the last position reading the node (`usize::MAX` pins it), and +/// `initial_high_water` offsets every allocation. +fn assign_slots_configured( + sizes: &[usize], + eval_order: &[NodeId], + last_use: &[usize], + initial_high_water: usize, +) -> (Vec<(usize, usize)>, usize) { + use std::collections::{BTreeMap, BTreeSet}; + + let n = sizes.len(); + let mut slots = vec![(0_usize, 0_usize); n]; + + let mut by_offset: BTreeMap = BTreeMap::new(); + let mut by_size: BTreeSet<(usize, usize)> = BTreeSet::new(); + let mut high_water: usize = initial_high_water; + + let mut release_at: Vec> = vec![Vec::new(); eval_order.len() + 1]; + for &node_id in eval_order { + let lu = last_use[node_id.index()]; + if lu != usize::MAX && lu < eval_order.len() { + release_at[lu + 1].push(node_id); + } + } + + for (pos, &node_id) in eval_order.iter().enumerate() { + for &n_to_free in &release_at[pos] { + let (off, len) = slots[n_to_free.index()]; + if len == 0 { + continue; + } + + let mut new_off = off; + let mut new_len = len; + if let Some(&right_len) = by_offset.get(&(off + len)) { + by_offset.remove(&(off + len)); + by_size.remove(&(right_len, off + len)); + new_len += right_len; + } + if let Some((&left_off, &left_len)) = by_offset.range(..off).next_back() + && left_off + left_len == off + { + by_offset.remove(&left_off); + by_size.remove(&(left_len, left_off)); + new_off = left_off; + new_len += left_len; + } + by_offset.insert(new_off, new_len); + by_size.insert((new_len, new_off)); + } + + let len = sizes[node_id.index()]; + if len == 0 { + slots[node_id.index()] = (0, 0); + continue; + } + + let chosen = by_size.range((len, 0)..).next().copied(); + let off = if let Some((region_len, region_off)) = chosen { + by_size.remove(&(region_len, region_off)); + by_offset.remove(®ion_off); + if region_len > len { + let rem_off = region_off + len; + let rem_len = region_len - len; + by_offset.insert(rem_off, rem_len); + by_size.insert((rem_len, rem_off)); + } + region_off + } else { + let off = high_water; + high_water += len; + off + }; + slots[node_id.index()] = (off, len); + } + + (slots, high_water) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::eval::CompiledExpr; + use crate::node::{ArrayData, Shape}; + + #[test] + #[should_panic(expected = "outside a branch block reads buffer element")] + fn assert_block_slots_private_catches_a_cross_block_read() { + let builder = IRBuilder::test_tape_reading_across_block_boundary(); + builder.assert_block_slots_private(7); + } + + #[test] + #[should_panic(expected = "which no earlier instruction defines")] + fn assert_block_slots_private_catches_an_undefined_selector() { + let builder = IRBuilder::test_tape_dispatching_on_an_undefined_selector(); + builder.assert_block_slots_private(7); + } + + #[test] + fn test_typed_ir_from_expression() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let two = arena.alloc(Node::Scalar(2.0)); + let expr = arena.alloc(Node::Mul(two, y)); + + let ir = TypedIr::from_dag(DagRef::new(&arena, expr)); + + assert_eq!(ir.output_len(), 3); + assert_eq!(ir.instructions().len(), 3); // Scalar, StateVector, Mul + } + + #[test] + fn test_slot_struct() { + let slot = Slot::from_usize(10, 5); + assert_eq!(slot.offset, BufSlot::new(10)); + assert_eq!(slot.len, Width::new(5)); + assert_eq!(slot.offset_usize(), 10); + assert_eq!(slot.len_usize(), 5); + assert!(!slot.is_scalar()); + + let scalar_slot = Slot::from_usize(0, 1); + assert!(scalar_slot.is_scalar()); + } + + #[test] + fn test_broadcast_kind_inference() { + assert_eq!(BroadcastKind::from_lens(1, 1), BroadcastKind::ScalarScalar); + assert_eq!(BroadcastKind::from_lens(1, 5), BroadcastKind::ScalarVector); + assert_eq!(BroadcastKind::from_lens(5, 1), BroadcastKind::VectorScalar); + assert_eq!(BroadcastKind::from_lens(5, 5), BroadcastKind::VectorVector); + } + + #[test] + fn test_ir_nested_expression() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let two = arena.alloc(Node::Scalar(2.0)); + let one = arena.alloc(Node::Scalar(1.0)); + let mul = arena.alloc(Node::Mul(x, two)); + let add = arena.alloc(Node::Add(mul, one)); + let sin = arena.alloc(Node::Sin(add)); + + let ir = TypedIr::from_dag(DagRef::new(&arena, sin)); + + assert_eq!(ir.output_len(), 3); + assert_eq!(ir.instructions().len(), 6); + } + + #[test] + fn test_ir_concat() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let b = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![2.0, 3.0], + shape: Shape::vector(2), + }))); + let c = arena.alloc(Node::Scalar(4.0)); + let concat = arena.alloc(Node::Concat(vec![a, b, c])); + + let ir = TypedIr::from_dag(DagRef::new(&arena, concat)); + + assert_eq!(ir.output_len(), 4); + assert_eq!(ir.consts().concat_sources.len(), 3); + } + + /// The message [`validate`] rejects `root` with as invalid. + fn invalid_message(arena: &Arena, root: NodeId) -> String { + match validate(arena, root) { + Err(GraphError::Invalid(message)) => message, + other => panic!("expected an invalid graph, got {other:?}"), + } + } + + #[test] + fn validate_flags_a_non_constant_matmul_lhs_as_unsupported() { + let mut arena = Arena::new(); + let s = arena.alloc(Node::Scalar(2.0)); + let v = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0], + shape: Shape::vector(1), + }))); + let mm = arena.alloc(Node::MatMul(s, v)); + assert_eq!( + validate(&arena, mm), + Err(GraphError::Unsupported( + "MatMul left operand must be a constant matrix, got Scalar(2.0)".to_string() + )) + ); + assert_eq!(validate(&arena, v), Ok(())); + } + + #[test] + fn validate_reports_unsupported_over_an_earlier_invalid_node() { + // The inverted extent comes first in topological order; its width is never computed. + let mut arena = Arena::new(); + let inverted = arena.alloc(Node::StateVector { start: 3, end: 1 }); + let s = arena.alloc(Node::Scalar(2.0)); + let v = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let mm = arena.alloc(Node::MatMul(s, v)); + let root = arena.alloc(Node::Concat(vec![inverted, mm])); + assert!(matches!( + validate(&arena, root), + Err(GraphError::Unsupported(message)) if message.contains("constant matrix") + )); + + let only_invalid = arena.alloc(Node::Concat(vec![inverted, v])); + assert!(invalid_message(&arena, only_invalid).contains("inverted extent")); + } + + #[test] + fn test_symbolic_jacobian_width_mismatch_rejected_before_eval_overlap() { + let mut arena = Arena::new(); + let short = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let long = arena.alloc(Node::StateVector { start: 2, end: 5 }); + let add = arena.alloc(Node::Add(short, long)); + let scalar_entry = arena.alloc(Node::Scalar(1.0)); + let symbolic_jacobian = arena.alloc(Node::Concat(vec![add, scalar_entry])); + + // Unrejected, the width-2 operand's read window would overlap `dst`. + let message = invalid_message(&arena, symbolic_jacobian); + assert!(message.contains("incompatible operand widths 2 and 3")); + } + + #[test] + fn validate_flags_other_evaluator_width_hazards() { + let mut matmul_arena = Arena::new(); + let matrix = matmul_arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0; 6], + shape: Shape::matrix(2, 3), + }))); + let short = matmul_arena.alloc(Node::StateVector { start: 0, end: 2 }); + let matmul = matmul_arena.alloc(Node::MatMul(matrix, short)); + assert!( + invalid_message(&matmul_arena, matmul) + .contains("has 3 columns but its vector operand has width 2") + ); + + let mut index_arena = Arena::new(); + let child = index_arena.alloc(Node::StateVector { start: 0, end: 2 }); + let index = index_arena.alloc(Node::Index { + child, + start: 1, + end: 3, + }); + assert!(invalid_message(&index_arena, index).contains("beyond child width 2")); + + let mut conditional_arena = Arena::new(); + let selector = conditional_arena.alloc(Node::Scalar(1.0)); + let short = conditional_arena.alloc(Node::StateVector { start: 0, end: 2 }); + let long = conditional_arena.alloc(Node::StateVector { start: 2, end: 5 }); + let conditional = conditional_arena.alloc(Node::Conditional { + selector, + branches: vec![short, long], + }); + assert!(invalid_message(&conditional_arena, conditional).contains("branch widths 2 and 3")); + + let mut interpolant_arena = Arena::new(); + let short = interpolant_arena.alloc(Node::StateVector { start: 0, end: 2 }); + let long = interpolant_arena.alloc(Node::StateVector { start: 2, end: 5 }); + let interpolant = interpolant_arena.alloc(Node::InterpolantNd { + data: Box::new(NdInterpolantData { + breakpoints: vec![vec![0.0, 1.0], vec![0.0, 1.0]], + coeffs: vec![0.0; 4], + order: 2, + }), + children: vec![short, long], + }); + assert!( + invalid_message(&interpolant_arena, interpolant).contains("cannot broadcast: 2 and 3") + ); + + // Only the rewrite passes build these, so lowering's own width check is what guards them. + let lowering_message = |arena: &Arena, root| { + infer_sizes_checked(arena, &arena.topological_order(root)) + .expect_err("lowering must refuse the graph") + }; + let mut reduce_arena = Arena::new(); + let basis = reduce_arena.alloc(Node::StateVector { start: 0, end: 2 }); + let picker = reduce_arena.alloc(Node::StateVector { start: 2, end: 5 }); + let reduce = reduce_arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: true, + }); + assert!( + lowering_message(&reduce_arena, reduce).contains("basis width 2 and picker width 3") + ); + + let mut empty_reduce_arena = Arena::new(); + let basis = empty_reduce_arena.alloc(Node::ZeroVector { len: 0 }); + let picker = empty_reduce_arena.alloc(Node::ZeroVector { len: 0 }); + let reduce = empty_reduce_arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: true, + }); + assert!( + lowering_message(&empty_reduce_arena, reduce).contains("requires non-empty operands") + ); + } + + #[test] + fn validate_refuses_extents_and_widths_past_the_tapes_u32_indices() { + // Unrefused, the first three lower into a panic narrowing an index to u32. + let mut arena = Arena::new(); + let late = arena.alloc(Node::StateVector { + start: 1 << 32, + end: (1 << 32) + 1, + }); + assert!(invalid_message(&arena, late).contains("reads states up to 4294967297")); + let late_dot = arena.alloc(Node::StateVectorDot { + start: 0, + end: (1 << 32) + 1, + }); + assert!(invalid_message(&arena, late_dot).contains("reads states up to 4294967297")); + let late_input = arena.alloc(Node::InputParameter { + name: "far".to_string(), + index: 0, + offset: 1 << 32, + width: 1, + }); + assert!(invalid_message(&arena, late_input).contains("offset 4294967296 with width 1")); + // A slot ending past u32::MAX, which the parameter registry refuses too. + let overrunning_input = arena.alloc(Node::InputParameter { + name: "long".to_string(), + index: 1, + offset: u32::MAX as usize, + width: 2, + }); + assert!(invalid_message(&arena, overrunning_input).contains("past the tape's u32")); + + let half = arena.alloc(Node::StateVector { + start: 0, + end: (1 << 31) + 1, + }); + assert_eq!(validate(&arena, half), Ok(())); + let doubled = arena.alloc(Node::Concat(vec![half, half])); + assert!(invalid_message(&arena, doubled).contains("has width 4294967298")); + } + + #[test] + fn validate_refuses_every_node_kind_only_the_rewrite_passes_build() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x = arena.alloc(Node::Scalar(0.5)); + let linear = InterpolantData::try_new(vec![0.0, 1.0], vec![0.0, 1.0]).expect("valid"); + let cubic = CubicInterpolantData::try_new(vec![0.0, 1.0], vec![[0.0; 4]]).expect("valid"); + let nd = + NdInterpolantData::try_new(vec![vec![0.0, 1.0]; 2], vec![0.0; 4], 2).expect("valid"); + for (kind, node) in [ + ("ZeroVector", Node::ZeroVector { len: 1 }), + ( + "ReduceArgSelect", + Node::ReduceArgSelect { + basis: y, + picker: y, + is_max: true, + }, + ), + ( + "TangentStateVector", + Node::TangentStateVector { start: 0, end: 1 }, + ), + ("TangentParameter", Node::TangentParameter { index: 0 }), + ( + "Interpolant1DLinearDeriv", + Node::Interpolant1DLinearDeriv { + data: Box::new(linear), + child: x, + }, + ), + ( + "Interpolant1DCubicDeriv", + Node::Interpolant1DCubicDeriv { + data: Box::new(cubic), + child: x, + }, + ), + ( + "InterpolantNdPartial", + Node::InterpolantNdPartial { + data: Box::new(nd), + children: vec![x, x], + axis: 0, + }, + ), + ] { + let internal = arena.alloc(node); + let root = arena.alloc(Node::Neg(internal)); + assert!( + invalid_message(&arena, root).contains(&format!("internal {kind} node")), + "{kind} must be refused" + ); + } + } + + #[test] + fn test_ir_dense_matmul_output_len() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], // row-major 2x3 + shape: Shape::matrix(2, 3), + }))); + let v = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 1.0, 1.0], + shape: Shape::vector(3), + }))); + let matmul = arena.alloc(Node::MatMul(a, v)); + let ir = TypedIr::from_dag(DagRef::new(&arena, matmul)); + assert_eq!(ir.output_len(), 2); + } + + #[test] + fn test_ir_interpolant() { + use crate::node::InterpolantData; + + let mut arena = Arena::new(); + let x = arena.alloc(Node::Scalar(1.5)); + let interp = arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + InterpolantData::try_new(vec![0.0, 1.0, 2.0], vec![0.0, 10.0, 20.0]) + .expect("valid interpolant"), + ), + child: x, + }); + + let ir = TypedIr::from_dag(DagRef::new(&arena, interp)); + + assert_eq!(ir.output_len(), 1); + assert_eq!(ir.consts().interpolants.len(), 1); + assert_eq!(ir.consts().interpolants[0].x_data, vec![0.0, 1.0, 2.0]); + } + + #[test] + fn slope_lookups_lower_against_the_value_pool() { + use crate::node::InterpolantData; + + let values = InterpolantData::try_new(vec![0.0, 1.0, 3.0], vec![0.0, 10.0, 40.0]) + .expect("valid interpolant"); + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Interpolant1DLinearDeriv { + data: Box::new(values), + child: x, + }); + let ir = TypedIr::from_dag(DagRef::new(&arena, root)); + + assert_eq!(ir.consts().interpolants.len(), 1); + let table = &ir.consts().interpolants[0]; + assert_eq!(table.x_data.len(), table.slopes.len() + 1); + assert_eq!(table.slopes, &[10.0, 15.0]); + assert!(matches!( + ir.instructions()[1], + Instruction::Interp1DLinearDeriv { interp_idx, .. } + if interp_idx == InterpolantIdx::new(0) + )); + } + + #[test] + fn test_const_pool() { + let mut pool = ConstPool::new(); + + let idx1 = pool.add_array(&[1.0, 2.0, 3.0]); + let idx2 = pool.add_array(&[4.0, 5.0]); + assert_eq!(idx1, ArrayIdx::new(0)); + assert_eq!(idx2, ArrayIdx::new(1)); + assert_eq!(pool.get_array(idx1, Width::new(3)), &[1.0, 2.0, 3.0]); + assert_eq!(pool.get_array(idx2, Width::new(2)), &[4.0, 5.0]); + + let data = + InterpolantData::try_new(vec![0.0, 1.0], vec![0.0, 10.0]).expect("valid interpolant"); + let interp_idx = pool.add_interpolant(&data); + assert_eq!(interp_idx, InterpolantIdx::new(0)); + assert_eq!(pool.interpolants[0].x_data, vec![0.0, 1.0]); + assert_eq!(pool.interpolants[0].slopes, vec![10.0]); + } + + #[test] + fn test_assign_slots_reuses_freed_slot_in_simple_chain() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let a = arena.alloc(Node::Sin(y)); + let b = arena.alloc(Node::Cos(a)); + let root = arena.alloc(Node::Neg(b)); + + let eval_order = arena.topological_order(root); + let sizes = infer_sizes(&arena, &eval_order); + let (slots, total) = assign_slots(&arena, &sizes, &eval_order, &[root]); + + assert!( + total <= 6, + "expected slot reuse to reduce buffer below 6, got {total}" + ); + + for &(off, len) in &slots { + assert!( + off + len <= total, + "invalid slot ({off}, {len}) > total {total}" + ); + } + } + + #[test] + fn test_assign_slots_coalesces_adjacent_free_regions() { + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y1 = arena.alloc(Node::StateVector { start: 2, end: 4 }); + let s0 = arena.alloc(Node::Sin(y0)); + let s1 = arena.alloc(Node::Sin(y1)); + let sum = arena.alloc(Node::Add(s0, s1)); + let cat = arena.alloc(Node::Concat(vec![y0, y1])); + let big = arena.alloc(Node::Sin(cat)); + let repeated_sum = arena.alloc(Node::Concat(vec![sum, sum])); + let root = arena.alloc(Node::Add(big, repeated_sum)); + + let eval_order = arena.topological_order(root); + let sizes = infer_sizes(&arena, &eval_order); + let (_slots, total) = assign_slots(&arena, &sizes, &eval_order, &[root]); + + let no_reuse: usize = eval_order.iter().map(|n| sizes[n.index()]).sum(); + assert!( + total < no_reuse, + "expected slot reuse total ({total}) < sequential sum ({no_reuse})" + ); + } + + #[test] + fn test_pinned_layout_disables_reuse_and_preserves_intermediates() { + // The product `a` is dead after the add, so only the reuse layout recycles its slot. + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let y2 = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let a = arena.alloc(Node::Mul(y0, y1)); + let root = arena.alloc(Node::Add(a, y2)); + + let reuse = TypedIr::from_dag(DagRef::new(&arena, root)); + let pinned = TypedIr::from_dag_pinned(DagRef::new(&arena, root)); + + assert!(pinned.buffer_size() > reuse.buffer_size()); + + let compiled = CompiledExpr::from_ir(pinned); + let mut s = vec![0.0; compiled.scratch_len()]; + let out = compiled.eval(&mut s, 0.0, &[2.0, 3.0, 5.0], &[], &[]); + assert_eq!(out, &[11.0]); + // SSA: each instruction writes a distinct slot. + let ir = compiled.ir(); + let mut offsets: Vec = Vec::new(); + for instr in ir.instructions() { + if let Some(dst) = instruction_dst_extent(instr, &ir.consts) { + offsets.push(dst.offset_usize()); + } + } + let unique: std::collections::BTreeSet = offsets.iter().copied().collect(); + assert_eq!(unique.len(), offsets.len(), "pinned layout reused a slot"); + } + + #[test] + fn a_split_tape_recycles_dead_tangent_slots() { + // d(sin^8 y) is a chain of `cos(x_k) * dx_k`: two width-4 tangents live at a time. + let mut arena = Arena::new(); + let mut x = arena.alloc(Node::StateVector { start: 0, end: 4 }); + for _ in 0..8 { + x = arena.alloc(Node::Sin(x)); + } + let mut tangent = crate::arena::Dag::new(arena, x); + tangent.rewrite_in_place(crate::tangent::tangent_wrt_states); + let (ir, split) = TypedIr::from_dag_split_eval(tangent.as_ref()); + assert_eq!(ir.buffer_size() - split.primal_buffer_size, 8); + + let (y, dy) = ([0.1, 0.2, 0.3, 0.4], [1.0, 0.0, -1.0, 2.0]); + let seeds = crate::eval::TangentInputs { + dy: Some(&dy), + dp: None, + }; + let monolithic = CompiledExpr::new(tangent.as_ref()); + let mut scratch = vec![0.0; monolithic.scratch_len()]; + let expected = monolithic.eval_with_tangent(&mut scratch, 0.0, &y, &[], &[], &seeds); + let expected: Vec = expected.iter().map(|v| v.to_bits()).collect(); + + let split = crate::eval::SplitCompiledExpr::compile(tangent.as_ref()); + let mut scratch = vec![0.0; split.scratch_len()]; + let mut primal = split.eval_primal(&mut scratch, 0.0, &y, &[], &[]); + let got: Vec = primal + .eval_tangent(&seeds) + .iter() + .map(|v| v.to_bits()) + .collect(); + assert_eq!(got, expected); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/lib.rs b/packages/pybamm-rust/pybamm-core/src/lib.rs new file mode 100644 index 0000000000..d6ba2f2405 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/lib.rs @@ -0,0 +1,38 @@ +//! Expression compiler and evaluator for `PyBaMM` models, filling the role +//! `CasADi` plays on `PyBaMM`'s other backends: +//! +//! 1. **Build**: the bindings allocate [`node`]s into an [`arena`], sharing subexpressions by id. +//! 2. **Rewrite**: [`simplify()`], [`zero_propagate()`], [`cse()`] and [`dce()`]. +//! 3. **Differentiate**: [`tangent`] emits a forward-mode DAG for JVPs. +//! 4. **Lower**: [`ir`] flattens the DAG into fixed-size instructions over one scratch buffer. +//! 5. **Evaluate**: [`eval`] runs a tape, [`CompiledExpr::eval_batch`] at several time points. + +#![deny(unused_must_use)] + +mod accumulate; +pub mod arena; +mod branch_regions; +mod error; +pub mod eval; +mod eval_batch; +pub mod ir; +mod newtype; +pub mod node; +mod params; +mod simplify; +pub mod tangent; +pub mod tangent_batch; +mod zero_propagate; + +pub use arena::{Arena, Dag, DagRef, NodeId, ReadExtents, resolve_state_count}; +pub use error::CoreError; +pub use eval::{CompiledExpr, SplitCompiledExpr, TangentInputs}; +pub use eval_batch::BatchEvalError; +pub use ir::{GraphError, TypedIr, infer_sizes, validate}; +pub use node::{ + ArrayData, CsrData, CubicInterpolantData, InterpolantData, NdInterpolantData, Node, Shape, +}; +pub use params::{PackedWidth, ParamIdx, ParamRegistry, ParamSlot, resolve_input_width}; +pub use simplify::{cse, dce, simplify, simplify_pipeline}; +pub use tangent::{DiffTarget, tangent_wrt_params, tangent_wrt_states}; +pub use zero_propagate::zero_propagate; diff --git a/packages/pybamm-rust/pybamm-core/src/newtype.rs b/packages/pybamm-rust/pybamm-core/src/newtype.rs new file mode 100644 index 0000000000..a439fd5da6 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/newtype.rs @@ -0,0 +1,33 @@ +//! Shared definition of the crate's index newtypes, so distinct integer roles cannot be swapped. + +/// Define a `#[repr(transparent)]` newtype over `$repr` with `new` and `get` accessors. +/// +/// `Default` is deliberately not derived: it would hand back a plausible index 0. +macro_rules! index_newtype { + ($(#[$attr:meta])* $name:ident($repr:ty)) => { + $(#[$attr])* + #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] + #[repr(transparent)] + pub struct $name($repr); + + impl $name { + /// Wrap a raw value. + pub const fn new(raw: $repr) -> Self { + Self(raw) + } + + /// The raw value. + pub const fn get(self) -> $repr { + self.0 + } + } + + impl ::std::fmt::Debug for $name { + fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result { + write!(f, "{}", self.0) + } + } + }; +} + +pub(crate) use index_newtype; diff --git a/packages/pybamm-rust/pybamm-core/src/node.rs b/packages/pybamm-rust/pybamm-core/src/node.rs new file mode 100644 index 0000000000..fa47f7a873 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/node.rs @@ -0,0 +1,1305 @@ +//! [`Node`], the expression DAG vocabulary handed over from Python. +//! +//! Nodes are not shape-checked on construction; [`validate`](crate::validate) +//! reports what lowering will reject. + +use crate::arena::NodeId; +use crate::error::CoreError; + +/// Declared `rows × cols` of a literal; only matmul reads it, elsewhere the flat length counts. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub struct Shape { + pub rows: usize, + pub cols: usize, +} + +/// Knots, values and per-segment slopes of a 1-D linear interpolant. +/// +/// `slopes` is derived only in [`try_new`](Self::try_new), which deserialising also uses. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr( + feature = "serialize", + derive(serde::Serialize, serde::Deserialize), + serde(try_from = "PersistedInterpolant") +)] +pub struct InterpolantData { + pub(crate) x_data: Vec, + pub(crate) y_data: Vec, + #[cfg_attr(feature = "serialize", serde(skip_serializing))] + pub(crate) slopes: Vec, +} + +/// Breakpoints and per-interval coefficients of a 1-D cubic interpolant. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr( + feature = "serialize", + derive(serde::Serialize, serde::Deserialize), + serde(try_from = "PersistedCubicInterpolant") +)] +pub struct CubicInterpolantData { + pub(crate) breakpoints: Vec, + /// Per-interval `[c0, c1, c2, c3]` of `c0 + c1*dx + c2*dx^2 + c3*dx^3`, + /// `dx = x - breakpoints[i]`. + pub(crate) coeffs: Vec<[f64; 4]>, +} + +/// Breakpoints and coefficients of a 2-D or 3-D tensor-product interpolant. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr( + feature = "serialize", + derive(serde::Serialize, serde::Deserialize), + serde(try_from = "PersistedNdInterpolant") +)] +pub struct NdInterpolantData { + pub(crate) breakpoints: Vec>, + /// Cell-major (axis 0 slowest) power-basis tensors, `order^ndim` per cell with + /// axis-0 power slowest: `p(dx) = Σ c[a0,..] · Π dx_i^a_i`. + pub(crate) coeffs: Vec, + /// Per-axis polynomial order: 2 = multilinear, 4 = tensor cubic. + pub(crate) order: u32, +} + +impl Shape { + /// `len × 1`, the column-vector orientation `PyBaMM` discretises into. + pub const fn vector(len: usize) -> Self { + Self { rows: len, cols: 1 } + } + + pub const fn matrix(rows: usize, cols: usize) -> Self { + Self { rows, cols } + } +} + +/// A constant sparse matrix in CSR form. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr( + feature = "serialize", + derive(serde::Serialize, serde::Deserialize), + serde(try_from = "PersistedCsr") +)] +pub struct CsrData { + pub(crate) indptr: Vec, + pub(crate) indices: Vec, + pub(crate) data: Vec, + pub(crate) shape: Shape, +} + +/// A constant dense array, stored row-major and flat. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr( + feature = "serialize", + derive(serde::Serialize, serde::Deserialize), + serde(try_from = "PersistedArray") +)] +pub struct ArrayData { + pub(crate) data: Vec, + pub(crate) shape: Shape, +} + +/// Define `$persisted`, the on-disk form of `$data`, which loads through the validating `$new`. +/// +/// Bincode is positional, so the fields must follow `$data`'s declaration order. +#[cfg(feature = "serialize")] +macro_rules! persisted_via_constructor { + ($persisted:ident => $data:ident::$new:ident($($field:ident: $ty:ty),+ $(,)?)) => { + #[derive(serde::Deserialize)] + struct $persisted { + $($field: $ty),+ + } + + impl TryFrom<$persisted> for $data { + type Error = CoreError; + + fn try_from(persisted: $persisted) -> Result { + Self::$new($(persisted.$field),+) + } + } + }; +} + +#[cfg(feature = "serialize")] +persisted_via_constructor!(PersistedCsr => CsrData::try_new( + indptr: Vec, + indices: Vec, + data: Vec, + shape: Shape, +)); +#[cfg(feature = "serialize")] +persisted_via_constructor!(PersistedArray => ArrayData::try_new(data: Vec, shape: Shape)); +#[cfg(feature = "serialize")] +persisted_via_constructor!(PersistedInterpolant => InterpolantData::try_new( + x_data: Vec, + y_data: Vec, +)); +#[cfg(feature = "serialize")] +persisted_via_constructor!(PersistedCubicInterpolant => CubicInterpolantData::try_new( + breakpoints: Vec, + coeffs: Vec<[f64; 4]>, +)); +#[cfg(feature = "serialize")] +persisted_via_constructor!(PersistedNdInterpolant => NdInterpolantData::try_new( + breakpoints: Vec>, + coeffs: Vec, + order: u32, +)); + +impl CsrData { + /// Build a CSR matrix, validating its structural invariants. + pub fn try_new( + indptr: Vec, + indices: Vec, + data: Vec, + shape: Shape, + ) -> Result { + let (rows, cols) = (shape.rows, shape.cols); + // Not `rows + 1`, which wraps for `rows == usize::MAX` and admits an empty indptr. + if indptr.len().checked_sub(1) != Some(rows) { + return Err(CoreError::Csr(format!( + "indptr length {} must equal rows + 1 for {rows} rows", + indptr.len(), + ))); + } + if indptr[0] != 0 { + return Err(CoreError::Csr(format!( + "indptr[0] must be 0, got {}", + indptr[0] + ))); + } + if indptr.windows(2).any(|w| w[1] < w[0]) { + return Err(CoreError::Csr("indptr must be non-decreasing".to_string())); + } + if indices.len() != data.len() { + return Err(CoreError::Csr(format!( + "indices length {} must equal data length {}", + indices.len(), + data.len() + ))); + } + let nnz = indptr[rows]; + if nnz != data.len() { + return Err(CoreError::Csr(format!( + "indptr tail {nnz} must equal nnz {}", + data.len() + ))); + } + if let Some(&max_col) = indices.iter().max() + && max_col >= cols + { + return Err(CoreError::Csr(format!( + "column index {max_col} out of range for {cols} columns" + ))); + } + Ok(Self { + indptr, + indices, + data, + shape, + }) + } + + /// Row-offset array, length `shape.rows + 1`. + pub fn indptr(&self) -> &[usize] { + &self.indptr + } + + /// Column index for each stored entry. + pub fn indices(&self) -> &[usize] { + &self.indices + } + + /// Stored values, parallel to [`indices`](Self::indices). + pub fn data(&self) -> &[f64] { + &self.data + } + + /// Matrix shape. + pub const fn shape(&self) -> Shape { + self.shape + } +} + +impl ArrayData { + /// Build a dense array, validating that `data.len()` equals `rows * cols`. + pub fn try_new(data: Vec, shape: Shape) -> Result { + let expected = shape.rows.checked_mul(shape.cols).ok_or_else(|| { + CoreError::Array(format!( + "shape {}x{} overflows usize", + shape.rows, shape.cols + )) + })?; + if data.len() != expected { + return Err(CoreError::Array(format!( + "data length {} must equal rows*cols = {expected}", + data.len() + ))); + } + Ok(Self { data, shape }) + } + + /// Row-major dense values. + pub fn data(&self) -> &[f64] { + &self.data + } + + /// Array shape. + pub const fn shape(&self) -> Shape { + self.shape + } +} + +impl InterpolantData { + /// Build a 1D linear table from finite, non-decreasing knots, deriving the slopes. + /// + /// A repeated interior knot is a step, right-continuous as `CasADi`'s; the end segments + /// extrapolate, so they must have nonzero width. + pub fn try_new(x_data: Vec, y_data: Vec) -> Result { + if x_data.is_empty() { + return Err(CoreError::Interpolant( + "x_data must be non-empty".to_string(), + )); + } + if x_data.len() != y_data.len() { + return Err(CoreError::Interpolant(format!( + "x_data length {} must equal y_data length {}", + x_data.len(), + y_data.len() + ))); + } + if !x_data.iter().all(|v| v.is_finite()) { + return Err(CoreError::Interpolant("x_data must be finite".to_string())); + } + if x_data.windows(2).any(|w| w[1] < w[0]) { + return Err(CoreError::Interpolant( + "x_data must be non-decreasing".to_string(), + )); + } + let n = x_data.len(); + if n > 1 && (x_data[1] <= x_data[0] || x_data[n - 1] <= x_data[n - 2]) { + return Err(CoreError::Interpolant( + "x_data must not repeat its first or last knot".to_string(), + )); + } + // The knot search never picks a zero-width segment, so its slope is never read. + let slopes = x_data + .windows(2) + .zip(y_data.windows(2)) + .map(|(x, y)| { + if x[1] > x[0] { + (y[1] - y[0]) / (x[1] - x[0]) + } else { + 0.0 + } + }) + .collect(); + Ok(Self { + x_data, + y_data, + slopes, + }) + } +} + +impl CubicInterpolantData { + /// Build a 1D cubic table from flat row-major coefficients `[c0,c1,c2,c3, c0,...]`. + pub fn try_from_flat(breakpoints: Vec, coeffs: &[f64]) -> Result { + let (tuples, rest) = coeffs.as_chunks::<4>(); + if !rest.is_empty() { + return Err(CoreError::Interpolant(format!( + "cubic interpolant coefficients must be a whole number of 4-tuples, got {}", + coeffs.len() + ))); + } + Self::try_new(breakpoints, tuples.to_vec()) + } + + /// Build a 1D cubic table, validating breakpoints and one coefficient tuple per segment. + pub fn try_new(breakpoints: Vec, coeffs: Vec<[f64; 4]>) -> Result { + if breakpoints.len() < 2 { + return Err(CoreError::Interpolant(format!( + "cubic interpolant needs at least 2 breakpoints, got {}", + breakpoints.len() + ))); + } + let nseg = breakpoints.len() - 1; + if coeffs.len() != nseg { + return Err(CoreError::Interpolant(format!( + "cubic interpolant needs {nseg} coefficient tuples (breakpoints - 1), got {}", + coeffs.len() + ))); + } + if !breakpoints.iter().all(|v| v.is_finite()) { + return Err(CoreError::Interpolant( + "cubic interpolant breakpoints must be finite".to_string(), + )); + } + if breakpoints.windows(2).any(|w| w[1] <= w[0]) { + return Err(CoreError::Interpolant( + "cubic interpolant breakpoints must be strictly increasing".to_string(), + )); + } + Ok(Self { + breakpoints, + coeffs, + }) + } +} + +impl NdInterpolantData { + /// Build an N-D tensor table, validating axes, order, knots and coefficient count. + pub fn try_new( + breakpoints: Vec>, + coeffs: Vec, + order: u32, + ) -> Result { + let ndim = breakpoints.len(); + if !(2..=3).contains(&ndim) { + return Err(CoreError::Interpolant(format!( + "N-D interpolant supports 2 or 3 axes, got {ndim}" + ))); + } + if order != 2 && order != 4 { + return Err(CoreError::Interpolant(format!( + "N-D interpolant order must be 2 or 4, got {order}" + ))); + } + // Checked: a table deserialised from caller bytes must not wrap to a small count. + let overflow = + || CoreError::Interpolant("N-D interpolant cell count overflows usize".to_string()); + let mut ncells = 1usize; + for (axis, knots) in breakpoints.iter().enumerate() { + if knots.len() < 2 { + return Err(CoreError::Interpolant(format!( + "N-D interpolant axis {axis} needs at least 2 breakpoints, got {}", + knots.len() + ))); + } + if !knots.iter().all(|v| v.is_finite()) { + return Err(CoreError::Interpolant(format!( + "N-D interpolant axis {axis} breakpoints must be finite" + ))); + } + if knots.windows(2).any(|w| w[1] <= w[0]) { + return Err(CoreError::Interpolant(format!( + "N-D interpolant axis {axis} breakpoints must be strictly increasing" + ))); + } + ncells = ncells.checked_mul(knots.len() - 1).ok_or_else(overflow)?; + } + let order_usize = order as usize; + let per_cell = match ndim { + 2 => order_usize * order_usize, + 3 => order_usize * order_usize * order_usize, + _ => unreachable!("ndim validated to 2..=3"), + }; + let expected = ncells.checked_mul(per_cell).ok_or_else(overflow)?; + if coeffs.len() != expected { + return Err(CoreError::Interpolant(format!( + "N-D interpolant needs {expected} coefficients (ncells * order^ndim), got {}", + coeffs.len() + ))); + } + Ok(Self { + breakpoints, + coeffs, + order, + }) + } +} + +/// One node of the expression DAG. +/// +/// Children are [`NodeId`]s into the owning arena. Arithmetic and unary variants are +/// element-wise, with a scalar operand broadcasting against a vector. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] +pub enum Node { + Scalar(f64), + Array(Box), + ZeroVector { + len: usize, + }, + SparseMatrix(Box), + /// Reads `y[start..end]` of the solver's global state vector. + StateVector { + start: usize, + end: usize, + }, + /// Reads `y'[start..end]`, in the same index space as [`StateVector`](Self::StateVector). + StateVectorDot { + start: usize, + end: usize, + }, + InputParameter { + name: String, + /// Registration order among distinct names; used for sensitivity indexing. + index: usize, + /// Offset into the packed `p` values array. + offset: usize, + /// Number of packed values this parameter occupies (>1 for vector inputs). + width: usize, + }, + Time, + + Add(NodeId, NodeId), + Sub(NodeId, NodeId), + Mul(NodeId, NodeId), + Div(NodeId, NodeId), + Pow(NodeId, NodeId), + /// Matrix-vector product; the left child must be a constant + /// [`SparseMatrix`](Self::SparseMatrix) or [`Array`](Self::Array). + MatMul(NodeId, NodeId), + Minimum(NodeId, NodeId), + Maximum(NodeId, NodeId), + Modulo(NodeId, NodeId), + Hypot(NodeId, NodeId), + EqualHeaviside(NodeId, NodeId), + NotEqualHeaviside(NodeId, NodeId), + Equality(NodeId, NodeId), + + /// Half-open slice `child[start..end]` of the child's own value. + Index { + child: NodeId, + start: usize, + end: usize, + }, + /// Children joined end to end; for a right-hand side this is the solver's equation order. + Concat(Vec), + + Interpolant1DLinear { + data: Box, + child: NodeId, + }, + /// 1D cubic/pchip interpolation (piecewise cubic, power-basis coeffs). + Interpolant1DCubic { + data: Box, + child: NodeId, + }, + /// N-D tensor-product interpolation, one equal-length child per axis. + InterpolantNd { + data: Box, + children: Vec, + }, + + Neg(NodeId), + Abs(NodeId), + Sqrt(NodeId), + Exp(NodeId), + Log(NodeId), + Sin(NodeId), + Cos(NodeId), + Tanh(NodeId), + Sinh(NodeId), + Cosh(NodeId), + Arcsinh(NodeId), + Arctan(NodeId), + Erf(NodeId), + Sign(NodeId), + Floor(NodeId), + Ceiling(NodeId), + MaxReduce(NodeId), + MinReduce(NodeId), + + /// Internal: `basis[k]`, `k` the first-occurrence argmax (or argmin) of `picker`; + /// the subgradient of `MaxReduce`/`MinReduce`. + ReduceArgSelect { + basis: NodeId, + picker: NodeId, + is_max: bool, + }, + + /// Selects `branches[i]` when `selector` is in the open window `(i + 0.5, i + 1.5)`; + /// an exact half-way value matches nothing and evaluates to zeros. + Conditional { + selector: NodeId, + branches: Vec, + }, + + /// Reads a slice of the tangent state vector. + TangentStateVector { + start: usize, + end: usize, + }, + /// Reads a single tangent parameter. + TangentParameter { + index: usize, + }, + /// Derivative of linear interpolation, from the table's pre-computed slopes. + Interpolant1DLinearDeriv { + data: Box, + child: NodeId, + }, + /// Derivative of 1D cubic/pchip interpolation. + Interpolant1DCubicDeriv { + data: Box, + child: NodeId, + }, + /// Partial derivative of N-D interpolation along `axis`. + InterpolantNdPartial { + data: Box, + children: Vec, + axis: u32, + }, +} + +/// Hand each child of `$node` to `$visit`, in declaration order. +/// +/// One match serves the shared and mutable visitors, so their orders cannot drift apart. +macro_rules! visit_children { + ($node:expr, $visit:ident) => { + match $node { + Node::Scalar(_) + | Node::Array(_) + | Node::ZeroVector { .. } + | Node::SparseMatrix(_) + | Node::StateVector { .. } + | Node::StateVectorDot { .. } + | Node::InputParameter { .. } + | Node::Time + | Node::TangentStateVector { .. } + | Node::TangentParameter { .. } => {}, + + Node::Add(l, r) + | Node::Sub(l, r) + | Node::Mul(l, r) + | Node::Div(l, r) + | Node::Pow(l, r) + | Node::MatMul(l, r) + | Node::Minimum(l, r) + | Node::Maximum(l, r) + | Node::Modulo(l, r) + | Node::Hypot(l, r) + | Node::EqualHeaviside(l, r) + | Node::NotEqualHeaviside(l, r) + | Node::Equality(l, r) => { + $visit(l); + $visit(r); + }, + + Node::Neg(c) + | Node::Abs(c) + | Node::Sqrt(c) + | Node::Exp(c) + | Node::Log(c) + | Node::Sin(c) + | Node::Cos(c) + | Node::Tanh(c) + | Node::Sinh(c) + | Node::Cosh(c) + | Node::Arcsinh(c) + | Node::Arctan(c) + | Node::Erf(c) + | Node::Sign(c) + | Node::Floor(c) + | Node::Ceiling(c) + | Node::MaxReduce(c) + | Node::MinReduce(c) => $visit(c), + + Node::Index { child, .. } + | Node::Interpolant1DLinear { child, .. } + | Node::Interpolant1DLinearDeriv { child, .. } + | Node::Interpolant1DCubic { child, .. } + | Node::Interpolant1DCubicDeriv { child, .. } => $visit(child), + + Node::Concat(children) + | Node::InterpolantNd { children, .. } + | Node::InterpolantNdPartial { children, .. } => { + for c in children { + $visit(c); + } + }, + + Node::Conditional { selector, branches } => { + $visit(selector); + for b in branches { + $visit(b); + } + }, + + Node::ReduceArgSelect { basis, picker, .. } => { + $visit(basis); + $visit(picker); + }, + } + }; +} + +/// Require `start <= end`; an inverted pair would wrap the width in release. +const fn check_extent(node: &'static str, start: usize, end: usize) -> Result<(), CoreError> { + if start > end { + return Err(CoreError::InvertedExtent { node, start, end }); + } + Ok(()) +} + +impl Node { + /// A zero of broadcast length `len`: a `Scalar` for 1, else a `ZeroVector` (0 included). + pub(crate) const fn zero(len: usize) -> Self { + if len == 1 { + Self::Scalar(0.0) + } else { + Self::ZeroVector { len } + } + } + + /// The value of a `Scalar` literal; constant-valued subtrees are not folded. + pub(crate) const fn as_scalar(&self) -> Option { + match self { + Self::Scalar(value) => Some(*value), + _ => None, + } + } + + /// The state slice `y[start..end]`. + pub fn state_vector(start: usize, end: usize) -> Result { + check_extent("StateVector", start, end)?; + Ok(Self::StateVector { start, end }) + } + + /// The time-derivative slice `y_dot[start..end]`. + pub fn state_vector_dot(start: usize, end: usize) -> Result { + check_extent("StateVectorDot", start, end)?; + Ok(Self::StateVectorDot { start, end }) + } + + /// `child[start..end]`; [`validate`](crate::validate) checks it fits the child. + pub fn index(child: NodeId, start: usize, end: usize) -> Result { + check_extent("Index", start, end)?; + Ok(Self::Index { child, start, end }) + } + + /// A branch selected by `selector`, 1-based. + pub fn conditional(selector: NodeId, branches: Vec) -> Result { + if branches.is_empty() { + return Err(CoreError::EmptyConditional); + } + Ok(Self::Conditional { selector, branches }) + } + + /// A tensor-product interpolant evaluated at one coordinate per axis. + pub fn interpolant_nd( + data: NdInterpolantData, + children: Vec, + ) -> Result { + let axes = data.breakpoints.len(); + if children.len() != axes { + return Err(CoreError::InterpolantNdArity { + axes, + children: children.len(), + }); + } + Ok(Self::InterpolantNd { + data: Box::new(data), + children, + }) + } + + /// Visit each child `NodeId` of this node, allocation-free. + pub(crate) fn for_each_child(&self, mut f: F) { + let mut visit = |child: &NodeId| f(*child); + visit_children!(self, visit); + } + + /// Visit each child `NodeId` in place, in [`for_each_child`](Self::for_each_child) order. + pub(crate) fn for_each_child_mut(&mut self, mut f: F) { + visit_children!(self, f); + } + + /// A copy of this node with `f` applied to each child `NodeId`, in + /// [`for_each_child`](Self::for_each_child) order. + #[must_use] + pub(crate) fn map_children NodeId>(&self, mut f: F) -> Self { + let mut mapped = self.clone(); + mapped.for_each_child_mut(|child| *child = f(*child)); + mapped + } +} + +fn hash_nd_interpolant_data(data: &NdInterpolantData, hasher: &mut H) { + use std::hash::Hash; + data.order.hash(hasher); + for knots in &data.breakpoints { + knots.len().hash(hasher); + for v in knots { + v.to_bits().hash(hasher); + } + } + for v in &data.coeffs { + v.to_bits().hash(hasher); + } +} + +/// Hash a node by structure for CSE, floats by bit pattern. +/// +/// Pairs with [`structurally_identical`]: identical nodes hash equally, but distinct ones +/// may collide, e.g. payloads differing in an even number of zero signs. +pub(crate) fn structural_hash(node: &Node, hasher: &mut H) { + use std::hash::Hash; + std::mem::discriminant(node).hash(hasher); + match node { + Node::Scalar(v) => v.to_bits().hash(hasher), + Node::Array(arr) => { + for v in &arr.data { + v.to_bits().hash(hasher); + } + arr.shape.rows.hash(hasher); + arr.shape.cols.hash(hasher); + }, + Node::ZeroVector { len } => len.hash(hasher), + Node::SparseMatrix(csr) => { + csr.indptr.hash(hasher); + csr.indices.hash(hasher); + for v in &csr.data { + v.to_bits().hash(hasher); + } + csr.shape.rows.hash(hasher); + csr.shape.cols.hash(hasher); + }, + Node::StateVector { start, end } + | Node::StateVectorDot { start, end } + | Node::TangentStateVector { start, end } + | Node::Index { start, end, .. } => { + start.hash(hasher); + end.hash(hasher); + }, + Node::InputParameter { + name, + index, + offset, + width, + } => { + name.hash(hasher); + index.hash(hasher); + offset.hash(hasher); + width.hash(hasher); + }, + Node::TangentParameter { index } => index.hash(hasher), + Node::Concat(children) => children.len().hash(hasher), + Node::Conditional { branches, .. } => branches.len().hash(hasher), + // Slopes derive from knots and values, so need no hashing. + Node::Interpolant1DLinear { data, .. } | Node::Interpolant1DLinearDeriv { data, .. } => { + for v in &data.x_data { + v.to_bits().hash(hasher); + } + for v in &data.y_data { + v.to_bits().hash(hasher); + } + }, + Node::Interpolant1DCubic { data, .. } | Node::Interpolant1DCubicDeriv { data, .. } => { + for v in &data.breakpoints { + v.to_bits().hash(hasher); + } + for c in &data.coeffs { + for v in c { + v.to_bits().hash(hasher); + } + } + }, + Node::InterpolantNd { data, .. } => hash_nd_interpolant_data(data, hasher), + Node::InterpolantNdPartial { data, axis, .. } => { + hash_nd_interpolant_data(data, hasher); + axis.hash(hasher); + }, + Node::Time + | Node::Add(_, _) + | Node::Sub(_, _) + | Node::Mul(_, _) + | Node::Div(_, _) + | Node::Pow(_, _) + | Node::MatMul(_, _) + | Node::Minimum(_, _) + | Node::Maximum(_, _) + | Node::Modulo(_, _) + | Node::Hypot(_, _) + | Node::EqualHeaviside(_, _) + | Node::NotEqualHeaviside(_, _) + | Node::Equality(_, _) + | Node::Neg(_) + | Node::Abs(_) + | Node::Sqrt(_) + | Node::Exp(_) + | Node::Log(_) + | Node::Sin(_) + | Node::Cos(_) + | Node::Tanh(_) + | Node::Sinh(_) + | Node::Cosh(_) + | Node::Arcsinh(_) + | Node::Arctan(_) + | Node::Erf(_) + | Node::Sign(_) + | Node::Floor(_) + | Node::Ceiling(_) + | Node::MaxReduce(_) + | Node::MinReduce(_) => {}, + Node::ReduceArgSelect { is_max, .. } => is_max.hash(hasher), + } + node.for_each_child(|c| c.hash(hasher)); +} + +/// Whether `a` and `b` are the same node, with floats compared by bit pattern. +/// +/// The equality [`structural_hash`] pairs with: derived `==` would merge `0.0` and `-0.0`. +pub(crate) fn structurally_identical(a: &Node, b: &Node) -> bool { + fn same_bits(x: &[f64], y: &[f64]) -> bool { + x.iter() + .map(|v| v.to_bits()) + .eq(y.iter().map(|v| v.to_bits())) + } + // `==` has already matched the variant and every non-float field. + a == b + && match (a, b) { + (Node::Scalar(x), Node::Scalar(y)) => x.to_bits() == y.to_bits(), + (Node::Array(x), Node::Array(y)) => same_bits(&x.data, &y.data), + (Node::SparseMatrix(x), Node::SparseMatrix(y)) => same_bits(&x.data, &y.data), + ( + Node::Interpolant1DLinear { data: x, .. } + | Node::Interpolant1DLinearDeriv { data: x, .. }, + Node::Interpolant1DLinear { data: y, .. } + | Node::Interpolant1DLinearDeriv { data: y, .. }, + ) => same_bits(&x.x_data, &y.x_data) && same_bits(&x.y_data, &y.y_data), + ( + Node::Interpolant1DCubic { data: x, .. } + | Node::Interpolant1DCubicDeriv { data: x, .. }, + Node::Interpolant1DCubic { data: y, .. } + | Node::Interpolant1DCubicDeriv { data: y, .. }, + ) => { + same_bits(&x.breakpoints, &y.breakpoints) + && same_bits(x.coeffs.as_flattened(), y.coeffs.as_flattened()) + }, + ( + Node::InterpolantNd { data: x, .. } | Node::InterpolantNdPartial { data: x, .. }, + Node::InterpolantNd { data: y, .. } | Node::InterpolantNdPartial { data: y, .. }, + ) => { + x.breakpoints + .iter() + .zip(&y.breakpoints) + .all(|(p, q)| same_bits(p, q)) + && same_bits(&x.coeffs, &y.coeffs) + }, + _ => true, + } +} + +const _: () = { + assert!(size_of::() <= 48); +}; + +#[cfg(test)] +mod tests { + use super::*; + use crate::arena::Arena; + + /// `result` is `Ok` exactly when `valid`, and otherwise an error `kind` accepts. + #[track_caller] + fn assert_validated( + result: Result, + valid: bool, + kind: fn(&CoreError) -> bool, + label: &str, + ) { + match result { + Ok(_) => assert!(valid, "{label} must be refused"), + Err(err) => assert!(!valid && kind(&err), "{label}: {err}"), + } + } + + fn is_interpolant(err: &CoreError) -> bool { + matches!(err, CoreError::Interpolant(_)) + } + + #[test] + fn csr_try_new_accepts_a_consistent_matrix_and_refuses_the_rest() { + let (square, tall) = (Shape::matrix(2, 2), Shape::matrix(usize::MAX, 1)); + let cases = [ + (vec![0, 1, 2], vec![0, 1], vec![1.0, 2.0], square, true), + (vec![0, 2], vec![0, 1], vec![1.0, 2.0], square, false), + (vec![], vec![], vec![], tall, false), + (vec![0, 2, 1], vec![0, 1], vec![1.0, 2.0], square, false), + (vec![0, 1, 2], vec![0, 5], vec![1.0, 2.0], square, false), + (vec![0, 1, 2], vec![0, 1], vec![1.0], square, false), + (vec![0, 1, 2], vec![0], vec![1.0], square, false), + ]; + for (indptr, indices, data, shape, valid) in cases { + let label = format!("{indptr:?} {indices:?} {data:?} {shape:?}"); + let result = CsrData::try_new(indptr, indices, data, shape); + assert_validated(result, valid, |e| matches!(e, CoreError::Csr(_)), &label); + } + } + + #[test] + fn array_try_new_accepts_matching_length() { + assert!(ArrayData::try_new(vec![1.0, 2.0, 3.0, 4.0], Shape::matrix(2, 2)).is_ok()); + } + + #[test] + fn array_try_new_rejects_length_mismatch() { + assert!(ArrayData::try_new(vec![1.0, 2.0, 3.0], Shape::matrix(2, 2)).is_err()); + } + + #[test] + fn interpolant_try_new_differences_a_valid_table_into_slopes_and_refuses_the_rest() { + /// Knots, values, and the slopes of a valid table or `None` for a refused one. + type Case = (&'static [f64], &'static [f64], Option<&'static [f64]>); + let cases: [Case; 10] = [ + (&[0.0, 1.0, 3.0], &[0.0, 10.0, 40.0], Some(&[10.0, 15.0])), + (&[1.0], &[2.0], Some(&[])), + // A repeated interior knot is a step, whose zero-width segment has no slope. + ( + &[0.0, 1.0, 1.0, 2.0], + &[0.0, 1.0, 3.0, 5.0], + Some(&[1.0, 0.0, 2.0]), + ), + (&[], &[], None), + (&[0.0, 1.0], &[10.0], None), + (&[0.0, 2.0, 1.0], &[1.0, 2.0, 3.0], None), + (&[0.0, f64::NAN], &[1.0, 2.0], None), + // The end segments extrapolate, so neither may have zero width. + (&[0.0, 0.0, 1.0], &[1.0, 2.0, 3.0], None), + (&[0.0, 1.0, 1.0], &[1.0, 2.0, 3.0], None), + (&[1.0, 1.0], &[1.0, 2.0], None), + ]; + for (x_data, y_data, slopes) in cases { + let label = format!("{x_data:?} -> {y_data:?}"); + let result = InterpolantData::try_new(x_data.to_vec(), y_data.to_vec()); + if let (Ok(data), Some(slopes)) = (&result, slopes) { + assert_eq!(data.slopes, slopes, "{label}"); + } + assert_validated(result, slopes.is_some(), is_interpolant, &label); + } + } + + #[test] + fn cubic_try_new_accepts_one_tuple_per_segment_and_refuses_the_rest() { + let segment = [1.0, 2.0, 3.0, 4.0]; + let cases = [ + (vec![0.0, 1.0], vec![segment], true), + (vec![0.0], vec![], false), + (vec![0.0, 1.0], vec![segment, segment], false), + (vec![1.0, 0.0], vec![segment], false), + ]; + for (breakpoints, coeffs, valid) in cases { + let label = format!("{breakpoints:?} with {} segment(s)", coeffs.len()); + let result = CubicInterpolantData::try_new(breakpoints, coeffs); + assert_validated(result, valid, is_interpolant, &label); + } + } + + #[test] + fn nd_try_new_accepts_a_consistent_table_and_refuses_the_rest() { + let square = || vec![vec![0.0, 1.0], vec![0.0, 1.0]]; + let cases = [ + (square(), vec![1.0, 2.0, 3.0, 4.0], 2, true), + (square(), vec![1.0, 2.0, 3.0, 4.0], 3, false), + (square(), vec![1.0, 2.0, 3.0], 2, false), + (vec![vec![0.0], vec![0.0, 1.0]], vec![1.0, 2.0], 2, false), + ]; + for (breakpoints, coeffs, order, valid) in cases { + let label = format!("{breakpoints:?}, {} coeff(s), order {order}", coeffs.len()); + let result = NdInterpolantData::try_new(breakpoints, coeffs, order); + assert_validated(result, valid, is_interpolant, &label); + } + } + + #[test] + fn test_for_each_child_collects_all_children() { + let mut arena = Arena::new(); + let a = arena.alloc(Node::Scalar(1.0)); + let b = arena.alloc(Node::Scalar(2.0)); + let add = Node::Add(a, b); + + let mut collected = Vec::new(); + add.for_each_child(|c| collected.push(c)); + assert_eq!(collected, vec![a, b]); + + let leaf = Node::Scalar(2.72); + let mut leaf_children = Vec::new(); + leaf.for_each_child(|c| leaf_children.push(c)); + assert_eq!(leaf_children, []); + + let concat = Node::Concat(vec![a, b, a]); + let mut concat_children = Vec::new(); + concat.for_each_child(|c| concat_children.push(c)); + assert_eq!(concat_children, vec![a, b, a]); + } + + #[test] + fn map_children_hands_out_children_in_for_each_child_order() { + let ids: Vec = (0..4u32).map(NodeId::from).collect(); + let nodes = [ + Node::Conditional { + selector: ids[3], + branches: vec![ids[0], ids[2]], + }, + Node::ReduceArgSelect { + basis: ids[1], + picker: ids[0], + is_max: true, + }, + Node::Sub(ids[2], ids[1]), + Node::Concat(vec![ids[0], ids[1], ids[2]]), + Node::Scalar(3.15), + Node::ZeroVector { len: 4 }, + ]; + for node in nodes { + assert_eq!( + node.map_children(|c| c), + node, + "the identity map keeps payloads" + ); + let mut children = Vec::new(); + node.for_each_child(|c| children.push(c)); + let mut handed = Vec::new(); + let mut calls = 0u32; + let mapped = node.map_children(|c| { + handed.push(c); + calls += 1; + NodeId::from(100 + calls) + }); + assert_eq!(handed, children, "{node:?}"); + let mut visited = Vec::new(); + mapped.for_each_child(|c| visited.push(c.raw())); + assert_eq!(visited, (101..=100 + calls).collect::>()); + } + } + + #[test] + fn test_structural_hash_distinguishes_payload_only_differences() { + use rustc_hash::FxHasher; + use std::hash::Hasher; + + let s1 = Node::StateVector { start: 0, end: 5 }; + let s2 = Node::StateVector { start: 1, end: 6 }; + + let mut h1 = FxHasher::default(); + let mut h2 = FxHasher::default(); + structural_hash(&s1, &mut h1); + structural_hash(&s2, &mut h2); + assert_ne!(h1.finish(), h2.finish()); + + let zv1 = Node::ZeroVector { len: 3 }; + let zv2 = Node::ZeroVector { len: 7 }; + let mut h3 = FxHasher::default(); + let mut h4 = FxHasher::default(); + structural_hash(&zv1, &mut h3); + structural_hash(&zv2, &mut h4); + assert_ne!(h3.finish(), h4.finish()); + + let ip1 = Node::InputParameter { + name: "p1".into(), + index: 0, + offset: 0, + width: 1, + }; + let ip2 = Node::InputParameter { + name: "p1".into(), + index: 1, + offset: 1, + width: 1, + }; + let mut h5 = FxHasher::default(); + let mut h6 = FxHasher::default(); + structural_hash(&ip1, &mut h5); + structural_hash(&ip2, &mut h6); + assert_ne!(h5.finish(), h6.finish()); + } + + #[test] + fn structural_identity_tells_signed_zeros_apart_in_every_payload() { + let child = NodeId::from(0u32); + let payloads = |zero: f64| { + [ + Node::Scalar(zero), + Node::Array(Box::new( + ArrayData::try_new(vec![zero], Shape::vector(1)).expect("valid array"), + )), + Node::SparseMatrix(Box::new( + CsrData::try_new(vec![0, 1], vec![0], vec![zero], Shape::matrix(1, 1)) + .expect("valid CSR"), + )), + Node::Interpolant1DLinearDeriv { + data: Box::new( + InterpolantData::try_new(vec![0.0, 1.0], vec![zero, 1.0]) + .expect("valid table"), + ), + child, + }, + Node::Interpolant1DCubic { + data: Box::new( + CubicInterpolantData::try_new(vec![0.0, 1.0], vec![[zero, 1.0, 0.0, 0.0]]) + .expect("valid cubic"), + ), + child, + }, + Node::InterpolantNdPartial { + data: Box::new( + NdInterpolantData::try_new( + vec![vec![0.0, 1.0], vec![0.0, 1.0]], + vec![zero, 1.0, 1.0, 1.0], + 2, + ) + .expect("valid N-D table"), + ), + children: vec![child, child], + axis: 1, + }, + ] + }; + for (positive, negative) in payloads(0.0).iter().zip(&payloads(-0.0)) { + assert!(structurally_identical(positive, &positive.clone())); + assert_eq!(positive, negative, "derived `==` equates the two zeros"); + assert!(!structurally_identical(positive, negative), "{positive:?}"); + } + } + + #[test] + fn a_sliced_node_refuses_an_inverted_extent() { + let child = NodeId::from(0u32); + assert!(matches!( + Node::state_vector(2, 1), + Err(CoreError::InvertedExtent { + node: "StateVector", + start: 2, + end: 1 + }) + )); + assert!(matches!( + Node::state_vector_dot(2, 1), + Err(CoreError::InvertedExtent { .. }) + )); + assert!(matches!( + Node::index(child, 3, 1), + Err(CoreError::InvertedExtent { .. }) + )); + + assert!(Node::state_vector(1, 1).is_ok()); + assert!(Node::index(child, 1, 3).is_ok()); + } + + #[test] + fn a_conditional_needs_a_branch_to_select() { + let selector = NodeId::from(0u32); + assert!(matches!( + Node::conditional(selector, vec![]), + Err(CoreError::EmptyConditional) + )); + assert!(Node::conditional(selector, vec![NodeId::from(1u32)]).is_ok()); + } + + #[test] + fn an_nd_interpolant_needs_one_child_per_axis() { + let data = + NdInterpolantData::try_new(vec![vec![0.0, 1.0], vec![0.0, 1.0]], vec![0.0; 4], 2) + .expect("a 2-axis multilinear table"); + assert!(matches!( + Node::interpolant_nd(data.clone(), vec![NodeId::from(0u32)]), + Err(CoreError::InterpolantNdArity { + axes: 2, + children: 1 + }) + )); + assert!(Node::interpolant_nd(data, vec![NodeId::from(0u32), NodeId::from(1u32)]).is_ok()); + } + + #[test] + fn flat_cubic_coefficients_must_be_whole_tuples() { + let breakpoints = vec![0.0, 1.0, 2.0]; + assert!(matches!( + CubicInterpolantData::try_from_flat(breakpoints.clone(), &[1.0; 7]), + Err(CoreError::Interpolant(_)) + )); + let table = CubicInterpolantData::try_from_flat(breakpoints, &[1.0; 8]) + .expect("two segments of four"); + assert_eq!(table.coeffs.len(), 2); + } +} + +#[cfg(all(test, feature = "serialize"))] +mod persisted_tests { + use super::*; + + fn table() -> InterpolantData { + InterpolantData::try_new(vec![0.0, 1.0, 3.0], vec![1.0, 2.0, 0.0]).expect("valid table") + } + + #[test] + fn a_linear_interpolant_round_trips_with_its_slopes_rederived() { + let data = table(); + let bytes = bincode::serialize(&data).expect("serialise"); + let back: InterpolantData = bincode::deserialize(&bytes).expect("deserialise"); + assert_eq!(back, data); + assert_eq!(back.slopes, data.slopes); + } + + /// Snapshots written before slopes were stored must still load. + #[test] + fn the_persisted_form_is_knots_and_values_only() { + let data = table(); + let bytes = bincode::serialize(&data).expect("serialise"); + let bare = bincode::serialize(&(&data.x_data, &data.y_data)).expect("serialise"); + assert_eq!(bytes, bare); + } + + #[test] + fn an_invalid_persisted_interpolant_is_refused() { + let repeated_last_knot = (vec![0.0, 1.0, 1.0], vec![1.0, 2.0, 3.0]); + let bytes = bincode::serialize(&repeated_last_knot).expect("serialise"); + assert!(bincode::deserialize::(&bytes).is_err()); + } + + fn round_trip(value: &T) + where + T: serde::Serialize + serde::de::DeserializeOwned + PartialEq + std::fmt::Debug, + { + let bytes = bincode::serialize(value).expect("serialise"); + assert_eq!( + &bincode::deserialize::(&bytes).expect("deserialise"), + value + ); + } + + #[test] + fn validated_tables_round_trip_field_for_field() { + round_trip( + &CsrData::try_new(vec![0, 1, 1], vec![2], vec![5.0], Shape::matrix(2, 3)) + .expect("valid CSR"), + ); + round_trip(&ArrayData::try_new(vec![1.0, 2.0], Shape::vector(2)).expect("valid array")); + round_trip( + &InterpolantData::try_new(vec![0.0, 1.0, 1.0, 2.0], vec![0.0, 1.0, 3.0, 5.0]) + .expect("valid step"), + ); + round_trip( + &CubicInterpolantData::try_new(vec![0.0, 1.0], vec![[1.0, 2.0, 3.0, 4.0]]) + .expect("valid cubic"), + ); + round_trip( + &NdInterpolantData::try_new(vec![vec![0.0, 1.0], vec![0.0, 2.0]], vec![0.5; 4], 2) + .expect("valid N-D table"), + ); + } + + #[test] + fn invalid_persisted_tables_are_refused() { + let csr_tail_past_data = ( + vec![0usize, 2], + vec![0usize], + vec![1.0], + Shape::matrix(1, 1), + ); + let bytes = bincode::serialize(&csr_tail_past_data).expect("serialise"); + assert!(bincode::deserialize::(&bytes).is_err()); + + let array_short = (vec![1.0], Shape::matrix(2, 2)); + let bytes = bincode::serialize(&array_short).expect("serialise"); + assert!(bincode::deserialize::(&bytes).is_err()); + + let cubic_short = (vec![0.0, 1.0, 2.0], vec![[0.0; 4]]); + let bytes = bincode::serialize(&cubic_short).expect("serialise"); + assert!(bincode::deserialize::(&bytes).is_err()); + + let nd_order_three = (vec![vec![0.0, 1.0], vec![0.0, 1.0]], vec![0.0; 9], 3u32); + let bytes = bincode::serialize(&nd_order_three).expect("serialise"); + assert!(bincode::deserialize::(&bytes).is_err()); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/params.rs b/packages/pybamm-rust/pybamm-core/src/params.rs new file mode 100644 index 0000000000..63b1dc4007 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/params.rs @@ -0,0 +1,602 @@ +//! The two index spaces a model's input parameters live in. +//! +//! A parameter's *registration index* and its *packed offset* into `inputs` +//! diverge once any input is vector-valued, so each is its own type. + +use std::collections::HashMap; + +use crate::arena::Arena; +use crate::error::CoreError; +use crate::ir::ids::{InputOffset, Width, tape_index}; +use crate::newtype::index_newtype; +use crate::node::Node; + +tape_index! { + /// Registration index of an input parameter; the space `dp` is indexed in. + ParamIdx +} + +index_newtype! { + /// Length of the packed `inputs` slice: the sum of every parameter's width. + #[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))] + PackedWidth(usize) +} + +/// Where one registered input parameter's values sit in the packed vector. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub struct ParamSlot { + /// Element offset into the packed `inputs` slice. + pub offset: InputOffset, + /// How many packed elements the parameter occupies; `> 1` for a vector input. + pub width: Width, +} + +impl ParamSlot { + /// The slot `[offset, offset + width)` of `name`, refused if it ends past `u32::MAX`. + fn try_new(name: &str, offset: usize, width: usize) -> Result { + if offset + .checked_add(width) + .is_none_or(|end| u32::try_from(end).is_err()) + { + return Err(CoreError::ParamSlotRange { + name: name.to_string(), + offset, + width, + }); + } + Ok(Self { + offset: InputOffset::from_usize(offset), + width: Width::from_usize(width), + }) + } +} + +/// The one place a parameter's registration index and packed slot are decided. +/// +/// Not owned by [`Arena`], since rewriters build fresh arenas and would drop it; +/// [`from_arena`](Self::from_arena) recovers it from the nodes. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct ParamRegistry { + by_name: HashMap, + /// Indexed by [`ParamIdx`]; `None` at an index naming no parameter. + entries: Vec>, + packed_width: PackedWidth, +} + +#[derive(Clone, PartialEq, Eq, Debug)] +struct ParamEntry { + name: String, + slot: ParamSlot, +} + +impl Default for ParamRegistry { + fn default() -> Self { + Self::new() + } +} + +impl ParamRegistry { + pub fn new() -> Self { + Self { + by_name: HashMap::new(), + entries: Vec::new(), + packed_width: PackedWidth::new(0), + } + } + + /// Register `name` at the end of the packed vector, or look it up if already registered. + /// + /// A repeat must repeat the width, since later offsets are running totals of widths. + pub fn register( + &mut self, + name: &str, + width: usize, + ) -> Result<(ParamIdx, ParamSlot), CoreError> { + if let Some(&index) = self.by_name.get(name) { + let slot = self + .slot(index) + .expect("a name in the registry indexes a registered parameter"); + if slot.width.index() != width { + return Err(CoreError::ParamWidthConflict { + name: name.to_string(), + width, + registered: slot.width.index(), + }); + } + return Ok((index, slot)); + } + let slot = ParamSlot::try_new(name, self.packed_width.get(), width)?; + let index = ParamIdx::from_usize(self.entries.len()); + self.insert(index, name, slot); + self.packed_width = PackedWidth::new(self.packed_width.get() + width); + Ok((index, slot)) + } + + /// Recover the registry an arena's nodes describe. + /// + /// Whole-arena, so a parameter only an output or event reads stays registered. + /// Conflicting or overlapping slots are refused; holes below the arena's length are + /// tolerated, and an index at or past it is refused before anything is allocated. + pub fn from_arena(arena: &Arena) -> Result { + // Each registration allocates a node, so a wider index can only come from forged bytes. + let in_range = |index: usize| { + if index < arena.len() { + Ok(index) + } else { + Err(CoreError::ParamIndexRange { + index, + n_nodes: arena.len(), + }) + } + }; + let mut registry = Self::new(); + for node in arena.nodes() { + match node { + Node::InputParameter { + name, + index, + offset, + width, + } => registry.absorb(name, in_range(*index)?, *offset, *width)?, + Node::TangentParameter { index } => registry.reserve(in_range(*index)?), + _ => {}, + } + } + registry.packed_width = registry.checked_packed_extent()?; + Ok(registry) + } + + /// The registration index `name` holds, or `None` if it holds none. + pub fn index_of(&self, name: &str) -> Option { + self.by_name.get(name).copied() + } + + /// Where parameter `idx` sits in the packed vector, if registered. + pub fn slot(&self, idx: ParamIdx) -> Option { + self.entry(idx).map(|entry| entry.slot) + } + + /// Every registered parameter in registration order, as `(index, name, slot)`. + /// + /// Unslotted indices are skipped; size tables by [`index_extent`](Self::index_extent). + pub fn registered(&self) -> impl Iterator { + self.entries + .iter() + .enumerate() + .filter_map(|(index, entry)| { + entry + .as_ref() + .map(|entry| (ParamIdx::from_usize(index), entry.name.as_str(), entry.slot)) + }) + } + + /// Registered names in registration order. + pub fn names(&self) -> Vec<&str> { + self.registered().map(|(_, name, _)| name).collect() + } + + /// One past the highest registration index; not the registered count. + pub const fn index_extent(&self) -> usize { + self.entries.len() + } + + /// How many indices actually carry a packed slot. + pub fn n_registered(&self) -> usize { + self.entries.iter().flatten().count() + } + + /// Length of the packed `inputs` slice these parameters need. + pub const fn packed_width(&self) -> PackedWidth { + self.packed_width + } + + fn entry(&self, idx: ParamIdx) -> Option<&ParamEntry> { + self.entries.get(idx.index())?.as_ref() + } + + /// Widen the index space so `index` is addressable, naming no parameter. + fn reserve(&mut self, index: usize) { + if self.entries.len() <= index { + self.entries.resize(index + 1, None); + } + } + + /// Take in one `InputParameter` node's fields, checking them against earlier nodes. + fn absorb( + &mut self, + name: &str, + index: usize, + offset: usize, + width: usize, + ) -> Result<(), CoreError> { + let slot = ParamSlot::try_new(name, offset, width)?; + self.reserve(index); + if let Some(existing) = self.entries[index].as_ref() { + if existing.name != name || existing.slot != slot { + return Err(CoreError::ParamSlotConflict { + index, + name: name.to_string(), + offset, + width, + other_name: existing.name.clone(), + other_offset: existing.slot.offset.index(), + other_width: existing.slot.width.index(), + }); + } + return Ok(()); + } + if let Some(&other_index) = self.by_name.get(name) { + return Err(CoreError::ParamNameConflict { + name: name.to_string(), + index, + other_index: other_index.index(), + }); + } + self.insert(ParamIdx::from_usize(index), name, slot); + Ok(()) + } + + /// Record `name` at `index`, widening the index space to reach it. + fn insert(&mut self, index: ParamIdx, name: &str, slot: ParamSlot) { + self.reserve(index.index()); + self.entries[index.index()] = Some(ParamEntry { + name: name.to_string(), + slot, + }); + self.by_name.insert(name.to_string(), index); + } + + /// One past the highest occupied packed element, refusing overlapping slots. + fn checked_packed_extent(&self) -> Result { + let mut ordered: Vec<(&str, ParamSlot)> = self + .registered() + .map(|(_, name, slot)| (name, slot)) + .collect(); + ordered.sort_by_key(|(_, slot)| slot.offset); + + let mut extent = 0usize; + let mut previous: Option<(&str, ParamSlot)> = None; + for (name, slot) in ordered { + if let Some((other_name, other)) = previous + && slot.offset.index() < other.offset.index() + other.width.index() + { + return Err(CoreError::ParamSlotOverlap { + name: name.to_string(), + offset: slot.offset.index(), + width: slot.width.index(), + other_name: other_name.to_string(), + other_offset: other.offset.index(), + other_width: other.width.index(), + }); + } + extent = slot.offset.index() + slot.width.index(); + previous = Some((name, slot)); + } + Ok(PackedWidth::new(extent)) + } +} + +/// Resolve a declared packed input width against `extent`; `None` adopts the extent. +/// +/// Below the extent is refused, not raised: the FFI borrows a buffer of the declared width. +pub const fn resolve_input_width( + declared: Option, + extent: PackedWidth, +) -> Result { + match declared { + Some(width) if width.get() < extent.get() => Err(CoreError::InputWidthTooSmall { + declared: width.get(), + extent: extent.get(), + }), + Some(width) => Ok(width), + None => Ok(extent), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A width-3 input then a scalar, so the two index spaces diverge. + fn wide_then_scalar() -> Arena { + let mut arena = Arena::new(); + let _wide = arena.alloc(Node::InputParameter { + name: "wide".to_string(), + index: 0, + offset: 0, + width: 3, + }); + let _scalar = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 1, + offset: 3, + width: 1, + }); + arena + } + + fn registry_of(arena: &Arena) -> ParamRegistry { + ParamRegistry::from_arena(arena).expect("a well-formed arena") + } + + fn input_node(registry: &mut ParamRegistry, name: &str, width: usize) -> Node { + let (index, slot) = registry.register(name, width).expect("a fresh name"); + Node::InputParameter { + name: name.to_string(), + index: index.index(), + offset: slot.offset.index(), + width: slot.width.index(), + } + } + + #[test] + fn a_vector_input_separates_the_two_spaces() { + let registry = registry_of(&wide_then_scalar()); + assert_eq!(registry.index_extent(), 2); + assert_eq!(registry.n_registered(), 2); + assert_eq!(registry.packed_width(), PackedWidth::new(4)); + assert_eq!(registry.slot(ParamIdx::new(2)), None); + + assert_eq!( + registry.slot(ParamIdx::new(1)).unwrap().offset, + InputOffset::new(3) + ); + assert_eq!( + registry.slot(ParamIdx::new(0)).unwrap().width, + Width::new(3) + ); + } + + #[test] + fn register_packs_each_name_after_the_last() { + let mut registry = ParamRegistry::new(); + let (wide, wide_slot) = registry.register("wide", 3).unwrap(); + let (k, k_slot) = registry.register("k", 1).unwrap(); + + assert_eq!((wide, k), (ParamIdx::new(0), ParamIdx::new(1))); + assert_eq!(wide_slot.offset, InputOffset::new(0)); + assert_eq!(k_slot.offset, InputOffset::new(3)); + assert_eq!(registry.packed_width(), PackedWidth::new(4)); + assert_eq!(registry.names(), vec!["wide", "k"]); + assert_eq!(registry.index_of("k"), Some(k)); + assert_eq!(registry.index_of("absent"), None); + } + + #[test] + fn re_registering_a_name_repeats_its_slot_and_refuses_another_width() { + let mut registry = ParamRegistry::new(); + let first = registry.register("k", 2).unwrap(); + assert_eq!(registry.register("k", 2).unwrap(), first); + assert_eq!(registry.packed_width(), PackedWidth::new(2)); + + assert!(matches!( + registry.register("k", 3), + Err(CoreError::ParamWidthConflict { + width: 3, + registered: 2, + .. + }) + )); + } + + #[test] + fn a_slot_ending_past_u32_max_is_an_error_not_a_panic() { + let max = u32::MAX as usize; + let mut registry = ParamRegistry::new(); + assert!(matches!( + registry.register("huge", max + 1), + Err(CoreError::ParamSlotRange { offset: 0, .. }) + )); + assert_eq!(registry, ParamRegistry::new()); + + registry + .register("fills", max) + .expect("ends exactly at u32::MAX"); + assert!(matches!( + registry.register("one_more", 1), + Err(CoreError::ParamSlotRange { offset, width: 1, .. }) if offset == max + )); + + let mut arena = Arena::new(); + let _ = arena.alloc(Node::InputParameter { + name: "far".to_string(), + index: 0, + offset: max, + width: 2, + }); + assert!(matches!( + ParamRegistry::from_arena(&arena), + Err(CoreError::ParamSlotRange { .. }) + )); + } + + #[test] + fn a_scan_recovers_what_register_built() { + let mut built = ParamRegistry::new(); + let mut arena = Arena::new(); + for (name, width) in [("wide", 3), ("k", 1), ("wide", 3)] { + let node = input_node(&mut built, name, width); + let _ = arena.alloc(node); + } + + assert_eq!(registry_of(&arena), built); + } + + #[test] + fn an_index_with_no_packed_slot_is_not_registered() { + let mut arena = Arena::new(); + let _a = arena.alloc(Node::InputParameter { + name: "a".to_string(), + index: 0, + offset: 0, + width: 1, + }); + for _ in 0..2 { + let _ = arena.alloc(Node::Time); + } + let _tangent_only = arena.alloc(Node::TangentParameter { index: 3 }); + + let registry = registry_of(&arena); + assert_eq!(registry.index_extent(), 4); + assert_eq!(registry.n_registered(), 1); + assert_eq!(registry.packed_width(), PackedWidth::new(1)); + + assert!(registry.slot(ParamIdx::new(0)).is_some()); + assert_eq!(registry.names(), ["a"]); + for unslotted in [1, 2, 3] { + let unslotted = ParamIdx::new(unslotted); + assert_eq!(registry.slot(unslotted), None, "{unslotted:?} has no slot"); + } + } + + #[test] + fn an_index_past_the_arena_is_an_error_not_an_allocation() { + for index in [5, usize::MAX] { + let tangent = Node::TangentParameter { index }; + let input = Node::InputParameter { + name: "k".to_string(), + index, + offset: 0, + width: 1, + }; + for node in [tangent, input] { + let mut arena = Arena::new(); + let _ = arena.alloc(Node::InputParameter { + name: "p".to_string(), + index: 0, + offset: 1, + width: 1, + }); + let _ = arena.alloc(node); + assert!(matches!( + ParamRegistry::from_arena(&arena), + Err(CoreError::ParamIndexRange { index: got, n_nodes: 2 }) if got == index + )); + } + } + } + + #[test] + fn two_slots_for_one_index_are_refused() { + let mut arena = Arena::new(); + for offset in [0, 7] { + let _ = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset, + width: 1, + }); + } + + assert!(matches!( + ParamRegistry::from_arena(&arena), + Err(CoreError::ParamSlotConflict { + index: 0, + offset: 7, + other_offset: 0, + .. + }) + )); + } + + #[test] + fn one_name_at_two_indices_is_refused() { + let mut arena = Arena::new(); + for (index, offset) in [(0, 0), (1, 1)] { + let _ = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index, + offset, + width: 1, + }); + } + + assert!(matches!( + ParamRegistry::from_arena(&arena), + Err(CoreError::ParamNameConflict { + index: 1, + other_index: 0, + .. + }) + )); + } + + #[test] + fn overlapping_slots_are_refused() { + let mut arena = Arena::new(); + let _ = arena.alloc(Node::InputParameter { + name: "wide".to_string(), + index: 0, + offset: 0, + width: 3, + }); + let _ = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 1, + offset: 2, + width: 1, + }); + + assert!(matches!( + ParamRegistry::from_arena(&arena), + Err(CoreError::ParamSlotOverlap { + offset: 2, + other_offset: 0, + other_width: 3, + .. + }) + )); + } + + #[test] + fn a_hole_in_the_packed_space_is_spanned() { + let mut arena = Arena::new(); + let _ = arena.alloc(Node::InputParameter { + name: "a".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let _ = arena.alloc(Node::InputParameter { + name: "b".to_string(), + index: 1, + offset: 4, + width: 1, + }); + + let registry = registry_of(&arena); + assert_eq!(registry.n_registered(), 2); + assert_eq!(registry.packed_width(), PackedWidth::new(5)); + } + + #[test] + fn an_arena_without_inputs_reports_both_spaces_empty() { + let registry = registry_of(&Arena::new()); + assert_eq!(registry, ParamRegistry::default()); + assert_eq!(registry.index_extent(), 0); + assert_eq!(registry.n_registered(), 0); + assert_eq!(registry.packed_width(), PackedWidth::new(0)); + assert_eq!(registry.names(), Vec::<&str>::new()); + } + + #[test] + fn a_declared_input_width_is_refused_below_the_extent_and_kept_above_it() { + let extent = PackedWidth::new(4); + assert!(matches!( + resolve_input_width(Some(PackedWidth::new(1)), extent), + Err(CoreError::InputWidthTooSmall { + declared: 1, + extent: 4 + }) + )); + assert_eq!( + resolve_input_width(Some(PackedWidth::new(4)), extent).unwrap(), + PackedWidth::new(4) + ); + assert_eq!( + resolve_input_width(Some(PackedWidth::new(6)), extent).unwrap(), + PackedWidth::new(6) + ); + assert_eq!(resolve_input_width(None, extent).unwrap(), extent); + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/simplify.rs b/packages/pybamm-rust/pybamm-core/src/simplify.rs new file mode 100644 index 0000000000..05b47af044 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/simplify.rs @@ -0,0 +1,1039 @@ +//! Expression simplification, CSE and DCE, chiefly to keep derivative expressions small. +//! +//! # Examples +//! +//! ``` +//! use pybamm_core::{Arena, Node, simplify}; +//! +//! let mut arena = Arena::new(); +//! let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); +//! let zero = arena.alloc(Node::Scalar(0.0)); +//! let expr = arena.alloc(Node::Add(x, zero)); +//! +//! let simplified = simplify(&mut arena, expr); +//! assert_eq!(simplified, x); +//! ``` + +use crate::arena::{Arena, Dag, DagRef, NodeId, NodeMap}; +use crate::ir::{BinaryOp, UnaryOp}; +use crate::node::{ArrayData, Node, Shape}; + +/// Constant-fold and apply algebraic identities, as `CasADi`'s SX simplifies on the fly. +/// +/// Like `CasADi`, it folds `0 * x`, `0 / x` and `x - x` to zero and `x / x` to one, so the +/// result matches the input only where every value is finite, up to the sign of zero. +/// Constant powers lower only where the rewrite takes `pow`'s value bitwise. +#[must_use] +pub fn simplify(arena: &mut Arena, root: NodeId) -> NodeId { + let mut simplified: NodeMap = NodeMap::new(arena.len()); + let mut lens: NodeMap = NodeMap::new(arena.len()); + for id in arena.topological_order(root) { + let result = simplify_node(arena, id, &simplified, &mut lens); + simplified.insert(id, result); + } + simplified.rebuilt(root) +} + +/// Exact literal match, so tiny nonzero coefficients (e.g. 1e-17) are never treated as 0 or 1. +fn is_scalar(arena: &Arena, id: NodeId, value: f64) -> bool { + arena.get(id).as_scalar() == Some(value) +} + +fn is_zero(arena: &Arena, id: NodeId) -> bool { + is_scalar(arena, id, 0.0) +} + +fn is_one(arena: &Arena, id: NodeId) -> bool { + is_scalar(arena, id, 1.0) +} + +fn is_neg_one(arena: &Arena, id: NodeId) -> bool { + is_scalar(arena, id, -1.0) +} + +/// [`node_width`](crate::ir::node_width) of a node, memoized via `lens`. +/// +/// Walks with an explicit stack, so graph depth is bounded by the heap. +pub fn node_len(arena: &Arena, id: NodeId, lens: &mut NodeMap) -> usize { + if let Some(&len) = lens.get(id) { + return len; + } + let mut pending = vec![id]; + while let Some(&top) = pending.last() { + if lens.contains_key(top) { + pending.pop(); + continue; + } + let n_pending = pending.len(); + let len = crate::ir::node_width(arena, arena.get(top), |child| { + lens.get(child).copied().unwrap_or_else(|| { + pending.push(child); + 0 + }) + }); + // Valid only if every child it read was already measured. + if pending.len() == n_pending { + pending.pop(); + lens.insert(top, len); + } + } + *lens.get(id).expect("the walk ends with `id` measured") +} + +/// Simplify `id`, whose children are already simplified in `simplified`. +/// +/// A node no rule rewrites is reused, or copied over its simplified children. +fn simplify_node( + arena: &mut Arena, + id: NodeId, + simplified: &NodeMap, + lens: &mut NodeMap, +) -> NodeId { + let s = |child: NodeId| simplified.rebuilt(child); + let rewritten = match *arena.get(id) { + Node::Add(lhs, rhs) => simplify_add(arena, s(lhs), s(rhs)), + Node::Sub(lhs, rhs) => simplify_sub(arena, s(lhs), s(rhs), lens), + Node::Mul(lhs, rhs) => simplify_mul(arena, s(lhs), s(rhs), lens), + Node::Div(lhs, rhs) => simplify_div(arena, s(lhs), s(rhs), lens), + Node::Pow(base, exp) => simplify_pow(arena, s(base), s(exp), lens), + Node::Minimum(lhs, rhs) => fold_binary(arena, s(lhs), s(rhs), BinaryOp::Minimum), + Node::Maximum(lhs, rhs) => fold_binary(arena, s(lhs), s(rhs), BinaryOp::Maximum), + Node::Modulo(lhs, rhs) => fold_binary(arena, s(lhs), s(rhs), BinaryOp::Modulo), + Node::Hypot(lhs, rhs) => fold_binary(arena, s(lhs), s(rhs), BinaryOp::Hypot), + + Node::Neg(child) => simplify_neg(arena, s(child)), + Node::Abs(child) => fold_unary(arena, s(child), UnaryOp::Abs), + Node::Sqrt(child) => fold_unary(arena, s(child), UnaryOp::Sqrt), + Node::Exp(child) => fold_unary(arena, s(child), UnaryOp::Exp), + Node::Log(child) => fold_unary(arena, s(child), UnaryOp::Log), + Node::Sin(child) => fold_unary(arena, s(child), UnaryOp::Sin), + Node::Cos(child) => fold_unary(arena, s(child), UnaryOp::Cos), + Node::Tanh(child) => fold_unary(arena, s(child), UnaryOp::Tanh), + Node::Sinh(child) => fold_unary(arena, s(child), UnaryOp::Sinh), + Node::Cosh(child) => fold_unary(arena, s(child), UnaryOp::Cosh), + Node::Arcsinh(child) => fold_unary(arena, s(child), UnaryOp::Arcsinh), + Node::Arctan(child) => fold_unary(arena, s(child), UnaryOp::Arctan), + Node::Erf(child) => fold_unary(arena, s(child), UnaryOp::Erf), + Node::Sign(child) => fold_unary(arena, s(child), UnaryOp::Sign), + Node::Floor(child) => fold_unary(arena, s(child), UnaryOp::Floor), + Node::Ceiling(child) => fold_unary(arena, s(child), UnaryOp::Ceiling), + + Node::ReduceArgSelect { basis, .. } + if is_zero(arena, s(basis)) + || matches!(arena.get(s(basis)), Node::ZeroVector { .. }) => + { + Some(arena.alloc(Node::Scalar(0.0))) + }, + + Node::Scalar(_) + | Node::Array(_) + | Node::ZeroVector { .. } + | Node::SparseMatrix(_) + | Node::StateVector { .. } + | Node::StateVectorDot { .. } + | Node::InputParameter { .. } + | Node::Time + | Node::TangentStateVector { .. } + | Node::TangentParameter { .. } + | Node::MatMul(..) + | Node::EqualHeaviside(..) + | Node::NotEqualHeaviside(..) + | Node::Equality(..) + | Node::MaxReduce(_) + | Node::MinReduce(_) + | Node::ReduceArgSelect { .. } + | Node::Index { .. } + | Node::Concat(_) + | Node::Interpolant1DLinear { .. } + | Node::Interpolant1DLinearDeriv { .. } + | Node::Interpolant1DCubic { .. } + | Node::Interpolant1DCubicDeriv { .. } + | Node::InterpolantNd { .. } + | Node::InterpolantNdPartial { .. } + | Node::Conditional { .. } => None, + }; + rewritten.unwrap_or_else(|| reuse_or_rebuild(arena, id, simplified)) +} + +/// `id` itself when no child changed, else a copy of it over the simplified children. +fn reuse_or_rebuild(arena: &mut Arena, id: NodeId, simplified: &NodeMap) -> NodeId { + let node = arena.get(id); + let mut changed = false; + node.for_each_child(|child| changed |= simplified.rebuilt(child) != child); + if !changed { + return id; + } + let copy = node.map_children(|child| simplified.rebuilt(child)); + arena.alloc(copy) +} + +/// `const op const -> const`, valued exactly as evaluation would. +fn fold_binary(arena: &mut Arena, lhs: NodeId, rhs: NodeId, op: BinaryOp) -> Option { + let (a, b) = (arena.get(lhs).as_scalar()?, arena.get(rhs).as_scalar()?); + Some(arena.alloc(Node::Scalar(op.apply(a, b)))) +} + +/// `op(const) -> const`, valued exactly as evaluation would. +fn fold_unary(arena: &mut Arena, child: NodeId, op: UnaryOp) -> Option { + let v = arena.get(child).as_scalar()?; + Some(arena.alloc(Node::Scalar(op.apply(v)))) +} + +/// `x + 0 -> x`, `0 + x -> x`, `const + const -> const`. +fn simplify_add(arena: &mut Arena, lhs: NodeId, rhs: NodeId) -> Option { + if let Some(sum) = fold_binary(arena, lhs, rhs, BinaryOp::Add) { + return Some(sum); + } + + if is_zero(arena, lhs) { + return Some(rhs); + } + + if is_zero(arena, rhs) { + return Some(lhs); + } + + None +} + +/// `x - 0 -> x`, `0 - x -> -x`, `x - x -> 0`, `const - const -> const`. +fn simplify_sub( + arena: &mut Arena, + lhs: NodeId, + rhs: NodeId, + lens: &mut NodeMap, +) -> Option { + if let Some(difference) = fold_binary(arena, lhs, rhs, BinaryOp::Sub) { + return Some(difference); + } + + if is_zero(arena, rhs) { + return Some(lhs); + } + + if is_zero(arena, lhs) { + return Some(arena.alloc(Node::Neg(rhs))); + } + + if lhs == rhs { + let len = node_len(arena, lhs, lens); + return Some(arena.alloc(Node::zero(len))); + } + + None +} + +/// `1 * x -> x`, `x * 1 -> x`, `-1 * x -> -x`, `0 * x -> 0`, `const * const -> const`. +fn simplify_mul( + arena: &mut Arena, + lhs: NodeId, + rhs: NodeId, + lens: &mut NodeMap, +) -> Option { + if let Some(product) = fold_binary(arena, lhs, rhs, BinaryOp::Mul) { + return Some(product); + } + + if is_one(arena, lhs) { + return Some(rhs); + } + + if is_one(arena, rhs) { + return Some(lhs); + } + + if is_neg_one(arena, lhs) { + return Some(arena.alloc(Node::Neg(rhs))); + } + + if is_neg_one(arena, rhs) { + return Some(arena.alloc(Node::Neg(lhs))); + } + + if is_zero(arena, lhs) || is_zero(arena, rhs) { + let len = node_len(arena, lhs, lens).max(node_len(arena, rhs, lens)); + return Some(arena.alloc(Node::zero(len))); + } + + None +} + +/// `x / 1 -> x`, `0 / x -> 0`, `x / x -> 1`, `const / const -> const`. +fn simplify_div( + arena: &mut Arena, + lhs: NodeId, + rhs: NodeId, + lens: &mut NodeMap, +) -> Option { + if let Some(quotient) = fold_binary(arena, lhs, rhs, BinaryOp::Div) { + return Some(quotient); + } + + if is_one(arena, rhs) { + return Some(lhs); + } + + if is_zero(arena, lhs) { + let len = node_len(arena, rhs, lens); + return Some(arena.alloc(Node::zero(len))); + } + + // Scalar only: a vector x / x is a ones vector. + if lhs == rhs && node_len(arena, lhs, lens) == 1 { + return Some(arena.alloc(Node::Scalar(1.0))); + } + + None +} + +/// `const^const -> const`, `1^x -> 1`, and `x^const` as [`lower_constant_pow`] rewrites it. +fn simplify_pow( + arena: &mut Arena, + base: NodeId, + exp: NodeId, + lens: &mut NodeMap, +) -> Option { + if let Some(power) = fold_binary(arena, base, exp, BinaryOp::Pow) { + return Some(power); + } + + // Scalar only: 1^x broadcasts to the exponent's length. + if is_one(arena, base) && node_len(arena, exp, lens) == 1 { + return Some(base); + } + + let e = arena.get(exp).as_scalar()?; + let len = node_len(arena, base, lens); + lower_constant_pow(arena, base, e, len) +} + +/// Rewrite `base^e`, `base` of length `len`, where the rewrite takes `pow`'s value bitwise. +/// +/// Only `e` in {0, 1, 2, -1} qualify: `1`, `x`, `x*x` and `1/x` round once from exact +/// operands, so equal a correctly rounded `pow` for every base, special values included, and +/// the evaluator's `pow` takes the same forms. A chain such as `x*(x*x)` rounds twice, and +/// `sqrt(x)` differs at `-0` and `-inf`. +#[allow(clippy::float_cmp)] +fn lower_constant_pow(arena: &mut Arena, base: NodeId, e: f64, len: usize) -> Option { + // A zero-width base against the scalar exponent is invalid; `x * x` would make it valid. + if len == 0 { + return None; + } + if e == 0.0 { + let ones = if len == 1 { + Node::Scalar(1.0) + } else { + let ones = ArrayData::try_new(vec![1.0; len], Shape::vector(len)) + .expect("a vector's shape matches its length"); + Node::Array(Box::new(ones)) + }; + return Some(arena.alloc(ones)); + } + if e == 1.0 { + return Some(base); + } + if e == 2.0 { + return Some(arena.alloc(Node::Mul(base, base))); + } + if e == -1.0 { + let one = arena.alloc(Node::Scalar(1.0)); + return Some(arena.alloc(Node::Div(one, base))); + } + None +} + +/// `-const -> const`, `--x -> x`. +fn simplify_neg(arena: &mut Arena, child: NodeId) -> Option { + if let Some(negated) = fold_unary(arena, child, UnaryOp::Neg) { + return Some(negated); + } + + match *arena.get(child) { + Node::Neg(inner) => Some(inner), + _ => None, + } +} + +/// Common subexpression elimination into a new arena of unique reachable nodes. +/// +/// The output is laid out in [`dce`]'s order, so a `dce` after it is the identity. +#[must_use] +pub fn cse(dag: DagRef<'_>) -> Dag { + use rustc_hash::FxHashMap; + use rustc_hash::FxHasher; + use std::hash::Hasher; + + use crate::node::{structural_hash, structurally_identical}; + + let (arena, root) = (dag.arena(), dag.root()); + let mut new_arena = Arena::new(); + let mut old_to_new: NodeMap = NodeMap::new(arena.len()); + // Newest node per hash; each node links to the one it displaced, by new id. + let mut newest_with_hash: FxHashMap = FxHashMap::default(); + let mut older_with_hash: Vec> = Vec::new(); + + for old_id in arena.topological_order(root) { + let remapped = arena.get(old_id).map_children(|c| old_to_new.rebuilt(c)); + let mut hasher = FxHasher::default(); + structural_hash(&remapped, &mut hasher); + let hash = hasher.finish(); + + let canonical = std::iter::successors(newest_with_hash.get(&hash).copied(), |id| { + older_with_hash[id.index()] + }) + .find(|&existing| structurally_identical(new_arena.get(existing), &remapped)); + let new_id = canonical.unwrap_or_else(|| { + let new_id = new_arena.alloc(remapped); + older_with_hash.push(newest_with_hash.insert(hash, new_id)); + new_id + }); + old_to_new.insert(old_id, new_id); + } + + Dag::new(new_arena, old_to_new.rebuilt(root)) +} + +/// Compile-prep pipeline: simplify -> zero propagation -> CSE. +/// +/// The result is dead-code free: CSE keeps only what the root reaches. +#[must_use] +pub fn simplify_pipeline(mut dag: Dag) -> Dag { + use crate::zero_propagate::zero_propagate; + + dag.rewrite_in_place(simplify); + let propagated = zero_propagate(dag.as_ref()); + // The largest of the three arenas, so free it before `cse` builds its own. + drop(dag); + cse(propagated.as_ref()) +} + +/// Dead code elimination into a new arena of nodes reachable from the root. +#[must_use] +pub fn dce(dag: DagRef<'_>) -> Dag { + let (arena, root) = (dag.arena(), dag.root()); + let mut new_arena = Arena::new(); + let mut old_to_new: NodeMap = NodeMap::new(arena.len()); + + for old_id in arena.topological_order(root) { + let copy = arena.get(old_id).map_children(|c| old_to_new.rebuilt(c)); + old_to_new.insert(old_id, new_arena.alloc(copy)); + } + + Dag::new(new_arena, old_to_new.rebuilt(root)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::node::Node; + + #[test] + fn simplify_pipeline_matches_manual_pass_order() { + use crate::zero_propagate::zero_propagate; + + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let s = arena.alloc(Node::Sin(y0)); + let one = arena.alloc(Node::Scalar(1.0)); + let zero = arena.alloc(Node::Scalar(0.0)); + let m = arena.alloc(Node::Mul(s, one)); + let decorated = arena.alloc(Node::Add(m, zero)); + + let simplified = simplify_pipeline(Dag::new(arena.clone(), decorated)); + + let mut manual = arena; + let r = simplify(&mut manual, decorated); + let propagated = zero_propagate(DagRef::new(&manual, r)); + let shared = cse(propagated.as_ref()); + let pruned = dce(shared.as_ref()); + + assert_eq!(simplified.arena().len(), pruned.arena().len()); + assert_eq!(simplified.root(), pruned.root()); + assert_eq!( + simplified.arena().get(simplified.root()), + pruned.arena().get(pruned.root()) + ); + } + + #[test] + fn cse_output_is_already_dead_code_free() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y_twin = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let two_twin = arena.alloc(Node::Scalar(2.0)); + let sin_a = arena.alloc(Node::Sin(y)); + let sin_b = arena.alloc(Node::Sin(y_twin)); + let exp_a = arena.alloc(Node::Exp(sin_a)); + let exp_b = arena.alloc(Node::Exp(sin_b)); + let _dead = arena.alloc(Node::Cos(y)); + let scaled_a = arena.alloc(Node::Mul(two, exp_a)); + let scaled_b = arena.alloc(Node::Mul(two_twin, exp_b)); + // The same operands under two ops, which must stay apart. + let sum = arena.alloc(Node::Add(scaled_a, sin_a)); + let difference = arena.alloc(Node::Sub(scaled_b, sin_b)); + let root = arena.alloc(Node::Concat(vec![sum, difference, sin_b])); + + let shared = cse(DagRef::new(&arena, root)); + let pruned = dce(shared.as_ref()); + + assert_eq!( + shared.arena().len(), + 8, + "duplicates merged, dead node dropped" + ); + let out = shared.arena(); + let Node::Concat(parts) = out.get(shared.root()) else { + panic!("the concat must survive") + }; + let (&Node::Add(scaled, sin), &Node::Sub(scaled_twin, sin_twin)) = + (out.get(parts[0]), out.get(parts[1])) + else { + panic!("the Add and the Sub must stay apart") + }; + assert_eq!((scaled, sin), (scaled_twin, sin_twin), "nested twins merge"); + assert_eq!(parts[2], sin); + assert_eq!(shared.arena().nodes(), pruned.arena().nodes()); + assert_eq!(shared.root(), pruned.root()); + } + + /// Deeper than a recursive walk survives on a 2 MiB test-thread stack. + const DEEP: usize = 200_000; + + #[test] + fn node_len_handles_a_deep_chain() { + let mut arena = Arena::new(); + let mut id = arena.alloc(Node::StateVector { start: 0, end: 3 }); + for _ in 0..DEEP { + id = arena.alloc(Node::Sin(id)); + } + let mut lens = NodeMap::new(arena.len()); + assert_eq!(node_len(&arena, id, &mut lens), 3); + } + + #[test] + fn a_zero_width_conditional_measures_zero_as_lowering_does() { + let mut arena = Arena::new(); + let selector = arena.alloc(Node::Scalar(1.0)); + let empty = arena.alloc(Node::StateVector { start: 0, end: 0 }); + let cond = arena.alloc(Node::Conditional { + selector, + branches: vec![empty, empty], + }); + let ratio = arena.alloc(Node::Div(cond, cond)); + + let lowered = crate::ir::infer_sizes(&arena, &arena.topological_order(ratio)); + let mut lens = NodeMap::new(arena.len()); + assert_eq!(node_len(&arena, cond, &mut lens), lowered[cond.index()]); + + let simplified = simplify(&mut arena, ratio); + assert!( + matches!(arena.get(simplified), Node::Div(..)), + "a zero-width x / x is not the scalar 1" + ); + } + + #[test] + fn simplify_handles_a_deep_chain() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let one = arena.alloc(Node::Scalar(1.0)); + let mut id = y; + for _ in 0..DEEP { + let sin = arena.alloc(Node::Sin(id)); + id = arena.alloc(Node::Mul(sin, one)); + } + + let mut root = simplify(&mut arena, id); + for _ in 0..DEEP { + let Node::Sin(child) = *arena.get(root) else { + panic!("every `* 1` must fold away, leaving the sines") + }; + root = child; + } + assert_eq!(root, y); + } + + /// `x^e` over a scalar state `x`, simplified and compiled once. + fn simplified_pow_of_state(e: f64) -> impl FnMut(f64) -> f64 { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let exp = arena.alloc(Node::Scalar(e)); + let expr = arena.alloc(Node::Pow(x, exp)); + let simplified = simplify(&mut arena, expr); + let compiled = crate::eval::CompiledExpr::new(DagRef::new(&arena, simplified)); + let mut scratch = vec![0.0; compiled.scratch_len()]; + move |y| compiled.eval(&mut scratch, 0.0, &[y], &[], &[])[0] + } + + /// `id` rendered infix, with the leaf `x` as `x`. + fn render(arena: &Arena, id: NodeId, x: NodeId) -> String { + let r = |child: NodeId| render(arena, child, x); + match arena.get(id) { + _ if id == x => "x".to_string(), + Node::Scalar(v) => v.to_string(), + Node::Array(ones) => format!("{:?}", ones.data), + Node::ZeroVector { len } => format!("zeros({len})"), + Node::Neg(a) => format!("-{}", r(*a)), + Node::Sqrt(a) => format!("sqrt({})", r(*a)), + Node::Add(a, b) => format!("({}+{})", r(*a), r(*b)), + Node::Sub(a, b) => format!("({}-{})", r(*a), r(*b)), + Node::Mul(a, b) => format!("({}*{})", r(*a), r(*b)), + Node::Div(a, b) => format!("({}/{})", r(*a), r(*b)), + Node::Pow(a, b) => format!("pow({},{})", r(*a), r(*b)), + other => panic!("unexpected {other:?}"), + } + } + + /// The node `build` makes over a state `x` of length `len`, simplified and rendered. + fn simplified(len: usize, build: impl FnOnce(&mut Arena, NodeId) -> Node) -> String { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: len }); + let node = build(&mut arena, x); + let expr = arena.alloc(node); + let root = simplify(&mut arena, expr); + render(&arena, root, x) + } + + /// `x` raised to a constant, simplified and rendered infix over the leaf `x`. + fn lowered_pow(len: usize, e: f64) -> String { + simplified(len, |arena, x| Node::Pow(x, arena.alloc(Node::Scalar(e)))) + } + + #[test] + fn identities_fold_as_casadi_folds_them() { + type Build = fn(&mut Arena, NodeId) -> Node; + fn lit(arena: &mut Arena, value: f64) -> NodeId { + arena.alloc(Node::Scalar(value)) + } + let cases: [(usize, Build, &str); 26] = [ + (1, |a, x| Node::Mul(x, lit(a, 1.0)), "x"), + (1, |a, x| Node::Mul(lit(a, 1.0), x), "x"), + (1, |a, x| Node::Add(x, lit(a, 0.0)), "x"), + (1, |a, x| Node::Add(lit(a, 0.0), x), "x"), + (1, |a, x| Node::Sub(x, lit(a, 0.0)), "x"), + (1, |a, x| Node::Sub(x, lit(a, -0.0)), "x"), + (1, |a, x| Node::Sub(lit(a, 0.0), x), "-x"), + (1, |a, x| Node::Sub(lit(a, -0.0), x), "-x"), + (1, |a, x| Node::Div(x, lit(a, 1.0)), "x"), + (1, |a, x| Node::Mul(lit(a, -1.0), x), "-x"), + (1, |a, x| Node::Mul(x, lit(a, -1.0)), "-x"), + (1, |a, x| Node::Pow(lit(a, 1.0), x), "1"), + (1, |a, x| Node::Neg(a.alloc(Node::Neg(x))), "x"), + (1, |a, x| Node::Mul(lit(a, 0.0), x), "0"), + (1, |a, x| Node::Mul(x, lit(a, 0.0)), "0"), + (1, |a, x| Node::Div(lit(a, 0.0), x), "0"), + (1, |_, x| Node::Sub(x, x), "0"), + (1, |_, x| Node::Div(x, x), "1"), + // An exact match only: a tiny coefficient is not zero. + ( + 1, + |a, x| Node::Mul(lit(a, 1e-17), x), + "(0.00000000000000001*x)", + ), + (3, |a, x| Node::Mul(lit(a, 0.0), x), "zeros(3)"), + (3, |a, x| Node::Mul(x, lit(a, 0.0)), "zeros(3)"), + (3, |_, x| Node::Sub(x, x), "zeros(3)"), + (3, |a, x| Node::Div(lit(a, 0.0), x), "zeros(3)"), + // A vector `x / x` is a ones vector and `1^x` broadcasts, so neither folds. + (3, |_, x| Node::Div(x, x), "(x/x)"), + (3, |a, x| Node::Pow(lit(a, 1.0), x), "pow(1,x)"), + (3, |a, x| Node::Sub(x, lit(a, -0.0)), "x"), + ]; + for (len, build, expected) in cases { + let mut input = Arena::new(); + let x = input.alloc(Node::StateVector { start: 0, end: len }); + let node = build(&mut input, x); + let expr = input.alloc(node); + let label = render(&input, expr, x); + assert_eq!( + simplified(len, build), + expected, + "{label}, x of length {len}" + ); + } + } + + #[test] + fn constant_powers_lower_only_where_exact() { + let cases = [ + (1, 0.0, "1"), + (1, 1.0, "x"), + (1, 2.0, "(x*x)"), + (1, -1.0, "(1/x)"), + (1, 3.0, "pow(x,3)"), + (1, -2.0, "pow(x,-2)"), + (1, 0.5, "pow(x,0.5)"), + (1, -0.5, "pow(x,-0.5)"), + (1, 0.25, "pow(x,0.25)"), + (3, 0.0, "[1.0, 1.0, 1.0]"), + (3, 1.0, "x"), + (3, 2.0, "(x*x)"), + (3, -1.0, "(1/x)"), + (3, 4.0, "pow(x,4)"), + (3, 0.5, "pow(x,0.5)"), + (0, 2.0, "pow(x,2)"), + ]; + for (len, e, expected) in cases { + assert_eq!(lowered_pow(len, e), expected, "x^{e}, x of length {len}"); + } + } + + /// Whatever `simplify` makes of a constant power takes `pow`'s value bit for bit, at + /// special values and across the binade range; a NaN need only meet a NaN. + #[test] + fn lowered_powers_take_pow_values_bitwise() { + let special = [ + 0.0, + -0.0, + f64::INFINITY, + f64::NEG_INFINITY, + f64::NAN, + f64::MIN_POSITIVE, + -f64::MIN_POSITIVE, + 5e-324, + -5e-324, + f64::MAX, + -f64::MAX, + 1.0, + -1.0, + 1.5e154, + 3e-162, + ]; + let mut state = 0x9E37_79B9_7F4A_7C15_u64; + let swept = std::iter::repeat_with(|| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + f64::from_bits(state) + }); + let bases: Vec = special.into_iter().chain(swept.take(20_000)).collect(); + let exponents = [ + 0.0, -0.0, 1.0, 2.0, -1.0, 3.0, 4.0, -2.0, 0.5, -0.5, 0.25, -0.25, 1.5, 100.0, -100.0, + ]; + for e in exponents { + let mut lowered = simplified_pow_of_state(e); + for &y in &bases { + let (got, want) = (lowered(y), BinaryOp::Pow.apply(y, e)); + assert!( + got.to_bits() == want.to_bits() || (got.is_nan() && want.is_nan()), + "y={y:e}, e={e}: lowered {got:e} != pow {want:e}" + ); + } + } + } + + #[test] + fn constant_operands_fold_to_their_value() { + type Binary = fn(NodeId, NodeId) -> Node; + type Unary = fn(NodeId) -> Node; + + let binary: [(Binary, f64, f64, f64); 5] = [ + (Node::Add, 2.0, 3.0, 5.0), + (Node::Sub, 5.0, 3.0, 2.0), + (Node::Mul, 3.0, 4.0, 12.0), + (Node::Div, 12.0, 4.0, 3.0), + (Node::Pow, 2.0, 3.0, 8.0), + ]; + for (op, a, b, expected) in binary { + let mut arena = Arena::new(); + let (lhs, rhs) = (arena.alloc(Node::Scalar(a)), arena.alloc(Node::Scalar(b))); + let expr = arena.alloc(op(lhs, rhs)); + let simplified = simplify(&mut arena, expr); + assert_eq!( + arena.get(simplified), + &Node::Scalar(expected), + "{:?}", + arena.get(expr) + ); + } + + // Bitwise against the interpreter's op too: `sign(0)` is 0, not `f64::signum`'s 1. + let unary: [(Unary, UnaryOp, f64, f64); 8] = [ + (Node::Neg, UnaryOp::Neg, 5.0, -5.0), + (Node::Sin, UnaryOp::Sin, 0.0, 0.0), + (Node::Cos, UnaryOp::Cos, 0.0, 1.0), + (Node::Exp, UnaryOp::Exp, 0.0, 1.0), + (Node::Sqrt, UnaryOp::Sqrt, 4.0, 2.0), + (Node::Erf, UnaryOp::Erf, 0.0, 0.0), + (Node::Erf, UnaryOp::Erf, 0.5, 0.520_499_877_813_046_5), + (Node::Sign, UnaryOp::Sign, 0.0, 0.0), + ]; + for (node, op, x, expected) in unary { + let mut arena = Arena::new(); + let operand = arena.alloc(Node::Scalar(x)); + let expr = arena.alloc(node(operand)); + let simplified = simplify(&mut arena, expr); + let label = format!("{:?}", arena.get(expr)); + let &Node::Scalar(folded) = arena.get(simplified) else { + panic!("{label} must fold to a scalar") + }; + assert_eq!( + folded.to_bits(), + op.apply(x).to_bits(), + "{label}: {folded:?}" + ); + assert_eq!(folded.to_bits(), expected.to_bits(), "{label}: {folded:?}"); + } + } + + #[test] + fn test_nested_simplification() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let one = arena.alloc(Node::Scalar(1.0)); + let x_plus_zero = arena.alloc(Node::Add(x, zero)); + let expr = arena.alloc(Node::Mul(x_plus_zero, one)); + + let simplified = simplify(&mut arena, expr); + + assert_eq!(simplified, x); + } + + #[test] + fn test_nested_simplification_complex() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let one = arena.alloc(Node::Scalar(1.0)); + + let x_times_one = arena.alloc(Node::Mul(x, one)); + let lhs = arena.alloc(Node::Add(x_times_one, zero)); + let y_times_zero = arena.alloc(Node::Mul(y, zero)); + let expr = arena.alloc(Node::Sub(lhs, y_times_zero)); + + assert_eq!(simplify(&mut arena, expr), x); + } + + #[test] + fn test_simplify_index() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 10 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let x_plus_zero = arena.alloc(Node::Add(x, zero)); + let expr = arena.alloc(Node::Index { + child: x_plus_zero, + start: 0, + end: 5, + }); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Index { child, start, end } => { + assert_eq!(*child, x); + assert_eq!(*start, 0); + assert_eq!(*end, 5); + }, + _ => panic!("Expected Index node"), + } + } + + #[test] + fn test_simplify_concat() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let y = arena.alloc(Node::StateVector { start: 5, end: 10 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let x_plus_zero = arena.alloc(Node::Add(x, zero)); + let expr = arena.alloc(Node::Concat(vec![x_plus_zero, y])); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Concat(children) => { + assert_eq!(children.len(), 2); + assert_eq!(children[0], x); + assert_eq!(children[1], y); + }, + _ => panic!("Expected Concat node"), + } + } + + #[test] + fn test_memoization() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let x_plus_zero = arena.alloc(Node::Add(x, zero)); + let expr = arena.alloc(Node::Add(x_plus_zero, x_plus_zero)); + + let simplified = simplify(&mut arena, expr); + + match arena.get(simplified) { + Node::Add(l, r) => { + assert_eq!(*l, x); + assert_eq!(*r, x); + }, + _ => panic!("Expected Add node"), + } + } + + #[test] + fn an_already_simple_graph_is_reused_not_copied() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let two = arena.alloc(Node::Scalar(2.0)); + let sin = arena.alloc(Node::Sin(y)); + let scaled = arena.alloc(Node::Mul(two, sin)); + let root = arena.alloc(Node::Sub(scaled, y)); + let len = arena.len(); + + assert_eq!(simplify(&mut arena, root), root); + assert_eq!(arena.len(), len, "nothing to rewrite, so nothing allocated"); + } + + #[test] + fn a_shared_binary_fold_is_memoised() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let three = arena.alloc(Node::Scalar(3.0)); + let four = arena.alloc(Node::Scalar(4.0)); + let hypot = arena.alloc(Node::Hypot(three, four)); + let left = arena.alloc(Node::Mul(y, hypot)); + let right = arena.alloc(Node::Mul(hypot, y)); + let root = arena.alloc(Node::Concat(vec![left, right])); + + let simplified = simplify(&mut arena, root); + + let Node::Concat(children) = arena.get(simplified) else { + panic!("the concat must survive") + }; + let (&Node::Mul(_, folded_left), &Node::Mul(folded_right, _)) = + (arena.get(children[0]), arena.get(children[1])) + else { + panic!("both products must survive") + }; + assert_eq!(arena.get(folded_left), &Node::Scalar(5.0)); + assert_eq!(folded_left, folded_right, "one fold, shared by both users"); + } + + #[test] + fn test_leaf_nodes_unchanged() { + let mut arena = Arena::new(); + + let scalar = arena.alloc(Node::Scalar(42.0)); + assert_eq!(simplify(&mut arena, scalar), scalar); + + let sv = arena.alloc(Node::StateVector { start: 0, end: 10 }); + assert_eq!(simplify(&mut arena, sv), sv); + + let time = arena.alloc(Node::Time); + assert_eq!(simplify(&mut arena, time), time); + + let param = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + assert_eq!(simplify(&mut arena, param), param); + } + + #[test] + fn cse_distinguishes_state_vectors_with_different_ranges() { + let mut arena = Arena::new(); + let y1 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y2 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let root = arena.alloc(Node::Add(y1, y2)); + + let shared = cse(DagRef::new(&arena, root)); + assert_eq!( + shared.arena().len(), + 3, + "StateVectors with different ranges must remain distinct" + ); + } + + /// Regression: `FxHash` cancels paired sign flips, and `==` equates `0.0` with `-0.0`. + #[test] + fn cse_keeps_arrays_that_differ_only_in_signed_zeros() { + use crate::eval::CompiledExpr; + use crate::node::{ArrayData, Shape}; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let one = arena.alloc(Node::Scalar(1.0)); + let mut reciprocal = |data: Vec| { + let array = ArrayData::try_new(data, Shape::vector(3)).expect("a valid array"); + let coeff = arena.alloc(Node::Array(Box::new(array))); + let scaled = arena.alloc(Node::Mul(coeff, y)); + arena.alloc(Node::Div(one, scaled)) + }; + let a = reciprocal(vec![0.0, -0.0, 3.0]); + let b = reciprocal(vec![-0.0, 0.0, 3.0]); + let root = arena.alloc(Node::Concat(vec![a, b])); + + let shared = cse(DagRef::new(&arena, root)); + assert_eq!( + shared.arena().len(), + arena.len(), + "no two nodes are identical" + ); + + let eval = |dag: DagRef<'_>| { + let compiled = CompiledExpr::new(dag); + let mut scratch = vec![0.0; compiled.scratch_len()]; + let values = compiled.eval(&mut scratch, 0.0, &[1.0; 3], &[], &[]); + values.iter().map(|v| v.to_bits()).collect::>() + }; + assert_eq!(eval(shared.as_ref()), eval(DagRef::new(&arena, root))); + } + + #[test] + fn dce_keeps_exactly_the_reachable_nodes_and_their_sharing() { + let mut arena = Arena::new(); + let _unreachable = arena.alloc(Node::Scalar(999.0)); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let _unreachable_too = arena.alloc(Node::Scalar(888.0)); + let one = arena.alloc(Node::Scalar(1.0)); + let sum = arena.alloc(Node::Add(x, one)); + let diamond = arena.alloc(Node::Mul(sum, sum)); + let negated = arena.alloc(Node::Neg(diamond)); + let root = arena.alloc(Node::Abs(negated)); + + let pruned = dce(DagRef::new(&arena, root)); + + let out = pruned.arena(); + assert_eq!(out.len(), arena.len() - 2, "only the unreachable nodes go"); + let &Node::Abs(abs_child) = out.get(pruned.root()) else { + panic!("expected Abs") + }; + let &Node::Neg(neg_child) = out.get(abs_child) else { + panic!("expected Neg") + }; + let &Node::Mul(l, r) = out.get(neg_child) else { + panic!("expected Mul") + }; + assert_eq!(l, r, "both operands reference the same Add node"); + let &Node::Add(state, constant) = out.get(l) else { + panic!("expected Add") + }; + assert_eq!(out.get(state), &Node::StateVector { start: 0, end: 1 }); + assert_eq!(out.get(constant), &Node::Scalar(1.0)); + } + + #[test] + fn a_reduce_arg_select_over_a_zero_basis_folds_to_zero() { + for (zero_basis, width) in [(Node::ZeroVector { len: 3 }, 3), (Node::Scalar(0.0), 1)] { + let mut arena = Arena::new(); + let picker = arena.alloc(Node::StateVector { + start: 0, + end: width, + }); + let basis = arena.alloc(zero_basis); + let node = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max: true, + }); + let simplified = simplify(&mut arena, node); + match arena.get(simplified) { + Node::Scalar(v) => assert!(v.abs() < f64::EPSILON), + other => panic!("expected Scalar(0.0), got {other:?}"), + } + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/tangent.rs b/packages/pybamm-rust/pybamm-core/src/tangent.rs new file mode 100644 index 0000000000..6fbaffd337 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/tangent.rs @@ -0,0 +1,1309 @@ +//! Symbolic forward-mode differentiation into Jacobian-vector product expressions. +//! +//! Tangent nodes read their seed at evaluation time, so one compiled tape serves every direction. + +use crate::arena::{Arena, NodeId, NodeMap}; +use crate::ir::infer_sizes; +use crate::node::Node; + +/// Which variables carry a tangent; every other leaf differentiates to zero. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DiffTarget { + /// `d/dy`, seeded from `TangentInputs::dy`. + States, + /// `d/dp`, seeded from `TangentInputs::dp`. + Params, +} + +/// Build the tangent expression for `(d expr/dy) @ dy`. +#[must_use] +pub fn tangent_wrt_states(arena: &mut Arena, expr: NodeId) -> NodeId { + differentiate_root(arena, expr, DiffTarget::States, None) +} + +/// Build the tangent expression for `(d expr/dp) @ dp`. +#[must_use] +pub fn tangent_wrt_params(arena: &mut Arena, expr: NodeId) -> NodeId { + differentiate_root(arena, expr, DiffTarget::Params, None) +} + +/// [`tangent_wrt_states`] seeding only `StateVector` nodes that overlap `active_indices`. +/// +/// `active_indices` must ascend. Overlap is all-or-nothing: a node straddling the +/// boundary is seeded over its whole range. +#[must_use] +pub fn tangent_wrt_subset(arena: &mut Arena, expr: NodeId, active_indices: &[usize]) -> NodeId { + assert!( + active_indices.is_sorted(), + "active state indices must ascend" + ); + differentiate_root(arena, expr, DiffTarget::States, Some(active_indices)) +} + +/// Differentiate `expr`, seeding only states in `state_filter` when one is given. +fn differentiate_root( + arena: &mut Arena, + expr: NodeId, + mode: DiffTarget, + state_filter: Option<&[usize]>, +) -> NodeId { + let mut pending = arena.topological_order(expr); + let widths = infer_sizes(arena, &pending); + // The order's buffer, one slot per reachable node, becomes the work stack. + pending.clear(); + pending.push(expr); + let mut walk = Differentiator { + mode, + state_filter, + widths, + memo: NodeMap::new(arena.len()), + pending, + parked_sums: Vec::new(), + }; + while let Some(&id) = walk.pending.last() { + if walk.memo.contains_key(id) { + walk.pending.pop(); + } else if let Some(derivative) = walk.differentiate(arena, id) { + walk.pending.pop(); + walk.memo.insert(id, derivative); + } + } + *walk + .memo + .get(expr) + .expect("the walk ends with `expr` differentiated") +} + +/// Forward-mode rules, walked with an explicit stack so graph depth is bounded by the heap. +/// +/// A rule missing a child's derivative queues the child and is retried after it, so nodes +/// are allocated in exactly the order a depth-first recursion would allocate them. +struct Differentiator<'a> { + mode: DiffTarget, + state_filter: Option<&'a [usize]>, + /// Primal output widths. + widths: Vec, + /// One derivative per primal node. + memo: NodeMap, + /// Nodes left to differentiate; each waits on those pushed above it. + pending: Vec, + /// `(node, next axis, sum so far)` of each `InterpolantNd` waiting on a child, innermost + /// last: a node's children finish before it, so the one retried is always on top. + parked_sums: Vec<(NodeId, usize, Option)>, +} + +impl Differentiator<'_> { + /// Queue the not-yet-differentiated `children`, first on top; `true` if there were none. + fn queue_missing(&mut self, children: &[NodeId]) -> bool { + let n_pending = self.pending.len(); + for &child in children.iter().rev() { + if !self.memo.contains_key(child) { + self.pending.push(child); + } + } + self.pending.len() == n_pending + } + + fn derivative(&self, id: NodeId) -> NodeId { + *self + .memo + .get(id) + .expect("a rule reads only children it has seen differentiated") + } + + /// The derivatives of `children`, or `None` once the missing ones are queued. + fn ready(&mut self, children: [NodeId; N]) -> Option<[NodeId; N]> { + self.queue_missing(&children) + .then(|| children.map(|child| self.derivative(child))) + } + + /// [`ready`](Self::ready) for a variadic node. + fn ready_all(&mut self, children: &[NodeId]) -> Option> { + self.queue_missing(children).then(|| { + children + .iter() + .map(|&child| self.derivative(child)) + .collect() + }) + } + + /// Where `id`'s axis-by-axis sum stopped, or its start if it never parked. + fn unpark_sum(&mut self, id: NodeId) -> (usize, Option) { + match self.parked_sums.last() { + Some(&(parked, axis, sum)) if parked == id => { + self.parked_sums.pop(); + (axis, sum) + }, + _ => (0, None), + } + } +} + +/// Whether the ascending `indices` hold any state in `start..end`. +fn overlaps(indices: &[usize], start: usize, end: usize) -> bool { + let first = indices.partition_point(|&index| index < start); + indices.get(first).is_some_and(|&index| index < end) +} + +/// Whether `id` is a literal, so its derivative is identically zero. +fn is_literal(arena: &Arena, id: NodeId) -> bool { + matches!( + arena.get(id), + Node::Scalar(_) | Node::Array(_) | Node::ZeroVector { .. } + ) +} + +/// Whether `id` is a zero [`Node::zero`] builds, so a product with it can be skipped. +fn is_literal_zero(arena: &Arena, id: NodeId) -> bool { + match arena.get(id) { + Node::Scalar(v) => v.to_bits() == 0.0_f64.to_bits(), + Node::ZeroVector { .. } => true, + _ => false, + } +} + +/// `d(a^b)/db`: `a^b * ln(a)` where `a > 0` and 0 elsewhere. +/// +/// `ln` reads `max(a, smallest subnormal)`, so the gated-off product is a finite zero, not +/// `0 * -inf`; for `a > 0` the gate and the `max` are exact, so the value is `a^b * ln(a)`. +fn exponent_partial(arena: &mut Arena, base: NodeId, exp: NodeId) -> NodeId { + let zero = arena.alloc(Node::Scalar(0.0)); + let positive = arena.alloc(Node::NotEqualHeaviside(zero, base)); + let a_pow_b = arena.alloc(Node::Pow(base, exp)); + let gated = arena.alloc(Node::Mul(positive, a_pow_b)); + let tiny = arena.alloc(Node::Scalar(f64::from_bits(1))); + let log_argument = arena.alloc(Node::Maximum(base, tiny)); + let log_a = arena.alloc(Node::Log(log_argument)); + arena.alloc(Node::Mul(gated, log_a)) +} + +/// `selector * da + (1 - selector) * db`, the subgradient `Minimum` and `Maximum` share. +fn blend_tangents(arena: &mut Arena, selector: NodeId, da: NodeId, db: NodeId) -> NodeId { + let one = arena.alloc(Node::Scalar(1.0)); + let one_minus_sel = arena.alloc(Node::Sub(one, selector)); + let sel_da = arena.alloc(Node::Mul(selector, da)); + let not_sel_db = arena.alloc(Node::Mul(one_minus_sel, db)); + arena.alloc(Node::Add(sel_da, not_sel_db)) +} + +impl Differentiator<'_> { + /// `id`'s derivative, or `None` once it has queued the work it waits on. + /// + /// Allocates nothing before returning `None`, so a retry starts clean; an `InterpolantNd` + /// parks what it allocated and resumes from there. + fn differentiate(&mut self, arena: &mut Arena, id: NodeId) -> Option { + let (mode, state_filter) = (self.mode, self.state_filter); + let width = self.widths[id.index()]; + + let result = match *arena.get(id) { + Node::StateVector { start, end } + if mode == DiffTarget::States + && state_filter.is_none_or(|filter| overlaps(filter, start, end)) => + { + arena.alloc(Node::TangentStateVector { start, end }) + }, + + Node::InputParameter { + index, + width: param_width, + .. + } if mode == DiffTarget::Params => { + // Replicated to the packed width so an `Index` slice stays inside the buffer slot. + let tangent = arena.alloc(Node::TangentParameter { index }); + if param_width == 1 { + tangent + } else { + arena.alloc(Node::Concat(vec![tangent; param_width])) + } + }, + + // Leaves that carry no tangent in this mode. d(y')/dy is the mass matrix, which the + // solver supplies separately. + Node::Scalar(_) + | Node::Array(_) + | Node::ZeroVector { .. } + | Node::SparseMatrix(_) + | Node::Time + | Node::StateVector { .. } + | Node::StateVectorDot { .. } + | Node::InputParameter { .. } + | Node::TangentStateVector { .. } + | Node::TangentParameter { .. } + // Step functions: zero away from the jumps, which are not modelled. + | Node::Sign(_) + | Node::Floor(_) + | Node::Ceiling(_) + | Node::EqualHeaviside(..) + | Node::NotEqualHeaviside(..) + | Node::Equality(..) + // Derivative-only nodes: the solver Jacobian never needs a second derivative. + | Node::Interpolant1DLinearDeriv { .. } + | Node::Interpolant1DCubicDeriv { .. } + | Node::InterpolantNdPartial { .. } + | Node::ReduceArgSelect { .. } => arena.alloc(Node::zero(width)), + + Node::Add(a, b) => { + let [da, db] = self.ready([a, b])?; + arena.alloc(Node::Add(da, db)) + }, + + Node::Sub(a, b) => { + let [da, db] = self.ready([a, b])?; + arena.alloc(Node::Sub(da, db)) + }, + + Node::Mul(a, b) => { + let [da, db] = self.ready([a, b])?; + let a_db = arena.alloc(Node::Mul(a, db)); + let da_b = arena.alloc(Node::Mul(da, b)); + arena.alloc(Node::Add(a_db, da_b)) + }, + + Node::Div(a, b) => { + let [da, db] = self.ready([a, b])?; + let da_b = arena.alloc(Node::Mul(da, b)); + let a_db = arena.alloc(Node::Mul(a, db)); + let numer = arena.alloc(Node::Sub(da_b, a_db)); + let two = arena.alloc(Node::Scalar(2.0)); + let b_sq = arena.alloc(Node::Pow(b, two)); + arena.alloc(Node::Div(numer, b_sq)) + }, + + // Partials times tangents; a constant exponent skips the `ln` term, which a + // zero base would turn into `0 * -inf` for every seed. + Node::Pow(base, exp) => { + let [da] = self.ready([base])?; + let db = if is_literal(arena, exp) { + None + } else { + Some(self.ready([exp])?[0]) + }; + let exp_minus_1 = if let Some(n) = arena.get(exp).as_scalar() { + arena.alloc(Node::Scalar(n - 1.0)) + } else { + let one = arena.alloc(Node::Scalar(1.0)); + arena.alloc(Node::Sub(exp, one)) + }; + let pow_nm1 = arena.alloc(Node::Pow(base, exp_minus_1)); + let d_base = arena.alloc(Node::Mul(exp, pow_nm1)); + let base_term = arena.alloc(Node::Mul(d_base, da)); + db.filter(|&db| !is_literal_zero(arena, db)) + .map_or(base_term, |db| { + let d_exp = exponent_partial(arena, base, exp); + let exp_term = arena.alloc(Node::Mul(d_exp, db)); + arena.alloc(Node::Add(base_term, exp_term)) + }) + }, + + Node::Neg(a) => { + let [da] = self.ready([a])?; + arena.alloc(Node::Neg(da)) + }, + + Node::Abs(a) => { + let [da] = self.ready([a])?; + let sign_a = arena.alloc(Node::Sign(a)); + arena.alloc(Node::Mul(sign_a, da)) + }, + + Node::Sqrt(a) => { + let [da] = self.ready([a])?; + let two = arena.alloc(Node::Scalar(2.0)); + let sqrt_a = arena.alloc(Node::Sqrt(a)); + let two_sqrt_a = arena.alloc(Node::Mul(two, sqrt_a)); + arena.alloc(Node::Div(da, two_sqrt_a)) + }, + + Node::Exp(a) => { + let [da] = self.ready([a])?; + let exp_a = arena.alloc(Node::Exp(a)); + arena.alloc(Node::Mul(exp_a, da)) + }, + + Node::Log(a) => { + let [da] = self.ready([a])?; + arena.alloc(Node::Div(da, a)) + }, + + Node::Sin(a) => { + let [da] = self.ready([a])?; + let cos_a = arena.alloc(Node::Cos(a)); + arena.alloc(Node::Mul(cos_a, da)) + }, + + Node::Cos(a) => { + let [da] = self.ready([a])?; + let sin_a = arena.alloc(Node::Sin(a)); + let neg_sin_a = arena.alloc(Node::Neg(sin_a)); + arena.alloc(Node::Mul(neg_sin_a, da)) + }, + + Node::Tanh(a) => { + let [da] = self.ready([a])?; + let tanh_a = arena.alloc(Node::Tanh(a)); + let one = arena.alloc(Node::Scalar(1.0)); + let two = arena.alloc(Node::Scalar(2.0)); + let tanh_sq = arena.alloc(Node::Pow(tanh_a, two)); + let one_minus_tanh_sq = arena.alloc(Node::Sub(one, tanh_sq)); + arena.alloc(Node::Mul(one_minus_tanh_sq, da)) + }, + + Node::Sinh(a) => { + let [da] = self.ready([a])?; + let cosh_a = arena.alloc(Node::Cosh(a)); + arena.alloc(Node::Mul(cosh_a, da)) + }, + + Node::Cosh(a) => { + let [da] = self.ready([a])?; + let sinh_a = arena.alloc(Node::Sinh(a)); + arena.alloc(Node::Mul(sinh_a, da)) + }, + + Node::Arcsinh(a) => { + let [da] = self.ready([a])?; + let one = arena.alloc(Node::Scalar(1.0)); + let two = arena.alloc(Node::Scalar(2.0)); + let a_sq = arena.alloc(Node::Pow(a, two)); + let one_plus_a_sq = arena.alloc(Node::Add(one, a_sq)); + let sqrt_denom = arena.alloc(Node::Sqrt(one_plus_a_sq)); + arena.alloc(Node::Div(da, sqrt_denom)) + }, + + Node::Arctan(a) => { + let [da] = self.ready([a])?; + let one = arena.alloc(Node::Scalar(1.0)); + let two = arena.alloc(Node::Scalar(2.0)); + let a_sq = arena.alloc(Node::Pow(a, two)); + let one_plus_a_sq = arena.alloc(Node::Add(one, a_sq)); + arena.alloc(Node::Div(da, one_plus_a_sq)) + }, + + Node::Erf(a) => { + let [da] = self.ready([a])?; + let two_over_sqrt_pi = arena.alloc(Node::Scalar(std::f64::consts::FRAC_2_SQRT_PI)); + let neg_one = arena.alloc(Node::Scalar(-1.0)); + let two = arena.alloc(Node::Scalar(2.0)); + let a_sq = arena.alloc(Node::Pow(a, two)); + let neg_a_sq = arena.alloc(Node::Mul(neg_one, a_sq)); + let exp_neg_a_sq = arena.alloc(Node::Exp(neg_a_sq)); + let coeff = arena.alloc(Node::Mul(two_over_sqrt_pi, exp_neg_a_sq)); + arena.alloc(Node::Mul(coeff, da)) + }, + + Node::Minimum(a, b) => { + // Subgradient: the (a <= b) indicator picks da. + let [da, db] = self.ready([a, b])?; + let selector = arena.alloc(Node::EqualHeaviside(a, b)); + blend_tangents(arena, selector, da, db) + }, + + Node::Maximum(a, b) => { + // EqualHeaviside(b, a) is (a >= b). + let [da, db] = self.ready([a, b])?; + let selector = arena.alloc(Node::EqualHeaviside(b, a)); + blend_tangents(arena, selector, da, db) + }, + + Node::Modulo(a, b) => { + // Ignores the discontinuities at integer multiples of b. + let [da] = self.ready([a])?; + let db = if is_literal(arena, b) { + None + } else { + Some(self.ready([b])?[0]) + }; + // A scalar `a` broadcasts against a vector `b`; replicate rather than add a zero, + // which would flip a -0.0 tangent. + let da = if self.widths[a.index()] == 1 && width > 1 { + arena.alloc(Node::Concat(vec![da; width])) + } else { + da + }; + db.map_or(da, |db| { + // d/db is -trunc(a / b); reading it off the remainder rounds as `%` does. + let shift = arena.alloc(Node::Sub(id, a)); + let neg_quotient = arena.alloc(Node::Div(shift, b)); + let db_term = arena.alloc(Node::Mul(neg_quotient, db)); + arena.alloc(Node::Add(da, db_term)) + }) + }, + + Node::Hypot(a, b) => { + let [da, db] = self.ready([a, b])?; + let a_da = arena.alloc(Node::Mul(a, da)); + let b_db = arena.alloc(Node::Mul(b, db)); + let numer = arena.alloc(Node::Add(a_da, b_db)); + let hypot_ab = arena.alloc(Node::Hypot(a, b)); + arena.alloc(Node::Div(numer, hypot_ab)) + }, + + Node::MatMul(a, b) => { + // The matrix is always constant, so it contributes no term of its own. + let [db] = self.ready([b])?; + arena.alloc(Node::MatMul(a, db)) + }, + + Node::Index { child, start, end } => { + let [d_child] = self.ready([child])?; + arena.alloc(Node::Index { + child: d_child, + start, + end, + }) + }, + + Node::Concat(ref children) => { + let d_children = self.ready_all(children)?; + arena.alloc(Node::Concat(d_children)) + }, + + Node::Interpolant1DLinear { ref data, child } => { + let [d_child] = self.ready([child])?; + let data = data.clone(); + + let deriv_interp = arena.alloc(Node::Interpolant1DLinearDeriv { data, child }); + + arena.alloc(Node::Mul(deriv_interp, d_child)) + }, + + Node::Interpolant1DCubic { ref data, child } => { + let [d_child] = self.ready([child])?; + let data = data.clone(); + let deriv_interp = arena.alloc(Node::Interpolant1DCubicDeriv { data, child }); + arena.alloc(Node::Mul(deriv_interp, d_child)) + }, + + // Each axis's term is allocated before the next child is differentiated, so a + // missing child parks the sum so far rather than discarding it. + Node::InterpolantNd { .. } => { + let (mut axis, mut sum) = self.unpark_sum(id); + loop { + let Node::InterpolantNd { data, children } = arena.get(id) else { + unreachable!("matched as an `InterpolantNd` above") + }; + let Some(&child) = children.get(axis) else { + break sum.expect("InterpolantNd has at least one child"); + }; + let Some([d_child]) = self.ready([child]) else { + self.parked_sums.push((id, axis, sum)); + return None; + }; + let partial = Node::InterpolantNdPartial { + data: data.clone(), + children: children.clone(), + axis: u32::try_from(axis).expect("axis index fits in u32"), + }; + let partial = arena.alloc(partial); + let term = arena.alloc(Node::Mul(partial, d_child)); + sum = Some(sum.map_or(term, |acc| arena.alloc(Node::Add(acc, term)))); + axis += 1; + } + }, + + Node::Conditional { + selector, + ref branches, + } => { + let d_branches = self.ready_all(branches)?; + arena.alloc(Node::Conditional { + selector, + branches: d_branches, + }) + }, + + // Ties go to the first occurrence, matching the primal. + Node::MaxReduce(a) => { + let [da] = self.ready([a])?; + arena.alloc(Node::ReduceArgSelect { + basis: da, + picker: a, + is_max: true, + }) + }, + Node::MinReduce(a) => { + let [da] = self.ready([a])?; + arena.alloc(Node::ReduceArgSelect { + basis: da, + picker: a, + is_max: false, + }) + }, + }; + + Some(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arena::DagRef; + use crate::eval::{CompiledExpr, TangentInputs}; + use crate::node::InterpolantData; + use crate::simplify::{cse, dce, simplify}; + use crate::zero_propagate::zero_propagate; + + /// Run the tangent pipeline end to end and return the compiled tape. + fn compile_tangent_pipeline(arena: &Arena, root: NodeId, wrt_params: bool) -> CompiledExpr { + let mut a = arena.clone(); + let root = if wrt_params { + tangent_wrt_params(&mut a, root) + } else { + tangent_wrt_states(&mut a, root) + }; + let root = simplify(&mut a, root); + let propagated = zero_propagate(DagRef::new(&a, root)); + let shared = cse(propagated.as_ref()); + let pruned = dce(shared.as_ref()); + CompiledExpr::new(pruned.as_ref()) + } + + /// `(d root/dy) @ dy` through the simplified tangent tape. + fn jvp(arena: &Arena, root: NodeId, y: &[f64], inputs: &[f64], dy: &[f64]) -> Vec { + let mut arena = arena.clone(); + let tangent = tangent_wrt_states(&mut arena, root); + let tangent = simplify(&mut arena, tangent); + let compiled = CompiledExpr::new(DagRef::new(&arena, tangent)); + let mut scratch = vec![0.0; compiled.scratch_len()]; + let seed = TangentInputs { + dy: Some(dy), + dp: None, + }; + compiled + .eval_with_tangent(&mut scratch, 0.0, y, &[], inputs, &seed) + .to_vec() + } + + /// The gradient of a scalar `root`, one forward sweep per state. + fn forward_gradient(arena: &Arena, root: NodeId, y: &[f64], inputs: &[f64]) -> Vec { + (0..y.len()) + .map(|state| { + let mut dy = vec![0.0; y.len()]; + dy[state] = 1.0; + jvp(arena, root, y, inputs, &dy)[0] + }) + .collect() + } + + #[test] + fn power_rule_with_an_input_exponent_holds_at_a_zero_or_negative_base() { + // d(y0^p)/dy0 = p * y0^(p - 1); a 0 * (1 / 0) form made it NaN at the origin. + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let p = arena.alloc(Node::InputParameter { + name: "p".into(), + index: 0, + offset: 0, + width: 1, + }); + let root = arena.alloc(Node::Pow(y0, p)); + for y in [0.0, -1.0, 2.0] { + let closed_form = [2.0 * y]; + assert_eq!(forward_gradient(&arena, root, &[y], &[2.0]), closed_form); + } + } + + #[test] + fn power_rule_with_a_state_exponent_is_finite_at_a_zero_base() { + // At (0, 2): d/dy0 = 2 * 0 and d/dy1 is 0 by the `a > 0` convention. + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let root = arena.alloc(Node::Pow(y0, y1)); + let y = [0.0, 2.0]; + assert_eq!(forward_gradient(&arena, root, &y, &[]), [0.0, 0.0]); + } + + #[test] + fn test_jvp_width_vector_plus_scalar_param() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let p = arena.alloc(Node::InputParameter { + name: "a".into(), + index: 0, + offset: 0, + width: 1, + }); + let root = arena.alloc(Node::Add(y, p)); + + let tp = compile_tangent_pipeline(&arena, root, true); + assert_eq!( + tp.output_len(), + 3, + "tangent_p tape collapsed below primal width" + ); + let mut s = vec![0.0; tp.scratch_len()]; + let tangent = TangentInputs { + dy: None, + dp: Some(&[1.0]), + }; + let r = tp.eval_with_tangent(&mut s, 0.0, &[0.0, 0.0, 0.0], &[], &[5.0], &tangent); + assert_eq!(r, &[1.0, 1.0, 1.0]); + } + + #[test] + fn test_jvp_width_wide_param_index_slice() { + // A width-1 tangent under `Index{1..2}` would read outside its buffer slot. + let mut arena = Arena::new(); + let b = arena.alloc(Node::InputParameter { + name: "b".into(), + index: 0, + offset: 0, + width: 2, + }); + let b1 = arena.alloc(Node::Index { + child: b, + start: 1, + end: 2, + }); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Mul(b1, y0)); + + let tp = compile_tangent_pipeline(&arena, root, true); + let mut s = vec![0.0; tp.scratch_len()]; + let seeded = TangentInputs { + dy: None, + dp: Some(&[1.0]), + }; + let r = tp.eval_with_tangent(&mut s, 0.0, &[3.0], &[], &[10.0, 20.0], &seeded); + assert_eq!(r, &[3.0]); + + let mut s2 = vec![0.0; tp.scratch_len()]; + let unseeded = TangentInputs { + dy: None, + dp: Some(&[0.0]), + }; + let r0 = tp.eval_with_tangent(&mut s2, 0.0, &[3.0], &[], &[10.0, 20.0], &unseeded); + assert_eq!(r0, &[0.0]); + } + + #[test] + fn test_jvp_width_zero_dfdy_keeps_length() { + use crate::node::ArrayData; + let mut arena = Arena::new(); + let cv = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0, 3.0], + shape: crate::node::Shape::vector(3), + }))); + let p = arena.alloc(Node::InputParameter { + name: "a".into(), + index: 0, + offset: 0, + width: 1, + }); + let root = arena.alloc(Node::Add(cv, p)); + + let ty = compile_tangent_pipeline(&arena, root, false); + assert_eq!( + ty.output_len(), + 3, + "tangent_y tape collapsed below primal width" + ); + let mut s = vec![0.0; ty.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[0.0]), + dp: None, + }; + let r = ty.eval_with_tangent(&mut s, 0.0, &[], &[], &[5.0], &tangent); + assert_eq!(r, &[0.0, 0.0, 0.0]); + } + + #[test] + fn test_jvp_width_concat_collapsed_child_offset() { + use crate::node::ArrayData; + let mut arena = Arena::new(); + let cv = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![1.0, 2.0], + shape: crate::node::Shape::vector(2), + }))); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Concat(vec![cv, y0])); + + let ty = compile_tangent_pipeline(&arena, root, false); + assert_eq!(ty.output_len(), 3, "concat tangent width must match primal"); + let mut s = vec![0.0; ty.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let r = ty.eval_with_tangent(&mut s, 0.0, &[7.0], &[], &[], &tangent); + assert_eq!(r, &[0.0, 0.0, 1.0]); + } + + #[test] + fn test_jvp_width_concat_empty_child_offset() { + use crate::node::ArrayData; + let mut arena = Arena::new(); + let empty = arena.alloc(Node::Array(Box::new(ArrayData { + data: vec![], + shape: crate::node::Shape::vector(0), + }))); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let one = arena.alloc(Node::Scalar(1.0)); + let expr = arena.alloc(Node::Add(y0, one)); + let root = arena.alloc(Node::Concat(vec![empty, expr])); + + let ty = compile_tangent_pipeline(&arena, root, false); + assert_eq!( + ty.output_len(), + 1, + "concat tangent width must match primal (empty child contributes 0)" + ); + let mut s = vec![0.0; ty.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let r = ty.eval_with_tangent(&mut s, 0.0, &[7.0], &[], &[], &tangent); + assert_eq!( + r, + &[1.0], + "d(y0 + 1)/dy0 must be 1, not swallowed by the empty child" + ); + } + + /// Deeper than a recursive walk survives on a 2 MiB test-thread stack. + const DEEP: usize = 200_000; + + #[test] + fn differentiates_a_deep_chain() { + use crate::node::NdInterpolantData; + // Cycles a one-child rule, a two-child rule and the axis-by-axis N-D rule. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let table = NdInterpolantData::try_new(vec![vec![0.0, 1.0]; 2], vec![1.0; 4], 2) + .expect("valid table"); + let mut id = y; + for level in 0..DEEP { + id = arena.alloc(match level % 3 { + 0 => Node::Sin(id), + 1 => Node::Mul(id, y), + _ => Node::InterpolantNd { + data: Box::new(table.clone()), + children: vec![id, y], + }, + }); + } + + let mut tangent = tangent_wrt_states(&mut arena, id); + for level in (0..DEEP).rev() { + tangent = match (level % 3, arena.get(tangent)) { + (0, &Node::Mul(_, inner)) => inner, + (1, &Node::Add(_, product)) => match *arena.get(product) { + Node::Mul(inner, _) => inner, + ref other => panic!("level {level}: expected da * b, got {other:?}"), + }, + (2, &Node::Add(first_axis, _)) => match *arena.get(first_axis) { + Node::Mul(_, inner) => inner, + ref other => panic!("level {level}: expected a partial term, got {other:?}"), + }, + (_, other) => panic!("level {level}: unexpected {other:?}"), + }; + } + assert!(matches!( + arena.get(tangent), + Node::TangentStateVector { start: 0, end: 1 } + )); + } + + #[test] + fn single_state_rules_match_their_closed_forms() { + use std::f64::consts::{FRAC_2_SQRT_PI, FRAC_PI_4}; + + use crate::node::CubicInterpolantData; + + type Build = fn(&mut Arena, NodeId) -> NodeId; + fn pow(arena: &mut Arena, x: NodeId, exponent: f64) -> NodeId { + let exponent = arena.alloc(Node::Scalar(exponent)); + arena.alloc(Node::Pow(x, exponent)) + } + fn linear(arena: &mut Arena, x: NodeId) -> NodeId { + let data = InterpolantData::try_new(vec![0.0, 1.0, 2.0], vec![0.0, 10.0, 30.0]); + arena.alloc(Node::Interpolant1DLinear { + data: Box::new(data.expect("valid interpolant")), + child: x, + }) + } + // p(dx) = 1 + 2*dx + 3*dx^2 + 4*dx^3 on [0, 5], extrapolated outside it. + fn cubic(arena: &mut Arena, x: NodeId) -> NodeId { + arena.alloc(Node::Interpolant1DCubic { + data: Box::new(CubicInterpolantData { + breakpoints: vec![0.0, 5.0], + coeffs: vec![[1.0, 2.0, 3.0, 4.0]], + }), + child: x, + }) + } + let cases: [(&str, Build, f64, f64); 30] = [ + ( + "sin", + |a, x| a.alloc(Node::Sin(x)), + FRAC_PI_4, + FRAC_PI_4.cos(), + ), + ( + "cos", + |a, x| a.alloc(Node::Cos(x)), + FRAC_PI_4, + -FRAC_PI_4.sin(), + ), + ( + "tanh", + |a, x| a.alloc(Node::Tanh(x)), + 1.0, + 1.0 - 1.0_f64.tanh().powi(2), + ), + ("sinh", |a, x| a.alloc(Node::Sinh(x)), 1.5, 1.5_f64.cosh()), + ("cosh", |a, x| a.alloc(Node::Cosh(x)), 1.5, 1.5_f64.sinh()), + ( + "arcsinh", + |a, x| a.alloc(Node::Arcsinh(x)), + 2.0, + 1.0 / 5.0_f64.sqrt(), + ), + ("arctan", |a, x| a.alloc(Node::Arctan(x)), 1.0, 0.5), + ( + "erf", + |a, x| a.alloc(Node::Erf(x)), + 1.0, + FRAC_2_SQRT_PI * (-1.0_f64).exp(), + ), + ("neg", |a, x| a.alloc(Node::Neg(x)), 5.0, -1.0), + ("abs", |a, x| a.alloc(Node::Abs(x)), 5.0, 1.0), + ("abs", |a, x| a.alloc(Node::Abs(x)), -5.0, -1.0), + ("sqrt", |a, x| a.alloc(Node::Sqrt(x)), 4.0, 0.25), + ("log", |a, x| a.alloc(Node::Log(x)), 2.0, 0.5), + ( + "exp chain", + |a, x| { + let two = a.alloc(Node::Scalar(2.0)); + let two_x = a.alloc(Node::Mul(two, x)); + a.alloc(Node::Exp(two_x)) + }, + 1.0, + 2.0 * 2.0_f64.exp(), + ), + ("x^3", |a, x| pow(a, x, 3.0), 2.0, 12.0), + ("x^2, negative base", |a, x| pow(a, x, 2.0), -3.0, -6.0), + ("x^1.5", |a, x| pow(a, x, 1.5), 4.0, 3.0), + ("x^0.5", |a, x| pow(a, x, 0.5), 9.0, 0.5 / 3.0), + ("x * x, memoised", |a, x| a.alloc(Node::Mul(x, x)), 3.0, 6.0), + ("linear interpolant", linear, 0.5, 10.0), + ("linear interpolant", linear, 1.5, 20.0), + ("linear interpolant, below", linear, -0.5, 10.0), + ("linear interpolant, above", linear, 2.5, 20.0), + ("cubic interpolant", cubic, 1.0, 20.0), + ("cubic interpolant", cubic, 0.0, 2.0), + ("cubic interpolant, below", cubic, -1.0, 8.0), + ("cubic interpolant, above", cubic, 6.0, 470.0), + ( + "conditional", + |a, x| { + let selector = a.alloc(Node::Scalar(1.0)); + let two = a.alloc(Node::Scalar(2.0)); + let three = a.alloc(Node::Scalar(3.0)); + let first = a.alloc(Node::Mul(two, x)); + let second = a.alloc(Node::Mul(three, x)); + a.alloc(Node::Conditional { + selector, + branches: vec![first, second], + }) + }, + 5.0, + 2.0, + ), + ("constant", |a, _| a.alloc(Node::Scalar(42.0)), 1.0, 0.0), + ("time", |a, _| a.alloc(Node::Time), 1.0, 0.0), + ]; + for (label, build, x, expected) in cases { + let mut arena = Arena::new(); + let state = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = build(&mut arena, state); + let got = jvp(&arena, root, &[x], &[], &[1.0])[0]; + assert!( + (got - expected).abs() < 1e-12, + "{label} at {x}: expected {expected}, got {got}" + ); + } + } + + #[test] + fn two_state_rules_match_their_closed_forms() { + use crate::node::NdInterpolantData; + + type Build = fn(&mut Arena, NodeId, NodeId) -> NodeId; + fn interpolant(arena: &mut Arena, x: NodeId, y: NodeId, data: NdInterpolantData) -> NodeId { + arena.alloc(Node::InterpolantNd { + data: Box::new(data), + children: vec![x, y], + }) + } + let cases: [(&str, Build, [f64; 2], [f64; 2]); 10] = [ + ( + "product", + |a, x, y| a.alloc(Node::Mul(x, y)), + [2.0, 4.0], + [4.0, 2.0], + ), + ( + "quotient", + |a, x, y| a.alloc(Node::Div(x, y)), + [1.0, 2.0], + [0.5, -0.25], + ), + ( + "minimum", + |a, x, y| a.alloc(Node::Minimum(x, y)), + [1.0, 3.0], + [1.0, 0.0], + ), + ( + "minimum", + |a, x, y| a.alloc(Node::Minimum(x, y)), + [5.0, 2.0], + [0.0, 1.0], + ), + ( + "maximum", + |a, x, y| a.alloc(Node::Maximum(x, y)), + [5.0, 2.0], + [1.0, 0.0], + ), + ( + "maximum", + |a, x, y| a.alloc(Node::Maximum(x, y)), + [1.0, 3.0], + [0.0, 1.0], + ), + ( + "hypot", + |a, x, y| a.alloc(Node::Hypot(x, y)), + [3.0, 4.0], + [0.6, 0.8], + ), + ( + "power", + |a, x, y| a.alloc(Node::Pow(x, y)), + [2.0, 3.0], + [12.0, 8.0 * std::f64::consts::LN_2], + ), + // p = 3*dx0^2 + 5*dx1^3 + dx0*dx1 on one cell [0,4]x[0,4]. + ( + "N-D interpolant", + |a, x, y| { + let mut coeffs = vec![0.0; 16]; + coeffs[2 * 4] = 3.0; + coeffs[3] = 5.0; + coeffs[4 + 1] = 1.0; + let breakpoints = vec![vec![0.0, 4.0], vec![0.0, 4.0]]; + let data = NdInterpolantData::try_new(breakpoints, coeffs, 4); + interpolant(a, x, y, data.expect("valid table")) + }, + [2.0, 1.0], + [13.0, 17.0], + ), + // p = 2 + 3*dx0 + 4*dx1 + 5*dx0*dx1 on one cell [0,2]x[0,2]. + ( + "bilinear interpolant", + |a, x, y| { + let breakpoints = vec![vec![0.0, 2.0], vec![0.0, 2.0]]; + let data = NdInterpolantData::try_new(breakpoints, vec![2.0, 4.0, 3.0, 5.0], 2); + interpolant(a, x, y, data.expect("valid table")) + }, + [1.0, 2.0], + [13.0, 9.0], + ), + ]; + for (label, build, y, expected) in cases { + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let root = build(&mut arena, x0, x1); + let got = forward_gradient(&arena, root, &y, &[]); + for (got, expected) in got.iter().zip(expected) { + assert!( + (got - expected).abs() < 1e-10, + "{label} at {y:?}: expected {expected}, got {got}" + ); + } + } + } + + #[test] + fn vector_rules_route_each_element() { + use crate::node::{CsrData, Shape}; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let max = arena.alloc(Node::MaxReduce(y)); + let min = arena.alloc(Node::MinReduce(y)); + let (at, dy) = ([0.3, 0.9, 0.1], [10.0, 20.0, 30.0]); + assert_eq!(jvp(&arena, max, &at, &[], &dy), [20.0], "argmax's tangent"); + assert_eq!(jvp(&arena, min, &at, &[], &dy), [30.0], "argmin's tangent"); + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let slice = arena.alloc(Node::Index { + child: y, + start: 1, + end: 3, + }); + let dy = [0.0, 1.0, 2.0, 0.0, 0.0]; + assert_eq!(jvp(&arena, slice, &[0.0; 5], &[], &dy), [1.0, 2.0]); + + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let y = arena.alloc(Node::StateVector { start: 2, end: 4 }); + let joined = arena.alloc(Node::Concat(vec![x, y])); + let dy = [1.0, 2.0, 3.0, 4.0]; + assert_eq!(jvp(&arena, joined, &[0.0; 4], &[], &dy), dy); + + let mut arena = Arena::new(); + let sparse = arena.alloc(Node::SparseMatrix(Box::new(CsrData { + indptr: vec![0, 1, 2], + indices: vec![0, 1], + data: vec![2.0, 3.0], + shape: Shape::matrix(2, 3), + }))); + let v = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let product = arena.alloc(Node::MatMul(sparse, v)); + let dy = [1.0, 0.0, 0.0]; + assert_eq!(jvp(&arena, product, &[0.0; 3], &[], &dy), [2.0, 0.0]); + } + + #[test] + fn test_parameter_differentiation() { + let mut arena = Arena::new(); + let p = arena.alloc(Node::InputParameter { + name: "k".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let expr = arena.alloc(Node::Mul(p, x)); + let compiled = compile_tangent_pipeline(&arena, expr, true); + let mut scratch = vec![0.0; compiled.scratch_len()]; + let seed = TangentInputs { + dy: None, + dp: Some(&[1.0]), + }; + let result = compiled.eval_with_tangent(&mut scratch, 0.0, &[3.0], &[], &[2.0], &seed); + assert_eq!(result, [3.0]); + } + + #[test] + #[allow(clippy::float_cmp)] + fn a_zero_width_input_has_a_zero_width_tangent() { + // A one-row tangent made `d[p, y0]/dp` two rows wide, its top row a spurious 1. + let mut arena = Arena::new(); + let p = arena.alloc(Node::InputParameter { + name: "p".to_string(), + index: 0, + offset: 0, + width: 0, + }); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Concat(vec![p, y0])); + let compiled = compile_tangent_pipeline(&arena, root, true); + let mut scratch = vec![0.0; compiled.scratch_len()]; + let seed = TangentInputs { + dy: None, + dp: Some(&[1.0]), + }; + let result = compiled.eval_with_tangent(&mut scratch, 0.0, &[3.0], &[], &[], &seed); + assert_eq!(result, [0.0]); + } + + #[test] + #[allow(clippy::float_cmp)] + fn modulo_tangent_keeps_a_scalar_dividends_broadcast_width() { + use crate::node::{ArrayData, CsrData, Shape}; + // y0 mod [2, 3, 5] is three wide; summing it through a matrix needs all three rows. + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let divisors = arena.alloc(Node::Array(Box::new( + ArrayData::try_new(vec![2.0, 3.0, 5.0], Shape::vector(3)).expect("valid array"), + ))); + let remainders = arena.alloc(Node::Modulo(y0, divisors)); + let ones = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new(vec![0, 3], vec![0, 1, 2], vec![1.0; 3], Shape::matrix(1, 3)) + .expect("valid row matrix"), + ))); + let root = arena.alloc(Node::MatMul(ones, remainders)); + + let jac = tangent_wrt_states(&mut arena, root); + let jac = simplify(&mut arena, jac); + let compiled = CompiledExpr::new(DagRef::new(&arena, jac)); + let mut scratch = vec![0.0; compiled.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let result = compiled.eval_with_tangent(&mut scratch, 0.0, &[7.5], &[], &[], &tangent); + assert_eq!(result, &[3.0]); + } + + #[test] + fn modulo_tangent_differentiates_an_input_divisor() { + // d(y0 mod p)/dp is -trunc(y0 / p), a positive step for a negative dividend. + let mut arena = Arena::new(); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let p = arena.alloc(Node::InputParameter { + name: "p".into(), + index: 0, + offset: 0, + width: 1, + }); + let root = arena.alloc(Node::Modulo(y0, p)); + + let tp = compile_tangent_pipeline(&arena, root, true); + let tangent = TangentInputs { + dy: None, + dp: Some(&[1.0]), + }; + for (y, expected) in [(2.0, -1.0), (3.0, -2.0), (0.5, 0.0), (-2.0, 1.0)] { + let mut s = vec![0.0; tp.scratch_len()]; + let r = tp.eval_with_tangent(&mut s, 0.0, &[y], &[], &[1.3], &tangent); + assert!( + (r[0] - expected).abs() < 1e-12, + "y0 = {y}: {} != {expected}", + r[0] + ); + } + } + + #[test] + fn test_tangent_wrt_subset() { + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let expr = arena.alloc(Node::Mul(x0, x1)); + + let deriv = tangent_wrt_subset(&mut arena, expr, &[1]); + let deriv = simplify(&mut arena, deriv); + + let topo = arena.topological_order(deriv); + for &nid in &topo { + if let Node::TangentStateVector { start, .. } = *arena.get(nid) { + assert_ne!(start, 0, "Should not differentiate w.r.t. x0"); + } + } + assert!( + topo.iter().any(|&nid| matches!( + arena.get(nid), + Node::TangentStateVector { start: 1, end: 2 } + )), + "Should contain TangentStateVector for x1" + ); + } + + #[test] + fn overlap_needs_one_active_index_inside_the_range() { + let active = [2, 5, 9]; + assert!(overlaps(&active, 0, 3), "2 lies in 0..3"); + assert!(overlaps(&active, 5, 6)); + assert!(!overlaps(&active, 3, 5), "end is exclusive"); + assert!(!overlaps(&active, 10, 12)); + assert!(!overlaps(&active, 5, 5), "an empty range holds nothing"); + assert!(!overlaps(&[], 0, 4)); + } + + #[test] + #[should_panic(expected = "active state indices must ascend")] + fn tangent_wrt_subset_rejects_unsorted_indices() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let _ = tangent_wrt_subset(&mut arena, x, &[1, 0]); + } + + #[test] + fn test_tangent_wrt_subset_numerical() { + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let x0_sq = arena.alloc(Node::Mul(x0, x0)); + let three = arena.alloc(Node::Scalar(3.0)); + let three_x1 = arena.alloc(Node::Mul(three, x1)); + let expr = arena.alloc(Node::Add(x0_sq, three_x1)); + + let deriv = tangent_wrt_subset(&mut arena, expr, &[1]); + let deriv = simplify(&mut arena, deriv); + let compiled = CompiledExpr::new(DagRef::new(&arena, deriv)); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + + let tangent = TangentInputs { + dy: Some(&[0.0, 1.0]), + dp: None, + }; + let result = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[5.0, 2.0], &[], &[], &tangent); + assert!( + (result[0] - 3.0).abs() < 1e-12, + "Expected 3.0, got {}", + result[0] + ); + + let tangent_x0 = TangentInputs { + dy: Some(&[1.0, 0.0]), + dp: None, + }; + let result_x0 = + compiled.eval_with_tangent(&mut s_compiled, 0.0, &[5.0, 2.0], &[], &[], &tangent_x0); + assert!( + result_x0[0].abs() < 1e-12, + "Expected 0.0, got {}", + result_x0[0] + ); + } + + #[test] + fn test_tangent_wrt_subset_all_active_matches_full() { + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let x1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let sum = arena.alloc(Node::Add(x0, x1)); + let expr = arena.alloc(Node::Mul(sum, x0)); + + let mut arena_full = arena.clone(); + let full = tangent_wrt_states(&mut arena_full, expr); + let full = simplify(&mut arena_full, full); + let compiled_full = CompiledExpr::new(DagRef::new(&arena_full, full)); + let mut s_compiled_full = vec![0.0; compiled_full.scratch_len()]; + + let subset = tangent_wrt_subset(&mut arena, expr, &[0, 1]); + let subset = simplify(&mut arena, subset); + let compiled_subset = CompiledExpr::new(DagRef::new(&arena, subset)); + let mut s_compiled_subset = vec![0.0; compiled_subset.scratch_len()]; + + for dy in &[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]] { + let tangent = TangentInputs { + dy: Some(dy), + dp: None, + }; + let r_full = compiled_full.eval_with_tangent( + &mut s_compiled_full, + 0.0, + &[3.0, 7.0], + &[], + &[], + &tangent, + ); + let r_sub = compiled_subset.eval_with_tangent( + &mut s_compiled_subset, + 0.0, + &[3.0, 7.0], + &[], + &[], + &tangent, + ); + assert!( + (r_full[0] - r_sub[0]).abs() < 1e-12, + "Mismatch for dy={dy:?}: full={}, subset={}", + r_full[0], + r_sub[0] + ); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/tangent_batch.rs b/packages/pybamm-rust/pybamm-core/src/tangent_batch.rs new file mode 100644 index 0000000000..ac888b59ea --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/tangent_batch.rs @@ -0,0 +1,386 @@ +//! Batched tangent sweeps: one walk of the tangent tape for `K` colour seeds. +//! +//! Primal slots are shared across lanes and tangent slots are lane-minor; each lane +//! accumulates in the scalar sweep's order, so results are bitwise identical. + +use crate::accumulate::{Accumulate, Product, dispatch}; +use crate::eval::{ + SplitCompiledExpr, csr_rows, split_read_write, split_two_read_write, unary_apply, + with_binary_op, with_unary_op, +}; +use crate::eval_batch::broadcast_lanes; +use crate::ir::ids::BufSlot; +use crate::ir::{BinaryOp, BroadcastKind, ConstPool, Instruction, TypedIr, UnaryOp}; +use crate::node::CsrData; + +#[derive(Clone, Copy)] +enum Operand { + Primal(usize), + /// Lane-minor tangent index, already scaled by `K`. + Tangent(usize), +} + +#[inline] +const fn operand(offset: usize, primal_len: usize) -> Operand { + if offset < primal_len { + Operand::Primal(offset) + } else { + Operand::Tangent((offset - primal_len) * K) + } +} + +/// Whether every instruction in the tangent section has a batched form. +/// +/// Only a `Concat` may read nothing but primal slots; any other reader of only primal +/// slots would be seed-independent, so a tape holding one sweeps scalar instead. +pub fn is_batchable(expr: &SplitCompiledExpr) -> bool { + let split = expr.split_info(); + let tangent = |slot: BufSlot| slot.index() >= split.primal_buffer_size; + expr.ir().instructions()[split.primal_end..] + .iter() + .all(|instr| match *instr { + Instruction::LoadTangentState { .. } | Instruction::Concat { .. } => true, + Instruction::Unary { src, .. } | Instruction::Index { src, .. } => tangent(src), + Instruction::MatMul { vec_src, .. } => tangent(vec_src), + Instruction::Binary { a, b, .. } => tangent(a) || tangent(b), + _ => false, + }) +} + +/// Whether the root is primal, i.e. the tangent folded to a seed-independent constant. +const fn root_is_primal(ir: &TypedIr, primal_len: usize) -> bool { + ir.root_slot().offset_usize() < primal_len +} + +/// Scratch length the tangent region needs for `lanes` lanes. +/// +/// A primal root gets an extra tail that the sweep broadcasts it into. +pub const fn tangent_scratch_len(expr: &SplitCompiledExpr, lanes: usize) -> usize { + let (ir, primal_len) = (expr.ir(), expr.split_info().primal_buffer_size); + let spill = if root_is_primal(ir, primal_len) { + ir.root_slot().len_usize() * lanes + } else { + 0 + }; + (ir.buffer_size() - primal_len) * lanes + spill +} + +/// Run the tangent section of `expr` for `lanes` seed vectors at once. +/// +/// `primal` holds the primal section's values; `seeds` is `(n_states, lanes)` lane-minor and +/// the result `(out_len, lanes)` lane-minor. `lanes` must be 8 or 4, and [`is_batchable`] must +/// hold. +pub fn run_tangent_batch<'t>( + expr: &SplitCompiledExpr, + lanes: usize, + primal: &[f64], + tan: &'t mut [f64], + seeds: &[f64], +) -> &'t [f64] { + match lanes { + 8 => sweep::<8>(expr, primal, tan, seeds), + 4 => sweep::<4>(expr, primal, tan, seeds), + other => panic!("no batched tangent sweep has {other} lanes"), + } +} + +fn sweep<'t, const K: usize>( + expr: &SplitCompiledExpr, + primal: &[f64], + tan: &'t mut [f64], + seeds: &[f64], +) -> &'t [f64] { + let (ir, split) = (expr.ir(), expr.split_info()); + let primal_len = split.primal_buffer_size; + let consts = ir.consts(); + + for instr in &ir.instructions()[split.primal_end..] { + if !writes_nothing(instr, consts) { + exec::(*instr, primal, tan, seeds, primal_len, consts); + } + } + + let root = ir.root_slot(); + if root_is_primal(ir, primal_len) { + let spill = (ir.buffer_size() - primal_len) * K; + broadcast_lanes(&mut tan[spill..], K, &primal[root.bounds()]); + return &tan[spill..spill + root.len_usize() * K]; + } + let base = (root.offset_usize() - primal_len) * K; + &tan[base..base + root.len_usize() * K] +} + +/// Whether `instr` writes no element: its zero-width slot sits at offset 0, in the primal region. +fn writes_nothing(instr: &Instruction, consts: &ConstPool) -> bool { + match *instr { + Instruction::LoadTangentState { width: len, .. } + | Instruction::Index { len, .. } + | Instruction::Unary { len, .. } + | Instruction::Binary { len, .. } => len.index() == 0, + Instruction::Concat { + sources_idx, + sources_len, + .. + } => consts.concat_sources[sources_idx.index()..][..sources_len.index()] + .iter() + .all(|slot| slot.len_usize() == 0), + Instruction::MatMul { csr_idx, .. } => consts.csr_data[csr_idx.index()].shape.rows == 0, + _ => false, + } +} + +fn exec( + instr: Instruction, + primal: &[f64], + tan: &mut [f64], + seeds: &[f64], + primal_len: usize, + consts: &ConstPool, +) { + match instr { + Instruction::LoadTangentState { start, width, dst } => { + let dst = tangent_base::(dst.index(), primal_len); + let len = width.index() * K; + tan[dst..dst + len].copy_from_slice(&seeds[start.index() * K..][..len]); + }, + + Instruction::Index { + src, + start, + dst, + len, + } => { + let src = tangent_base::(src.offset_by(start), primal_len); + let dst = tangent_base::(dst.index(), primal_len); + tan.copy_within(src..src + len.index() * K, dst); + }, + + Instruction::Concat { + sources_idx, + sources_len, + dst, + } => { + let mut write = tangent_base::(dst.index(), primal_len); + for i in 0..sources_len.index() { + let slot = consts.concat_sources[sources_idx.index() + i]; + let len = slot.len_usize(); + match operand::(slot.offset_usize(), primal_len) { + Operand::Tangent(src) => tan.copy_within(src..src + len * K, write), + Operand::Primal(src) => { + broadcast_lanes(&mut tan[write..], K, &primal[src..src + len]); + }, + } + write += len * K; + } + }, + + Instruction::Unary { op, src, dst, len } => unary::( + op, + tangent_base::(src.index(), primal_len), + tangent_base::(dst.index(), primal_len), + len.index(), + tan, + ), + + Instruction::Binary { + op, + a, + b, + dst, + len, + kind, + } => binary_op::( + op, + operand::(a.index(), primal_len), + operand::(b.index(), primal_len), + tangent_base::(dst.index(), primal_len), + len.index(), + kind, + primal, + tan, + ), + + Instruction::MatMul { + csr_idx, + vec_src, + dst, + } => { + let csr = &consts.csr_data[csr_idx.index()]; + let (x, out) = split_read_write( + tan, + tangent_base::(vec_src.index(), primal_len), + csr.shape.cols * K, + tangent_base::(dst.index(), primal_len), + csr.shape.rows * K, + ); + dispatch(CsrTangents:: { csr, x, out }); + }, + + other => unreachable!("instruction {other:?} is not batchable; check is_batchable first"), + } +} + +/// `out = A @ x` for `K` tangent directions, each accumulated as the scalar `CsrProduct` does. +struct CsrTangents<'a, const K: usize> { + csr: &'a CsrData, + x: &'a [f64], + out: &'a mut [f64], +} + +impl Product for CsrTangents<'_, K> { + type Output = (); + + // Must inline into `dispatch`'s FMA-enabled wrapper. + #[allow(clippy::inline_always)] + #[inline(always)] + fn run(self) { + let Self { csr, x, out } = self; + for (out, (cols, vals)) in out.as_chunks_mut::().0.iter_mut().zip(csr_rows(csr)) { + let mut acc = [0.0_f64; K]; + for (&col, &value) in cols.iter().zip(vals) { + for (a, &v) in acc.iter_mut().zip(&x[col * K..][..K]) { + *a = A::mul_acc(*a, value, v); + } + } + out.copy_from_slice(&acc); + } + } +} + +/// The lane-minor index of tangent slot `offset`. +#[inline] +fn tangent_base(offset: usize, primal_len: usize) -> usize { + debug_assert!( + offset >= primal_len, + "is_batchable admits only tangent slots here" + ); + (offset - primal_len) * K +} + +#[allow(clippy::too_many_arguments)] +fn binary_op( + op: BinaryOp, + a: Operand, + b: Operand, + dst: usize, + len: usize, + kind: BroadcastKind, + primal: &[f64], + tan: &mut [f64], +) { + with_binary_op!(op, f => binary::(f, a, b, dst, len, kind, primal, tan)); +} + +fn unary(op: UnaryOp, src: usize, dst: usize, len: usize, tan: &mut [f64]) { + // Lane-minor, an element-wise op over `len` elements is one over `len * K` values. + with_unary_op!(op, f => unary_apply(tan, f, src, dst, len * K)); +} + +/// `out[e] = f(a(e), b(e))` lane by lane, for each element `e` of `out`. +// Like `eval`'s op helpers, it must inline so the operand closures fuse into one loop. +#[allow(clippy::inline_always)] +#[inline(always)] +fn zip_lanes( + out: &mut [[f64; K]], + f: &impl Fn(f64, f64) -> f64, + a: impl Fn(usize) -> [f64; K], + b: impl Fn(usize) -> [f64; K], +) { + for (e, out) in out.iter_mut().enumerate() { + for ((o, x), y) in out.iter_mut().zip(a(e)).zip(b(e)) { + *o = f(x, y); + } + } +} + +#[allow(clippy::too_many_arguments)] +fn binary f64>( + f: F, + a: Operand, + b: Operand, + dst: usize, + len: usize, + kind: BroadcastKind, + primal: &[f64], + tan: &mut [f64], +) { + use BroadcastKind::{ScalarScalar, ScalarVector, VectorScalar, VectorVector}; + if len == 1 { + // Splitting windows costs more than one element's op, which dominates scalar-row tapes. + let mut out = [0.0_f64; K]; + match (a, b) { + (Operand::Tangent(a), Operand::Tangent(b)) => { + for ((o, &x), &y) in out.iter_mut().zip(&tan[a..][..K]).zip(&tan[b..][..K]) { + *o = f(x, y); + } + }, + (Operand::Primal(a), Operand::Tangent(b)) => { + let x = primal[a]; + for (o, &y) in out.iter_mut().zip(&tan[b..][..K]) { + *o = f(x, y); + } + }, + (Operand::Tangent(a), Operand::Primal(b)) => { + let y = primal[b]; + for (o, &x) in out.iter_mut().zip(&tan[a..][..K]) { + *o = f(x, y); + } + }, + (Operand::Primal(_), Operand::Primal(_)) => { + unreachable!("is_batchable admits no binary over two primal operands") + }, + } + tan[dst..][..K].copy_from_slice(&out); + return; + } + // Lane-minor, a vector operand is `len` chunks of `K` lanes and a broadcast one a single chunk. + let a_len = if matches!(kind, VectorScalar | VectorVector) { + len + } else { + 1 + }; + let b_len = if matches!(kind, ScalarVector | VectorVector) { + len + } else { + 1 + }; + match (a, b) { + (Operand::Tangent(a), Operand::Tangent(b)) => { + let (x, y, out) = + split_two_read_write(tan, (a, a_len * K), (b, b_len * K), (dst, len * K)); + let (x, y, out) = ( + x.as_chunks::().0, + y.as_chunks::().0, + out.as_chunks_mut::().0, + ); + match kind { + ScalarScalar | VectorVector => zip_lanes(out, &f, |e| x[e], |e| y[e]), + ScalarVector => zip_lanes(out, &f, |_| x[0], |e| y[e]), + VectorScalar => zip_lanes(out, &f, |e| x[e], |_| y[0]), + } + }, + (Operand::Primal(a), Operand::Tangent(b)) => { + let (y, out) = split_read_write(tan, b, b_len * K, dst, len * K); + let (y, out) = (y.as_chunks::().0, out.as_chunks_mut::().0); + let x = &primal[a..a + a_len]; + match kind { + ScalarScalar | VectorVector => zip_lanes(out, &f, |e| [x[e]; K], |e| y[e]), + ScalarVector => zip_lanes(out, &f, |_| [x[0]; K], |e| y[e]), + VectorScalar => zip_lanes(out, &f, |e| [x[e]; K], |_| y[0]), + } + }, + (Operand::Tangent(a), Operand::Primal(b)) => { + let (x, out) = split_read_write(tan, a, a_len * K, dst, len * K); + let (x, out) = (x.as_chunks::().0, out.as_chunks_mut::().0); + let y = &primal[b..b + b_len]; + match kind { + ScalarScalar | VectorVector => zip_lanes(out, &f, |e| x[e], |e| [y[e]; K]), + ScalarVector => zip_lanes(out, &f, |_| x[0], |e| [y[e]; K]), + VectorScalar => zip_lanes(out, &f, |e| x[e], |_| [y[0]; K]), + } + }, + (Operand::Primal(_), Operand::Primal(_)) => { + unreachable!("is_batchable admits no binary over two primal operands") + }, + } +} diff --git a/packages/pybamm-rust/pybamm-core/src/zero_propagate.rs b/packages/pybamm-rust/pybamm-core/src/zero_propagate.rs new file mode 100644 index 0000000000..2919a9c4d7 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/src/zero_propagate.rs @@ -0,0 +1,373 @@ +//! Shape-aware zero propagation: folds every subtree it proves identically zero. +//! +//! The rules are `CasADi`'s on-the-fly ones (`0 * x`, `0 / x`, `0 @ v` and `M @ 0` are zero), +//! so the result matches the input up to the sign of zero wherever every intermediate of the +//! input is finite. At a `0 * inf` or `0 / 0` it yields zero where the input yields NaN. + +use crate::arena::{Arena, Dag, DagRef, NodeId, NodeMap}; +use crate::ir::node_width; +use crate::node::Node; + +/// A node's output length, and whether its value is proved identically zero. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct ShapeInfo { + len: usize, + all_zero: bool, +} + +fn infer_shape(arena: &Arena, id: NodeId, shapes: &[ShapeInfo]) -> ShapeInfo { + let node = arena.get(id); + let is_zero = |child: &NodeId| shapes[child.index()].all_zero; + let all_zero = match node { + Node::Scalar(v) => *v == 0.0, + Node::ZeroVector { .. } => true, + Node::Array(arr) => arr.data.iter().all(|&v| v == 0.0), + + Node::Add(a, b) | Node::Sub(a, b) => is_zero(a) && is_zero(b), + Node::Mul(a, b) | Node::MatMul(a, b) => is_zero(a) || is_zero(b), + Node::Div(a, _) => is_zero(a), + + Node::Neg(c) + | Node::Abs(c) + | Node::Index { child: c, .. } + | Node::Sqrt(c) + | Node::Sin(c) + | Node::Tanh(c) + | Node::Sinh(c) + | Node::Arcsinh(c) + | Node::Arctan(c) + | Node::Erf(c) + | Node::Sign(c) + | Node::Floor(c) + | Node::Ceiling(c) + | Node::MaxReduce(c) + | Node::MinReduce(c) + | Node::ReduceArgSelect { basis: c, .. } => is_zero(c), + // A conditional yields one of its branches or zeros, whatever the selector. + Node::Concat(children) + | Node::Conditional { + branches: children, .. + } => children.iter().all(is_zero), + + // log(0) = -inf, and exp, cos and cosh of 0 are 1. + Node::Log(_) + | Node::Exp(_) + | Node::Cos(_) + | Node::Cosh(_) + | Node::SparseMatrix(_) + | Node::StateVector { .. } + | Node::StateVectorDot { .. } + | Node::TangentStateVector { .. } + | Node::InputParameter { .. } + | Node::TangentParameter { .. } + | Node::Time + | Node::Pow(..) + | Node::Minimum(..) + | Node::Maximum(..) + | Node::Hypot(..) + | Node::Modulo(..) + | Node::EqualHeaviside(..) + | Node::NotEqualHeaviside(..) + | Node::Equality(..) + | Node::Interpolant1DLinear { .. } + | Node::Interpolant1DLinearDeriv { .. } + | Node::Interpolant1DCubic { .. } + | Node::Interpolant1DCubicDeriv { .. } + | Node::InterpolantNd { .. } + | Node::InterpolantNdPartial { .. } => false, + }; + let len = node_width(arena, node, |child| shapes[child.index()].len); + ShapeInfo { len, all_zero } +} + +/// Shape of every node in `order`, indexed by `NodeId`. +fn analyze_shapes(arena: &Arena, order: &[NodeId]) -> Vec { + let unvisited = ShapeInfo { + len: 0, + all_zero: false, + }; + let mut shapes = vec![unvisited; arena.len()]; + for &id in order { + shapes[id.index()] = infer_shape(arena, id, &shapes); + } + shapes +} + +/// Rebuild `dag` with every subtree proved identically zero folded to a zero literal. +/// +/// Exact up to the sign of zero wherever every intermediate of `dag` is finite; see the module +/// docs for what happens elsewhere. +pub fn zero_propagate(dag: DagRef<'_>) -> Dag { + let (arena, root) = (dag.arena(), dag.root()); + let order = arena.topological_order(root); + let shapes = analyze_shapes(arena, &order); + // `zero + other` folds to `other` only when no broadcast is lost. + let absorbed = |zero: NodeId, other: NodeId| { + shapes[zero.index()].all_zero && shapes[zero.index()].len == shapes[other.index()].len + }; + + let mut new_arena = Arena::new(); + let mut old_to_new: NodeMap = NodeMap::new(arena.len()); + + for old_id in order { + let shape = shapes[old_id.index()]; + let node = arena.get(old_id); + let new_id = match *node { + _ if shape.all_zero => new_arena.alloc(Node::zero(shape.len)), + Node::Add(a, b) if absorbed(a, b) => old_to_new.rebuilt(b), + Node::Add(a, b) | Node::Sub(a, b) if absorbed(b, a) => old_to_new.rebuilt(a), + Node::Sub(a, b) if absorbed(a, b) => new_arena.alloc(Node::Neg(old_to_new.rebuilt(b))), + _ => new_arena.alloc(node.map_children(|c| old_to_new.rebuilt(c))), + }; + old_to_new.insert(old_id, new_id); + } + + Dag::new(new_arena, old_to_new.rebuilt(root)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::eval::CompiledExpr; + use crate::node::{ArrayData, CsrData, Shape}; + + fn array(data: Vec, shape: Shape) -> Node { + Node::Array(Box::new( + ArrayData::try_new(data, shape).expect("a valid array"), + )) + } + + /// `root`'s shape as the pass infers it. + fn shape_of(arena: &Arena, root: NodeId) -> ShapeInfo { + analyze_shapes(arena, &arena.topological_order(root))[root.index()] + } + + fn eval(dag: DagRef<'_>, y: &[f64]) -> Vec { + let compiled = CompiledExpr::new(dag); + let mut scratch = vec![0.0; compiled.scratch_len()]; + compiled.eval(&mut scratch, 0.0, y, &[], &[]).to_vec() + } + + #[test] + fn a_literal_is_zero_only_when_every_entry_is() { + let cases = [ + (Node::Scalar(0.0), 1, true), + (Node::Scalar(-0.0), 1, true), + (Node::Scalar(5.0), 1, false), + (Node::ZeroVector { len: 100 }, 100, true), + (Node::StateVector { start: 0, end: 10 }, 10, false), + (array(vec![0.0, -0.0, 0.0], Shape::vector(3)), 3, true), + // Tiny nonzero coefficients must not be erased as zeros. + (array(vec![1.0e-17, -2.0e-17], Shape::vector(2)), 2, false), + ]; + for (node, len, all_zero) in cases { + let mut arena = Arena::new(); + let id = arena.alloc(node); + let expected = ShapeInfo { len, all_zero }; + assert_eq!(shape_of(&arena, id), expected, "{:?}", arena.get(id)); + } + } + + #[test] + fn operators_fold_by_casadis_on_the_fly_rules() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let selector = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let zero = arena.alloc(Node::ZeroVector { len: 2 }); + let matrix = arena.alloc(array(vec![1.0, 2.0, 3.0, 4.0], Shape::matrix(2, 2))); + let zero_matrix = arena.alloc(array(vec![0.0; 4], Shape::matrix(2, 2))); + let conditional = |branches| Node::Conditional { selector, branches }; + let cases = [ + (Node::Mul(zero, x), true), + (Node::Mul(x, zero), true), + (Node::Add(zero, x), false), + (Node::Add(zero, zero), true), + (Node::Sub(zero, zero), true), + (Node::Div(zero, x), true), + (Node::Div(x, zero), false), + (Node::MatMul(matrix, zero), true), + (Node::MatMul(zero_matrix, x), true), + (Node::MatMul(matrix, x), false), + (Node::Sin(zero), true), + (Node::Exp(zero), false), + (Node::Cos(zero), false), + (Node::Log(zero), false), + (Node::Concat(vec![zero, zero]), true), + (Node::Concat(vec![zero, x]), false), + (conditional(vec![zero, zero]), true), + (conditional(vec![zero, x]), false), + ]; + for (node, all_zero) in cases { + let id = arena.alloc(node); + assert_eq!( + shape_of(&arena, id).all_zero, + all_zero, + "{:?}", + arena.get(id) + ); + } + } + + #[test] + fn the_rebuild_drops_proven_zeros_and_keeps_the_rest() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let y = arena.alloc(Node::StateVector { start: 5, end: 10 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let zero_times_x = arena.alloc(Node::Mul(zero, x)); + let csr = CsrData::try_new( + vec![0, 2, 4], + vec![0, 1, 0, 1], + vec![1.0, 2.0, 3.0, 4.0], + Shape::matrix(2, 2), + ) + .expect("a valid matrix"); + let matrix = arena.alloc(Node::SparseMatrix(Box::new(csr))); + let zero_vector = arena.alloc(Node::ZeroVector { len: 2 }); + let rebuilt = |arena: &mut Arena, node| { + let root = arena.alloc(node); + let folded = zero_propagate(DagRef::new(arena, root)); + folded.arena().get(folded.root()).clone() + }; + + let y_alone = rebuilt(&mut arena, Node::Add(zero_times_x, y)); + assert_eq!(y_alone, Node::StateVector { start: 5, end: 10 }); + let product = rebuilt(&mut arena, Node::MatMul(matrix, zero_vector)); + assert_eq!(product, Node::ZeroVector { len: 2 }); + let kept = rebuilt(&mut arena, Node::Add(x, y)); + assert!( + matches!(kept, Node::Add(..)), + "a nonzero sum is kept, got {kept:?}" + ); + } + + #[test] + fn test_zero_propagate_preserves_tiny_nonzero_derivative_coefficients() { + use crate::eval::TangentInputs; + use crate::tangent::tangent_wrt_states; + + let mut arena = Arena::new(); + let coeff = arena.alloc(array(vec![1.0e-17, -2.0e-17], Shape::vector(2))); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let expr = arena.alloc(Node::Mul(coeff, y)); + + let tangent = tangent_wrt_states(&mut arena, expr); + let folded = zero_propagate(DagRef::new(&arena, tangent)); + let compiled = CompiledExpr::new(folded.as_ref()); + let mut s_compiled = vec![0.0; compiled.scratch_len()]; + let tangent_inputs = TangentInputs { + dy: Some(&[1.0, 1.0]), + dp: None, + }; + + let result = compiled.eval_with_tangent( + &mut s_compiled, + 0.0, + &[3.0, 4.0], + &[], + &[], + &tangent_inputs, + ); + assert_eq!(result.len(), 2); + assert!((result[0] - 1.0e-17).abs() < f64::EPSILON); + assert!((result[1] - (-2.0e-17)).abs() < f64::EPSILON); + } + + /// Its own length is never read, so a huge declared shape must not overflow. + #[test] + fn a_sparse_matrix_literal_measures_zero_as_lowering_does() { + let mut arena = Arena::new(); + let matrix = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new(vec![0, 0, 0], vec![], vec![], Shape::matrix(2, usize::MAX)) + .expect("an empty CSR matrix is valid"), + ))); + assert_eq!(shape_of(&arena, matrix).len, 0); + } + + /// Regression: an all-zero dense matrix used to fold to a `ZeroVector` and stay a + /// `MatMul` operand, which lowering rejects; now the product folds with it. + #[test] + fn a_zero_matrix_literal_folds_its_product() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let sin_y = arena.alloc(Node::Sin(y)); + let matrix = arena.alloc(array(vec![0.0; 4], Shape::matrix(2, 2))); + let product = arena.alloc(Node::MatMul(matrix, sin_y)); + let root = arena.alloc(Node::Add(product, y)); + + let folded = zero_propagate(DagRef::new(&arena, root)); + assert_eq!(folded.arena().get(folded.root()), arena.get(y)); + assert_eq!(eval(folded.as_ref(), &[0.5, -1.5]), [0.5, -1.5]); + } + + /// The motivating case: a parameter tangent of an input-free quotient is `0 / x`. + #[test] + fn a_parameter_tangent_of_an_input_free_quotient_folds_away() { + use crate::simplify::simplify_pipeline; + use crate::tangent::tangent_wrt_params; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let log = arena.alloc(Node::Log(y)); + let root = arena.alloc(Node::Div(log, y)); + + let mut tangent = Dag::new(arena, root); + tangent.rewrite_in_place(tangent_wrt_params); + let folded = simplify_pipeline(tangent); + assert_eq!(folded.arena().nodes(), [Node::ZeroVector { len: 3 }]); + } + + #[test] + fn test_pow_shape_uses_node_width() { + let mut arena = Arena::new(); + let base = arena.alloc(Node::Scalar(2.0)); + let exp = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let pow = arena.alloc(Node::Pow(base, exp)); + + assert_eq!(shape_of(&arena, pow).len, 3); + } + + #[test] + fn test_zero_times_pow_vector_exponent_preserves_shape() { + let mut arena = Arena::new(); + let zero = arena.alloc(Node::Scalar(0.0)); + let base = arena.alloc(Node::Scalar(2.0)); + let exp = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let pow = arena.alloc(Node::Pow(base, exp)); + let expr = arena.alloc(Node::Mul(zero, pow)); + + let folded = zero_propagate(DagRef::new(&arena, expr)); + + assert_eq!( + folded.arena().get(folded.root()), + &Node::ZeroVector { len: 3 } + ); + } + + #[test] + #[allow(clippy::float_cmp)] + fn zero_propagate_may_normalise_the_sign_of_zero() { + // `-y + 0` and `0 - y` fold to `-y`; at y = 0 the unfolded sum and difference are +0.0. + let builds: [fn(&mut Arena, NodeId, NodeId) -> Node; 2] = [ + |arena, y, zero| Node::Add(arena.alloc(Node::Neg(y)), zero), + |_, y, zero| Node::Sub(zero, y), + ]; + for build in builds { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zero = arena.alloc(Node::Scalar(0.0)); + let node = build(&mut arena, y, zero); + let root = arena.alloc(node); + + let a = eval(DagRef::new(&arena, root), &[0.0])[0]; + let b = eval(zero_propagate(DagRef::new(&arena, root)).as_ref(), &[0.0])[0]; + + let label = format!("{:?}", arena.get(root)); + assert_eq!(a, b, "{label}: values must be equal"); + assert!(a == 0.0 && b == 0.0, "{label}: both must be zero"); + // The sign is not guaranteed; pin it so a contract change fails here. + assert!(a.is_sign_positive(), "{label}: unfolded is +0.0"); + assert!(b.is_sign_negative(), "{label}: the folded form keeps -0.0"); + } + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/common/cases.rs b/packages/pybamm-rust/pybamm-core/tests/common/cases.rs new file mode 100644 index 0000000000..f85bb1313d --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/common/cases.rs @@ -0,0 +1,1277 @@ +use proptest::prelude::*; +use pybamm_core::{Arena, CompiledExpr, DagRef, Node, NodeId, TypedIr}; + +#[derive(Clone, Debug)] +#[allow(dead_code)] +pub struct DagCase { + pub arena: Arena, + pub root: NodeId, + pub y: Vec, + pub n_states: usize, + /// Packed input-parameter values; empty unless [`Extra::Inputs`] is on. + pub inputs: Vec, +} + +#[derive(Clone, Debug)] +#[allow(dead_code)] +pub struct TangentCase { + pub arena: Arena, + pub root: NodeId, + pub y: Vec, + pub seeds: Vec>, + pub n_states: usize, + /// Packed input-parameter values; empty unless [`Extra::Inputs`] is on. + pub inputs: Vec, +} + +/// An opt-in generator extension; with none, every mode draws exactly its original cases. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[allow(dead_code)] +pub enum Extra { + /// `InputParameter` leaves of widths 1, 2 and 3, read from the case's `inputs`. + Inputs, + /// `Pow` of a positive base by a literal or an input-parameter exponent. + Pow, + /// `Minimum` and `Maximum`. + MinMax, + /// `0.0`, `-0.0` and `ZeroVector` leaves, never a `Div` denominator. + Zeros, + /// A sparse `MatMul` of random pattern, a `Conditional` on a literal or input selector, + /// a 2-knot linear interpolant, and `MaxReduce`/`MinReduce`. + Structural, + /// `Log` of a positive argument and `Hypot` with a literal leg; away from the smooth mode, + /// also `Modulo` by a positive divisor, `Sign` and `Floor`, and in the full mode `<` and `==`. + Elementwise, +} + +#[allow(dead_code)] +pub const ALL_EXTRAS: &[Extra] = &[ + Extra::Inputs, + Extra::Pow, + Extra::MinMax, + Extra::Zeros, + Extra::Structural, + Extra::Elementwise, +]; + +/// Widths of the [`Extra::Inputs`] leaves, packed in order. +const INPUT_WIDTHS: [usize; 3] = [1, 2, 3]; +const N_INPUTS: usize = 6; +/// Literal exponents for [`Extra::Pow`]: some `simplify` lowers, some it leaves as `pow`. +const POW_EXPONENTS: [f64; 6] = [2.0, 3.0, 4.0, -1.0, -2.0, 0.5]; + +#[allow(dead_code)] +pub fn eval_dag(arena: &Arena, root: NodeId, t: f64, y: &[f64], inputs: &[f64]) -> Vec { + let ir = TypedIr::from_dag(DagRef::new(arena, root)); + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + expr.eval(&mut s, t, y, &[], inputs).to_vec() +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ValueClass { + AnyFinite, + PositiveFinite, + NonZeroFinite, + BoundedSmall, + Zero, +} + +#[derive(Clone, Copy, Debug)] +enum UnarySpec { + Neg, + Abs, + Sqrt, + Exp, + Sin, + Cos, + Tanh, + Sinh, + Cosh, + Arcsinh, + Arctan, + Erf, + Log, + Sign, + Floor, +} + +#[derive(Clone, Copy, Debug)] +enum BinarySpec { + Add, + Sub, + Mul, + Div, + EqualHeaviside, + Pow, + Minimum, + Maximum, + Modulo, + Hypot, + NotEqualHeaviside, + Equality, +} + +#[derive(Clone, Debug)] +enum ExprSpec { + Scalar(f64), + Input { + index: usize, + offset: usize, + width: usize, + }, + ZeroVector(usize), + StateSlice { + start: usize, + end: usize, + }, + Unary { + op: UnarySpec, + child: Box, + }, + Binary { + op: BinarySpec, + lhs: Box, + rhs: Box, + }, + Concat(Vec), + Index { + child: Box, + start: usize, + end: usize, + }, + Conditional { + selector: Box, + branches: Vec, + }, + Interpolant1DLinear { + x_data: Vec, + y_data: Vec, + child: Box, + }, + SparseMatMul { + indptr: Vec, + indices: Vec, + data: Vec, + nrows: usize, + ncols: usize, + rhs: Box, + }, + Reduce { + max: bool, + child: Box, + }, +} + +#[derive(Clone, Debug)] +struct TypedEntry { + spec: ExprSpec, + len: usize, + class: ValueClass, +} + +const fn can_broadcast(a_len: usize, b_len: usize) -> bool { + a_len == b_len || a_len == 1 || b_len == 1 +} + +fn lower_expr(arena: &mut Arena, spec: &ExprSpec) -> NodeId { + match spec { + ExprSpec::Scalar(v) => arena.alloc(Node::Scalar(*v)), + &ExprSpec::Input { + index, + offset, + width, + } => arena.alloc(Node::InputParameter { + name: format!("p{index}"), + index, + offset, + width, + }), + &ExprSpec::ZeroVector(len) => arena.alloc(Node::ZeroVector { len }), + ExprSpec::StateSlice { start, end, .. } => arena.alloc(Node::StateVector { + start: *start, + end: *end, + }), + ExprSpec::Unary { op, child, .. } => { + let child = lower_expr(arena, child); + match op { + UnarySpec::Neg => arena.alloc(Node::Neg(child)), + UnarySpec::Abs => arena.alloc(Node::Abs(child)), + UnarySpec::Sqrt => arena.alloc(Node::Sqrt(child)), + UnarySpec::Exp => arena.alloc(Node::Exp(child)), + UnarySpec::Sin => arena.alloc(Node::Sin(child)), + UnarySpec::Cos => arena.alloc(Node::Cos(child)), + UnarySpec::Tanh => arena.alloc(Node::Tanh(child)), + UnarySpec::Sinh => arena.alloc(Node::Sinh(child)), + UnarySpec::Cosh => arena.alloc(Node::Cosh(child)), + UnarySpec::Arcsinh => arena.alloc(Node::Arcsinh(child)), + UnarySpec::Arctan => arena.alloc(Node::Arctan(child)), + UnarySpec::Erf => arena.alloc(Node::Erf(child)), + UnarySpec::Log => arena.alloc(Node::Log(child)), + UnarySpec::Sign => arena.alloc(Node::Sign(child)), + UnarySpec::Floor => arena.alloc(Node::Floor(child)), + } + }, + ExprSpec::Binary { op, lhs, rhs, .. } => { + let lhs = lower_expr(arena, lhs); + let rhs = lower_expr(arena, rhs); + match op { + BinarySpec::Add => arena.alloc(Node::Add(lhs, rhs)), + BinarySpec::Sub => arena.alloc(Node::Sub(lhs, rhs)), + BinarySpec::Mul => arena.alloc(Node::Mul(lhs, rhs)), + BinarySpec::Div => arena.alloc(Node::Div(lhs, rhs)), + BinarySpec::EqualHeaviside => arena.alloc(Node::EqualHeaviside(lhs, rhs)), + BinarySpec::Pow => arena.alloc(Node::Pow(lhs, rhs)), + BinarySpec::Minimum => arena.alloc(Node::Minimum(lhs, rhs)), + BinarySpec::Maximum => arena.alloc(Node::Maximum(lhs, rhs)), + BinarySpec::Modulo => arena.alloc(Node::Modulo(lhs, rhs)), + BinarySpec::Hypot => arena.alloc(Node::Hypot(lhs, rhs)), + BinarySpec::NotEqualHeaviside => arena.alloc(Node::NotEqualHeaviside(lhs, rhs)), + BinarySpec::Equality => arena.alloc(Node::Equality(lhs, rhs)), + } + }, + ExprSpec::Concat(children) => { + let children: Vec = children.iter().map(|c| lower_expr(arena, c)).collect(); + arena.alloc(Node::Concat(children)) + }, + ExprSpec::Index { + child, start, end, .. + } => { + let child = lower_expr(arena, child); + arena.alloc(Node::Index { + child, + start: *start, + end: *end, + }) + }, + ExprSpec::Conditional { + selector, branches, .. + } => { + let selector = lower_expr(arena, selector); + let branches: Vec = branches.iter().map(|b| lower_expr(arena, b)).collect(); + arena.alloc(Node::Conditional { selector, branches }) + }, + ExprSpec::Interpolant1DLinear { + x_data, + y_data, + child, + .. + } => { + let child = lower_expr(arena, child); + arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + pybamm_core::InterpolantData::try_new(x_data.clone(), y_data.clone()) + .expect("valid interpolant"), + ), + child, + }) + }, + ExprSpec::SparseMatMul { + indptr, + indices, + data, + nrows, + ncols, + rhs, + .. + } => { + let mat = arena.alloc(Node::SparseMatrix(Box::new( + pybamm_core::CsrData::try_new( + indptr.clone(), + indices.clone(), + data.clone(), + pybamm_core::Shape::matrix(*nrows, *ncols), + ) + .expect("valid test matrix"), + ))); + let rhs = lower_expr(arena, rhs); + arena.alloc(Node::MatMul(mat, rhs)) + }, + ExprSpec::Reduce { max, child } => { + let child = lower_expr(arena, child); + if *max { + arena.alloc(Node::MaxReduce(child)) + } else { + arena.alloc(Node::MinReduce(child)) + } + }, + } +} + +#[allow(clippy::cast_possible_truncation)] +fn seed_set_from_entropy(n_states: usize, entropy: &[u8]) -> Vec> { + let mut seeds: Vec> = (0..n_states) + .map(|j| { + let mut v = vec![0.0; n_states]; + v[j] = 1.0; + v + }) + .collect(); + + seeds.push(vec![1.0; n_states]); + seeds.push( + (0..n_states) + .map(|i| if i % 2 == 0 { 1.0 } else { -1.0 }) + .collect(), + ); + + let mut dense = vec![0.0; n_states]; + for (i, slot) in dense.iter_mut().enumerate() { + let b = entropy + .get(i) + .copied() + .unwrap_or_else(|| (i as u8).wrapping_mul(17)); + *slot = (f64::from(b) / 127.5) - 1.0; + } + let max_abs = dense + .iter() + .fold(0.0_f64, |acc, v| acc.max(v.abs())) + .max(1e-12); + for slot in &mut dense { + *slot /= max_abs; + } + seeds.push(dense); + seeds +} + +const fn unary_input_class(op: UnarySpec, mode: GenMode) -> ValueClass { + match op { + UnarySpec::Exp => ValueClass::BoundedSmall, + // Sinh/Cosh grow like Exp, so unbounded inputs break the FD oracle. + UnarySpec::Sinh | UnarySpec::Cosh => match mode { + GenMode::Smooth => ValueClass::BoundedSmall, + GenMode::Full | GenMode::TangentSafe => ValueClass::AnyFinite, + }, + UnarySpec::Sqrt | UnarySpec::Log => ValueClass::PositiveFinite, + UnarySpec::Sign + | UnarySpec::Floor + | UnarySpec::Arcsinh + | UnarySpec::Arctan + | UnarySpec::Erf + | UnarySpec::Sin + | UnarySpec::Cos + | UnarySpec::Tanh + | UnarySpec::Neg + | UnarySpec::Abs => ValueClass::AnyFinite, + } +} + +const fn unary_output_class(op: UnarySpec, child_class: ValueClass) -> ValueClass { + match op { + UnarySpec::Exp | UnarySpec::Sqrt | UnarySpec::Abs => ValueClass::PositiveFinite, + UnarySpec::Sin + | UnarySpec::Cos + | UnarySpec::Tanh + | UnarySpec::Erf + | UnarySpec::Arctan + | UnarySpec::Sign => ValueClass::BoundedSmall, + UnarySpec::Neg => match child_class { + ValueClass::BoundedSmall => ValueClass::BoundedSmall, + ValueClass::NonZeroFinite => ValueClass::NonZeroFinite, + _ => ValueClass::AnyFinite, + }, + UnarySpec::Sinh + | UnarySpec::Cosh + | UnarySpec::Arcsinh + | UnarySpec::Log + | UnarySpec::Floor => ValueClass::AnyFinite, + } +} + +const fn binary_output_class(op: BinarySpec, a: ValueClass, b: ValueClass) -> ValueClass { + match op { + BinarySpec::Add | BinarySpec::Sub | BinarySpec::Minimum | BinarySpec::Maximum => { + if matches!(a, ValueClass::BoundedSmall) && matches!(b, ValueClass::BoundedSmall) { + ValueClass::BoundedSmall + } else { + ValueClass::AnyFinite + } + }, + BinarySpec::Div => { + if matches!(a, ValueClass::PositiveFinite) && matches!(b, ValueClass::PositiveFinite) { + ValueClass::PositiveFinite + } else { + ValueClass::AnyFinite + } + }, + // A power is never a base again, so towers cannot overflow. + BinarySpec::Mul + | BinarySpec::EqualHeaviside + | BinarySpec::Pow + | BinarySpec::Modulo + | BinarySpec::Hypot + | BinarySpec::NotEqualHeaviside + | BinarySpec::Equality => ValueClass::AnyFinite, + } +} + +const fn class_satisfies(actual: ValueClass, required: ValueClass) -> bool { + use ValueClass::*; + match required { + AnyFinite => true, + BoundedSmall => matches!(actual, BoundedSmall | Zero), + PositiveFinite => matches!(actual, PositiveFinite), + NonZeroFinite => matches!(actual, NonZeroFinite | PositiveFinite), + Zero => matches!(actual, Zero), + } +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum GenMode { + Full, + TangentSafe, + Smooth, +} + +#[allow(clippy::cast_possible_truncation)] +fn build_random_spec(entropy: &[u8], n_states: usize, mode: GenMode, extras: &[Extra]) -> ExprSpec { + let mut pool: Vec = Vec::new(); + let state_class = if mode == GenMode::Smooth { + ValueClass::BoundedSmall + } else { + ValueClass::PositiveFinite + }; + + pool.push(TypedEntry { + spec: ExprSpec::StateSlice { + start: 0, + end: n_states, + }, + len: n_states, + class: state_class, + }); + + for i in 0..n_states.min(4) { + pool.push(TypedEntry { + spec: ExprSpec::StateSlice { + start: i, + end: i + 1, + }, + len: 1, + class: state_class, + }); + } + + let scalar_vals = [0.5, 1.0, 2.0, -1.0]; + for &v in &scalar_vals { + let class = if v > 0.0 { + ValueClass::PositiveFinite + } else { + ValueClass::NonZeroFinite + }; + pool.push(TypedEntry { + spec: ExprSpec::Scalar(v), + len: 1, + class, + }); + } + + add_extra_leaves(&mut pool, n_states, extras); + + let smooth_unary = [ + UnarySpec::Neg, + UnarySpec::Sin, + UnarySpec::Cos, + UnarySpec::Tanh, + UnarySpec::Sinh, + UnarySpec::Cosh, + UnarySpec::Arcsinh, + UnarySpec::Arctan, + UnarySpec::Erf, + ]; + let extra_unary = [UnarySpec::Abs, UnarySpec::Sqrt]; + let exp_unary = [UnarySpec::Exp]; + let elementwise = extras.contains(&Extra::Elementwise); + let structural = extras.contains(&Extra::Structural); + + let smooth_binary = [BinarySpec::Add, BinarySpec::Sub, BinarySpec::Mul]; + + let target_successful_ops = if mode == GenMode::Smooth { 12 } else { 16 }; + let max_attempts = target_successful_ops * 3; + let mut successful_ops = 0; + let mut attempts = 0; + let mut last_good_idx = 0; + let mut byte_cursor = 0; + + let next_byte = |cursor: &mut usize| -> u8 { + let b = entropy + .get(*cursor) + .copied() + .unwrap_or_else(|| (*cursor as u8).wrapping_mul(31)); + *cursor += 1; + b + }; + + while successful_ops < target_successful_ops && attempts < max_attempts { + attempts += 1; + let op_kind = next_byte(&mut byte_cursor) % if structural { 6 } else { 5 }; + + match op_kind { + 0 | 1 => { + let mut all_unary: Vec = if mode == GenMode::Smooth { + smooth_unary + .iter() + .chain(exp_unary.iter()) + .copied() + .collect() + } else { + smooth_unary + .iter() + .chain(exp_unary.iter()) + .chain(extra_unary.iter()) + .copied() + .collect() + }; + if elementwise { + all_unary.push(UnarySpec::Log); + // Steps would put a jump inside a finite-difference stencil. + if mode != GenMode::Smooth { + all_unary.extend([UnarySpec::Sign, UnarySpec::Floor]); + } + } + let op = all_unary[next_byte(&mut byte_cursor) as usize % all_unary.len()]; + let required = unary_input_class(op, mode); + + let child_idx = next_byte(&mut byte_cursor) as usize % pool.len(); + if class_satisfies(pool[child_idx].class, required) { + let child = &pool[child_idx]; + let out_class = unary_output_class(op, child.class); + let out_len = child.len; + pool.push(TypedEntry { + spec: ExprSpec::Unary { + op, + child: Box::new(child.spec.clone()), + }, + len: out_len, + class: out_class, + }); + last_good_idx = pool.len() - 1; + successful_ops += 1; + } + }, + 2 | 3 => { + let mut all_binary: Vec = match mode { + GenMode::Smooth => smooth_binary.to_vec(), + GenMode::TangentSafe => { + let mut v = smooth_binary.to_vec(); + v.push(BinarySpec::Div); + v + }, + GenMode::Full => { + let mut v = smooth_binary.to_vec(); + v.push(BinarySpec::Div); + v.push(BinarySpec::EqualHeaviside); + v + }, + }; + if extras.contains(&Extra::Pow) { + all_binary.push(BinarySpec::Pow); + } + if extras.contains(&Extra::MinMax) { + all_binary.extend([BinarySpec::Minimum, BinarySpec::Maximum]); + } + if elementwise { + all_binary.push(BinarySpec::Hypot); + if mode != GenMode::Smooth { + all_binary.push(BinarySpec::Modulo); + } + if mode == GenMode::Full { + all_binary.extend([BinarySpec::NotEqualHeaviside, BinarySpec::Equality]); + } + } + let op = all_binary[next_byte(&mut byte_cursor) as usize % all_binary.len()]; + + let a_idx = next_byte(&mut byte_cursor) as usize % pool.len(); + let b_idx = next_byte(&mut byte_cursor) as usize % pool.len(); + if matches!(op, BinarySpec::Pow) { + let base = &pool[a_idx]; + if !class_satisfies(base.class, ValueClass::PositiveFinite) { + continue; + } + let exponent = pow_exponent(next_byte(&mut byte_cursor), extras); + pool.push(TypedEntry { + spec: ExprSpec::Binary { + op, + lhs: Box::new(base.spec.clone()), + rhs: Box::new(exponent), + }, + len: base.len, + class: ValueClass::AnyFinite, + }); + last_good_idx = pool.len() - 1; + successful_ops += 1; + continue; + } + if matches!(op, BinarySpec::Hypot) { + pool.push(hypot_entry(&pool[a_idx], next_byte(&mut byte_cursor))); + last_good_idx = pool.len() - 1; + successful_ops += 1; + continue; + } + let a = &pool[a_idx]; + let b = &pool[b_idx]; + + if !can_broadcast(a.len, b.len) { + continue; + } + + let ok = match op { + BinarySpec::Div => class_satisfies(b.class, ValueClass::NonZeroFinite), + BinarySpec::Modulo => class_satisfies(b.class, ValueClass::PositiveFinite), + _ => true, + }; + + if ok { + let out_len = a.len.max(b.len); + let out_class = binary_output_class(op, a.class, b.class); + pool.push(TypedEntry { + spec: ExprSpec::Binary { + op, + lhs: Box::new(a.spec.clone()), + rhs: Box::new(b.spec.clone()), + }, + len: out_len, + class: out_class, + }); + last_good_idx = pool.len() - 1; + successful_ops += 1; + } + }, + 4 => { + let do_index = next_byte(&mut byte_cursor) % 2 == 0; + if do_index { + let child_idx = next_byte(&mut byte_cursor) as usize % pool.len(); + let child = &pool[child_idx]; + if child.len > 1 { + let start = next_byte(&mut byte_cursor) as usize % child.len; + let max_end = child.len; + let end = + start + 1 + (next_byte(&mut byte_cursor) as usize % (max_end - start)); + let end = end.min(max_end); + pool.push(TypedEntry { + spec: ExprSpec::Index { + child: Box::new(child.spec.clone()), + start, + end, + }, + len: end - start, + class: child.class, + }); + last_good_idx = pool.len() - 1; + successful_ops += 1; + } + } else { + let n_parts = 2 + (next_byte(&mut byte_cursor) as usize % 3); + let parts: Vec = (0..n_parts) + .map(|_| { + let idx = next_byte(&mut byte_cursor) as usize % pool.len(); + pool[idx].clone() + }) + .collect(); + let total_len: usize = parts.iter().map(|p| p.len).sum(); + let specs: Vec = parts.iter().map(|p| p.spec.clone()).collect(); + pool.push(TypedEntry { + spec: ExprSpec::Concat(specs), + len: total_len, + class: ValueClass::AnyFinite, + }); + last_good_idx = pool.len() - 1; + successful_ops += 1; + } + }, + 5 => { + if let Some(entry) = + structural_entry(&pool, &mut || next_byte(&mut byte_cursor), extras) + { + pool.push(entry); + last_good_idx = pool.len() - 1; + successful_ops += 1; + } + }, + _ => {}, + } + } + + pool[last_good_idx].spec.clone() +} + +/// `Hypot` of `leg` and a positive literal, in the order `byte` picks. +/// +/// A literal leg keeps it off the origin, where both modes' partials are `0 / 0`. +fn hypot_entry(leg: &TypedEntry, byte: u8) -> TypedEntry { + const LITERAL_LEGS: [f64; 3] = [0.5, 1.0, 2.0]; + let literal = ExprSpec::Scalar(LITERAL_LEGS[usize::from(byte / 2) % LITERAL_LEGS.len()]); + let leg_spec = Box::new(leg.spec.clone()); + let (lhs, rhs) = if byte.is_multiple_of(2) { + (leg_spec, Box::new(literal)) + } else { + (Box::new(literal), leg_spec) + }; + TypedEntry { + spec: ExprSpec::Binary { + op: BinarySpec::Hypot, + lhs, + rhs, + }, + len: leg.len, + class: ValueClass::PositiveFinite, + } +} + +/// One [`Extra::Structural`] node over `pool`, or `None` when the drawn operands do not fit. +fn structural_entry( + pool: &[TypedEntry], + next_byte: &mut impl FnMut() -> u8, + extras: &[Extra], +) -> Option { + fn pick(pool: &[TypedEntry], byte: u8) -> &TypedEntry { + &pool[usize::from(byte) % pool.len()] + } + match next_byte() % 4 { + 0 => { + let (first, second) = (pick(pool, next_byte()), pick(pool, next_byte())); + if first.len != second.len { + return None; + } + let class = if [first.class, second.class] + .iter() + .all(|&class| class_satisfies(class, ValueClass::BoundedSmall)) + { + ValueClass::BoundedSmall + } else { + ValueClass::AnyFinite + }; + // A literal window centre always picks a branch; the input does unless it lands on + // a half-way value, where the zeros it gives are still bounded. + let selector = match next_byte() % 3 { + 0 => ExprSpec::Scalar(1.0), + 1 => ExprSpec::Scalar(2.0), + _ if extras.contains(&Extra::Inputs) => ExprSpec::Input { + index: 0, + offset: 0, + width: 1, + }, + _ => ExprSpec::Scalar(2.0), + }; + Some(TypedEntry { + spec: ExprSpec::Conditional { + selector: Box::new(selector), + branches: vec![first.spec.clone(), second.spec.clone()], + }, + len: first.len, + class, + }) + }, + 1 => { + let child = pick(pool, next_byte()); + let x0 = f64::from(next_byte() % 4) - 2.0; + let (y0, y1) = ( + f64::from(next_byte() % 5) - 2.0, + f64::from(next_byte() % 5) - 2.0, + ); + Some(TypedEntry { + spec: ExprSpec::Interpolant1DLinear { + x_data: vec![x0, x0 + 1.5], + y_data: vec![y0, y1], + child: Box::new(child.spec.clone()), + }, + len: child.len, + class: ValueClass::AnyFinite, + }) + }, + 2 => { + const COEFFICIENTS: [f64; 5] = [0.5, -1.0, 1.0, 2.0, -0.25]; + let rhs = pick(pool, next_byte()); + let (nrows, ncols) = (1 + usize::from(next_byte() % 4), rhs.len); + let (mut indptr, mut indices, mut data) = (vec![0], Vec::new(), Vec::new()); + for _ in 0..nrows { + for col in 0..ncols { + let byte = next_byte(); + if byte.is_multiple_of(2) { + indices.push(col); + data.push(COEFFICIENTS[usize::from(byte / 2) % COEFFICIENTS.len()]); + } + } + indptr.push(indices.len()); + } + Some(TypedEntry { + spec: ExprSpec::SparseMatMul { + indptr, + indices, + data, + nrows, + ncols, + rhs: Box::new(rhs.spec.clone()), + }, + len: nrows, + class: ValueClass::AnyFinite, + }) + }, + _ => { + let child = pick(pool, next_byte()); + Some(TypedEntry { + spec: ExprSpec::Reduce { + max: next_byte().is_multiple_of(2), + child: Box::new(child.spec.clone()), + }, + len: 1, + class: child.class, + }) + }, + } +} + +/// Push the leaves `extras` add: input parameters, then zeros of each length present. +fn add_extra_leaves(pool: &mut Vec, n_states: usize, extras: &[Extra]) { + let mut lens = vec![n_states]; + if extras.contains(&Extra::Inputs) { + let mut offset = 0; + for (index, width) in INPUT_WIDTHS.into_iter().enumerate() { + pool.push(TypedEntry { + spec: ExprSpec::Input { + index, + offset, + width, + }, + len: width, + class: ValueClass::PositiveFinite, + }); + offset += width; + lens.push(width); + } + } + if extras.contains(&Extra::Zeros) { + for zero in [0.0, -0.0] { + pool.push(TypedEntry { + spec: ExprSpec::Scalar(zero), + len: 1, + class: ValueClass::Zero, + }); + } + lens.sort_unstable(); + lens.dedup(); + for len in lens.into_iter().filter(|&len| len > 1) { + pool.push(TypedEntry { + spec: ExprSpec::ZeroVector(len), + len, + class: ValueClass::Zero, + }); + } + } +} + +/// A literal exponent, or the width-1 input parameter when inputs are on. +fn pow_exponent(byte: u8, extras: &[Extra]) -> ExprSpec { + let choice = usize::from(byte) % (POW_EXPONENTS.len() + 1); + match POW_EXPONENTS.get(choice) { + Some(&exponent) => ExprSpec::Scalar(exponent), + None if extras.contains(&Extra::Inputs) => ExprSpec::Input { + index: 0, + offset: 0, + width: 1, + }, + None => ExprSpec::Scalar(2.0), + } +} + +/// Lower a random `mode` spec drawn from `entropy`. +fn random_dag(entropy: &[u8], n_states: usize, mode: GenMode, extras: &[Extra]) -> (Arena, NodeId) { + let spec = build_random_spec(entropy, n_states, mode, extras); + let mut arena = Arena::new(); + let root = lower_expr(&mut arena, &spec); + (arena, root) +} + +#[allow(clippy::needless_pass_by_value)] +fn build_random_eval_case(entropy: Vec, n_states: usize, y: Vec) -> DagCase { + build_eval_case(&entropy, n_states, y, Vec::new(), &[]) +} + +fn build_eval_case( + entropy: &[u8], + n_states: usize, + y: Vec, + inputs: Vec, + extras: &[Extra], +) -> DagCase { + let (arena, root) = random_dag(entropy, n_states, GenMode::Full, extras); + DagCase { + arena, + root, + y, + n_states, + inputs, + } +} + +#[allow(clippy::needless_pass_by_value)] +fn build_random_smooth_tangent_case(entropy: Vec, n_states: usize, y: Vec) -> TangentCase { + build_tangent_case(&entropy, n_states, y, Vec::new(), GenMode::Smooth, &[]) +} + +#[allow(clippy::needless_pass_by_value)] +fn build_random_split_eval_case(entropy: Vec, n_states: usize, y: Vec) -> TangentCase { + build_tangent_case(&entropy, n_states, y, Vec::new(), GenMode::TangentSafe, &[]) +} + +fn build_tangent_case( + entropy: &[u8], + n_states: usize, + y: Vec, + inputs: Vec, + mode: GenMode, + extras: &[Extra], +) -> TangentCase { + let (arena, root) = random_dag(entropy, n_states, mode, extras); + TangentCase { + arena, + root, + y, + seeds: seed_set_from_entropy(n_states, entropy), + n_states, + inputs, + } +} + +fn tangent_case_from_dag(case: DagCase, entropy: &[u8]) -> TangentCase { + TangentCase { + arena: case.arena, + root: case.root, + y: case.y, + seeds: seed_set_from_entropy(case.n_states, entropy), + n_states: case.n_states, + inputs: case.inputs, + } +} + +/// Deep linear chain: y[0..n] + 1 + 1 + ... (depth additions). +#[allow(dead_code)] +pub fn deep_chain_case(n_states: usize, depth: usize) -> DagCase { + let mut arena = Arena::new(); + let mut cur = arena.alloc(Node::StateVector { + start: 0, + end: n_states, + }); + for _ in 0..depth { + let one = arena.alloc(Node::Scalar(1.0)); + cur = arena.alloc(Node::Add(cur, one)); + } + let y = vec![1.0; n_states]; + DagCase { + arena, + root: cur, + y, + n_states, + inputs: Vec::new(), + } +} + +/// Wide fan-out: Concat(2*y[0], 2*y[1], ..., 2*y[fanout-1]). +#[allow(dead_code)] +pub fn wide_concat_case(fanout: usize) -> DagCase { + let mut arena = Arena::new(); + let mut terms = Vec::with_capacity(fanout); + for i in 0..fanout { + let yi = arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }); + let two = arena.alloc(Node::Scalar(2.0)); + terms.push(arena.alloc(Node::Mul(yi, two))); + } + let root = arena.alloc(Node::Concat(terms)); + let y: Vec = (0..fanout) + .map(|i| (i as f64).mul_add(0.03, 0.05)) + .collect(); + DagCase { + arena, + root, + y, + n_states: fanout, + inputs: Vec::new(), + } +} + +#[allow(dead_code)] +pub fn broadcast_mix_case(n_states: usize) -> DagCase { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { + start: 0, + end: n_states, + }); + let one = arena.alloc(Node::Scalar(1.0)); + let shifted = arena.alloc(Node::Add(y, one)); + let scale = arena.alloc(Node::Scalar(0.5)); + let root = arena.alloc(Node::Mul(shifted, scale)); + DagCase { + arena, + root, + y: vec![1.0; n_states], + n_states, + inputs: Vec::new(), + } +} + +#[allow(dead_code)] +pub fn index_slice_case() -> DagCase { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let root = arena.alloc(Node::Index { + child: y, + start: 1, + end: 3, + }); + DagCase { + arena, + root, + y: vec![0.5, 1.0, 1.5, 2.0], + n_states: 4, + inputs: Vec::new(), + } +} + +#[allow(dead_code)] +pub fn conditional_case() -> DagCase { + let mut arena = Arena::new(); + let selector = arena.alloc(Node::Scalar(1.0)); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let s2 = arena.alloc(Node::Scalar(2.0)); + let b1 = arena.alloc(Node::Add(y0, s2)); + let s3 = arena.alloc(Node::Scalar(3.0)); + let b2 = arena.alloc(Node::Mul(y1, s3)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![b1, b2], + }); + DagCase { + arena, + root, + y: vec![1.0, 2.0], + n_states: 2, + inputs: Vec::new(), + } +} + +#[allow(dead_code)] +pub fn interpolant_case() -> DagCase { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Interpolant1DLinear { + data: Box::new( + pybamm_core::InterpolantData::try_new(vec![0.0, 1.0, 2.0], vec![0.0, 10.0, 20.0]) + .expect("valid interpolant"), + ), + child: x, + }); + DagCase { + arena, + root, + y: vec![1.5], + n_states: 1, + inputs: Vec::new(), + } +} + +#[allow(dead_code)] +pub fn sparse_matmul_case() -> DagCase { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let mat = arena.alloc(Node::SparseMatrix(Box::new( + pybamm_core::CsrData::try_new( + vec![0, 2, 4], + vec![0, 1, 2, 3], + vec![1.0, -1.0, 2.0, 3.0], + pybamm_core::Shape::matrix(2, 4), + ) + .expect("valid test matrix"), + ))); + let root = arena.alloc(Node::MatMul(mat, y)); + DagCase { + arena, + root, + y: vec![0.5, 1.0, 1.5, 2.0], + n_states: 4, + inputs: Vec::new(), + } +} + +/// Intentional structural duplicate: (3*y[0]) + (3*y[0]) built twice. +#[allow(dead_code)] +pub fn duplicate_subexpr_case() -> DagCase { + let mut arena = Arena::new(); + let y1 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let s1 = arena.alloc(Node::Scalar(3.0)); + let m1 = arena.alloc(Node::Mul(y1, s1)); + let y2 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let s2 = arena.alloc(Node::Scalar(3.0)); + let m2 = arena.alloc(Node::Mul(y2, s2)); + let root = arena.alloc(Node::Add(m1, m2)); + DagCase { + arena, + root, + y: vec![1.25], + n_states: 1, + inputs: Vec::new(), + } +} + +/// Deep smooth composition: repeatedly apply tanh(0.5*x + 0.25). +#[allow(dead_code)] +pub fn smooth_composition_case(n_states: usize, depth: usize) -> TangentCase { + let mut arena = Arena::new(); + let mut cur = arena.alloc(Node::StateVector { + start: 0, + end: n_states, + }); + for _ in 0..depth { + let half = arena.alloc(Node::Scalar(0.5)); + let scaled = arena.alloc(Node::Mul(cur, half)); + let bias = arena.alloc(Node::Scalar(0.25)); + let shifted = arena.alloc(Node::Add(scaled, bias)); + cur = arena.alloc(Node::Tanh(shifted)); + } + TangentCase { + arena, + root: cur, + y: vec![0.2; n_states], + seeds: seed_set_from_entropy(n_states, &[31, 32, 33, 34]), + n_states, + inputs: Vec::new(), + } +} + +#[allow(dead_code)] +pub fn targeted_eval_cases() -> Vec { + vec![ + deep_chain_case(1, 64), + deep_chain_case(8, 100), + wide_concat_case(32), + wide_concat_case(64), + broadcast_mix_case(8), + index_slice_case(), + conditional_case(), + interpolant_case(), + sparse_matmul_case(), + duplicate_subexpr_case(), + ] +} + +#[allow(dead_code)] +pub fn targeted_smooth_tangent_cases() -> Vec { + vec![ + smooth_composition_case(1, 64), + smooth_composition_case(4, 32), + tangent_case_from_dag(wide_concat_case(8), &[5, 6, 7, 8]), + tangent_case_from_dag(broadcast_mix_case(6), &[9, 10, 11, 12]), + tangent_case_from_dag(index_slice_case(), &[13, 14, 15, 16]), + ] +} + +#[allow(dead_code)] +pub fn targeted_split_eval_cases() -> Vec { + vec![ + tangent_case_from_dag(wide_concat_case(8), &[17, 18, 19, 20]), + tangent_case_from_dag(index_slice_case(), &[21, 22, 23, 24]), + tangent_case_from_dag(conditional_case(), &[25, 26, 27, 28]), + tangent_case_from_dag(sparse_matmul_case(), &[29, 30, 31, 32]), + ] +} + +fn arb_random_eval_case() -> impl Strategy { + (2_usize..=8).prop_flat_map(|n_states| { + ( + prop::collection::vec(any::(), 8..48), + prop::collection::vec(0.1_f64..10.0, n_states), + ) + .prop_map(move |(entropy, y)| build_random_eval_case(entropy, n_states, y)) + }) +} + +#[allow(dead_code)] +pub fn arb_eval_case() -> impl Strategy { + arb_random_eval_case() +} + +fn arb_random_smooth_tangent_case() -> impl Strategy { + (2_usize..=6).prop_flat_map(|n_states| { + ( + prop::collection::vec(any::(), 8..36), + prop::collection::vec(-1.5_f64..1.5, n_states), + ) + .prop_map(move |(entropy, y)| build_random_smooth_tangent_case(entropy, n_states, y)) + }) +} + +#[allow(dead_code)] +pub fn arb_smooth_tangent_case() -> impl Strategy { + arb_random_smooth_tangent_case() +} + +#[allow(dead_code)] +pub fn arb_split_eval_case() -> impl Strategy { + (2_usize..=6).prop_flat_map(|n_states| { + ( + prop::collection::vec(any::(), 8..36), + prop::collection::vec(0.1_f64..10.0, n_states), + ) + .prop_map(move |(entropy, y)| build_random_split_eval_case(entropy, n_states, y)) + }) +} + +/// Input values for `extras`: positive, so an input can be a `Pow` exponent or a divisor. +fn arb_inputs(extras: &'static [Extra]) -> impl Strategy> { + let len = if extras.contains(&Extra::Inputs) { + N_INPUTS + } else { + 0 + }; + prop::collection::vec(0.5_f64..2.0, len) +} + +/// [`arb_eval_case`] with `extras` switched on. +#[allow(dead_code)] +pub fn arb_eval_case_with(extras: &'static [Extra]) -> impl Strategy { + (2_usize..=8).prop_flat_map(move |n_states| { + ( + prop::collection::vec(any::(), 8..48), + prop::collection::vec(0.1_f64..10.0, n_states), + arb_inputs(extras), + ) + .prop_map(move |(entropy, y, inputs)| { + build_eval_case(&entropy, n_states, y, inputs, extras) + }) + }) +} + +/// [`arb_smooth_tangent_case`] with `extras` switched on. +#[allow(dead_code)] +pub fn arb_smooth_tangent_case_with( + extras: &'static [Extra], +) -> impl Strategy { + (2_usize..=6).prop_flat_map(move |n_states| { + ( + prop::collection::vec(any::(), 8..36), + prop::collection::vec(-1.5_f64..1.5, n_states), + arb_inputs(extras), + ) + .prop_map(move |(entropy, y, inputs)| { + build_tangent_case(&entropy, n_states, y, inputs, GenMode::Smooth, extras) + }) + }) +} + +/// [`arb_split_eval_case`] with `extras` switched on. +#[allow(dead_code)] +pub fn arb_split_eval_case_with(extras: &'static [Extra]) -> impl Strategy { + (2_usize..=6).prop_flat_map(move |n_states| { + ( + prop::collection::vec(any::(), 8..36), + prop::collection::vec(0.1_f64..10.0, n_states), + arb_inputs(extras), + ) + .prop_map(move |(entropy, y, inputs)| { + build_tangent_case(&entropy, n_states, y, inputs, GenMode::TangentSafe, extras) + }) + }) +} diff --git a/packages/pybamm-rust/pybamm-core/tests/common/mod.rs b/packages/pybamm-rust/pybamm-core/tests/common/mod.rs new file mode 100644 index 0000000000..5ba0a9b804 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/common/mod.rs @@ -0,0 +1,2 @@ +pub mod cases; +pub mod numeric_eq; diff --git a/packages/pybamm-rust/pybamm-core/tests/common/numeric_eq.rs b/packages/pybamm-rust/pybamm-core/tests/common/numeric_eq.rs new file mode 100644 index 0000000000..d6f17dcd28 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/common/numeric_eq.rs @@ -0,0 +1,73 @@ +#[allow(dead_code)] +#[track_caller] +pub fn assert_bitwise_eq(lhs: &[f64], rhs: &[f64]) { + assert_eq!( + lhs.len(), + rhs.len(), + "length mismatch: {} vs {}", + lhs.len(), + rhs.len() + ); + for (i, (l, r)) in lhs.iter().zip(rhs.iter()).enumerate() { + assert_eq!( + l.to_bits(), + r.to_bits(), + "bit mismatch at index {i}: {l} vs {r}" + ); + } +} + +/// Bitwise equality, except a ±0 difference: folding `0 * x` and the like may +/// normalise the sign of a zero result. +#[allow(dead_code)] +#[track_caller] +pub fn assert_eq_up_to_zero_sign(lhs: &[f64], rhs: &[f64]) { + assert_eq!( + lhs.len(), + rhs.len(), + "length mismatch: {} vs {}", + lhs.len(), + rhs.len() + ); + for (i, (l, r)) in lhs.iter().zip(rhs.iter()).enumerate() { + if *l == 0.0 && *r == 0.0 { + continue; + } + assert_eq!( + l.to_bits(), + r.to_bits(), + "mismatch at index {i} (beyond the sign-of-zero carve-out): {l} vs {r}" + ); + } +} + +#[allow(dead_code)] +#[track_caller] +pub fn assert_close(lhs: &[f64], rhs: &[f64], rtol: f64, atol: f64) { + assert_eq!( + lhs.len(), + rhs.len(), + "length mismatch: {} vs {}", + lhs.len(), + rhs.len() + ); + for (i, (l, r)) in lhs.iter().zip(rhs.iter()).enumerate() { + assert!( + !l.is_nan() && !r.is_nan(), + "NaN mismatch at index {i}: {l} vs {r}" + ); + if !l.is_finite() || !r.is_finite() { + assert!( + l.to_bits() == r.to_bits(), + "non-finite mismatch at index {i}: {l} vs {r}" + ); + continue; + } + let tol = rtol.mul_add(l.abs().max(r.abs()), atol); + assert!( + (l - r).abs() <= tol, + "mismatch at index {i}: {l} vs {r} (diff={}, tol={tol})", + (l - r).abs() + ); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_batch_eval.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_batch_eval.rs new file mode 100644 index 0000000000..76dee061e5 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_batch_eval.rs @@ -0,0 +1,527 @@ +//! Differential test: `eval_batch` must match scalar `eval` bitwise (same float ops, same order). + +mod common; + +use common::cases::{DagCase, arb_eval_case, targeted_eval_cases}; +use proptest::prelude::*; +use pybamm_core::ir::{BinaryOp, UnaryOp}; +use pybamm_core::{ + Arena, ArrayData, BatchEvalError, CompiledExpr, CsrData, CubicInterpolantData, DagRef, + InterpolantData, NdInterpolantData, Node, NodeId, Shape, tangent_wrt_states, +}; + +/// Assert `eval_batch` over `k` lanes equals `k` scalar `eval`s, bit for bit. +#[track_caller] +fn assert_batch_matches_scalar( + arena: &Arena, + root: NodeId, + n_states: usize, + k: usize, + ts: &[f64], + y_cols: &[f64], + inputs: &[f64], +) { + let expr = CompiledExpr::new(DagRef::new(arena, root)); + let out_len = expr.output_len(); + + let mut scalar = vec![0.0_f64; out_len * k]; + let mut s = vec![0.0_f64; expr.scratch_len()]; + for l in 0..k { + let y = &y_cols[l * n_states..(l + 1) * n_states]; + let res = expr.eval(&mut s, ts[l], y, &[], inputs); + scalar[l * out_len..(l + 1) * out_len].copy_from_slice(res); + } + + let mut batch_scratch = vec![0.0_f64; expr.scratch_len() * k]; + let root_slice = expr + .eval_batch(&mut batch_scratch, k, ts, y_cols, n_states, inputs) + .expect("primal tape must batch-evaluate"); + + for l in 0..k { + for e in 0..out_len { + let got = root_slice[e * k + l]; + let want = scalar[l * out_len + e]; + assert_eq!( + got.to_bits(), + want.to_bits(), + "lane {l}, elem {e}: batch {got} != scalar {want} (k={k})" + ); + } + } +} + +/// `(case, k, ts[k], y_cols[n_states*k])` with `k` sampled from {1, 2, 7, 32}. +fn arb_batch_case() -> impl Strategy, Vec)> { + arb_eval_case() + .prop_flat_map(|case| (Just(case), prop::sample::select(vec![1_usize, 2, 7, 32]))) + .prop_flat_map(|(case, k)| { + let n_states = case.n_states; + ( + Just(case), + Just(k), + prop::collection::vec(-2.0_f64..5.0, k), + prop::collection::vec(0.1_f64..10.0, n_states * k), + ) + }) +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(400))] + + #[test] + fn batch_matches_scalar_random_dag((case, k, ts, y_cols) in arb_batch_case()) { + assert_batch_matches_scalar(&case.arena, case.root, case.n_states, k, &ts, &y_cols, &[]); + } +} + +#[test] +fn targeted_cases_batch_match_scalar() { + for case in targeted_eval_cases() { + let n = case.n_states; + for &k in &[1_usize, 2, 7, 32] { + // Stay near the case's own `y` so guarded ops (sqrt/div) stay in range. + let mut y_cols = vec![0.0_f64; n * k]; + for l in 0..k { + for i in 0..n { + y_cols[l * n + i] = + case.y[i].mul_add(0.03_f64.mul_add(l as f64, 1.0), 0.01 * i as f64); + } + } + let ts: Vec = (0..k).map(|l| l as f64 * 0.1).collect(); + assert_batch_matches_scalar(&case.arena, case.root, n, k, &ts, &y_cols, &[]); + } + } +} + +/// Build a column-major `(n_states, k)` state matrix from a closure. +fn state_cols(n_states: usize, k: usize, f: impl Fn(usize, usize) -> f64) -> Vec { + let mut cols = vec![0.0_f64; n_states * k]; + for l in 0..k { + for i in 0..n_states { + cols[l * n_states + i] = f(i, l); + } + } + cols +} + +#[test] +fn all_binary_ops_match_scalar() { + let ops = [ + BinaryOp::Add, + BinaryOp::Sub, + BinaryOp::Mul, + BinaryOp::Div, + BinaryOp::Pow, + BinaryOp::Minimum, + BinaryOp::Maximum, + BinaryOp::Modulo, + BinaryOp::Hypot, + BinaryOp::EqualHeaviside, + BinaryOp::NotEqualHeaviside, + BinaryOp::Equality, + ]; + let k = 7; + let ts: Vec = (0..k).map(|l| l as f64 * 0.3).collect(); + let y_cols = state_cols(2, k, |i, l| 0.5 + i as f64 + l as f64 * 0.11); + for op in ops { + let mut arena = Arena::new(); + let a = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let b = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let root = alloc_binary(&mut arena, op, a, b); + assert_batch_matches_scalar(&arena, root, 2, k, &ts, &y_cols, &[]); + } +} + +#[test] +fn all_unary_ops_match_scalar() { + let ops = [ + UnaryOp::Neg, + UnaryOp::Abs, + UnaryOp::Sqrt, + UnaryOp::Exp, + UnaryOp::Log, + UnaryOp::Sin, + UnaryOp::Cos, + UnaryOp::Tanh, + UnaryOp::Sinh, + UnaryOp::Cosh, + UnaryOp::Arcsinh, + UnaryOp::Arctan, + UnaryOp::Erf, + UnaryOp::Sign, + UnaryOp::Floor, + UnaryOp::Ceiling, + ]; + let k = 5; + let ts: Vec = (0..k).map(|l| l as f64 * 0.2).collect(); + // Positive so Sqrt/Log stay in-domain. + let y_cols = state_cols(1, k, |_, l| 0.3 + l as f64 * 0.37); + for op in ops { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = alloc_unary(&mut arena, op, x); + assert_batch_matches_scalar(&arena, root, 1, k, &ts, &y_cols, &[]); + } +} + +#[test] +fn broadcast_scalar_vector_per_lane() { + let mut arena = Arena::new(); + let s = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let v = arena.alloc(Node::StateVector { start: 1, end: 4 }); + let root = arena.alloc(Node::Mul(s, v)); + let k = 4; + let ts = vec![0.0; k]; + let y_cols = state_cols(4, k, |i, l| 1.0 + i as f64 * 0.5 + l as f64); + assert_batch_matches_scalar(&arena, root, 4, k, &ts, &y_cols, &[]); +} + +#[test] +fn reduce_argselect_tie_breaking_per_lane() { + // Lane 0: [5,5,3] (max tie at 0), lane 1: [1,5,5] (max tie at 1), lane 2: [3,1,1] + // (min tie at 1), lane 3: [NaN,2,2] (NaN skipped, tie at 1). + let picks = [ + [5.0, 5.0, 3.0], + [1.0, 5.0, 5.0], + [3.0, 1.0, 1.0], + [f64::NAN, 2.0, 2.0], + ]; + let k = picks.len(); + let y_cols = state_cols( + 6, + k, + |i, l| if i < 3 { picks[l][i] } else { 10.0 * i as f64 }, + ); + for is_max in [true, false] { + let mut arena = Arena::new(); + let picker = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let basis = arena.alloc(Node::StateVector { start: 3, end: 6 }); + let root = arena.alloc(Node::ReduceArgSelect { + basis, + picker, + is_max, + }); + assert_batch_matches_scalar(&arena, root, 6, k, &vec![0.0; k], &y_cols, &[]); + } +} + +#[test] +fn max_and_min_reduce_match_scalar() { + let mut arena = Arena::new(); + let v = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let max = arena.alloc(Node::MaxReduce(v)); + let min = arena.alloc(Node::MinReduce(v)); + let root = arena.alloc(Node::Concat(vec![max, min])); + let k = 5; + let y_cols = state_cols(4, k, |i, l| match (i, l) { + (0, 3) => f64::NAN, + (2, 4) => f64::NEG_INFINITY, + _ => ((i * 7 + l * 3) % 5) as f64 - 2.0, + }); + assert_batch_matches_scalar(&arena, root, 4, k, &vec![0.0; k], &y_cols, &[]); +} + +#[test] +fn dense_matmul_and_constant_loads_match_scalar() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let matrix = arena.alloc(Node::Array(Box::new( + ArrayData::try_new(vec![1.5, -2.0, 0.25, 3.0, 0.1, -0.7], Shape::matrix(2, 3)) + .expect("valid array"), + ))); + let product = arena.alloc(Node::MatMul(matrix, y)); + let offset = arena.alloc(Node::Array(Box::new( + ArrayData::try_new(vec![0.3, -0.9], Shape::vector(2)).expect("valid array"), + ))); + let shifted = arena.alloc(Node::Add(product, offset)); + let zeros = arena.alloc(Node::ZeroVector { len: 2 }); + let root = arena.alloc(Node::Concat(vec![shifted, zeros, offset])); + let k = 4; + let y_cols = state_cols(3, k, |i, l| 0.2 + i as f64 * 0.7 - l as f64 * 0.45); + assert_batch_matches_scalar(&arena, root, 3, k, &vec![0.0; k], &y_cols, &[]); +} + +#[test] +fn interp_1d_variants_match_scalar() { + let linear = InterpolantData::try_new(vec![0.0, 0.1, 0.3, 0.7], vec![0.2, -1.3, 2.9, 0.05]) + .expect("valid interpolant"); + let cubic = CubicInterpolantData::try_new( + vec![0.0, 0.4, 1.0], + vec![[1.0, -2.0, 0.5, 3.0], [0.4, 1.5, -1.0, 0.25]], + ) + .expect("valid interpolant"); + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let value_linear = arena.alloc(Node::Interpolant1DLinear { + data: Box::new(linear.clone()), + child: x, + }); + let slope = arena.alloc(Node::Interpolant1DLinearDeriv { + data: Box::new(linear), + child: x, + }); + let value = arena.alloc(Node::Interpolant1DCubic { + data: Box::new(cubic.clone()), + child: x, + }); + let derivative = arena.alloc(Node::Interpolant1DCubicDeriv { + data: Box::new(cubic), + child: x, + }); + let root = arena.alloc(Node::Concat(vec![value_linear, slope, value, derivative])); + // Below, inside and above both tables' domains. + let k = 6; + let y_cols = state_cols(3, k, |i, l| -0.3 + l as f64 * 0.27 + i as f64 * 0.1); + assert_batch_matches_scalar(&arena, root, 3, k, &vec![0.0; k], &y_cols, &[]); +} + +#[test] +fn sparse_matmul_matches_scalar() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let mat = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, 2, 4], + vec![0, 1, 2, 3], + vec![1.0, -1.0, 2.0, 3.0], + Shape::matrix(2, 4), + ) + .expect("valid matrix"), + ))); + let root = arena.alloc(Node::MatMul(mat, y)); + let k = 6; + let ts = vec![0.0; k]; + let y_cols = state_cols(4, k, |i, l| 0.5 + i as f64 + l as f64 * 0.13); + assert_batch_matches_scalar(&arena, root, 4, k, &ts, &y_cols, &[]); +} + +#[test] +fn nd_interpolant_and_partial_match_scalar() { + // Axis 1's length-1 source broadcasts over axis 0's three-element one. + let data = NdInterpolantData::try_new( + vec![vec![0.0, 1.0, 3.0], vec![-1.0, 2.0]], + vec![1.0, -2.0, 0.5, 3.0, -1.5, 0.25, 2.0, 4.0], + 2, + ) + .expect("valid interpolant"); + let mut arena = Arena::new(); + let x0 = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let x1 = arena.alloc(Node::StateVector { start: 3, end: 4 }); + let value = arena.alloc(Node::InterpolantNd { + data: Box::new(data.clone()), + children: vec![x0, x1], + }); + let partial = arena.alloc(Node::InterpolantNdPartial { + data: Box::new(data), + children: vec![x0, x1], + axis: 1, + }); + let root = arena.alloc(Node::Concat(vec![value, partial])); + let k = 5; + let ts = vec![0.0; k]; + let y_cols = state_cols(4, k, |i, l| -0.5 + i as f64 * 0.9 + l as f64 * 0.45); + assert_batch_matches_scalar(&arena, root, 4, k, &ts, &y_cols, &[]); +} + +#[test] +fn time_and_concat_match_scalar() { + let mut arena = Arena::new(); + let t = arena.alloc(Node::Time); + let y0 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let root = arena.alloc(Node::Concat(vec![t, y0])); + let k = 4; + let ts: Vec = (0..k).map(|l| 0.1 + l as f64 * 0.9).collect(); + let y_cols = state_cols(1, k, |_, l| l as f64 * 2.0 - 1.0); + assert_batch_matches_scalar(&arena, root, 1, k, &ts, &y_cols, &[]); +} + +#[test] +#[should_panic(expected = "eval_batch needs n_states = 2 values for each of 2 lanes")] +fn a_state_matrix_of_the_wrong_length_is_refused() { + // Three values for two lanes used to stride both lanes over one state each. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let expr = CompiledExpr::new(DagRef::new(&arena, y)); + let mut scratch = vec![0.0; expr.scratch_len() * 2]; + let _ = expr.eval_batch(&mut scratch, 2, &[0.0, 0.0], &[1.0, 2.0, 3.0], 2, &[]); +} + +#[test] +fn tangent_load_returns_error() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let two = arena.alloc(Node::Scalar(2.0)); + let root = arena.alloc(Node::Mul(y, two)); + let troot = tangent_wrt_states(&mut arena, root); + let expr = CompiledExpr::new(DagRef::new(&arena, troot)); + let mut scratch = vec![0.0_f64; expr.scratch_len() * 2]; + let err = expr.eval_batch(&mut scratch, 2, &[0.0, 0.0], &[1.0, 2.0], 1, &[]); + assert_eq!(err, Err(BatchEvalError::NonPrimalInstruction)); +} + +/// An `InputParameter` selector, so every lane agrees on the branch. +fn input_selector(arena: &mut Arena, name: &str, index: usize) -> NodeId { + arena.alloc(Node::InputParameter { + name: name.to_string(), + index, + offset: index, + width: 1, + }) +} + +#[test] +fn dispatch_skips_blocks_no_lane_selects() { + // An `InputParameter` selector agrees across lanes, so only one block runs. + let mut arena = Arena::new(); + let selector = input_selector(&mut arena, "s", 0); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let b1 = arena.alloc(Node::Sin(y)); + let mut b2 = y; + for _ in 0..6 { + b2 = arena.alloc(Node::Exp(b2)); + } + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![b1, b2], + }); + let k = 4; + let ts = vec![0.0; k]; + let y_cols = state_cols(1, k, |_, l| 0.1 + l as f64 * 0.05); + for sel in [1.0_f64, 2.0, 0.0] { + assert_batch_matches_scalar(&arena, root, 1, k, &ts, &y_cols, &[sel]); + } +} + +#[test] +fn dispatch_skips_unneeded_block_even_when_it_would_error() { + // Branch 2 would error if evaluated, so passing proves its block is skipped. + let mut arena = Arena::new(); + let selector = input_selector(&mut arena, "s", 0); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let b1 = arena.alloc(Node::Sin(y)); + let y_dot = arena.alloc(Node::StateVectorDot { start: 0, end: 1 }); + let b2 = arena.alloc(Node::Exp(y_dot)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![b1, b2], + }); + let expr = CompiledExpr::new(DagRef::new(&arena, root)); + let k = 4; + let ts = vec![0.0; k]; + let y_cols = state_cols(1, k, |_, l| 0.1 + l as f64 * 0.05); + let mut scratch = vec![0.0_f64; expr.scratch_len() * k]; + let result = expr.eval_batch(&mut scratch, k, &ts, &y_cols, 1, &[1.0]); + assert!( + result.is_ok(), + "branch 2 must never run when no lane selects it: {result:?}" + ); +} + +#[test] +fn two_dispatches_in_one_tape_match_scalar() { + // The union-dispatch loop must resume after the first `Dispatch`. + let mut arena = Arena::new(); + let sel1 = input_selector(&mut arena, "s1", 0); + let sel2 = input_selector(&mut arena, "s2", 1); + let y1 = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y2 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let a1 = arena.alloc(Node::Sin(y1)); + let a2 = arena.alloc(Node::Cos(y1)); + let cond1 = arena.alloc(Node::Conditional { + selector: sel1, + branches: vec![a1, a2], + }); + let b1 = arena.alloc(Node::Exp(y2)); + let b2 = arena.alloc(Node::Sqrt(y2)); + let cond2 = arena.alloc(Node::Conditional { + selector: sel2, + branches: vec![b1, b2], + }); + let root = arena.alloc(Node::Add(cond1, cond2)); + let expr = CompiledExpr::new(DagRef::new(&arena, root)); + assert_eq!( + expr.ir().branch_block_lens().len(), + 4, + "two two-branch dispatches, or this test covers nothing" + ); + let k = 4; + let ts = vec![0.0; k]; + let y_cols = state_cols(2, k, |i, l| match i { + 0 => 0.2 + l as f64 * 0.1, + _ => 0.5 + l as f64 * 0.2, + }); + for sels in [[1.0, 1.0], [2.0, 2.0], [1.0, 2.0], [0.0, 1.0]] { + assert_batch_matches_scalar(&arena, root, 2, k, &ts, &y_cols, &sels); + } +} + +#[test] +fn dispatch_on_boundary_and_nan_selectors_matches_scalar() { + // Half-integer boundaries, NaN, negative and infinite selectors match no branch. + let mut arena = Arena::new(); + let selector = input_selector(&mut arena, "s", 0); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Cos(y)); + let root = arena.alloc(Node::Conditional { + selector, + branches: vec![b1, b2], + }); + let k = 4; + let ts = vec![0.0; k]; + let y_cols = state_cols(1, k, |_, l| 0.1 + l as f64 * 0.05); + for sel in [0.5_f64, 1.5, 2.5, -1.0, f64::NAN, f64::INFINITY] { + assert_batch_matches_scalar(&arena, root, 1, k, &ts, &y_cols, &[sel]); + } + + let lane_selector = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let y1 = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let lb1 = arena.alloc(Node::Sin(y1)); + let lb2 = arena.alloc(Node::Cos(y1)); + let lane_root = arena.alloc(Node::Conditional { + selector: lane_selector, + branches: vec![lb1, lb2], + }); + let weird = [f64::NAN, 0.5, 1.5, 2.5, -1.0, f64::INFINITY, 1.0, 2.0]; + let kw = weird.len(); + let y_cols = state_cols(2, kw, |i, l| if i == 0 { weird[l] } else { 0.3 }); + assert_batch_matches_scalar(&arena, lane_root, 2, kw, &vec![0.0; kw], &y_cols, &[]); +} + +fn alloc_binary(arena: &mut Arena, op: BinaryOp, a: NodeId, b: NodeId) -> NodeId { + arena.alloc(match op { + BinaryOp::Add => Node::Add(a, b), + BinaryOp::Sub => Node::Sub(a, b), + BinaryOp::Mul => Node::Mul(a, b), + BinaryOp::Div => Node::Div(a, b), + BinaryOp::Pow => Node::Pow(a, b), + BinaryOp::Minimum => Node::Minimum(a, b), + BinaryOp::Maximum => Node::Maximum(a, b), + BinaryOp::Modulo => Node::Modulo(a, b), + BinaryOp::Hypot => Node::Hypot(a, b), + BinaryOp::EqualHeaviside => Node::EqualHeaviside(a, b), + BinaryOp::NotEqualHeaviside => Node::NotEqualHeaviside(a, b), + BinaryOp::Equality => Node::Equality(a, b), + }) +} + +fn alloc_unary(arena: &mut Arena, op: UnaryOp, x: NodeId) -> NodeId { + arena.alloc(match op { + UnaryOp::Neg => Node::Neg(x), + UnaryOp::Abs => Node::Abs(x), + UnaryOp::Sqrt => Node::Sqrt(x), + UnaryOp::Exp => Node::Exp(x), + UnaryOp::Log => Node::Log(x), + UnaryOp::Sin => Node::Sin(x), + UnaryOp::Cos => Node::Cos(x), + UnaryOp::Tanh => Node::Tanh(x), + UnaryOp::Sinh => Node::Sinh(x), + UnaryOp::Cosh => Node::Cosh(x), + UnaryOp::Arcsinh => Node::Arcsinh(x), + UnaryOp::Arctan => Node::Arctan(x), + UnaryOp::Erf => Node::Erf(x), + UnaryOp::Sign => Node::Sign(x), + UnaryOp::Floor => Node::Floor(x), + UnaryOp::Ceiling => Node::Ceiling(x), + }) +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_eval.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_eval.rs new file mode 100644 index 0000000000..83545bd09c --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_eval.rs @@ -0,0 +1,52 @@ +mod common; + +use common::cases::{arb_eval_case, targeted_eval_cases}; +use common::numeric_eq::assert_bitwise_eq; +use proptest::prelude::*; +use pybamm_core::{Arena, CompiledExpr, DagRef, NodeId, TypedIr}; + +/// The slot-reusing tape must match the no-reuse (SSA) one bitwise, an independent +/// oracle for the liveness-driven slot allocator. +#[track_caller] +fn assert_reuse_layout_matches_pinned(arena: &Arena, root: NodeId, y: &[f64]) { + let eval = |ir: TypedIr| { + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + expr.eval(&mut s, 0.0, y, &[], &[]).to_vec() + }; + let dag = DagRef::new(arena, root); + assert_bitwise_eq( + &eval(TypedIr::from_dag(dag)), + &eval(TypedIr::from_dag_pinned(dag)), + ); +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(250))] + + #[test] + fn reuse_layout_matches_pinned_layout(case in arb_eval_case()) { + assert_reuse_layout_matches_pinned(&case.arena, case.root, &case.y); + } + + #[test] + fn repeated_eval_on_same_compiled_expr(case in arb_eval_case()) { + let expr = CompiledExpr::new(DagRef::new(&case.arena, case.root)); + let mut s = vec![0.0; expr.scratch_len()]; + + let _ = expr.eval(&mut s, 0.0, &case.y, &[], &[]).to_vec(); + + // Eval with shifted input, then repeat — catches buffer state leaks + let y2: Vec = case.y.iter().map(|v| v + 0.1).collect(); + let first = expr.eval(&mut s, 0.0, &y2, &[], &[]).to_vec(); + let second = expr.eval(&mut s, 0.0, &y2, &[], &[]).to_vec(); + assert_bitwise_eq(&first, &second); + } +} + +#[test] +fn targeted_cases_reuse_layout_matches_pinned_layout() { + for case in targeted_eval_cases() { + assert_reuse_layout_matches_pinned(&case.arena, case.root, &case.y); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.proptest-regressions b/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.proptest-regressions new file mode 100644 index 0000000000..f36b9df85b --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.proptest-regressions @@ -0,0 +1,8 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 9f844515f49dea055f1a6f428f32fa1800f5339048d39418b6e41e46f1093b6e # shrinks to case = DagCase { arena: Arena { nodes: [Scalar(2.0), Scalar(2.0), Sub(NodeId(0), NodeId(1)), Arcsinh(NodeId(2)), Scalar(2.0), Scalar(2.0), Sub(NodeId(4), NodeId(5)), StateVector { start: 0, end: 1 }, StateVector { start: 0, end: 1 }, Sub(NodeId(7), NodeId(8)), Sub(NodeId(6), NodeId(9)), Sin(NodeId(10)), StateVector { start: 0, end: 2 }, StateVector { start: 0, end: 2 }, EqualHeaviside(NodeId(12), NodeId(13)), StateVector { start: 0, end: 2 }, Concat([NodeId(3), NodeId(11), NodeId(14), NodeId(15)])] }, root: NodeId(16), y: [0.1, 0.1], n_states: 2 } +cc 98f136e05db4b8b6f59d73cfb7396ac04a85dddc73f0154b004b8a4bc6f75f57 # shrinks to case = DagCase { arena: Arena { nodes: [Scalar(0.0), StateVector { start: 0, end: 2 }, Scalar(3.0), Pow(NodeId(1), NodeId(2)), StateVector { start: 0, end: 2 }, Add(NodeId(3), NodeId(4)), StateVector { start: 0, end: 2 }, Scalar(3.0), Pow(NodeId(6), NodeId(7)), StateVector { start: 0, end: 2 }, Add(NodeId(8), NodeId(9)), Scalar(-1.0), Cos(NodeId(11)), Concat([NodeId(0), NodeId(5), NodeId(10), NodeId(12)])] }, root: NodeId(13), y: [0.1, 8.442679952532439], n_states: 2, inputs: [] } diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.rs new file mode 100644 index 0000000000..4bc2bb0266 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_rewrite_semantics.rs @@ -0,0 +1,191 @@ +mod common; + +use common::cases::{ + DagCase, Extra, arb_eval_case_with, duplicate_subexpr_case, eval_dag, targeted_eval_cases, +}; +use common::numeric_eq::{assert_bitwise_eq, assert_eq_up_to_zero_sign}; +use proptest::prelude::*; +use pybamm_core::{Arena, ArrayData, DagRef, Node, NodeId, Shape, cse, simplify, zero_propagate}; + +/// Zero literals and `ZeroVector`s, which the passes fold, and literal powers, which +/// `simplify` lowers. +const ZEROS_AND_POWERS: &[Extra] = &[Extra::Zeros, Extra::Pow]; + +fn eval(case: &DagCase) -> Vec { + eval_dag(&case.arena, case.root, 0.0, &case.y, &case.inputs) +} + +/// Whether every node `case.root` reaches evaluates finite at the case's point. +fn every_intermediate_finite(case: &DagCase) -> bool { + case.arena + .topological_order(case.root) + .into_iter() + .all(|id| { + eval_dag(&case.arena, id, 0.0, &case.y, &case.inputs) + .iter() + .all(|v| v.is_finite()) + }) +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(250))] + + /// CSE merges only bitwise-identical nodes, so it is exact, signed zeros and all. + #[test] + fn cse_is_bitwise_exact(case in arb_eval_case_with(ZEROS_AND_POWERS)) { + let shared = cse(DagRef::new(&case.arena, case.root)); + let after = eval_dag(shared.arena(), shared.root(), 0.0, &case.y, &case.inputs); + assert_bitwise_eq(&eval(&case), &after); + } + + /// Its documented contract: exact up to the sign of zero where every intermediate is finite. + #[test] + fn zero_propagate_is_exact_up_to_the_sign_of_zero( + case in arb_eval_case_with(ZEROS_AND_POWERS), + ) { + prop_assume!(every_intermediate_finite(&case)); + let folded = zero_propagate(DagRef::new(&case.arena, case.root)); + let after = eval_dag(folded.arena(), folded.root(), 0.0, &case.y, &case.inputs); + assert_eq_up_to_zero_sign(&eval(&case), &after); + } + + /// Exact up to the sign of zero where every intermediate is finite, literal powers + /// included: `simplify` lowers one only where the rewrite takes `pow`'s value bitwise. + #[test] + fn simplify_is_exact_up_to_the_sign_of_zero(case in arb_eval_case_with(ZEROS_AND_POWERS)) { + prop_assume!(every_intermediate_finite(&case)); + let mut arena = case.arena.clone(); + let root = simplify(&mut arena, case.root); + let after = eval_dag(&arena, root, 0.0, &case.y, &case.inputs); + assert_eq_up_to_zero_sign(&eval(&case), &after); + } +} + +/// Graphs whose rewrites once went wrong, each with the points to check them at. +fn targeted_rewrite_cases() -> Vec<(Arena, NodeId, Vec>)> { + let mut cases = Vec::new(); + + // Arrays differing only in the signs of their zeros, which CSE once merged. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let one = arena.alloc(Node::Scalar(1.0)); + let quotients = [[0.0, -0.0, 3.0], [-0.0, 0.0, 3.0]].map(|data| { + let data = ArrayData::try_new(data.to_vec(), Shape::vector(3)).expect("valid array"); + let array = arena.alloc(Node::Array(Box::new(data))); + let product = arena.alloc(Node::Mul(array, y)); + arena.alloc(Node::Div(one, product)) + }); + let root = arena.alloc(Node::Concat(quotients.to_vec())); + cases.push((arena, root, vec![vec![1.0, 1.0, 1.0], vec![-2.0, 0.5, 4.0]])); + + // A single-branch conditional nested in a privatised cone, which lowering once panicked on. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let selector = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let always = arena.alloc(Node::Scalar(1.0)); + let sin = arena.alloc(Node::Sin(y)); + let inner = arena.alloc(Node::Conditional { + selector: always, + branches: vec![sin], + }); + let branches = vec![ + arena.alloc(Node::Neg(y)), + arena.alloc(Node::Exp(inner)), + arena.alloc(Node::Cos(inner)), + ]; + let root = arena.alloc(Node::Conditional { selector, branches }); + let points = [1.0, 2.0, 3.0, 0.5, f64::NAN].map(|selector| vec![0.7, selector]); + cases.push((arena, root, points.to_vec())); + + // A conditional of zeros is zero for every selector, NaN included. + let mut arena = Arena::new(); + let selector = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let zeros = arena.alloc(Node::ZeroVector { len: 2 }); + let cond = arena.alloc(Node::Conditional { + selector, + branches: vec![zeros, zeros], + }); + let y = arena.alloc(Node::StateVector { start: 1, end: 3 }); + let root = arena.alloc(Node::Add(cond, y)); + let points = [1.0, 2.0, 0.0, f64::NAN].map(|selector| vec![selector, -1.5, 0.25]); + cases.push((arena, root, points.to_vec())); + + cases +} + +#[test] +fn rewrite_passes_preserve_eval_on_targeted_cases() { + let generated = targeted_eval_cases() + .into_iter() + .map(|case| (case.arena, case.root, vec![case.y])); + for (arena, root, points) in generated.chain(targeted_rewrite_cases()) { + let shared = cse(DagRef::new(&arena, root)); + let folded = zero_propagate(DagRef::new(&arena, root)); + let mut simplified = arena.clone(); + let simplified_root = simplify(&mut simplified, root); + for y in &points { + let original = eval_dag(&arena, root, 0.0, y, &[]); + let after = [ + eval_dag(shared.arena(), shared.root(), 0.0, y, &[]), + eval_dag(folded.arena(), folded.root(), 0.0, y, &[]), + eval_dag(&simplified, simplified_root, 0.0, y, &[]), + ]; + for after in after { + assert_bitwise_eq(&original, &after); + } + } + } +} + +#[test] +fn cse_reduces_known_duplicate_graph() { + let case = duplicate_subexpr_case(); + let shared = cse(DagRef::new(&case.arena, case.root)); + assert!( + shared.arena().len() < case.arena.len(), + "CSE failed to reduce known duplicates: {} -> {}", + case.arena.len(), + shared.arena().len(), + ); +} + +/// Exact up to the sign of zero where every intermediate is finite; at a `0 / 0` the fold +/// gives the zero `CasADi` would, where evaluation gives NaN. +#[test] +fn zero_propagate_is_exact_where_every_intermediate_is_finite() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let selector = arena.alloc(Node::StateVector { start: 2, end: 3 }); + let zero = arena.alloc(Node::ZeroVector { len: 2 }); + let zero_matrix = ArrayData::try_new(vec![0.0; 4], Shape::matrix(2, 2)).expect("valid array"); + let zero_matrix = arena.alloc(Node::Array(Box::new(zero_matrix))); + let quotient = arena.alloc(Node::Div(zero, y)); + let product = arena.alloc(Node::MatMul(zero_matrix, y)); + let branch = arena.alloc(Node::Conditional { + selector, + branches: vec![zero, quotient], + }); + let terms = [quotient, product, branch].map(|term| arena.alloc(Node::Add(term, y))); + let root = arena.alloc(Node::Concat(terms.to_vec())); + + let folded = zero_propagate(DagRef::new(&arena, root)); + assert!( + folded.arena().len() < arena.len(), + "every zero term must fold away" + ); + for state in [[0.5, -2.0, 1.0], [3.0, 0.25, 2.0], [-1.0, 1.0, 1.5]] { + let original = eval_dag(&arena, root, 0.0, &state, &[]); + let after = eval_dag(folded.arena(), folded.root(), 0.0, &state, &[]); + assert_eq_up_to_zero_sign(&original, &after); + } + + let singular = [0.0, 1.0, 2.0]; + let original = eval_dag(&arena, root, 0.0, &singular, &[]); + let after = eval_dag(folded.arena(), folded.root(), 0.0, &singular, &[]); + assert!(original[0].is_nan(), "0 / 0 + 0 evaluates to NaN"); + assert_eq!( + after[0].to_bits(), + 0.0_f64.to_bits(), + "the fold leaves y = 0" + ); +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.proptest-regressions b/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.proptest-regressions new file mode 100644 index 0000000000..6feff31ab2 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.proptest-regressions @@ -0,0 +1,9 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 3cbce9a9d1b06125e5f5ffec4e785214e6278679d98aab41ffba0925fad13d52 # shrinks to case = TangentCase { arena: Arena { nodes: [StateVector { start: 0, end: 5 }, StateVector { start: 0, end: 5 }, Neg(NodeId(1)), Tanh(NodeId(2)), EqualHeaviside(NodeId(0), NodeId(3)), Index { child: NodeId(4), start: 2, end: 4 }] }, root: NodeId(5), y: [0.1, 0.1, 0.1, 0.1, 0.1], seeds: [[1.0, 0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, -1.0, 1.0, -1.0, 1.0], [-1.0, -1.0, -1.0, -1.0, -1.0]], n_states: 5 } +cc f4b85e8b1c0824b4a35cc97b1c38db3de0b8eaec65bbf677203beddc9aaad31b # shrinks to case = TangentCase { arena: Arena { nodes: [StateVector { start: 0, end: 4 }, Abs(NodeId(0)), StateVector { start: 0, end: 4 }, EqualHeaviside(NodeId(1), NodeId(2)), Index { child: NodeId(3), start: 2, end: 4 }] }, root: NodeId(4), y: [0.1, 0.1, 0.1, 0.1], seeds: [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, -1.0, 1.0, -1.0], [-0.4893617021276595, -1.0, -0.39574468085106385, 0.5063829787234041]], n_states: 4 } +cc dc5ce2ed538db747c102ea1eed62013153aeb0a83f15896620df7515e8690fdf # shrinks to case = TangentCase { arena: Arena { nodes: [StateVector { start: 0, end: 4 }, StateVector { start: 0, end: 4 }, Add(NodeId(0), NodeId(1)), StateVector { start: 0, end: 4 }, StateVector { start: 0, end: 1 }, EqualHeaviside(NodeId(3), NodeId(4)), StateVector { start: 0, end: 1 }, Concat([NodeId(2), NodeId(5), NodeId(6)]), Sinh(NodeId(7))] }, root: NodeId(8), y: [0.1, 0.1, 0.1, 0.1], seeds: [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, -1.0, 1.0, -1.0], [-0.24110671936758896, -0.9762845849802372, -0.6521739130434782, -1.0]], n_states: 4 } diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.rs new file mode 100644 index 0000000000..c948535c7b --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_split_eval.rs @@ -0,0 +1,217 @@ +mod common; + +use common::cases::{TangentCase, arb_split_eval_case, targeted_split_eval_cases}; +use common::numeric_eq::assert_bitwise_eq; +use proptest::prelude::*; +use pybamm_core::{ + Arena, CompiledExpr, CsrData, DagRef, Node, NodeId, Shape, SplitCompiledExpr, TangentInputs, + TypedIr, cse, dce, simplify, tangent_wrt_states, zero_propagate, +}; + +/// Every seed's sweep on one cached primal matches a monolithic tangent tape. +fn check_split_eval_matches_monolithic(case: &TangentCase) { + let mut arena = case.arena.clone(); + let tangent_root = tangent_wrt_states(&mut arena, case.root); + + let mono = CompiledExpr::new(DagRef::new(&arena, tangent_root)); + let mut s_mono = vec![0.0; mono.scratch_len()]; + let split = SplitCompiledExpr::compile(DagRef::new(&arena, tangent_root)); + let mut s_split = vec![0.0; split.scratch_len()]; + let mut cache = split.eval_primal(&mut s_split, 0.0, &case.y, &[], &[]); + + for seed in &case.seeds { + let tangent_inputs = TangentInputs { + dy: Some(seed), + dp: None, + }; + let mono_result = + mono.eval_with_tangent(&mut s_mono, 0.0, &case.y, &[], &[], &tangent_inputs); + assert_bitwise_eq(mono_result, cache.eval_tangent(&tangent_inputs)); + } +} + +fn check_repeated_tangent_stable(case: &TangentCase) { + let mut arena = case.arena.clone(); + let tangent_root = tangent_wrt_states(&mut arena, case.root); + + let compiled = SplitCompiledExpr::compile(DagRef::new(&arena, tangent_root)); + let mut s = vec![0.0; compiled.scratch_len()]; + let mut cache = compiled.eval_primal(&mut s, 0.0, &case.y, &[], &[]); + + for seed in &case.seeds { + let tangent_inputs = TangentInputs { + dy: Some(seed), + dp: None, + }; + + let first = cache.eval_tangent(&tangent_inputs).to_vec(); + let second = cache.eval_tangent(&tangent_inputs).to_vec(); + assert_bitwise_eq(&first, &second); + } +} + +fn check_split_eval_after_pipeline(case: &TangentCase) { + let mut arena = case.arena.clone(); + let root = tangent_wrt_states(&mut arena, case.root); + let root = simplify(&mut arena, root); + let folded = zero_propagate(DagRef::new(&arena, root)); + let shared = cse(folded.as_ref()); + let pruned = dce(shared.as_ref()); + + for seed in &case.seeds { + let tangent_inputs = TangentInputs { + dy: Some(seed), + dp: None, + }; + + let ir_mono = TypedIr::from_dag(pruned.as_ref()); + let mono = CompiledExpr::from_ir(ir_mono); + let mut s_mono = vec![0.0; mono.scratch_len()]; + let mono_result = mono + .eval_with_tangent(&mut s_mono, 0.0, &case.y, &[], &[], &tangent_inputs) + .to_vec(); + + let split = SplitCompiledExpr::compile(pruned.as_ref()); + let mut s_split = vec![0.0; split.scratch_len()]; + let mut cache = split.eval_primal(&mut s_split, 0.0, &case.y, &[], &[]); + let split_result = cache.eval_tangent(&tangent_inputs).to_vec(); + + assert_bitwise_eq(&mono_result, &split_result); + } +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(200))] + + #[test] + fn split_eval_matches_monolithic(case in arb_split_eval_case()) { + check_split_eval_matches_monolithic(&case); + } + + #[test] + fn split_eval_repeated_tangent_is_stable(case in arb_split_eval_case()) { + check_repeated_tangent_stable(&case); + } + + #[test] + fn split_eval_after_production_pipeline(case in arb_split_eval_case()) { + check_split_eval_after_pipeline(&case); + } +} + +/// Unit seeds over `n` states, then `extra` mixed ones. +fn unit_seeds(n: usize, extra: &[&[f64]]) -> Vec> { + let units = (0..n).map(|i| (0..n).map(|j| if i == j { 1.0 } else { 0.0 }).collect()); + units + .chain(extra.iter().map(|seed| seed.to_vec())) + .collect() +} + +/// A `TangentCase` for `build`'s graph over `y.len()` states. +fn case( + y: Vec, + seeds: Vec>, + build: impl FnOnce(&mut Arena) -> NodeId, +) -> TangentCase { + let mut arena = Arena::new(); + let root = build(&mut arena); + TangentCase { + arena, + root, + n_states: y.len(), + y, + seeds, + inputs: Vec::new(), + } +} + +/// Hand-built graphs: products, transcendentals, a stencil and a sparse product. +fn hand_built_split_eval_cases() -> Vec { + let y_squared = case( + vec![2.0, 3.0, 4.0], + unit_seeds(3, &[&[1.0, 1.0, 1.0], &[0.3, -0.7, 1.2]]), + |arena| { + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + arena.alloc(Node::Mul(y, y)) + }, + ); + let exp_sin = case( + vec![0.1, 0.5, 1.0], + unit_seeds(3, &[&[0.5, -0.3, 0.8]]), + |arena| { + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let exp_y = arena.alloc(Node::Exp(y)); + let sin_y = arena.alloc(Node::Sin(y)); + arena.alloc(Node::Mul(exp_y, sin_y)) + }, + ); + let n = 10; + let tridiagonal = case( + (0..n).map(|i| (i as f64 * 0.3).sin()).collect(), + unit_seeds(n, &[]), + |arena| { + let states: Vec<_> = (0..n) + .map(|i| { + arena.alloc(Node::StateVector { + start: i, + end: i + 1, + }) + }) + .collect(); + let two = arena.alloc(Node::Scalar(2.0)); + let rows = (0..n) + .map(|i| { + let two_mid = arena.alloc(Node::Mul(two, states[i])); + let left = states[i.saturating_sub(1)]; + let right = states[(i + 1).min(n - 1)]; + let sum_lr = arena.alloc(Node::Add(left, right)); + arena.alloc(Node::Add(two_mid, sum_lr)) + }) + .collect(); + arena.alloc(Node::Concat(rows)) + }, + ); + let sparse_matmul = case( + vec![0.1, 0.2, 0.3, 0.4], + unit_seeds(4, &[&[1.0, 0.0, 1.0, 0.0], &[0.2, -0.4, 0.6, -0.8]]), + |arena| { + let y = arena.alloc(Node::StateVector { start: 0, end: 4 }); + let mat = arena.alloc(Node::SparseMatrix(Box::new( + CsrData::try_new( + vec![0, 2, 5, 8, 10], + vec![0, 1, 0, 1, 2, 1, 2, 3, 2, 3], + vec![1.0; 10], + Shape::matrix(4, 4), + ) + .expect("valid test matrix"), + ))); + arena.alloc(Node::MatMul(mat, y)) + }, + ); + // Many seeds against one cached primal. + let scaled_exp = case( + vec![0.1, 0.2, 0.3, 0.4, 0.5], + (0..20_u32) + .map(|s| (0..5_u32).map(|i| (f64::from(s + i) * 0.7).sin()).collect()) + .collect(), + |arena| { + let y = arena.alloc(Node::StateVector { start: 0, end: 5 }); + let exp_y = arena.alloc(Node::Exp(y)); + let two = arena.alloc(Node::Scalar(2.0)); + arena.alloc(Node::Mul(two, exp_y)) + }, + ); + vec![y_squared, exp_sin, tridiagonal, sparse_matmul, scaled_exp] +} + +#[test] +fn targeted_split_eval_cases_pass() { + for case in targeted_split_eval_cases() + .into_iter() + .chain(hand_built_split_eval_cases()) + { + check_split_eval_matches_monolithic(&case); + check_repeated_tangent_stable(&case); + check_split_eval_after_pipeline(&case); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.proptest-regressions b/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.proptest-regressions new file mode 100644 index 0000000000..fb82b7cc26 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.proptest-regressions @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc fc6afd0373771469220a9e47b281f78c5e79208224feb66daf251380e3107990 # shrinks to case = TangentCase { arena: Arena { nodes: [Scalar(0.5), Neg(NodeId(0)), Arctan(NodeId(1))] }, root: NodeId(2), y: [0.1, 0.1], seeds: [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [1.0, -1.0], [-0.2784313725490196, -1.0]], n_states: 2 } diff --git a/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.rs b/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.rs new file mode 100644 index 0000000000..325e419690 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/proptest_tangent_batch.rs @@ -0,0 +1,131 @@ +mod common; + +use common::cases::{ + ALL_EXTRAS, TangentCase, arb_split_eval_case, arb_split_eval_case_with, + targeted_split_eval_cases, +}; +use common::numeric_eq::assert_bitwise_eq; +use proptest::prelude::*; +use pybamm_core::tangent_batch::{is_batchable, run_tangent_batch, tangent_scratch_len}; +use pybamm_core::{Arena, DagRef, Node, SplitCompiledExpr, TangentInputs, tangent_wrt_states}; + +/// A batched tangent sweep reproduces its scalar sweeps bit for bit, at every lane count. +fn check_batched_matches_scalar(case: &TangentCase) { + for lanes in [4, 8] { + check_lanes(case, lanes); + } +} + +fn check_lanes(case: &TangentCase, lanes: usize) { + let mut arena = case.arena.clone(); + let tangent_root = tangent_wrt_states(&mut arena, case.root); + let compiled = SplitCompiledExpr::compile(DagRef::new(&arena, tangent_root)); + if !is_batchable(&compiled) { + return; + } + let primal_len = compiled.split_info().primal_buffer_size; + let n_states = case.y.len(); + + // Fill every lane, reusing the case's seeds rescaled, so no lane rides along as zeros. + let seeds: Vec> = (0..lanes) + .map(|lane| match case.seeds.len() { + 0 => vec![0.0; n_states], + len => { + let scale = (lane / len + 1) as f64; + case.seeds[lane % len].iter().map(|v| v * scale).collect() + }, + }) + .collect(); + + let mut scratch = vec![0.0; compiled.scratch_len()]; + + let mut cache = compiled.eval_primal(&mut scratch, 0.0, &case.y, &[], &case.inputs); + let scalar: Vec> = seeds + .iter() + .map(|seed| { + cache + .eval_tangent(&TangentInputs { + dy: Some(seed), + dp: None, + }) + .to_vec() + }) + .collect(); + + let mut lane_seeds = vec![0.0; n_states * lanes]; + for (lane, seed) in seeds.iter().enumerate() { + for (state, &value) in seed.iter().enumerate() { + lane_seeds[state * lanes + lane] = value; + } + } + + // Scalar sweeps write only tangent slots, so the primal region is still this point's. + let mut tan = vec![0.0; tangent_scratch_len(&compiled, lanes)]; + let primal = &scratch[..primal_len]; + let batched = run_tangent_batch(&compiled, lanes, primal, &mut tan, &lane_seeds); + + for (lane, want) in scalar.iter().enumerate() { + let got: Vec = (0..want.len()).map(|e| batched[e * lanes + lane]).collect(); + assert_bitwise_eq(want, &got); + } +} + +#[test] +fn targeted_batched_tangent_matches_scalar() { + for case in targeted_split_eval_cases() { + check_batched_matches_scalar(&case); + } +} + +/// Zero-width loads, slices and concatenations sit at offset 0, inside the primal region. +#[test] +fn zero_width_tangent_nodes_sweep_batched() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let sin = arena.alloc(Node::Sin(y)); + let empty_state = arena.alloc(Node::StateVector { start: 2, end: 2 }); + let empty_slice = arena.alloc(Node::Index { + child: sin, + start: 1, + end: 1, + }); + let empty_concat = arena.alloc(Node::Concat(vec![empty_state])); + let root = arena.alloc(Node::Concat(vec![ + sin, + empty_state, + empty_slice, + empty_concat, + ])); + let case = TangentCase { + arena, + root, + y: vec![0.3, 0.7], + seeds: vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![0.5, -2.0]], + n_states: 2, + inputs: Vec::new(), + }; + + let mut tangent_arena = case.arena.clone(); + let tangent_root = tangent_wrt_states(&mut tangent_arena, case.root); + let compiled = SplitCompiledExpr::compile(DagRef::new(&tangent_arena, tangent_root)); + assert!( + is_batchable(&compiled), + "the case must take the batched sweep" + ); + check_batched_matches_scalar(&case); +} + +proptest! { + #[test] + fn batched_tangent_matches_scalar(case in arb_split_eval_case()) { + check_batched_matches_scalar(&case); + } + + /// Steps, interpolants and reductions, whose tangents fold to primal roots and constants. + #[test] + fn batched_tangent_matches_scalar_with_every_extra( + case in arb_split_eval_case_with(ALL_EXTRAS), + ) { + check_batched_matches_scalar(&case); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/test_branch_shortcircuit.rs b/packages/pybamm-rust/pybamm-core/tests/test_branch_shortcircuit.rs new file mode 100644 index 0000000000..cf5ab1c75e --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_branch_shortcircuit.rs @@ -0,0 +1,724 @@ +//! Only-active-branch execution, asserted from the executed instruction count. + +use pybamm_core::arena::{Arena, NodeId}; +use pybamm_core::eval::{CompiledExpr, TangentInputs}; +use pybamm_core::ir::{Instruction, TypedIr}; +use pybamm_core::node::{CsrData, Node, Shape}; +use pybamm_core::{DagRef, SplitCompiledExpr, simplify, tangent_wrt_states}; + +/// `cond(sel, [chain(y, 1), chain(y, 8), chain(y, 16)])` over an `InputParameter` selector. +fn uneven_branches() -> (Arena, NodeId, Vec) { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let mut branches = Vec::new(); + let mut lens = Vec::new(); + for depth in [1_usize, 8, 16] { + let mut node = y; + for _ in 0..depth { + node = arena.alloc(Node::Sin(node)); + } + branches.push(node); + lens.push(depth); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches, + }); + (arena, cond, lens) +} + +#[test] +fn only_the_active_branch_executes() { + let (arena, root, lens) = uneven_branches(); + let expr = CompiledExpr::new(DagRef::new(&arena, root)); + let mut scratch = vec![0.0; expr.scratch_len()]; + + // common = y load + selector load + Conditional + Dispatch = 4. + let common = 4; + assert_eq!(expr.ir().common_instruction_count(), common); + assert_eq!( + expr.ir().instructions().len(), + common + lens.iter().sum::() + ); + for (i, &len) in lens.iter().enumerate() { + let sel = (i + 1) as f64; + let no_seed = TangentInputs::default(); + let (executed, _) = + expr.eval_counted_with_tangent(&mut scratch, 0.0, &[0.3], &[], &[sel], &no_seed); + assert_eq!( + executed, + common + len, + "selector {sel} executed {executed} instructions, expected {}", + common + len + ); + } + + // No match: the dispatch runs, no block does. + let no_seed = TangentInputs::default(); + let (executed, out) = + expr.eval_counted_with_tangent(&mut scratch, 0.0, &[0.3], &[], &[0.0], &no_seed); + assert_eq!(executed, common); + assert_eq!(out, &[0.0]); +} + +#[test] +fn executed_count_is_independent_of_the_other_branches() { + let cost_of_branch_one = |depth_of_last: usize| { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b1 = arena.alloc(Node::Sin(y)); + let mut b2 = y; + for _ in 0..depth_of_last { + b2 = arena.alloc(Node::Cos(b2)); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + let expr = CompiledExpr::new(DagRef::new(&arena, cond)); + let mut scratch = vec![0.0; expr.scratch_len()]; + let no_seed = TangentInputs::default(); + expr.eval_counted_with_tangent(&mut scratch, 0.0, &[0.3], &[], &[1.0], &no_seed) + .0 + }; + assert_eq!(cost_of_branch_one(2), cost_of_branch_one(64)); +} + +#[test] +fn short_circuited_results_are_bitwise_identical() { + let (arena, root, _) = uneven_branches(); + let short = CompiledExpr::new(DagRef::new(&arena, root)); + let mut s = vec![0.0; short.scratch_len()]; + for sel in [ + 0.0_f64, + 0.49, + 0.5, + 0.51, + 1.0, + 1.49, + 1.5, + 2.0, + 2.5, + 3.0, + 3.49, + 3.5, + 4.0, + -1.0, + f64::NAN, + f64::INFINITY, + ] { + let got = short.eval(&mut s, 0.0, &[0.3], &[], &[sel])[0]; + let expected = expected_conditional(0.3, sel); + assert_eq!( + got.to_bits(), + expected.to_bits(), + "selector {sel}: got {got}, expected {expected}" + ); + } +} + +/// The semantics contract, written out independently of the evaluator. +fn expected_conditional(y: f64, sel: f64) -> f64 { + branch_window(sel, 3).map_or(0.0, |i| { + let mut v = y; + for _ in 0..[1_usize, 8, 16][i] { + v = v.sin(); + } + v + }) +} + +#[test] +fn nested_conditionals_degrade_without_miscompiling() { + // The inner cone is forced common, so only the outer conditional short-circuits. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let outer_sel = arena.alloc(Node::InputParameter { + name: "o".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let inner_sel = arena.alloc(Node::InputParameter { + name: "i".to_string(), + index: 1, + offset: 1, + width: 1, + }); + let ia = arena.alloc(Node::Sin(y)); + let ib = arena.alloc(Node::Cos(y)); + let inner = arena.alloc(Node::Conditional { + selector: inner_sel, + branches: vec![ia, ib], + }); + let ob = arena.alloc(Node::Neg(y)); + let outer = arena.alloc(Node::Conditional { + selector: outer_sel, + branches: vec![inner, ob], + }); + + let ir = TypedIr::from_dag(DagRef::new(&arena, outer)); + assert_eq!( + ir.dispatch_count(), + 1, + "only the outer conditional is blockable" + ); + assert_eq!(ir.branch_block_lens(), vec![1, 1]); + + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + let y0 = 0.3_f64; + for (o, i, expected) in [ + (1.0, 1.0, y0.sin()), + (1.0, 2.0, y0.cos()), + (1.0, 0.0, 0.0), + (2.0, 1.0, -y0), + (0.0, 1.0, 0.0), + ] { + let got = expr.eval(&mut s, 0.0, &[y0], &[], &[o, i])[0]; + assert_eq!(got.to_bits(), expected.to_bits(), "outer {o} inner {i}"); + } +} + +#[test] +fn a_branch_owning_no_nodes_gets_an_empty_block() { + // Branch 1 is a bare shared state load (common), branch 2 owns one node. + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b2 = arena.alloc(Node::Sin(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![y, b2], + }); + let ir = TypedIr::from_dag(DagRef::new(&arena, cond)); + assert_eq!(ir.branch_block_lens(), vec![0, 1]); + + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + assert_eq!(expr.eval(&mut s, 0.0, &[0.3], &[], &[1.0]), &[0.3]); + assert_eq!( + expr.eval(&mut s, 0.0, &[0.3], &[], &[2.0]), + &[0.3_f64.sin()] + ); +} + +/// Reuse-slot and pinned layouts vs a closed-form oracle on a graph mixing every block hazard. +#[test] +fn scheduled_tape_matches_the_pinned_tape_and_a_closed_form_oracle() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 3 }); + let sel = arena.alloc(Node::StateVector { start: 3, end: 4 }); + let shared_inner = arena.alloc(Node::Exp(y)); + let shared = arena.alloc(Node::Sqrt(shared_inner)); + + let b0 = arena.alloc(Node::Sin(shared)); + let csr = arena.alloc(Node::SparseMatrix( + CsrData::try_new( + vec![0, 1, 2, 3], + vec![1, 2, 0], + vec![2.0, -1.5, 0.75], + Shape { rows: 3, cols: 3 }, + ) + .expect("valid csr") + .into(), + )); + let mm = arena.alloc(Node::MatMul(csr, shared)); + let b1 = arena.alloc(Node::Cos(mm)); + + let inner_sel = arena.alloc(Node::Time); + let ia = arena.alloc(Node::Tanh(y)); + let ib = arena.alloc(Node::Abs(y)); + let b2 = arena.alloc(Node::Conditional { + selector: inner_sel, + branches: vec![ia, ib], + }); + + let outer = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b0, b1, b2], + }); + // `shared_inner` is also read outside the conditional, so it must stay `Common`. + let root = arena.alloc(Node::Add(outer, shared_inner)); + + let scheduled = CompiledExpr::new(DagRef::new(&arena, root)); + let pinned = CompiledExpr::from_ir(TypedIr::from_dag_pinned(DagRef::new(&arena, root))); + // Block lengths count instructions, not nodes: branch 1's `SparseMatrix` emits none. + assert_eq!(scheduled.ir().dispatch_count(), 1); + assert_eq!(scheduled.ir().branch_block_lens(), vec![2, 3, 1]); + assert_eq!(pinned.ir().dispatch_count(), 1); + assert_eq!(pinned.ir().branch_block_lens(), vec![2, 3, 1]); + + let mut sched_scratch = vec![0.0; scheduled.scratch_len()]; + let mut pinned_scratch = vec![0.0; pinned.scratch_len()]; + for sel_val in [0.0_f64, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, f64::NAN] { + for t in [1.0_f64, 2.0, 0.0] { + let state = [0.35, -0.2, 0.9, sel_val]; + let got = scheduled.eval(&mut sched_scratch, t, &state, &[], &[]); + let want = pinned.eval(&mut pinned_scratch, t, &state, &[], &[]); + let oracle = expected_hazard_graph([state[0], state[1], state[2]], sel_val, t); + assert_eq!(got.len(), want.len()); + assert_eq!(got.len(), oracle.len()); + for (i, ((g, w), o)) in got.iter().zip(want).zip(&oracle).enumerate() { + assert_eq!( + g.to_bits(), + w.to_bits(), + "selector {sel_val}, t {t}, element {i}: scheduled {g}, pinned {w}" + ); + assert!( + (g - o).abs() <= 1e-15 * o.abs(), + "selector {sel_val}, t {t}, element {i}: got {g}, oracle {o}" + ); + } + } + } +} + +/// The hazard graph's semantics, written out independently of the evaluator. +fn expected_hazard_graph(y: [f64; 3], sel: f64, t: f64) -> Vec { + let exp_y = y.map(f64::exp); + let shared = exp_y.map(f64::sqrt); + let outer = match branch_window(sel, 3) { + Some(0) => shared.map(f64::sin), + Some(1) => [2.0 * shared[1], -1.5 * shared[2], 0.75 * shared[0]].map(f64::cos), + Some(2) => match branch_window(t, 2) { + Some(0) => y.map(f64::tanh), + Some(1) => y.map(f64::abs), + _ => [0.0; 3], + }, + _ => [0.0; 3], + }; + (0..3).map(|i| outer[i] + exp_y[i]).collect() +} + +/// The 1-based round-to-nearest branch window of the semantics contract. +fn branch_window(selector: f64, n_branches: usize) -> Option { + (0..n_branches).find(|&i| { + let idx = (i + 1) as f64; + selector > idx - 0.5 && selector < idx + 0.5 + }) +} + +/// Lanes select different branches, so every dispatched block stays live. +#[test] +fn batch_eval_over_dispatched_blocks_matches_per_lane_eval() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + // A state-derived selector: `InputParameter` broadcasts to every lane. + let sel = arena.alloc(Node::StateVector { start: 1, end: 2 }); + let mut branches = Vec::new(); + for depth in [1_usize, 4, 9] { + let mut node = y; + for _ in 0..depth { + node = arena.alloc(Node::Sin(node)); + } + branches.push(node); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches, + }); + + let expr = CompiledExpr::new(DagRef::new(&arena, cond)); + assert_eq!( + expr.ir().branch_block_lens(), + vec![1, 4, 9], + "the tape must actually carry blocks for this to test anything" + ); + + // Four lanes: branch 1, branch 3, branch 2, and no match. + let k = 4; + let ts = vec![0.0; k]; + let y_cols = [0.3, 1.0, 0.4, 3.0, 0.5, 2.0, 0.6, 0.0]; + let n_states = 2; + + let mut per_lane = vec![0.0_f64; k]; + let mut s = vec![0.0; expr.scratch_len()]; + for l in 0..k { + let y_l = &y_cols[l * n_states..(l + 1) * n_states]; + per_lane[l] = expr.eval(&mut s, ts[l], y_l, &[], &[])[0]; + } + + let mut batch_scratch = vec![0.0; expr.scratch_len() * k]; + let batched = expr + .eval_batch(&mut batch_scratch, k, &ts, &y_cols, n_states, &[]) + .expect("primal tape batches"); + for l in 0..k { + assert_eq!( + batched[l].to_bits(), + per_lane[l].to_bits(), + "lane {l}: batch {} vs scalar {}", + batched[l], + per_lane[l] + ); + } +} + +/// Two sibling top-level conditionals, the shape `PyBaMM`'s unified experiment model builds. +#[test] +fn two_independent_conditionals_get_one_dispatch_each() { + const DEPTHS_A: [usize; 2] = [1, 3]; + const DEPTHS_B: [usize; 2] = [2, 5]; + + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel_a = arena.alloc(Node::InputParameter { + name: "a".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let sel_b = arena.alloc(Node::InputParameter { + name: "b".to_string(), + index: 1, + offset: 1, + width: 1, + }); + let mut a_branches = Vec::new(); + for depth in DEPTHS_A { + let mut node = y; + for _ in 0..depth { + node = arena.alloc(Node::Sin(node)); + } + a_branches.push(node); + } + let mut b_branches = Vec::new(); + for depth in DEPTHS_B { + let mut node = y; + for _ in 0..depth { + node = arena.alloc(Node::Cos(node)); + } + b_branches.push(node); + } + let cond_a = arena.alloc(Node::Conditional { + selector: sel_a, + branches: a_branches, + }); + let cond_b = arena.alloc(Node::Conditional { + selector: sel_b, + branches: b_branches, + }); + let root = arena.alloc(Node::Add(cond_a, cond_b)); + + let ir = TypedIr::from_dag(DagRef::new(&arena, root)); + assert_eq!(ir.dispatch_count(), 2); + // `b`'s dispatch is emitted first: the topological walk reaches its cone first. + assert_eq!(ir.branch_block_lens(), vec![2, 5, 1, 3]); + // y + two selectors + two Conditionals + Add + two Dispatches. + let common = 8; + assert_eq!(ir.common_instruction_count(), common); + assert_eq!(ir.instructions().len(), 19); + + let expr = CompiledExpr::from_ir(ir); + let mut scratch = vec![0.0; expr.scratch_len()]; + let y0 = 0.3_f64; + for (a_index, sel_a_val) in [(None, 0.0_f64), (Some(0), 1.0), (Some(1), 2.0), (None, 3.0)] { + for (b_index, sel_b_val) in [(None, 0.0_f64), (Some(0), 1.0), (Some(1), 2.0), (None, 3.0)] { + let a_cost = a_index.map_or(0, |i| DEPTHS_A[i]); + let b_cost = b_index.map_or(0, |i| DEPTHS_B[i]); + let no_seed = TangentInputs::default(); + let inputs = [sel_a_val, sel_b_val]; + let (executed, out) = + expr.eval_counted_with_tangent(&mut scratch, 0.0, &[y0], &[], &inputs, &no_seed); + assert_eq!( + executed, + common + a_cost + b_cost, + "selectors ({sel_a_val}, {sel_b_val}) executed {executed}" + ); + + let chain = |depth: usize, f: fn(f64) -> f64| { + let mut v = y0; + for _ in 0..depth { + v = f(v); + } + v + }; + let expected = a_index.map_or(0.0, |i| chain(DEPTHS_A[i], f64::sin)) + + b_index.map_or(0.0, |i| chain(DEPTHS_B[i], f64::cos)); + assert_eq!( + out[0].to_bits(), + expected.to_bits(), + "selectors ({sel_a_val}, {sel_b_val}): got {}, expected {expected}", + out[0] + ); + } + } +} + +#[test] +fn jvp_over_dispatched_blocks_matches_the_split_eval_tape() { + let mut arena = Arena::new(); + let x = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + // Nonlinear branches, so each adjoint genuinely needs its own primal value. + let b0 = arena.alloc(Node::Sin(x)); + let b1 = arena.alloc(Node::Exp(x)); + let b2 = arena.alloc(Node::Tanh(x)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b0, b1, b2], + }); + let jac = tangent_wrt_states(&mut arena, cond); + let jac = simplify(&mut arena, jac); + + let scheduled = CompiledExpr::new(DagRef::new(&arena, jac)); + let split = SplitCompiledExpr::compile(DagRef::new(&arena, jac)); + assert_eq!(scheduled.ir().dispatch_count(), 1); + assert_eq!(scheduled.ir().branch_block_lens(), vec![2, 2, 5]); + assert_eq!( + split.ir().dispatch_count(), + 2, + "split-eval splits this cone into one block per partition half" + ); + + let mut sched_scratch = vec![0.0; scheduled.scratch_len()]; + let mut split_scratch = vec![0.0; split.scratch_len()]; + let tangent = TangentInputs { + dy: Some(&[1.0]), + dp: None, + }; + let x0 = 0.4_f64; + for sel_val in [ + 0.0_f64, + 0.5, + 1.0, + 1.5, + 2.0, + 2.5, + 3.0, + 3.5, + 4.0, + -1.0, + f64::NAN, + f64::INFINITY, + ] { + let got = + scheduled.eval_with_tangent(&mut sched_scratch, 0.0, &[x0], &[], &[sel_val], &tangent) + [0]; + let want = + split.eval_with_tangent(&mut split_scratch, 0.0, &[x0], &[], &[sel_val], &tangent)[0]; + assert_eq!( + got.to_bits(), + want.to_bits(), + "selector {sel_val}: scheduled {got}, split {want}" + ); + } + + // Closed form, so the two tapes can't pass by a shared mistake. + for (sel_val, expected) in [ + (1.0_f64, x0.cos()), + (2.0, x0.exp()), + (3.0, x0.tanh().mul_add(-x0.tanh(), 1.0)), + (0.0, 0.0), + ] { + let got = + scheduled.eval_with_tangent(&mut sched_scratch, 0.0, &[x0], &[], &[sel_val], &tangent) + [0]; + assert!( + (got - expected).abs() < 1e-12, + "selector {sel_val}: got {got}, expected {expected}" + ); + } +} + +#[test] +fn split_eval_blocks_respect_the_primal_tangent_partition() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let mut b1 = y; + for _ in 0..3 { + b1 = arena.alloc(Node::Sin(b1)); + } + let mut b2 = y; + for _ in 0..12 { + b2 = arena.alloc(Node::Exp(b2)); + } + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + let tangent_root = tangent_wrt_states(&mut arena, cond); + + let compiled = SplitCompiledExpr::compile(DagRef::new(&arena, tangent_root)); + let (ir, split) = (compiled.ir(), compiled.split_info()); + // One block per half, so the cone really is split rather than degraded. + assert_eq!(ir.dispatch_count(), 2); + + for (i, instr) in ir.instructions().iter().enumerate() { + if let Instruction::Dispatch { + blocks_idx, + blocks_len, + span_len, + .. + } = *instr + { + // The blocks tile the span after the `Dispatch`, which evaluation then skips whole. + let mut next = i + 1; + for b in 0..blocks_len.index() { + let block = ir.consts().branch_blocks[blocks_idx.index() + b].bounds(i); + assert_eq!(block.start, next, "block {b} of the Dispatch at {i}"); + next = block.end; + assert!( + split.primal_end <= block.start || split.primal_end >= block.end, + "primal_end {} splits block [{}, {})", + split.primal_end, + block.start, + block.end + ); + } + assert_eq!( + i + span_len.index(), + next, + "the Dispatch at {i} spans its blocks" + ); + } + } + + let expr = CompiledExpr::from_ir(ir.clone()); + let mut s = vec![0.0; expr.scratch_len()]; + let seed = [1.0, 0.0]; + let tangent = TangentInputs { + dy: Some(&seed), + dp: None, + }; + let mut counts = Vec::new(); + for sel_val in [1.0_f64, 2.0] { + let (executed, _) = + expr.eval_counted_with_tangent(&mut s, 0.0, &[0.3, 0.4], &[], &[sel_val], &tangent); + counts.push(executed); + } + assert!( + counts[0] < counts[1], + "branch 1 ({}) should cost less than branch 2 ({})", + counts[0], + counts[1] + ); +} + +#[test] +fn split_eval_matches_monolithic_for_every_selector() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 2 }); + let sel = arena.alloc(Node::InputParameter { + name: "s".to_string(), + index: 0, + offset: 0, + width: 1, + }); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Exp(y)); + let b3 = arena.alloc(Node::Cos(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2, b3], + }); + let tangent_root = tangent_wrt_states(&mut arena, cond); + + let mono = CompiledExpr::new(DagRef::new(&arena, tangent_root)); + let split = SplitCompiledExpr::compile(DagRef::new(&arena, tangent_root)); + let seed = [1.0, 1.0]; + let tangent = TangentInputs { + dy: Some(&seed), + dp: None, + }; + let mut sm = vec![0.0; mono.scratch_len()]; + let mut ss = vec![0.0; split.scratch_len()]; + for sel_val in [0.0_f64, 1.0, 2.0, 3.0, 4.0, 1.5, f64::NAN] { + let a = mono + .eval_with_tangent(&mut sm, 0.0, &[0.3, 0.4], &[], &[sel_val], &tangent) + .to_vec(); + let mut cache = split.eval_primal(&mut ss, 0.0, &[0.3, 0.4], &[], &[sel_val]); + let b = cache.eval_tangent(&tangent).to_vec(); + assert_eq!(a.len(), b.len()); + for (x, z) in a.iter().zip(&b) { + assert_eq!(x.to_bits(), z.to_bits(), "selector {sel_val}"); + } + } +} + +/// A tangent-tainted selector must degrade: a primal `Dispatch` would read an unwritten slot. +#[test] +fn a_tangent_selector_degrades_rather_than_dispatching_on_an_unwritten_slot() { + let mut arena = Arena::new(); + let y = arena.alloc(Node::StateVector { start: 0, end: 1 }); + let dy = arena.alloc(Node::TangentStateVector { start: 0, end: 1 }); + let one = arena.alloc(Node::Scalar(1.0)); + let sel = arena.alloc(Node::Add(dy, one)); + let b1 = arena.alloc(Node::Sin(y)); + let b2 = arena.alloc(Node::Exp(y)); + let cond = arena.alloc(Node::Conditional { + selector: sel, + branches: vec![b1, b2], + }); + + let mono = CompiledExpr::new(DagRef::new(&arena, cond)); + let split = SplitCompiledExpr::compile(DagRef::new(&arena, cond)); + let mut sm = vec![0.0; mono.scratch_len()]; + let mut ss = vec![0.0; split.scratch_len()]; + let y0 = 0.4_f64; + for seed in [-1.0_f64, 0.0, 0.5, 1.0, 2.0, f64::NAN] { + let seeds = [seed]; + let tangent = TangentInputs { + dy: Some(&seeds), + dp: None, + }; + let want = mono.eval_with_tangent(&mut sm, 0.0, &[y0], &[], &[], &tangent)[0]; + let mut cache = split.eval_primal(&mut ss, 0.0, &[y0], &[], &[]); + let got = cache.eval_tangent(&tangent)[0]; + let expected = match seed + 1.0 { + s if s > 0.5 && s < 1.5 => y0.sin(), + s if s > 1.5 && s < 2.5 => y0.exp(), + _ => 0.0, + }; + assert_eq!(got.to_bits(), want.to_bits(), "seed {seed}"); + assert_eq!(got.to_bits(), expected.to_bits(), "seed {seed}"); + } +} + +#[test] +fn pinned_layout_short_circuits_and_stays_ssa() { + let (arena, root, lens) = uneven_branches(); + let ir = TypedIr::from_dag_pinned(DagRef::new(&arena, root)); + assert_eq!(ir.branch_block_lens(), vec![1, 8, 16]); + // Still SSA (reverse AD needs it): `y`, the selector and the `Conditional` are the 3. + assert_eq!(ir.buffer_size(), 3 + lens.iter().sum::()); + + let expr = CompiledExpr::from_ir(ir); + let mut s = vec![0.0; expr.scratch_len()]; + for (i, &len) in lens.iter().enumerate() { + let sel = (i + 1) as f64; + let no_seed = TangentInputs::default(); + let (executed, out) = + expr.eval_counted_with_tangent(&mut s, 0.0, &[0.3], &[], &[sel], &no_seed); + assert_eq!(executed, 4 + len); + assert_eq!(out[0].to_bits(), expected_conditional(0.3, sel).to_bits()); + } +} diff --git a/packages/pybamm-rust/pybamm-core/tests/test_primitive_values.rs b/packages/pybamm-rust/pybamm-core/tests/test_primitive_values.rs new file mode 100644 index 0000000000..e5182e5c92 --- /dev/null +++ b/packages/pybamm-rust/pybamm-core/tests/test_primitive_values.rs @@ -0,0 +1,435 @@ +//! Every primitive takes its reference value on every evaluation path: `f64`'s for the +//! transcendentals, exact forms for `pow` at `0`, `1`, `2` and `-1`, and `libm`'s for `erf`, +//! `asinh`, `minimum` and `maximum`. Where IEEE 754 specifies the result exactly (`sqrt`, +//! `floor`, `ceil`, `abs`, `%`) the tapes use `f64` and the references `libm`, so agreement +//! checks both against that result. +//! +//! IEEE 754 leaves NaN payloads unspecified, so a NaN need only meet a NaN; every other +//! result, signed zeros and infinities included, must match bit for bit. + +use std::f64::consts::{FRAC_PI_2, PI}; + +use pybamm_core::ir::{BinaryOp, UnaryOp}; +use pybamm_core::{Arena, CompiledExpr, DagRef, Node, NodeId, simplify}; + +type UnaryNode = fn(NodeId) -> Node; +type BinaryNode = fn(NodeId, NodeId) -> Node; +type UnaryFn = fn(f64) -> f64; +type BinaryFn = fn(f64, f64) -> f64; + +const UNARY_OPS: [UnaryOp; 16] = [ + UnaryOp::Neg, + UnaryOp::Abs, + UnaryOp::Sqrt, + UnaryOp::Exp, + UnaryOp::Log, + UnaryOp::Sin, + UnaryOp::Cos, + UnaryOp::Tanh, + UnaryOp::Sinh, + UnaryOp::Cosh, + UnaryOp::Arcsinh, + UnaryOp::Arctan, + UnaryOp::Erf, + UnaryOp::Sign, + UnaryOp::Floor, + UnaryOp::Ceiling, +]; + +const BINARY_OPS: [BinaryOp; 12] = [ + BinaryOp::Add, + BinaryOp::Sub, + BinaryOp::Mul, + BinaryOp::Div, + BinaryOp::Pow, + BinaryOp::Minimum, + BinaryOp::Maximum, + BinaryOp::Modulo, + BinaryOp::Hypot, + BinaryOp::EqualHeaviside, + BinaryOp::NotEqualHeaviside, + BinaryOp::Equality, +]; + +/// `pow` as the tapes define it: the platform's, but `1`, `x`, `x*x` and `1/x` for the +/// exponents `0`, `1`, `2` and `-1`. +#[allow(clippy::float_cmp)] +fn pow(x: f64, y: f64) -> f64 { + if y == 0.0 { + 1.0 + } else if y == 1.0 { + x + } else if y == 2.0 { + x * x + } else if y == -1.0 { + 1.0 / x + } else { + x.powf(y) + } +} + +/// The node that lowers to `op`, and the value it must take. +/// +/// Exhaustive, so a new primitive does not compile until it has a reference here. +fn unary(op: UnaryOp) -> (UnaryNode, UnaryFn) { + match op { + UnaryOp::Neg => (Node::Neg, |x| -x), + UnaryOp::Abs => (Node::Abs, libm::fabs), + UnaryOp::Sqrt => (Node::Sqrt, libm::sqrt), + UnaryOp::Exp => (Node::Exp, f64::exp), + UnaryOp::Log => (Node::Log, f64::ln), + UnaryOp::Sin => (Node::Sin, f64::sin), + UnaryOp::Cos => (Node::Cos, f64::cos), + UnaryOp::Tanh => (Node::Tanh, f64::tanh), + UnaryOp::Sinh => (Node::Sinh, f64::sinh), + UnaryOp::Cosh => (Node::Cosh, f64::cosh), + UnaryOp::Arcsinh => (Node::Arcsinh, libm::asinh), + UnaryOp::Arctan => (Node::Arctan, f64::atan), + UnaryOp::Erf => (Node::Erf, libm::erf), + // `sign(±0) = +0`, unlike `copysign(1, x)`. + UnaryOp::Sign => (Node::Sign, |x| { + if x.is_nan() { + x + } else if x == 0.0 { + 0.0 + } else { + libm::copysign(1.0, x) + } + }), + UnaryOp::Floor => (Node::Floor, libm::floor), + UnaryOp::Ceiling => (Node::Ceiling, libm::ceil), + } +} + +/// The node that lowers to `op`, and the value it must take; the arithmetic is IEEE's own. +/// +/// Exhaustive, so a new primitive does not compile until it has a reference here. +fn binary(op: BinaryOp) -> (BinaryNode, BinaryFn) { + match op { + BinaryOp::Add => (Node::Add, |x, y| x + y), + BinaryOp::Sub => (Node::Sub, |x, y| x - y), + BinaryOp::Mul => (Node::Mul, |x, y| x * y), + BinaryOp::Div => (Node::Div, |x, y| x / y), + BinaryOp::Pow => (Node::Pow, pow), + BinaryOp::Minimum => (Node::Minimum, libm::fminimum_num), + BinaryOp::Maximum => (Node::Maximum, libm::fmaximum_num), + BinaryOp::Modulo => (Node::Modulo, libm::fmod), + BinaryOp::Hypot => (Node::Hypot, f64::hypot), + BinaryOp::EqualHeaviside => (Node::EqualHeaviside, |x, y| f64::from(u8::from(x <= y))), + BinaryOp::NotEqualHeaviside => (Node::NotEqualHeaviside, |x, y| f64::from(u8::from(x < y))), + BinaryOp::Equality => (Node::Equality, |x, y| { + f64::from(u8::from(libm::fabs(x - y) < 1e-14)) + }), + } +} + +/// Bitwise equal, or both NaN; with `any_zero_sign`, also `+0` against `-0`. +fn same(got: f64, want: f64, any_zero_sign: bool) -> bool { + got.to_bits() == want.to_bits() + || (got.is_nan() && want.is_nan()) + || (any_zero_sign && got == 0.0 && want == 0.0) +} + +#[track_caller] +fn assert_matches(label: &str, args: &[f64], got: f64, want: f64) { + assert!( + same(got, want, false), + "{label}{args:?} = {got:e} ({:#018x}), reference gives {want:e} ({:#018x})", + got.to_bits(), + want.to_bits(), + ); +} + +/// xorshift64*, so the sweep is reproducible without a rand dependency. +struct XorShift(u64); + +impl XorShift { + const fn next_u64(&mut self) -> u64 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + self.0.wrapping_mul(0x2545_F491_4F6C_DD1D) + } + + fn unit(&mut self) -> f64 { + (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64 + } + + /// Any bit pattern, so every exponent appears, infinities and NaNs among them. + const fn any_bits(&mut self) -> f64 { + f64::from_bits(self.next_u64()) + } + + /// A magnitude log-uniform over `1e-12..1e12`, either sign: the range model values span. + fn log_uniform(&mut self) -> f64 { + let sign = if self.next_u64() & 1 == 0 { 1.0 } else { -1.0 }; + // `libm`'s, so the sweep is the same on every target. + sign * libm::pow(10.0, 24.0 * self.unit() - 12.0) + } + + /// Uniform over `-50..50`, where the functions do most of their switching. + fn moderate(&mut self) -> f64 { + 100.0 * self.unit() - 50.0 + } + + /// An integer or half-integer in `-10..=10`, where `pow` takes its special paths. + fn exponent(&mut self) -> f64 { + (self.next_u64() % 41) as f64 / 2.0 - 10.0 + } +} + +/// Signed zeros, subnormals, the extremes, infinities and NaN, plus the points where the +/// functions switch branches, overflow or saturate. +fn special_values() -> Vec { + let magnitudes = [ + 0.0, + f64::from_bits(1), + f64::MIN_POSITIVE, + 1e-300, + 1e-8, + 0.1, + 0.5, + 0.84375, + 1.0 - f64::EPSILON / 2.0, + 1.0, + 1.0 + f64::EPSILON, + 1.25, + 1.5, + 2.0, + 1.0 / 0.35, + 3.0, + FRAC_PI_2, + PI, + 6.0, + 19.0, + 22.0, + 709.782_712_893_384, + 710.475_860_073_943_9, + 745.133_219_101_941_1, + 1e22, + 1e300, + f64::MAX, + f64::INFINITY, + ]; + // Opaque, so the compiler cannot fold a reference call on a literal with the host's libm. + std::hint::black_box( + magnitudes + .iter() + .flat_map(|&v| [v, -v]) + .chain([f64::NAN]) + .collect(), + ) +} + +const SWEEP: usize = 1 << 14; + +fn unary_inputs() -> Vec { + let mut rng = XorShift(0x9E37_79B9_7F4A_7C15); + let mut xs = special_values(); + for _ in 0..SWEEP { + xs.extend([rng.any_bits(), rng.log_uniform(), rng.moderate()]); + } + xs +} + +fn binary_inputs() -> (Vec, Vec) { + let special = special_values(); + let (mut xs, mut ys): (Vec, Vec) = special + .iter() + .flat_map(|&x| special.iter().map(move |&y| (x, y))) + .unzip(); + let mut rng = XorShift(0xD1B5_4A32_D192_ED03); + for _ in 0..SWEEP { + let pairs = [ + (rng.any_bits(), rng.any_bits()), + (rng.log_uniform(), rng.log_uniform()), + (rng.moderate(), rng.moderate()), + (rng.log_uniform(), rng.exponent()), + (rng.moderate(), rng.exponent()), + ]; + for (x, y) in pairs { + xs.push(x); + ys.push(y); + } + } + (xs, ys) +} + +fn state(arena: &mut Arena, start: usize, end: usize) -> NodeId { + arena.alloc(Node::StateVector { start, end }) +} + +fn eval(arena: &Arena, root: NodeId, y: &[f64]) -> Vec { + let expr = CompiledExpr::new(DagRef::new(arena, root)); + let mut scratch = vec![0.0; expr.scratch_len()]; + expr.eval(&mut scratch, 0.0, y, &[], &[]).to_vec() +} + +/// Two lanes, the second holding `lane_two`; returns each lane's root values. +fn eval_two_lanes(arena: &Arena, root: NodeId, y: &[f64], lane_two: &[f64]) -> [Vec; 2] { + let expr = CompiledExpr::new(DagRef::new(arena, root)); + let mut scratch = vec![0.0; expr.scratch_len() * 2]; + let y_cols = [y, lane_two].concat(); + let out = expr + .eval_batch(&mut scratch, 2, &[0.0, 0.0], &y_cols, y.len(), &[]) + .expect("a primal tape batches"); + // Lane-minor: element `e` of lane `l` sits at `2 * e + l`. + [0, 1].map(|lane| out.iter().skip(lane).step_by(2).copied().collect()) +} + +/// The literal `simplify` folds `root` to, if it folds. +fn folded(arena: &mut Arena, root: NodeId) -> Option { + let simplified = simplify(arena, root); + match *arena.get(simplified) { + Node::Scalar(v) => Some(v), + _ => None, + } +} + +#[test] +fn unary_primitives_take_reference_values() { + let xs = unary_inputs(); + let negated: Vec = xs.iter().map(|x| -x).collect(); + for op in UNARY_OPS { + let (node, reference) = unary(op); + let label = format!("{op:?}"); + let mut arena = Arena::new(); + let y = state(&mut arena, 0, xs.len()); + let root = arena.alloc(node(y)); + + for (&x, got) in xs.iter().zip(eval(&arena, root, &xs)) { + assert_matches(&label, &[x], got, reference(x)); + } + let lanes = eval_two_lanes(&arena, root, &xs, &negated); + for (lane, inputs) in lanes.iter().zip([&xs, &negated]) { + for (&x, &got) in inputs.iter().zip(lane) { + assert_matches(&format!("{label} (batched)"), &[x], got, reference(x)); + } + } + for x in special_values() { + let mut arena = Arena::new(); + let literal = arena.alloc(Node::Scalar(x)); + let root = arena.alloc(node(literal)); + if let Some(got) = folded(&mut arena, root) { + assert_matches(&format!("{label} (folded)"), &[x], got, reference(x)); + } + } + } +} + +#[test] +fn binary_primitives_take_reference_values() { + let (xs, ys) = binary_inputs(); + let n = xs.len(); + let (forward, swapped) = ([xs.as_slice(), &ys].concat(), [ys.as_slice(), &xs].concat()); + for op in BINARY_OPS { + let (node, reference) = binary(op); + let label = format!("{op:?}"); + let mut arena = Arena::new(); + let (lhs, rhs) = (state(&mut arena, 0, n), state(&mut arena, n, 2 * n)); + let root = arena.alloc(node(lhs, rhs)); + + for ((&x, &y), got) in xs.iter().zip(&ys).zip(eval(&arena, root, &forward)) { + assert_matches(&label, &[x, y], got, reference(x, y)); + } + let [lane_one, lane_two] = eval_two_lanes(&arena, root, &forward, &swapped); + for (i, (&x, &y)) in xs.iter().zip(&ys).enumerate() { + let batched = format!("{label} (batched)"); + assert_matches(&batched, &[x, y], lane_one[i], reference(x, y)); + assert_matches(&batched, &[y, x], lane_two[i], reference(y, x)); + } + } +} + +/// A one-element operand against a vector, either side, as the broadcast kernels run it. +#[test] +fn broadcast_binary_primitives_take_reference_values() { + let vector: Vec = unary_inputs().into_iter().step_by(24).collect(); + let n = vector.len(); + for op in BINARY_OPS { + let (node, reference) = binary(op); + for s in special_values() { + let y = [&[s], vector.as_slice()].concat(); + let mut arena = Arena::new(); + let (scalar, wide) = (state(&mut arena, 0, 1), state(&mut arena, 1, n + 1)); + let scalar_left = arena.alloc(node(scalar, wide)); + let scalar_right = arena.alloc(node(wide, scalar)); + let left = eval(&arena, scalar_left, &y); + let right = eval(&arena, scalar_right, &y); + for (i, &v) in vector.iter().enumerate() { + assert_matches( + &format!("{op:?} (scalar left)"), + &[s, v], + left[i], + reference(s, v), + ); + assert_matches( + &format!("{op:?} (scalar right)"), + &[v, s], + right[i], + reference(v, s), + ); + } + } + } +} + +#[test] +fn constant_folding_takes_reference_values() { + let special = special_values(); + for op in BINARY_OPS { + let (node, reference) = binary(op); + for &x in &special { + for &y in &special { + let mut arena = Arena::new(); + let (lhs, rhs) = (arena.alloc(Node::Scalar(x)), arena.alloc(Node::Scalar(y))); + let root = arena.alloc(node(lhs, rhs)); + if let Some(got) = folded(&mut arena, root) { + assert_matches(&format!("{op:?} (folded)"), &[x, y], got, reference(x, y)); + } + } + } + } +} + +/// C99 Annex F's `pow` special values, which `simplify` assumes in folding `1^x` to `1`. +/// Run on each CI target, this checks that target's C library. +#[test] +fn pow_takes_c99_special_values() { + let (inf, nan) = (f64::INFINITY, f64::NAN); + let cases = [ + (nan, 0.0, 1.0), + (nan, -0.0, 1.0), + (1.0, nan, 1.0), + (1.0, -inf, 1.0), + (-1.0, inf, 1.0), + (-1.0, -inf, 1.0), + (-0.0, 3.0, -0.0), + (-0.0, 1.5, 0.0), + (-0.0, -3.0, -inf), + (0.0, -1.5, inf), + (0.5, inf, 0.0), + (0.5, -inf, inf), + (2.0, inf, inf), + (2.0, -inf, 0.0), + (-inf, 3.0, -inf), + (-inf, 0.5, inf), + (-inf, -3.0, -0.0), + (inf, -0.5, 0.0), + (-8.0, 1.0 / 3.0, nan), + ]; + let n = cases.len(); + let mut arena = Arena::new(); + let (lhs, rhs) = (state(&mut arena, 0, n), state(&mut arena, n, 2 * n)); + let root = arena.alloc(Node::Pow(lhs, rhs)); + // Bases in the first `n` states, exponents in the next `n`. + let y: Vec = cases + .iter() + .map(|c| c.0) + .chain(cases.iter().map(|c| c.1)) + .collect(); + let got = eval(&arena, root, &y); + for (&(x, y, want), got) in cases.iter().zip(got) { + assert_matches("Pow", &[x, y], got, want); + } +} diff --git a/packages/pybamm-rust/rustfmt.toml b/packages/pybamm-rust/rustfmt.toml new file mode 100644 index 0000000000..bc3133b257 --- /dev/null +++ b/packages/pybamm-rust/rustfmt.toml @@ -0,0 +1,14 @@ +edition = "2024" + +max_width = 100 +use_small_heuristics = "Default" +tab_spaces = 4 +newline_style = "Auto" + +reorder_imports = true +reorder_modules = true + +use_field_init_shorthand = true +use_try_shorthand = true + +match_block_trailing_comma = true diff --git a/pyproject.toml b/pyproject.toml index 85f13ee7c8..9cb5630ff5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,6 +7,8 @@ [tool.uv.workspace] members = ["packages/*"] +# uv would glob the Rust crate and error on its absent pyproject. +exclude = ["packages/pybamm-rust"] [tool.uv.sources] pybammsolvers = { workspace = true }