Fix precompiledTargetModule tests (#6455)

* Fix precompiledTargetModule tests

Add SPIRV-Tool linker support to gfx unit tests
and use the linker in precompileModule tests
that use precompiled modules to reconstitute
SPIRV shaders that were modularly compiled.

Fix a Slang reference count bug in the
precompile service.

* Use sm_6_6

New DXC requires higher version
for linkability.

* Rename helper function, pass by reference

* Link through slang-glslang

* Add missing files

* Fix metal

* format code

---------

Co-authored-by: slangbot <186143334+slangbot@users.noreply.github.com>
Co-authored-by: Yong He <yonghe@outlook.com>

6 files changed, 159 insertions, 4 deletions

diff --git a/tools/gfx/vulkan/glslang-module.cpp b/tools/gfx/vulkan/glslang-module.cpp
new file mode 100644
index 000000000..03afef794
--- /dev/null
+++ b/tools/gfx/vulkan/glslang-module.cpp
@@ -0,0 +1,108 @@
+// glslang-module.cpp
+#include "glslang-module.h"
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#if SLANG_WINDOWS_FAMILY
+#include <windows.h>
+#else
+#include <dlfcn.h>
+#endif
+
+#include "../renderer-shared.h"
+
+namespace gfx
+{
+using namespace Slang;
+
+// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! GlslangModule
+// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
+Slang::Result GlslangModule::init()
+{
+    if (isInitialized())
+    {
+        destroy();
+    }
+
+    const char* dynamicLibraryName = "Unknown";
+
+#if SLANG_WINDOWS_FAMILY
+    dynamicLibraryName = "slang-glslang.dll";
+    HMODULE module = ::LoadLibraryA(dynamicLibraryName);
+    m_module = (void*)module;
+#elif SLANG_APPLE_FAMILY
+    dynamicLibraryName = "libslang_glslang.dylib";
+    m_module = dlopen(dynamicLibraryName, RTLD_NOW | RTLD_GLOBAL);
+#else
+    dynamicLibraryName = "libslang_glslang.so";
+    m_module = dlopen(dynamicLibraryName, RTLD_NOW);
+#endif
+
+    if (!m_module)
+    {
+        return SLANG_FAIL;
+    }
+
+    // Load functions
+#if SLANG_WINDOWS_FAMILY
+    m_linkSPIRVFunc = (glslang_LinkSPIRVFunc)GetProcAddress((HMODULE)m_module, "glslang_linkSPIRV");
+#else
+    m_linkSPIRVFunc = (glslang_LinkSPIRVFunc)dlsym(m_module, "glslang_linkSPIRV");
+#endif
+    if (!m_linkSPIRVFunc)
+    {
+        return SLANG_FAIL;
+    }
+
+    return SLANG_OK;
+}
+
+void GlslangModule::destroy()
+{
+    if (!isInitialized())
+    {
+        return;
+    }
+
+#if SLANG_WINDOWS_FAMILY
+    ::FreeLibrary((HMODULE)m_module);
+#else
+    dlclose(m_module);
+#endif
+    m_module = nullptr;
+}
+
+ComPtr<ISlangBlob> GlslangModule::linkSPIRV(List<ComPtr<ISlangBlob>> spirvModules)
+{
+
+    if (!m_linkSPIRVFunc)
+    {
+        return nullptr;
+    }
+
+    glslang_LinkRequest request = {};
+
+    std::vector<const uint32_t*> moduleCodePtrs(spirvModules.getCount());
+    std::vector<uint32_t> moduleSizes(spirvModules.getCount());
+    for (Index i = 0; i < spirvModules.getCount(); ++i)
+    {
+        moduleCodePtrs[i] = (const uint32_t*)spirvModules[i]->getBufferPointer();
+        moduleSizes[i] = spirvModules[i]->getBufferSize() / sizeof(uint32_t);
+        SLANG_ASSERT(spirvModules[i]->getBufferSize() % sizeof(uint32_t) == 0);
+    }
+    request.modules = moduleCodePtrs.data();
+    request.moduleSizes = moduleSizes.data();
+    request.moduleCount = spirvModules.getCount();
+    request.linkResult = nullptr;
+
+    m_linkSPIRVFunc(&request);
+
+    ComPtr<ISlangBlob> linkedSPIRV;
+    linkedSPIRV = RawBlob::create(request.linkResult, request.linkResultSize * sizeof(uint32_t));
+    return linkedSPIRV;
+}
+
+} // namespace gfx
diff --git a/tools/gfx/vulkan/glslang-module.h b/tools/gfx/vulkan/glslang-module.h
new file mode 100644
index 000000000..cda903577
--- /dev/null
+++ b/tools/gfx/vulkan/glslang-module.h
@@ -0,0 +1,36 @@
+// glslang-module.h
+#pragma once
+
+#include "core/slang-list.h"
+#include "external/spirv-tools/include/spirv-tools/linker.hpp"
+#include "slang-com-helper.h"
+#include "slang-com-ptr.h"
+#include "slang-glslang/slang-glslang.h"
+#include "slang.h"
+
+namespace gfx
+{
+
+struct GlslangModule
+{
+    /// true if has been initialized
+    SLANG_FORCE_INLINE bool isInitialized() const { return m_module != nullptr; }
+
+    /// Initialize
+    Slang::Result init();
+
+    /// Destroy
+    void destroy();
+
+    /// Dtor
+    ~GlslangModule() { destroy(); }
+
+    Slang::ComPtr<ISlangBlob> linkSPIRV(Slang::List<Slang::ComPtr<ISlangBlob>> spirvModules);
+
+protected:
+    void* m_module = nullptr;
+
+    glslang_LinkSPIRVFunc m_linkSPIRVFunc = nullptr;
+};
+
+} // namespace gfx
diff --git a/tools/gfx/vulkan/vk-device.cpp b/tools/gfx/vulkan/vk-device.cpp
index cffa094ae..981b5abac 100644
--- a/tools/gfx/vulkan/vk-device.cpp
+++ b/tools/gfx/vulkan/vk-device.cpp
@@ -1019,6 +1019,7 @@ SlangResult DeviceImpl::initialize(const Desc& desc)
     SLANG_RETURN_ON_FAIL(RendererBase::initialize(desc));
     SlangResult initDeviceResult = SLANG_OK;
 
+    m_glslang.init();
     for (int forceSoftware = 0; forceSoftware <= 1; forceSoftware++)
     {
         initDeviceResult = m_module.init(forceSoftware != 0);
diff --git a/tools/gfx/vulkan/vk-device.h b/tools/gfx/vulkan/vk-device.h
index 3b6c83103..27cc4c3ce 100644
--- a/tools/gfx/vulkan/vk-device.h
+++ b/tools/gfx/vulkan/vk-device.h
@@ -1,6 +1,7 @@
 // vk-device.h
 #pragma once
 
+#include "glslang-module.h"
 #include "vk-base.h"
 #include "vk-framebuffer.h"
 
@@ -196,6 +197,7 @@ public:
 
     VulkanModule m_module;
     VulkanApi m_api;
+    GlslangModule m_glslang;
 
     VulkanDeviceQueue m_deviceQueue;
     uint32_t m_queueFamilyIndex;
diff --git a/tools/gfx/vulkan/vk-shader-program.cpp b/tools/gfx/vulkan/vk-shader-program.cpp
index 43a295786..1627c95a7 100644
--- a/tools/gfx/vulkan/vk-shader-program.cpp
+++ b/tools/gfx/vulkan/vk-shader-program.cpp
@@ -1,6 +1,7 @@
 // vk-shader-program.cpp
 #include "vk-shader-program.h"
 
+#include "external/spirv-tools/include/spirv-tools/linker.hpp"
 #include "vk-device.h"
 #include "vk-util.h"
 
@@ -71,15 +72,22 @@ VkPipelineShaderStageCreateInfo ShaderProgramImpl::compileEntryPoint(
 
 Result ShaderProgramImpl::createShaderModule(
     slang::EntryPointReflection* entryPointInfo,
-    ComPtr<ISlangBlob> kernelCode)
+    List<ComPtr<ISlangBlob>>& kernelCodes)
 {
-    m_codeBlobs.add(kernelCode);
+    ComPtr<ISlangBlob> linkedKernel = m_device->m_glslang.linkSPIRV(kernelCodes);
+    if (!linkedKernel)
+    {
+        return SLANG_FAIL;
+    }
+
+    m_codeBlobs.add(linkedKernel);
+
     VkShaderModule shaderModule;
     auto realEntryPointName = entryPointInfo->getNameOverride();
     const char* spirvBinaryEntryPointName = "main";
     m_stageCreateInfos.add(compileEntryPoint(
         spirvBinaryEntryPointName,
-        kernelCode,
+        linkedKernel,
         (VkShaderStageFlagBits)VulkanUtil::getShaderStage(entryPointInfo->getStage()),
         shaderModule));
     m_entryPointNames.add(realEntryPointName);
diff --git a/tools/gfx/vulkan/vk-shader-program.h b/tools/gfx/vulkan/vk-shader-program.h
index 3fd56669a..56b354612 100644
--- a/tools/gfx/vulkan/vk-shader-program.h
+++ b/tools/gfx/vulkan/vk-shader-program.h
@@ -37,7 +37,7 @@ public:
 
     virtual Result createShaderModule(
         slang::EntryPointReflection* entryPointInfo,
-        ComPtr<ISlangBlob> kernelCode) override;
+        List<ComPtr<ISlangBlob>>& kernelCodes) override;
 };