[Test] Gate extension-requiring tests on actual device capabilities instead of a static arch table.

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)
 
 

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/test_utils.py

@@ -247,33 +247,32 @@ def exclude_arch_platform(arch, system, exclude):
             if exclude_arch_platform(req_arch, curr_system, exclude):
                 continue
 
-            if not all(_qd_core.is_extension_supported(req_arch, e) for e in require):
-                continue
-
             if qd.extension.adstack in require:
                 options["ad_stack_experimental_enabled"] = True
 
+            # Extension support can depend on the capabilities of the device, which does not exist yet at collection
+            # time: the requirements are carried along the parametrization and resolved against the initialized
+            # program by the `wanted_arch` fixture, which skips when one is unsupported.
             current_options = copy.deepcopy(options)
+            required_extensions = list(require)
             for feature, param in zip(_test_features, req_params):
                 value = param.value
-                required_extensions = param.required_extensions
-                if current_options.setdefault(feature, value) != value or any(
-                    not _qd_core.is_extension_supported(req_arch, e) for e in required_extensions
-                ):
+                if current_options.setdefault(feature, value) != value:
                     break
-            else:  # no break occurs, required extensions are supported
-                parameters.append((req_arch, current_options))
+                required_extensions.extend(param.required_extensions)
+            else:  # no break occurs, the feature options are consistent
+                parameters.append((req_arch, current_options, tuple(required_extensions)))
 
         if not parameters:
-            marks.append(pytest.mark.skip(reason="No all required extensions are supported"))
+            marks.append(pytest.mark.skip(reason="No supported archs after exclusions"))
         else:
             marks.append(
                 pytest.mark.parametrize(
-                    "req_arch,req_options",
+                    "req_arch,req_options,req_extensions",
                     parameters,
                     ids=[
                         f"arch={arch.name}-{i}" if len(parameters) > 1 else f"arch={arch.name}"
-                        for i, (arch, _) in enumerate(parameters)
+                        for i, (arch, _, _) in enumerate(parameters)
                     ],
                 )
             )