Repository navigation
Metal: read byref kernel arguments in place - #1013
Merged
Merged
Conversation
Byref kernel arguments arrive in device memory, and were copied to a stack slot for the body, which expects AS 0 pointers. Optimization normally removes that copy, but not when the argument is indexed dynamically, and the AGX back-end miscompiles some kernels where it survives: a closure capturing arrays and a dynamically indexed `SVector` wrote nothing for most threads (JuliaGPU/Metal.jl#1018). Pass arguments that are only loaded from as casts of the device pointer instead, which address-space inference turns into device-memory loads. Arguments with other uses, e.g. passed to a call that also receives stack memory, are still copied.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## tb/metal-pointer-words #1013 +/- ##
==========================================================
+ Coverage 89.52% 89.54% +0.01%
==========================================================
Files 31 31
Lines 6924 6935 +11
==========================================================
+ Hits 6199 6210 +11
Misses 725 725 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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.
add_parameter_address_spaces!copies every byref kernel argument from device memory into a stack slot, because the kernel body expects AS 0 pointers. Optimization usually removes that copy. When the argument is indexed dynamically, the copy survives, and Apple's back-end miscompiles some of these kernels. In JuliaGPU/Metal.jl#1018, a closure capturing two arrays and a 4-element tuple (orSVector) indexed in a loop, and called at two sites, left 511 of 512 outputs unwritten:With
MTL_SHADER_VALIDATION, or without any one of the dynamic indexing, the bounds checks or the second call site, the same kernel is correct.Arguments whose only uses are plain loads, directly or through GEPs and bitcasts, are now passed as an
addrspacecastof the device pointer. TheInferAddressSpacesPassinfinish_ir!turns those into device-memory loads, so no copy is made. This also matches what MSL does with buffer arguments. Arguments with other uses keep the copy. An example is a tuple passed to a non-inlined function that is also called with a stack tuple: its address space can't be inferred, and a device pointer left in AS 0 would be truncated like the pointers in #1012. Volatile and atomic loads also keep the copy.This builds on #1012: a Julia object reference inside a byref argument is now loaded from device memory, which that PR makes safe.
Tested on Julia 1.10, 1.11 and 1.12 with the GPUCompiler Metal tests; the new IR test fails without the change. The Metal.jl reproducer is fixed and tested in JuliaGPU/Metal.jl#1012.