Address glslang ordering requirments for 'derivative_group_*NV' (#4323)

* Address glslang ordering requirments for 'derivative_group_*NV'

fixes: #4305

The solution is to emit some `layout`s after a module source is emitted.

Added to slangs gfx backend code to enable the compute shader derivative extension for testing purposes.

* address review

* enable removed test

---------

Co-authored-by: Yong He <yonghe@outlook.com>

8 files changed, 74 insertions, 11 deletions

diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index d85c72ae5..6412ff730 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -129,9 +129,12 @@ void CLikeSourceEmitter::emitPreModuleImpl()
         m_writer->emit(prelude->getStringSlice());
         m_writer->emit("\n");
     }
-    for (auto prelude : m_requiredPreludesRaw)
+}
+void CLikeSourceEmitter::emitPostModuleImpl()
+{   
+    if(m_requiredAfter.requireComputeDerivatives.getLength() > 0)
     {
-        m_writer->emit(prelude);
+        m_writer->emit(m_requiredAfter.requireComputeDerivatives);
         m_writer->emit("\n");
     }
 }
diff --git a/source/slang/slang-emit-c-like.h b/source/slang/slang-emit-c-like.h
index 9b39aa6f4..ecff57296 100644
--- a/source/slang/slang-emit-c-like.h
+++ b/source/slang/slang-emit-c-like.h
@@ -431,7 +431,7 @@ public:
     void emitFrontMatter(TargetRequest* targetReq) { emitFrontMatterImpl(targetReq); }
 
     void emitPreModule() { emitPreModuleImpl(); }
-
+    void emitPostModule() { emitPostModuleImpl(); }
     void emitModule(IRModule* module, DiagnosticSink* sink)
         { m_irModule = module; emitModuleImpl(module, sink); }
 
@@ -476,6 +476,7 @@ public:
         /// For example on targets that don't have built in vector/matrix support, this is where
         /// the appropriate generated declarations occur.
     virtual void emitPreModuleImpl();
+    virtual void emitPostModuleImpl();
 
     virtual void beforeComputeEmitActions(IRModule* module) { SLANG_UNUSED(module); };
 
@@ -590,8 +591,11 @@ public:
     // to use for it when emitting code.
     Dictionary<IRInst*, String> m_mapInstToName;
 
-    OrderedHashSet<String> m_requiredPreludesRaw;
     OrderedHashSet<IRStringLit*> m_requiredPreludes;
+    struct RequiredAfter
+    {
+        String requireComputeDerivatives;
+    }m_requiredAfter;
 };
 
 }
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index 8063fcc1d..d13bc96d1 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -2208,9 +2208,7 @@ void GLSLSourceEmitter::handleRequiredCapabilitiesImpl(IRInst* inst)
         {
             // only allowed 1 of derivative_group_quadsNV or derivative_group_linearNV
             if (m_entryPointStage != Stage::Compute
-                || m_requiredPreludesRaw.contains("layout(derivative_group_quadsNV) in;")
-                || m_requiredPreludesRaw.contains("layout(derivative_group_linearNV) in;")
-                )
+                || m_requiredAfter.requireComputeDerivatives.getLength() > 0)
                 return;
 
             _requireGLSLExtension(UnownedStringSlice("GL_NV_compute_shader_derivatives"));
