[matthiaskoenig]

A smaller prototype for the open question of #1418:

A sparse difference quotient with the same coloring would also feed KLU, without Enzyme: one model evaluation per color, plus the finite difference error. It might be the cheaper route to most of the KLU speedup, if exactness turns out not to matter.

On the four benchmark models of #1418, it is the whole KLU speedup. KLU with the sparse difference quotient is as fast as KLU with the Enzyme Jacobian, takes the same steps, and gives the same error against a reference. Exactness does not matter for these models. The models load as fast and with as much memory as before, KLU saves the dense matrices, and the build needs no Enzyme.

It is a draft for discussion, stacked on #1425 (the counters of the integrators, which the benchmark needs); the first two commits are those of #1425.

What it does

It reuses the parts of #1418 that do not need Enzyme: the sparsity pattern and the column coloring from the SBML math (JacobianSparsity), the assembly J = cf * N * dv/dy plus the rate rule rows, the CVODE settings and the CMake option RR_USE_KLU. The Enzyme tangents are replaced by forward differences:

• LLVMExecutableModel::getStateVectorJacobianDQ evaluates the reaction rates and rate rule rates at y, and once per color with the states of that color perturbed.
• The CVODE callback uses the increments of CVODE's own dense difference quotient (cvLsDenseDQJac), so every entry is the forward difference CVODE would compute, with fewer model evaluations.
• The pattern and the coloring are computed for every model at load (0.2 ms for hctz_body, 3.8 ms for Chen_MSB2009), so there is no load option; the integrator settings alone switch it on.
• A model without a pattern falls back to the dense difference quotient: conserved moiety analysis, a stoichiometry set by a rule, a species conversion factor, or a conversion factor that depends on the state. Functions that Enzyme cannot differentiate (factorial, rateOf, ...) are no problem for differences.

r = roadrunner.RoadRunner("model.xml")
r.integrator.jacobian = "sparse_difference_quotient"
r.integrator.linear_solver = "klu"     # needs RR_USE_KLU
r.simulate(0, 100, 101)
r.integrator.getStatistics()           # #1425, plus num_jac_colors, num_jac_model_evals, sparse_linear_solver

Results

The models and simulations of #1418: simulate(0, end, 201) at rtol = atol = 1e-10, LLJit at DEFAULT. The Enzyme numbers come from the #1418 branch, with a small patch that sums its counters over the reinitializations of CVODE like #1425 does (enzyme-counter-totals.patch). Every number is the median over 5 rounds, in which the two builds and the models take turns, each in a process of its own; a simulation time is the median of 10 runs, in which the configurations take turns. The machine was shared with other jobs.

• speedup: against dense + DQ of the same build.
• model evaluations for Jacobians: of the state vector rate for CVODE's difference quotient, of the reaction rates for the sparse one (colors + 1 per Jacobian).
• one Jacobian: for CVODE's difference quotient, n times one evaluation of the state vector rate; for the sparse one, the time in the Jacobian callback (collect_timing); for Enzyme, the timing tool of Prototype: exact Jacobian from Enzyme and the sparse KLU solver for CVODE (#1412, #1409) #1418 at the middle of the timecourse.
• error / tol: the largest |x - x_ref| / (atol + rtol |x_ref|) over the output, against dense + DQ at rtol = 1e-12, atol = 1e-14.

model	states	colors	configuration	simulate (ms)	speedup	steps	error test fails	Jacobians	model evaluations for Jacobians	one Jacobian (µs)	error / tol
hctz_body	28	5	dense + DQ	1.5	1.00x	878	2	15	420	8.74	29.7
			dense + sparse DQ	1.4	1.07x	878	2	15	90	1.81	29.7
			KLU + sparse DQ	1.21	1.24x	878	2	15	90	1.81	29.7
			dense + DQ (Enzyme build)	1.45	1.00x	878	2	15	420	7.79	29.7
			KLU + Enzyme	1.21	1.20x	878	2	15	0	3.96	29.8
glimepiride_body	51	2	dense + DQ	1.1	1.00x	338	0	6	306	10.1	38.4
			dense + sparse DQ	1.05	1.05x	338	0	6	18	0.963	38.4
			KLU + sparse DQ	0.783	1.41x	338	0	6	18	0.963	38.4
			dense + DQ (Enzyme build)	1.03	1.00x	338	0	6	306	8.74	38.4
			KLU + Enzyme	0.744	1.39x	338	0	6	0	0.967	38.4
