coop: rewire WGSL support using references

kvark · kvark · commit 056a5c1e307a · 2025-09-27T17:47:26.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -166,7 +166,7 @@ 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 with METAL on the output are supported. By @kvark in [#8251](https://github.com/gfx-rs/wgpu/issues/8251).
+- 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
 
diff --git a/naga/src/back/mod.rs b/naga/src/back/mod.rs
@@ -311,6 +311,16 @@ pub const fn binary_operation_str(op: crate::BinaryOperator) -> &'static str {
     }
 }
 
+impl crate::TypeInner {
+    /// Returns true if a variable of this type is a handle.
+    pub const fn is_handle(&self) -> bool {
+        match *self {
+            Self::Image { .. } | Self::Sampler { .. } | Self::AccelerationStructure { .. } => true,
+            _ => false,
+        }
+    }
+}
+
 impl crate::Statement {
     /// Returns true if the statement directly terminates the current block.
     ///
diff --git a/naga/src/back/msl/writer.rs b/naga/src/back/msl/writer.rs
@@ -6316,16 +6316,22 @@ template <typename A>
         b: Handle<crate::Expression>,
     ) -> BackendResult {
         let (a_c, a_r, scalar) = match *func_ctx.resolve_type(a, &module.types) {
-            crate::TypeInner::CooperativeMatrix {
-                columns,
-                rows,
-                scalar,
-                ..
-            } => (columns, rows, scalar),
+            crate::TypeInner::Pointer { base, space: _ } => match module.types[base].inner {
+                crate::TypeInner::CooperativeMatrix {
+                    columns,
+                    rows,
+                    scalar,
+                    ..
+                } => (columns, rows, scalar),
+                _ => unreachable!(),
+            },
             _ => unreachable!(),
         };
         let (b_c, b_r) = match *func_ctx.resolve_type(b, &module.types) {
-            crate::TypeInner::CooperativeMatrix { columns, rows, .. } => (columns, rows),
+            crate::TypeInner::Pointer { base, space: _ } => match module.types[base].inner {
+                crate::TypeInner::CooperativeMatrix { columns, rows, .. } => (columns, rows),
+                _ => unreachable!(),
+            },
             _ => unreachable!(),
         };
         let wrapped = WrappedFunction::CooperativeMultiplyAdd {
diff --git a/naga/src/back/spv/block.rs b/naga/src/back/spv/block.rs
@@ -3717,7 +3717,13 @@ impl BlockContext<'_> {
                             self.cached[stride],
                         ));
                     } else {
-                        let result_type_id = self.get_expression_type_id(&self.fun_info[target].ty);
+                        let result_type_id =
+                            match *self.fun_info[target].ty.inner_with(&self.ir_module.types) {
+                                crate::TypeInner::Pointer { base, space: _ } => {
+                                    self.get_handle_type_id(base)
+                                }
+                                _ => unreachable!(),
+                            };
                         let id = self.gen_id();
                         block.body.push(Instruction::coop_load(
                             result_type_id,
diff --git a/naga/src/back/spv/writer.rs b/naga/src/back/spv/writer.rs
@@ -971,14 +971,13 @@ impl Writer {
                 }
             }
 
-            // Handle globals are pre-emitted and should be loaded automatically.
-            //
-            // Any that are binding arrays we skip as we cannot load the array, we must load the result after indexing.
             match ir_module.types[var.ty].inner {
+                // Any that are binding arrays we skip as we cannot load the array, we must load the result after indexing.
                 crate::TypeInner::BindingArray { .. } => {
                     gv.access_id = gv.var_id;
                 }
                 _ => {
+                    // Handle globals are pre-emitted and should be loaded automatically.
                     if var.space == crate::AddressSpace::Handle {
                         let var_type_id = self.get_handle_type_id(var.ty);
                         let id = self.id_gen.next();
@@ -1064,6 +1063,7 @@ impl Writer {
                     }
                 }),
             );
+
             context
                 .function
                 .variables
diff --git a/naga/src/back/wgsl/writer.rs b/naga/src/back/wgsl/writer.rs
@@ -993,13 +993,13 @@ impl<W: Write> Writer<W> {
             } => {
                 let op_str = if store { "Store" } else { "Load" };
                 let suffix = if row_major { "T" } else { "" };
-                write!(self.out, "coop{op_str}{suffix}(")?;
-                self.write_expr(module, target, func_ctx)?;
+                write!(self.out, "{level}coop{op_str}{suffix}(")?;
+                self.write_expr_with_indirection(module, target, func_ctx, Indirection::Reference)?;
                 write!(self.out, ", ")?;
                 self.write_expr(module, pointer, func_ctx)?;
                 write!(self.out, ", ")?;
                 self.write_expr(module, stride, func_ctx)?;
-                write!(self.out, ")")?
+                writeln!(self.out, ");")?
             }
         }
 
@@ -1714,11 +1714,11 @@ impl<W: Write> Writer<W> {
             | Expression::WorkGroupUniformLoadResult { .. } => {}
             Expression::CooperativeMultiplyAdd { a, b, c } => {
                 write!(self.out, "coopMultiplyAdd(")?;
-                self.write_expr(module, a, func_ctx)?;
+                self.write_expr_with_indirection(module, a, func_ctx, Indirection::Reference)?;
                 write!(self.out, ", ")?;
-                self.write_expr(module, b, func_ctx)?;
+                self.write_expr_with_indirection(module, b, func_ctx, Indirection::Reference)?;
                 write!(self.out, ", ")?;
-                self.write_expr(module, c, func_ctx)?;
+                self.write_expr_with_indirection(module, c, func_ctx, Indirection::Reference)?;
                 write!(self.out, ")")?;
             }
         }
diff --git a/naga/src/front/wgsl/error.rs b/naga/src/front/wgsl/error.rs
@@ -412,6 +412,7 @@ pub(crate) enum Error<'a> {
     TypeTooLarge {
         span: Span,
     },
+    InvalidCooperativeMatrix,
     UnderspecifiedCooperativeMatrix,
     UnsupportedCooperativeScalar(Span),
 }
