Skip to content

Commit 15960ba

Browse files
committed
Fix cooperative matrix tests on RADV and llvmpipe.
In `examples/features/src/cooperative_matrix`: - Fix `shader_f16_16x16.wgsl` to properly load and store matrices from the row-major buffers. - Split the `COOPERATIVE_MATRIX` test into two separate tests, `COOPERATIVE_MATRIX_F32` and `COOPERATIVE_MATRIX_F16`, and require `wgpu::Features::SHADER_F16` for the latter. - If the hardware doesn't support any of the configurations we have a shader for, log a warning and return (that is, pass). Comments in the code explain why we do this instead of putting a `skip` on the test.
1 parent 4e58833 commit 15960ba

3 files changed

Lines changed: 75 additions & 26 deletions

File tree

examples/features/src/cooperative_matrix/shader_f16_16x16.wgsl

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -42,24 +42,27 @@ fn main(@builtin(workgroup_id) workgroup_id: vec3<u32>) {
4242
let tile_row = workgroup_id.x * TILE_SIZE;
4343
let tile_col = workgroup_id.y * TILE_SIZE;
4444

45-
// Load the C tile (accumulator)
45+
// Load the C tile (accumulator). `matrix_c` is row-major, so this must
46+
// use the row-major (`...T`) load/store variants -- the plain
47+
// `coopLoad`/`coopStore` are column-major and would silently read/write
48+
// a transposed tile.
4649
let c_offset = tile_row * stride + tile_col;
47-
var c_tile = coopLoad<coop_mat16x16<f16, C>>(&matrix_c[c_offset], stride);
50+
var c_tile = coopLoadT<coop_mat16x16<f16, C>>(&matrix_c[c_offset], stride);
4851

4952
// Iterate over K dimension in tiles
5053
for (var k: u32 = 0u; k < K; k += TILE_SIZE) {
5154
// Load A tile: rows [tile_row, tile_row+16), cols [k, k+16)
5255
let a_offset = tile_row * K + k;
53-
let a_tile = coopLoad<coop_mat16x16<f16, A>>(&matrix_a[a_offset], K);
56+
let a_tile = coopLoadT<coop_mat16x16<f16, A>>(&matrix_a[a_offset], K);
5457

5558
// Load B tile: rows [k, k+16), cols [tile_col, tile_col+16)
5659
let b_offset = k * stride + tile_col;
57-
let b_tile = coopLoad<coop_mat16x16<f16, B>>(&matrix_b[b_offset], stride);
60+
let b_tile = coopLoadT<coop_mat16x16<f16, B>>(&matrix_b[b_offset], stride);
5861

5962
// Multiply and accumulate: C += A * B
6063
c_tile = coopMultiplyAdd(a_tile, b_tile, c_tile);
6164
}
6265

6366
// Store the result back to C
64-
coopStore(c_tile, &matrix_c[c_offset], stride);
67+
coopStoreT(c_tile, &matrix_c[c_offset], stride);
6568
}

examples/features/src/cooperative_matrix/tests.rs

Lines changed: 65 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,79 @@
11
use super::*;
22
use wgpu_test::{gpu_test, GpuTestConfiguration, TestParameters};
33