Smith_BMCSystBiol2013	104	3	dense + DQ	14.6	1.00x	2162	19	39	4056	51.3	1.36e+03
			dense + sparse DQ	14.4	1.01x	2322	33	41	164	3.68	1.36e+03
			KLU + sparse DQ	8.58	1.70x	2162	19	39	156	3.68	1.36e+03
			dense + DQ (Enzyme build)	14.4	1.00x	2162	19	39	4056	49.2	1.36e+03
			KLU + Enzyme	8.59	1.67x	2162	19	39	0	2.85	1.36e+03
Chen_MSB2009	500	7	dense + DQ	1.22e+03	1.00x	3100	275	94	47000	1.85e+03	1.9e+05
			dense + sparse DQ	1.06e+03	1.16x	3013	278	95	760	83.8	1.1e+05
			KLU + sparse DQ	260	4.70x	1840	68	63	504	83.8	1.84e+05
			dense + DQ (Enzyme build)	1.24e+03	1.00x	3100	275	94	47000	1.8e+03	1.9e+05
			KLU + Enzyme	267	4.63x	1874	70	62	0	36	1.62e+05

Load time, with recompilation:

model	load, this build (s)	load with Enzyme (s)
hctz_body	0.12	0.26
glimepiride_body	0.18	0.44
Smith_BMCSystBiol2013	0.33	0.80
Chen_MSB2009	1.14	3.80

