[Bug] Fix reverse-mode ndarray gradients on Vulkan by binding the grad buffer directly.

python/quadrants/lang/misc.py

@@ -836,13 +836,33 @@ def get_host_arch_list():
 
 def is_extension_enabled(ext: Extension) -> bool:
     """
-    Directly returns whether extension is enabled, without needing to
-    pass in current architecture. Also takes into account config, in the case
-    of adstack.
+    Directly returns whether extension is enabled on the CURRENT program, without needing to pass in the current
+    architecture.
+
+    Unlike the arch-level `is_extension_supported`, this also accounts for the config (for adstack) and for the
+    capabilities of the initialized device: on the SPIR-V archs (vulkan, metal), extension support depends on
+    per-device capabilities that only the running device can report.
     """
     arch = impl.current_cfg().arch
     if ext == extension.adstack:
-        return is_extension_supported(arch, ext) and impl.current_cfg().ad_stack_experimental_enabled
+        if not (is_extension_supported(arch, ext) and impl.current_cfg().ad_stack_experimental_enabled):
+            return False
+        if arch in (vulkan, metal):
+            # The on-device adstack size-expression interpreter reads through raw 64-bit buffer addresses, so it
+            # requires physical storage buffers and 64-bit integer arithmetic (see the launch gate in
+            # runtime/gfx/adstack_sizer_launch.cpp).
+            caps = impl.get_runtime().prog.get_device_caps()
+            return bool(caps.get(_qd_core.DeviceCapability.spirv_has_physical_storage_buffer)) and bool(
+                caps.get(_qd_core.DeviceCapability.spirv_has_int64)
+            )
+        return True
+    if ext == extension.data64 and arch in (vulkan, metal):
+        # 64-bit data support on the SPIR-V archs is a property of the device, so the static arch-level table (which
+        # cannot advertise it) is superseded by the device capabilities.
+        caps = impl.get_runtime().prog.get_device_caps()
+        return bool(caps.get(_qd_core.DeviceCapability.spirv_has_int64)) and bool(
+            caps.get(_qd_core.DeviceCapability.spirv_has_float64)
+        )
     return is_extension_supported(arch, ext)
 
 

quadrants/codegen/spirv/spirv_codegen.cpp

@@ -900,10 +900,16 @@ void TaskCodegen::visit(ExternalPtrStmt *stmt) {
       ir_->decorate(spv::OpDecorate, linear_offset, spv::DecorationNoSignedWrap);
     }
   }
