FlashPDE

中文说明

FlashPDE is a PyTorch-integrated library of fused Triton operators for structured-grid PDE residuals in scientific machine learning. It provides 14 hand-derived differentiable operators used by 17 benchmark configurations in 1D--3D. Each operator is exposed through torch.autograd.Function, so existing models can replace eager finite-difference residual evaluation without changing the neural architecture or optimizer.

Highlights

• 14 fused PDE operators and 17 runnable benchmark configurations.
• Forward stencils, discrete-adjoint backward passes, and boundary-gradient corrections implemented in Triton.
• MLP and CNN field generators through a common PyTorch workflow.
• Eager PyTorch FD, compiled-model, and coordinate-autograd baselines.
• Reproducible gradient verification, throughput, fixed-epoch, T2S, memory, and kernel-scaling experiments.

Installation

FlashPDE requires an NVIDIA GPU. The paper environment uses Python 3.10, PyTorch 2.9.1, and Triton 3.5.1.

Conda environment used for reproduction

git clone https://github.com/factnn/FlashPDE.git
cd FlashPDE
conda env create -f environment.yml
conda activate flashpde

Existing PyTorch environment

Install a CUDA-enabled PyTorch build suitable for your system, then run:

git clone https://github.com/factnn/FlashPDE.git
cd FlashPDE
pip install -e .

The package also provides two command-line entry points:

flashpde-run --help
flashpde-verify --help

Operator and Case Coverage

One operator can serve multiple cases when their discrete residual structure is identical.

Operator ID	Representative equation	Benchmark case(s)
1d_steady	Steady Burgers	burgers_1d_steady
1d_unsteady	Unsteady Burgers	burgers_1d_unsteady
1d_heat	1D heat/diffusion	diffusion_1d
2d_compressible	Compressible Euler	sod_1d
2d_poisson	2D Poisson	poisson_2d
2d_transport	Advection--diffusion	transport_2d, diffusion_2d, heat_2d
2d_wave	2D wave	wave_2d
2d_allencahn	Allen--Cahn	allen_cahn
2d_ns_steady	Steady Navier--Stokes	ldc_2d
2d_ns_unsteady	Unsteady Navier--Stokes	tgv_2d
3d_ns_steady	Steady Navier--Stokes	ldc_3d
3d_ns_unsteady	Unsteady Navier--Stokes	tgv_3d, tgv_3d_smooth
3d_poisson	3D Poisson	poisson_3d
3d_heat	3D heat	heat_3d

Quick Start

from cases.ldc_2d.physics import loss_triton, make_context, make_mlp

ctx = make_context("cuda")
model = make_mlp().cuda()
loss = loss_triton(model, ctx)
loss.backward()

Representative MLP backends are mlp_vanilla, mlp_canpinn, mlp_compile, and mlp_triton. CNN variants are also available for compatible cases. Run python run.py --help for the complete backend list.

Reproducing the Experiments

All commands below should be run from the repository root. Use --gpu N to select a device.

Gradient verification

Compare all 14 Triton operators with eager PyTorch FD in float64:

python verify_gradients.py --gpu 0

This regenerates:

• VERIFICATION.md, a human-readable table;
• verification_results.json, the same results in machine-readable form.

The committed verification report records 14/14 passing operators on an NVIDIA A100.

Forward--backward throughput

python run.py --case ldc_2d --backend all --track 1 --gpu 0

Track 1 performs warm-up iterations followed by repeated timed iterations and reports latency, throughput, and peak allocated GPU memory. The timed region includes model forward, PDE residual, loss, and backward evaluation, but excludes the optimizer update.

Fixed-epoch training

python run.py \
  --case ldc_2d \
  --backend mlp_canpinn,mlp_triton \
  --track 2 \
  --max-epochs 10000 \
  --gpu 0

Time to solution (T2S)

Each case defines its default convergence threshold. Multi-seed evaluation can be launched with:

python run.py \
  --case ldc_2d \
  --backend mlp_canpinn,mlp_triton \
  --track 2 \
  --runs 5 \
  --max-epochs 300000 \
  --gpu 0

Use --threshold VALUE to override the case threshold. Track 2 stores histories, checkpoints, aggregate statistics, and plots under output/<case>/ unless --out-dir is specified.

Kernel scaling

python scaling.py --case ldc_2d --gpu 0

Scaling uses random field tensors without a neural network and stores results in output/scaling/scaling_<case>.npy. Run the command once per case; the grid sequence is defined in SCALING_GRIDS in scaling.py.

All registered cases

python run.py --case all --backend mlp_canpinn,mlp_triton --track 1 --gpu 0

This iterates over all 17 cases. Full convergence runs are expensive; use case-specific commands or the supplied multi-GPU scripts for long experiments.

Repository Layout

cases/                 PDE definitions, models, thresholds, and metrics
kernels/               14 Triton operator implementations
engine/                throughput and convergence experiment runners
plots/                 solution plotting utilities
solvers/               reference numerical solvers used by selected cases
verify_gradients.py    unified float64 operator verification
scaling.py             kernel-only scaling experiments
run.py                 unified benchmark entry point

Results

• Gradient verification
• Benchmark summary

Performance depends on the GPU, software versions, grid shape, and first-run Triton autotuning. Compilation and autotuning overhead should not be mixed with steady-state throughput measurements.

Scope and Limitations

FlashPDE currently targets regular Cartesian grids and NVIDIA GPUs supported by Triton. Adding a new PDE residual requires a derived discrete adjoint and a corresponding operator implementation. The library accelerates PDE evaluation; end-to-end gains are smaller when neural-network execution dominates.

Citation

If you use FlashPDE, please cite our paper (arXiv:2607.18020):

@misc{zang2026flashpde,
  title         = {{FlashPDE: A Drop-in Fused Triton Operator Library for Neural PDE Solvers}},
  author        = {Peiyu Zang and Bosen Xie and Ruoxiang Xu and Yongqiang Cai},
  year          = {2026},
  eprint        = {2607.18020},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2607.18020}
}

GitHub can also generate citation metadata from CITATION.cff.

License

FlashPDE is released under the MIT License.