Fix cooperative matrix tests on RADV and llvmpipe. (gfx-rs#9892)

Author: jimblandy

examples/features/src/cooperative_matrix/shader_f16_16x16.wgsl

@@ -42,24 +42,27 @@ fn main(@builtin(workgroup_id) workgroup_id: vec3<u32>) {
     let tile_row = workgroup_id.x * TILE_SIZE;
     let tile_col = workgroup_id.y * TILE_SIZE;
 
-    // Load the C tile (accumulator)
+    // Load the C tile (accumulator). `matrix_c` is row-major, so this must
+    // use the row-major (`...T`) load/store variants -- the plain
+    // `coopLoad`/`coopStore` are column-major and would silently read/write
+    // a transposed tile.
     let c_offset = tile_row * stride + tile_col;
-    var c_tile = coopLoad<coop_mat16x16<f16, C>>(&matrix_c[c_offset], stride);
+    var c_tile = coopLoadT<coop_mat16x16<f16, C>>(&matrix_c[c_offset], stride);
 
     // Iterate over K dimension in tiles
     for (var k: u32 = 0u; k < K; k += TILE_SIZE) {
         // Load A tile: rows [tile_row, tile_row+16), cols [k, k+16)
         let a_offset = tile_row * K + k;
-        let a_tile = coopLoad<coop_mat16x16<f16, A>>(&matrix_a[a_offset], K);
+        let a_tile = coopLoadT<coop_mat16x16<f16, A>>(&matrix_a[a_offset], K);
 
         // Load B tile: rows [k, k+16), cols [tile_col, tile_col+16)
         let b_offset = k * stride + tile_col;
-        let b_tile = coopLoad<coop_mat16x16<f16, B>>(&matrix_b[b_offset], stride);
+        let b_tile = coopLoadT<coop_mat16x16<f16, B>>(&matrix_b[b_offset], stride);
 
         // Multiply and accumulate: C += A * B
         c_tile = coopMultiplyAdd(a_tile, b_tile, c_tile);
     }
 
     // Store the result back to C
-    coopStore(c_tile, &matrix_c[c_offset], stride);
+    coopStoreT(c_tile, &matrix_c[c_offset], stride);
 }

examples/features/src/cooperative_matrix/tests.rs

@@ -1,34 +1,79 @@
 use super::*;
 use wgpu_test::{gpu_test, GpuTestConfiguration, TestParameters};
 
+// These tests are split into `COOPERATIVE_MATRIX_F16` and
+// `COOPERATIVE_MATRIX_F32` because the latter can run without
+// `wgpu::Features::SHADER_F16`.
+
 #[gpu_test]
-pub static COOPERATIVE_MATRIX: GpuTestConfiguration = GpuTestConfiguration::new()
+pub static COOPERATIVE_MATRIX_F32: 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();
+        // `shader.wgsl` hardcodes 8x8 f32 tiles (`coop_mat8x8<f32, ...>`),
+        // so only an exact match is usable here.
+        let config = 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
+        });
+        let Some(config) = config else {
+            // Not every adapter that supports EXPERIMENTAL_COOPERATIVE_MATRIX
+            // exposes an 8x8x8 f32/f32 configuration -- e.g. tensor/matrix-core
+            // hardware commonly multiplies in a reduced-precision input type
+            // and only optionally accumulates at f32, so plain f32 inputs are
+            // often unsupported. We can't `.skip()` this per-adapter without
+            // that list growing without bound across every GPU architecture
+            // contributors happen to test on, so we log and move on instead.
+            log::warn!(
+                "No 8x8x8 f32 cooperative matrix configuration found among: \
+                 {coop_props:?}; skipping test"
+            );
+            return;
+        };
+        let ExecuteResults {
+            max_error,
+            tolerance,
+            matrix: _,
+        } = execute(&ctx.device, &ctx.queue, config).await;
+        assert!(max_error < tolerance);
+    });
+
+#[gpu_test]
+pub static COOPERATIVE_MATRIX_F16: GpuTestConfiguration = GpuTestConfiguration::new()
+    .parameters(
+        TestParameters::default()
+            .features(wgpu::Features::EXPERIMENTAL_COOPERATIVE_MATRIX | wgpu::Features::SHADER_F16)
+            .limits(wgpu::Limits::default()),
+    )
+    .run_async(|ctx| async move {
+        let coop_props = ctx.adapter.cooperative_matrix_properties();
+        // `shader_f16_16x16.wgsl` hardcodes 16x16 f16 tiles
+        // (`coop_mat16x16<f16, ...>`), so only an exact match is usable here.
+        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
+        });
+        let Some(config) = config else {
+            // See the comment in COOPERATIVE_MATRIX_F32 above: not every adapter
+            // exposes a 16x16x16 f16/f16 configuration (e.g. some only pair
+            // smaller tile sizes with f16), and we don't want a per-adapter
+            // `.skip()` list to grow without bound, so we log and move on.
+            log::warn!(
+                "No 16x16x16 f16 cooperative matrix configuration found among: \
+                 {coop_props:?}; skipping test"
+            );
+            return;
+        };
         let ExecuteResults {
             max_error,
             tolerance,

examples/features/src/lib.rs

@@ -84,7 +84,8 @@ fn all_tests() -> Vec<wgpu_test::GpuTestInitializer> {
     #[cfg(not(wasm_test))]
     {
         test_list.push(big_compute_buffers::tests::TWO_BUFFERS);
-        test_list.push(cooperative_matrix::tests::COOPERATIVE_MATRIX);
+        test_list.push(cooperative_matrix::tests::COOPERATIVE_MATRIX_F32);
+        test_list.push(cooperative_matrix::tests::COOPERATIVE_MATRIX_F16);
     }
 
     test_list