@@ -1388,6 +1389,11 @@ impl<'a> Error<'a> {
                     crate::valid::MAX_TYPE_SIZE
                 )],
             },
+            Error::InvalidCooperativeMatrix => ParseError {
+                message: "given type is not a cooperative matrix".into(),
+                labels: vec![],
+                notes: vec![format!("must be coop_mat")],
+            },
             Error::UnderspecifiedCooperativeMatrix => ParseError {
                 message: "cooperative matrix constructor is underspecified".into(),
                 labels: vec![],
diff --git a/naga/src/front/wgsl/lower/mod.rs b/naga/src/front/wgsl/lower/mod.rs
@@ -846,6 +846,15 @@ impl<'source, 'temp, 'out> ExpressionContext<'source, 'temp, 'out> {
     fn ensure_type_exists(&mut self, inner: ir::TypeInner) -> Handle<ir::Type> {
         self.as_global().ensure_type_exists(None, inner)
     }
+
+    fn _get_runtime_expression(&self, expr: Handle<ir::Expression>) -> &ir::Expression {
+        match self.expr_type {
+            ExpressionContextType::Runtime(ref ctx) => &ctx.function.expressions[expr],
+            ExpressionContextType::Constant(_) | ExpressionContextType::Override => {
+                unreachable!()
+            }
+        }
+    }
 }
 
 struct ArgumentContext<'ctx, 'source> {
@@ -955,6 +964,13 @@ impl<T> Typed<T> {
             Self::Plain(expr) => Typed::Plain(f(expr)?),
         })
     }
+
+    fn ref_or<E>(self, error: E) -> core::result::Result<T, E> {
+        match self {
+            Self::Reference(v) => Ok(v),
+            Self::Plain(_) => Err(error),
+        }
+    }
 }
 
 /// A single vector component or swizzle.
