Repository navigation
fix(rocm): let strided copies run past 2^32 elements - #2237
Merged
Merged
Conversation
The ROCm general copy kernels (copy_g_byval, copy_gg_byval, copy_gg_dynamic_nd, copy_gg_dynamic) launched one thread per element on an uncapped int grid, so HIP rejected any strided copy of 2^32 elements or more with hipErrorInvalidConfiguration. copy_general_dynamic also sized its launch by the whole destination, so a one-block DynamicSliceUpdate into a slab that large failed the same way. The grid is now capped at 65535 blocks (rocm::copy_grid_blocks, computed in 64 bits), the dynamic launch covers the copied shape, and each kernel takes a kLoop template argument: below the cap it is the old body, past it a grid-stride loop on a 64-bit counter. One kernel with both paths, or a loop for every size, measured 1.2 to 1.7% slower on a 4M-element transpose, so the split keeps typical copies unchanged within noise. copy_contiguous uses the helper and moves to 64-bit indices above INT32_MAX, so its uint32 counter cannot wrap. tests/rocm_strided_copy_grid.rs adds fast tests past the cap and four ignored 8 GiB tests that failed before this change and pass after it on gfx1151. Recorded as LOCAL_FIXES.md item 42. Closes #2184
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
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.
Summary
copy_g_byval,copy_gg_byval,copy_gg_dynamic_nd,copy_gg_dynamic) now cap the grid at 65535 blocks (rocm::copy_grid_blocks, 64-bit block count) and loop past the cap, so strided copies of 2^32 elements or more no longer fail withhipErrorInvalidConfiguration.copy_general_dynamicsizes its launch by the copied shape instead ofout.size().DynamicSliceUpdatepasses the whole destination asout, so a one-block update into a 2^32-element slab used to launch 2^32 threads.kLooptemplate argument. Below the cap the instantiation is the old one-thread-per-element body; past it the threads stride on a 64-bit counter. A single kernel holding both paths, and a loop for every size, each measured 1.2 to 1.7% slower on a 4M-element transpose, so the split is what keeps typical copies unchanged. It doubles those four kernels' instantiations.copy_contiguoususes the same helper and switches to 64-bit indices aboveINT32_MAX: itsuint32_tcounter adds up to 2^26 per pass and would wrap just under 2^32. This was found by reading the code and was not reproduced.LOCAL_FIXES.mditem 42. It stays in mlxcelverse under the fork policy.Other ROCm launchers with the same uncapped grid (binary, unary/ternary 2-D, reduce init/col, arange, gather/scatter, quantize, sort iota,
copy_col_row) are not on the copy path. They are filed as #2234.Verification (gfx1151)
tests/rocm_strided_copy_grid.rs, fast tests: 42.6M elements, about 2.5 passes of the capped grid, throughcontiguous,concatenateandslice_update_dynamic, compared with a host reference. All 3 pass. With the loop cut to one pass, all 3 fail.--ignoredtests (8 GiB each, 2^32 + 32,768 f16 elements, rows checked on both sides of 2^31 and 2^32 and at the end): all 4 failed on main with "invalid configuration argument" and pass with the fix. Each ran insidescripts/rocm_gpu_guard.sh.AMD_LOG_LEVEL=3confirms each op reaches the kernel it targets.contiguousof a transposed view:[1, 8, 4096, 128]: 147.82 vs 147.99 us[1, 2048, 32, 128]: 388.85 vs 387.40 usconcatenateof two[1, 32, 2048, 128]views: 836.75 vs 833.55 us[1, 1, 32, 128](12 to 14 us) and[1, 8192, 32, 128](past the cap; best run 5455 vs 5423 us) vary more between rounds than between builds.make verify-rocmon e5f4419 (this branch rebased on 2aa5211): OK; 163 test suites, 12194 passed, 0 failed, 403 ignored, smoke included.verify-versions,verify-kernel-dtype-keys,verify-kernel-port-dispatch,verify-llama-compat,verify-fmt,verify-rocm-overlay. Clippy on the new test target passes.Not verified: Metal and CUDA are not available on this host. The change touches only
patches-rocm/and a test that is gated onfeature = "rocm", so the Metal and CUDA paths are not built from it.Closes #2184