Pure JAX implementation of ormqr for faster QR solves

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``
 

desc/backend.py

@@ -92,17 +92,114 @@
     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 (

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())

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"]