From c5db04b8ba72d14d6bf2e20abf9c1b8d8466abc6 Mon Sep 17 00:00:00 2001
From: jsmall-nvidia <jsmall@nvidia.com>
Date: Thu, 2 Apr 2020 19:52:12 -0400
Subject: Fix WaveGetLaneIndex for glsl  (#1306)

* Fix typo in stdlib around WaveGetLaneIndex and WaveGetLaneCount

* Reorder emit so #extensions come before layout

* Added wave-get-lane-index.slang test.
---
 source/slang/hlsl.meta.slang                          |  4 ++--
 source/slang/slang-emit.cpp                           | 16 +++++++++-------
 tests/hlsl-intrinsic/wave-get-lane-index.slang        | 19 +++++++++++++++++++
 .../wave-get-lane-index.slang.expected.txt            |  4 ++++
 4 files changed, 34 insertions(+), 9 deletions(-)
 create mode 100644 tests/hlsl-intrinsic/wave-get-lane-index.slang
 create mode 100644 tests/hlsl-intrinsic/wave-get-lane-index.slang.expected.txt

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index f82f7b5f4..483c2299f 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -2659,13 +2659,13 @@ uint WaveActiveCountBits(bool value);
 
 __glsl_extension(GL_KHR_shader_subgroup_basic)
 __spirv_version(1.3)
-__target_intrinsic(glsl, "gl_SubgroupSize")
+__target_intrinsic(glsl, "(gl_SubgroupSize)")
 __target_intrinsic(cuda, "(warpSize)")
 uint WaveGetLaneCount();
 
 __glsl_extension(GL_KHR_shader_subgroup_basic)
 __spirv_version(1.3)
-__target_intrinsic(glsl, "gl_SubgroupInvocationID")
+__target_intrinsic(glsl, "(gl_SubgroupInvocationID)")
 __target_intrinsic(cuda, "_getLaneId()")
 uint WaveGetLaneIndex();
 
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index e613f5462..3caef0a9f 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -641,6 +641,15 @@ SlangResult emitEntryPointSourceFromIR(
     // There may be global-scope modifiers that we should emit now
     sourceEmitter->emitPreprocessorDirectives();
 
+    RefObject* extensionTracker = sourceEmitter->getExtensionTracker();
+
+    if (auto glslExtensionTracker = as<GLSLExtensionTracker>(extensionTracker))
+    {
+        StringBuilder builder;
+        glslExtensionTracker->appendExtensionRequireLines(builder);
+        sourceWriter.emit(builder.getUnownedSlice());
+    }
+
     sourceEmitter->emitLayoutDirectives(targetRequest);
 
     String prefix = sourceWriter.getContent();
@@ -648,13 +657,6 @@ SlangResult emitEntryPointSourceFromIR(
     StringBuilder finalResultBuilder;
     finalResultBuilder << prefix;
 
-    RefObject* extensionTracker = sourceEmitter->getExtensionTracker();
-
-    if (auto glslExtensionTracker = as<GLSLExtensionTracker>(extensionTracker))
-    {
-        glslExtensionTracker->appendExtensionRequireLines(finalResultBuilder);
-    }
-
     finalResultBuilder << code;
 
     // Write out the result
diff --git a/tests/hlsl-intrinsic/wave-get-lane-index.slang b/tests/hlsl-intrinsic/wave-get-lane-index.slang
new file mode 100644
index 000000000..50cf9822c
--- /dev/null
+++ b/tests/hlsl-intrinsic/wave-get-lane-index.slang
@@ -0,0 +1,19 @@
+//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute 
+//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-slang -compute
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx12 -use-dxil -profile cs_6_0
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute
+//TEST(compute):COMPARE_COMPUTE_EX:-cuda -compute
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    int idx = dispatchThreadID.x;
+    int laneId = WaveGetLaneIndex();
+    // The laneCount will be dependent on target hardware. It seems a count of 1 is valid in spec. 
+    // For now we'll just check it's not 0.
+    int laneCount = WaveGetLaneCount();
+    outputBuffer[idx] = ((laneCount > 0) ? 0x100 : 0) + laneId;
+}
\ No newline at end of file
diff --git a/tests/hlsl-intrinsic/wave-get-lane-index.slang.expected.txt b/tests/hlsl-intrinsic/wave-get-lane-index.slang.expected.txt
new file mode 100644
index 000000000..f94894bb2
--- /dev/null
+++ b/tests/hlsl-intrinsic/wave-get-lane-index.slang.expected.txt
@@ -0,0 +1,4 @@
+100
+101
+102
+103
-- 
cgit v1.2.3