Add cooperative matrix 1 support (#6565)

* initial wip for spirv

* working tiled example

* clean up store and load

* minor fixes

* fix loadAny name

* add initial tests, including broken/unimplemented intrinsics

* fix subscript

* run tests at 16x16, remove not supported arithmetic tests

* minor fixups on implementation

* rename CoopMatMatrixUse

* Update tests to pass validation layers locally

* Add mat-mul-add test and minor fixes

* Add more tests

* Remove dead code

* Add coopMatLoad function and tests, enforce constexpr for matrix layout

* Use getVectorOrCoopMatrixElementType in place of getVectorElementType

2 files changed, 40 insertions, 0 deletions

diff --git a/tests/cooperative-matrix/diagnostics/mat-mul-add-different-scope.slang b/tests/cooperative-matrix/diagnostics/mat-mul-add-different-scope.slang
new file mode 100644
index 000000000..0c4308308
--- /dev/null
+++ b/tests/cooperative-matrix/diagnostics/mat-mul-add-different-scope.slang
@@ -0,0 +1,20 @@
+//DIAGNOSTIC_TEST(compute):SIMPLE(filecheck=CHECK): -entry computeMain -stage compute -target spirv
+
+RWStructuredBuffer<float> outputBuffer;
+
+typealias CoopMatAType = CoopMat<float16_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixA>;
+typealias CoopMatBType = CoopMat<float16_t, CoopMatScope::Workgroup, 16, 16, CoopMatMatrixUse::MatrixB>;
+typealias CoopMatCType = CoopMat<float32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+
+// CHECK: error 39999: could not specialize generic for arguments of type
+
+[numthreads(32, 1, 1)]
+void computeMain()
+{
+    let matA = CoopMatAType(3.0);
+    let matB = CoopMatBType(5.0);
+    let matC = CoopMatCType(1.0);
+
+    const let result = coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::None);
+    result.store(outputBuffer, 0, 16, CoopMatMatrixLayout::RowMajor);
+}
diff --git a/tests/cooperative-matrix/diagnostics/mat-mul-add-incorrect-matrix-use.slang b/tests/cooperative-matrix/diagnostics/mat-mul-add-incorrect-matrix-use.slang
new file mode 100644
index 000000000..5b7dc7a5b
--- /dev/null
+++ b/tests/cooperative-matrix/diagnostics/mat-mul-add-incorrect-matrix-use.slang
@@ -0,0 +1,20 @@
+//DIAGNOSTIC_TEST(compute):SIMPLE(filecheck=CHECK): -entry computeMain -stage compute -target spirv
+
+RWStructuredBuffer<float> outputBuffer;
+
+typealias CoopMatAType = CoopMat<float16_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixA>;
+typealias CoopMatBType = CoopMat<float16_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixA>;
+typealias CoopMatCType = CoopMat<float32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+
+// CHECK: error 41400: static assertion failed, matrix uses for `coopMatMulAdd` matrix parameters must be `MatrixA`, `MatrixB` and `MatrixAccumulator`
+
+[numthreads(32, 1, 1)]
+void computeMain()
+{
+    let matA = CoopMatAType(3.0);
+    let matB = CoopMatBType(5.0);
+    let matC = CoopMatCType(1.0);
+
+    let result = coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::None);
+    result.store(outputBuffer, 0, 16, CoopMatMatrixLayout::RowMajor);
+}