From 2580bb02f7a079ab1c0106b5960a21ed1627bca0 Mon Sep 17 00:00:00 2001
From: Jay Kwak <82421531+jkwak-work@users.noreply.github.com>
Date: Thu, 15 May 2025 03:47:43 +0000
Subject: Add new coopmat2 functions: Reduce and Transpose (#7027)

This commit adds three new functions for CoopMat as described in the proposal document,

Cooperative matrix 2 proposal spec#12
The new functions are:

CoopMat<T,S,M,N,R>::Transpose
CoopMat<T,S,M,N,R>::ReduceRow
CoopMat<T,S,M,N,R>::ReduceColumn
CoopMat<T,S,M,N,R>::ReduceRowAndColumn
CoopMat<T,S,M,N,R>::Reduce2x2
---
 tests/cooperative-matrix/add.slang                 | 10 ++--
 tests/cooperative-matrix/array.slang               | 12 ++--
 tests/cooperative-matrix/comparison.slang          |  8 ++-
 tests/cooperative-matrix/conversion.slang          | 13 +++--
 tests/cooperative-matrix/copyFrom.slang            |  8 ++-
 .../diagnostics/mat-mul-add-different-scope.slang  | 20 -------
 .../mat-mul-add-incorrect-matrix-use.slang         | 20 -------
 tests/cooperative-matrix/div.slang                 | 10 ++--
 tests/cooperative-matrix/fill.slang                |  6 +-
 tests/cooperative-matrix/inout.slang               |  8 ++-
 tests/cooperative-matrix/length.slang              | 21 +++++++
 .../load-store-arbitrary-array-vec.slang           | 14 ++---
 .../load-store-arbitrary-array.slang               | 13 +++--
 .../load-store-groupshared.slang                   | 13 +++--
 .../load-store-rwbyteaddressbuffer.slang           | 13 +++--
 .../load-store-rwstructuredbuffer.slang            | 13 +++--
 tests/cooperative-matrix/map-element-single.slang  | 50 ++++++++++++++++
 .../mat-mul-add-spirv-matrix-operands.slang        | 67 +++++++++++++++-------
 tests/cooperative-matrix/mat-mul-add.slang         | 12 ++--
 tests/cooperative-matrix/mod.slang                 | 21 +++----
 tests/cooperative-matrix/mul.slang                 | 10 ++--
 tests/cooperative-matrix/out.slang                 |  8 ++-
 tests/cooperative-matrix/parameter.slang           |  8 ++-
 tests/cooperative-matrix/reduce.slang              | 52 +++++++++++++++++
 tests/cooperative-matrix/return.slang              |  8 ++-
 tests/cooperative-matrix/scalar-mul.slang          |  8 ++-
 tests/cooperative-matrix/struct.slang              | 12 ++--
 tests/cooperative-matrix/sub.slang                 | 10 ++--
 tests/cooperative-matrix/subscript-in-func.slang   |  8 ++-
 tests/cooperative-matrix/subscript.slang           |  6 +-
 tests/cooperative-matrix/transpose.slang           | 36 ++++++++++++
 tests/cooperative-matrix/unary_neg.slang           |  8 ++-
 32 files changed, 359 insertions(+), 167 deletions(-)
 delete mode 100644 tests/cooperative-matrix/diagnostics/mat-mul-add-different-scope.slang
 delete mode 100644 tests/cooperative-matrix/diagnostics/mat-mul-add-incorrect-matrix-use.slang
 create mode 100644 tests/cooperative-matrix/length.slang
 create mode 100644 tests/cooperative-matrix/map-element-single.slang
 create mode 100644 tests/cooperative-matrix/reduce.slang
 create mode 100644 tests/cooperative-matrix/transpose.slang

(limited to 'tests')

diff --git a/tests/cooperative-matrix/add.slang b/tests/cooperative-matrix/add.slang
index ecf2f16ba..c904da0f7 100644
--- a/tests/cooperative-matrix/add.slang
+++ b/tests/cooperative-matrix/add.slang
@@ -15,7 +15,9 @@ ByteAddressBuffer input1;
 //TEST_INPUT:ubuffer(data=[0 1 2 3], stride=4, count=256),name=input2
 ByteAddressBuffer input2;
 
-typealias CoopMatType = CoopMat<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
@@ -23,10 +25,10 @@ void computeMain()
     let stride = 16;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat1 = CoopMatType.load(input1, 0, stride, matrixLayout);
-    let mat2 = CoopMatType.load(input2, 0, stride, matrixLayout);
+    let mat1 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input1, 0, stride);
+    let mat2 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input2, 0, stride);
     let result = mat1 + mat2;
 
-    result.store(outputBuffer, 0, stride, matrixLayout);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, stride);
 }
 
diff --git a/tests/cooperative-matrix/array.slang b/tests/cooperative-matrix/array.slang
index ab1f92a99..eb26bacbf 100644
--- a/tests/cooperative-matrix/array.slang
+++ b/tests/cooperative-matrix/array.slang
@@ -19,7 +19,9 @@ ByteAddressBuffer input2;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
-typealias CoopMatType = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
@@ -28,9 +30,9 @@ void computeMain()
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
     CoopMatType coopMatArray[2];
-    coopMatArray[0] = CoopMatType.load(input1, 0, stride, matrixLayout);
-    coopMatArray[1] = CoopMatType.load(input2, 0, stride, matrixLayout);
+    coopMatArray[0] = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input1, 0, stride);
+    coopMatArray[1] = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input2, 0, stride);
 
-    coopMatArray[0].store(outputBuffer, 0, stride, matrixLayout);
-    coopMatArray[1].store(outputBuffer, 4, stride, matrixLayout);
+    coopMatArray[0].Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, stride);
+    coopMatArray[1].Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 4, stride);
 }
