diff --git a/CHANGELOG.md b/CHANGELOG.md index 389fc8d989..90d81f632f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,8 @@ Changelog Performance Improvements -- Speeds up the ``"qr"`` trust-region subproblem and Newton-step solves in the least-squares optimizers by reusing the Jacobian QR factorization across the Levenberg-Marquardt parameter sweep. On ``jax >= 0.10.0`` this uses ``qr_multiply`` to additionally avoid forming ``Q`` explicitly; on older versions a fallback preserves the same results. +- Speeds up the ``"qr"`` trust-region subproblem and Newton-step solves in the least-squares optimizers by reusing the Jacobian QR factorization across the Levenberg-Marquardt parameter sweep, and by using ``qr_multiply`` to apply ``Q`` without forming it explicitly. +- Adds a pure-JAX ``qr_multiply`` fallback (a blocked Householder / compact-WY implementation) for ``jax < 0.10.0``, so the above ``Q``-avoidance speedup is available on the currently pinned JAX with no jaxlib rebuild (~1.5-2x faster than forming ``Q`` on CPU for tall Jacobians, larger on GPU). - Adds ``surf_batch_size`` kwarg to Boozer and omnigenous field compute variables, to allow for tuning of the memory usage when computing these quantities by choosing how many surfaces to simultaneously compute. - Also adds ``surf_batch_size`` as an additional kwarg to ``make_boozmn_output``, as well as the objectives ``Omnigenity`` and ``QuasisymmetryBoozer`` diff --git a/desc/backend.py b/desc/backend.py index e893d540f5..362ad901a8 100644 --- a/desc/backend.py +++ b/desc/backend.py @@ -92,17 +92,114 @@ def _diag_to_full(d, e): if Version(jax.__version__) >= Version("0.10.0"): from jax.scipy.linalg import qr_multiply else: + # Implements `jax.scipy.linalg.qr_multiply` without the ``ormqr`` primitive + # (which requires a jaxlib rebuild), so it runs on any installed jaxlib. The + # Householder reflectors ``Q`` are applied to ``c`` via the blocked CWY/UT + # transform of Puglisi (1992) and Joffrain & Low (2006) without ever forming + # ``Q``, on large tall systems this beats forming ``Q`` (``orgqr``) and, on GPU, + # cuSOLVER's ``ormqr``. + + # Ported from https://github.com/jax-ml/jax/pull/36575. The `DotAlgorithmPreset` + # used there only mattered for float32 throughput, which DESC does not use, so + # plain matmuls are used here. Block sizes are the upstream A100-tuned values. + + def _householder_multiply(a, taus, c, *, transpose=False): + """Apply the reflectors in ``a``/``taus`` to ``c`` from the left, one block. + + Forms ``Q = I - V T^{-1} V^H`` via the identity ``T^{-1} + T^{-H} = V^H V``, + recovering the triangular ``T^{-1}`` by correcting the diagonal. + """ + m, k = a.shape + # V: unit lower-trapezoidal (reflectors below the diagonal, unit diagonal). + V = jnp.where( + jnp.tril(jnp.ones((m, k), bool), -1), a, jnp.eye(m, k, dtype=a.dtype) + ) + diag_correction = (1 / taus) if transpose else (1 / taus).conj() + diag_correction = jnp.expand_dims(diag_correction, -1) * jnp.eye( + k, dtype=a.dtype + ) + Vh = V.conj().swapaxes(-1, -2) + # solve_triangular reads only the relevant triangle, so passing the full + # Gram matrix V^H V (minus the diagonal correction) recovers T^{-1}. + T_inv = Vh @ V - diag_correction + z = solve_triangular(T_inv, Vh @ c, lower=transpose) + with jax.default_matmul_precision("highest"): + return c - V @ z + + def _blocked_householder_multiply(a, taus, c, *, left, transpose): + """Apply Q (or Q^H) to c in blocks (block sizes tuned on A100).""" + if not left: # c @ Q == (Q^H @ c^H)^H + ct = c.conj().swapaxes(-1, -2) + out = _blocked_householder_multiply( + a, taus, ct, left=True, transpose=not transpose + ) + return out.conj().swapaxes(-1, -2) + + if a.ndim > 2: # batch dims via vmap, keep the core logic 2-D + fn = functools.partial( + _blocked_householder_multiply, left=True, transpose=transpose + ) + return vmap(fn)(a, taus, c) + + m = a.shape[0] + k = taus.shape[0] + if k == 0: # no reflectors -> Q is the identity + return c + # Balances the per-block V^H V cost against the number of sequential blocks. + esize = a.dtype.itemsize + hi_limit = 4096 * max(1, 8 // esize) + mid_limit = 4096 * esize + nb = min(k, 256 if m <= hi_limit else 128 if m <= mid_limit else 64) + + blocks = range(0, k, nb) + for j0 in blocks if transpose else reversed(blocks): + c = c.at[j0:, :].set( + _householder_multiply( + a[j0:, j0 : j0 + nb], + taus[j0 : j0 + nb], + c[j0:, :], + transpose=transpose, + ) + ) + return c + @functools.partial(jit, static_argnames="mode") def qr_multiply(a, c, mode="right"): - """Fallback for ``jax.scipy.linalg.qr_multiply`` (added in JAX 0.10.0).""" - Q, R = qr(a, mode="economic") + """Pure-JAX drop-in for ``jax.scipy.linalg.qr_multiply``; returns `(CQ, R)`. + + ``a = Q @ R`` is the (economic) QR factorization. For ``mode="right"`` + returns ``c @ Q`` (with 1-D ``c`` treated as a length-``M`` row vector); for + ``mode="left"`` returns ``Q @ c``. ``R`` has shape ``(min(M, N), N)``. + """ + m, n = a.shape + k = min(m, n) + # mode="raw" returns the packed reflectors (transposed) plus tau factors, + # via the existing geqrf primitive -- no new primitive / jaxlib rebuild. + h, taus = jnp.linalg.qr(a, mode="raw") + packed = h.swapaxes( + -1, -2 + ) # (M, N): lower triangle = reflectors, upper = R + R = jnp.triu(packed)[:k, :] + # When m <= n geqrf's last reflector (row m-1) is the identity (tau == 0). + # Drop it statically -- the economic Q is unchanged and we avoid 1/0. + n_refl = k - 1 if m <= n else k + V = packed[:, :n_refl] + taus = taus[:n_refl] + c1d = c.ndim == 1 + if mode == "right": - # 1-D c (all DESC uses) matches the old Q.T @ c; c @ Q keeps - # higher-dim c consistent with qr_multiply rather than silently wrong - cq = Q.T @ c if c.ndim == 1 else c @ Q - else: - cq = Q @ c - return cq, R + cq = _blocked_householder_multiply( + V, taus, c[None, :] if c1d else c, left=False, transpose=False + ) + cq = cq[:, :k] # economic Q has min(M, N) columns + return (cq[0] if c1d else cq), R + + C = c[:, None] if c1d else c + pad = jnp.zeros((m - k, C.shape[1]), C.dtype) + cq = _blocked_householder_multiply( + V, taus, jnp.vstack([C, pad]), left=True, transpose=False + ) + return (cq[:, 0] if c1d else cq), R from jax.scipy.special import gammaln from jax.tree_util import ( diff --git a/tests/test_backend.py b/tests/test_backend.py index 3cf96ec130..2584497859 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -3,7 +3,19 @@ import numpy as np import pytest -from desc.backend import _lstsq, jax, jnp, put, root, root_scalar, sign, vmap +from desc.backend import ( + _lstsq, + jax, + jnp, + put, + qr, + qr_multiply, + root, + root_scalar, + sign, + solve_triangular, + vmap, +) @pytest.mark.unit @@ -140,3 +152,79 @@ def test_lstsq(): np.testing.assert_allclose( _lstsq(A, b), np.linalg.lstsq(A, b, rcond=None)[0], rtol=1e-6 ) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "m, n", + [ + (100, 20), # tall + (20, 100), # wide + (50, 50), # square + (600, 260), # tall, and has more than nb columns + ], +) +@pytest.mark.parametrize("cond", [None, 1e8]) +def test_qr_multiply(m, n, cond): + """Test qr_multiply matches forming Q explicitly.""" + + def _qr_multiply_ref(a, c, mode="right"): + """Reference qr_multiply that forms Q explicitly.""" + Q, R = qr(a, mode="economic") + if mode == "right": + cq = Q.T @ c if c.ndim == 1 else c @ Q + else: + cq = Q @ c + return cq, R + + rng = np.random.default_rng(seed=0) + k = min(m, n) + + if cond is None: + # gaussian matrices are usually well conditioned + A = rng.standard_normal((m, n)) + else: + # create some ill conditioned matrix using reverse SVD + # this is still full rank + U = np.linalg.qr(rng.standard_normal((m, k)))[0] + V = np.linalg.qr(rng.standard_normal((n, k)))[0] + A = (U * np.logspace(0, -np.log10(cond), k)) @ V.T + b = rng.standard_normal(m) + print( + f"Running {m=} {n=} {cond=}, actual condition number is {np.linalg.cond(A):.3e}" + ) + + # mode="right" with 1D c is Q.T@b + Qtb, R = qr_multiply(A, b, mode="right") + Qtb_ref, R_ref = _qr_multiply_ref(A, b, mode="right") + assert R.shape == (k, n) + np.testing.assert_allclose(R, R_ref, rtol=1e-12, atol=1e-12 * np.abs(A).max()) + np.testing.assert_allclose(Qtb, Qtb_ref, rtol=1e-10, atol=1e-10) + + # mode="right" with 2D c is c@Q + C = rng.standard_normal((3, m)) + CQ, _ = qr_multiply(A, C, mode="right") + np.testing.assert_allclose(CQ, _qr_multiply_ref(A, C, "right")[0], atol=1e-10) + + # mode="left" is Q@c, with c=I recovers Q + Q, _ = qr_multiply(A, np.eye(k), mode="left") + np.testing.assert_allclose(Q, _qr_multiply_ref(A, np.eye(k), "left")[0], atol=1e-10) + np.testing.assert_allclose(Q.T @ Q, np.eye(k), atol=1e-10) + np.testing.assert_allclose(Q @ R, A, atol=1e-10 * np.abs(A).max()) + y = rng.standard_normal(k) + Qy, _ = qr_multiply(A, y, mode="left") + np.testing.assert_allclose(Qy, _qr_multiply_ref(A, y, "left")[0], atol=1e-10) + + # solve A@x = b + if m >= n: + x = solve_triangular(R, Qtb) + x_ref = solve_triangular(R_ref, Qtb_ref) + else: + # for wide A, use the QR of A.T + Q1, R1 = qr_multiply(A.T, np.eye(k), mode="left") + Q1_ref, R1_ref = _qr_multiply_ref(A.T, np.eye(k), mode="left") + x = Q1 @ solve_triangular(R1.T, b, lower=True) + x_ref = Q1_ref @ solve_triangular(R1_ref.T, b, lower=True) + x_np = np.linalg.lstsq(A, b, rcond=None)[0] + np.testing.assert_allclose(x, x_ref, rtol=1e-8, atol=1e-8 * np.abs(x_np).max()) + np.testing.assert_allclose(x, x_np, rtol=1e-8, atol=1e-8 * np.abs(x_np).max()) diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 233e47cda7..fad69672fb 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -1079,7 +1079,7 @@ def test_constrained_AL_lsq(): ctol=ctol, x_scale="auto", copy=True, - options={}, + options={"tr_method": "svd"}, ) V2 = eq2.compute("V")["V"] AR2 = eq2.compute("R0/a")["R0/a"]