Skip to content

Override Eigen's product evaluators - #3456

Open
SteveBronder wants to merge 8 commits into
developfrom
feature/eigen-var-product-hooks
Open

Override Eigen's product evaluators#3456
SteveBronder wants to merge 8 commits into
developfrom
feature/eigen-var-product-hooks

Conversation

@SteveBronder

@SteveBronder SteveBronder commented Oct 13, 2026

Copy link
Copy Markdown
Collaborator

Summary

This makes Eigen's operator* on var matrices call our rev/fwd functions instead of Eigen's Product evaluators. So A * B records one callback through multiply instead of O(n^3) varis. This works for matrix and array like types.

In order to be able to overload Eigen's Assignment and binary evaluator types this PR also adds a few other compound matrix multiplication operators that

  • scale_multiply(scalar, Mat1, Mat2) for s * (A * B)
  • scale_multiply(s, A, B) for s * (A * B)
  • multiply_add(A, B, C) for C.noalias() += A * B
  • multiply_sub(A, B, C) for C.noalias() -= A * B

Getting this to work required specializing Eigen's product_evaluator, binary_evaluator, and Assignment structs so that if they saw an output or input with a scalar var type the compiler would select our specializations instead of Eigens. This is kind of a brutal task because we essentially need to find the set of template parameters we can fix in our struct specialization such that the compiler will choose ours instead of Eigens.

To get this to work I had to build a bunch of test cases of basic matrix multiplication and then have a check that, if our overload was used, the resulting expression would have an extra tag that shows our specialization was used. Then I threw the problem at a bunch of claudes because I've tried fiddling with this before and it was horrible. But claude can try many insane combinations of specializations for the structs till they find a set that works correctly for our code. After it found a set that worked, I did a bunch of refinement to figure out what the smallest set of specializations was needed such that we need as few specializations of the above structs as possible.

We end up needing one specialization for product_evaluator, binary_evaluator, Assignment per assignment op (=, +=, -=), respectively. So not that bad actually!

Tests

./runTests.py test/unit/math/rev/fun/eigen_hook_dispatch_test.cpp \
  test/unit/math/rev/core/eigen_product_hook_test.cpp \
  test/unit/math/rev/core/eigen_*_mre_test.cpp \
  test/unit/math/mix/core/eigen_product_hook_kernels_tape_test.cpp \
  test/unit/math/rev/fun/scale_multiply_test.cpp \
  test/unit/math/rev/fun/multiply_add_test.cpp

./runTests.py test/unit/math/rev/core -f mre_test

Side Effects

  • Any Eigen product with a var operand now goes through Stan's functions, including the ones Eigen forms inside its own kernels. I need to do a deeper analysis of Eigen's code to make sure they are not creating a Product expression anywhere and then accessing it element by element.

  • Anytime we update Eigen versions we are going to be fixing even more internals. I think in the long run what we need to do is have Eigen do the same struct trick we do where they have an explicit template parameter that you can use to overload evaluators.

    Release notes

    Eigen expressions that multiply var matrices (A * B, element-wise products, scalar times matrix, and C.noalias() += A * B) now use Stan's reverse-mode functions, recording one node per product instead of one per scalar operation. Adds scale_multiply and multiply_add.

Checklist

  • Copyright holder: Simons Foundation

    The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
    - Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
    - Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

  • the basic tests are passing

    • unit tests pass (to run, use: ./runTests.py test/unit)
    • header checks pass, (make test-headers)
    • dependencies checks pass, (make test-math-dependencies)
    • docs build, (make doxygen)
    • code passes the built in C++ standards checks (make cpplint)

  • the code is written in idiomatic C++ and changes are documented in the doxygen

  • the new changes are tested

@SteveBronder
Evaluate Eigen products of autodiff scalars with Stan kernels
9cbbe42 
Eigen `operator*` on matrices and arrays with `var`, `fvar<T>`,
`std::complex<var>` or `std::complex<fvar<T>>` scalars now dispatches to
Stan kernels instead of Eigen's per-scalar kernels:

- Matrix products (every product tag, lazyProduct, noalias, =, +=, -=)
  through a prim-level specialization of Eigen's dense product_evaluator
  that forwards autodiff products to eigen_matrix_product_impl and keeps
  Eigen's own evaluation for all other scalars, plus Assignment
  specializations for same-scalar destinations.
- Coefficient-wise products and products with a scalar through
  binary_evaluator specializations for scalar_product_op.
- Scalar factors are folded into the product they scale, so s * (A * B),
  (A * B) * s and (s * A) * B each record one callback.

The reverse-mode kernel records one callback plus one nochain vari per
result coefficient. The forward-mode kernel splits fvar products into
products of values and tangents (one concatenated product for the
tangent), so fvar<var> matrix products record two callbacks instead of
several varis per scalar operation; tape-free fvar<double> element-wise
products stay lazy.
@SteveBronder
Keep reverse-mode product callbacks small
dcaea9c 
Scalar factors of a product kernel are stored on the arena at the first
autodiff factor instead of in a fixed array inside every callback, result
adjoints are read lazily, and var operand adjoints are updated with one
vectorized expression.
@SteveBronder
Route Eigen var products to Stan's functions from Eigen's own structs
7e84424 
Replace the evaluator/customization-point layers and the rev and fwd
product kernels with specializations of Eigen's structs that call Stan's
free functions directly:

