(Naga) Cooperative Matrix Support

CHANGELOG.md

@@ -44,7 +44,7 @@ Bottom level categories:
 
 #### Deferred command buffer actions: `map_buffer_on_submit` and `on_submitted_work_done`
 
-You may schedule buffer mapping and a submission-complete callback to run automatically after you submit, directly from encoders, command buffers, and passes. 
+You may schedule buffer mapping and a submission-complete callback to run automatically after you submit, directly from encoders, command buffers, and passes.
 
 ```rust
 // Record some GPU work so the submission isn't empty and touches `buffer`.
@@ -150,7 +150,7 @@ By @cwfitzgerald in [#8163](https://github.com/gfx-rs/wgpu/pull/8163).
 
 #### Multi-draw indirect is now unconditionally supported when indirect draws are supported
 
-We have removed `Features::MULTI_DRAW_INDIRECT` as it was unconditionally available on all platforms. 
+We have removed `Features::MULTI_DRAW_INDIRECT` as it was unconditionally available on all platforms.
 `RenderPass::multi_draw_indirect` is now available if the device supports downlevel flag `DownlevelFlags::INDIRECT_EXECUTION`.
 
 If you are using spirv-passthrough with multi-draw indirect and `gl_DrawID`, you can know if `MULTI_DRAW_INDIRECT` is being emulated
@@ -166,6 +166,8 @@ By @cwfitzgerald in [#8162](https://github.com/gfx-rs/wgpu/pull/8162).
 
 - Added support for external textures based on WebGPU's [`GPUExternalTexture`](https://www.w3.org/TR/webgpu/#gpuexternaltexture). These allow shaders to transparently operate on potentially multiplanar source texture data in either RGB or YCbCr formats via WGSL's `texture_external` type. This is gated behind the `Features::EXTERNAL_TEXTURE` feature, which is currently only supported on DX12. By @jamienicol in [#4386](https://github.com/gfx-rs/wgpu/issues/4386).
 
+- Added support for cooperative load/store operations in shaders. Currently only WGSL on the input and SPIR-V,METAL, and WGSL on the output are supported. By @kvark in [#8251](https://github.com/gfx-rs/wgpu/issues/8251).
+
 ### Changes
 
 #### General

Cargo.toml

@@ -254,6 +254,8 @@ ndk-sys = "0.6"
 # These overrides allow our examples to explicitly depend on release crates
 [patch.crates-io]
 wgpu = { path = "./wgpu" }
+rspirv = { git = "https://github.com/gfx-rs/rspirv", rev = "89ce4d0e64c91b0635f617409dc57cb031749a39" }
+spirv = { git = "https://github.com/gfx-rs/rspirv", rev = "89ce4d0e64c91b0635f617409dc57cb031749a39" }
 
 [profile.release]
 lto = "thin"

naga/src/back/dot/mod.rs

@@ -403,6 +403,16 @@ impl StatementGraph {
                         },
                     }
                 }
+                S::CooperativeStore { target, data } => {
+                    self.dependencies.push((id, target, "target"));
+                    self.dependencies.push((id, data.pointer, "pointer"));
+                    self.dependencies.push((id, data.stride, "stride"));
+                    if data.row_major {
+                        "CoopStoreT"
+                    } else {
+                        "CoopStore"
+                    }
+                }
             };
             // Set the last node to the merge node
             last_node = merge_id;
@@ -742,6 +752,18 @@ fn write_function_expressions(
                 let ty = if committed { "Committed" } else { "Candidate" };
                 (format!("get{ty}HitVertexPositions").into(), 4)
             }
+            E::CooperativeLoad { ref data, .. } => {
+                edges.insert("pointer", data.pointer);
+                edges.insert("stride", data.stride);
+                let suffix = if data.row_major { "T " } else { "" };
+                (format!("coopLoad{suffix}").into(), 4)
+            }
+            E::CooperativeMultiplyAdd { a, b, c } => {
+                edges.insert("a", a);
+                edges.insert("b", b);
+                edges.insert("c", c);
+                ("cooperativeMultiplyAdd".into(), 4)
+            }
         };
 
         // give uniform expressions an outline

naga/src/back/glsl/mod.rs

@@ -1107,7 +1107,8 @@ impl<'a, W: Write> Writer<'a, W> {
             TypeInner::Array { base, size, .. } => self.write_array_size(base, size)?,
             // Write all variants instead of `_` so that if new variants are added a
             // no exhaustiveness error is thrown
-            TypeInner::Pointer { .. }
+            TypeInner::CooperativeMatrix { .. }
+            | TypeInner::Pointer { .. }
             | TypeInner::Struct { .. }
             | TypeInner::Image { .. }
             | TypeInner::Sampler { .. }
@@ -2804,6 +2805,7 @@ impl<'a, W: Write> Writer<'a, W> {
                 }
                 writeln!(self.out, ");")?;
             }
+            Statement::CooperativeStore { .. } => unimplemented!(),
         }
 
         Ok(())
@@ -4340,7 +4342,9 @@ impl<'a, W: Write> Writer<'a, W> {
             }
             // not supported yet
             Expression::RayQueryGetIntersection { .. }
-            | Expression::RayQueryVertexPositions { .. } => unreachable!(),
+            | Expression::RayQueryVertexPositions { .. }
+            | Expression::CooperativeLoad { .. }
+            | Expression::CooperativeMultiplyAdd { .. } => unreachable!(),
         }
 
         Ok(())

naga/src/back/hlsl/writer.rs

@@ -2747,6 +2747,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 }
                 writeln!(self.out, ");")?;
             }
+            Statement::CooperativeStore { .. } => unimplemented!(),
         }
 
         Ok(())
@@ -4275,7 +4276,11 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 }
             }
             // Not supported yet
-            Expression::RayQueryVertexPositions { .. } => unreachable!(),
+            Expression::RayQueryVertexPositions { .. }
+            | Expression::CooperativeLoad { .. }
+            | Expression::CooperativeMultiplyAdd { .. } => {
+                unreachable!()
+            }
             // Nothing to do here, since call expression already cached
             Expression::CallResult(_)
             | Expression::AtomicResult { .. }

naga/src/back/mod.rs

@@ -312,12 +312,10 @@ pub const fn binary_operation_str(op: crate::BinaryOperator) -> &'static str {
 }
 
 impl crate::TypeInner {
-    /// Returns true if this is a handle to a type rather than the type directly.
+    /// Returns true if a variable of this type is a handle.
     pub const fn is_handle(&self) -> bool {
         match *self {
-            crate::TypeInner::Image { .. }
-            | crate::TypeInner::Sampler { .. }
-            | crate::TypeInner::AccelerationStructure { .. } => true,
+            Self::Image { .. } | Self::Sampler { .. } | Self::AccelerationStructure { .. } => true,
             _ => false,
         }
     }

naga/src/back/msl/mod.rs

@@ -228,8 +228,10 @@ pub enum Error {
     UnsupportedArrayOf(String),
     #[error("array of type '{0:?}' is not supported")]
     UnsupportedArrayOfType(Handle<crate::Type>),
-    #[error("ray tracing is not supported prior to MSL 2.3")]
+    #[error("ray tracing is not supported prior to MSL 2.4")]
     UnsupportedRayTracing,
+    #[error("cooperative matrix is not supported prior to MSL 2.3")]
+    UnsupportedCooperativeMatrix,
     #[error("overrides should not be present at this stage")]
     Override,
     #[error("bitcasting to {0:?} is not supported")]