yum-mirror/slang

Making it easier to work with shaders

git clone https://git.yummers.dev/yum-mirror/slang

Jay KwakImplement MapElement for CoopMat (#7159)984d7f22f

master
2.5 KiB64 linesraw
1//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -Xslang -DTEST_MODE=0 -render-feature cooperative-matrix-per-element-operations
2//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -Xslang -DTEST_MODE=1 -render-feature cooperative-matrix-per-element-operations
3//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -Xslang -DTEST_MODE=2 -render-feature cooperative-matrix-per-element-operations
4//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -Xslang -DTEST_MODE=3 -render-feature cooperative-matrix-per-element-operations
5//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -Xslang -DTEST_MODE=4 -render-feature cooperative-matrix-per-element-operations
6
7//CHECK: type: int32_t
8//CHECK-NEXT: 8
9//CHECK-NEXT: 10
10//CHECK-NEXT: 12
11//CHECK-NEXT: 14
12
13//TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4),name=input1
14StructuredBuffer<int> input1;
15
16//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer
17RWStructuredBuffer<int32_t> outputBuffer;
18
19using namespace linalg;
20
21typealias CoopMatType = CoopMat<int, MemoryScope.Subgroup, 16, 16, CoopMatMatrixUse.MatrixAccumulator>;
22
23int MapOp(uint32_t row, uint32_t col, int value)
24{
25    return value * 2 + 1 + 2 + 3;
26}
27
28[numthreads(32, 1, 1)]
29void computeMain()
30{
31    let stride = 16;
32    CoopMatType mat1 = CoopMatType.Load<CoopMatMatrixLayout.RowMajor>(input1, 0, stride);
33
34    // Testing the capturing lambda
35    int c0 = 1;
36    int c1 = 2;
37    int c2 = 3;
38
39    CoopMatType result;
40
41#if TEST_MODE == 0
42    result = mat1.MapElement(MapOp);
43
44#elif TEST_MODE == 1
45    // Lambda via a temp variable (no capture)
46    let func = ((uint32_t row, uint32_t column, int value) => value * 2 + 1 + 2 + 3);
47    result = mat1.MapElement(func);
48
49#elif TEST_MODE == 2
50    // Directly use lambda (no capture)
51    result = mat1.MapElement((uint32_t row, uint32_t column, int value) => value * 2 + 1 + 2 + 3);
52
53#elif TEST_MODE == 3
54    // Lambda via a temp variable (capture)
55    let func = ((uint32_t row, uint32_t column, int value) => value * 2 + c0 + c1 + c2);
56    result = mat1.MapElement(func);
57
58#elif TEST_MODE == 4
59    // Directly use lambda (capture)
60    result = mat1.MapElement((uint32_t row, uint32_t column, int value) => value * 2 + c0 + c1 + c2);
61#endif
62
63    result.Store<CoopMatMatrixLayout.RowMajor>(outputBuffer, 0, stride);
64}