diff --git a/tests/cooperative-matrix/comparison.slang b/tests/cooperative-matrix/comparison.slang
index ce99f1550..8b29876df 100644
--- a/tests/cooperative-matrix/comparison.slang
+++ b/tests/cooperative-matrix/comparison.slang
@@ -14,7 +14,9 @@ ByteAddressBuffer input2;
 //TEST_INPUT:ubuffer(data=[0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<uint> outputBuffer;
 
-typealias CoopMatType = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain(uint3 threadIndex : SV_DispatchThreadID)
@@ -22,8 +24,8 @@ void computeMain(uint3 threadIndex : SV_DispatchThreadID)
     let stride = 4;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat1 = CoopMatType.load(input1, 0, stride, matrixLayout);
-    let mat2 = CoopMatType.load(input2, 0, stride, matrixLayout);
+    let mat1 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input1, 0, stride);
+    let mat2 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input2, 0, stride);
 
     uint32_t equals = mat1 == mat2 ? 1 : 0;
     uint32_t lessThan = mat1 < mat2 ? 1 : 0;
diff --git a/tests/cooperative-matrix/conversion.slang b/tests/cooperative-matrix/conversion.slang
index fbc422b7e..24fba3cc6 100644
--- a/tests/cooperative-matrix/conversion.slang
+++ b/tests/cooperative-matrix/conversion.slang
@@ -12,6 +12,7 @@ RWStructuredBuffer<float> outputBuffer;
 //TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4, count=256),name=input
 ByteAddressBuffer input;
 
+using namespace linalg;
 
 [numthreads(32, 1, 1)]
 void computeMain()
@@ -19,12 +20,12 @@ void computeMain()
     let stride = 16;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let intMat = CoopMat<int, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>.load(input, 0, stride, matrixLayout);
-    let floatMat = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(intMat);
-    let uintMat = CoopMat<uint, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(intMat);
-    let halfMat = CoopMat<half, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(uintMat);
-    let floatMat2 = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(halfMat);
+    let intMat = CoopMat<int, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>.Load<CoopMatMatrixLayout::RowMajor>(input, 0, stride);
+    let floatMat = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(intMat);
+    let uintMat = CoopMat<uint, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(intMat);
+    let halfMat = CoopMat<half, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(uintMat);
+    let floatMat2 = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(halfMat);
 
     let result = floatMat + floatMat2;
-    result.store(outputBuffer, 0, stride, matrixLayout);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, stride);
 }
diff --git a/tests/cooperative-matrix/copyFrom.slang b/tests/cooperative-matrix/copyFrom.slang
index 4be537489..9ab8acd4a 100644
--- a/tests/cooperative-matrix/copyFrom.slang
+++ b/tests/cooperative-matrix/copyFrom.slang
@@ -6,12 +6,14 @@
 //TEST_INPUT:ubuffer(stride=4, count = 256):out,name=outputBuffer
 RWStructuredBuffer<int32_t> outputBuffer;
 
+using namespace linalg;
+
 [numthreads(32, 1, 1)]
 void computeMain()
 {
-    let mat = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(4.0);
-    var result = CoopMat<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(0);
+    let mat = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(4.0);
+    var result = CoopMat<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(0);
     result.copyFrom(mat);
-    result.store(outputBuffer, 0, 16, CoopMatMatrixLayout::RowMajor);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, 16);
 }
 
diff --git a/tests/cooperative-matrix/diagnostics/mat-mul-add-different-scope.slang b/tests/cooperative-matrix/diagnostics/mat-mul-add-different-scope.slang
deleted file mode 100644
index 0c4308308..000000000
--- a/tests/cooperative-matrix/diagnostics/mat-mul-add-different-scope.slang
+++ /dev/null
@@ -1,20 +0,0 @@
-//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
deleted file mode 100644
index 5b7dc7a5b..000000000
--- a/tests/cooperative-matrix/diagnostics/mat-mul-add-incorrect-matrix-use.slang
+++ /dev/null
@@ -1,20 +0,0 @@
-//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);
-}
diff --git a/tests/cooperative-matrix/div.slang b/tests/cooperative-matrix/div.slang
index 4f697fb4e..17b237280 100644
--- a/tests/cooperative-matrix/div.slang
+++ b/tests/cooperative-matrix/div.slang
@@ -15,7 +15,9 @@ ByteAddressBuffer input1;
 //TEST_INPUT:ubuffer(data=[2 3 4 5], stride=4, count=256),name=input2
 ByteAddressBuffer input2;
 
-typealias CoopMatType = CoopMat<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
@@ -23,9 +25,9 @@ void computeMain()
     let stride = 16;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat1 = CoopMatType.load(input1, 0, stride, matrixLayout);
-    let mat2 = CoopMatType.load(input2, 0, stride, matrixLayout);
+    let mat1 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input1, 0, stride);
+    let mat2 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input2, 0, stride);
     let result = mat1 / mat2;
 
-    result.store(outputBuffer, 0, stride, matrixLayout);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, stride);
 }
diff --git a/tests/cooperative-matrix/fill.slang b/tests/cooperative-matrix/fill.slang
index e5c7c4765..f9068577b 100644
--- a/tests/cooperative-matrix/fill.slang
+++ b/tests/cooperative-matrix/fill.slang
@@ -6,11 +6,13 @@
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<int32_t> outputBuffer;
 
+using namespace linalg;
+
 [numthreads(32, 1, 1)]
 void computeMain()
 {
-    var result : CoopMat<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+    var result : CoopMat<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
     result.fill(10);
-    result.store(outputBuffer, 0, 16, CoopMatMatrixLayout::RowMajor);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, 16);
 }
 
