coop: rewire WGSL support using references

Author: kvark

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
 

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, ")")?;
             }
         }

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![],

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

@@ -955,6 +955,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 +1684,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 +1980,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 +2137,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 +3136,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 +3176,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 }

naga/tests/out/wgsl/wgsl-cooperative-matrix.wgsl

@@ -8,6 +8,8 @@ fn main() {
     var c: coop_mat8x8<f32,C> = coop_mat8x8<f32,C>();
     var d: coop_mat8x8<f32,C>;
 
-coopLoad((&c), (&ext[4]), 8u)    d = coopMultiplyAdd((&a), (&b), (&c));
-coopStore((&c), (&ext[0]), 8u)    return;
+    coopLoad(c, (&ext[4]), 8u);
+    d = coopMultiplyAdd(a, b, c);
+    coopStore(c, (&ext[0]), 8u);
+    return;
 }