coop: rewire WGSL support using references

Author: kvark

CHANGELOG.md

@@ -376,7 +376,7 @@ By @wumpf in [#8282](https://github.com/gfx-rs/wgpu/pull/8282), [#8285](https://
 
 - Expose `naga::front::wgsl::UnimplementedEnableExtension`. By @ErichDonGubler in [#8237](https://github.com/gfx-rs/wgpu/pull/8237).
 
-- 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/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.
     ///

naga/src/back/msl/writer.rs

@@ -6327,16 +6327,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 {

naga/src/back/spv/block.rs

@@ -3719,7 +3719,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,

naga/src/back/spv/writer.rs

@@ -970,14 +970,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();
@@ -1063,6 +1062,7 @@ impl Writer {
                     }
                 }),
             );
+
             context
                 .function
                 .variables

naga/src/back/wgsl/writer.rs

@@ -994,13 +994,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, ");")?
             }
         }
 
@@ -1715,11 +1715,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

@@ -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.
@@ -1679,12 +1695,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 {
@@ -1979,12 +1991,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) {
@@ -2141,10 +2148,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),
                         }
                     }
@@ -3142,7 +3147,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)?
@@ -3180,9 +3187,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/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> {

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));
+                    }
+                }
+            }
         })
     }
 }