diff --git a/tests/cooperative-matrix/inout.slang b/tests/cooperative-matrix/inout.slang
index 8e04a9f7a..3ff9a16d6 100644
--- a/tests/cooperative-matrix/inout.slang
+++ b/tests/cooperative-matrix/inout.slang
@@ -12,7 +12,9 @@ ByteAddressBuffer input;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
-typealias CoopMatType = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 void doubleCoopMat(inout CoopMatType mat)
 {
@@ -25,7 +27,7 @@ void computeMain()
     let stride = 16;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    var mat = CoopMatType.load(input, 0, stride, matrixLayout);
+    var mat = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input, 0, stride);
     doubleCoopMat(mat);
-    mat.store(outputBuffer, 0, stride, matrixLayout);
+    mat.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, stride);
 }
diff --git a/tests/cooperative-matrix/length.slang b/tests/cooperative-matrix/length.slang
new file mode 100644
index 000000000..580b713f3
--- /dev/null
+++ b/tests/cooperative-matrix/length.slang
@@ -0,0 +1,21 @@
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK):-vk -output-using-type -emit-spirv-directly -skip-spirv-validation
+
+// Note the length is NOT row * column.
+// When the memory scope is set to subgroup, each thread gets 16 * 16 / 32 = 8 where 32 is the value used in `numthreads`.
+
+//CHK:8
+
+//TEST_INPUT:ubuffer(stride=4, count=1):out,name=outputBuffer
+RWStructuredBuffer<int32_t> outputBuffer;
+
+using namespace linalg;
+
+// It appears that only Subgroup is supposed at the moment
+typealias CoopMatSubgroup    = CoopMat<int32_t, MemoryScope.Subgroup,    16, 16, CoopMatMatrixUse.MatrixAccumulator>;
+
+[numthreads(32, 1, 1)]
+void computeMain()
+{
+    outputBuffer[0] = CoopMatSubgroup.GetLength();
+}
+
diff --git a/tests/cooperative-matrix/load-store-arbitrary-array-vec.slang b/tests/cooperative-matrix/load-store-arbitrary-array-vec.slang
index 51af16ada..d34494fd8 100644
--- a/tests/cooperative-matrix/load-store-arbitrary-array-vec.slang
+++ b/tests/cooperative-matrix/load-store-arbitrary-array-vec.slang
@@ -23,7 +23,9 @@ RWByteAddressBuffer input;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<uint32_t> outputBuffer;
 
-typealias CoopMatType = CoopMat<uint32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<uint32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator>;
 
 groupshared float3[128] tempShared;
 
@@ -31,11 +33,9 @@ groupshared float3[128] tempShared;
 void computeMain()
 {
     let stride = 16;
-    let matrixLayout = CoopMatMatrixLayout::RowMajor;
-
-    let mat = coopMatLoad<uint32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(input, 0, stride, matrixLayout);
-    mat.storeAny(tempShared, 0, stride, matrixLayout);
+    let mat = coopMatLoad<uint32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator, CoopMatMatrixLayout.RowMajor>(input, 0, stride);
+    mat.Store<CoopMatMatrixLayout.RowMajor>(tempShared, 0, stride);
 
-    let result = CoopMatType.loadAny(tempShared, 0, stride, matrixLayout);
-    result.store(outputBuffer, 0, stride, matrixLayout);
+    let result = CoopMatType.Load<CoopMatMatrixLayout.RowMajor>(tempShared, 0, stride);
+    result.Store<CoopMatMatrixLayout.RowMajor>(outputBuffer, 0, stride);
 }
diff --git a/tests/cooperative-matrix/load-store-arbitrary-array.slang b/tests/cooperative-matrix/load-store-arbitrary-array.slang
index 160aad9da..8de4e84cc 100644
--- a/tests/cooperative-matrix/load-store-arbitrary-array.slang
+++ b/tests/cooperative-matrix/load-store-arbitrary-array.slang
@@ -23,7 +23,9 @@ RWByteAddressBuffer input;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<uint32_t> outputBuffer;
 
-typealias CoopMatType = CoopMat<uint32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<uint32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator>;
 
 groupshared float[256] tempShared;
 
@@ -31,11 +33,10 @@ groupshared float[256] tempShared;
 void computeMain()
 {
     let stride = 16;
-    let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat = coopMatLoad<uint32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(input, 0, stride, matrixLayout);
-    mat.storeAny(tempShared, 0, stride, matrixLayout);
+    let mat = coopMatLoad<uint32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator, CoopMatMatrixLayout.RowMajor>(input, 0, stride);
+    mat.Store<CoopMatMatrixLayout.RowMajor>(tempShared, 0, stride);
 
-    let result = CoopMatType.loadAny(tempShared, 0, stride, matrixLayout);
-    result.store(outputBuffer, 0, stride, matrixLayout);
+    let result = CoopMatType.Load<CoopMatMatrixLayout.RowMajor>(tempShared, 0, stride);
+    result.Store<CoopMatMatrixLayout.RowMajor>(outputBuffer, 0, stride);
 }
diff --git a/tests/cooperative-matrix/load-store-groupshared.slang b/tests/cooperative-matrix/load-store-groupshared.slang
index 8d867abb7..db32e85fd 100644
--- a/tests/cooperative-matrix/load-store-groupshared.slang
+++ b/tests/cooperative-matrix/load-store-groupshared.slang
@@ -10,22 +10,23 @@
 // CHECK-NEXT: 8
 
 //TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8], stride=4, count=256):name=input
-RWByteAddressBuffer input;
+ByteAddressBuffer input;
 
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<uint32_t> outputBuffer;
 
