Merge PR #2905

crates/sam2/examples/gpu_matmul.rs

@@ -67,10 +67,10 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) li:
 "#;
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
-    let adapter = pollster::block_on(
-        instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() }),
-    )?;
+    // wgpu's environment variables apply (`WGPU_BACKEND`, `WGPU_POWER_PREF`, ...), as for SAM 2's own instance
+    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle().with_env());
+    let power_preference = wgpu::PowerPreference::from_env().unwrap_or(wgpu::PowerPreference::HighPerformance);
+    let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference, ..Default::default() }))?;
     println!("adapter: {:?}", adapter.get_info().name);
     let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { required_limits: adapter.limits(), ..Default::default() }))?;
     let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: None, source: wgpu::ShaderSource::Wgsl(SHADER.into()) });

crates/sam2/src/gpu.rs

@@ -99,12 +99,37 @@ pub fn half_weights() -> bool {
 /// The process's one GPU device, once something has asked for it.
 static DEVICE: std::sync::OnceLock<std::result::Result<Arc<Gpu>, String>> = std::sync::OnceLock::new();
 
+/// The adapter SAM 2 refused and its instance, kept for the life of the process, as the device made on it used to be.
+/// Dropping them crashed the tests with a segmentation fault on WSL2's OpenGL driver (Mesa 23.2 on D3D12). At most one:
+/// [`Gpu::new`] runs once.
+static REFUSED: std::sync::OnceLock<(wgpu::Instance, wgpu::Adapter)> = std::sync::OnceLock::new();
+
+/// SAM 2's wgpu instance reads wgpu's environment variables (`WGPU_BACKEND`, `WGPU_VALIDATION`,
+/// `WGPU_ALLOW_UNDERLYING_NONCOMPLIANT_ADAPTER`…), as the render context's does, so both choose among the same adapters.
+fn instance_descriptor() -> wgpu::InstanceDescriptor {
+    wgpu::InstanceDescriptor::new_without_display_handle().with_env()
+}
+
+/// Why SAM 2 does not use an adapter, if it does not: a software one (the CPU path is faster), or one without compute
+/// shaders (OpenGL ES 3.0, for one), where every kernel would fail validation. Checked before a device is made there.
+fn refusal(name: &str, backend: wgpu::Backend, device_type: wgpu::DeviceType, downlevel: wgpu::DownlevelFlags) -> Option<String> {
+    if device_type == wgpu::DeviceType::Cpu {
+        return Some(format!("only a software adapter ({name}) is available"));
+    }
+    (!downlevel.contains(wgpu::DownlevelFlags::COMPUTE_SHADERS)).then(|| format!("the GPU adapter {name} ({backend:?}) has no compute shaders"))
+}
+
 impl Gpu {
     /// The process's one GPU device, created on first use. Several devices at once upset some Vulkan loaders
     /// (a crash when two threads create or drop them together), and one is all SAM 2 needs: models come and go,
     /// the device stays. An error (no GPU) is remembered too.
     pub fn shared() -> Result<Arc<Gpu>> {
-        DEVICE.get_or_init(|| Gpu::new().map(Arc::new).map_err(|e| e.to_string())).clone().map_err(Error::Compute)
+        // The message alone is kept, so the error says "SAM 2 compute:" once.
+        let keep = |e: Error| match e {
+            Error::Compute(m) => m,
+            e => e.to_string(),
+        };
+        DEVICE.get_or_init(|| Gpu::new().map(Arc::new).map_err(keep)).clone().map_err(Error::Compute)
     }
 
     /// Whether the shared device has been looked for yet, and found: `None` before anything has asked
@@ -114,16 +139,17 @@ impl Gpu {
         DEVICE.get().map(|d| d.is_ok())
     }
 
-    /// A device on the best real GPU (a software adapter is refused: the CPU path is faster).
+    /// A device on the best real GPU (high performance unless `WGPU_POWER_PREF` says otherwise; see [`refusal`] for the
+    /// adapters the CPU path runs instead of).
     fn new() -> Result<Gpu> {
-        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
-        let adapter = pollster::block_on(
-            instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() }),
-        )
-        .map_err(|e| Error::Compute(format!("no GPU adapter: {e}")))?;
+        let instance = wgpu::Instance::new(instance_descriptor());
+        let power_preference = wgpu::PowerPreference::from_env().unwrap_or(wgpu::PowerPreference::HighPerformance);
+        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference, ..Default::default() }))
+            .map_err(|e| Error::Compute(format!("no GPU adapter: {e}")))?;
         let info = adapter.get_info();
