solvers: real spectral FFT + TPP space-time passes, retargeted onto the D1 arbiter - #307
Merged
Merged
Conversation
…alect port
Make the ts-spectral-opt solver real and buildable end-to-end.
Kernels (TargetHooks/{CPU,AMD,NVIDIA}/StockhamRadix4.*): replace the
twiddle-less single-butterfly placeholders with a complete mixed-radix
(radix-4 + radix-2 tail) Stockham autosort FFT, identical math across all
three backends. CPU verified vs naive fp64 DFT; AMD proven on gfx1151
(benchmarks/spectral/hip/stockham_gfx_harness.hip, verdict=pass N=64..1M,
~96 GFLOP/s). Symbols now ts_stockham_r{4,2}_{cpu,nvidia,amd} + a runtime
driver ts_fft_stockham_{...}; LowerToTargetIR emits them (dropped the wrong
gfx94x/sm90 arch-suffixed names) and tags tessera.target_ir.arbiter_op.
LegalizeSpectral: dynamic axes no longer fabricate a bogus radix-4 stage —
they defer to the runtime driver via dynamic_shape/dynamic_axes; removed a
stray std::reverse that ordered stages tail-first.
Dialect ported to LLVM 22 so ts-spectral-opt builds standalone: fix the
CMake tablegen (LLVM_TARGET_DEFINITIONS, binary include dir, MLIROptLib/
Func link), the .td (StrAttr, Pure, drop invalid trait/assemblyFormats, add
a Plan TypeDef), and the dialect C++ glue (generated classes at global
scope). Rewrote the lit fixtures to use a real plan op; 10/10 lit pass.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…22 build Turn the tpp-space-time stub passes into real transforms and make the solver buildable/testable standalone. - HaloInfer: derive per-dim halo from the actual access pattern (grad order -> radius; stencil radius from the kernel operand's shape); structured tpp.halo array (was a hardcoded "1,1,0" string). - DistributeHalo: materialise real tpp.halo.exchange ops (new dialect op) in front of each halo consumer, carrying widths/mesh-axes/overlap; fusion-aware. - VectorizeTPP: compute real vector width + tile shape from field shape/dtype. - LowerTPPToTargetIR: annotate ops with hardware-free Target-IR call symbols (cpu/nvidia/amd) + tessera.target_ir.arbiter_op; keep lowered.bc.masked. - LegalizeSpaceTime: default+validate stencil scheme/order and time.step scheme -> stages/order/dt; unknown scheme is a hard error. - FuseStencilTime: group sibling stencils reading the same field (the %Hx/%Hy case) into one shared halo exchange (union halo). - time.step gets NoTerminator so it parses (its region needed an undefined tpp.yield); added a CPU stencil target hook (ts_stencil_grad_cpu). Numerical sentinel: benchmarks/correctness_microbench.cpp (tpp-correctness) implements the linearised shallow-water semantics and checks gradient 2nd-order convergence, periodic-BC correctness, and mass/energy conservation + a traveling-wave match (verdict=pass). Standalone tessera-tpp-opt driver added; 13/13 lit pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…biter Point both solvers' lower-*-to-target-ir seams at the Workstream-C D1 arbiter (Decision #28): the Target-IR call symbols the passes emit (ts_fft_stockham_cpu, ts_stencil_grad_cpu) become first-class F4-gated arbiter candidates. - candidate.py: additive register_op_kind() + _OP_KIND_VERIFY + a verify_by_reference() helper so solver op-kinds (with their own numpy reference, not a fusion_core region) plug into the same enumerate -> F4-gate -> select pipeline. Existing op-kinds unchanged. - emit/spectral_candidates.py: OP_SPECTRAL_FFT, SpectralFFTRegion (reference = numpy.fft). CPU + ROCm candidates run the real shipped Stockham kernels via ctypes (CPU host-portable; ROCm on gfx1151 via the new ts_fft_stockham_amd_hostptr wrapper). - emit/tpp_candidates.py: OP_TPP_STENCIL, StencilGradRegion (reference = periodic central-difference). CPU candidate runs the shipped stencil kernel. Proven: CPU FFT + stencil match numpy through the arbiter; the ROCm FFT lane runs the real gfx1151 kernel and F4-passes vs numpy.fft; a wrong higher-tier candidate is F4-rejected. 13 new tests + 16 existing arbiter tests pass; mypy clean on the changed files. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 85c6c25e3c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+50
to
+54
| if (backend == "nvidia") | ||
| return r4 ? "ts_stockham_r4_nvidia" : "ts_stockham_r2_nvidia"; | ||
| if (backend == "amd") | ||
| return r4 ? "ts_stockham_r4_amd" : "ts_stockham_r2_amd"; | ||
| return r4 ? "ts_stockham_r4_cpu" : "ts_stockham_r2_cpu"; |
There was a problem hiding this comment.
When a static FFT length has an odd factor, e.g. 192 = 4^3 * 3, pickRadixSequence still emits a radix-3/5/7/residual stage, but this mapping annotates every non-4 stage as ts_stockham_r2_*. The shipped r2 kernels execute only two-point butterflies, so the lowered pipeline produces an incorrect partial FFT instead of rejecting/defering that size or calling a real radix-3/5/7 implementation.
Useful? React with 👍 / 👎.
| // Collect first — we mutate operands while walking. | ||
| SmallVector<Operation *, 8> consumers; | ||
| m.walk([&](Operation *op) { | ||
| if (op->hasAttr("tpp.halo") && !op->hasAttr("tpp.halo.distributed")) |
There was a problem hiding this comment.
If -tpp-distribute-halo runs more than once, the tpp.halo.exchange ops inserted by the first run also have tpp.halo and lack tpp.halo.distributed, so this collection treats them as consumers on the next run. Because their operand is the original field, the later “already fed by exchange” check does not fire, and the pass stacks another exchange in front of the generated one, duplicating halo communication despite the pass being intended to be idempotent.
Useful? React with 👍 / 👎.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Makes the spectral (FFT) and TPP (space-time stencil) solvers real and buildable end-to-end, then retargets both onto the Workstream-C D1 candidate arbiter (Decision #28) so their Target-IR call symbols become first-class, F4-gated arbiter candidates.
Commits
spectral: real Stockham FFT kernels + dynamic-N legalize + LLVM-22 dialect portstockham_gfx_harness.hip, N=64..1M, ~96 GFLOP/s).LegalizeSpectral: dynamic axes defer to the runtime driver instead of fabricating a bogus radix-4 stage; removed a straystd::reverse.ts-spectral-optbuilds standalone; 10/10 lit pass.tpp: make all 7 space-time passes real + correctness sentinel + LLVM-22 buildHaloInfer(real per-dim halo from access pattern),DistributeHalo(materialisestpp.halo.exchange),VectorizeTPP,LowerTPPToTargetIR(hardware-free Target-IR call symbols +tessera.target_ir.arbiter_op),LegalizeSpaceTime(scheme/order validation),FuseStencilTime(shared union-halo exchange).correctness_microbench.cpp(tpp-correctness): linearised shallow-water — 2nd-order gradient convergence, periodic-BC, mass/energy conservation, traveling-wave match (verdict=pass). Standalonetessera-tpp-opt; 13/13 lit pass.arbiter: retarget spectral FFT + TPP stencil onto the D1 candidate arbitercandidate.py: additiveregister_op_kind()+verify_by_reference()so solver op-kinds (with their own numpy reference, not a fusion_core region) plug into the same enumerate → F4-gate → select pipeline. Existing op-kinds unchanged.emit/spectral_candidates.py(OP_SPECTRAL_FFT, reference =numpy.fft; CPU + ROCm run the real Stockham kernels via ctypes) andemit/tpp_candidates.py(OP_TPP_STENCIL, reference = periodic central-difference).Proof
numpy.fft; a wrong higher-tier candidate is F4-rejected.🤖 Generated with Claude Code