Add option to use original entrypoint in spirv output. (#3047)

* Add option to use original entrypoint in spirv output.

* Fix.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

6 files changed, 126 insertions, 23 deletions

diff --git a/source/compiler-core/slang-glslang-compiler.cpp b/source/compiler-core/slang-glslang-compiler.cpp
index 3c5f7a2b8..4c0bc74d7 100644
--- a/source/compiler-core/slang-glslang-compiler.cpp
+++ b/source/compiler-core/slang-glslang-compiler.cpp
@@ -56,11 +56,12 @@ public:
     
 protected:
 
-    SlangResult _invoke(glslang_CompileRequest_1_1& request);
+    SlangResult _invoke(glslang_CompileRequest_1_2& request);
 
     glslang_CompileFunc_1_0 m_compile_1_0 = nullptr; 
     glslang_CompileFunc_1_1 m_compile_1_1 = nullptr; 
-    
+    glslang_CompileFunc_1_2 m_compile_1_2 = nullptr;
+
     ComPtr<ISlangSharedLibrary> m_sharedLibrary;
 };
 
@@ -68,8 +69,10 @@ SlangResult GlslangDownstreamCompiler::init(ISlangSharedLibrary* library)
 {
     m_compile_1_0 = (glslang_CompileFunc_1_0)library->findFuncByName("glslang_compile");
     m_compile_1_1 = (glslang_CompileFunc_1_1)library->findFuncByName("glslang_compile_1_1");
+    m_compile_1_2 = (glslang_CompileFunc_1_2)library->findFuncByName("glslang_compile_1_2");
+
 
-    if (m_compile_1_0 == nullptr && m_compile_1_1 == nullptr)
+    if (m_compile_1_0 == nullptr && m_compile_1_1 == nullptr && m_compile_1_2 == nullptr)
     {
         return SLANG_FAIL;
     }
@@ -80,7 +83,11 @@ SlangResult GlslangDownstreamCompiler::init(ISlangSharedLibrary* library)
     m_desc = Desc(SLANG_PASS_THROUGH_GLSLANG);
 
     Slang::String filename;
-    if (m_compile_1_1)
+    if (m_compile_1_2)
+    {
+        filename = Slang::SharedLibraryUtils::getSharedLibraryFileName((void*)m_compile_1_2);
+    }
+    else if (m_compile_1_1)
     {
         filename = Slang::SharedLibraryUtils::getSharedLibraryFileName((void*)m_compile_1_1);
     }
@@ -96,17 +103,27 @@ SlangResult GlslangDownstreamCompiler::init(ISlangSharedLibrary* library)
     return SLANG_OK;
 }
 