@@ -2225,12 +2223,12 @@ void GLSLSourceEmitter::handleRequiredCapabilitiesImpl(IRInst* inst)
                 if (isQuad)
                 {
                     verifyComputeDerivativeGroupModifiers(getSink(), inst->sourceLoc, true, false, numThreadsDecor);
-                    m_requiredPreludesRaw.add("layout(derivative_group_quadsNV) in;");
+                    m_requiredAfter.requireComputeDerivatives = "layout(derivative_group_quadsNV) in;";
                 }
                 else
                 {
                     verifyComputeDerivativeGroupModifiers(getSink(), inst->sourceLoc, false, true, numThreadsDecor);
-                    m_requiredPreludesRaw.add("layout(derivative_group_linearNV) in;");
+                    m_requiredAfter.requireComputeDerivatives = "layout(derivative_group_linearNV) in;";
                 }
             }
         }
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index da348fa64..6d0ab1daa 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -1587,6 +1587,10 @@ SlangResult CodeGenContext::emitEntryPointsSourceFromIR(ComPtr<IArtifact>& outAr
     // Append the modules output code
     finalResult.append(code);
 
+    // Append all content that should be at the end of a module
+    sourceEmitter->emitPostModule();
+    finalResult.append(sourceWriter.getContentAndClear());
+
     // Write out the result
 
     auto artifact = ArtifactUtil::createArtifactForCompileTarget(asExternal(target));
diff --git a/tests/bugs/gh-4305.slang b/tests/bugs/gh-4305.slang
new file mode 100644
index 000000000..dd09bda1e
--- /dev/null
+++ b/tests/bugs/gh-4305.slang
@@ -0,0 +1,40 @@
+//TEST:SIMPLE(filecheck=CHECK): -stage compute  -entry computeMain -target glsl -allow-glsl -DQUAD -DMODIFIER
+//TEST:SIMPLE(filecheck=CHECK): -stage compute  -entry computeMain -target glsl -allow-glsl -DMODIFIER
+//TEST:SIMPLE(filecheck=CHECK): -stage compute  -entry computeMain -target glsl -allow-glsl -DQUAD
+//TEST:SIMPLE(filecheck=CHECK): -stage compute  -entry computeMain -target glsl -allow-glsl
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -compute -entry computeMain -output-using-type -emit-spirv-via-glsl -allow-glsl -xslang -DQUAD
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -compute -entry computeMain -output-using-type -emit-spirv-via-glsl -allow-glsl
+
+//TEST_INPUT: ubuffer(data=[0], stride=4):out,name outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+//TEST_INPUT: Texture2D(size=4, content = one):name t2D
+Texture2D<float> t2D;
+//TEST_INPUT: Sampler:name samplerState
+SamplerState samplerState;
+
+// "local_size_x" must appear before "derivative_group_"
+//CHECK:local_size_x
+//CHECK:derivative_group_{{.*}}NV
+
+#ifndef MODIFIER
+#ifdef QUAD
+layout(derivative_group_quadsNV) in;
+#else
+layout(derivative_group_linearNV) in;
+#endif // #ifdef QUAD
+layout(local_size_x = 2, local_size_y = 2, local_size_z = 1) in;
+
+#else
+
+#ifdef QUAD
+[DerivativeGroupQuad]
+#else
+[DerivativeGroupLinear]
+#endif // #ifdef QUAD
+[numthreads(2, 2, 1)]
+#endif //#ifndef MODIFIER
+void computeMain()
+{
+    //BUF:1
+    outputBuffer[0] = t2D.Sample(samplerState, float2(0.5, 0.5));
+}
\ No newline at end of file
diff --git a/tests/glsl-intrinsic/intrinsic-texture.slang b/tests/glsl-intrinsic/intrinsic-texture.slang
index 11cf2c5d0..98d93e09f 100644
--- a/tests/glsl-intrinsic/intrinsic-texture.slang
+++ b/tests/glsl-intrinsic/intrinsic-texture.slang
@@ -7,8 +7,7 @@
 //TEST:SIMPLE(filecheck=CUDA): -allow-glsl -stage compute  -entry computeMain -target cuda
 //TEST:SIMPLE(filecheck=CUDA): -allow-glsl -stage fragment -entry fragMain -target cuda
 
-// Enable this test when an issue#4305 is addressed.
-//T-EST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl -output-using-type
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl -output-using-type
 
 
 //TEST_INPUT:ubuffer(data=[0], stride=4):out,name=outputBuffer
diff --git a/tools/gfx/vulkan/vk-api.h b/tools/gfx/vulkan/vk-api.h
index 79d4f8cfd..dd3681537 100644
--- a/tools/gfx/vulkan/vk-api.h
+++ b/tools/gfx/vulkan/vk-api.h
@@ -287,6 +287,10 @@ struct VulkanExtendedFeatureProperties
         VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VARIABLE_POINTER_FEATURES_KHR
     };
 
+    VkPhysicalDeviceComputeShaderDerivativesFeaturesNV computeShaderDerivativeFeatures = {
+        VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COMPUTE_SHADER_DERIVATIVES_FEATURES_NV
+    };
+
     // Clock features
     VkPhysicalDeviceShaderClockFeaturesKHR clockFeatures = { 
         VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_CLOCK_FEATURES_KHR
diff --git a/tools/gfx/vulkan/vk-device.cpp b/tools/gfx/vulkan/vk-device.cpp
index ad0924294..3583e2a1f 100644
--- a/tools/gfx/vulkan/vk-device.cpp
+++ b/tools/gfx/vulkan/vk-device.cpp
@@ -464,6 +464,10 @@ Result DeviceImpl::initVulkanInstanceAndDevice(
         extendedFeatures.variablePointersFeatures.pNext = deviceFeatures2.pNext;
         deviceFeatures2.pNext = &extendedFeatures.variablePointersFeatures;
         
+        // Compute shader derivative features.
+        extendedFeatures.computeShaderDerivativeFeatures.pNext = deviceFeatures2.pNext;
+        deviceFeatures2.pNext = &extendedFeatures.computeShaderDerivativeFeatures;
+
         // Extended dynamic states
         extendedFeatures.extendedDynamicStateFeatures.pNext = deviceFeatures2.pNext;
         deviceFeatures2.pNext = &extendedFeatures.extendedDynamicStateFeatures;
@@ -696,6 +700,13 @@ Result DeviceImpl::initVulkanInstanceAndDevice(
             VK_KHR_VARIABLE_POINTERS_EXTENSION_NAME,
             "variable-pointer"
         );
+        
+        SIMPLE_EXTENSION_FEATURE(
+            extendedFeatures.computeShaderDerivativeFeatures,
+            computeDerivativeGroupLinear,
+            VK_NV_COMPUTE_SHADER_DERIVATIVES_EXTENSION_NAME,
+            "computeDerivativeGroupLinear"
+        );
 
 #undef SIMPLE_EXTENSION_FEATURE