fix(rocm): get_launch_args clamps the grid without a grid-stride contract
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #1801
## Problem / Background
`get_launch_args` in `src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/kernel_utils.hpp:227` computes `num_blocks` and then caps it at `src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/kernel_utils.hpp:236` with `num_blocks = std::min(num_blocks, 65535)`. It returns that capped grid with nothing in its name, signature, or comment saying the caller's kernel must be grid-stride to cover the rest. A kernel written the ordinary way, one index per thread with an early return past the end, silently leaves every element past `65535 * 256` unwritten, and the output buffer keeps whatever the allocation held. That is the failure shape of lablup/mlxcel#1823, where a branch returned without launching and `quantized_matmul` returned the contents of a fresh allocation.
Nothing calls the helper today, which is why this has not bitten. The kernels in this backend write their own grid-stride loops (`binary.hip:21`, `unary.hip:24`). The two kernels added in lablup/mlxcel#1856 compute their own clamped geometry instead of calling it, and both carry a comment saying why: `hadamard.hip:112` and `sort.hip:418-420` ("the launch geometry below caps the grid, so a one-index-per-thread kernel would silently leave the tail of a large output unwritten"). The next person who reaches for the obvious-looking helper gets the bug.
## Current Behavior
The fork's helper is not upstream's. Upstream CUDA's definition at `mlx/backend/cuda/kernel_utils.cu:33-50` does not clamp x at all (`num_blocks.x = cuda::ceil_div(nthreads, block_dim)`; CUDA's x limit is 2^31 - 1), takes a `max_block_dim` parameter, and honours the `large` flag through `get_2d_grid_dims`. The 65535 constant appears there only as `max_grid_yz_dim` in `get_launch_args_general` (`kernel_utils.cu:53-62`), which folds the y overflow into z rather than dropping it. The fork's version ignores `shape`, `strides`, and `large`, fixes `block_size` at 256, and drops the overflow. Upstream's `hadamard.cu:150` does clamp to 65535, but its kernel is grid-stride (`mlx/backend/cuda/device/hadamard.cuh:61`, `transform += gridDim.x`), so clamp and loop are a matched pair. Metal has no equivalent helper.
## Scope
**In scope:** `src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/kernel_utils.hpp` and an entry in `src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md`.
**Out of scope:** changing any existing kernel's launch geometry. `binary.hip`, `unary.hip`, `hadamard.hip`, and `sort.hip` are already correct and must not be touched.
## Proposed Solution
Make the helper safe to use, or remove it. Three options to weigh rather than a prescription, since the file's future depends on how closely the overlay tracks upstream:
1. Delete both overloads. Nothing calls them, and every kernel here writes its own loop. Cheapest, and it cannot regress.
2. Keep them and make the contract impossible to miss: rename to something that names the requirement (for example `get_grid_stride_launch_args`) and add a comment stating that the returned grid is capped and the kernel must be grid-stride.
3. Fold the overflow into y the way `get_launch_args_general` does, so a non-grid-stride kernel is also correct. Closest to upstream, most code.
Record the decision and its reason so it is not relitigated.
## Implementation Notes
- **Reuse**: whichever option is taken, follow the pattern already used by `hadamard.hip:110-115` and `sort.hip:726-728`, where the clamp and the grid-stride loop sit next to each other with a comment tying them together.
- **Constraints**: this file is a fork overlay copied over the ROCm backend, so a rename must not collide with an upstream symbol the fork may later introduce. Metal and CUDA builds never copy `patches-rocm/`.
- **Edge cases**: `size == 0` yields `num_blocks == 0` and a zero-extent launch; whichever option is taken must not make that throw. `work_per_thread > size` already collapses `adjusted_size` to 1.
## Acceptance Criteria
- [ ] No caller can obtain a silently truncated grid from `kernel_utils.hpp`: either the helper is gone, or its name and comment state the grid-stride requirement, or it covers the full range.
- [ ] `grep -rn "get_launch_args" src/lib/mlx-cpp/patches-rocm/` shows no call site that pairs a capped grid with a non-grid-stride kernel.
- [ ] The decision and its reason are recorded as a new numbered entry in `src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md` (entries currently run to 18), ending with the file's standard "Applies to the fork; to be proposed there" marker.
- [ ] The entry is listed as an upstreaming candidate under #1813, since the defect is fork-side and not present in upstream CUDA.
- [ ] The ROCm build still compiles and the gfx1151 correctness matrix is unchanged.
## Verification
```
cargo build --release --features rocm
make verify
```
Then run the gfx1151 correctness gate added in lablup/mlxcel#1826 and confirm no regression. A pass is a clean build plus an unchanged matrix result; there is no behavior change to observe, because the helper has no callers.
## Technical Considerations
The defect is fork-side, so it belongs in the upstreaming list maintained by #1813 rather than in an upstream MLX report. Related: #1823 (same silent-uninitialized-output failure shape), #1856 (the two kernels that deliberately avoided this helper).
Contributor guide
Research direction
Start with src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/kernel_utils.hpp:227-236, then compare the overflow handling in get_launch_args_general and the grid-stride patterns in hadamard.hip:110-115 and sort.hip:726-728. Record the chosen approach in src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md, run cargo build --release --features rocm and make verify, and confirm the gfx1151 correctness matrix is unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100