Skip to content

Repository files navigation

ForgeStencil: Autonomous Agents for Stencil Optimization & Deployment

An AI system that closes the loop from automatic research to automatic deployment of stencil optimizations on real scientific and industrial software.

License CUDA Python Reproducible

English | 中文

ForgeStencil pairs two LLM agents that together turn a stencil optimization idea into a verified speedup inside real software — with no human in the loop. It has been run on 100 end-to-end validated applications (oil & gas, electromagnetics, medical imaging, CFD, climate, astrophysics, materials, HPC benchmarks), delivering a median 1.41× / geometric-mean 2.05× end-to-end speedup, every number reproducible against the application's own production GPU code.

🧭 Architecture

flowchart LR
  subgraph KA["Kernel Agent"]
    direction LR
    P[Plan] --> C[Code] --> Pr[Profile] --> P
  end
  subgraph AA["App Agent"]
    direction LR
    H[Locate hotspot] --> F["Forge app-specific operator<br/>(knowledge-base-driven)"] --> V[Verify correctness] --> I[Integrate]
  end
  KA -->|optimized operators| LIB[("kernel/<br/>operator library =<br/>knowledge base")]
  LIB -->|techniques & operators| AA
  AA -->|"single USE_OURS switch"| HR["Harness:<br/>interleaved measure vs<br/>the app's own GPU code"]
  HR -->|validated speedup| RES[("results<br/>registry")]
Loading

📑 Contents

Highlights · Why ForgeStencil? · Quick Start · Results · How it works · Repository layout · FAQ · Limitations · Contributing · Roadmap · Acknowledgments · License · Citation


✨ Highlights

  • 🤖 Two-agent loop — a Kernel Agent researches & generates CUDA stencil operators (Plan → Code → Profile), an App Agent forges app-specific operators on top of this knowledge base and integrates them into real applications.
  • 🔬 Research, not just codegen — the agent discovers optimization strategies (tiling, fusion, layout, occupancy, host-side restructuring), it doesn't just search a fixed space.
  • 🏭 Real applications, not benchmarks — the baseline is the app's own production GPU code, not a simplified reference.
  • 🔒 Auditable measurement protocol — single in-program switch, program's own correctness check, interleaved median, geometric mean over all standard cases. Faking a speedup becomes a system-level error.
  • 📦 Patch-mode, license-clean — no third-party source is bundled; each app ships a provenance record + fetch script + our integration patch.
  • 🧬 Multi-architecture — one operator library with runtime dispatch for A100 (sm_80), H100 (sm_90) and B200 (sm_100); generation-specific re-forge wins are arch-gated so other paths stay byte-identical. See the cross-generation study.
  • ♻️ Reproducible — every reported number is re-runnable from a fresh clone (verified end-to-end on A100).

🆚 Why ForgeStencil?

Stencil optimization has been automated for two decades — but only the code generation. Humans still design the optimization strategy and do the integration. ForgeStencil automates the strategy discovery and the deployment.

Capability DSL / Codegen
(Halide, Devito)
Autotuners
(AN5D, EBISU, DRStencil)
ForgeStencil
Generate optimized kernel code
Discover new optimization strategies (human-designed) ⚠️ (search within a fixed space) (agent explores)
Deploy into a real application (hotspot → forge → integrate)
Verify correctness inside the real program
Auditable e2e speedup vs the app's own production GPU code

🚀 Quick Start

Requirements: an NVIDIA GPU (validated on A100-SXM4-80GB, sm_80), CUDA 12.x, a C++17 host compiler, Python 3.9+ with numpy + cupy, and a git with HTTPS support (a minimal conda git may lack the https helper — use /usr/bin/git).

git clone <this-repo> forge-stencil && cd forge-stencil
python -c "import numpy, cupy; print('deps OK')"     # verify install
🤖 Using a coding agent (e.g. Claude Code)?

Paste this to your agent to set up and run the first demo automatically:

Read the README of <this-repo> (raw URL), then from a clone run the Quick Start operator demo python tools/run.py --stencil star_1 --shape 256 on an idle GPU and report the measured speedup vs the Halide baseline.

1. Measure an operator in ~1 minute (no app download needed)

