coop: add GPU test, API doc

Author: kvark

docs/api-specs/cooperative_matrix.md

@@ -78,7 +78,7 @@ fn render(/*whatever args you need to render*/) {
   /* do other preparations on the TlasInstance.*/
   encoder.build_acceleration_structures(iter::empty(), iter::once(&tlas_package));
   /* more render code */
-  queue.submit(Some(encoder.finish()));
+  queue.submit([encoder.finish()]);
 }
 ```
 

docs/api-specs/ray_tracing.md

@@ -2,70 +2,26 @@
 
 This example demonstrates how to use cooperative matrix operations (also known as tensor cores on NVIDIA GPUs) to perform efficient matrix multiplication on the GPU.
 
-## Overview
+For the full description of the cooperative matrix feature (supported configurations, WGSL types and operations, validation rules, and backend support), see the central API spec:
 
-Cooperative matrices allow a workgroup to collectively load, store, and perform matrix multiply-accumulate operations on small tiles of data. This enables hardware-accelerated matrix math that can be significantly faster than traditional element-wise approaches.
+- `docs/api-specs/cooperative_matrix.md`
 
-The example computes `C = A * B + C` where:
+## Example specifics
+
+This example computes `C = A * B + C` where:
 - A is a 64×64 matrix
 - B is a 64×64 matrix
 - C is a 64×64 matrix (accumulator/result)
 
-## Querying Supported Configurations
-
-Before using cooperative matrices, you should query what configurations your hardware supports:
-
-```rust
-let coop_props = adapter.cooperative_matrix_properties();
-for prop in &coop_props {
-    println!(
-        "{:?}x{:?}x{:?} - AB: {:?}, CR: {:?}",
-        prop.m_size, prop.n_size, prop.k_size,
-        prop.ab_type, prop.cr_type
-    );
-}
-```
-
-Each `CooperativeMatrixProperties` describes a supported configuration with:
-- `m_size`, `n_size`, `k_size`: Matrix dimensions as `naga::CooperativeSize` (M×K × K×N → M×N)
-- `ab_type`: Element type for input matrices A and B (as `naga::Scalar`)
-- `cr_type`: Element type for accumulator matrix C and the result
-- `saturating_accumulation`: Whether overflow clamping is supported
-
-## Key Concepts
-
-### Cooperative Matrix Types
-
-In WGSL, cooperative matrices are declared with a specific size, element type, and role:
-
-```wgsl
-coop_mat8x8<f32, A>  // Matrix A (left operand)
-coop_mat8x8<f32, B>  // Matrix B (right operand)
-coop_mat8x8<f32, C>  // Matrix C (accumulator)
-```
-
-The role (A, B, or C) determines how the matrix is used in multiply-accumulate operations.
-
-### Operations
-
-- `coopLoad<T>(pointer, stride)` - Cooperatively load a tile from memory
-- `coopStore(matrix, pointer, stride)` - Cooperatively store a tile to memory
-- `coopMultiplyAdd(a, b, c)` - Compute `a * b + c`
-
-### Workgroup Cooperation
-
-All threads in a workgroup must participate in cooperative matrix operations together. The workgroup size should match the cooperative matrix dimensions (8×8 in this example).
+The example:
+- Tiles the 64×64 matrices into cooperative matrix tiles (e.g. 8×8) and performs a tiled matmul
+- Uses a compute shader and compares GPU results against a CPU reference implementation
 
 ## Requirements
 
-- GPU with cooperative matrix support:
-  - Metal: Apple7+ (A14 chip) or Mac2+ (M1 chip) with MSL 2.3+
-    - Supports 8x8 f32, 8x8 f16, and mixed precision (f16 inputs, f32 accumulator)
-  - Vulkan: Requires VK_KHR_cooperative_matrix extension
-    - Most NVIDIA/AMD GPUs support f16 at 16x16 sizes
-    - 8x8 f32 support varies by hardware
-- `Features::EXPERIMENTAL_COOPERATIVE_MATRIX` must be enabled
-- Use `adapter.cooperative_matrix_properties()` to check available configurations
+- A GPU and backend that expose `Features::EXPERIMENTAL_COOPERATIVE_MATRIX`
+- A configuration returned from `adapter.cooperative_matrix_properties()` that matches the tile size and element types used by this example
+- See `docs/api-specs/cooperative_matrix.md` for details on hardware / backend support
 
 ## Running
 
@@ -76,8 +32,5 @@ cargo run --bin wgpu-examples -- cooperative_matrix
 ## Notes
 
 - This is an experimental feature and may not work on all hardware
-- Always query `adapter.cooperative_matrix_properties()` to check what's supported
-- The 8x8 f32 matrix format is well supported on Metal (simdgroup matrix operations)
-- Vulkan support depends on hardware - most GPUs (NVIDIA, AMD) support f16 inputs at 16x16 sizes
 - The shader uses the standard `create_shader_module` with full validation
 - Results are verified against a CPU reference implementation

examples/features/src/cooperative_matrix/README.md

@@ -163,6 +163,49 @@ async fn run() {
             .expect("Failed to create device")
     };
 
+    let results = execute(&device, &queue, config).await;
+
+    log::info!(
+        "Matrix multiplication {M}x{K}x{N} completed using {} precision!",
+        if use_f16 { "f16" } else { "f32" }
+    );
+    log::info!("Max error vs CPU reference: {:.6}", results.max_error);
+
+    if results.max_error < results.tolerance {
+        log::info!(
+            "✓ Results match CPU reference within tolerance ({})",
+            results.tolerance
+        );
+    } else {
+        log::warn!(
+            "✗ Results differ from CPU reference (tolerance: {})",
+            results.tolerance
+        );
+    }
+
+    // Print a small sample of the result
+    log::info!("Sample of result matrix C (top-left 4x4):");
+    for i in 0..4 {
+        let row: Vec<String> = (0..4)
+            .map(|j| format!("{:6.2}", results.matrix[i * N as usize + j]))
+            .collect();
+        log::info!("  [{}]", row.join(", "));
+    }
+}
+
+struct ExecuteResults {
+    max_error: f32,
+    tolerance: f32,
+    matrix: Vec<f32>,
+}
+
+async fn execute(
+    device: &wgpu::Device,
+    queue: &wgpu::Queue,
+    config: &wgpu::CooperativeMatrixProperties,
+) -> ExecuteResults {
+    let use_f16 = config.ab_type == wgpu::CooperativeScalarType::F16;
+
     // Select the appropriate shader based on configuration
     let shader_source = if use_f16 {
         include_str!("shader_f16_16x16.wgsl")
@@ -347,7 +390,7 @@ async fn run() {
         compute_pass.set_pipeline(&pipeline);
         compute_pass.set_bind_group(0, &bind_group, &[]);
         // Dispatch one workgroup per tile of the output
-        compute_pass.dispatch_workgroups(M / tile_size, N / tile_size, 1);
+        compute_pass.dispatch_workgroups(M / config.m_size, N / config.m_size, 1);
     }
 
     // Copy result to staging buffer
@@ -398,29 +441,11 @@ async fn run() {
         max_error = max_error.max(error);
     }
 
-    log::info!(
-        "Matrix multiplication {M}x{K}x{N} completed using {} precision!",
-        if use_f16 { "f16" } else { "f32" }
-    );
-    log::info!("Max error vs CPU reference: {max_error:.6}");
-
-    if max_error < tolerance {
-        log::info!("✓ Results match CPU reference within tolerance ({tolerance})");
-    } else {
-        log::warn!("✗ Results differ from CPU reference (tolerance: {tolerance})");
-    }
-
-    // Print a small sample of the result
-    log::info!("Sample of result matrix C (top-left 4x4):");
-    for i in 0..4 {
-        let row: Vec<String> = (0..4)
-            .map(|j| format!("{:6.2}", result[i * N as usize + j]))
-            .collect();
-        log::info!("  [{}]", row.join(", "));
+    ExecuteResults {
+        max_error,
+        tolerance,
+        matrix: result,
     }
-
-    drop(data);
-    staging_buffer.unmap();
 }
 
 pub fn main() {
@@ -440,3 +465,6 @@ pub fn main() {
         wasm_bindgen_futures::spawn_local(run());
     }
 }
+
+#[cfg(test)]
+pub mod tests;

examples/features/src/cooperative_matrix/mod.rs

@@ -7,7 +7,7 @@
 // Each workgroup cooperatively loads tiles of A and B, multiplies them,
 // and accumulates the result into C.
 
-enable experimental_cooperative_matrix;
+enable wgpu_cooperative_matrix;
 
 // Matrix dimensions (8x8 tiles)
 const TILE_SIZE: u32 = 8u;

examples/features/src/cooperative_matrix/shader.wgsl

@@ -11,7 +11,7 @@
 // and accumulates the result into C.
 
 enable f16;
-enable experimental_cooperative_matrix;
+enable wgpu_cooperative_matrix;
 
 // Matrix dimensions (16x16 tiles)
 const TILE_SIZE: u32 = 16u;

examples/features/src/cooperative_matrix/shader_f16_16x16.wgsl

@@ -0,0 +1,38 @@
+use super::*;
+use wgpu_test::{gpu_test, GpuTestConfiguration, TestParameters};
+
+#[gpu_test]
+pub static COOPERATIVE_MATRIX: GpuTestConfiguration = GpuTestConfiguration::new()
+    .parameters(
+        TestParameters::default()
+            .features(wgpu::Features::EXPERIMENTAL_COOPERATIVE_MATRIX)
+            .limits(wgpu::Limits::default()),
+    )
+    .run_async(|ctx| async move {
+        let coop_props = ctx.adapter.cooperative_matrix_properties();
+        let config = coop_props
+            .iter()
+            .find(|prop| {
+                prop.m_size == 16
+                    && prop.n_size == 16
+                    && prop.k_size == 16
+                    && prop.ab_type == wgpu::CooperativeScalarType::F16
+                    && prop.cr_type == wgpu::CooperativeScalarType::F16
+            })
+            .or_else(|| {
+                coop_props.iter().find(|prop| {
+                    prop.m_size == 8
+                        && prop.n_size == 8
+                        && prop.k_size == 8
+                        && prop.ab_type == wgpu::CooperativeScalarType::F32
+                        && prop.cr_type == wgpu::CooperativeScalarType::F32
+                })
+            })
+            .unwrap();
+        let ExecuteResults {
+            max_error,
+            tolerance,
+            matrix: _,
+        } = execute(&ctx.device, &ctx.queue, config).await;
+        assert!(max_error < tolerance);
+    });

examples/features/src/cooperative_matrix/tests.rs

@@ -81,6 +81,7 @@ fn all_tests() -> Vec<wgpu_test::GpuTestInitializer> {
     #[cfg(not(target_arch = "wasm32"))]
     {
         test_list.push(big_compute_buffers::tests::TWO_BUFFERS);
+        test_list.push(cooperative_matrix::tests::COOPERATIVE_MATRIX);
     }
 
     test_list

examples/features/src/lib.rs

@@ -531,7 +531,6 @@ impl<'source, 'temp, 'out> ExpressionContext<'source, 'temp, 'out> {
         span: Span,
     ) -> Result<'source, Handle<ir::Expression>> {
         let mut eval = self.as_const_evaluator();
-        log::debug!("appending {expr:?}");
         eval.try_eval_and_append(expr, span)
             .map_err(|e| Box::new(Error::ConstantEvaluatorError(e.into(), span)))
     }

naga/src/front/wgsl/lower/mod.rs

@@ -17,7 +17,7 @@ pub struct EnableExtensions {
     /// Whether `enable f16;` was written earlier in the shader module.
     f16: bool,
     clip_distances: bool,
-    experimental_cooperative_matrix: bool,
+    wgpu_cooperative_matrix: bool,
 }
 
 impl EnableExtensions {
@@ -29,7 +29,7 @@ impl EnableExtensions {
             f16: false,
             dual_source_blending: false,
             clip_distances: false,
-            experimental_cooperative_matrix: false,
+            wgpu_cooperative_matrix: false,
         }
     }
 
@@ -44,9 +44,7 @@ impl EnableExtensions {
             ImplementedEnableExtension::DualSourceBlending => &mut self.dual_source_blending,
             ImplementedEnableExtension::F16 => &mut self.f16,
             ImplementedEnableExtension::ClipDistances => &mut self.clip_distances,
-            ImplementedEnableExtension::ExperimentalCooperativeMatrix => {
-                &mut self.experimental_cooperative_matrix
-            }
+            ImplementedEnableExtension::WgpuCooperativeMatrix => &mut self.wgpu_cooperative_matrix,
         };
         *field = true;
     }
@@ -62,9 +60,7 @@ impl EnableExtensions {
             ImplementedEnableExtension::DualSourceBlending => self.dual_source_blending,
             ImplementedEnableExtension::F16 => self.f16,
             ImplementedEnableExtension::ClipDistances => self.clip_distances,
-            ImplementedEnableExtension::ExperimentalCooperativeMatrix => {
-                self.experimental_cooperative_matrix
-            }
+            ImplementedEnableExtension::WgpuCooperativeMatrix => self.wgpu_cooperative_matrix,
         }
     }
 }
@@ -97,7 +93,7 @@ impl EnableExtension {
     const MESH_SHADER: &'static str = "wgpu_mesh_shader";
     const RAY_QUERY: &'static str = "wgpu_ray_query";
     const RAY_QUERY_VERTEX_RETURN: &'static str = "wgpu_ray_query_vertex_return";
-    const EXPERIMENTAL_COOPERATIVE_MATRIX: &'static str = "experimental_cooperative_matrix";
+    const COOPERATIVE_MATRIX: &'static str = "wgpu_cooperative_matrix";
     const SUBGROUPS: &'static str = "subgroups";
     const PRIMITIVE_INDEX: &'static str = "primitive_index";
 
@@ -114,8 +110,8 @@ impl EnableExtension {
             Self::RAY_QUERY_VERTEX_RETURN => {
                 Self::Implemented(ImplementedEnableExtension::WgpuRayQueryVertexReturn)
             }
-            Self::EXPERIMENTAL_COOPERATIVE_MATRIX => {
-                Self::Implemented(ImplementedEnableExtension::ExperimentalCooperativeMatrix)
+            Self::COOPERATIVE_MATRIX => {
+                Self::Implemented(ImplementedEnableExtension::WgpuCooperativeMatrix)
             }
             Self::SUBGROUPS => Self::Unimplemented(UnimplementedEnableExtension::Subgroups),
             Self::PRIMITIVE_INDEX => {
@@ -134,9 +130,7 @@ impl EnableExtension {
                 ImplementedEnableExtension::WgpuRayQueryVertexReturn => {
                     Self::RAY_QUERY_VERTEX_RETURN
                 }
-                ImplementedEnableExtension::ExperimentalCooperativeMatrix => {
-                    Self::EXPERIMENTAL_COOPERATIVE_MATRIX
-                }
+                ImplementedEnableExtension::WgpuCooperativeMatrix => Self::COOPERATIVE_MATRIX,
                 ImplementedEnableExtension::DualSourceBlending => Self::DUAL_SOURCE_BLENDING,
                 ImplementedEnableExtension::F16 => Self::F16,
                 ImplementedEnableExtension::ClipDistances => Self::CLIP_DISTANCES,
@@ -176,8 +170,8 @@ pub enum ImplementedEnableExtension {
     WgpuRayQuery,
     /// Enables the `wgpu_ray_query_vertex_return` extension, native only.
     WgpuRayQueryVertexReturn,
-    /// Enables the `experimental_cooperative_matrix` extension, native only.
-    ExperimentalCooperativeMatrix,
+    /// Enables the `wgpu_cooperative_matrix` extension, native only.
+    WgpuCooperativeMatrix,
 }
 
 /// A variant of [`EnableExtension::Unimplemented`].