From d0b6a0b1ab49b5958015f31364c5ad73d9cd03eb Mon Sep 17 00:00:00 2001 From: Darren Wihandi <65404740+fairywreath@users.noreply.github.com> Date: Tue, 15 Apr 2025 15:57:45 -0600 Subject: 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 --- tests/cooperative-matrix/add.slang | 32 +++++++++++++ tests/cooperative-matrix/array.slang | 36 +++++++++++++++ tests/cooperative-matrix/comparison.slang | 35 ++++++++++++++ tests/cooperative-matrix/conversion.slang | 30 ++++++++++++ tests/cooperative-matrix/copyFrom.slang | 17 +++++++ .../diagnostics/mat-mul-add-different-scope.slang | 20 ++++++++ .../mat-mul-add-incorrect-matrix-use.slang | 20 ++++++++ tests/cooperative-matrix/div.slang | 31 +++++++++++++ tests/cooperative-matrix/fill.slang | 16 +++++++ tests/cooperative-matrix/inout.slang | 31 +++++++++++++ .../load-store-arbitrary-array-vec.slang | 41 +++++++++++++++++ .../load-store-arbitrary-array.slang | 41 +++++++++++++++++ .../load-store-groupshared.slang | 31 +++++++++++++ .../load-store-rwbyteaddressbuffer.slang | 26 +++++++++++ .../load-store-rwstructuredbuffer.slang | 27 +++++++++++ .../mat-mul-add-spirv-matrix-operands.slang | 45 ++++++++++++++++++ tests/cooperative-matrix/mat-mul-add.slang | 23 ++++++++++ tests/cooperative-matrix/mod.slang | 53 ++++++++++++++++++++++ tests/cooperative-matrix/mul.slang | 31 +++++++++++++ tests/cooperative-matrix/out.slang | 34 ++++++++++++++ tests/cooperative-matrix/parameter.slang | 32 +++++++++++++ tests/cooperative-matrix/return.slang | 32 +++++++++++++ tests/cooperative-matrix/scalar-mul.slang | 27 +++++++++++ tests/cooperative-matrix/struct.slang | 42 +++++++++++++++++ tests/cooperative-matrix/sub.slang | 31 +++++++++++++ tests/cooperative-matrix/subscript-in-func.slang | 34 ++++++++++++++ tests/cooperative-matrix/subscript.slang | 23 ++++++++++ tests/cooperative-matrix/unary_neg.slang | 27 +++++++++++ 28 files changed, 868 insertions(+) create mode 100644 tests/cooperative-matrix/add.slang create mode 100644 tests/cooperative-matrix/array.slang create mode 100644 tests/cooperative-matrix/comparison.slang create mode 100644 tests/cooperative-matrix/conversion.slang create mode 100644 tests/cooperative-matrix/copyFrom.slang create mode 100644 tests/cooperative-matrix/diagnostics/mat-mul-add-different-scope.slang create mode 100644 tests/cooperative-matrix/diagnostics/mat-mul-add-incorrect-matrix-use.slang create mode 100644 tests/cooperative-matrix/div.slang create mode 100644 tests/cooperative-matrix/fill.slang create mode 100644 tests/cooperative-matrix/inout.slang create mode 100644 tests/cooperative-matrix/load-store-arbitrary-array-vec.slang create mode 100644 tests/cooperative-matrix/load-store-arbitrary-array.slang create mode 100644 tests/cooperative-matrix/load-store-groupshared.slang create mode 100644 tests/cooperative-matrix/load-store-rwbyteaddressbuffer.slang create mode 100644 tests/cooperative-matrix/load-store-rwstructuredbuffer.slang create mode 100644 tests/cooperative-matrix/mat-mul-add-spirv-matrix-operands.slang create mode 100644 tests/cooperative-matrix/mat-mul-add.slang create mode 100644 tests/cooperative-matrix/mod.slang create mode 100644 tests/cooperative-matrix/mul.slang create mode 100644 tests/cooperative-matrix/out.slang create mode 100644 tests/cooperative-matrix/parameter.slang create mode 100644 tests/cooperative-matrix/return.slang create mode 100644 tests/cooperative-matrix/scalar-mul.slang create mode 100644 tests/cooperative-matrix/struct.slang create mode 100644 tests/cooperative-matrix/sub.slang create mode 100644 tests/cooperative-matrix/subscript-in-func.slang create mode 100644 tests/cooperative-matrix/subscript.slang create mode 100644 tests/cooperative-matrix/unary_neg.slang (limited to 'tests') diff --git a/tests/cooperative-matrix/add.slang b/tests/cooperative-matrix/add.slang new file mode 100644 index 000000000..3d8348d13 --- /dev/null +++ b/tests/cooperative-matrix/add.slang @@ -0,0 +1,32 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: int32_t +// CHECK-NEXT: 1 +// CHECK-NEXT: 3 +// CHECK-NEXT: 5 +// CHECK-NEXT: 7 + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +//TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4, count=256),name=input1 +ByteAddressBuffer input1; + +//TEST_INPUT:ubuffer(data=[0 1 2 3], stride=4, count=256),name=input2 +ByteAddressBuffer input2; + +typealias CoopMatType = CoopMat; + +[numthreads(32, 1, 1)] +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 result = mat1 + mat2; + + result.store(outputBuffer, 0, stride, matrixLayout); +} + diff --git a/tests/cooperative-matrix/array.slang b/tests/cooperative-matrix/array.slang new file mode 100644 index 000000000..b46c0f66b --- /dev/null +++ b/tests/cooperative-matrix/array.slang @@ -0,0 +1,36 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: float +// CHECK: 1.000000 +// CHECK-NEXT: 2.000000 +// CHECK-NEXT: 3.000000 +// CHECK-NEXT: 4.000000 +// CHECK-NEXT: 5.000000 +// CHECK-NEXT: 6.000000 +// CHECK-NEXT: 7.000000 +// CHECK-NEXT: 8.000000 + +//TEST_INPUT:ubuffer(data=[1.0 2.0 3.0 4.0], stride=256),name=input1 +ByteAddressBuffer input1; + +//TEST_INPUT:ubuffer(data=[5.0 6.0 7.0 8.0], stride=256),name=input1 +ByteAddressBuffer input2; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + 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].store(outputBuffer, 0, stride, matrixLayout); + coopMatArray[1].store(outputBuffer, 4, stride, matrixLayout); +} diff --git a/tests/cooperative-matrix/comparison.slang b/tests/cooperative-matrix/comparison.slang new file mode 100644 index 000000000..bcf0c90ae --- /dev/null +++ b/tests/cooperative-matrix/comparison.slang @@ -0,0 +1,35 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: uint32_t +// CHECK-NEXT: 0 +// CHECK-NEXT: 1 +// CHECK-NEXT: 1 + +//TEST_INPUT:ubuffer(data=[1.0 2.0 3.0 4.0], stride=4, count=256),name=input1 +ByteAddressBuffer input1; + +//TEST_INPUT:ubuffer(data=[1.0 3.0 2.0 4.0], stride=4, count=256),name=input2 +ByteAddressBuffer input2; + +//TEST_INPUT:ubuffer(data=[0 0 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +[numthreads(32, 1, 1)] +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); + + uint32_t equals = mat1 == mat2 ? 1 : 0; + uint32_t lessThan = mat1 < mat2 ? 1 : 0; + uint32_t lessThanOrEquals = mat1 <= mat2 ? 1 : 0; + + outputBuffer[0] = equals; + outputBuffer[1] = lessThan; + outputBuffer[2] = lessThanOrEquals; +} diff --git a/tests/cooperative-matrix/conversion.slang b/tests/cooperative-matrix/conversion.slang new file mode 100644 index 000000000..745882ab8 --- /dev/null +++ b/tests/cooperative-matrix/conversion.slang @@ -0,0 +1,30 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: float +// CHECK-NEXT: 2.000000 +// CHECK-NEXT: 4.000000 +// CHECK-NEXT: 6.000000 +// CHECK-NEXT: 8.000000 + +//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +//TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4, count=256),name=input +ByteAddressBuffer input; + + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + let intMat = CoopMat.load(input, 0, stride, matrixLayout); + let floatMat = CoopMat(intMat); + let uintMat = CoopMat(intMat); + let halfMat = CoopMat(uintMat); + let floatMat2 = CoopMat(halfMat); + + let result = floatMat + floatMat2; + result.store(outputBuffer, 0, stride, matrixLayout); +} diff --git a/tests/cooperative-matrix/copyFrom.slang b/tests/cooperative-matrix/copyFrom.slang new file mode 100644 index 000000000..f7270545e --- /dev/null +++ b/tests/cooperative-matrix/copyFrom.slang @@ -0,0 +1,17 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: int32_t +// CHECK-COUNT-256: 4 + +//TEST_INPUT:ubuffer(stride=4, count = 256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +[numthreads(32, 1, 1)] +void computeMain() +{ + let mat = CoopMat(4.0); + var result = CoopMat(0); + result.copyFrom(mat); + result.store(outputBuffer, 0, 16, CoopMatMatrixLayout::RowMajor); +} + 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 outputBuffer; + +typealias CoopMatAType = CoopMat; +typealias CoopMatBType = CoopMat; +typealias CoopMatCType = CoopMat; + +// 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 outputBuffer; + +typealias CoopMatAType = CoopMat; +typealias CoopMatBType = CoopMat; +typealias CoopMatCType = CoopMat; + +// 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 new file mode 100644 index 000000000..29207e0e4 --- /dev/null +++ b/tests/cooperative-matrix/div.slang @@ -0,0 +1,31 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: int32_t +// CHECK-NEXT: 2 +// CHECK-NEXT: 1 +// CHECK-NEXT: 1 +// CHECK-NEXT: 0 + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +//TEST_INPUT:ubuffer(data=[4 3 5 2], stride=4, count=256),name=input1 +ByteAddressBuffer input1; + +//TEST_INPUT:ubuffer(data=[2 3 4 5], stride=4, count=256),name=input2 +ByteAddressBuffer input2; + +typealias CoopMatType = CoopMat; + +[numthreads(32, 1, 1)] +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 result = mat1 / mat2; + + result.store(outputBuffer, 0, stride, matrixLayout); +} diff --git a/tests/cooperative-matrix/fill.slang b/tests/cooperative-matrix/fill.slang new file mode 100644 index 000000000..d1a46d053 --- /dev/null +++ b/tests/cooperative-matrix/fill.slang @@ -0,0 +1,16 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: int32_t +// CHECK-COUNT-256: 10 + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +[numthreads(32, 1, 1)] +void computeMain() +{ + var result : CoopMat; + result.fill(10); + result.store(outputBuffer, 0, 16, CoopMatMatrixLayout::RowMajor); +} + diff --git a/tests/cooperative-matrix/inout.slang b/tests/cooperative-matrix/inout.slang new file mode 100644 index 000000000..7284953b4 --- /dev/null +++ b/tests/cooperative-matrix/inout.slang @@ -0,0 +1,31 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: float +// CHECK-NEXT: 2.000000 +// CHECK-NEXT: 4.000000 +// CHECK-NEXT: 6.000000 +// CHECK-NEXT: 8.000000 + +//TEST_INPUT:ubuffer(data=[1.0 2.0 3.0 4.0], stride=4, count=256),name=input1 +ByteAddressBuffer input; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +void doubleCoopMat(inout CoopMatType mat) +{ + mat = mat * 2.0; +} + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + var mat = CoopMatType.load(input, 0, stride, matrixLayout); + doubleCoopMat(mat); + mat.store(outputBuffer, 0, stride, matrixLayout); +} diff --git a/tests/cooperative-matrix/load-store-arbitrary-array-vec.slang b/tests/cooperative-matrix/load-store-arbitrary-array-vec.slang new file mode 100644 index 000000000..0afad3284 --- /dev/null +++ b/tests/cooperative-matrix/load-store-arbitrary-array-vec.slang @@ -0,0 +1,41 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -emit-spirv-directly + +// CHECK: 1 +// CHECK-NEXT: 2 +// CHECK-NEXT: 3 +// CHECK-NEXT: 4 +// CHECK-NEXT: 5 +// CHECK-NEXT: 6 +// CHECK-NEXT: 7 +// CHECK-NEXT: 8 +// CHECK-NEXT: 9 +// CHECK-NEXT: A +// CHECK-NEXT: B +// CHECK-NEXT: C +// CHECK-NEXT: D +// CHECK-NEXT: E +// CHECK-NEXT: F +// CHECK-NEXT: 10 + +//TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16], stride=4, count=256):name=input +RWByteAddressBuffer input; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +groupshared float3[128] tempShared; + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + let mat = coopMatLoad(input, 0, stride, matrixLayout); + mat.storeAny(tempShared, 0, stride, matrixLayout); + + let result = CoopMatType.loadAny(tempShared, 0, stride, matrixLayout); + result.store(outputBuffer, 0, stride, matrixLayout); +} diff --git a/tests/cooperative-matrix/load-store-arbitrary-array.slang b/tests/cooperative-matrix/load-store-arbitrary-array.slang new file mode 100644 index 000000000..496e62387 --- /dev/null +++ b/tests/cooperative-matrix/load-store-arbitrary-array.slang @@ -0,0 +1,41 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -emit-spirv-directly + +// CHECK: 1 +// CHECK-NEXT: 2 +// CHECK-NEXT: 3 +// CHECK-NEXT: 4 +// CHECK-NEXT: 5 +// CHECK-NEXT: 6 +// CHECK-NEXT: 7 +// CHECK-NEXT: 8 +// CHECK-NEXT: 9 +// CHECK-NEXT: A +// CHECK-NEXT: B +// CHECK-NEXT: C +// CHECK-NEXT: D +// CHECK-NEXT: E +// CHECK-NEXT: F +// CHECK-NEXT: 10 + +//TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16], stride=4, count=256):name=input +RWByteAddressBuffer input; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +groupshared float[256] tempShared; + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + let mat = coopMatLoad(input, 0, stride, matrixLayout); + mat.storeAny(tempShared, 0, stride, matrixLayout); + + let result = CoopMatType.loadAny(tempShared, 0, stride, matrixLayout); + result.store(outputBuffer, 0, stride, matrixLayout); +} diff --git a/tests/cooperative-matrix/load-store-groupshared.slang b/tests/cooperative-matrix/load-store-groupshared.slang new file mode 100644 index 000000000..c2334c0ce --- /dev/null +++ b/tests/cooperative-matrix/load-store-groupshared.slang @@ -0,0 +1,31 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -emit-spirv-directly + +// CHECK: 1 +// CHECK-NEXT: 2 +// CHECK-NEXT: 3 +// CHECK-NEXT: 4 +// CHECK-NEXT: 5 +// CHECK-NEXT: 6 +// CHECK-NEXT: 7 +// CHECK-NEXT: 8 + +//TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8], stride=4, count=256):name=input +RWByteAddressBuffer input; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +groupshared uint32_t[256] tempShared; + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + let mat = coopMatLoad(input, 0, stride, matrixLayout); + mat.store(tempShared, 0, stride, matrixLayout); + + let result = coopMatLoad(tempShared, 0, stride, matrixLayout); + result.store(outputBuffer, 0, stride, matrixLayout); +} diff --git a/tests/cooperative-matrix/load-store-rwbyteaddressbuffer.slang b/tests/cooperative-matrix/load-store-rwbyteaddressbuffer.slang new file mode 100644 index 000000000..08c90992a --- /dev/null +++ b/tests/cooperative-matrix/load-store-rwbyteaddressbuffer.slang @@ -0,0 +1,26 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: 1 +// CHECK-NEXT: 2 +// CHECK-NEXT: 3 +// CHECK-NEXT: 4 +// CHECK-NEXT: 5 +// CHECK-NEXT: 6 +// CHECK-NEXT: 7 +// CHECK-NEXT: 8 + +//TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8], stride=4, count=256):name=inputBuffer +RWByteAddressBuffer inputBuffer; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWByteAddressBuffer outputBuffer; + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + let mat = coopMatLoad(inputBuffer, 0, stride, matrixLayout); + mat.store(outputBuffer, 0, 16, matrixLayout); +} diff --git a/tests/cooperative-matrix/load-store-rwstructuredbuffer.slang b/tests/cooperative-matrix/load-store-rwstructuredbuffer.slang new file mode 100644 index 000000000..d71634082 --- /dev/null +++ b/tests/cooperative-matrix/load-store-rwstructuredbuffer.slang @@ -0,0 +1,27 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: int32_t +// CHECK-NEXT: 1 +// CHECK-NEXT: 2 +// CHECK-NEXT: 3 +// CHECK-NEXT: 4 +// CHECK-NEXT: 5 +// CHECK-NEXT: 6 +// CHECK-NEXT: 7 +// CHECK-NEXT: 8 + +//TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8], stride=4, count=256),name=buf +RWStructuredBuffer inputBuffer; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + let mat = coopMatLoad(inputBuffer, 0, stride, matrixLayout); + mat.store(outputBuffer, 0, stride, matrixLayout); +} diff --git a/tests/cooperative-matrix/mat-mul-add-spirv-matrix-operands.slang b/tests/cooperative-matrix/mat-mul-add-spirv-matrix-operands.slang new file mode 100644 index 000000000..934104a28 --- /dev/null +++ b/tests/cooperative-matrix/mat-mul-add-spirv-matrix-operands.slang @@ -0,0 +1,45 @@ +//TEST(compute):SIMPLE(filecheck=CHECK): -entry computeMain -stage compute -target spirv + +// This test checks that the correct SPIRV Cooperative Matrix Operands are emitted for OpCooperativeMatrixMulAddKHR operaions +RWStructuredBuffer outputBuffer1; +RWStructuredBuffer outputBuffer2; +RWStructuredBuffer outputBuffer3; +RWStructuredBuffer outputBuffer4; +RWStructuredBuffer outputBuffer5; + +typealias CoopMatAType = CoopMat; +typealias CoopMatBType = CoopMat; +typealias CoopMatCType = CoopMat; + +[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); + + // CHECK: OpCooperativeMatrixMulAddKHR {{.*}} MatrixASignedComponentsKHR|MatrixBSignedComponentsKHR + coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::MatrixASigned | CoopMatMatrixOperands::MatrixBSigned).store(outputBuffer2, 0, 16, matrixLayout); + + + // CHECK: OpCooperativeMatrixMulAddKHR {{.*}} MatrixCSignedComponentsKHR + coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::MatrixCSigned).store(outputBuffer2, 0, 16, matrixLayout); + + + // CHECK: OpCooperativeMatrixMulAddKHR {{.*}} MatrixResultSignedComponentsKHR + coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::MatrixResultSigned).store(outputBuffer3, 0, 16, matrixLayout); + + // CHECK: OpCooperativeMatrixMulAddKHR {{.*}} SaturatingAccumulationKHR + coopMatMulAdd(matA, matB, matC, CoopMatMatrixOperands::SaturatingAccumulation).store(outputBuffer4, 0, 16, matrixLayout); + + 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); +} + + diff --git a/tests/cooperative-matrix/mat-mul-add.slang b/tests/cooperative-matrix/mat-mul-add.slang new file mode 100644 index 000000000..aea61989d --- /dev/null +++ b/tests/cooperative-matrix/mat-mul-add.slang @@ -0,0 +1,23 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: float +// CHECK-COUNT-256: 241.0 + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatAType = CoopMat; +typealias CoopMatBType = CoopMat; +typealias CoopMatCType = CoopMat; + +[numthreads(32, 1, 1)] +void computeMain() +{ + // ( 3.0 * 5.0 ) * 16 + 1.0 = 241.0 + 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/mod.slang b/tests/cooperative-matrix/mod.slang new file mode 100644 index 000000000..116713481 --- /dev/null +++ b/tests/cooperative-matrix/mod.slang @@ -0,0 +1,53 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -emit-spirv-directly + +// CHECK: 0 +// CHECK-NEXT: 0 +// CHECK-NEXT: 1 +// CHECK-NEXT: 2 + +// CHECK: 0 +// CHECK-NEXT: 0 +// CHECK: 1 +// CHECK-NEXT: 2 + +// CHECK: 0 +// CHECK-NEXT: 0 +// CHECK: 1 +// CHECK-NEXT: 2 + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWByteAddressBuffer outputBuffer; + +//TEST_INPUT:ubuffer(data=[4 3 5 7], stride=4, count=256),name=input1 +ByteAddressBuffer input1; + +//TEST_INPUT:ubuffer(data=[2 3 4 5], stride=4, count=256),name=input2 +ByteAddressBuffer input2; + +typealias CoopMatIntType = CoopMat; +typealias CoopMatUintType = CoopMat; +typealias CoopMatFloatType = CoopMat; + +[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 mat3 = CoopMatFloatType(mat1); + let mat4 = CoopMatFloatType(mat2); + + let mat5 = CoopMatUintType(mat1); + let mat6 = CoopMatUintType(mat2); + + let result = mat1 % mat2; + 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); +} diff --git a/tests/cooperative-matrix/mul.slang b/tests/cooperative-matrix/mul.slang new file mode 100644 index 000000000..9b9fd67af --- /dev/null +++ b/tests/cooperative-matrix/mul.slang @@ -0,0 +1,31 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: int32_t +// CHECK-NEXT: 2 +// CHECK-NEXT: 6 +// CHECK-NEXT: 12 +// CHECK-NEXT: 20 + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +//TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4, count=256),name=input1 +ByteAddressBuffer input1; + +//TEST_INPUT:ubuffer(data=[2 3 4 5], stride=4, count=256),name=input2 +ByteAddressBuffer input2; + +typealias CoopMatType = CoopMat; + +[numthreads(32, 1, 1)] +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 result = mat1 * mat2; + result.store(outputBuffer, 0, 4, matrixLayout); +} diff --git a/tests/cooperative-matrix/out.slang b/tests/cooperative-matrix/out.slang new file mode 100644 index 000000000..147c2dd77 --- /dev/null +++ b/tests/cooperative-matrix/out.slang @@ -0,0 +1,34 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: float +// CHECK-NEXT: 2.000000 +// CHECK-NEXT: 4.000000 +// CHECK-NEXT: 6.000000 +// CHECK-NEXT: 8.000000 + +// XXX FW: having out instead of in below actually properly creates two output buffers, nice +//TEST_INPUT:ubuffer(data=[1.0 2.0 3.0 4.0], stride=4, count=256):name=inputBuffer +StructuredBuffer inputBuffer; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +void doubleCoopMat(CoopMatType mat, out CoopMatType result) +{ + result = mat * 2.0; +} + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout); + + CoopMatType result; + doubleCoopMat(mat, result); + result.store(outputBuffer, 0, stride, matrixLayout); +} diff --git a/tests/cooperative-matrix/parameter.slang b/tests/cooperative-matrix/parameter.slang new file mode 100644 index 000000000..19adf4177 --- /dev/null +++ b/tests/cooperative-matrix/parameter.slang @@ -0,0 +1,32 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: float +// CHECK-NEXT: 3.000000 +// CHECK-NEXT: 6.000000 +// CHECK-NEXT: 9.000000 +// CHECK-NEXT: 12.000000 + +//TEST_INPUT:ubuffer(data=[1.0 2.0 3.0 4.0], stride=4, count=256):name=inputBuffer +StructuredBuffer inputBuffer; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +static let stride = 16; +//static let matrixLayout = CoopMatMatrixLayout::RowMajor; +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); +} + +[numthreads(32, 1, 1)] +void computeMain() +{ + let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout); + processCoopMat(mat); +} diff --git a/tests/cooperative-matrix/return.slang b/tests/cooperative-matrix/return.slang new file mode 100644 index 000000000..339c9d04d --- /dev/null +++ b/tests/cooperative-matrix/return.slang @@ -0,0 +1,32 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type + +// CHECK: type: float +// CHECK-NEXT: 3.000000 +// CHECK-NEXT: 6.000000 +// CHECK-NEXT: 9.000000 +// CHECK-NEXT: 12.000000 + +//TEST_INPUT:ubuffer(data=[1.0 2.0 3.0 4.0], stride=4, count=256):name=inputBuffer +StructuredBuffer inputBuffer; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +CoopMatType doubleCoopmat(CoopMatType mat) +{ + return mat * 3.0; +} + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout); + + let result = doubleCoopmat(mat); + result.store(outputBuffer, 0, 4, matrixLayout); +} diff --git a/tests/cooperative-matrix/scalar-mul.slang b/tests/cooperative-matrix/scalar-mul.slang new file mode 100644 index 000000000..73f0fbbfc --- /dev/null +++ b/tests/cooperative-matrix/scalar-mul.slang @@ -0,0 +1,27 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: float +// CHECK-NEXT: 4.500000 +// CHECK-NEXT: 9.000000 +// CHECK-NEXT: 13.500000 +// CHECK-NEXT: 18.000000 + +//TEST_INPUT:ubuffer(data=[1.0 2.0 3.0 4.0], stride=4, count=256),name=inputBuffer +ByteAddressBuffer inputBuffer; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout); + + let result = mat * 4.5; + result.store(outputBuffer, 0, 4, matrixLayout); +} diff --git a/tests/cooperative-matrix/struct.slang b/tests/cooperative-matrix/struct.slang new file mode 100644 index 000000000..24bc2c367 --- /dev/null +++ b/tests/cooperative-matrix/struct.slang @@ -0,0 +1,42 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: float +// CHECK-NEXT: 1.000000 +// CHECK-NEXT: 2.000000 +// CHECK-NEXT: 3.000000 +// CHECK-NEXT: 4.000000 +// CHECK-NEXT: 5.000000 +// CHECK-NEXT: 6.000000 +// CHECK-NEXT: 7.000000 +// CHECK-NEXT: 8.000000 + +//TEST_INPUT:ubuffer(data=[1.0 2.0 3.0 4.0], stride=4, count=256),name=input1 +ByteAddressBuffer input1; + +//TEST_INPUT:ubuffer(data=[5.0 6.0 7.0 8.0], stride=4, count=256),name=input2 +ByteAddressBuffer input2; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +struct MyStruct +{ + CoopMatType mat1; + CoopMatType mat2; +}; + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + MyStruct s; + s.mat1 = CoopMatType.load(input1, 0, stride, matrixLayout); + s.mat2 = CoopMatType.load(input2, 0, stride, matrixLayout); + + s.mat1.store(outputBuffer, 0, stride, matrixLayout); + s.mat2.store(outputBuffer, 4, stride, matrixLayout); +} diff --git a/tests/cooperative-matrix/sub.slang b/tests/cooperative-matrix/sub.slang new file mode 100644 index 000000000..7b20d7c11 --- /dev/null +++ b/tests/cooperative-matrix/sub.slang @@ -0,0 +1,31 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: int32_t +// CHECK-NEXT: 1 +// CHECK-NEXT: 1 +// CHECK-NEXT: 1 +// CHECK-NEXT: 1 + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +//TEST_INPUT:ubuffer(data=[2 3 4 5], stride=4, count=256),name=input1 +ByteAddressBuffer input1; + +//TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4, count=256),name=input2 +ByteAddressBuffer input2; + +typealias CoopMatType = CoopMat; + +[numthreads(32, 1, 1)] +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 result = mat1 - mat2; + result.store(outputBuffer, 0, 4, matrixLayout); +} diff --git a/tests/cooperative-matrix/subscript-in-func.slang b/tests/cooperative-matrix/subscript-in-func.slang new file mode 100644 index 000000000..ef63d62de --- /dev/null +++ b/tests/cooperative-matrix/subscript-in-func.slang @@ -0,0 +1,34 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: float +// CHECK-NEXT: 1.000000 +// CHECK-NEXT: 4.000000 +// CHECK-NEXT: 9.000000 +// CHECK-NEXT: 16.000000 + +//TEST_INPUT:ubuffer(data=[1.0 2.0 3.0 4.0], stride=4, count=256):name=inputBuffer +StructuredBuffer inputBuffer; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +static const int stride = 16; +static const CoopMatMatrixLayout matrixLayout = CoopMatMatrixLayout::RowMajor; + +void squareCoopMatElements(CoopMatType mat) +{ + for (int i = 0; i < 4; ++i) + { + mat[i] = mat[i] * mat[i]; + } + mat.store(outputBuffer, 0, stride, matrixLayout); +} + +[numthreads(32, 1, 1)] +void computeMain() +{ + let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout); + squareCoopMatElements(mat); +} diff --git a/tests/cooperative-matrix/subscript.slang b/tests/cooperative-matrix/subscript.slang new file mode 100644 index 000000000..731edee82 --- /dev/null +++ b/tests/cooperative-matrix/subscript.slang @@ -0,0 +1,23 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: int32_t +// CHECK-NEXT: 2 +// CHECK-NEXT: 4 +// CHECK: 7 +// CHECK-NEXT: 11 + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +[numthreads(32, 1, 1)] +void computeMain() +{ + CoopMatType mat; + mat[0] = 2; + mat[1] = mat[0]+2; + mat[2] = mat[1]+3; + mat[3] = mat[2]+4; + mat.store(outputBuffer, 0, 16, CoopMatMatrixLayout::RowMajor); +} diff --git a/tests/cooperative-matrix/unary_neg.slang b/tests/cooperative-matrix/unary_neg.slang new file mode 100644 index 000000000..8c6436caf --- /dev/null +++ b/tests/cooperative-matrix/unary_neg.slang @@ -0,0 +1,27 @@ +//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly + +// CHECK: type: int32_t +// CHECK-NEXT: -1 +// CHECK-NEXT: -2 +// CHECK-NEXT: -3 +// CHECK-NEXT: -4 + +//TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4, count=256),name=inputBuffer +ByteAddressBuffer inputBuffer; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typealias CoopMatType = CoopMat; + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 4; + let matrixLayout = CoopMatMatrixLayout::RowMajor; + + let mat = CoopMatType.load(inputBuffer, 0, stride, matrixLayout); + + let result = -mat; + result.store(outputBuffer, 0, 4, matrixLayout); +} -- cgit v1.2.3