-  if (caps_->get(DeviceCapability::spirv_has_physical_storage_buffer)) {
+  if (caps_->get(DeviceCapability::spirv_has_physical_storage_buffer) && !stmt->is_grad) {
     std::vector<int> indices = arg_id;
     // Pick the data or gradient pointer slot of the ndarray argument struct. Without this, reverse-mode AD kernels
     // accumulate into x.data instead of x.grad and host-side gradients stay at zero.
+    //
+    // Gradient slots deliberately do NOT take this physical_storage_buffer path: MoltenVK misreads the second
+    // buffer_device_address member of the args struct (the grad pointer), so the reverse accumulate lands on a
+    // garbage address and host gradients stay zero. Route the gradient through the bound-storage-buffer path
+    // instead (the same path used on backends without physical_storage_buffer, and the reason field gradients work
+    // on Vulkan), which binds the grad ndarray as its own descriptor.
     indices.push_back(stmt->is_grad ? TypeFactory::GRAD_PTR_POS_IN_NDARRAY : TypeFactory::DATA_PTR_POS_IN_NDARRAY);
     spirv::Value addr_ptr = ir_->make_access_chain(ir_->get_pointer_type(ir_->u64_type(), spv::StorageClassUniform),
                                                    get_buffer_value(BufferType::Args, PrimitiveType::i32), indices);

tests/python/conftest.py

@@ -96,7 +96,7 @@ def _warmup() -> qd.i8:
 
 
 @pytest.fixture(autouse=True)
-def wanted_arch(request, req_arch, req_options):
+def wanted_arch(request, req_arch, req_options, req_extensions):
     if req_arch is not None:
         if req_arch in (qd.cuda, qd.amdgpu):
             if not request.node.get_closest_marker("run_in_serial"):
@@ -113,6 +113,13 @@ def wanted_arch(request, req_arch, req_options):
         if "print_full_traceback" not in req_options:
             req_options["print_full_traceback"] = True
         qd.init(arch=req_arch, enable_fallback=False, **req_options)
+
+        # Extension support can depend on the capabilities of the device that `qd.init` just created, so the test
+        # requirements resolve against the initialized program rather than a static arch table.
+        for ext in req_extensions:
+            if not qd.is_extension_enabled(ext):
+                qd.reset()
+                pytest.skip(f"Extension '{ext.name}' is unsupported by the '{req_arch.name}' device on this machine.")
     yield
     if req_arch is not None:
         qd.reset()
@@ -122,7 +129,7 @@ def pytest_generate_tests(metafunc):
     if not getattr(metafunc.function, "__qd_test__", False):
         # For test functions not wrapped with @test_utils.test(),
         # fill with empty values to avoid undefined fixtures
-        metafunc.parametrize("req_arch,req_options", [(None, None)], ids=["none"])
+        metafunc.parametrize("req_arch,req_options,req_extensions", [(None, None, ())], ids=["none"])
 
 
 def _exit_marker_dir():

tests/python/test_ad_dataclass.py

@@ -23,15 +23,12 @@
 
 from tests import test_utils
 
-archs_support_ndarray_ad = [qd.cpu, qd.cuda, qd.amdgpu, qd.metal]
-
-
 # ----------------------------------------------------------------------------
 # qd.ndarray members
 # ----------------------------------------------------------------------------
 
 
-@test_utils.test(arch=archs_support_ndarray_ad, require=qd.extension.adstack)
+@test_utils.test(require=qd.extension.adstack)
 def test_ad_dataclass_ndarray_typed_annotation():
     """dataclass holding qd.ndarrays, passed via typed-dataclass kernel-arg annotation."""
     N = 5
@@ -65,7 +62,7 @@ def compute(s: State):
     np.testing.assert_allclose(a.grad.to_numpy(), b.to_numpy())
 
 
-@test_utils.test(arch=archs_support_ndarray_ad, require=qd.extension.adstack)
+@test_utils.test(require=qd.extension.adstack)
 def test_ad_dataclass_ndarray_template():
     """dataclass holding qd.ndarrays, passed via qd.template()."""
     N = 5
@@ -143,7 +140,7 @@ def compute(s: qd.template()):
 # ----------------------------------------------------------------------------
 
 
-@test_utils.test(arch=archs_support_ndarray_ad, require=qd.extension.adstack)
+@test_utils.test(require=qd.extension.adstack)
 def test_ad_dataclass_tensor_ndarray_backend():
     """dataclass holding qd.tensor(..., backend=NDARRAY) members; ndarray-AD via kernel.grad()."""
     N = 5
@@ -219,7 +216,7 @@ def compute(s: qd.template()):
 # ----------------------------------------------------------------------------
 
 
-@test_utils.test(arch=archs_support_ndarray_ad, require=qd.extension.adstack)
+@test_utils.test(require=qd.extension.adstack)
 def test_ad_dataclass_mixed_ndarray_and_tensor_ndarray_backend():
     """Single dataclass holds one qd.ndarray member and one qd.tensor(NDARRAY) member; verify the kernel can read/write
     both and that gradients flow through both."""
@@ -259,7 +256,7 @@ def compute(s: State):
     np.testing.assert_allclose(a_tens.grad.to_numpy(), b.to_numpy())
 
 
-@test_utils.test(arch=archs_support_ndarray_ad, require=qd.extension.adstack)
+@test_utils.test(require=qd.extension.adstack)
 def test_ad_dataclass_mixed_ndarray_and_field_in_same_class():
     """Single dataclass holds both a qd.ndarray member and a qd.field member. The kernel reads and writes both;
     differentiation is checked through the ndarray path via ``kernel.grad()`` (the field is along for the ride; its