The fastest way to see ForgeStencil work — builds our kernel and measures it (baselines are optional; absent ones show as null):

python tools/run.py --stencil star_1 --shape 256 --gpu 0
Expected output
  [1/1] star_1 256³ ...
  0.074 ms  [PASS]                     # our kernel, verified vs NumPy reference
  === Halide (single-step) ===
  Stencil   Shape     AMReX(ms)  Ours(ms)   BL(ms)   BL(GC/s)   Status
  star_1    256³           N/A     0.074     0.144    116.63    [PASS]
  Correctness: ALL PASSED

(Here our kernel is 0.074 ms vs Halide 0.144 ms; AMReX/other baselines are optional — see baselines/.)

2. Reproduce a full end-to-end application result

Fetch the real upstream, apply our patch, and re-measure with the app's own correctness check and timer:

cd applications/haccmk
./vendor.sh                                 # fetch upstream at the recorded version
git apply patches/integration.patch         # our USE_OURS injection (from this app dir)
python ../../harness/run_e2e.py --app haccmk --gpu auto
Expected output
[e2e] build [orig] ...  build [ours] ...
[e2e]   round 1/7: orig=0.0027s  ours=0.0001s
[e2e]   ...
[e2e] app-level correctness: True  (builtin check PASS)
[e2e] MEASURED end-to-end GEOMEAN speedup over 1 case(s): 36.703x

(Verified from a fresh clone on A100. The absolute ratio varies with GPU/node state; correctness is invariant.)

Run any of the 100 apps the same way — see applications/README.md. To drive the agents themselves, see agents/README.md.


📊 Results (audited)

All figures come directly from the result registry (results/integration_registry.json) and are reproducible via the harness. We report the full distribution.

End-to-end, across 100 validated applications

Metric Value
Median end-to-end speedup 1.41×
Geometric mean 2.05×
Range 0.998× – 97.7×
≥ 1.05× 89%
≥ 1.20× 73%
≥ 1.50× 43%
≥ 3× 21%
≥ 10× 7%

Operator-level performance

Roofline (baseline-independent). Measured with Nsight Compute, ForgeStencil's memory-bound stencil kernels reach 74–85% of the A100's 2039 GB/s peak DRAM bandwidth — near the practical hardware ceiling for bandwidth-bound stencils. (box_2 and diamond_ts4 are compute/reuse-bound and sit in a different regime.)

vs the best baseline at each shape — at the same precision. At every (stencil, shape) we take the strongest applicable public result, comparing like with like: single-step vs single-step (Halide, Devito), temporal-blocking vs temporal-blocking (EBISU), and fp16 vs fp16 — where star_1/star_2 use fp16 on cubic shapes, Halide is run at fp16 too.

Stencil (precision) Cases Geomean vs best baseline Leads on
star_1 (7-pt, fp16/fp32) 9 1.63× 9 / 9†
star_2 (13-pt, fp16/fp32) 7 1.92× 7 / 7
box_1 (27-pt, fp32) 5 3.08× 5 / 5
box_2 (fp32) 3 4.60× 3 / 3
diamond_1 (fp32) 3 3.08× 3 / 3
varcoeff_star_1 (fp32) 5 1.51× 5 / 5
Overall 32 ≈2.16× 32 / 32

Numbers are the median of 7 robust passes; the bar is set by both Halide and Devito. Crucially, star_1/star_2 win even when Halide runs at the same fp16 precision (fp16-vs-fp16: 1.75× over the 8 cubic shapes) — the advantage is a genuine same-precision implementation win, not a precision artifact (Halide barely benefits from fp16; our kernels exploit the halved bandwidth). On temporal-blocking stencils vs EBISU (temporal blocking): diamond_ts4 1.82× (win), diamond_ts8 0.73× (loss). †star_1 128³ is an exact tie (1.00×); all other shapes lead. Reproduce with tools/run.py after fetching baselines via baselines/vendor_baselines.sh.

Across GPU generations (A100 → H100 → B200)

The operator library carries runtime architecture dispatch (sm_80 / sm_90 / sm_100): the same kernel/ sources pick arch-gated variants where a generation-specific re-forge won, and the A100 paths stay byte-identical otherwise. The cross-generation study asked: do agent-forged kernels transfer, and when does re-forging pay?

