[Bug] Fix reverse-mode ndarray gradients on Vulkan by binding the grad buffer directly.#801
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 160f30c446
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…d buffer directly. Reverse-mode autodiff over an `ndarray` argument produced all-zero gradients on the Vulkan/MoltenVK backend: the forward pass and the gradient value were computed correctly, but the accumulate never reached the host `.grad` buffer. Root cause: on backends with `spirv_has_physical_storage_buffer` the ndarray data and gradient pointers are passed as buffer_device_address `u64`s embedded in the args struct, and MoltenVK misreads the SECOND such pointer member (the gradient slot) - it resolves to a garbage address, so the reverse accumulate no-ops. The data slot (first pointer) reads fine, and quadrants' own Metal backend reads the identical SPIR-V correctly; only MoltenVK's SPIRV-Cross mishandles it. This is independent of the member's offset or alignment (verified: a 16-aligned grad slot still fails) and of the buffer type (uniform vs storage). Route the gradient ndarray through the bound-storage-buffer descriptor path instead of the physical-address path - the same path already used on backends without physical_storage_buffer, and the reason field gradients already work on Vulkan. The data pointer keeps the physical-address path. This sidesteps the MoltenVK struct-member defect with no change to the args layout or the data path, and no regression on Metal/CUDA/CPU. With this fix every backend supports ndarray reverse-mode autodiff, so the ndarray autodiff tests no longer need the `archs_support_ndarray_ad` arch restriction: drop the list and let the tests run on the full expected-arch matrix (the tests fail on Vulkan without this fix).
160f30c to
e9ba11a
Compare
…nd skip tests accordingly.
5c3237d to
6d2021a
Compare
|
Closing: the bound-buffer route this PR takes is capped by Metal's 31 buffer binding slots and cannot compile a real reverse-mode kernel (Genesis's rigid backward already overflows it). Root cause and both rejected attempts documented in #802. A real fix requires MoltenVK to read the gradient device address correctly on the physical-storage-buffer path. |
Issue: #
Stacked on #800 -> #799 -> #796 -> #795. Review/merge those first; this PR is based on #800 (which adds Metal to the ndarray-autodiff test matrix) and additionally enables Vulkan once the gradient-buffer bug below is fixed. Base retargets down the stack as each lands.
Brief Summary
copilot:summary
Walkthrough
copilot:walkthrough