Performance: kron-based sparse product kernels, factorization reuse in mldivide, new blkdiag method

[gregkaplan]

This PR was prepared at the request of Greg Kaplan, following a performance exploration of the package together with Claude (Anthropic's Claude Code). The motivating application is computing large NxN transition-path Jacobians for continuous-time heterogeneous-agent GE models (backward HJB / forward KFE recursions built from large sparse matrices, sparse linear solves, and embedded nonlinear solvers), where the package is used to differentiate the stacked equilibrium conditions. Profiling that application against current master identified a small number of hot spots; every change below was benchmarked against master and validated to produce identical derivatives.

Changes

1. private/matvalXmatder.m and private/matderXmatval.m: single kron-based sparse product

The 2018 implementations loop over the inner dimension of the matrix product with logical masks over all nonzeros, which is O(ninter*(nnz(A)+nnz(B))). The identities

• vec(A*dB) = kron(I_ncol, A) * vec(dB)
• vec(dA*B) = kron(B', I_nrow) * vec(dA)

turn each kernel into one sparse matrix product, letting MATLAB's multithreaded sparse BLAS do the work.

Kernel-level timings (R2025b, Apple Silicon; A 5 nnz/row, ncol=20, 100 derivative directions):

inner dim	matvalXmatder master	kron	matderXmatval master	kron
50	0.0029s	0.0001s	0.0031s	0.0001s
500	0.0178s	0.0003s	0.0426s	0.0007s
2000	0.0860s	0.0010s	1.1474s	0.0074s

Package-level effect on mtimes for (1000x1000 sparse) * (1000x20) with 100 directions: ADAD 0.196s -> 0.040s (4.9x), doubleAD 0.125s -> 0.0008s (156x), AD*double 0.077s -> 0.040s (1.9x). Derivative fingerprints identical to master.

2. mldivide.m: reuse one factorization per call

Each matrix branch previously factorized the same matrix multiple times: twice for a single right-hand side (value solve + derivative solve), and once per column for matrix right-hand sides. The patch creates dec = decomposition(sparse(...)) once per call and reuses it for all solves. The binary_ext handling and all branch semantics are unchanged.

On a 3600x3600 five-point-stencil matrix with 100 derivative directions: 13% faster with 1 RHS, 27% faster with 10 RHS; identical results. Gains grow with factorization cost relative to triangular-solve cost, i.e. precisely in the large sparse systems the package targets. Requires R2017b+ (decomposition), which the package effectively already assumes (implicit expansion is used in private/valXder.m).

3. New method: blkdiag.m

Building block-diagonal AD matrices by repeated concatenation, e.g.

A = [A; sparse(n, n*(k-1)), Amat{k}, sparse(n, n*(K-k))];

is O(nblocks^2) and permutes the full stacked derivative matrix at every step. In the HANK application above (assembling the intensity matrix from per-income-state blocks inside the HJB loop), this concatenation pattern was 63% of total Jacobian runtime. The new blkdiag method assembles the values and the stacked derivatives with one sparse() triplet call each, accepts any mix of myAD and numeric blocks, and matches the concatenation result bit-for-bit. For 30 blocks of 40x40 with 100 directions: 0.225s -> 0.003s (68x).

4. compile_mex_files.m: comment out mex compilation, with an explanatory note

On recent MATLAB the pure-MATLAB kernels are as fast as or faster than the C mex kernels, which predate multithreaded sparse ops and implicit expansion: valXder mex is at parity with the one-line .m, and the matdrivXvecval mex measured 4-37x slower than its 2018 .m implementation (N = 500 to 5000, K = 198). Because a compiled mex silently shadows the .m file of the same name, running the compile script can make the package slower on modern MATLAB. Compilation lines are kept (commented) for old releases, with a note to delete stale binaries. (Incidentally, the matvalXmatder.c mex produced results differing from both .m implementations in our tests, so its compilation remaining commented out seems wise.)

private/matdrivXvecval.m is deliberately left unchanged: a kron reformulation was faster at moderate sizes but slower at N=5000, so the existing implementation stands.

Validation

• All changed code paths (mtimes ADAD / ADdouble / double*AD, all six matrix branches of mldivide, blkdiag with mixed AD/numeric blocks) validated against central finite differences: max error ~1e-10.
• blkdiag output compared element-by-element against the concatenation construction: identical values and derivatives.
• End-to-end application test: a 198x198 GE transition Jacobian of a continuous-time HANK model (100-period backward HJB / forward KFE, 1500-state sparse systems, upwinding, embedded fixed-point iterations) computed with master and with this branch: Jacobians identical to the last bit.
• EXAMPLE_AutoDiff_syntax.m runs unchanged.

Benchmarks were run on MATLAB R2025b (Apple Silicon, maca64). The kernel changes are pure sparse linear algebra, so relative gains should carry over to other platforms.

🤖 Generated with Claude Code

[gregkaplan]

Follow-up commit (7150b1e) — please review this branch at 7150b1e rather than the initial commit. Two problems with the kron formulation I originally proposed turned up on further testing, and fixing them properly also made the kernels faster.

1. Memory. kron(I,A) and kron(B',I) materialize ncol*nnz(A) and nrow*nnz(B) entries respectively, regardless of how sparse the derivative payload is. For a dense 2000x2000 value matrix with 100 columns that is 4e8 stored entries (~5 GB) even when the payload has a handful of nonzeros. The benchmarks in the original commit message were all in the sparse regime and did not expose this.

Both kernels now avoid forming any kron intermediate:

• matvalXmatder: the contraction runs over the leading index of the column-wise stacking, so reshaping the payload to (ninter x ncol*nderiv) lays every direction side by side and one sparse product handles them all. This is strictly better than kron — same result, O(nnz) peak memory, and faster everywhere measured (dense-A / sparse-payload 0.0067s -> 0.0001s; HANK-like 0.0007s -> 0.0003s).
• matderXmatval: the contraction runs over the trailing index, so a plain reshape does not line the directions up. The payload is reindexed once (O(nnz)) to move the derivative direction into the row index, one sparse product contracts all directions, and the result is reindexed back. 5-6x faster than kron in the wide/dense regimes, within ~30% of it in the moderate-sparse regime, with bounded memory.

Package-level mtimes for (1000x1000 sparse)*(1000x20) with 100 directions, versus master: AD*AD 0.196s -> 0.0021s (93x), AD*double 0.077s -> 0.0013s (59x), double*AD 0.125s -> 0.0005s (250x). Derivative fingerprints identical to master.

2. Complex correctness. Neither kernel takes a transpose now, so both are correct for complex input by construction. My kron version of matderXmatval used B' where vec(dA*B) = kron(B.', I)*vec(dA) requires the nonconjugate transpose — a one-character error that passes every real-valued test.

Worth flagging that master has an analogous pre-existing bug in both kernels: Aval = Aval(:)' conjugates the stored derivative values. Max derivative error for complex input against a per-direction reference: master 5.1 (matderXmatval) and 7.3 (matvalXmatder), my kron-with-B' 4.2, both new kernels ~5e-16. So this branch fixes a complex-input bug that predates it. If the package's intended scope is real-valued derivatives only, it may be worth saying so explicitly in the README instead — but the kernels are now correct either way, at no cost.

Tests. Added test_matrix_kernels.m: checks mtimes (AD*AD, AD*double, double*AD) and mldivide against the per-direction product rule across five shapes for both real and complex inputs, plus zero-payload, single-direction, and blkdiag cases. All pass on this branch; the complex cases fail on master, which is how the conjugation bug surfaced.

Re-validated after the change: finite-difference checks on all touched paths (~1e-10), EXAMPLE_AutoDiff_syntax.m output unchanged, and the continuous-time HANK transition Jacobian (198x198, 100-period backward HJB / forward KFE, 1500-state sparse solves) still bit-identical to master.

[gregkaplan]

One correction to something I wrote in the PR description, since it was imprecise in a way that matters.

I said the matvalXmatder.c mex "produced results differing from both .m implementations". That is true as stated but the diagnosis was wrong, and the real finding is more serious: matvalXmatder.c implements an older calling convention than the current mtimes.m.

The C header documents it — "To get dB/dx to dB'/dx, you can call dertransp(dB/dx,m) prior to calling matvalXmatder" — and the pre-2018 mtimes.m did exactly that. The 2018 matvalXmatder.m instead takes the column-wise stacked derivative directly, and today's mtimes.m calls it that way. Verified numerically against a per-direction reference:

upstream .m                      err 0.00e+00
mex(A, dB)                       err 1.79e+00     <- how mtimes.m calls it today
mex(A, dertransp(dB, ninter))    err 0.00e+00     <- the convention the .c expects

So the kernel is not buggy in itself; it is correct through the old API. But since a compiled mex silently shadows the same-named .m, compiling matvalXmatder.c against the current mtimes.m makes AD matrix-matrix products return wrong derivatives with no warning. That is a sharper reason to leave its compilation commented out than "it might be slower", and it is now recorded in compile_mex_files.m (commit e161ed2) along with the note that reviving the C path requires either changing mtimes.m back to pass dertransp(dB, n) or updating the C source to the column-wise convention.

Worth flagging for anyone with a stale binary in their working copy: this would have been silently wrong since the 2018 kernel rewrite, not just under this branch. It may be worth deleting the checked-in *.mexmaci64/*.mexa64 files from the repository for the same reason.

Also fixed in that commit: the note still described matvalXmatder.m/matderXmatval.m as kron-based, which stopped being true with 7150b1e.