• KLU brings the speedup, and any Jacobian with the sparsity pattern makes it possible. KLU + sparse DQ and KLU + Enzyme are within 6 % of each other on all four models. With the dense solver, the cheaper Jacobian gains at most 16 %.
• The steps do not depend on the exactness of the Jacobian. On Chen_MSB2009, KLU + sparse DQ takes 1840 steps with 68 error test failures, KLU + Enzyme 1874 with 70, and dense + DQ 3100 with 275. The other three models have identical counters. So the fewer steps with KLU come from the linear solver, as Prototype: exact Jacobian from Enzyme and the sparse KLU solver for CVODE (#1412, #1409) #1418 supposed. The 587 against 1349 steps in Prototype: exact Jacobian from Enzyme and the sparse KLU solver for CVODE (#1412, #1409) #1418 were the steps after the last reinitialization only; Chen_MSB2009 reinitializes CVODE 4 times, at piecewise switches.
• The error against the reference is the same for every configuration: it is the integration error.
• One sparse difference quotient costs about as much as one Enzyme Jacobian: faster on hctz_body, as fast on glimepiride_body, slower on Smith_BMCSystBiol2013 and Chen_MSB2009 (84 µs against 36 µs, and 22 times faster than CVODE's 1.85 ms). Over the 63 Jacobians of a simulation of Chen_MSB2009, the difference is 3 ms of 260 ms.

Memory

Each model, build and configuration in a fresh process, median of 5 processes (memory.py):

• RSS after the load and peak RSS of the load: minus the RSS before the load; they hold the compiled code, which LLVM does not allocate with malloc.
• heap after the simulation: the memory allocated with malloc and in use after simulate(0, end, 201) (glibc mallinfo2): the CVODE memory, the matrices and factorizations of the linear solver and the workspace of the Jacobian. Unlike the RSS, it drops when the dense matrices are replaced by sparse ones.
• heap, change by the simulation: the part allocated during the simulation (the saved copy of the Jacobian, the KLU factorization, the workspace of the sparse difference quotient, the result).

model	build and configuration	RSS after the load (MiB)	peak RSS of the load (MiB)	heap after the simulation (MiB)	heap, change by the simulation (MiB)
hctz_body	dense + DQ	30.9	30.9	22.32	0.05
	dense + sparse DQ	30.9	30.9	22.33	0.06
	KLU + sparse DQ	30.9	30.9	22.32	0.06
	#1418 build without Enzyme: dense + DQ	32.4	32.4	22.40	0.05
	#1418 build with Enzyme: dense + DQ	39.3	39.3	22.46	0.05
	#1418 build with Enzyme: KLU + Enzyme	39.3	39.3	22.45	0.06
glimepiride_body	dense + DQ	35.1	35.1	23.72	0.11
	dense + sparse DQ	35.1	35.1	23.73	0.11
	KLU + sparse DQ	35.1	35.1	23.70	0.10
	#1418 build without Enzyme: dense + DQ	36.5	36.5	23.81	0.11
	#1418 build with Enzyme: dense + DQ	47.7	47.7	23.85	0.11
	#1418 build with Enzyme: KLU + Enzyme	47.7	47.7	23.82	0.10
Smith_BMCSystBiol2013	dense + DQ	51.9	51.9	32.97	0.29
	dense + sparse DQ	51.9	51.9	33.00	0.32
	KLU + sparse DQ	51.9	51.9	32.90	0.29
	#1418 build without Enzyme: dense + DQ	53.4	53.4	33.04	0.29
	#1418 build with Enzyme: dense + DQ	64.7	64.7	33.11	0.29
	#1418 build with Enzyme: KLU + Enzyme	64.6	64.6	33.05	0.30
Chen_MSB2009	dense + DQ	110.3	119.8	55.98	2.68
	dense + sparse DQ	110.3	119.8	56.14	2.81
	KLU + sparse DQ	110.3	119.8	53.02	1.53
	#1418 build without Enzyme: dense + DQ	111.2	120.4	56.03	2.68
	#1418 build with Enzyme: dense + DQ	210.7	210.7	56.32	2.68
	#1418 build with Enzyme: KLU + Enzyme	181.7	208.3	53.36	1.54

• The sparse difference quotient costs no memory worth mentioning: 0.16 MiB of heap on Chen_MSB2009, less than 0.04 MiB on the other models, and the same RSS after the load.
• KLU saves memory on the large model: the dense solver keeps two dense n x n matrices (2 x 1.9 MiB for 500 states), KLU keeps them sparse (4400 nonzeros) plus its factorization. The heap of Chen_MSB2009 is 3.0 MiB smaller with KLU, with either Jacobian; on the small models the difference is below 0.1 MiB.
• The Enzyme Jacobian costs memory at load: 7 to 11 MiB of RSS on the small models, and 70 to 100 MiB on Chen_MSB2009 (its RSS after the load varies between processes, 182 to 211 MiB; the peak of the load is 208 to 211 MiB against 120 MiB without Enzyme).

Validation

• SparseDQJacobianTests.SbmlTestSuite compares the sparse difference quotient with the dense forward difference of the state vector rate with the same increments on the semantic cases of the SBML test suite, at the start and at t = 2, with MCJit and LLJit. They agree to rounding in all 1465 cases with a pattern, and the dense difference has no nonzero outside the pattern. 150 cases fall back (104 with a stoichiometry set by a rule, 46 with species conversion factors), and 208 do not load. The test runs when RR_SBML_TEST_SUITE points to test/sbml-test-suite/semantic.
• SparseDQJacobianTests (6 tests, both JITs) also checks that CVODE gives the result of its dense difference quotient with the dense solver and with KLU, including events, the fallbacks, and the settings in any order.
• The rest of ctest passes with RR_USE_KLU=ON, apart from the tests that collide under ctest -j (ModelEditingTests and StateSavingTests fail under ctest -j because they share one output directory #1417).

How to reproduce

prototype/sparse_dq_jacobian/ has RESULTS.md, benchmark.py (writes benchmark.csv), memory.py (writes memory.csv), table.py (prints the tables above from them), models.py and enzyme-counter-totals.patch. Both scripts need this build with RR_USE_KLU=ON and a build of the #1418 branch with the patch.

Not covered

• The fallbacks could be handled by differencing the state vector rate itself, with a coloring of the pattern of J instead of dv/dy. Conserved moiety analysis needs the Jacobian with respect to the independent species.
• Models restored with loadState have no pattern and fall back.
• getFullJacobian, the steady-state solvers and forward sensitivities could use it too.
• KLU still needs SuiteSparse, and SUNDIALS built with it, in libroadrunner-deps and the wheels.
• prototype/sparse_dq_jacobian/ would not go upstream as it is.

[matthiaskoenig]

@luciansmith This is a smaller solution for the sparse matrix with KLU which gives the same speedups. I.e. the numerical Jacobian methods can be used to figure out the sparsity pattern which then speeds up simulation (and probably also memory usage because the complete matrices don't have to be in memory, only the sparse parts).
This should speedup almost all simulations by at least 1.2x to 5x for larger model (10-20x for very large models).

This does not require to calculate the Jacobian, just the sparsity pattern. If we could bring this into roadrunner it would be great, the full Jacobian does not help too much at this point.

[matthiaskoenig]

Also the memory savings here in addition to the speedup is very big for us. We are running ~50 000 roadrunner instances in FEM meshes, if we save 3 MB per instance that is a huge thing for the servers ;)