+using namespace linalg;
+
 groupshared uint32_t[256] tempShared;
 
 [numthreads(32, 1, 1)]
 void computeMain()
 {
     let stride = 16;
-    let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat = coopMatLoad<uint32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(input, 0, stride, matrixLayout);
-    mat.store(tempShared, 0, stride, matrixLayout);
+    let mat = coopMatLoad<uint32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator, CoopMatMatrixLayout.RowMajor>(input, 0, stride);
+    mat.Store<CoopMatMatrixLayout.RowMajor>(tempShared, 0, stride);
 
-    let result = coopMatLoad<uint32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(tempShared, 0, stride, matrixLayout);
-    result.store(outputBuffer, 0, stride, matrixLayout);
+    let result = coopMatLoad<uint32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator, CoopMatMatrixLayout.RowMajor>(tempShared, 0, stride);
+    result.Store<CoopMatMatrixLayout.RowMajor>(outputBuffer, 0, stride);
 }
diff --git a/tests/cooperative-matrix/load-store-rwbyteaddressbuffer.slang b/tests/cooperative-matrix/load-store-rwbyteaddressbuffer.slang
index 6bba8331e..6894bdfe5 100644
--- a/tests/cooperative-matrix/load-store-rwbyteaddressbuffer.slang
+++ b/tests/cooperative-matrix/load-store-rwbyteaddressbuffer.slang
@@ -1,4 +1,5 @@
 //TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -skip-spirv-validation
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -skip-spirv-validation -Xslang -DRWBAB
 
 // CHECK: 1
 // CHECK-NEXT: 2
@@ -10,17 +11,21 @@
 // CHECK-NEXT: 8
 
 //TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8], stride=4, count=256):name=inputBuffer
+#if defined(RWBAB)
 RWByteAddressBuffer inputBuffer;
+#else
+ByteAddressBuffer inputBuffer;
+#endif
 
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWByteAddressBuffer outputBuffer;
 
+using namespace linalg;
+
 [numthreads(32, 1, 1)]
 void computeMain()
 {
     let stride = 16;
-    let matrixLayout = CoopMatMatrixLayout::RowMajor;
-
-    let mat = coopMatLoad<uint32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(inputBuffer, 0, stride, matrixLayout);
-    mat.store(outputBuffer, 0, 16, matrixLayout);
+    let mat = coopMatLoad<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator, CoopMatMatrixLayout.RowMajor>(inputBuffer, 0, stride);
+    mat.Store<CoopMatMatrixLayout.RowMajor>(outputBuffer, 0, 16);
 }
diff --git a/tests/cooperative-matrix/load-store-rwstructuredbuffer.slang b/tests/cooperative-matrix/load-store-rwstructuredbuffer.slang
index e161fb7b2..6a94fd30e 100644
--- a/tests/cooperative-matrix/load-store-rwstructuredbuffer.slang
+++ b/tests/cooperative-matrix/load-store-rwstructuredbuffer.slang
@@ -1,4 +1,5 @@
 //TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -skip-spirv-validation
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -skip-spirv-validation -Xslang -DRWSB
 
 // CHECK: type: int32_t
 // CHECK-NEXT: 1
@@ -11,17 +12,21 @@
 // CHECK-NEXT: 8
 
 //TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8], stride=4, count=256),name=buf
+#if defined(RWSB)
 RWStructuredBuffer<int32_t> inputBuffer;
+#else
+StructuredBuffer<int32_t> inputBuffer;
+#endif
 
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<int32_t> outputBuffer;
 
+using namespace linalg;
+
 [numthreads(32, 1, 1)]
 void computeMain()
 {
     let stride = 16;
-    let matrixLayout = CoopMatMatrixLayout::RowMajor;
-
-    let mat = coopMatLoad<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>(inputBuffer, 0, stride, matrixLayout);
-    mat.store(outputBuffer, 0, stride, matrixLayout);
+    let mat = coopMatLoad<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator, CoopMatMatrixLayout.RowMajor>(inputBuffer, 0, stride);
+    mat.Store<CoopMatMatrixLayout.RowMajor>(outputBuffer, 0, stride);
 }
diff --git a/tests/cooperative-matrix/map-element-single.slang b/tests/cooperative-matrix/map-element-single.slang
new file mode 100644
index 000000000..583630d14
--- /dev/null
+++ b/tests/cooperative-matrix/map-element-single.slang
@@ -0,0 +1,50 @@
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -render-feature cooperative-matrix-per-element-operations -skip-spirv-validation -Xslang -DTEST_MODE=0
+//DISABLE_TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -render-feature cooperative-matrix-per-element-operations -skip-spirv-validation -Xslang -DTEST_MODE=1
+//DISABLE_TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -render-feature cooperative-matrix-per-element-operations -skip-spirv-validation -Xslang -DTEST_MODE=2
+
+//CHECK: type: int32_t
+//CHECK-NEXT: 2
+//CHECK-NEXT: 4
+//CHECK-NEXT: 6
+//CHECK-NEXT: 8
+
+//TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4),name=input1
+StructuredBuffer<int> input1;
+
+//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
+RWStructuredBuffer<int32_t> outputBuffer;
+
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<int, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator>;
+
+int MapOp(uint32_t row, uint32_t col, int value)
+{
+    return value * 2;
+}
+
+[numthreads(32, 1, 1)]
+void computeMain()
+{
+    let stride = 16;
+    CoopMatType mat1 = CoopMatType.Load<CoopMatMatrixLayout.RowMajor>(input1, 0, stride);
+
+    CoopMatType result;
+
+#if TEST_MODE == 0
+    result = mat1.MapElement(MapOp);
+
+#elif TEST_MODE == 1
+    // Lambda through IFunc.
+    // TODO: Not working due to issue #7024
+    IFunc<int, uint32_t, uint32_t, int> func = ((uint32_t row, uint32_t column, int value) => value * 2);
+    result = mat1.MapElement(func);
+
+#elif TEST_MODE == 2
+    // Directly use lambda.
+    // TODO: Not working due to issue #7024
+    result = mat1.MapElement((uint32_t row, uint32_t column, int value) => (int)(value));
+#endif
+
+    result.Store<CoopMatMatrixLayout.RowMajor>(outputBuffer, 0, stride);
+}
diff --git a/tests/cooperative-matrix/mat-mul-add-spirv-matrix-operands.slang b/tests/cooperative-matrix/mat-mul-add-spirv-matrix-operands.slang
index 7ad6c639f..a93efb4b9 100644
--- a/tests/cooperative-matrix/mat-mul-add-spirv-matrix-operands.slang
+++ b/tests/cooperative-matrix/mat-mul-add-spirv-matrix-operands.slang
@@ -1,45 +1,68 @@
 //TEST(compute):SIMPLE(filecheck=CHECK): -entry computeMain -stage compute -target spirv -skip-spirv-validation
 
 // This test checks that the correct SPIRV Cooperative Matrix Operands are emitted for OpCooperativeMatrixMulAddKHR operaions
