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>

4 files changed, 77 insertions, 2 deletions

diff --git a/source/slang-glslang/CMakeLists.txt b/source/slang-glslang/CMakeLists.txt
index c457ad472..e1ced542d 100644
--- a/source/slang-glslang/CMakeLists.txt
+++ b/source/slang-glslang/CMakeLists.txt
@@ -6,7 +6,7 @@ if(SLANG_ENABLE_SLANG_GLSLANG)
         .
         MODULE
         USE_FEWER_WARNINGS
-        LINK_WITH_PRIVATE glslang SPIRV SPIRV-Tools-opt
+        LINK_WITH_PRIVATE glslang SPIRV SPIRV-Tools-opt SPIRV-Tools-link
         INCLUDE_DIRECTORIES_PRIVATE ${slang_SOURCE_DIR}/include
         INSTALL
         EXPORT_SET_NAME SlangTargets
diff --git a/source/slang-glslang/slang-glslang.cpp b/source/slang-glslang/slang-glslang.cpp
index 8d17afc79..bbb3f6afc 100644
--- a/source/slang-glslang/slang-glslang.cpp
+++ b/source/slang-glslang/slang-glslang.cpp
@@ -6,6 +6,7 @@
 #include "glslang/Public/ShaderLang.h"
 #include "slang.h"
 #include "spirv-tools/libspirv.h"
+#include "spirv-tools/linker.hpp"
 #include "spirv-tools/optimizer.hpp"
 
 #ifdef _WIN32
@@ -979,3 +980,67 @@ extern "C"
     request.set(*inRequest);
     return glslang_compile_1_1(&request);
 }
+
+extern "C"
+#ifdef _MSC_VER
+    _declspec(dllexport)
+#else
+    __attribute__((__visibility__("default")))
+#endif
+        int glslang_linkSPIRV(glslang_LinkRequest* request)
+{
+    if (!request || !request->modules || request->linkResult)
+        return false;
+
+    try
+    {
+        spvtools::Context context(SPV_ENV_UNIVERSAL_1_5);
+        spvtools::LinkerOptions options = {};
+
+        spvtools::MessageConsumer consumer = [](spv_message_level_t level,
+                                                const char* source,
+                                                const spv_position_t& position,
+                                                const char* message)
+        {
+            printf("SPIRV-TOOLS: %s\n", message);
+            printf("SPIRV-TOOLS: %s\n", source);
+            printf("SPIRV-TOOLS: %zu:%zu\n", position.index, position.column);
+        };
+        context.SetMessageConsumer(consumer);
+
+        std::vector<std::vector<uint32_t>> moduleVecs(request->moduleCount);
+        std::vector<const uint32_t*> moduleData(request->moduleCount);
+        std::vector<size_t> moduleSizes(request->moduleCount);
+
+        for (size_t i = 0; i < request->moduleCount; ++i)
+        {
+            moduleData[i] = request->modules[i];
+            moduleSizes[i] = request->moduleSizes[i];
+        }
+
+        std::vector<uint32_t> linkedBinary;
+        spv_result_t success = spvtools::Link(
+            context,
+            moduleData.data(),
+            moduleSizes.data(),
+            request->moduleCount,
+            &linkedBinary,
+            options);
+
+        if (success == SPV_SUCCESS)
+        {
+            request->linkResult = new uint32_t[linkedBinary.size()];
+            memcpy(
+                (void*)request->linkResult,
+                linkedBinary.data(),
+                linkedBinary.size() * sizeof(uint32_t));
+            request->linkResultSize = linkedBinary.size();
+        }
+
+        return success;
+    }
+    catch (...)
+    {
+        return false;
+    }
+}
diff --git a/source/slang-glslang/slang-glslang.h b/source/slang-glslang/slang-glslang.h
index 7b955c5af..8343b4c8e 100644
--- a/source/slang-glslang/slang-glslang.h
+++ b/source/slang-glslang/slang-glslang.h
@@ -152,10 +152,19 @@ inline void glslang_CompileRequest_1_2::set(const glslang_CompileRequest_1_1& in
     memcpy(this, &in, sizeof(in));
 }
 
+typedef struct glslang_LinkRequest_t
+{
+    const uint32_t** modules;    // Input: array of pointers to SPIR-V modules
+    const uint32_t* moduleSizes; // Input: array of sizes of SPIR-V modules in 32-bit words
+    int moduleCount;             // Input: number of modules in the array
+    const uint32_t* linkResult;  // Output: pointer to linked SPIR-V module
+    size_t linkResultSize;       // Output: size of the linked SPIR-V module in 32-bit words
+} glslang_LinkRequest;
+
 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);
 typedef bool (*glslang_ValidateSPIRVFunc)(const uint32_t* contents, int contentsSize);
 typedef bool (*glslang_DisassembleSPIRVFunc)(const uint32_t* contents, int contentsSize);
-
+typedef bool (*glslang_LinkSPIRVFunc)(glslang_LinkRequest* request);
 #endif
diff --git a/source/slang/slang-compiler-tu.cpp b/source/slang/slang-compiler-tu.cpp
index c20fc9a80..8bfe9a9ab 100644
--- a/source/slang/slang-compiler-tu.cpp
+++ b/source/slang/slang-compiler-tu.cpp
@@ -291,6 +291,7 @@ SLANG_NO_THROW SlangResult SLANG_MCALL ComponentType::getModuleDependency(
     {
         return SLANG_E_INVALID_ARG;
     }
+    getModuleDependencies()[dependencyIndex]->addRef();
     *outModule = getModuleDependencies()[dependencyIndex];
     return SLANG_OK;
 }