coop: support generic argument on coopLoad

Author: kvark

naga/src/back/wgsl/writer.rs

@@ -1712,9 +1712,29 @@ impl<W: Write> Writer<W> {
             | Expression::SubgroupBallotResult
             | Expression::SubgroupOperationResult { .. }
             | Expression::WorkGroupUniformLoadResult { .. } => {}
-            Expression::CooperativeLoad { ref data, .. } => {
+            Expression::CooperativeLoad {
+                columns,
+                rows,
+                role,
+                ref data,
+            } => {
                 let suffix = if data.row_major { "T" } else { "" };
-                write!(self.out, "coopLoad{suffix}(")?;
+                let scalar = func_ctx.info[data.pointer]
+                    .ty
+                    .inner_with(&module.types)
+                    .pointer_base_type()
+                    .unwrap()
+                    .inner_with(&module.types)
+                    .scalar()
+                    .unwrap();
+                write!(
+                    self.out,
+                    "coopLoad{suffix}<coop_mat{}x{}<{},{:?}>>(",
+                    columns as u32,
+                    rows as u32,
+                    scalar.try_to_wgsl().unwrap(),
+                    role,
+                )?;
                 self.write_expr(module, data.pointer, func_ctx)?;
                 write!(self.out, ", ")?;
                 self.write_expr(module, data.stride, func_ctx)?;

naga/src/front/wgsl/error.rs

@@ -413,6 +413,7 @@ pub(crate) enum Error<'a> {
         span: Span,
     },
     UnderspecifiedCooperativeMatrix,
+    InvalidCooperativeLoadType(Span),
     UnsupportedCooperativeScalar(Span),
 }
 
@@ -1393,6 +1394,11 @@ impl<'a> Error<'a> {
                 labels: vec![],
                 notes: vec![format!("must be F32")],
             },
+            Error::InvalidCooperativeLoadType(span) => ParseError {
+                message: "cooperative load should have a generic type for coop_mat".into(),
+                labels: vec![(span, "type needs the coop_mat<...>".into())],
+                notes: vec![format!("must be a valid cooperative type")],
+            },
             Error::UnsupportedCooperativeScalar(span) => ParseError {
                 message: "cooperative scalar type is not supported".into(),
                 labels: vec![(span, "type needs the scalar type specified".into())],

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

@@ -1905,6 +1905,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     stmt.span,
                     function,
                     arguments,
+                    None,
                     &mut ctx.as_expression(block, &mut emitter),
                     true,
                 )?;
@@ -2227,9 +2228,10 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             ast::Expression::Call {
                 ref function,
                 ref arguments,
+                result_ty,
             } => {
                 let handle = self
-                    .call(span, function, arguments, ctx, false)?
+                    .call(span, function, arguments, result_ty, ctx, false)?
                     .ok_or(Error::FunctionReturnsVoid(function.span))?;
                 return Ok(Typed::Plain(handle));
             }
@@ -2424,6 +2426,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         span: Span,
         function: &ast::Ident<'source>,
         arguments: &[Handle<ast::Expression<'source>>],
+        result_ty: Option<(Handle<ast::Type<'source>>, Span)>,
         ctx: &mut ExpressionContext<'source, '_, '_>,
         is_statement: bool,
     ) -> Result<'source, Option<Handle<ir::Expression>>> {
@@ -3145,9 +3148,20 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                             let row_major = function.name.ends_with("T");
                             let mut args = ctx.prepare_args(arguments, 1, span);
                             let pointer = self.expression(args.next()?, ctx)?;
-                            //TODO: read from generic argument
-                            let columns = crate::CooperativeSize::Eight;
-                            let rows = crate::CooperativeSize::Eight;
+                            let (matrix_ty, matrix_span) = result_ty.expect("generic argument");
+                            let (columns, rows, role) = match ctx.types[matrix_ty] {
+                                ast::Type::CooperativeMatrix {
+                                    columns,
+                                    rows,
+                                    role,
+                                    ..
+                                } => (columns, rows, role),
+                                _ => {
+                                    return Err(Box::new(Error::InvalidCooperativeLoadType(
+                                        matrix_span,
+                                    )))
+                                }
+                            };
                             let stride = if args.total_args > 1 {
                                 self.expression(args.next()?, ctx)?
                             } else {
@@ -3167,7 +3181,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                             crate::Expression::CooperativeLoad {
                                 columns,
                                 rows,
-                                role: crate::CooperativeRole::C, //TODO
+                                role,
                                 data: crate::CooperativeData {
                                     pointer,
                                     stride,

naga/src/front/wgsl/parse/ast.rs

@@ -487,6 +487,7 @@ pub enum Expression<'a> {
     Call {
         function: Ident<'a>,
         arguments: Vec<Handle<Expression<'a>>>,
+        result_ty: Option<(Handle<Type<'a>>, Span)>,
     },
     Index {
         base: Handle<Expression<'a>>,

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

@@ -800,6 +800,11 @@ impl Parser {
             }
             // everything else must be handled later, since they can be hidden by user-defined functions.
             _ => {
+                let result_ty = if lexer.peek().0 == Token::Paren('<') {
+                    Some(self.singular_generic(lexer, ctx)?)
+                } else {
+                    None
+                };
                 let arguments = self.arguments(lexer, ctx)?;
                 ctx.unresolved.insert(ast::Dependency {
                     ident: name,
@@ -811,6 +816,7 @@ impl Parser {
                         span: name_span,
                     },
                     arguments,
+                    result_ty,
                 }
             }
         };
@@ -959,7 +965,7 @@ impl Parser {
                 } else if let Token::Paren('(') = lexer.peek().0 {
                     self.pop_rule_span(lexer);
                     return self.function_call(lexer, word, span, ctx);
-                } else if word == "bitcast" {
+                } else if ["bitcast", "coopLoad"].contains(&word) {
                     self.pop_rule_span(lexer);
                     return self.function_call(lexer, word, span, ctx);
                 } else {

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

@@ -5,7 +5,7 @@ var<storage, read_write> ext: array<f32>;
 
 @compute @workgroup_size(8, 8, 1)
 fn main() {
-    var c = coopLoad(&ext[4]);
+    var c = coopLoad<coop_mat8x8<f32, C>>(&ext[4]);
     var d = coopMultiplyAdd(a, b, c);
     coopStore(d, &ext[0]);
     c = d;

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

@@ -19,7 +19,7 @@ var<storage, read_write> ext: array<f32>;
 
 @compute @workgroup_size(8, 8, 1)
 fn main() {
-    var c = coopLoad(&ext[4]);
+    var c = coopLoad<coop_mat8x8<f32, C>>(&ext[4]);
     var d = coopMultiplyAdd(a, b, c);
     coopStore(d, &ext[0]);
     c = d;
@@ -71,7 +71,7 @@ OpMemberDecorate %22 0 Offset 0
 %28 = OpAccessChain  %27  %21 %9
 OpBranch %34
 %34 = OpLabel
-OpLine %3 8 23
+OpLine %3 8 44
 OpLine %3 8 13
 %37 = OpAccessChain  %35  %28 %36
 %38 = OpCooperativeMatrixLoadKHR  %13  %37 %11 %8

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

@@ -8,7 +8,7 @@ fn main() {
     var c: coop_mat8x8<f32,C>;
     var d: coop_mat8x8<f32,C>;
 
-    c = coopLoad((&ext[4]), 8u);
+    c = coopLoad<coop_mat8x8<f32,C>>((&ext[4]), 8u);
     let _e6 = a;
     let _e8 = b;
     let _e9 = c;