Transition What the HW relaxed Outcome Re-forge capture
A100 → H100 (BW +1.68×) compute, barely bandwidth kernels arrive still bandwidth-saturated (DRAM 77–90%) → near-optimal transfer box_2 (compute-bound): +6–7% via config re-tuning (kernel routing + launch bounds)
H100 → B200 (BW +2.4×) bandwidth ≫ memory-issue rate kernels fall off the roofline (DRAM 78–90% → 44–65%), B200/H100 geomean 1.51× vs the 2.39× bandwidth ratio box_1: structural deep-prefetch, +7.0% @512³ / +5.1% @1024³

The vs-Halide advantage holds across generations (geomean 2.16× on A100 → 2.38× on H100). Full story, per-round forge logs, and raw data: docs/CROSS_GEN_REFORGE.md, docs/FORGE_LOG_H100.md, docs/FORGE_LOG_B200.md, results/cross_gen/. Measurement caveat: B200 numbers come from a preemptible cloud pod — same-pod A/B ratios are the result; absolute milliseconds are node/clock-dependent.

How to read the end-to-end numbers

The speedup mostly reflects how much optimization headroom the upstream application had — not the raw strength of a single kernel.

  • When the application's own GPU code has structural bottlenecks (too many tiny kernel launches, low occupancy, per-step host synchronization), an end-to-end restructuring can yield large gains. Examples (win source labeled): pathfinder 97.7×, bspline_vgh (QMCPACK B-spline) 58.4×, haccmk (HACC short-range force) 36.2× — mostly structural / launch-fusion wins.
  • When the application is already highly tuned (vendor SDK kernels, hand-rolled supergrid kernels), results are typically 1.0–1.5×. Examples: sw4lite 1.00× (parity), hpgmg_fv 1.02×, cloverleaf 1.01×.
  • Operator-level forging gains land in between, e.g. gprMax/FDTD Maxwell 2.47×.

A 1.0× parity result indicates the upstream was already at the hardware limit, and our integration preserved correctness without regressing.


🔧 How it works

ForgeStencil never bundles third-party source. For each integrated application it ships:

