autojit-graspa

name

autojit-graspa

description

Automated GPU optimization for gRASPA Monte Carlo molecular simulator. Given a test case (simulation.input + force field + CIF), runs a 3-phase pipeline that produces a specialized, optimized binary with bit-for-bit identical results. Phase 1 removes dead code deterministically, Phase 2 proposes CUDA kernel optimizations one at a time, Phase 3 tackles structural multi-file changes. All changes are gated by an automated correctness checker. Use when a user asks to "optimize gRASPA", "speed up this simulation", or provides a gRASPA test case and wants better performance.

AutoJIT-gRASPA: GPU Optimization Skill for gRASPA

You are an agent helping optimize gRASPA, an NVIDIA GPU-accelerated Monte Carlo simulator for molecular adsorption in nanoporous materials. Your job: given a specific simulation, produce a faster binary that gives bit-for-bit identical results.

When to Use This Skill

Trigger this skill when the user:

• Provides a gRASPA simulation directory (with simulation.input, .def files, .cif) and asks to optimize it
• Says things like "optimize gRASPA", "speed up this simulation", "run the JIT pipeline", "specialize this case"
• Has a gRASPA test case and wants to measure baseline vs optimized performance

Prerequisites Check

Before starting, verify:

1. GPU compiler available — ls /opt/nvidia/hpc_sdk/Linux_x86_64/*/compilers/bin/nvc++ should succeed
2. Original gRASPA source — either the user provides it, or you find it at a path like /path/to/gRASPA/src_clean/ containing axpy.cu, VDW_Coulomb.cu, main.cpp, data_struct.cpp, read_data.cpp, and ~30 header files
3. Test case files — directory containing simulation.input, a molecule .def file, pseudo_atoms.def, force_field_mixing_rules.def, and a framework .cif

If any is missing, ask the user for the location before proceeding.

Workflow Overview

Phase 1 (no AI, deterministic)    Phase 2 (AI, per-kernel)        Phase 3 (AI, structural)
──────────────────────────────    ────────────────────────        ────────────────────────
analyze.py  →  feature_manifest   phase2_cli.sh loops:            run_phase3.py for:
build_graph.py → strip_plan         propose 1 CUDA change             - CUDA graphs
specialize.py → run all strips      apply with Edit                    - kernel fusion
  + correctness-verify each step    compile → run → score              - on-the-fly updates
                                    keep if faster, revert else      each multi-turn debug
Typical: +0-3%                    Typical: +5-10%                   Typical: +0-10%

Expected total speedup: 5-20% depending on case. Best result observed so far: Ar_MgMOF74_UFF at +17.9%.

Setup (One-Time)

Copy the skill's contents to a working directory:

WORK_DIR="$HOME/autojit-graspa-run"          # or user-specified path
mkdir -p "$WORK_DIR"
cp -r <SKILL_ROOT>/scripts/*      "$WORK_DIR/"
cp -r <SKILL_ROOT>/src            "$WORK_DIR/src"         # original gRASPA source
mkdir -p "$WORK_DIR/baselines" "$WORK_DIR/test_cases"
cd "$WORK_DIR"

The scripts expect this directory layout:

<WORK_DIR>/
├── src/                         # Original gRASPA source (read-only)
├── baselines/                   # Reference outputs for correctness checks
├── test_cases/                  # User's simulation inputs
│   └── <CASE_NAME>/
│       ├── simulation.input
│       ├── <molecule>.def
│       ├── pseudo_atoms.def
│       ├── force_field_mixing_rules.def
│       └── <framework>.cif
├── analyze.py, build_graph.py, strip.py, specialize.py    # Phase 1
├── score.py, benchmark.sh                                  # Correctness + timing
├── phase2_cli.sh, run_agent.py                             # Phase 2
├── run_phase3.py                                           # Phase 3
└── (generated during runs:)
    ├── specialized_<CASE>/src/          # Phase 1 output
    ├── phase2_experiment_log_<CASE>.jsonl
    ├── phase3_log_<CASE>.jsonl
    └── best_src_<CASE>/                 # Best version found

Phase 1: Deterministic Dead Code Stripping

No LLM needed. Three sub-steps:

Step 1a: Detect features

python3 analyze.py test_cases/<CASE_NAME>

Produces test_cases/<CASE_NAME>/feature_manifest.json listing which gRASPA features the simulation actually uses (Ewald, DNN, TMMC, CBCFC, Gibbs, tail correction, block pockets, rotation, etc.).

Step 1b: Build code graph + strip plan

python3 build_graph.py src/ test_cases/<CASE_NAME>/feature_manifest.json

Produces src/code_graph.json (function call graph) and test_cases/<CASE_NAME>/strip_plan.json (ordered list of safe code removal operations).

Step 1c: Apply strips with per-step verification

python3 specialize.py test_cases/<CASE_NAME>

For each strip in the plan:

1. Back up source
2. Apply transformation (gut function / replace conditional / remove if-block)
3. Compile with NVC_COMPILE
4. Run the binary
5. Compare output against a freshly-generated baseline via score.py
6. If pass → keep. If fail → revert.

Important: The first time specialize.py runs on a new case, it generates the baseline output automatically and saves it to baselines/<CASE_NAME>.txt. Subsequent verification uses this baseline.

Expected Phase 1 Outcome

Consistently, 9 of 16 strip steps succeed across typical test cases. The same 7 steps usually fail:

• CB/CFC gut (kernel linkage issues)
• Identity swap gut (compile error)
• Dead move branches for Widom / CBCFC / Gibbs volume (RNG or energy drift)
• Block pocket conditional stubs (pattern not present)
• Runtime conditional replacement (pattern not present)

Phase 1 alone typically gives 0-3% speedup — the hot GPU kernels aren't touched. Its real value is shrinking the search space for Phase 2.

Configuring the baseline cycle count

Before running Phase 1, set NumberOfInitializationCycles in simulation.input to a value that gives a measurable baseline (usually 5-10 seconds):

• For small systems (single-atom adsorbate, small framework): try 100,000 cycles
• For medium systems: 5,000-10,000 cycles
• For large systems with Ewald: 5,000 cycles

The original research configs often have millions of cycles — too slow for optimization iteration. Reduce them.

Phase 2: LLM-Driven Micro-Optimizations

This phase is agent-agnostic — the orchestration is in bash + Python, the LLM is interchangeable. Multiple options provided:

Option A: API-based (requires ANTHROPIC_API_KEY)

export ANTHROPIC_API_KEY="sk-ant-..."
GRASPA_SRC=specialized_<CASE>/src python3 run_agent.py \
    --test-case <CASE_NAME> --max-experiments 30

Option B: Claude Code CLI (no API key, uses claude command)

bash phase2_cli.sh --cycles 30 --test-case <CASE_NAME>

Option C: Codex CLI / Other Agents (manual or scripted loop)

The examples/Henrys_coefficient/ case was produced by OpenAI Codex CLI, proving the pipeline doesn't depend on Claude specifically. The recipe for any agent that can read source and apply edits:

# Per experiment (you can wrap this in a loop):

# 1. Backup
cp -r specialized_<CASE>/src .phase2_backup

# 2. Ask your agent to propose ONE optimization. Pass it the experiment log
#    and SKILL.md context. The agent must apply changes via its edit tool
#    and emit `DESCRIPTION: <one sentence>` on a single line. Example with codex:
codex exec --workdir specialized_<CASE>/src "$(cat <<EOF
Read SKILL.md and the recent phase2_experiment_log_<CASE>.jsonl.
Current best time: <T_best>s.
Propose ONE focused CUDA optimization (rules from SKILL.md).
Apply it via the edit tool. Then print exactly one line:
DESCRIPTION: <what you changed>
EOF
)"

# 3. Compile + run + score
bash benchmark.sh <CASE>     # exits 0 on pass, returns JSON with time

# 4. If correct + faster: keep, save best_src_<CASE>/, log JSON entry.
#    Otherwise: restore from .phase2_backup, log a failed entry.

The key contract: any LLM/agent that can produce valid edits and DESCRIPTION: lines is compatible. The correctness gate (score.py) and benchmark loop (benchmark.sh) don't care which model produced the change.

Each cycle does:

1. Back up source
2. Call claude -p "<prompt>" with context: current best time, recent experiment log, optimization priorities, pointers to source files
3. Claude reads source, proposes ONE focused change, applies it via Edit, prints DESCRIPTION: <what>
4. Bash script: compiles (NVC_COMPILE), runs binary, scores via score.py
5. If correct + faster → keep, save to best_src_<CASE>/
6. Otherwise → revert from backup
7. Log JSON entry to phase2_experiment_log_<CASE>.jsonl

Optimization priorities for Phase 2

These are in the prompt. The ones that have historically worked:

Optimization	Where	Effect
Remove BlockEnergy dead init	VDW kernel prologue	+1-5%
cudaStreamSynchronize(0) instead of cudaDeviceSynchronize()	After kernel launches	+0.5-2.5%
OnTheFly VDW kernel (register-based displacement)	For TRANSLATION moves	+3-5%
Shared-memory → warp-shuffle reduction tail	Hottest kernel only	+1%
__launch_bounds__(128, 4)	Hottest kernel only	+0.5-1%
__ldg() for broadcast reads	Where all threads read same address	+0.25%
Dead variable removal	Everywhere	neutral (but clean)

Things that have historically NOT worked (avoid unless new evidence):

• __restrict__ hints (nvc++ often already deduces)
• DEFAULTTHREAD changes (case-dependent; 64 helps CO2-MFI but hurts Ar)
• Reduction tail rewrites on non-hottest kernels
• rsqrt replacements for sqrt (compiler already does this)
• Removing always-true guards like !FF.noCharges

What to tell Claude in each cycle

The phase2_cli.sh script already includes a structured prompt. Key rules it enforces:

• ONE change per experiment — never combine ideas
• Do NOT alter the RNG sequence or control flow logic
• Do NOT modify read_data.cpp, write_data.h, print_statistics.cuh
• Apply changes directly via Edit tool
• Print DESCRIPTION: <one sentence> after applying
• If out of ideas, print DESCRIPTION: NO_MORE_IDEAS

Phase 3: Structural Optimizations

Riskier, multi-file changes with iterative debug loops.

python3 run_phase3.py --task <task-id> --test-case <CASE_NAME>

Pre-defined tasks:

• on-the-fly-translation — eliminate get_new_position<<<1,1>>> kernel launch
• cuda-graphs — stream-capture the kernel sequence
• kernel-fusion-vdw-ewald-init — merge Ewald init kernel into VDW
• dual-stream-overlap — overlap VDW and Ewald on separate streams

When Phase 3 pays off

Only when kernel launch overhead is the dominant bottleneck. Diagnose this by:

• Running the binary under Nsight Systems OR
• Checking if the hot path has many small <<<1, Molsize>>> kernels (Molsize < 8)

For Ar_MgMOF74_UFF (single-atom adsorbate), Phase 3 gave +10% by eliminating deferred get_new_position launches. For CO2_MgMOF74_UFF (large compute grid), Phase 3 gave only +0.06% — compute dominates launches.

The Correctness Gate: score.py

This is the safety net that makes the whole pipeline work. After every optimization, score.py <new_output.txt> <baseline.txt> checks:

1. Exact move counts — translation/rotation/insertion/deletion counts must match baseline (same RNG seed → same moves)
2. Exact molecule counts — final atom counts per pseudo-atom type
3. Energy tolerance — relative 1e-5 on VDW / Coulomb / Ewald / tail
4. Energy drift — per-component absolute value < 3e-5
5. GPU vs CPU drift — not worse than baseline
6. Structure factor consistency — CPU vs GPU max diff < 1e-4

Any failure → auto-revert. Exit codes: 0=pass, 1=fail, 2=parse error.

The pipeline's power comes from this gate — it lets the LLM propose hundreds of changes; only correct AND faster ones stick.

Complete Workflow Example

Given a user who supplies a test case in /some/path/MyCase/:

# 1. Setup working dir
WORK=$HOME/autojit-work
mkdir -p $WORK && cd $WORK
cp -r <SKILL_ROOT>/scripts/* .
cp -r <SKILL_ROOT>/src .
mkdir -p test_cases baselines
cp -r /some/path/MyCase test_cases/

# 2. Reduce cycle count for fast iteration
# Edit test_cases/MyCase/simulation.input:
#   NumberOfInitializationCycles <reduce to ~5000-100000>
#   NumberOfProductionCycles 0

# 3. Phase 1 (all three sub-steps via specialize.py)
python3 specialize.py test_cases/MyCase
# Check output: should say "X/16 succeeded"
# Artifacts: specialized_MyCase/src/, baselines/MyCase.txt

# 4. Phase 2 (20-30 cycles typical)
bash phase2_cli.sh --cycles 30 --test-case MyCase
# Monitor phase2_experiment_log_MyCase.jsonl
# Best source saved to best_src_MyCase/

# 5. (Optional) Phase 3
python3 run_phase3.py --task on-the-fly-translation --test-case MyCase

# 6. Report
# Compare initial baseline time to best_src time
# Show user the delta and which optimizations stuck

Reporting to the User

After each phase, give the user a concise summary:

After Phase 1:

Phase 1 complete: 9/16 strip steps succeeded. Baseline time: . Post-strip time: . Speedup: <X%>. Specialized source at specialized_<CASE>/src/.

After Phase 2 (per cycle):

Cycle N/30: . Compiled: yes. Correct: yes. Time: s (<±X>%). <KEPT|REVERTED>. Best: <T_best>.

After Phase 2 (summary):

Phase 2 complete: <N_kept>/<N_total> experiments kept. Baseline: s → Best: <T_best>s. Total speedup: <X%>. Top-contributing changes: .

Final report:

Total optimization: s → <T_best>s (<X%> speedup). Optimized source at best_src_<CASE>/. Verified correct against baseline across all tests.

Case Studies (See examples/)

Five fully-worked examples are included in this skill, spanning the optimization space and demonstrating that the pipeline is agent-agnostic (4 driven by Claude Code, 1 driven by OpenAI Codex):

Example	Driver	Speedup	Bottleneck	Winning Phase
Ar_MgMOF74_UFF/	Claude Code	+17.9%	Kernel-launch overhead (single-atom adsorbate)	Phase 3 kernel elimination
CO2_MFI/	Claude Code	+12%	Mid-size Ewald compute	Phase 2 micro-opts
CO2_MgMOF74_UFF/	Claude Code	+9.1%	Large Ewald compute grid (5×3×3 supercell)	Phase 2 micro-opts
Henrys_coefficient/	Codex CLI	+8.8%	Widom multiple-trial kernel	Phase 2 micro-opts
NU2000_pX_LinkerRotations/	Claude Code	+3.9%	Special-rotation moves on flexible linker	Phase 3 CUDA Graph capture

• Ar_MgMOF74_UFF — Argon in MgMOF-74 (VDW-only, single atom, no rotation). Phase 3 was the big winner: eliminated deferred <<<1,1>>> kernel launches.
• CO2_MFI — CO2 in MFI zeolite (3-atom rigid molecule, charges, smaller framework). Phase 2 micro-optimizations were dominant.
• CO2_MgMOF74_UFF — CO2 in MgMOF-74 (charges + Ewald, large 5×3×3 supercell). Phase 2 won; Phase 3 stalled because compute dominates launch overhead.
• Henrys_coefficient — Henry's-law calculation (Widom-only, no swaps) on CO2 + MgMOF-74. Driven by OpenAI Codex CLI, not Claude. Validates that any code-execution-capable agent can drive Phase 2 — Codex independently rediscovered the dead-init class of optimizations on the multiple-trial VDW kernel.
• NU2000_pX_LinkerRotations — p-xylene in NU-2000 with linker rotations (the only case with SPECIAL_ROTATION moves and multi-atom rigid linkers). Phase 3 CUDA Graph capture was the key win — this case was where naive Phase 2 optimizations historically failed (catastrophic energy drift), so it's a useful negative example for what NOT to apply.

Each example directory contains:

• input/ — original test case files
• feature_manifest.json, strip_plan.json — Phase 1 artifacts
• phase1_strip_log.jsonl — Phase 1 step-by-step log
• phase2_experiment_log.jsonl — Phase 2 experiments (if run)
• phase3_log.jsonl — Phase 3 tasks (if run)
• baseline_output.txt — reference simulation output
• optimized_src/ — final optimized source code
• REPORT.md — summary of what happened and why

Study these before running on a new case — the same optimizations tend to work.

Common Failure Modes

1. Compile fails after strip — usually because a function was gutted whose signature is referenced elsewhere (e.g., CBCFC kernel). Auto-reverts. Not your problem.

2. Energy drift after strip — moves that are "unused" per the manifest turn out to be still reachable via some edge path. Auto-reverts. Log it.

3. Move counts differ after LLM change — the LLM touched the RNG sequence (e.g., reordered Random.Update() or changed branch conditions). Auto-reverts.

4. Not faster — common and fine. ~70% of LLM proposals don't speed things up or produce measurement noise. Auto-reverts. Keep iterating.

5. Claude CLI returns empty — network/timeout. Backoff and retry.

6. Baseline is stale — if the user changes simulation.input cycle counts, regenerate the baseline: delete baselines/<CASE>.txt and re-run specialize.py.

Compiler Details

gRASPA uses NVIDIA HPC SDK nvc++ (NOT standard CUDA toolkit):

CXX="/opt/nvidia/hpc_sdk/Linux_x86_64/24.5/compilers/bin/nvc++"
NVCFLAG="-O3 -std=c++20 -target=gpu -Minline -fopenmp -cuda -stdpar=multicore"
LINKFLAG="-L/opt/nvidia/hpc_sdk/Linux_x86_64/24.5/cuda/lib64 -L/usr/lib64/ -L/opt/local/lib/gcc11/"

The NVC_COMPILE script in src/ compiles 5 files in parallel then links to nvc_main.x. If the HPC SDK version differs, update both -L paths in NVC_COMPILE.

File Reference

Python scripts (in scripts/):

• analyze.py — Phase 1a feature detection from simulation.input
• build_graph.py — Phase 1b static analysis + strip plan generation
• strip.py — brace-counting transformer used by specialize.py
• specialize.py — Phase 1c orchestrator
• score.py — correctness checker (6 checks, returns JSON)
• run_agent.py — Phase 2 via API (alternative to CLI)
• run_phase3.py — Phase 3 multi-turn agent
• analyze_logs.py — extract ideas from Phase 2 log for Phase 3

Bash scripts (in scripts/):

• benchmark.sh — one-shot compile+run+score
• phase2_cli.sh — Phase 2 via claude CLI (no API key)
• setup_specialize.sh — creates working directory from template
• run_all.sh — batch runner for multiple test cases

Reading the Source

Key hot-path files in src/:

• VDW_Coulomb.cu — the GPU energy kernels (~2000 lines, main optimization target)
• mc_single_particle.h — translation/rotation dispatch + acceptance
• axpy.cu — MC main loop and move selection
• Ewald_Energy_Functions.h — Ewald summation kernels (if charges present)
• data_struct.h — contains #define DEFAULTTHREAD 128 near the top

Files you should almost never edit:

• read_data.cpp — input parser (3500 lines, fragile)
• write_data.h, print_statistics.cuh — output formatting
• main.cpp, main.h — program entry

Tips for Efficient Agent Work

1. Start by reading the existing experiment log if one exists. Don't re-propose ideas that already failed.

2. Use the feature_manifest to understand what the case actually needs. For VDW-only cases, any Ewald-related optimization is pointless. For charge-bearing cases, never disable Ewald.

3. Measure on a reasonable baseline — too short (<3 seconds) gives noisy measurements; too long (>60s) slows iteration. 5-15 seconds is the sweet spot.

4. The hottest kernel dominates — for most cases it's Calculate_Single_Body_Energy_VDWReal in VDW_Coulomb.cu. Optimize that before touching anything else.

5. Respect the correctness gate — if a plausible-looking optimization fails, trust the score.py output. Don't try to relax the tolerances.