jejjohnson / jejjohnson/spectraldiffx
perf: scan-based matrix-vector products for Chebyshev time-stepping
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
When time-stepping a PDE with Chebyshev spatial discretization, the inner loop applies `D @ u` (or `D2 @ u`) hundreds to thousands of times. The current implementation uses a Python loop or manual unrolling, which prevents XLA from fusing the operations across time steps. Using `jax.lax.scan` enables the compiler to optimize the entire trajectory computation as a single fused kernel.
## Current Pattern (Unoptimized)
```python
# Typical time-stepping loop
def rhs(u, t):
du_dx = D @ u # O(N^2) matmul
d2u_dx2 = D2 @ u # O(N^2) matmul
return nu * d2u_dx2 - u * du_dx # Burgers' equation
# Python loop — each iteration is a separate XLA dispatch
u = u0
for _ in range(n_steps):
k1 = rhs(u, t)
k2 = rhs(u + 0.5*dt*k1, t + 0.5*dt)
u = u + dt * k2
```
**Problems:**
- Each `D @ u` is a separate XLA kernel launch
- No cross-step fusion (e.g., the IFFT of step k and FFT of step k+1 could fuse)
- Python loop overhead (~10-100 us per iteration)
- Cannot efficiently differentiate through the loop (for adjoints/gradients)
## Proposed: scan-based time-stepping
```python
def step(u, _):
"""Single RK2 step — compiled as part of the scan body."""
k1 = rhs(u, t)
k2 = rhs(u + 0.5*dt*k1, t + 0.5*dt)
u_new = u + dt * k2
return u_new, u_new # (carry, output)
# Single XLA compilation for the entire trajectory
u_final, u_trajectory = jax.lax.scan(step, u0, xs=None, length=n_steps)
```
**Benefits:**
- **Single XLA compilation**: the entire n_steps trajectory is one fused kernel
- **Cross-step fusion**: XLA can fuse memory-bound operations across steps
- **Reverse-mode AD**: `jax.grad` through `scan` uses O(1) memory (checkpointing)
- **Eliminates Python dispatch overhead**: no per-step host-device sync
### scan for the matrix-vector product itself
For the `D @ u` operation specifically, when D is applied repeatedly in a recurrence, `scan` can express this:
```python
def apply_D_k_times(u, k):
"""Compute D^k @ u via scan (k successive applications)."""
def body(carry, _):
return D @ carry, None
result, _ = jax.lax.scan(body, u, xs=None, length=k)
return result
```
This is useful for higher-order derivatives without precomputing D^k:
- Avoids O(N^3) cost of computing D^k
- Each step is O(N^2), total O(k*N^2)
- XLA fuses the k matmuls into one kernel
### Combining with FFT-based derivatives (#64)
When FFT-based derivatives are available, the scan body becomes:
```python
def step(u, _):
du_dx = cheb_diff_fft(u) # O(N log N) via DCT + recurrence
d2u_dx2 = cheb_diff2_fft(u) # O(N log N) via DCT + recurrence
return u + dt * (nu * d2u_dx2 - u * du_dx), u
```
The scan compiles DCT → recurrence → IDCT → nonlinear term → DCT → ... as a single fused pipeline. This is where the real performance gain comes from: the intermediate arrays never leave GPU registers.
## Back-of-Envelope: Overhead Reduction
| Component | Python loop | scan-based | Savings |
|-----------|------------|------------|---------|
| Kernel launch per step | ~10 us | 0 (fused) | 10 us/step |
| Host-device sync | ~5 us | 0 (fused) | 5 us/step |
| Memory allocation | Per-step | Preallocated | Variable |
| **Total overhead (1000 steps)** | **~15 ms** | **~0 ms** | **15 ms** |
For N=64, the actual matmul takes ~0.5 us, so the 15 us overhead per step is **30x larger than the computation**. Scan eliminates this entirely.
## Implementation Plan
1. Add `solve_ivp_scan` utility that wraps a Chebyshev RHS function in `jax.lax.scan`
2. Provide example time-steppers: Euler, RK2, RK4 as scan-compatible step functions
3. Demonstrate on Burgers' equation (1D Chebyshev + nonlinear advection)
4. Benchmark: Python loop vs scan vs diffrax
## JAX-Specific Notes
- `jax.lax.scan` supports reverse-mode AD via implicit checkpointing
- For very long trajectories, use `jax.checkpoint` to trade compute for memory
- `scan` with `unroll=k` can partially unroll for better instruction scheduling
- The scan body must be a pure function (no side effects) — fits the JAX model
## Acceptance Criteria
- [ ] scan-based time-stepping utility for Chebyshev PDEs
- [ ] Benchmark showing reduced overhead vs Python loop
- [ ] Example: Burgers' equation with Chebyshev + scan
- [ ] Compatible with `jax.grad` for adjoint-based optimization
- [ ] Works with both matrix-based and FFT-based derivatives
## Part of
- Epic #61
- Complements #64 (FFT-based derivatives)
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.