Summary
Reverse-mode automatic-differentiation (AD) into an ndarray produces zero (garbage) host gradients on
MoltenVK (the only Vulkan implementation on Apple platforms). The straightforward workaround -- bind each
gradient ndarray as its own storage buffer instead of reading its device address -- cannot support any real
reverse-mode kernel, because bound buffers are capped at Metal's 31 buffer binding slots. This issue records
the root cause and the two rejected attempts so the dead end is not re-explored.
Symptom
On the vulkan arch on Apple (MoltenVK), a qd.kernel that accumulates into arr.grad during the reverse
pass leaves the host-visible gradient at zero. The forward pass and the primal (arr) reads are correct.
CPU / CUDA / AMDGPU / native Metal are all correct.
Root cause
Each ndarray kernel argument is passed as a struct whose members include two buffer_device_address
(physical-storage-buffer) pointers: the data pointer (1st address member) and the gradient pointer (2nd
address member). SPIR-V codegen loads the chosen pointer from the args struct and does OpConvertUToPtr.
MoltenVK's SPIRV-Cross -> MSL translator misreads the second buffer_device_address member of that
struct. The data pointer (1st member) translates correctly; the gradient pointer (2nd member) resolves to a
garbage address, so the reverse accumulate lands nowhere and the host gradient stays zero.
Verified independent of:
- struct member offset (relaid the args struct; grad member still misread),
- member alignment (forced 16-byte alignment; still misread, and desynced host/device layouts),
- buffer type (routed via a plain SSBO instead of physical-storage-buffer for the read; no change).
Native Metal advertises spirv_has_physical_storage_buffer too, but its (separate, newer) SPIRV-Cross reads
the 2nd member correctly -- native Metal gradients are fine on the physical path.
Why binding the gradient as its own buffer does not work
The natural workaround is to skip the physical path for gradients and instead bind each gradient ndarray as
its own storage buffer (a descriptor), the same path backends without physical-storage-buffer already use.
This produces correct gradients on MoltenVK for small kernels, but each bound gradient ndarray consumes
one Metal buffer binding slot, and Metal caps at 31 slots ([[buffer(0)]] .. [[buffer(30)]]).
Real reverse-mode kernels blow past 31. Proven with Genesis's rigid-body backward kernel, which fails MSL
compilation outright and aborts the process:
[metal_device.mm:get_mtl_library] cannot compile metal library from source (msl_bytes=1897122):
program_source:419:2767: error: 'buffer' attribute parameter is out of bounds: must be between 0 and 30
(This surfaced only on a cold cache; a warm offline cache loads the pre-compiled metallib and never
recompiles from MSL, which is why it looked intermittent / parallel-only at first.)
So the bound-buffer route is a ceiling, not a fix: it works for toy AD tests and cannot scale to a real
reverse kernel. The physical-storage-buffer path is precisely the mechanism that removes the 31-slot limit
(device addresses live inside the args struct and cost zero binding slots); abandoning it for gradients
re-introduces the very limit it exists to avoid.
Rejected attempts
Attempt A (over-broad). In TaskCodegen::visit(ExternalPtrStmt), gate the physical path off for every
gradient on every physical-storage-buffer backend:
if (caps_->get(DeviceCapability::spirv_has_physical_storage_buffer) && !stmt->is_grad) { /* physical */ }
else { /* bound buffer */ }
Rejected: besides the 31-slot ceiling above, this also regresses native Metal, which read the gradient
correctly on the physical path -- it needlessly moved Metal gradients onto bound buffers and overflowed the
31-slot limit on large kernels.
Attempt B (scoped to MoltenVK). Add a device capability spirv_grad_ptr_needs_bound_buffer, set only on
MoltenVK (in vulkan_device_creator.cpp, #if defined(__APPLE__) where PSB is enabled), and gate the
detour on it so native Metal and native Vulkan keep the correct physical path:
const bool grad_via_bound_buffer = stmt->is_grad && caps_->get(DeviceCapability::spirv_grad_ptr_needs_bound_buffer);
if (caps_->get(DeviceCapability::spirv_has_physical_storage_buffer) && !grad_via_bound_buffer) { /* physical */ }
Rejected: this only prevents the collateral native-Metal regression. On MoltenVK itself the gradient still
rides the bound-buffer path, still capped at 31 slots, so it still cannot compile a real reverse kernel. It
does not fix MoltenVK; it just localizes the failure.
What a real fix requires
Make MoltenVK read the gradient device address (2nd buffer_device_address struct member) correctly on the
physical-storage-buffer path -- unlimited, zero binding slots. That is a MoltenVK / SPIRV-Cross defect we do
not build and cannot patch directly. Options: emit SPIR-V that MoltenVK's translator handles correctly
(none found so far -- offset / alignment / buffer-type all ruled out), or an upstream MoltenVK fix.
Until then, reverse-mode ndarray AD is unsupported on MoltenVK (vulkan arch on Apple). Native Metal,
native Vulkan, CPU, CUDA, and AMDGPU are unaffected.
PR #801 (which took the bound-buffer route) is closed in favor of this record.
Summary
Reverse-mode automatic-differentiation (AD) into an ndarray produces zero (garbage) host gradients on
MoltenVK (the only Vulkan implementation on Apple platforms). The straightforward workaround -- bind each
gradient ndarray as its own storage buffer instead of reading its device address -- cannot support any real
reverse-mode kernel, because bound buffers are capped at Metal's 31 buffer binding slots. This issue records
the root cause and the two rejected attempts so the dead end is not re-explored.
Symptom
On the
vulkanarch on Apple (MoltenVK), aqd.kernelthat accumulates intoarr.gradduring the reversepass leaves the host-visible gradient at zero. The forward pass and the primal (
arr) reads are correct.CPU / CUDA / AMDGPU / native Metal are all correct.
Root cause
Each ndarray kernel argument is passed as a struct whose members include two
buffer_device_address(physical-storage-buffer) pointers: the data pointer (1st address member) and the gradient pointer (2nd
address member). SPIR-V codegen loads the chosen pointer from the args struct and does
OpConvertUToPtr.MoltenVK's SPIRV-Cross -> MSL translator misreads the second
buffer_device_addressmember of thatstruct. The data pointer (1st member) translates correctly; the gradient pointer (2nd member) resolves to a
garbage address, so the reverse accumulate lands nowhere and the host gradient stays zero.
Verified independent of:
Native Metal advertises
spirv_has_physical_storage_buffertoo, but its (separate, newer) SPIRV-Cross readsthe 2nd member correctly -- native Metal gradients are fine on the physical path.
Why binding the gradient as its own buffer does not work
The natural workaround is to skip the physical path for gradients and instead bind each gradient ndarray as
its own storage buffer (a descriptor), the same path backends without physical-storage-buffer already use.
This produces correct gradients on MoltenVK for small kernels, but each bound gradient ndarray consumes
one Metal buffer binding slot, and Metal caps at 31 slots (
[[buffer(0)]]..[[buffer(30)]]).Real reverse-mode kernels blow past 31. Proven with Genesis's rigid-body backward kernel, which fails MSL
compilation outright and aborts the process:
(This surfaced only on a cold cache; a warm offline cache loads the pre-compiled metallib and never
recompiles from MSL, which is why it looked intermittent / parallel-only at first.)
So the bound-buffer route is a ceiling, not a fix: it works for toy AD tests and cannot scale to a real
reverse kernel. The physical-storage-buffer path is precisely the mechanism that removes the 31-slot limit
(device addresses live inside the args struct and cost zero binding slots); abandoning it for gradients
re-introduces the very limit it exists to avoid.
Rejected attempts
Attempt A (over-broad). In
TaskCodegen::visit(ExternalPtrStmt), gate the physical path off for everygradient on every physical-storage-buffer backend:
Rejected: besides the 31-slot ceiling above, this also regresses native Metal, which read the gradient
correctly on the physical path -- it needlessly moved Metal gradients onto bound buffers and overflowed the
31-slot limit on large kernels.
Attempt B (scoped to MoltenVK). Add a device capability
spirv_grad_ptr_needs_bound_buffer, set only onMoltenVK (in
vulkan_device_creator.cpp,#if defined(__APPLE__)where PSB is enabled), and gate thedetour on it so native Metal and native Vulkan keep the correct physical path:
Rejected: this only prevents the collateral native-Metal regression. On MoltenVK itself the gradient still
rides the bound-buffer path, still capped at 31 slots, so it still cannot compile a real reverse kernel. It
does not fix MoltenVK; it just localizes the failure.
What a real fix requires
Make MoltenVK read the gradient device address (2nd
buffer_device_addressstruct member) correctly on thephysical-storage-buffer path -- unlimited, zero binding slots. That is a MoltenVK / SPIRV-Cross defect we do
not build and cannot patch directly. Options: emit SPIR-V that MoltenVK's translator handles correctly
(none found so far -- offset / alignment / buffer-type all ruled out), or an upstream MoltenVK fix.
Until then, reverse-mode ndarray AD is unsupported on MoltenVK (
vulkanarch on Apple). Native Metal,native Vulkan, CPU, CUDA, and AMDGPU are unaffected.
PR #801 (which took the bound-buffer route) is closed in favor of this record.