-SlangResult GlslangDownstreamCompiler::_invoke(glslang_CompileRequest_1_1& request)
+SlangResult GlslangDownstreamCompiler::_invoke(glslang_CompileRequest_1_2& request)
 {
     int err = 1;
-    if (m_compile_1_1)
+    if (m_compile_1_2)
     {
-        err = m_compile_1_1(&request);
+        err = m_compile_1_2(&request);
+    }
+    else if (m_compile_1_1)
+    {
+        glslang_CompileRequest_1_1 request_1_1;
+        memcpy(&request_1_1, &request, sizeof(request_1_1));
+        request_1_1.sizeInBytes = sizeof(request_1_1);
+        err = m_compile_1_1(&request_1_1);
     }
     else if (m_compile_1_0)
     {
+        glslang_CompileRequest_1_1 request_1_1;
+        memcpy(&request_1_1, &request, sizeof(request_1_1));
+        request_1_1.sizeInBytes = sizeof(request_1_1);
         glslang_CompileRequest_1_0 request_1_0;
-        request_1_0.set(request);
+        request_1_0.set(request_1_1);
         err = m_compile_1_0(&request_1_0);
     }
 
@@ -178,7 +195,7 @@ SlangResult GlslangDownstreamCompiler::compile(const CompileOptions& inOptions,
 
     String sourcePath = ArtifactUtil::findPath(sourceArtifact);
 
-    glslang_CompileRequest_1_1 request;
+    glslang_CompileRequest_1_2 request;
     memset(&request, 0, sizeof(request));
     request.sizeInBytes = sizeof(request);
 
@@ -217,6 +234,8 @@ SlangResult GlslangDownstreamCompiler::compile(const CompileOptions& inOptions,
     request.optimizationLevel = (unsigned)options.optimizationLevel;
     request.debugInfoType = (unsigned)options.debugInfoType;
 
+    request.entryPointName = options.entryPointName.begin();
+
     const SlangResult invokeResult = _invoke(request);
 
     auto artifact = ArtifactUtil::createArtifactForCompileTarget(options.targetType);
@@ -271,7 +290,7 @@ SlangResult GlslangDownstreamCompiler::convert(IArtifact* from, const ArtifactDe
         (*(StringBuilder*)userData).append((char const*)data, (char const*)data + size);
     };
 
-    glslang_CompileRequest_1_1 request;
+    glslang_CompileRequest_1_2 request;
     memset(&request, 0, sizeof(request));
     request.sizeInBytes = sizeof(request);
 
diff --git a/source/slang-glslang/slang-glslang.cpp b/source/slang-glslang/slang-glslang.cpp
index ffe3161ea..b23a34e65 100644
--- a/source/slang-glslang/slang-glslang.cpp
+++ b/source/slang-glslang/slang-glslang.cpp
@@ -97,7 +97,7 @@ static void dump(
 }
 
 static void dumpDiagnostics(
-    const glslang_CompileRequest_1_1& request,
+    const glslang_CompileRequest_1_2& request,
     std::string const&      log)
 {
     dump(log.c_str(), log.length(), request.diagnosticFunc, request.diagnosticUserData, stderr);
@@ -147,7 +147,7 @@ struct SPIRVOptimizationDiagnostic
 
 // Apply the SPIRV-Tools optimizer to generated SPIR-V based on the desired optimization level
 // TODO: add flag for optimizing SPIR-V size as well
-static void glslang_optimizeSPIRV(spv_target_env targetEnv, const glslang_CompileRequest_1_1& request, std::vector<SPIRVOptimizationDiagnostic>& outDiags, std::vector<unsigned int>& ioSpirv)
+static void glslang_optimizeSPIRV(spv_target_env targetEnv, const glslang_CompileRequest_1_2& request, std::vector<SPIRVOptimizationDiagnostic>& outDiags, std::vector<unsigned int>& ioSpirv)
 {
     const auto optimizationLevel = request.optimizationLevel;
 
@@ -509,7 +509,7 @@ static spv_target_env _getUniversalTargetEnv(glslang::EShTargetLanguageVersion i
     return SPV_ENV_UNIVERSAL_1_2;
 }
 
-static int glslang_compileGLSLToSPIRV(const glslang_CompileRequest_1_1& request)
+static int glslang_compileGLSLToSPIRV(const glslang_CompileRequest_1_2& request)
 {
     // Check that the encoding matches
     assert(glslang::EShTargetSpv_1_4 == _makeTargetLanguageVersion(1, 4));
@@ -607,6 +607,9 @@ static int glslang_compileGLSLToSPIRV(const glslang_CompileRequest_1_1& request)
             return 1;
         }
 
+        if (request.entryPointName)
+            shader->setEntryPoint(request.entryPointName);
+
         program->addShader(shader);
 
         if (!program->link(messages))
@@ -691,7 +694,7 @@ static int glslang_compileGLSLToSPIRV(const glslang_CompileRequest_1_1& request)
     return 0;
 }
 
-static int glslang_dissassembleSPIRV(const glslang_CompileRequest_1_1& request)
+static int glslang_dissassembleSPIRV(const glslang_CompileRequest_1_2& request)
 {
     typedef unsigned int SPIRVWord;
 
@@ -744,7 +747,7 @@ public:
     bool m_isInitialized = false;
 };
 
-static int _compile(const glslang_CompileRequest_1_1& request)
+static int _compile(const glslang_CompileRequest_1_2& request)
 {
     int result = 0;
     switch (request.action)
@@ -771,7 +774,7 @@ _declspec(dllexport)
 #else
 __attribute__((__visibility__("default")))
 #endif
-int glslang_compile_1_1(glslang_CompileRequest_1_1* inRequest)
+int glslang_compile_1_2(glslang_CompileRequest_1_2 * inRequest)
 {
     static ProcessInitializer g_processInitializer;
     if (!g_processInitializer.init())
@@ -786,7 +789,7 @@ int glslang_compile_1_1(glslang_CompileRequest_1_1* inRequest)
     }
 
     // If it's the right size just use it
-    if (inRequest->sizeInBytes == sizeof(glslang_CompileRequest_1_1))
+    if (inRequest->sizeInBytes == sizeof(glslang_CompileRequest_1_2))
     {
         return _compile(*inRequest);
     }
@@ -796,8 +799,8 @@ int glslang_compile_1_1(glslang_CompileRequest_1_1* inRequest)
 
         // Try to ensure some binary compatibility, by using sizeInBytes member, and copying
 
-        glslang_CompileRequest_1_1 request;
-        
+        glslang_CompileRequest_1_2 request;
+
         // Copy into request
         const size_t copySize = (inRequest->sizeInBytes > sizeof(request)) ? sizeof(request) : inRequest->sizeInBytes;
         ::memcpy(&request, inRequest, copySize);
@@ -814,6 +817,21 @@ _declspec(dllexport)
 #else
 __attribute__((__visibility__("default")))
 #endif
+int glslang_compile_1_1(glslang_CompileRequest_1_1* inRequest)
+{
+    glslang_CompileRequest_1_2 request;
+    memset(&request, 0, sizeof(request));
+    request.sizeInBytes = sizeof(request);
+    request.set(*inRequest);
+    return glslang_compile_1_2(&request);
+}
+
+extern "C"
+#ifdef _MSC_VER
+_declspec(dllexport)
+#else
+__attribute__((__visibility__("default")))
+#endif
 int glslang_compile(glslang_CompileRequest_1_0* inRequest)
 {
     glslang_CompileRequest_1_1 request;
diff --git a/source/slang-glslang/slang-glslang.h b/source/slang-glslang/slang-glslang.h
index 88cab0bf9..1fb76f1f3 100644
--- a/source/slang-glslang/slang-glslang.h
+++ b/source/slang-glslang/slang-glslang.h
@@ -3,6 +3,7 @@
 #define SLANG_GLSLANG_H_INCLUDED
 
 #include <stddef.h>
+#include <memory>
 
 typedef void (*glslang_OutputFunc)(void const* data, size_t size, void* userData);
 
@@ -91,6 +92,41 @@ struct glslang_CompileRequest_1_1
     glsl_SPIRVVersion   spirvVersion;               ///< The SPIR-V version. If all are 0 will use the default which is 1.2 currently
 };
 
+// 1.2 version
+struct glslang_CompileRequest_1_2
+{
+    /// Set from 1.1 
+    void set(const glslang_CompileRequest_1_1& in);
+
+    size_t              sizeInBytes;            ///< Size in bytes of this structure
+
+    // START! Embed the glslang_CompileRequest_1_0 fields
+    char const* sourcePath;
+
+    void const* inputBegin;
+    void const* inputEnd;
+
+    glslang_OutputFunc  diagnosticFunc;
+    void* diagnosticUserData;
+
+    glslang_OutputFunc  outputFunc;
+    void* outputUserData;
+
+    int                 slangStage;
+
+    unsigned            action;
+
+    unsigned            optimizationLevel;
+    unsigned            debugInfoType;
+    // END! Embed the glslang_CompileRequest_1_0 fields
+
+    const char* spirvTargetName;            /// A valid TargetName. If null will use universal based on the spirVersion.
+    glsl_SPIRVVersion   spirvVersion;               ///< The SPIR-V version. If all are 0 will use the default which is 1.2 currently
+
+    // glslang_CompileRequest_1_2 fields
+    const char* entryPointName; // The name of the entrypoint that will appear in output spirv.
+};
+
 void glslang_CompileRequest_1_0::set(const glslang_CompileRequest_1_1& in)
 {
     SLANG_GLSLANG_COMPILE_REQUEST_1_0(SLANG_GLSLANG_FIELD_COPY)
@@ -101,7 +137,14 @@ void glslang_CompileRequest_1_1::set(const glslang_CompileRequest_1_0& in)
     SLANG_GLSLANG_COMPILE_REQUEST_1_0(SLANG_GLSLANG_FIELD_COPY)
 }
 
+void glslang_CompileRequest_1_2::set(const glslang_CompileRequest_1_1& in)
+{
+    memcpy(this, &in, sizeof(in));
+}
+
 typedef int (*glslang_CompileFunc_1_0)(glslang_CompileRequest_1_0* request);
 typedef int (*glslang_CompileFunc_1_1)(glslang_CompileRequest_1_1* request);
+typedef int (*glslang_CompileFunc_1_2)(glslang_CompileRequest_1_2* request);
+
 
 #endif
diff --git a/source/slang/slang-compiler.cpp b/source/slang/slang-compiler.cpp
index 4e1ab8e98..12e96d7bb 100644
--- a/source/slang/slang-compiler.cpp
+++ b/source/slang/slang-compiler.cpp
@@ -1222,11 +1222,18 @@ namespace Slang
                 auto entryPoint = getEntryPoint(entryPointIndex);
                 profile = getEffectiveProfile(entryPoint, targetReq);
 
-                options.entryPointName = allocator.allocate(getText(entryPoint->getName()));
-                auto entryPointNameOverride = getProgram()->getEntryPointNameOverride(entryPointIndex);
-                if (entryPointNameOverride.getLength() != 0)
+                if (getTargetReq()->getHLSLToVulkanLayoutOptions() && !getTargetReq()->getHLSLToVulkanLayoutOptions()->getUseOriginalEntryPointName())
                 {
-                    options.entryPointName = allocator.allocate(entryPointNameOverride);
+                }
+                else
+                {
+
+                    options.entryPointName = allocator.allocate(getText(entryPoint->getName()));
+                    auto entryPointNameOverride = getProgram()->getEntryPointNameOverride(entryPointIndex);
+                    if (entryPointNameOverride.getLength() != 0)
+                    {
+                        options.entryPointName = allocator.allocate(entryPointNameOverride);
+                    }
                 }
             }
             else 
diff --git a/source/slang/slang-hlsl-to-vulkan-layout-options.h b/source/slang/slang-hlsl-to-vulkan-layout-options.h
index 4ae169d90..2d24dfccb 100644
--- a/source/slang/slang-hlsl-to-vulkan-layout-options.h
+++ b/source/slang/slang-hlsl-to-vulkan-layout-options.h
@@ -119,6 +119,8 @@ public:
         /// True if the compiler should invert the Y coordinate of any SV_Position output.
     bool shouldInvertY() const { return m_invertY; }
 
+    bool getUseOriginalEntryPointName() const { return m_useOriginalEntryPointName; }
+
         /// Given an kind and a binding infer the vulkan binding.
         /// Will return an invalid binding if one is not found
     Binding inferBinding(Kind kind, const Binding& inBinding) const;
@@ -145,6 +147,8 @@ public:
 
     void setInvertY(bool value) { m_invertY = value; }
 
+    void setUseOriginalEntryPointName(bool value) { m_useOriginalEntryPointName = value; }
+
         /// Ctor
     HLSLToVulkanLayoutOptions();
     
@@ -171,6 +175,9 @@ protected:
 
         /// Whether to invert the Y coordinate of SV_Position output.
     bool m_invertY = false;
+
+        /// If set, will use the original entry point name in the generated SPIRV instead of "main".
+    bool m_useOriginalEntryPointName = false;
 };
 
 } // namespace Slang
diff --git a/source/slang/slang-options.cpp b/source/slang/slang-options.cpp
index e59109b37..b31a502e7 100644
--- a/source/slang/slang-options.cpp
+++ b/source/slang/slang-options.cpp
@@ -93,6 +93,8 @@ enum class OptionKind
     VulkanBindShift,
     VulkanBindGlobals,
     VulkanInvertY,
+    VulkanUseEntryPointName,
+
 
     GLSLForceScalarLayout,
     EnableEffectAnnotations,
@@ -498,6 +500,7 @@ void initCommandOptions(CommandOptions& options)
         { OptionKind::VulkanBindGlobals, "-fvk-bind-globals", "-fvk-bind-globals <N> <descriptor-set>",
         "Places the $Globals cbuffer at descriptor set <descriptor-set> and binding <N>."},
         { OptionKind::VulkanInvertY, "-fvk-invert-y", nullptr, "Negates (additively inverts) SV_Position.y before writing to stage output."},
+        { OptionKind::VulkanUseEntryPointName, "-fvk-use-entrypoint-name", nullptr, "Uses the entrypoint name from the source instead of 'main' in the spirv output."},
         { OptionKind::EnableEffectAnnotations,
          "-enable-effect-annotations", nullptr, 
          "Enables support for legacy effect annotation syntax."},
@@ -2016,6 +2019,12 @@ SlangResult OptionsParser::_parse(
                 m_hlslToVulkanLayoutOptions->setInvertY(true);
                 break;
             }
+            case OptionKind::VulkanUseEntryPointName:
+            {
+                // -fvk-use-entrypoint-name
+                m_hlslToVulkanLayoutOptions->setUseOriginalEntryPointName(true);
+                break;
+            }
             case OptionKind::Profile: SLANG_RETURN_ON_FAIL(_parseProfile(arg)); break;
             case OptionKind::Capability:
             {