-RWStructuredBuffer<int> outputBuffer1;
-RWStructuredBuffer<int> outputBuffer2;
-RWStructuredBuffer<int> outputBuffer3;
+RWStructuredBuffer<uint32_t> outputBuffer1;
+RWStructuredBuffer<uint32_t> outputBuffer2;
+RWStructuredBuffer<uint32_t> outputBuffer3;
 RWStructuredBuffer<int> outputBuffer4;
-RWStructuredBuffer<int> outputBuffer5;
+RWStructuredBuffer<uint32_t> outputBuffer5;
+RWStructuredBuffer<int> outputBuffer6;
 
-typealias CoopMatAType = CoopMat<int16_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixA>;
-typealias CoopMatBType = CoopMat<int16_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixB>;
-typealias CoopMatCType = CoopMat<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+__generic<T : __BuiltinArithmeticType>
+typealias CoopMatAType = CoopMat<T, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixA>;
+
+__generic<T : __BuiltinArithmeticType>
+typealias CoopMatBType = CoopMat<T, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixB>;
+
+__generic<T : __BuiltinArithmeticType>
+typealias CoopMatCType = CoopMat<T, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
 {
-    let matA = CoopMatAType(3);
-    let matB = CoopMatBType(5);
-    let matC = CoopMatCType(1);
-
-    let matrixLayout = CoopMatMatrixLayout::RowMajor;
-
     // CHECK: OpCooperativeMatrixMulAddKHR {{.*}} NoneKHR
-    coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::None).store(outputBuffer1, 0, 16, matrixLayout);
+    coopMatMulAdd<uint32_t, false>(
+        CoopMatAType<uint16_t>(2),
+        CoopMatBType<uint32_t>(3),
+        CoopMatCType<uint16_t>(4)
+    ).Store<CoopMatMatrixLayout::RowMajor>(outputBuffer1, 0, 16);
 
     // CHECK: OpCooperativeMatrixMulAddKHR {{.*}} MatrixASignedComponentsKHR|MatrixBSignedComponentsKHR
-    coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::MatrixASigned | CoopMatMatrixOperands::MatrixBSigned).store(outputBuffer2, 0, 16, matrixLayout);
-
+    coopMatMulAdd<uint32_t, false>(
+        CoopMatAType<int16_t>(2),
+        CoopMatBType<int32_t>(3),
+        CoopMatCType<uint16_t>(4)
+    ).Store<CoopMatMatrixLayout::RowMajor>(outputBuffer2, 0, 16);
 
     // CHECK: OpCooperativeMatrixMulAddKHR {{.*}} MatrixCSignedComponentsKHR
-    coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::MatrixCSigned).store(outputBuffer2, 0, 16, matrixLayout);
-
+    coopMatMulAdd<uint32_t, false>(
+        CoopMatAType<uint16_t>(2),
+        CoopMatBType<uint32_t>(3),
+        CoopMatCType<int16_t>(4)
+    ).Store<CoopMatMatrixLayout::RowMajor>(outputBuffer3, 0, 16);
 
     // CHECK: OpCooperativeMatrixMulAddKHR {{.*}} MatrixResultSignedComponentsKHR
-    coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::MatrixResultSigned).store(outputBuffer3, 0, 16, matrixLayout);
+    coopMatMulAdd<int, false>(
+        CoopMatAType<uint16_t>(2),
+        CoopMatBType<uint32_t>(3),
+        CoopMatCType<uint16_t>(4)
+    ).Store<CoopMatMatrixLayout::RowMajor>(outputBuffer4, 0, 16);
 
     // CHECK: OpCooperativeMatrixMulAddKHR {{.*}} SaturatingAccumulationKHR
-    coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::SaturatingAccumulation).store(outputBuffer4, 0, 16, matrixLayout);
+    coopMatMulAdd<uint32_t, true>(
+        CoopMatAType<uint16_t>(2),
+        CoopMatBType<uint32_t>(3),
+        CoopMatCType<uint16_t>(4)
+    ).Store<CoopMatMatrixLayout::RowMajor>(outputBuffer5, 0, 16);
 
-    let allOperands = CoopMatMatrixOperands::MatrixASigned | CoopMatMatrixOperands::MatrixBSigned | CoopMatMatrixOperands::MatrixCSigned | CoopMatMatrixOperands::MatrixResultSigned | CoopMatMatrixOperands::SaturatingAccumulation;
     // CHECK: OpCooperativeMatrixMulAddKHR {{.*}} MatrixASignedComponentsKHR|MatrixBSignedComponentsKHR|MatrixCSignedComponentsKHR|MatrixResultSignedComponentsKHR|SaturatingAccumulationKHR
