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/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/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")]

naga/src/back/msl/writer.rs

@@ -236,6 +236,7 @@ impl Display for TypeContext<'_> {
                 rows,
                 scalar,
             } => put_numeric_type(out, scalar, &[rows, columns]),
+            // Requires Metal-2.3
             crate::TypeInner::CooperativeMatrix {
                 columns,
                 rows,
@@ -244,8 +245,7 @@ impl Display for TypeContext<'_> {
             } => {
                 write!(
                     out,
-                    "{}::simdgroup_{}{}x{}",
-                    NAMESPACE,
+                    "{NAMESPACE}::simdgroup_{}{}x{}",
                     scalar.to_msl_name(),
                     columns as u32,
                     rows as u32,
@@ -485,6 +485,7 @@ enum WrappedFunction {
         class: crate::ImageClass,
     },
     CooperativeMultiplyAdd {
+        space: crate::AddressSpace,
         columns: crate::CooperativeSize,
         rows: crate::CooperativeSize,
         intermediate: crate::CooperativeSize,
@@ -2842,6 +2843,9 @@ impl<W: Write> Writer<W> {
                 write!(self.out, "}}")?;
             }
             crate::Expression::CooperativeMultiplyAdd { a, b, c } => {
+                if context.lang_version < (2, 3) {
+                    return Err(Error::UnsupportedCooperativeMatrix);
+                }
                 write!(self.out, "{COOPERATIVE_MULTIPLY_ADD_FUNCTION}(")?;
                 self.put_expression(a, context, true)?;
                 write!(self.out, ", ")?;
@@ -4239,10 +4243,14 @@ impl<W: Write> Writer<W> {
                     row_major,
                 } => {
                     let op_str = if store { "store" } else { "load" };
-                    write!(self.out, "{level}{NAMESPACE}::simdgroup_{op_str}(")?;
+                    write!(self.out, "{level}simdgroup_{op_str}(")?;
                     self.put_expression(target, &context.expression, true)?;
-                    write!(self.out, ", ")?;
-                    self.put_expression(pointer, &context.expression, true)?;
+                    write!(self.out, ", &")?;
+                    self.put_access_chain(
+                        pointer,
+                        context.expression.policies.index,
+                        &context.expression,
+                    )?;
                     write!(self.out, ", ")?;
                     self.put_expression(stride, &context.expression, true)?;
                     if row_major {
@@ -6312,23 +6320,31 @@ template <typename A>
         &mut self,
         module: &crate::Module,
         func_ctx: &back::FunctionCtx,
+        space: crate::AddressSpace,
         a: Handle<crate::Expression>,
         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 {
+            space,
             columns: b_c,
             rows: a_r,
             intermediate: a_c,
@@ -6337,15 +6353,11 @@ template <typename A>
         if !self.wrapped_functions.insert(wrapped) {
             return Ok(());
         }
-        let scalar_name = match scalar.width {
-            2 => "half",
-            4 => "float",
-            8 => "double",
-            _ => unreachable!(),
-        };
+        let space_name = space.to_msl_name().unwrap_or_default();
+        let scalar_name = scalar.to_msl_name();
         writeln!(
             self.out,
-            "{NAMESPACE}::simdgroup_{scalar_name}{}x{} {COOPERATIVE_MULTIPLY_ADD_FUNCTION}(const {NAMESPACE}::simdgroup_{scalar_name}{}x{}& a, const {NAMESPACE}::simdgroup_{scalar_name}{}x{}& b, const {NAMESPACE}::simdgroup_{scalar_name}{}x{}& c) {{",
+            "{NAMESPACE}::simdgroup_{scalar_name}{}x{} {COOPERATIVE_MULTIPLY_ADD_FUNCTION}(const {space_name} {NAMESPACE}::simdgroup_{scalar_name}{}x{}& a, const {space_name} {NAMESPACE}::simdgroup_{scalar_name}{}x{}& b, const {space_name} {NAMESPACE}::simdgroup_{scalar_name}{}x{}& c) {{",
             b_c as u32, a_r as u32, a_c as u32, a_r as u32, b_c as u32, b_r as u32, b_c as u32, a_r as u32,
         )?;
         let l1 = back::Level(1);
@@ -6354,10 +6366,7 @@ template <typename A>
             "{l1}{NAMESPACE}::simdgroup_{scalar_name}{}x{} d;",
             b_c as u32, a_r as u32
         )?;
-        writeln!(
-            self.out,
-            "{l1}{NAMESPACE}::simdgroup_multiply_accumulate(d,a,b,c);"
-        )?;
+        writeln!(self.out, "{l1}simdgroup_multiply_accumulate(d,a,b,c);")?;
         writeln!(self.out, "{l1}return d;")?;
         writeln!(self.out, "}}")?;
         writeln!(self.out)?;
@@ -6439,7 +6448,8 @@ template <typename A>
                     self.write_wrapped_image_query(module, func_ctx, image, query)?;
                 }
                 crate::Expression::CooperativeMultiplyAdd { a, b, c: _ } => {
-                    self.write_wrapped_cooperative_multiply_add(module, func_ctx, a, b)?;
+                    let space = crate::AddressSpace::Private;
+                    self.write_wrapped_cooperative_multiply_add(module, func_ctx, space, a, b)?;
                 }
                 _ => {}
             }
@@ -6632,7 +6642,6 @@ template <typename A>
                     names: &self.names,
                     handle,
                     usage: fun_info[handle],
-
                     reference: true,
                 };
                 let separator =

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,

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

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

@@ -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 }

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> {