4+
// These tests are split into `COOPERATIVE_MATRIX_F16` and
5+
// `COOPERATIVE_MATRIX_F32` because the latter can run without
6+
// `wgpu::Features::SHADER_F16`.
7+
48
#[gpu_test]
5-
pub static COOPERATIVE_MATRIX: GpuTestConfiguration = GpuTestConfiguration::new()
9+
pub static COOPERATIVE_MATRIX_F32: GpuTestConfiguration = GpuTestConfiguration::new()
610
.parameters(
711
TestParameters::default()
812
.features(wgpu::Features::EXPERIMENTAL_COOPERATIVE_MATRIX)
913
.limits(wgpu::Limits::default()),
1014
)
1115
.run_async(|ctx| async move {
1216
let coop_props = ctx.adapter.cooperative_matrix_properties();
13-
let config = coop_props
14-
.iter()
15-
.find(|prop| {
16-
prop.m_size == 16
17-
&& prop.n_size == 16
18-
&& prop.k_size == 16
19-
&& prop.ab_type == wgpu::CooperativeScalarType::F16
20-
&& prop.cr_type == wgpu::CooperativeScalarType::F16
21-
})
22-
.or_else(|| {
23-
coop_props.iter().find(|prop| {
24-
prop.m_size == 8
25-
&& prop.n_size == 8
26-
&& prop.k_size == 8
27-
&& prop.ab_type == wgpu::CooperativeScalarType::F32
28-
&& prop.cr_type == wgpu::CooperativeScalarType::F32
29-
})
30-
})
31-
.unwrap();
17+
// `shader.wgsl` hardcodes 8x8 f32 tiles (`coop_mat8x8<f32, ...>`),
18+
// so only an exact match is usable here.
19+
let config = coop_props.iter().find(|prop| {
20+
prop.m_size == 8
21+
&& prop.n_size == 8
22+
&& prop.k_size == 8
23+
&& prop.ab_type == wgpu::CooperativeScalarType::F32
24+
&& prop.cr_type == wgpu::CooperativeScalarType::F32
25+
});
26+
let Some(config) = config else {
27+
// Not every adapter that supports EXPERIMENTAL_COOPERATIVE_MATRIX
28+
// exposes an 8x8x8 f32/f32 configuration -- e.g. tensor/matrix-core
29+
// hardware commonly multiplies in a reduced-precision input type
30+
// and only optionally accumulates at f32, so plain f32 inputs are
31+
// often unsupported. We can't `.skip()` this per-adapter without
32+
// that list growing without bound across every GPU architecture
33+
// contributors happen to test on, so we log and move on instead.
34+
log::warn!(
35+
"No 8x8x8 f32 cooperative matrix configuration found among: \
36+
{coop_props:?}; skipping test"
37+
);
38+
return;
39+
};
40+
let ExecuteResults {
41+
max_error,
42+
tolerance,
43+
matrix: _,
44+
} = execute(&ctx.device, &ctx.queue, config).await;
45+
assert!(max_error < tolerance);
46+
});
47+
48+
#[gpu_test]
49+
pub static COOPERATIVE_MATRIX_F16: GpuTestConfiguration = GpuTestConfiguration::new()
50+
.parameters(
51+
TestParameters::default()
52+
.features(wgpu::Features::EXPERIMENTAL_COOPERATIVE_MATRIX | wgpu::Features::SHADER_F16)
53+
.limits(wgpu::Limits::default()),
54+
)
55+
.run_async(|ctx| async move {
56+
let coop_props = ctx.adapter.cooperative_matrix_properties();
57+
// `shader_f16_16x16.wgsl` hardcodes 16x16 f16 tiles
58+
// (`coop_mat16x16<f16, ...>`), so only an exact match is usable here.
59+
let config = coop_props.iter().find(|prop| {
60+
prop.m_size == 16
61+
&& prop.n_size == 16
62+
&& prop.k_size == 16
63+
&& prop.ab_type == wgpu::CooperativeScalarType::F16
64+
&& prop.cr_type == wgpu::CooperativeScalarType::F16
65+
});
66+
let Some(config) = config else {
67+
// See the comment in COOPERATIVE_MATRIX_F32 above: not every adapter
68+
// exposes a 16x16x16 f16/f16 configuration (e.g. some only pair
69+
// smaller tile sizes with f16), and we don't want a per-adapter
70+
// `.skip()` list to grow without bound, so we log and move on.
71+
log::warn!(
72+
"No 16x16x16 f16 cooperative matrix configuration found among: \
73+
{coop_props:?}; skipping test"
74+
);
75+
return;
76+
};
3277
let ExecuteResults {
3378
max_error,
3479
tolerance,

examples/features/src/lib.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,8 @@ fn all_tests() -> Vec<wgpu_test::GpuTestInitializer> {
8484
#[cfg(not(wasm_test))]
8585
{
8686
test_list.push(big_compute_buffers::tests::TWO_BUFFERS);
87-
test_list.push(cooperative_matrix::tests::COOPERATIVE_MATRIX);
87+
test_list.push(cooperative_matrix::tests::COOPERATIVE_MATRIX_F32);
88+
test_list.push(cooperative_matrix::tests::COOPERATIVE_MATRIX_F16);
8889
}
8990

9091
test_list

0 commit comments

Comments
 (0)