Skip to content

fix(gpu): FFT_Bundle host-side FFT fallback to fix Wannier90 segfault on GPU - #8150

Open
chengleizheng wants to merge 6 commits into
deepmodeling:developfrom
chengleizheng:fix/wannier90-gpu-bug
Open

chengleizheng wants to merge 6 commits into
deepmodeling:developfrom
chengleizheng:fix/wannier90-gpu-bug

Conversation

@chengleizheng

Copy link
Copy Markdown
Collaborator

Fix #7770

This is the follow-up announced in #7847, which fixed the LDOS and wfc-read crashes of this issue; this PR fixes the remaining one, the Wannier90 GPU segfault (101_PW_W90, 100_PW_W90 in the issue).

Root cause. In a GPU build only the wavefunction basis (pw_wfc) uses a GPU FFT backend; all rho-family bases always keep a CPU backend. The Wannier90 output runs on the host CPU mirror (stp.sync_cpu() in ctrl_output_pw.cpp) and performs its transforms through the host-memory FFT interface (recip2real/real2recip -> fftxyfor/fftzbac/get_auxg_data/...). The GPU backends (FFT_CUDA/FFT_ROCM) implement only the device-memory 3D-FFT interface; the host-memory virtuals are inherited from FFT_BASE, where they are declared __attribute__((weak)) without a definition, so on ELF their vtable slots are null. The first call of this kind (e.g. get_auxg_data() in PW_Basis_K::recip2real, pw_transform_k.cpp) jumps to address 0 -- the null-pointer segfault reported in the issue.

Fix. When device != "cpu", FFT_Bundle::initfft now creates an additional FFT_CPU backend per active precision (host_fft_float/host_fft_double) that serves the host-memory interface, and all host-memory entry points are routed through get_host_fft_float()/get_host_fft_double(). The device-memory 3D-FFT path (fft3D_forward/backward, get_auxr_3d_data) still uses the device backend, and when device == "cpu" the selectors return the original backends, so CPU builds are unchanged (no extra object is created). The same fallback also covers the identical latent crash in the Berry-phase output (unk_overlap_pw), the ML-KEDF descriptors and DFPT, which all use the host-memory interface of the wfc bundle.

We kept the host fallback instead of routing these calls through the recip2real_gpu device path used in #7847 because these consumers work on host-resident arrays and hold no device copy of the spectra: the device path would add H2D/D2H traffic and buffer-layout work, and it currently requires gamma_only == false and poolnproc == 1. Serving the host interface from an FFT_CPU instance -- exactly what a CPU bundle has always done -- is the minimal fix; the only cost is one extra host FFTW plan set (plus its two work buffers) per wavefunction basis in GPU builds.

Verification:

  • CUDA build and CPU build of this branch compile; CPU results unchanged.
  • GPU: 101_PW_W90 previously segfaulted during the Wannier90 output; it now completes. Against the CPU reference, eigenvalues agree to ~1e-12 eV, |amn| to 5.7e-8, |mmn| exactly. amn/mmn differ only by the intrinsic per-eigenstate gauge phases.
  • 101_PW_W90 stays disabled in CASES_GPU.txt: catch_properties.sh compares amn/mmn element-wise (Re/Im), which cannot accept a gauge rotation; re-enabling the case needs a phase-invariant comparison for these two files as a follow-up.

No INPUT or docs change; the device 3D-FFT path is untouched.

@chengleizheng

chengleizheng commented Oct 15, 2026 •

Copy link
Copy Markdown
Collaborator Author

before the fix, when running the case, we get:

[ RUN      ] 101_PW_W90/
[LAPTOP-G5PS9KT3:42553] *** Process received signal ***
[LAPTOP-G5PS9KT3:42553] Signal: Segmentation fault (11)
[LAPTOP-G5PS9KT3:42553] Signal code: Address not mapped (1)
[LAPTOP-G5PS9KT3:42553] Failing at address: (nil)
[LAPTOP-G5PS9KT3:42553] [ 0] /lib/x86_64-linux-gnu/libc.so.6(+0x42520)[0x76ad36842520]
[LAPTOP-G5PS9KT3:42553] *** End of error message ***
--------------------------------------------------------------------------
Primary job  terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.
--------------------------------------------------------------------------
--------------------------------------------------------------------------
mpirun noticed that process rank 0 with PID 0 on node LAPTOP-G5PS9KT3 exited on signal 11 (Segmentation fault).
......

after the fix, we have the result:

[ RUN      ] 101_PW_W90/
Compare Error: line 2, column 4
 diamond.amn: -0.538560137556
 OUT.autotest/diamond.amn: -0.768120379683
 [ERROR     ]  Fatal Error in catch_properties.sh
[----------] Time elapsed: 13.004 seconds

../integrate/Autotest.sh: line 294: /tmp/abacus-autotest.hVQpii/101_PW_W90/.res: No such file or directory
[ PASSED   ]  0 test cases passed.

Running the GPU version, ETOT = -308.0061503011086 vs -308.0061503011041 in result.ref.
This PR fixes the GPU crash. The remaining W90 ERROR is an expected behaviour(gauge freedom in amn/mmn), not a bug; enabling the test on GPU would need a gauge-invariant comparison.

@mohanchen mohanchen added Bugs Bugs that only solvable with sufficient knowledge of DFT GPU & DCU & HPC GPU and DCU and HPC related any issues labels Oct 15, 2026

@ABACUS-PR-BOT ABACUS-PR-BOT left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Static review of 8e4299c08dcff5237b7010ebf4e5ba7a85e5499e, independently checked by a second reviewer. The remaining single-precision output gap and a regression-test suggestion are inline.

Non-blocking: please add the exact CPU/CUDA configure/build, 101_PW_W90 run and comparison commands used for the reported results, including precision, ENABLE_FLOAT_FFTW, rank count and executable revision, so the manual verification can be reproduced.

git diff --check passed. The pinned governance check exited 0; its warnings were reviewed individually. No builds or runtime tests were run under this static-review request. CI diagnostics were unavailable through the connector (sign-in required).

Comment thread source/source_base/module_fft/fft_bundle.cpp Outdated
Comment thread source/source_base/module_fft/fft_bundle.cpp Outdated
@chengleizheng

Copy link
Copy Markdown
Collaborator Author

At the suggestion of ABACUS-PR-BOT, host_fft_double is now created whenever device != "cpu", independently of double_flag, and its setupFFT() / cleanFFT() / clear() calls are null-guarded instead of being nested under if (double_flag).

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Bugs Bugs that only solvable with sufficient knowledge of DFT GPU & DCU & HPC GPU and DCU and HPC related any issues

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[GPU] PW segfaults: 057_PW_SO_IW (SOC init_wfc), 100_PW_W90 (Wannier90), scf_out_ldos (post-SCF)

3 participants