-    coopMatMulAdd(matA, matB, matC, allOperands).store(outputBuffer5, 0, 16, matrixLayout);
+    coopMatMulAdd<int, true>(
+        CoopMatAType<int16_t>(2),
+        CoopMatBType<int32_t>(3),
+        CoopMatCType<int16_t>(4)
+    ).Store<CoopMatMatrixLayout::RowMajor>(outputBuffer6, 0, 16);
 }
 
 
diff --git a/tests/cooperative-matrix/mat-mul-add.slang b/tests/cooperative-matrix/mat-mul-add.slang
index 417f06d10..3dc472de7 100644
--- a/tests/cooperative-matrix/mat-mul-add.slang
+++ b/tests/cooperative-matrix/mat-mul-add.slang
@@ -6,9 +6,11 @@
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
-typealias CoopMatAType = CoopMat<float16_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixA>;
-typealias CoopMatBType = CoopMat<float16_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixB>;
-typealias CoopMatCType = CoopMat<float32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatAType = CoopMat<float16_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixA>;
+typealias CoopMatBType = CoopMat<float16_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixB>;
+typealias CoopMatCType = CoopMat<float32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
@@ -18,6 +20,6 @@ void computeMain()
     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);
+    let result = coopMatMulAdd<float, false>(matA, matB, matC);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, 16);
 }
diff --git a/tests/cooperative-matrix/mod.slang b/tests/cooperative-matrix/mod.slang
index 5167c8ef8..015dd227a 100644
--- a/tests/cooperative-matrix/mod.slang
+++ b/tests/cooperative-matrix/mod.slang
@@ -16,7 +16,7 @@
 // CHECK-NEXT: 2
 
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
-RWByteAddressBuffer outputBuffer;
+RWStructuredBuffer<int> outputBuffer;
 
 //TEST_INPUT:ubuffer(data=[4 3 5 7], stride=4, count=256),name=input1
 ByteAddressBuffer input1;
@@ -24,18 +24,19 @@ ByteAddressBuffer input1;
 //TEST_INPUT:ubuffer(data=[2 3 4 5], stride=4, count=256),name=input2
 ByteAddressBuffer input2;
 
-typealias CoopMatIntType = CoopMat<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
-typealias CoopMatUintType = CoopMat<uint32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
-typealias CoopMatFloatType = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatIntType = CoopMat<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+typealias CoopMatUintType = CoopMat<uint32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+typealias CoopMatFloatType = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
 {
     let stride = 16;
-    let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat1 = CoopMatIntType.load(input1, 0, stride, matrixLayout);
-    let mat2 = CoopMatIntType.load(input2, 0, stride, matrixLayout);
+    let mat1 = CoopMatIntType.Load<CoopMatMatrixLayout::RowMajor>(input1, 0, stride);
+    let mat2 = CoopMatIntType.Load<CoopMatMatrixLayout::RowMajor>(input2, 0, stride);
 
     let mat3 = CoopMatFloatType(mat1);
     let mat4 = CoopMatFloatType(mat2);
@@ -47,7 +48,7 @@ void computeMain()
     let result2 = CoopMatIntType(mat3 % mat4);
     let result3 = CoopMatIntType(mat5 % mat6);
 
-    result.store(outputBuffer, 0, stride, matrixLayout);
-    result2.store(outputBuffer, 16, stride, matrixLayout);
-    result3.store(outputBuffer, 32, stride, matrixLayout);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, stride);
+    result2.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 16, stride);
+    result3.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 32, stride);
 }
diff --git a/tests/cooperative-matrix/mul.slang b/tests/cooperative-matrix/mul.slang
index 0ac332698..4f5751c6e 100644
--- a/tests/cooperative-matrix/mul.slang
+++ b/tests/cooperative-matrix/mul.slang
@@ -15,7 +15,9 @@ ByteAddressBuffer input1;
 //TEST_INPUT:ubuffer(data=[2 3 4 5], stride=4, count=256),name=input2
 ByteAddressBuffer input2;
 
-typealias CoopMatType = CoopMat<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
@@ -23,9 +25,9 @@ void computeMain()
     let stride = 16;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat1 = CoopMatType.load(input1, 0, stride, matrixLayout);
-    let mat2 = CoopMatType.load(input2, 0, stride, matrixLayout);
+    let mat1 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input1, 0, stride);
+    let mat2 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input2, 0, stride);
 
     let result = mat1 * mat2;
-    result.store(outputBuffer, 0, 4, matrixLayout);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, 4);
 }
diff --git a/tests/cooperative-matrix/out.slang b/tests/cooperative-matrix/out.slang
index 5b342afc0..7d4a83fc2 100644
--- a/tests/cooperative-matrix/out.slang
+++ b/tests/cooperative-matrix/out.slang
@@ -13,7 +13,9 @@ StructuredBuffer<float> inputBuffer;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
-typealias CoopMatType = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 void doubleCoopMat(CoopMatType mat, out CoopMatType result)
 {
@@ -26,9 +28,9 @@ void computeMain()
     let stride = 16;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout);
+    let mat = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(inputBuffer, 0, stride);
 
     CoopMatType result;
     doubleCoopMat(mat, result);
-    result.store(outputBuffer, 0, stride, matrixLayout);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, stride);
 }
diff --git a/tests/cooperative-matrix/parameter.slang b/tests/cooperative-matrix/parameter.slang
index 8a4bb3315..eb6823aa1 100644
--- a/tests/cooperative-matrix/parameter.slang
+++ b/tests/cooperative-matrix/parameter.slang
@@ -12,7 +12,9 @@ StructuredBuffer<float> inputBuffer;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
-typealias CoopMatType = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 static let stride = 16;
 //static let matrixLayout = CoopMatMatrixLayout::RowMajor;