applications/<app>/
├── PROVENANCE.md         # upstream repo URL + sha256 of the exact bytes + result
├── vendor.sh             # fetches the upstream source from its original location
├── patches/*.patch       # OUR integration diff: adds a USE_OUR_<APP> switch
└── e2e/{manifest.json, build.sh}
kernel/<app>_ours.cuh     # OUR optimized operator

The harness runs the original and optimized paths interleaved inside the same binary (toggled by a single env switch), checks the program's own correctness, reads the program's own timer, and reports the geometric mean speedup over all of the application's standard cases.

The auditable measurement protocol

  1. Baseline is the application's own, production-used GPU implementation.
  2. Original and optimized paths differ by exactly one switch; the original path is byte-for-byte identical to upstream.
  3. Correctness comes from the program's own check (bit-exact or its own tolerance).
  4. Timing uses the program's own timer; original/optimized are interleaved over multiple rounds and the median is taken.
  5. The reported number is the geometric mean over all of the application's standard cases.

If any condition fails, the result is marked not validated / blocked and does not count.


📁 Repository layout

Path Contents
kernel/ Operator library — optimized CUDA stencil kernels (*_ours.cuh, stencil_kernel.cu) and the Python dispatch layer.
tools/ run.py — the operator-level benchmark driver (builds our kernel, compares vs baselines).
harness/ The end-to-end measurement protocol (run_e2e.py, measure_*.py, provenance.py, manifest_schema.json).
agents/ Kernel Agent + App Agent orchestration loops and their operating rules — see agents/README.md.
applications/ 100 per-application reproduction packages (provenance + fetch script + patches + harness config). No upstream source. See applications/README.md.
operator_bench/ Per-application operator-proxy benchmarks (bench_ours.cu) used by the App Agent workflow.
baselines/ Runner scripts + vendor_baselines.sh fetch script for third-party stencil baselines. No baseline source bundled. See baselines/README.md.
results/ The audited integration result registry; results/cross_gen/ holds the H100/B200 cross-generation measurement data.
docs/ Studies: cross-generation re-forge (CROSS_GEN_REFORGE.md) + distilled per-round forge logs for H100/B200.

❓ FAQ

vendor.sh fails with "remote-https is not a git command"

Your git is a minimal conda build without the HTTPS remote helper. Use the system git: PATH=/usr/bin:$PATH ./vendor.sh, or conda install -c conda-forge git.

The upstream clone times out / disconnects

External git can be flaky. vendor.sh retries automatically; just re-run it. For large repos, a warm DNS/proxy helps. Fetches are shallow where possible.

My measured speedup differs from the README / registry

Absolute speedups depend on GPU model, clocks, and node load (a shared, busy node inflates the baseline). Correctness is invariant; the ratio is reproducible on an idle, clock-locked A100. Lock clocks with nvidia-smi -lgc <freq> and run on an unloaded GPU for tight numbers.

How do I pick the GPU?

run_e2e.py --gpu auto picks the idlest GPU; tools/run.py --gpu <id> takes an explicit index. Use nvidia-smi to find a free device.

Do I need to build the baselines / AMReX to get a result?

No. tools/run.py degrades gracefully — absent baselines show as null and the NumPy CPU reference is used for correctness. Fetch baselines only when you want the comparison numbers.


🚧 Limitations

  • Validation coverage per architecture. The operator library carries runtime dispatch for A100 (sm_80), H100 (sm_90) and B200 (sm_100), with operator-level results measured on all three (see the cross-generation study). The 100 end-to-end application packages are validated on A100; re-running them on H100/B200 is on the roadmap.
  • Speedup ≠ kernel strength. Large end-to-end gains usually reflect structural headroom in the upstream app, not a universally faster kernel; against already-tuned code results are typically 1.0–1.5×.
  • CPU-only upstreams. A few apps have no GPU version; those are cross-arch comparisons (hardware-dominated) and are marked research, never validated.
  • LLM non-determinism. The agents are driven by an LLM; a given run's path is not bit-reproducible, though the measured artifacts (patches, results) are.
  • Correctness is the program's own check. Where an app ships a weak self-check, our correctness guarantee is only as strong as that check.
  • Responsible use. ForgeStencil optimizes performance only; it preserves the application's numerical semantics and never alters scientific outputs beyond the program's own correctness tolerance.

🤝 Contributing

Contributions are welcome — bug reports, new application integrations (patch-mode only, never bundle upstream source), and verified upstream URLs for the few apps still marked UNKNOWN. See CONTRIBUTING.md.


🎯 Roadmap

  • Cross-generation study — A100 → H100 → B200 operator-level measurement + re-forge (docs)
  • Operator-level roofline / achievable-bandwidth analysis (Nsight Compute, results)
  • End-to-end application re-validation on H100/B200
  • Broaden per-case operator baselines (Devito, EBISU, DRStencil, ConvStencil, AN5D)
  • Publish agent autonomy & cost metrics (iterations, interventions, tokens per app)
  • More application integrations (community-driven — open an issue to propose one)
  • Bilingual README (English / 中文)

Priorities are shaped by community needs — open an issue to weigh in.


🙏 Acknowledgments

ForgeStencil stands on the upstream applications and stencil-optimization projects it measures and compares against, including HeCBench, AMReX, Halide, Devito, EBISU, ConvStencil, DRStencil, FlashFFTStencil, and the many scientific codes listed in applications/. Each is used under its own license; none is redistributed here.


📄 License

Apache-2.0 (permissive, includes a patent grant) — see LICENSE. ForgeStencil does not redistribute third-party upstream source code; see THIRD_PARTY_LICENSES.md for how upstream applications and benchmark baselines are referenced.


📚 Citation

If you use ForgeStencil in your research, please cite it (see CITATION.cff):

BibTeX
@software{forgestencil,
  title        = {ForgeStencil: Autonomous Agents for Stencil Optimization and Deployment},
  author       = {The ForgeStencil Authors},
  year         = {2026},
  url          = {https://github.com/OpenBMB/ForgeStencil}
}

About

No description, website, or topics provided.

Resources

Contributing

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages