From 984d7f22f8a0909dc870c65bb927094c54f55402 Mon Sep 17 00:00:00 2001 From: Jay Kwak <82421531+jkwak-work@users.noreply.github.com> Date: Thu, 29 May 2025 16:36:49 -0700 Subject: Implement MapElement for CoopMat (#7159) With this PR, MapElement works for the following signatures: - CoopMat<...>::MapElement(functype(...)); - CoopMat<...>::MapElement(capturing-lambda); - CoopMat<...>::MapElement(not-capturing-lambda); - Tuple,...>::MapElement(functype(...)); - Tuple,...>::MapElement(capturing-lambda); - Tuple,...>::MapElement(not-capturing-lambda); --- tests/cooperative-matrix/map-element-single.slang | 40 ++++++++----- tests/cooperative-matrix/map-element-tuple.slang | 68 +++++++++++++++++++++++ 2 files changed, 95 insertions(+), 13 deletions(-) create mode 100644 tests/cooperative-matrix/map-element-tuple.slang (limited to 'tests/cooperative-matrix') diff --git a/tests/cooperative-matrix/map-element-single.slang b/tests/cooperative-matrix/map-element-single.slang index 1661ee105..ecf35953e 100644 --- a/tests/cooperative-matrix/map-element-single.slang +++ b/tests/cooperative-matrix/map-element-single.slang @@ -1,12 +1,14 @@ -//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -emit-spirv-directly -render-feature cooperative-matrix-per-element-operations -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 -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 -Xslang -DTEST_MODE=2 +//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 +//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 +//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 +//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 +//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 //CHECK: type: int32_t -//CHECK-NEXT: 2 -//CHECK-NEXT: 4 -//CHECK-NEXT: 6 //CHECK-NEXT: 8 +//CHECK-NEXT: 10 +//CHECK-NEXT: 12 +//CHECK-NEXT: 14 //TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4),name=input1 StructuredBuffer input1; @@ -20,7 +22,7 @@ typealias CoopMatType = CoopMat(input1, 0, stride); + // Testing the capturing lambda + int c0 = 1; + int c1 = 2; + int c2 = 3; + 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 func = ((uint32_t row, uint32_t column, int value) => value * 2); + // Lambda via a temp variable (no capture) + let func = ((uint32_t row, uint32_t column, int value) => value * 2 + 1 + 2 + 3); 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)); + // Directly use lambda (no capture) + result = mat1.MapElement((uint32_t row, uint32_t column, int value) => value * 2 + 1 + 2 + 3); + +#elif TEST_MODE == 3 + // Lambda via a temp variable (capture) + let func = ((uint32_t row, uint32_t column, int value) => value * 2 + c0 + c1 + c2); + result = mat1.MapElement(func); + +#elif TEST_MODE == 4 + // Directly use lambda (capture) + result = mat1.MapElement((uint32_t row, uint32_t column, int value) => value * 2 + c0 + c1 + c2); #endif result.Store(outputBuffer, 0, stride); diff --git a/tests/cooperative-matrix/map-element-tuple.slang b/tests/cooperative-matrix/map-element-tuple.slang new file mode 100644 index 000000000..06ab99d8f --- /dev/null +++ b/tests/cooperative-matrix/map-element-tuple.slang @@ -0,0 +1,68 @@ +//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 +//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 +//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 +//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 +//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 + +//CHECK:type: int32_t +//CHECK-NEXT:9 +//CHECK-NEXT:12 +//CHECK-NEXT:15 +//CHECK-NEXT:14 + +//TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4),name=input1 +StructuredBuffer input1; + +//TEST_INPUT:ubuffer(data=[0 1 2 3], stride=4),name=input2 +StructuredBuffer input2; + +//TEST_INPUT:ubuffer(data=[2 3 4 1], stride=4),name=input3 +StructuredBuffer input3; + +//TEST_INPUT:ubuffer(stride=4, count=256):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +using namespace linalg; + +typealias CoopMatType = CoopMat; + +int MapOp(uint32_t row, uint32_t col, int a, int b, int c) +{ + return a + b + c + 1 + 2 + 3; +} + +[numthreads(32, 1, 1)] +void computeMain() +{ + let stride = 16; + let mat1 = CoopMatType.Load(input1, 0, stride); + let mat2 = CoopMatType.Load(input2, 0, stride); + let mat3 = CoopMatType.Load(input3, 0, stride); + + // Testing the capturing lambda + int c0 = 1; + int c1 = 2; + int c2 = 3; + + CoopMatType result; + +#if TEST_MODE == 0 + result = makeTuple(mat1, mat2, mat3).MapElement(MapOp); + +#elif TEST_MODE == 1 + let f = ((uint32_t x, uint32_t y, int a, int b, int c) => a + b + c + 1 + 2 + 3); + result = makeTuple(mat1, mat2, mat3).MapElement(f); + +#elif TEST_MODE == 2 + result = makeTuple(mat1, mat2, mat3).MapElement((uint32_t x, uint32_t y, int a, int b, int c) => a + b + c + 1 + 2 + 3); + +#elif TEST_MODE == 3 + let f = ((uint32_t x, uint32_t y, int a, int b, int c) => a + b + c + c0 + c1 + c2); + result = makeTuple(mat1, mat2, mat3).MapElement(f); + +#elif TEST_MODE == 4 + result = makeTuple(mat1, mat2, mat3).MapElement((uint32_t x, uint32_t y, int a, int b, int c) => a + b + c + c0 + c1 + c2); +#endif + + result.Store(outputBuffer, 0, stride); +} -- cgit v1.2.3