lablup / lablup/mlxcel

fix(rocm): hipFFT blocks when too many plans are alive, root cause unknown

Open
#1,876 0 comments 0 reactions 0 assignees View on GitHub
area:core platform:linux priority:medium status:ready type:bug
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

Part of #1801

## Problem / Background

The ROCm FFT port caps its hipFFT plan cache at 8 entries, against the CUDA file's 128, because past roughly a dozen live plans the next `hipfftMakePlanMany64` never returns. The GPU is idle and the calling thread sleeps on a pipe read. The cap bounds the symptom; the cause is not identified, so 8 is an empirical limit rather than a considered design choice, and every FFT-heavy workload (Kokoro text to speech, the Phi-4-multimodal audio front end) pays repeated plan rebuilds for it.

## Current Behavior

`src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/fft.hip:128-141` sets `default_capacity` to 8 on the `LRUBytesKeyCache`, with a comment recording the measurement. Upstream's `mlx/backend/cuda/fft.cu:90-92` uses 128 at the same call site. `LOCAL_FIXES.md` entry 18 carries the same record.

Measured on a Radeon 8060S (gfx1151), ROCm 10.0.0, with a probe cycling 26 transform shapes (rfft, irfft, round trips, and complex fft at lengths 2, 8, 400, 512, 1024, 1200, 2048, batches 3, 4, 512): capacity 8 completes every run, three runs in a row; capacities 16, 24, 32, and 128 all hang, and always at the same position in the sequence rather than on a particular shape. Reordering the cases moves the hang to whatever runs last, and 60 repeats of a single shape never hang, which is what isolates it to the number of live plans rather than the number of executions.

It is not a plain plan limit in hipFFT. A standalone program outside MLX creates 64 plans without trouble, both with hipFFT's own work-area allocation and with `hipfftSetAutoAllocation(handle, 0)` plus a manual work area and an execution per plan. So it is an interaction with how MLX holds or executes them. One relevant detail: a plan outlives its cache entry while a pending execution still holds its `shared_ptr` (the cache stores `std::shared_ptr`, `fft.hip:138`), so the live count can exceed the configured capacity, which is why the cap has to leave headroom. Note also that this port sets `hipfftSetAutoAllocation(handle, 0)` at `fft.hip:180` and allocates the work area per execution from MLX's pool, which upstream's CUDA path does not do in the same way.

`MLX_ROCM_FFT_CACHE_SIZE` raises the cap for anyone measuring it.

## Scope

**In scope:** investigating the block, and `src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/fft.hip` plus `LOCAL_FIXES.md` entry 18 if the cap turns out to be avoidable.

**Out of scope:** changing the transform math or the typed-API port itself, which is verified correct to 1e-5 relative against the CPU stream. Changing `LRUBytesKeyCache`.

## Proposed Solution

Find the cause and, if it is ours, remove the cap. Starting points, in order:

1. Capture where the blocked thread actually is: attach with `gdb -p` and take a backtrace of all threads. The pipe read suggests a runtime waiting on something rather than a GPU stall, so the frame above it is the answer.
2. Check whether rocFFT's runtime kernel compilation is involved, and whether its cache path or a lock is the thing being waited on. `ROCFFT_RTC_CACHE_PATH` and the rocFFT RTC log knobs are the levers.
3. Re-run the probe with `hipfftSetAutoAllocation` left on, so hipFFT owns the work area, to see whether the manual work area at `fft.hip:180` is implicated.
4. Check whether releasing plans eagerly on eviction, rather than when the last `shared_ptr` drops, changes the threshold. That separates "too many handles alive" from "too many handles in the cache".

If the limit is genuinely inside rocFFT, stop at a minimal standalone reproducer and report it.

## Implementation Notes

- **Reuse**: the existing 26-shape probe and `MLX_ROCM_FFT_CACHE_SIZE` are the measurement apparatus; do not write a new harness. `examples/wht_numeric_probe.rs` is the in-tree pattern for a probe binary if one needs to be checked in.
- **Constraints**: gfx1151 and ROCm 10.0.0 are the only measured configuration. Any conclusion should say whether it was checked on another GPU or ROCm version. The hang is a hard hang with no timeout, so every experiment needs an external kill.
- **Edge cases**: the headroom point above means a raised cap is only safe if the live-plan count, not the cache size, is what is bounded. Measure the live count, not the configured capacity.
- **Error handling**: if the cap stays, `check_hipfft_error` cannot help, because the call never returns. Any mitigation has to be a bound on plan count rather than an error path.

## Acceptance Criteria

- [ ] The cause is identified and written down, or, if it is not found, the dead ends are recorded (what was checked, what it ruled out) so the next person does not repeat them.
- [ ] If the block turns out to be avoidable on our side, the cap is raised or removed and the 26-shape probe passes three consecutive runs at the new capacity on gfx1151.
- [ ] If the limit is genuinely in rocFFT, a minimal standalone reproducer exists, the report is filed upstream, and the comment at `fft.hip:128` cites that report by URL instead of saying the cause is unidentified.
- [ ] `LOCAL_FIXES.md` entry 18 is updated to match whatever the outcome is.
- [ ] FFT numerical results are unchanged: rfft, irfft, round trips, and complex fft within 1e-5 relative against the CPU stream at the lengths and batches already covered.

## Verification

```
MLX_ROCM_FFT_CACHE_SIZE=128 cargo test --release --features rocm fft
```

A pass is three consecutive completions of the 26-shape probe at the raised capacity with no hang, plus the existing FFT correctness checks still within 1e-5 relative. If the cap stays at 8, a pass is instead the written-down cause or dead ends plus the upstream report link in the source comment.

## Technical Considerations

Related: #1825 (the FFT implementation this came out of), PR #1861 (the port), #1813 (upstreaming the fork-side local fixes). The hang is a hard hang, so this should not be worked on a machine anyone else is using.

Contributor guide

Open the contributing guide

Research direction

Start with src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/fft.hip, LOCAL_FIXES.md entry 18, and the existing 26-shape probe. Capture all-thread backtraces with gdb, then test the listed auto-allocation, plan-lifetime, and rocFFT runtime hypotheses, externally killing each hang. Done means the cause or ruled-out paths are documented, the cap outcome and local-fix note are updated, and cargo test --release --features rocm fft plus three probe runs verifies the acceptance criteria.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
backend, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.