Skip to content

Commit 50952c5

Browse files
albopclaude
albop
and
claude
committed
docs: computational engines (vectorized and JAX), how they are implemented and optimized
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent c5e72bc commit 50952c5

2 files changed

Lines changed: 212 additions & 0 deletions

File tree

‎docs/engines.md‎

Lines changed: 211 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,211 @@
1+
# Computational Engines
2+
3+
The global solution algorithms (`time_iteration`, `newton`) and the residual functions
4+
(`F`, `dF_1`, `F_and_dF_1`, `dF_2`) take an `engine` argument that selects how the Euler
5+
residuals and their derivatives are computed:
6+
7+
| `engine` | implementation | device |
8+
|---|---|---|
9+
| `"auto"` (default) | JAX if it can be imported, else vectorized | JAX's default device |
10+
| `"jax"` | JAX | JAX's default device (a GPU if one is visible) |
11+
| `"gpu"` (or `"cuda"`) | JAX | the first GPU (error if there is none) |
12+
| `"cpu"` | JAX | the CPU |
13+
| `"vectorized"` (or `"numpy"`) | NumPy / numba | the host |
14+
15+
Both implementations compute the same objects, those of Dolo.jl. For controls $x$ on the grid
16+
and a decision rule $\varphi$ (a spline per exogenous regime interpolating the controls):
17+
18+
$$F(x, \varphi)_{i,n} = \sum_k w_{ik}\, f\big(m_i, s_n, x_{in}, M_{ik}, S_{ikn}, \varphi_{r(k)}(S_{ikn})\big),
19+
\qquad S_{ikn} = g(m_i, s_n, x_{in}, M_{ik})$$
20+
21+
where $i$ is the exogenous regime today, $n$ the grid point, and $k$ the node of the process
22+
tomorrow, with weight $w_{ik}$, value $M_{ik}$ and regime $r(k)$. The algorithms need:
23+
24+
- $F$, and $dF_1 = \partial F / \partial x$: block diagonal, one $n_x \times n_x$ block per
25+
point;
26+
- $dF_2 = \partial F / \partial \varphi$: a linear operator
27+
$L.v = \sum_k F_x[i,k,n] \cdot v_{r(k)}(S_{ikn})$ acting on variations $v$ of the decision
28+
rule; it is stored as the blocks $F_x$ and the points $S$;
29+
- $J \backslash L$ (preconditioning by $dF_1$), and solutions of $(I - L)\,dx = b$ (Neumann
30+
series or GMRES).
31+
32+
All engines use the same algorithms, with the same defaults and stopping rules as Dolo.jl
33+
(see [Defaults](#defaults-identical-to-dolojl)).
34+
35+
---
36+
37+
## Which code runs
38+
39+
| | JAX engines | vectorized engine |
40+
|---|---|---|
41+
| no complementarity conditions (or `ignore_constraints=True`), all regimes with the same number of nodes | **device loops**: everything on the device | host loops |
42+
| complementarity conditions | host loops (Fischer–Burmeister smoothing, `ncpsolve` for the inner solver of time iteration): `time_iteration` evaluates the residuals and Jacobians with the JAX kernels and copies them to the host, `newton` uses the vectorized residuals | host loops |
43+
44+
The *device loops* (`_time_iteration_device` in `algos/time_iteration.py`, `_newton_device` in
45+
`algos/newton.py`) keep the controls, the spline coefficients, the Jacobians and the operator
46+
$L$ on the device. Only error norms (scalars) go to the host during the iterations; the host
47+
`DecisionRule` returned in the result is fitted once, at the end. The other loops work on host
48+
(NumPy) arrays.
49+
50+
---
51+
52+
## Vectorized engine (NumPy / numba)
53+
54+
Code: `numeric/tau.py` (`euler_residuals_tau`), `algos/newton.py` (`euler_residuals`,
55+
`Operator`, `neumann_solve`), `algos/time_iteration.py`.
56+
57+
- **Model equations** are evaluated on `(N, n)` arrays, one call per (regime, node): NumPy
58+
expressions emitted from the equations (or the numba functions of the legacy compiler).
59+
- **Decision rule**: cubic splines, coefficients by prefiltering. The approximators use the JAX
60+
implementation (`interpolation_jax`) by default, on JAX's default device: the values are
61+
copied there and back at each evaluation.
62+
- **dF_1** by finite differences of $F$ (one evaluation of the residuals per control);
63+
**dF_2** by finite differences of $f$ with respect to the next-period controls
64+
(`euler_residuals(..., with_jres=True)`), stored as dense blocks.
65+
- **Linear solves** on the host: `scipy.sparse.linalg.gmres` (with Dolo.jl's settings: from
66+
$x_0 = b$, one cycle, absolute tolerance) or a Neumann series; each product $L.v$ refits a
67+
spline to $v$.
68+
69+
It is the reference for the checks of the JAX engine, and `newton` uses it for problems with
70+
complementarity conditions whatever the engine. It is not tuned for speed: on integration_B (20³×9 points, 8
71+
controls), $F$ takes about 140 ms and the finite-difference $dF_1$ about 1.3 s (RTX 2080 host,
72+
2026-10-20).
73+
74+
---
75+
76+
## JAX engine
77+
78+
Code:
79+
80+
| module | content |
81+
|---|---|
82+
| `numeric/residual_jax.py` | `CompiledResidual`: entry point (`F`, `dF_1`, `F_and_dF_1`, `dF_2`, `prefilter`), point functions, the previous vmapped kernels (`fast=False`) |
83+
| `numeric/residual_kernels_jax.py` | `GridKernels`: staged kernels for F, (F, dF_1), dF_2; structural sparsity and colored Jacobians; small jitted solver steps |
84+
| `numeric/policy_operator_jax.py` | the operator $L$: location of the points, prefilter, spline evaluation, contraction, $J\backslash L$, GMRES |
85+
| `numeric/residuals.py` | `PolicyLinearOperator` (both engines) and the residual wrappers |
86+
| `numeric/interpolation_jax.py` | splines (prefilter matrices, basis) |
87+
88+
All arrays are `(n_regimes, N, n_x)`, one regime being a regime axis of length one. Float64 only
89+
for now.
90+
91+
### Kernels for F, dF_1 and dF_2 (`GridKernels`)
92+
93+
1. **Model equations as JAX point functions.** The equations are emitted as `jax.numpy`
94+
expressions of one point; they are `vmap`ped over the grid points and the next-period nodes.
95+
Terms that depend only on today's variables are computed once per point, not once per node
96+
(`vmap` does not batch what does not depend on the node).
97+
2. **Separate compiled stages.** $S = g(\cdot)$, the decision rule at $S$, and the arbitrage part
98+
are three jitted functions. Compiled as one function, XLA fuses the spline evaluation (gathers)
99+
into its consumers and recomputes it for each output: several times slower.
100+
3. **Spline evaluation.** For each point: index of the cell (one flat index into the coefficients
101+
of all regimes) and the 1-D basis weights; the value is the nested sum
102+
$\sum_p w_{0p} \sum_q w_{1q} \sum_r w_{2r}\, c_{pqr}$ (84 multiply-adds per output in 3-D) on
103+
coefficients gathered by flat index. No $4\times4\times4$ patches are materialized.
104+
4. **Only the used columns.** The decision rule is interpolated only for the controls that $f$
105+
depends on next period (6 of 8 on integration_B).
106+
5. **Linearized decision rule** (as Dolo.jl's `Linearized`). For $dF_1$, the decision rule at
107+
$S(x)$ is replaced by $X_0 + G\,(S(x) - S_0)$, with the value $X_0$ and gradient $G$ of the
108+
spline at $S_0$ computed by the spline code (the gradient: the same nested sum with the
109+
derivative of the basis in one dimension). Autodiff never differentiates the gathers.
110+
6. **Sparse (colored) forward-mode Jacobians.** Columns of a Jacobian that never share a row
111+
are differentiated together, with one `jvp` per group of columns ("color"): 5 instead of 8
112+
for $dF_1$, 4 instead of 8 for $\partial f / \partial X$ (which gives $dF_2$) on integration_B.
113+
The sparsity pattern is **structural**: it is read from the traced program (`jaxpr`), by
114+
propagating which inputs each value depends on, with the values of the parameters (so that
115+
$x^0$ and $0 \cdot x$ do not depend on $x$). It is never inferred from numerical zeros: on
116+
integration_B, $\partial f / \partial p_f'$ is zero at the steady-state guess (because
117+
$db_f = 0$ there) but not elsewhere.
118+
7. **Bounded memory.** At most $2^{20}$ (regime, node, point) triples are processed at once;
119+
beyond, regimes go in sequence (`lax.map` with `batch_size`), sums over nodes are done per
120+
regime. 20³×9 runs in one batch; 40³×9 fits in 6 GB.
121+
8. **dF_2 compressed.** `dF_2_compressed` returns only the columns of the blocks $F_x$ that can
122+
be nonzero; the dense blocks are built only on request (`CompiledResidual.dF_2`,
123+
`PolicyLinearOperator.F_x`).
124+
125+
### The operator L (`PolicyLinearOperator`, `policy_operator_jax`)
126+
127+
1. **Points located once.** $S$ does not change while $L$ is applied (in a linear solve): the
128+
cell index and the basis weights of each $S_{ikn}$ are computed when the operator is built.
129+
2. **Column compression.** Columns of $F_x$ that are zero everywhere are dropped: $L.v$
130+
prefilters and interpolates only these columns of $v$. $J^{-1} F_x$ keeps them zero.
131+
3. **$L.v$ in three jitted stages**: prefilter (three small matrix products per regime), nested
132+
spline evaluation, and the contraction $\sum_k F_x \cdot v(S)$ written elementwise (faster
133+
than `einsum`'s batched products).
134+
4. **$J \backslash L$**: each $n_x \times n_x$ block of $J$ is inverted once
135+
(`jnp.linalg.inv`), the products are elementwise (a batched solve with $K \cdot n_c$
136+
right-hand sides needs transposed copies); one regime at a time beyond $2^{20}$ blocks.
137+
Scaling the operator multiplies a factor, not the blocks.
138+
5. **Linear solves on the device.** Neumann series (one scalar to the host per term), or GMRES
139+
written as Dolo.jl's `gmres!`: one cycle of at most `maxit` Krylov vectors from $x_0 = b$,
140+
modified Gram–Schmidt on the device (the basis is updated in place), Givens rotations of the
141+
small Hessenberg matrix on the host (one column copied per iteration), absolute tolerance on
142+
the 2-norm of the residual. The number of iterations is kept in `last_iterations`.
143+
144+
### Solvers (device loops)
145+
146+
- `CompiledResidual.prefilter` is one jitted computation for all regimes; device arrays stay on
147+
the device.
148+
- Small steps of the loops (block Newton direction $J^{-1} r$, updates $x + a\,y$, norms) are
149+
jitted functions: one dispatch each, no constants copied from the host.
150+
- Only the error norms are copied to the host (checked with `jax.transfer_guard`). The scalar
151+
arguments of the jitted functions (step sizes, tolerances) are the only other transfers.
152+
153+
### Things that did not work (on these GPUs)
154+
155+
- Pallas (Triton) kernels need a GPU of compute capability ≥ 8.0; the RTX 2080 and GTX 1660 Ti
156+
are 7.5.
157+
- XLA computes each output element independently: the spline value with its 3 derivatives costs
158+
about 4 evaluations, whatever the formulation; multi-output fusions or materialized tensor-product
159+
weights were slower.
160+
- An unrolled Gauss–Jordan inverse of the 8×8 blocks is faster than `jnp.linalg.inv` (2.6 vs
161+
5.9 ms) but takes minutes to compile; written as a loop over pivots it is slower (11.5 ms).
162+
- `jax.scipy.sparse.linalg.gmres` is as fast as the GMRES above, but its tolerance is relative
163+
and it reorthogonalizes: it does not reproduce Dolo.jl's iterations.
164+
165+
---
166+
167+
## Defaults (identical to Dolo.jl)
168+
169+
| | option | default |
170+
|---|---|---|
171+
| `time_iteration` | `maxit`, `tol_eps`, `tol_eta` | 500, 1e-8, 1e-6 |
172+
| | `inner` (Newton solver of $F(x_1, \varphi_0) = 0$) | `maxit=10`, `max_backsteps=5`, `tol=tol_eps` |
173+
| | `accelerate` (when given) | `wait=10`, `linsolve="neumann"`, `tol=1e-10`, `maxit=1000` (Neumann terms) or 50 (GMRES Krylov vectors), `safeguarded=False` |
174+
| `newton` | `maxit`, `tol_eps`, `tol_eta` | 20, 1e-8, None (no step criterion) |
175+
| | `linsolve` | `"neumann"`, `tol=1e-10`, `maxit=1000` (Neumann) or 50 (GMRES) |
176+
| | `linesearch` | None (full Newton steps); `"backtracking"`: `max_backsteps=10`; `"golden_section"`: `maxiter=5` |
177+
178+
Stopping rules, as in Dolo.jl: time iteration stops before a step when the residual is below
179+
`tol_eps` (or not finite) and after a step when the step is below `tol_eta`; Newton stops when
180+
the residual is below `tol_eps` (or not finite), and goes back to the previous point and stops if
181+
no step decreases the residual. Iterations count the steps taken. `improved_time_iteration`
182+
(legacy name of `newton`) keeps GMRES and backtracking.
183+
184+
Full Newton steps (the default) need a good starting point: from the default initial guess they
185+
diverge on `rbc` and `rbc_mc`, where `linesearch="backtracking"` converges.
186+
187+
---
188+
189+
## Performance
190+
191+
integration_B (two countries, Epstein–Zin, 3 endogenous states, VAR(1) discretized to 9 Markov
192+
states, 8 controls), 20³×9 = 72 000 points, Float64, same options as Dolo.jl's benchmark (GMRES,
193+
linear tolerance 1e-10, backtracking for Newton), 2026-10-21:
194+
195+
| | Dolo.jl, RTX 2080 | dolo.py JAX, RTX 2080 | Dolo.jl, GTX 1660 Ti | dolo.py JAX, GTX 1660 Ti |
196+
|---|---|---|---|---|
197+
| time iteration (243 iterations) | 53.8 s | 42.5 s | 91 s | 70.9 s |
198+
| accelerated time iteration (15) | 8.1 s | 6.7 s | 13.3 s | 11.0 s |
199+
| Newton (6) | 3.62 s | 3.66 s | 5.8 s | 5.76 s |
200+
201+
Kernels on the RTX 2080 (20³×9): $F$ 12.8 ms (Dolo.jl 19.1), $F$ and $dF_1$ 53 ms (Dolo.jl
202+
$F + dF_1$: 85), $dF_2$ 22 ms (51.5), $L.v$ 7.1 ms (5.8), $J\backslash L$ 10.2 ms (13.6),
203+
prefilter 1.5 ms (0.5).
204+
205+
When measuring JAX code: calls are asynchronous (wait with `jax.block_until_ready` before
206+
reading a clock); compilation happens once per array shape (warm up on the grid that is timed);
207+
`jax.transfer_guard("log")` shows the copies between host and device.
208+
209+
Tests: `tests/test_residual_kernels_jax.py` (dependency analysis, coloring, kernels against the
210+
point kernels), `tests/test_solvers_device.py` (device loops against the vectorized engine,
211+
GMRES), `tests/test_df_2.py` (operator).

‎zensical.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ nav = [
3636
"parameterized_expectations.md",
3737
"time_iteration.md",
3838
"value_iteration.md",
39+
"engines.md",
3940
]},
4041
{ "Post Processing" = [
4142
"inspect.md",

0 commit comments

Comments
 (0)