@@ -1677,12 +1693,8 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                         .as_expression(block, &mut emitter)
                         .interrupt_emitter(ir::Expression::LocalVariable(var), Span::UNDEFINED)?;
                     block.extend(emitter.finish(&ctx.function.expressions));
-                    let typed = if ctx.module.types[ty].inner.is_handle() {
-                        Typed::Plain(handle)
-                    } else {
-                        Typed::Reference(handle)
-                    };
-                    ctx.local_table.insert(v.handle, Declared::Runtime(typed));
+                    ctx.local_table
+                        .insert(v.handle, Declared::Runtime(Typed::Reference(handle)));
 
                     match initializer {
                         Some(initializer) => ir::Statement::Store {
@@ -1977,12 +1989,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 let value_span = ctx.ast_expressions.get_span(value);
                 let target = self
                     .expression_for_reference(value, &mut ctx.as_expression(block, &mut emitter))?;
-                let target_handle = match target {
-                    Typed::Reference(handle) => handle,
-                    Typed::Plain(_) => {
-                        return Err(Box::new(Error::BadIncrDecrReferenceType(value_span)))
-                    }
-                };
+                let target_handle = target.ref_or(Error::BadIncrDecrReferenceType(value_span))?;
 
                 let mut ectx = ctx.as_expression(block, &mut emitter);
                 let scalar = match *resolve_inner!(ectx, target_handle) {
@@ -2139,10 +2146,8 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     LoweredGlobalDecl::Var(handle) => {
                         let expr = ir::Expression::GlobalVariable(handle);
                         let v = &ctx.module.global_variables[handle];
-                        let force_value = ctx.module.types[v.ty].inner.is_handle();
                         match v.space {
                             ir::AddressSpace::Handle => Typed::Plain(expr),
-                            _ if force_value => Typed::Plain(expr),
                             _ => Typed::Reference(expr),
                         }
                     }
@@ -3140,7 +3145,9 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                             let row_major = function.name.ends_with("T");
 
                             let mut args = ctx.prepare_args(arguments, 2, span);
-                            let target = self.expression(args.next()?, ctx)?;
+                            let target = self
+                                .expression_for_reference(args.next()?, ctx)?
+                                .ref_or(Error::InvalidCooperativeMatrix)?;
                             let pointer = self.expression(args.next()?, ctx)?;
                             let stride = if args.total_args > 2 {
                                 self.expression(args.next()?, ctx)?
@@ -3178,9 +3185,15 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                         }
                         "coopMultiplyAdd" => {
                             let mut args = ctx.prepare_args(arguments, 3, span);
-                            let a = self.expression(args.next()?, ctx)?;
-                            let b = self.expression(args.next()?, ctx)?;
-                            let c = self.expression(args.next()?, ctx)?;
+                            let a = self
+                                .expression_for_reference(args.next()?, ctx)?
+                                .ref_or(Error::InvalidCooperativeMatrix)?;
+                            let b = self
+                                .expression_for_reference(args.next()?, ctx)?
+                                .ref_or(Error::InvalidCooperativeMatrix)?;
+                            let c = self
+                                .expression_for_reference(args.next()?, ctx)?
+                                .ref_or(Error::InvalidCooperativeMatrix)?;
                             args.finish()?;
 
                             ir::Expression::CooperativeMultiplyAdd { a, b, c }
diff --git a/naga/src/proc/type_methods.rs b/naga/src/proc/type_methods.rs
@@ -191,17 +191,6 @@ impl crate::TypeInner {
         }
     }
 
-    /// Returns true if a variable of this type is a handle.
-    pub const fn is_handle(&self) -> bool {
-        match *self {
-            Self::Image { .. }
-            | Self::Sampler { .. }
-            | Self::AccelerationStructure { .. }
-            | Self::CooperativeMatrix { .. } => true,
-            _ => false,
-        }
-    }
-
     /// Attempt to calculate the size of this type. Returns `None` if the size
     /// exceeds the limit of [`crate::valid::MAX_TYPE_SIZE`].
     pub fn try_size(&self, gctx: super::GlobalCtx) -> Option<u32> {
diff --git a/naga/src/proc/typifier.rs b/naga/src/proc/typifier.rs
@@ -454,8 +454,7 @@ impl<'a> ResolveContext<'a> {
             }
             crate::Expression::GlobalVariable(h) => {
                 let var = &self.global_vars[h];
-                let ty = &types[var.ty].inner;
-                if var.space == crate::AddressSpace::Handle || ty.is_handle() {
+                if var.space == crate::AddressSpace::Handle {
                     TypeResolution::Handle(var.ty)
                 } else {
                     TypeResolution::Value(Ti::Pointer {
@@ -466,15 +465,10 @@ impl<'a> ResolveContext<'a> {
             }
             crate::Expression::LocalVariable(h) => {
                 let var = &self.local_vars[h];
-                let ty = &types[var.ty].inner;
-                if ty.is_handle() {
-                    TypeResolution::Handle(var.ty)
-                } else {
-                    TypeResolution::Value(Ti::Pointer {
-                        base: var.ty,
-                        space: crate::AddressSpace::Function,
-                    })
-                }
+                TypeResolution::Value(Ti::Pointer {
+                    base: var.ty,
+                    space: crate::AddressSpace::Function,
+                })
             }
             crate::Expression::Load { pointer } => match *past(pointer)?.inner_with(types) {
                 Ti::Pointer { base, space: _ } => {
@@ -807,7 +801,15 @@ impl<'a> ResolveContext<'a> {
                 scalar: crate::Scalar::U32,
                 size: crate::VectorSize::Quad,
             }),
-            crate::Expression::CooperativeMultiplyAdd { a: _, b: _, c } => past(c)?.clone(),
+            crate::Expression::CooperativeMultiplyAdd { a: _, b: _, c } => {
+                match *past(c)?.inner_with(types) {
+                    Ti::Pointer { base, space: _ } => TypeResolution::Handle(base),
+                    ref other => {
+                        log::error!("Pointer type {other:?}");
+                        return Err(ResolveError::InvalidPointer(c));
+                    }
+                }
+            }
         })
     }
 }
diff --git a/naga/src/valid/expression.rs b/naga/src/valid/expression.rs
@@ -1256,34 +1256,23 @@ impl super::Validator {
             },
             E::SubgroupBallotResult | E::SubgroupOperationResult { .. } => self.subgroup_stages,
             E::CooperativeMultiplyAdd { a, b, c } => {
-                match resolver[a] {
-                    Ti::CooperativeMatrix {
-                        role: crate::CooperativeRole::A,
-                        ..
-                    } => {}
-                    ref other => {
-                        log::error!("A operand type: {other:?}");
-                        return Err(ExpressionError::InvalidCooperativeOperand(a));
-                    }
-                }
-                match resolver[b] {
-                    Ti::CooperativeMatrix {
-                        role: crate::CooperativeRole::B,
-                        ..
-                    } => {}
-                    ref other => {
-                        log::error!("B operand type: {other:?}");
-                        return Err(ExpressionError::InvalidCooperativeOperand(b));
-                    }
-                }
-                match resolver[c] {
-                    Ti::CooperativeMatrix {
-                        role: crate::CooperativeRole::C,
-                        ..
-                    } => {}
-                    ref other => {
-                        log::error!("C operand type: {other:?}");
-                        return Err(ExpressionError::InvalidCooperativeOperand(c));
+                let roles = [
+                    crate::CooperativeRole::A,
+                    crate::CooperativeRole::B,
+                    crate::CooperativeRole::C,
+                ];
+                for (operand, expected_role) in [a, b, c].into_iter().zip(roles) {
+                    match resolver[operand] {
+                        Ti::Pointer { base, space: _ } => match module.types[base].inner {
+                            Ti::CooperativeMatrix { role, .. } if role == expected_role => {}
+                            ref other => {
+                                log::error!("{expected_role:?} operand type: {other:?}");
+                                return Err(ExpressionError::InvalidCooperativeOperand(a));
+                            }
+                        },
+                        _ => {
+                            return Err(ExpressionError::InvalidPointerType(operand));
+                        }
                     }
                 }
                 ShaderStages::COMPUTE
diff --git a/naga/src/valid/function.rs b/naga/src/valid/function.rs
@@ -1078,7 +1078,7 @@ impl super::Validator {
                     } else if let Some(tr) = pointer_base_tr {
                         context.compare_types(value_tr, &tr)
                     } else {
-                        value_ty.is_handle()
+                        false
                     };
 
                     if !good {
@@ -1631,15 +1631,22 @@ impl super::Validator {
                 } => {
                     stages &= super::ShaderStages::COMPUTE;
 
-                    let target_scalar =
-                        match *context.resolve_type_inner(target, &self.valid_expression_set)? {
+                    let target_scalar = match *context.resolve_pointer_type(target) {
+                        crate::TypeInner::Pointer { base, space: _ } => match context.types[base]
+                            .inner
+                        {
                             Ti::CooperativeMatrix { scalar, .. } => scalar,
                             ref other => {
                                 log::error!("Target operand type: {other:?}");
                                 return Err(FunctionError::InvalidCooperativeStoreTarget(target)
                                     .with_span_handle(target, context.expressions));
                             }
-                        };
+                        },
+                        _ => {
+                            return Err(FunctionError::InvalidCooperativeStoreTarget(target)
+                                .with_span_handle(target, context.expressions));
+                        }
+                    };
 
                     let ty_inner = context.resolve_pointer_type(pointer);
                     //TODO: validate stride
diff --git a/naga/tests/in/wgsl/cooperative-matrix.wgsl b/naga/tests/in/wgsl/cooperative-matrix.wgsl
diff --git a/naga/tests/out/ir/wgsl-cooperative-matrix.compact.ron b/naga/tests/out/ir/wgsl-cooperative-matrix.compact.ron
diff --git a/naga/tests/out/ir/wgsl-cooperative-matrix.ron b/naga/tests/out/ir/wgsl-cooperative-matrix.ron
diff --git a/naga/tests/out/msl/wgsl-cooperative-matrix.msl b/naga/tests/out/msl/wgsl-cooperative-matrix.msl
diff --git a/naga/tests/out/spv/wgsl-cooperative-matrix.spvasm b/naga/tests/out/spv/wgsl-cooperative-matrix.spvasm
diff --git a/naga/tests/out/wgsl/wgsl-cooperative-matrix.wgsl b/naga/tests/out/wgsl/wgsl-cooperative-matrix.wgsl

Original file line number	Diff line number	Diff line change
`@@ -971,14 +971,13 @@ impl Writer {`
`971`	`971`	`}`
`972`	`972`	`}`
`973`	`973`
`974`		`- // Handle globals are pre-emitted and should be loaded automatically.`
`975`		`- //`
`976`		`- // Any that are binding arrays we skip as we cannot load the array, we must load the result after indexing.`
`977`	`974`	`match ir_module.types[var.ty].inner {`
	`975`	`+ // Any that are binding arrays we skip as we cannot load the array, we must load the result after indexing.`
`978`	`976`	`crate::TypeInner::BindingArray { .. } => {`
`979`	`977`	`gv.access_id = gv.var_id;`
`980`	`978`	`}`
`981`	`979`	`_ => {`
	`980`	`+ // Handle globals are pre-emitted and should be loaded automatically.`
`982`	`981`	`if var.space == crate::AddressSpace::Handle {`
`983`	`982`	`let var_type_id = self.get_handle_type_id(var.ty);`
`984`	`983`	`let id = self.id_gen.next();`
`@@ -1064,6 +1063,7 @@ impl Writer {`
`1064`	`1063`	`}`
`1065`	`1064`	`}),`
`1066`	`1065`	`);`
	`1066`	`+`
`1067`	`1067`	`context`
`1068`	`1068`	`.function`
`1069`	`1069`	`.variables`