- product_evaluator (dense and per-tag lazy) calls multiply, or
  scale_multiply when an operand is Eigen's `s * A` / `A * s`.
- binary_evaluator for scalar_product_op of var with var, double or int
  calls elt_multiply, multiply(s, X) or scale_multiply.
- Assignment of a product to a var destination calls multiply,
  multiply_add (`noalias() +=`, `C + A * B`) or subtract (`noalias() -=`).

Add scale_multiply(s, A, B) and multiply_add(A, B, C) (prim and rev), each
recording one callback. This pass handles real var only: products with
fvar or complex scalars use Eigen's evaluation, and the tests that need
them are reverse-only or removed until those are routed as well.
@SteveBronder
Simplify the Eigen specializations that route var products to Stan
9a6a82b 
Move the hooks header to stan/math/rev/meta and replace the macros with
plain specializations of Eigen's structs:

- binary_evaluator: one specialization. Its var condition sits in the
  CwiseBinaryOp functor argument, which Eigen's specialization leaves
  free, so it covers var with var, double and int in one place.
- Assignment: one specialization per functor (=, +=, -=). The condition
  sits in both scalars of the functor, so products of non-dense operands
  run Eigen's own specialization and the copied Eigen bodies are gone.
  Inner products are now computed by multiply as well.
- product_evaluator and binary_evaluator share one evaluator body,
  ad_product_evaluator, which no longer allocates a result it discards.
- Lazy products (lazyProduct()) are left to Eigen.

The traits strip references and const themselves, with
require_eigen_product_t and require_eigen_cwise_product_t aliases.
Every Stan specialization derives from eigen_ad_hook so tests can check
which struct Eigen picks.

scale_multiply and multiply_add return their arena result instead of
copying it to the heap.

Tests: eigen_hook_dispatch_test checks which specialization Eigen picks
for each scalar pairing, operator and destination; MRE tests instantiate
each struct on its own; eigen_product_hook_test gains products nested in
larger expressions, Stan functions and block/transpose destinations.
@SteveBronder
Remove fixed-size product tests until the arena alignment fix lands
94dbdbb 
Fixed-size var products now go through multiply, whose callback captures
fixed-size Eigen matrices inside an 8 byte aligned vari. They are fixed
on fix/aligned-vari-alloc; add these tests back once that is merged.
@SteveBronder
Add multiply_sub and take Stan product arguments by forwarding reference
ede91ae 
- Add multiply_sub(A, B, C) = C - A * B (prim and rev), one reverse pass
  callback like multiply_add. `C.noalias() -= A * B` now uses it instead
  of subtract(C, multiply(A, B)), so it records one callback instead of
  two.
- scale_multiply and multiply_add take their arguments by forwarding
  reference. The rev versions read operand values through value_of instead
  of caching copies, evaluating the result adjoint once per callback. The
  prim versions return a lazy expression through make_holder.
- Fix rev scale_multiply computing its value with the var scale, which
  recorded a vari per coefficient on top of the callback.
- Remove test/unit/math/ad_tape_util.hpp and read the reverse-mode stack
  sizes directly in the tests that used it; glm_forward_mode_tape_test is
  back to develop's version.
@SteveBronder SteveBronder changed the title Feature/eigen var product hooks Override Eigen's product evaluators Oct 13, 2026
@SteveBronder SteveBronder self-assigned this Oct 13, 2026
@jachymb

jachymb commented Oct 14, 2026

Copy link
Copy Markdown
Contributor

Ah, is this the proper fix in place of #3418 ? Could perhaps #3418 even be reverted after this to make code a bit less verbose?

@SteveBronder

SteveBronder commented Oct 14, 2026

Copy link
Copy Markdown
Collaborator Author

Yes if this PR works we could revert that PR, but I think there is still a lot of work to get this in shape

@SteveBronder
Copy var Maps passed from Eigen expressions to Stan's functions
da45b61 
Stan's functions assume an Eigen::Map points at arena memory and keep the
Map in their reverse-pass callbacks (arena_matrix(const Map&) does not
copy), so a callback reads whatever vars the mapped memory holds during the
reverse pass. Eigen's own algorithms map workspaces they reuse and free:
HouseholderQR's reflections (Householder.h applyHouseholderOnTheLeft) map a
var workspace that every reflector overwrites and that is freed when the QR
object or HouseholderSequence::evalTo returns. Callbacks then sent adjoints
to the wrong varis or read freed memory, which segfaulted
MathMixMatFun.qrQ on arm64 (and in optimized x86 builds).

copy_var_map copies the vars (not their varis) of a contiguous Map of vars
into the arena and passes everything else through. The hooks apply it to
every operand they pass to multiply, scale_multiply, elt_multiply,
multiply_add and multiply_sub, including the destination of += / -=.
Double Maps and arena_matrix are not copied.

Tests: Map operands and Map destinations whose buffer is refilled with new
vars before the reverse pass must still propagate the gradient to the vars
the forward pass used.

Also adds is_eigen_array_v and has is_eigen_array decay its argument, and
moves the dispatch tests onto the AgradRev fixture.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

@SteveBronder SteveBronder

Labels

feature performance

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@SteveBronder @jachymb