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<CoopMat<...>,...>::MapElement(functype(...));
 - Tuple<CoopMat<...>,...>::MapElement(capturing-lambda);
 - Tuple<CoopMat<...>,...>::MapElement(not-capturing-lambda);
---
 .../tuple/tuple-expand-multiple.slang              | 37 ++++++++++++++++++++++
 tests/language-feature/tuple/tuple-expand.slang    | 25 +++++++++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 tests/language-feature/tuple/tuple-expand-multiple.slang
 create mode 100644 tests/language-feature/tuple/tuple-expand.slang

(limited to 'tests/language-feature')

diff --git a/tests/language-feature/tuple/tuple-expand-multiple.slang b/tests/language-feature/tuple/tuple-expand-multiple.slang
new file mode 100644
index 000000000..9392e2d66
--- /dev/null
+++ b/tests/language-feature/tuple/tuple-expand-multiple.slang
@@ -0,0 +1,37 @@
+//TEST:COMPARE_COMPUTE(filecheck-buffer=CHECK): -output-using-type
+
+//TEST_INPUT: set outputBuffer = out ubuffer(data=[0 0 0 0 0 0], stride=4)
+RWStructuredBuffer<int> outputBuffer;
+
+extension<
+    T0 : __BuiltinArithmeticType,
+    each Ts0 : __BuiltinArithmeticType,
+    each Ts1 : __BuiltinArithmeticType
+> Tuple<T0, T0, expand each Ts0, expand each Ts1>
+{
+    static int getSize_Ts0() { return countof(Ts0); }
+    static int getSize_Ts1() { return countof(Ts1); }
+}
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    int i = 2;
+    float f0 = 3, f1 = 5;
+    uint ui0 = 4, ui1 = 6;
+
+    let s0 = makeTuple(i, i); // T, T
+    let s1 = makeTuple(i, i, f0, ui0); // T, T, Ts0, Ts1
+    let s2 = makeTuple(i, i, f0, ui0, f1, ui1); // T, T, Ts0, Ts0, Ts1, Ts1
+
+    outputBuffer[0] = s0.getSize_Ts0();
+    outputBuffer[1] = s0.getSize_Ts1();
+    outputBuffer[2] = s1.getSize_Ts0();
+    outputBuffer[3] = s1.getSize_Ts1();
+    outputBuffer[4] = s2.getSize_Ts0();
+    outputBuffer[5] = s2.getSize_Ts1();
+
+    // CHECK-COUNT-2:0
+    // CHECK-COUNT-2:1
+    // CHECK-COUNT-2:2
+}
diff --git a/tests/language-feature/tuple/tuple-expand.slang b/tests/language-feature/tuple/tuple-expand.slang
new file mode 100644
index 000000000..c9ea26161
--- /dev/null
+++ b/tests/language-feature/tuple/tuple-expand.slang
@@ -0,0 +1,25 @@
+//TEST:COMPARE_COMPUTE(filecheck-buffer=CHECK): -output-using-type
+
+//TEST_INPUT: set outputBuffer = out ubuffer(data=[0 0], stride=4)
+RWStructuredBuffer<int> outputBuffer;
+
+struct X<B, each T, each U>
+{
+    int getTSize() { return countof(T); }
+    int getUSize() { return countof(U); }
+}
+
+func foo<each T, each U>() -> X<bool, expand Ptr<each T>, int, expand Ptr<each U>, float> // unify
+{
+    return {};
+}
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    let x = foo<int, float>();
+
+    outputBuffer[0] = x.getTSize();
+    outputBuffer[1] = x.getUSize();
+    // CHECK-COUNT-2: 2
+}
-- 
cgit v1.2.3