feat: add accelerated compiled graph backends and WGPU support#1
Merged
Conversation
There was a problem hiding this comment.
Pull request overview
This PR significantly expands petite-ad beyond the legacy tuple-graph APIs by introducing reusable graph/tape/compiled-IR execution, adding forward-mode differentiation utilities, extending the mono/multi op sets and Hessian support (exact and finite-difference), and wiring in feature-gated backend infrastructure (including an initial WGPU skeleton).
Changes:
- Added a reusable
Graph/TapeAPI plus compiled IR/workspace abstractions and backend capability reporting. - Added forward-mode AD utilities (
ForwardAD) and expanded mono/multi operation support (incl. checked-domain evaluation and higher-order Hessian methods). - Added optimizer utilities (Gradient Descent, Adam), new examples/benchmarks, and CI/dev tooling updates.
Reviewed changes
Copilot reviewed 43 out of 47 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
src/optim.rs |
Adds simple optimizers (GD/Adam) with basic validation and tests. |
src/multi/tests.rs |
Adds Hessian and Hessian-row tests for multivariate tuple graphs. |
src/multi/parser.rs |
Introduces a small expression parser that builds reusable Graphs. |
src/multi/op_rules.rs |
Centralizes scalar + first/second-derivative local rules for MultiAD ops. |
src/multi/multi_hessian_common.rs |
Shared dual-number Hessian algorithm for FR/RF exact methods. |
src/multi/multi_ad_rf.rs |
Implements/exposes exact RF Hessian API + tests. |
src/multi/multi_ad_fr.rs |
Implements/exposes exact FR Hessian API + tests. |
src/multi/builder.rs |
Extends builder with node-handle graph construction + legacy compatibility. |
src/multi.rs |
Re-exports compiled/graph/backends and wires in new modules. |
src/mono/types.rs |
Adds a Dual type for forward-mode/higher-order internals. |
src/mono/tests.rs |
Adds tests for finite-difference Hessian + new mono ops. |
src/mono/tests_ho.rs |
Adds higher-order (RR/FR/RF) test coverage and cross-method consistency checks. |
src/mono/mono_hessian_common.rs |
Shared utilities for mono FR/RF Hessian computations. |
src/mono/mono_ad.rs |
Extends mono op set + adds checked evaluation and finite-difference Hessian helpers. |
src/mono/mono_ad_rf.rs |
Adds mono RF exact Hessian API built on shared utilities. |
src/mono/mono_ad_fr.rs |
Adds mono FR exact Hessian API built on shared utilities. |
src/mono.rs |
Wires mono module exports, including higher-order APIs and tests. |
src/main.rs |
Updates demo binary to showcase Hessian computations. |
src/macros.rs |
Extends macros for new mono ops and higher-order mono op enums; extends multi ops list. |
src/lib.rs |
Updates crate-level docs and re-exports for new APIs (graph/tape/compiled/forward/optim). |
src/forward.rs |
Adds public forward-mode AD API for mono expressions, tuple graphs, and reusable graphs. |
src/error.rs |
Extends error model with InvalidGraph and DomainError helpers. |
examples/wasm_demo.html |
Adds a sketch HTML file documenting a potential WASM demo shape. |
examples/optimizer_adam.rs |
Adds an Adam optimization example using Graph. |
examples/newton_method.rs |
Adds a Newton’s method example using exact Hessians. |
examples/hessian_performance.rs |
Adds a simple timing comparison example for exact Hessian methods. |
examples/hessian_demo.rs |
Demonstrates approximate vs exact Hessian methods and their accuracy differences. |
examples/graph_api.rs |
Demonstrates the new reusable Graph and ExprGraph APIs. |
examples/gradient_descent.rs |
Adds a basic gradient descent example using Graph. |
Cargo.toml |
Adds feature flags (serde, backend-wgpu) + new deps and profile tuning. |
benches/hessian_benchmark.rs |
Adds benchmarks for RR/FR/RF Hessian computation and scalability. |
benches/compute_benchmark.rs |
Extends benchmarks to cover graphs/tapes/compiled IR, SIMD batch paths, and checked mode. |
AGENTS.md |
Adds contributor/agent guidance, commands, and conventions. |
.pre-commit-config.yaml |
Updates pre-commit hooks and adds generic file hygiene checks. |
.gitignore |
Ignores Cargo.lock for the library repository. |
.github/workflows/ci.yml |
Adds CI workflow for formatting/clippy/tests/bench compile + coverage + audit. |
.cargo/config.toml |
Adds mold linker config and registry settings. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request significantly expands the library's capabilities by introducing exact second-order automatic differentiation (Hessian computation) using Reverse-over-Reverse, Forward-over-Reverse, and Reverse-over-Forward methods for both univariate and multivariate functions. It also adds a public forward-mode AD API, a reusable graph and tape API for repeated evaluations, compiled IR for batch execution with SIMD and prototype GPU backends, and built-in optimizers like Adam and Gradient Descent. Extensive documentation and benchmarks have been added to support these features. Feedback focuses on optimizing performance by reducing heap allocations in the forward-mode and Hessian accumulation loops, avoiding redundant arithmetic in value-only passes and Hessian matrix construction, and handling numerical edge cases such as power operations with a zero base.
Add exact second-order automatic differentiation using Reverse-over-Reverse, Forward-over-Reverse, and Reverse-over-Forward methods for both univariate and multivariate functions.
Add reusable Graph and Tape APIs, forward-mode differentiation utilities, checked-domain evaluation, multi-output graphs, tape workspace reuse, and expanded mono/multi operation support.
Add compiled instruction IR for acceleration-ready execution, batch compute/gradient APIs, reusable batch buffers, graph serialization and parser utilities, optimizers (GD, Adam), and stable math helpers (Log1pExp, LogAddExp).
Add flat backend IR with OpCode/FlatInstruction, SIMD backends (f64x2 SSE2, f64x4 AVX), scalar-lane fallback coverage for all opcodes, automatic backend dispatch, support diagnostics, device batch planning, and mock-device execution harness.
Add real WGPU backend with device initialization, GPU buffer allocation/upload/download, native WGSL compute kernel for exact-safe batch value computation, execution tracing, and host fallback for unsupported graphs and gradients.
Add 863 tests achieving 89.7% line coverage (LLVM engine). Fix arity validation, Adam overflow, pow zero-base NaN, Hessian perf, CI --all-features, value-only forward pass, mold linker config, and pre-commit hook cleanliness.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Details
backend-wgpusupport with real adapter/device initialization, GPU buffers/transfers, restricted exact-safe native batch compute, and host fallback for unsupported graphs/batches and gradientsValidation