@@ -21,12 +23,12 @@ static const CoopMatMatrixLayout matrixLayout = CoopMatMatrixLayout::RowMajor;
 void processCoopMat(CoopMatType mat)
 {
     // XXX: hmmm, some error when matrixLAyout is static let
-    (mat * 3.0).store(outputBuffer, 0, stride, matrixLayout);
+    (mat * 3.0).Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, stride);
 }
 
 [numthreads(32, 1, 1)]
 void computeMain()
 {
-    let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout);
+    let mat = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(inputBuffer, 0, stride);
     processCoopMat(mat);
 }
diff --git a/tests/cooperative-matrix/reduce.slang b/tests/cooperative-matrix/reduce.slang
new file mode 100644
index 000000000..bbba587ce
--- /dev/null
+++ b/tests/cooperative-matrix/reduce.slang
@@ -0,0 +1,52 @@
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK_ROW):-vk -output-using-type -emit-spirv-directly -render-feature cooperative-matrix-reductions -skip-spirv-validation -Xslang -DTEST_MODE=0
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK_COLUMN):-vk -output-using-type -emit-spirv-directly -render-feature cooperative-matrix-reductions -skip-spirv-validation -Xslang -DTEST_MODE=1
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK_ROW_AND_COLUMN):-vk -output-using-type -emit-spirv-directly -render-feature cooperative-matrix-reductions -skip-spirv-validation -Xslang -DTEST_MODE=2
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHK_2X2):-vk -output-using-type -emit-spirv-directly -render-feature cooperative-matrix-reductions -skip-spirv-validation -Xslang -DTEST_MODE=3
+
+//CHK_ROW-COUNT-8: 36
+
+//CHK_COLUMN:1
+//CHK_COLUMN-NEXT:2
+//CHK_COLUMN-NEXT:3
+//CHK_COLUMN-NEXT:4
+//CHK_COLUMN-NEXT:5
+//CHK_COLUMN-NEXT:6
+//CHK_COLUMN-NEXT:7
+//CHK_COLUMN-NEXT:8
+
+//CHK_ROW_AND_COLUMN: 36
+
+//CHK_2X2:3
+//CHK_2X2:7
+//CHK_2X2:11
+//CHK_2X2:15
+
+//TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8], stride=4, count=256),name=buf
+StructuredBuffer<int32_t> inputBuffer;
+
+//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
+RWStructuredBuffer<int32_t> outputBuffer;
+
+using namespace linalg;
+
+int32_t CombineOp(int32_t lhs, int32_t rhs)
+{
+    return lhs + rhs;
+}
+
+[numthreads(32, 1, 1)]
+void computeMain()
+{
+    let stride = 16;
+    let mat = coopMatLoad<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator, CoopMatMatrixLayout.RowMajor>(inputBuffer, 0, stride);
+#if TEST_MODE == 0
+    let result = mat.ReduceRow<8>(CombineOp);
+#elif TEST_MODE == 1
+    let result = mat.ReduceColumn<8>(CombineOp);
+#elif TEST_MODE == 2
+    let result = mat.ReduceRowAndColumn<8, 8>(CombineOp);
+#elif TEST_MODE == 3
+    let result = mat.Reduce2x2(CombineOp);
+#endif
+    result.Store<CoopMatMatrixLayout.RowMajor>(outputBuffer, 0, stride);
+}
diff --git a/tests/cooperative-matrix/return.slang b/tests/cooperative-matrix/return.slang
index 722a31b8b..b67117892 100644
--- a/tests/cooperative-matrix/return.slang
+++ b/tests/cooperative-matrix/return.slang
@@ -12,7 +12,9 @@ StructuredBuffer<float> inputBuffer;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
-typealias CoopMatType = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 CoopMatType doubleCoopmat(CoopMatType mat)
 {
@@ -25,8 +27,8 @@ void computeMain()
     let stride = 16;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout);
+    let mat = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(inputBuffer, 0, stride);
 
     let result = doubleCoopmat(mat);
-    result.store(outputBuffer, 0, 4, matrixLayout);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, 4);
 }
diff --git a/tests/cooperative-matrix/scalar-mul.slang b/tests/cooperative-matrix/scalar-mul.slang
index 9d266920e..29d8f03f0 100644
--- a/tests/cooperative-matrix/scalar-mul.slang
+++ b/tests/cooperative-matrix/scalar-mul.slang
@@ -12,7 +12,9 @@ ByteAddressBuffer inputBuffer;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
-typealias CoopMatType = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
@@ -20,8 +22,8 @@ void computeMain()
     let stride = 16;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout);
+    let mat = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(inputBuffer, 0, stride);
 
     let result = mat * 4.5;
-    result.store(outputBuffer, 0, 4, matrixLayout);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, 4);
 }
diff --git a/tests/cooperative-matrix/struct.slang b/tests/cooperative-matrix/struct.slang
index 38d98b44f..592e553ae 100644
--- a/tests/cooperative-matrix/struct.slang
+++ b/tests/cooperative-matrix/struct.slang
@@ -19,7 +19,9 @@ ByteAddressBuffer input2;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
-typealias CoopMatType = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 struct MyStruct
 {
@@ -34,9 +36,9 @@ void computeMain()
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
     MyStruct s;
-    s.mat1 = CoopMatType.load(input1, 0, stride, matrixLayout);
-    s.mat2 = CoopMatType.load(input2, 0, stride, matrixLayout);
+    s.mat1 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input1, 0, stride);
+    s.mat2 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input2, 0, stride);
 