-        if info.device_type == wgpu::DeviceType::Cpu {
-            return Err(Error::Compute(format!("only a software adapter ({}) is available", info.name)));
+        if let Some(why) = refusal(&info.name, info.backend, info.device_type, adapter.get_downlevel_capabilities().flags) {
+            let _ = REFUSED.set((instance, adapter));
+            return Err(Error::Compute(why));
         }
         let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
             label: Some("effectcraft sam2"),

crates/sam2/src/gpu/tests.rs

@@ -452,3 +452,38 @@ fn half_precision_weights_give_the_same_product_to_half_precision() {
     assert_close(&got, &want, 2e-3 * (k as f32).sqrt(), "f16 vs f32 weights");
     assert!(got != want, "the weights really were rounded to half precision");
 }
+
+/// The instance reads wgpu's environment variables, as the render context's does (#2900): with
+/// `WGPU_ALLOW_UNDERLYING_NONCOMPLIANT_ADAPTER=1` wgpu lists Mesa's Dozen, the only Vulkan driver on WSL2. Run in a
+/// child process, which has them set.
+#[test]
+fn the_instance_reads_the_wgpu_environment() {
+    const CHILD: &str = "EFFECTCRAFT_TEST_SAM2_WGPU_ENV_CHILD";
+    if std::env::var_os(CHILD).is_none() {
+        let out = std::process::Command::new(std::env::current_exe().unwrap())
+            .args(["--exact", "gpu::tests::the_instance_reads_the_wgpu_environment", "--nocapture"])
+            .env(CHILD, "1")
+            .env("WGPU_ALLOW_UNDERLYING_NONCOMPLIANT_ADAPTER", "1")
+            .env("WGPU_BACKEND", "vulkan")
+            .output()
+            .unwrap();
+        let stdout = String::from_utf8_lossy(&out.stdout);
+        assert!(out.status.success() && stdout.contains("1 passed"), "{stdout}\n{}", String::from_utf8_lossy(&out.stderr));
+        return;
+    }
+    let desc = instance_descriptor();
+    assert!(desc.flags.contains(wgpu::InstanceFlags::ALLOW_UNDERLYING_NONCOMPLIANT_ADAPTER), "{:?}", desc.flags);
+    assert_eq!(desc.backends, wgpu::Backends::VULKAN);
+}
+
+/// An adapter without compute shaders (a GL one, here) is refused before a device is made on it, with one error
+/// instead of a validation error per kernel; so is a software one.
+#[test]
+fn adapters_without_compute_shaders_are_refused() {
+    use wgpu::{Backend, DeviceType, DownlevelFlags};
+    assert_eq!(refusal("GPU", Backend::Vulkan, DeviceType::DiscreteGpu, DownlevelFlags::all()), None);
+    let gl = refusal("GPU", Backend::Gl, DeviceType::Other, DownlevelFlags::all() - DownlevelFlags::COMPUTE_SHADERS);
+    assert_eq!(gl.as_deref(), Some("the GPU adapter GPU (Gl) has no compute shaders"));
+    let software = refusal("llvmpipe", Backend::Vulkan, DeviceType::Cpu, DownlevelFlags::all());
+    assert_eq!(software.as_deref(), Some("only a software adapter (llvmpipe) is available"));
+}

docs/roto-brush-sam2.md

@@ -155,7 +155,9 @@ of clicks, and measured confidence collapses with many (0.99 with 1–7 points,
 ## Speed (this PC: AMD RX 7900 GRE and a 16-thread CPU, release build)
 SAM 2 runs on the GPU when there is one (wgpu compute kernels, so AMD, NVIDIA and Intel alike; candle itself has no AMD
 backend) and on the CPU otherwise. The model loads on the GPU automatically; `EFFECTCRAFT_SAM2_DEVICE=cpu` (or `gpu`)
-overrides it. A GPU that fails (no adapter, out of memory, a lost device) makes the CPU take over for the rest of the
+overrides it. wgpu's environment variables apply to the adapter it picks (`WGPU_BACKEND`, `WGPU_POWER_PREF`,
+`WGPU_ALLOW_UNDERLYING_NONCOMPLIANT_ADAPTER`…); a software adapter, or one without compute shaders, leaves SAM 2 on
+the CPU. A GPU that fails (no adapter, out of memory, a lost device) makes the CPU take over for the rest of the
 session, and `roto.status` reports `tracker.device` (the GPU's name, or `CPU`).
 
 | Step | CPU | GPU |