jejjohnson / jejjohnson/spectraldiffx

Epic: Chebyshev Performance Optimizations

Open
#61 0 comments 0 reactions 0 assignees View on GitHub
epic
Dominant language
Python
Stars
1
Forks
0
PR merge metrics
No merged PRs in 30d

Description

## Summary

The current Chebyshev implementation uses dense matrices and dense linear algebra throughout. While mathematically clean and correct, several well-known techniques can dramatically improve performance — especially for the Helmholtz solver (O(N^3) per call) and differentiation (O(N^2) per call).

## Sub-Issues

| # | Title | Speedup | Effort | Priority |
|---|-------|---------|--------|----------|
| #62 | Pre-factored LU for Helmholtz solver | 2-10x for repeated solves | Small | P0 |
| #64 | FFT-based Chebyshev differentiation | O(N^2) -> O(N log N) | Medium | P1 |
| #65 | Scan-based matrix-vector products for time-stepping | Reduced overhead, XLA fusion | Small | P1 |
| #66 | Ultra-spherical / banded Chebyshev solvers | O(N^3) -> O(N) solves | Large | P2 |

## Current Bottlenecks

### Differentiation: O(N^2) dense matmul
\`\`\`python
# Current: D @ u where D is (N+1) x (N+1) fully dense
du_dx = D @ u # O(N^2) per derivative
d2u_dx2 = D2 @ u # O(N^2), D2 precomputed
\`\`\`

### Helmholtz solver: O(N^3) dense LU per call
\`\`\`python
# Current: fresh LU decomposition every call
A = D @ D - alpha * I # O(N^3) matmul
u = jnp.linalg.solve(A, b) # O(N^3) LU
\`\`\`

| N | Matrix size | Dense LU flops | Time @ 1 TFLOP/s |
|---|------------|----------------|-------------------|
| 32 | 33 x 33 | ~72K | 72 ns |
| 64 | 65 x 65 | ~549K | 0.5 us |
| 128 | 129 x 129 | ~4.3M | 4.3 us |
| 256 | 257 x 257 | ~34M | 34 us |
| 512 | 513 x 513 | ~270M | 270 us |

The O(N^3) scaling becomes dominant for N > 128, especially in time-stepping loops where solve() is called thousands of times.

## References

- Trefethen, "Spectral Methods in MATLAB" (2000), Chapters 6-8
- Olver & Townsend, "A fast and well-conditioned spectral method" (2013) — ultra-spherical
- Driscoll, Hale & Trefethen, "Chebfun and numerical quadrature" (2014)
- Burns et al., "Dedalus: A flexible framework for numerical simulations with spectral methods" (2020)

## Depends On

- #28 (Chebyshev extensions) — these optimizations complement the feature additions

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.