coop: first bits of Vulkan support for the type

Author: kvark

naga/src/back/spv/instructions.rs

@@ -281,6 +281,22 @@ impl super::Instruction {
         instruction
     }
 
+    pub(super) fn type_coop_matrix(
+        id: Word,
+        scalar_type_id: Word,
+        row_count: crate::CooperativeVectorSize,
+        column_count: crate::CooperativeVectorSize,
+    ) -> Self {
+        let mut instruction = Self::new(Op::TypeCooperativeMatrixKHR);
+        instruction.set_result(id);
+        instruction.add_operand(scalar_type_id);
+        instruction.add_operand(spirv::Scope::Subgroup as u32);
+        instruction.add_operand(column_count as u32);
+        instruction.add_operand(row_count as u32);
+        instruction.add_operand(spirv::CooperativeMatrixUse::MatrixAKHR as u32); //TODO: configure or expose
+        instruction
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub(super) fn type_image(
         id: Word,

naga/src/back/spv/mod.rs

@@ -338,6 +338,33 @@ impl NumericType {
     }
 }
 
+/// A cooperative type, for use in [`LocalType`].
+#[derive(Debug, PartialEq, Hash, Eq, Copy, Clone)]
+enum CooperativeType {
+    Matrix {
+        columns: crate::CooperativeVectorSize,
+        rows: crate::CooperativeVectorSize,
+        scalar: crate::CooperativeScalar,
+    },
+}
+
+impl CooperativeType {
+    const fn from_inner(inner: &crate::TypeInner) -> Option<Self> {
+        match *inner {
+            crate::TypeInner::CooperativeMatrix {
+                columns,
+                rows,
+                scalar,
+            } => Some(Self::Matrix {
+                columns,
+                rows,
+                scalar,
+            }),
+            _ => None,
+        }
+    }
+}
+
 /// A SPIR-V type constructed during code generation.
 ///
 /// This is the variant of [`LookupType`] used to represent types that might not
@@ -387,6 +414,7 @@ impl NumericType {
 enum LocalType {
     /// A numeric type.
     Numeric(NumericType),
+    Cooperative(CooperativeType),
     Pointer {
         base: Word,
         class: spirv::StorageClass,

naga/src/back/spv/writer.rs

@@ -6,10 +6,11 @@ use spirv::Word;
 use super::{
     block::DebugInfoInner,
     helpers::{contains_builtin, global_needs_wrapper, map_storage_class},
-    Block, BlockContext, CachedConstant, CachedExpressions, DebugInfo, EntryPointContext, Error,
-    Function, FunctionArgument, GlobalVariable, IdGenerator, Instruction, LocalImageType,
-    LocalType, LocalVariable, LogicalLayout, LookupFunctionType, LookupType, NumericType, Options,
-    PhysicalLayout, PipelineOptions, ResultMember, Writer, WriterFlags, BITS_PER_BYTE,
+    Block, BlockContext, CachedConstant, CachedExpressions, CooperativeType, DebugInfo,
+    EntryPointContext, Error, Function, FunctionArgument, GlobalVariable, IdGenerator, Instruction,
+    LocalImageType, LocalType, LocalVariable, LogicalLayout, LookupFunctionType, LookupType,
+    NumericType, Options, PhysicalLayout, PipelineOptions, ResultMember, Writer, WriterFlags,
+    BITS_PER_BYTE,
 };
 use crate::{
     arena::{Handle, HandleVec, UniqueArena},
@@ -373,6 +374,12 @@ impl Writer {
         })
     }
 
+    pub(super) fn get_cooperative_type_id(&mut self, scalar: crate::CooperativeScalar) -> Word {
+        match scalar {
+            crate::CooperativeScalar::F32 => self.get_f32_type_id(),
+        }
+    }
+
     pub(super) fn get_f32_pointer_type_id(&mut self, class: spirv::StorageClass) -> Word {
         let f32_id = self.get_f32_type_id();
         self.get_pointer_type_id(f32_id, class)
@@ -434,7 +441,9 @@ impl Writer {
                 // these cases, so unwrap.
                 LocalType::Numeric(NumericType::from_inner(inner).unwrap())
             }
-            crate::TypeInner::CooperativeMatrix { .. } => return None,
+            crate::TypeInner::CooperativeMatrix { .. } => {
+                LocalType::Cooperative(CooperativeType::from_inner(inner).unwrap())
+            }
             crate::TypeInner::Pointer { base, space } => {
                 let base_type_id = self.get_handle_type_id(base);
                 LocalType::Pointer {
@@ -1331,6 +1340,14 @@ impl Writer {
                 self.require_any("16 bit floating-point", &[spirv::Capability::Float16])?;
                 self.use_extension("SPV_KHR_16bit_storage");
             }
+            // Cooperative types and ops
+            crate::TypeInner::CooperativeMatrix { .. } => {
+                self.require_any(
+                    "cooperative matrix",
+                    &[spirv::Capability::CooperativeMatrixKHR],
+                )?;
+                self.use_extension("SPV_KHR_cooperative_matrix");
+            }
             _ => {}
         }
         Ok(())
@@ -1357,12 +1374,31 @@ impl Writer {
         instruction.to_words(&mut self.logical_layout.declarations);
     }
 
+    fn write_cooperative_type_declaration_local(&mut self, id: Word, coop: CooperativeType) {
+        let instruction = match coop {
+            CooperativeType::Matrix {
+                columns,
+                rows,
+                scalar,
+            } => {
+                let scalar_id = self.get_cooperative_type_id(scalar);
+                Instruction::type_coop_matrix(id, scalar_id, rows, columns)
+            }
+        };
+
+        instruction.to_words(&mut self.logical_layout.declarations);
+    }
+
     fn write_type_declaration_local(&mut self, id: Word, local_ty: LocalType) {
         let instruction = match local_ty {
             LocalType::Numeric(numeric) => {
                 self.write_numeric_type_declaration_local(id, numeric);
                 return;
             }
+            LocalType::Cooperative(coop) => {
+                self.write_cooperative_type_declaration_local(id, coop);
+                return;
+            }
             LocalType::Pointer { base, class } => Instruction::type_pointer(id, class, base),
             LocalType::Image(image) => {
                 let local_type = LocalType::Numeric(NumericType::Scalar(image.sampled_type));

naga/src/valid/mod.rs

@@ -186,6 +186,8 @@ bitflags::bitflags! {
         /// Support for `quantizeToF16`, `pack2x16float`, and `unpack2x16float`, which store
         /// `f16`-precision values in `f32`s.
         const SHADER_FLOAT16_IN_FLOAT32 = 1 << 28;
+        /// Support for cooperative matrix types and operations
+        const COOPERATIVE_MATRIX = 1 << 29;
     }
 }
 

naga/src/valid/type.rs

@@ -420,6 +420,7 @@ impl super::Validator {
                 rows: _,
                 scalar,
             } => {
+                self.require_type_capability(Capabilities::COOPERATIVE_MATRIX)?;
                 if scalar != crate::CooperativeScalar::F32 {
                     return Err(TypeError::MatrixElementNotFloat);
                 }