-    s.mat1.store(outputBuffer, 0, stride, matrixLayout);
-    s.mat2.store(outputBuffer, 4, stride, matrixLayout);
+    s.mat1.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, stride);
+    s.mat2.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 4, stride);
 }
diff --git a/tests/cooperative-matrix/sub.slang b/tests/cooperative-matrix/sub.slang
index 7fe6f3fea..a42fee9e8 100644
--- a/tests/cooperative-matrix/sub.slang
+++ b/tests/cooperative-matrix/sub.slang
@@ -15,7 +15,9 @@ ByteAddressBuffer input1;
 //TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4, count=256),name=input2
 ByteAddressBuffer input2;
 
-typealias CoopMatType = CoopMat<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
@@ -23,9 +25,9 @@ void computeMain()
     let stride = 16;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat1 = CoopMatType.load(input1, 0, stride, matrixLayout);
-    let mat2 = CoopMatType.load(input2, 0, stride, matrixLayout);
+    let mat1 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input1, 0, stride);
+    let mat2 = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(input2, 0, stride);
 
     let result = mat1 - mat2;
-    result.store(outputBuffer, 0, 4, matrixLayout);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, 4);
 }
diff --git a/tests/cooperative-matrix/subscript-in-func.slang b/tests/cooperative-matrix/subscript-in-func.slang
index 585eabf92..1bb8df433 100644
--- a/tests/cooperative-matrix/subscript-in-func.slang
+++ b/tests/cooperative-matrix/subscript-in-func.slang
@@ -12,7 +12,9 @@ StructuredBuffer<float> inputBuffer;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
-typealias CoopMatType = CoopMat<float, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<float, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 static const int stride = 16;
 static const CoopMatMatrixLayout matrixLayout = CoopMatMatrixLayout::RowMajor;
@@ -23,12 +25,12 @@ void squareCoopMatElements(CoopMatType mat)
     {
         mat[i] = mat[i] * mat[i];
     }
-    mat.store(outputBuffer, 0, stride, matrixLayout);
+    mat.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, stride);
 }
 
 [numthreads(32, 1, 1)]
 void computeMain()
 {
-    let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout);
+    let mat = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(inputBuffer, 0, stride);
     squareCoopMatElements(mat);
 }
diff --git a/tests/cooperative-matrix/subscript.slang b/tests/cooperative-matrix/subscript.slang
index cfe164f04..0d765b6b3 100644
--- a/tests/cooperative-matrix/subscript.slang
+++ b/tests/cooperative-matrix/subscript.slang
@@ -9,7 +9,9 @@
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<int32_t> outputBuffer;
 
-typealias CoopMatType = CoopMat<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
@@ -19,5 +21,5 @@ void computeMain()
     mat[1] = mat[0]+2;
     mat[2] = mat[1]+3;
     mat[3] = mat[2]+4;
-    mat.store(outputBuffer, 0, 16, CoopMatMatrixLayout::RowMajor);
+    mat.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, 16);
 }
diff --git a/tests/cooperative-matrix/transpose.slang b/tests/cooperative-matrix/transpose.slang
new file mode 100644
index 000000000..cd317cac2
--- /dev/null
+++ b/tests/cooperative-matrix/transpose.slang
@@ -0,0 +1,36 @@
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -skip-spirv-validation -render-feature cooperative-matrix-conversions
+
+//CHECK: type: int32_t
+//CHECK: 1
+//CHECK-COUNT-15: 0
+//CHECK: 2
+//CHECK-COUNT-15: 0
+//CHECK: 3
+//CHECK-COUNT-15: 0
+//CHECK: 4
+//CHECK-COUNT-15: 0
+//CHECK: 5
+//CHECK-COUNT-15: 0
+//CHECK: 6
+//CHECK-COUNT-15: 0
+//CHECK: 7
+//CHECK-COUNT-15: 0
+//CHECK: 8
+//CHECK-COUNT-15: 0
+
+//TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8], stride=4, count=256),name=buf
+StructuredBuffer<int32_t> inputBuffer;
+
+//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
+RWStructuredBuffer<int32_t> outputBuffer;
+
+using namespace linalg;
+
+[numthreads(32, 1, 1)]
+void computeMain()
+{
+    let stride = 16;
+    let mat = coopMatLoad<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator, CoopMatMatrixLayout.RowMajor>(inputBuffer, 0, stride);
+    let result = mat.Transpose();
+    result.Store<CoopMatMatrixLayout.RowMajor>(outputBuffer, 0, stride);
+}
diff --git a/tests/cooperative-matrix/unary_neg.slang b/tests/cooperative-matrix/unary_neg.slang
index 5d58f1395..be803d8c6 100644
--- a/tests/cooperative-matrix/unary_neg.slang
+++ b/tests/cooperative-matrix/unary_neg.slang
@@ -12,7 +12,9 @@ ByteAddressBuffer inputBuffer;
 //TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
 RWStructuredBuffer<int32_t> outputBuffer;
 
-typealias CoopMatType = CoopMat<int32_t, CoopMatScope::Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
+using namespace linalg;
+
+typealias CoopMatType = CoopMat<int32_t, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse::MatrixAccumulator>;
 
 [numthreads(32, 1, 1)]
 void computeMain()
@@ -20,8 +22,8 @@ void computeMain()
     let stride = 4;
     let matrixLayout = CoopMatMatrixLayout::RowMajor;
 
-    let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout);
+    let mat = CoopMatType.Load<CoopMatMatrixLayout::RowMajor>(inputBuffer, 0, stride);
 
     let result = -mat;
-    result.store(outputBuffer, 0, 4, matrixLayout);
+    result.Store<CoopMatMatrixLayout::RowMajor>(outputBuffer, 0, 4);
 }
-- 
cgit v1.2.3