From 4dd41de7558b2b3b7fe99b4b624dc91969031abe Mon Sep 17 00:00:00 2001
From: ArielG-NV <159081215+ArielG-NV@users.noreply.github.com>
Date: Thu, 27 Jun 2024 07:53:00 -0400
Subject: Remove returned-array-legalization pass for metal (#4478)

* disable return array optimization pass for metal targets

fixes: #4468
---
 source/slang/slang-emit.cpp                        |  3 +-
 tests/bugs/gh-2959.slang                           |  2 +-
 .../groupshared-threadlocal-same-parameter.slang   | 53 ++++++++++++++++++++++
 3 files changed, 56 insertions(+), 2 deletions(-)
 create mode 100644 tests/metal/groupshared-threadlocal-same-parameter.slang

diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index bb6bab0ab..d8f0686d5 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -1269,7 +1269,8 @@ Result linkAndOptimizeIR(
 
     // Rewrite functions that return arrays to return them via `out` parameter,
     // since our target languages doesn't allow returning arrays.
-    legalizeArrayReturnType(irModule);
+    if(!isMetalTarget(targetRequest))
+        legalizeArrayReturnType(irModule);
 
     if (isKhronosTarget(targetRequest) || target == CodeGenTarget::HLSL)
     {
diff --git a/tests/bugs/gh-2959.slang b/tests/bugs/gh-2959.slang
index f6c58d1ee..056900b95 100644
--- a/tests/bugs/gh-2959.slang
+++ b/tests/bugs/gh-2959.slang
@@ -1,5 +1,5 @@
 //TEST(compute):COMPARE_COMPUTE(filecheck-buffer=BUF): -shaderobj -output-using-type
-//DISABLE_TEST(compute):COMPARE_COMPUTE:-slang -shaderobj -mtl
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=BUF):-slang -shaderobj -mtl
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name outputBuffer
 RWStructuredBuffer<uint> outputBuffer;
diff --git a/tests/metal/groupshared-threadlocal-same-parameter.slang b/tests/metal/groupshared-threadlocal-same-parameter.slang
new file mode 100644
index 000000000..ae4cac8df
--- /dev/null
+++ b/tests/metal/groupshared-threadlocal-same-parameter.slang
@@ -0,0 +1,53 @@
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=BUF): -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=BUF): -vk -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=BUF): -vk -emit-spirv-directly -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=BUF):-slang -shaderobj -mtl
+
+//TEST_INPUT:ubuffer(data=[0 0], stride=4):out,name outputBuffer
+RWStructuredBuffer<uint> outputBuffer;
+
+static groupshared uint g_values[2] = { 1, 0 };
+static uint g_altValues[2] = { 2, 3 };
+
+static groupshared uint g_valuesReturned[2];
+static uint g_altValuesReturned[2];
+
+uint[2] maybeGroupSharedReturn(uint id)
+{
+    if (id == 0)
+    {
+        return g_values;
+    }
+    else
+    {
+        return g_altValues;
+    }
+}
+
+[numthreads(2, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    AllMemoryBarrierWithGroupSync();
+    uint tid = dispatchThreadID.x;
+    if (tid == 0)
+    {
+        g_valuesReturned[tid] = maybeGroupSharedReturn(tid)[tid];
+    }
+    else
+    {
+        g_altValuesReturned[tid] = maybeGroupSharedReturn(tid)[tid];
+    }
+
+    AllMemoryBarrierWithGroupSync();
+    if (tid == 0)
+    {
+        outputBuffer[tid] = g_valuesReturned[tid];
+    }
+    else
+    {
+        outputBuffer[tid] = g_altValuesReturned[tid];
+    }
+
+    // BUF: 1
+    // BUF: 3
+}
-- 
cgit v1.2.3