From 02706dfc5f0526f4f8ca337be16d7b00640ba168 Mon Sep 17 00:00:00 2001
From: cheneym2 <acheney@nvidia.com>
Date: Wed, 26 Feb 2025 21:20:29 -0500
Subject: 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>
---
 tools/gfx/renderer-shared.cpp | 114 ++++++++++++++++++++++++++++++++++--------
 1 file changed, 94 insertions(+), 20 deletions(-)

(limited to 'tools/gfx/renderer-shared.cpp')

diff --git a/tools/gfx/renderer-shared.cpp b/tools/gfx/renderer-shared.cpp
index 9f5574dcc..ae4ea3eef 100644
--- a/tools/gfx/renderer-shared.cpp
+++ b/tools/gfx/renderer-shared.cpp
@@ -1103,32 +1103,106 @@ void ShaderProgramBase::init(const IShaderProgram::Desc& inDesc)
 
 Result ShaderProgramBase::compileShaders(RendererBase* device)
 {
+    auto compileTarget = device->slangContext.compileTarget;
     // For a fully specialized program, read and store its kernel code in `shaderProgram`.
     auto compileShader = [&](slang::EntryPointReflection* entryPointInfo,
                              slang::IComponentType* entryPointComponent,
                              SlangInt entryPointIndex)
     {
         auto stage = entryPointInfo->getStage();
-        ComPtr<ISlangBlob> kernelCode;
-        ComPtr<ISlangBlob> diagnostics;
-        auto compileResult = device->getEntryPointCodeFromShaderCache(
-            entryPointComponent,
-            entryPointIndex,
-            0,
-            kernelCode.writeRef(),
-            diagnostics.writeRef());
-        if (diagnostics)
+        List<ComPtr<ISlangBlob>> kernelCodes;
         {
-            DebugMessageType msgType = DebugMessageType::Warning;
-            if (compileResult != SLANG_OK)
-                msgType = DebugMessageType::Error;
-            getDebugCallback()->handleMessage(
-                msgType,
-                DebugMessageSource::Slang,
-                (char*)diagnostics->getBufferPointer());
+            ComPtr<ISlangBlob> downstreamIR;
+            ComPtr<ISlangBlob> diagnostics;
+            auto compileResult = device->getEntryPointCodeFromShaderCache(
+                entryPointComponent,
+                entryPointIndex,
+                0,
+                downstreamIR.writeRef(),
+                diagnostics.writeRef());
+            if (diagnostics)
+            {
+                DebugMessageType msgType = DebugMessageType::Warning;
+                if (compileResult != SLANG_OK)
+                    msgType = DebugMessageType::Error;
+                getDebugCallback()->handleMessage(
+                    msgType,
+                    DebugMessageSource::Slang,
+                    (char*)diagnostics->getBufferPointer());
+            }
+            kernelCodes.add(downstreamIR);
         }
-        SLANG_RETURN_ON_FAIL(compileResult);
-        SLANG_RETURN_ON_FAIL(createShaderModule(entryPointInfo, kernelCode));
+
+        // If target precompilation was used, kernelCode may only represent the
+        // glue code holding together the bits of precompiled target IR.
+        // Collect those dependency target IRs too.
+        ComPtr<slang::IModulePrecompileService_Experimental> componentPrecompileService;
+        if (entryPointComponent->queryInterface(
+                slang::IModulePrecompileService_Experimental::getTypeGuid(),
+                (void**)componentPrecompileService.writeRef()) == SLANG_OK)
+        {
+            SlangInt dependencyCount = componentPrecompileService->getModuleDependencyCount();
+            if (dependencyCount > 0)
+            {
+                for (int dependencyIndex = 0; dependencyIndex < dependencyCount; dependencyIndex++)
+                {
+                    ComPtr<slang::IModule> dependencyModule;
+                    {
+                        ComPtr<slang::IBlob> diagnosticsBlob;
+                        auto result = componentPrecompileService->getModuleDependency(
+                            dependencyIndex,
+                            dependencyModule.writeRef(),
+                            diagnosticsBlob.writeRef());
+                        if (diagnosticsBlob)
+                        {
+                            DebugMessageType msgType = DebugMessageType::Warning;
+                            if (result != SLANG_OK)
+                                msgType = DebugMessageType::Error;
+                            getDebugCallback()->handleMessage(
+                                msgType,
+                                DebugMessageSource::Slang,
+                                (char*)diagnosticsBlob->getBufferPointer());
+                        }
+                        SLANG_RETURN_ON_FAIL(result);
+                    }
+
+                    ComPtr<slang::IBlob> downstreamIR;
+                    {
+                        ComPtr<slang::IBlob> diagnosticsBlob;
+                        SlangResult result = SLANG_OK;
+                        ComPtr<slang::IModulePrecompileService_Experimental> precompileService;
+                        result = dependencyModule->queryInterface(
+                            slang::IModulePrecompileService_Experimental::getTypeGuid(),
+                            (void**)precompileService.writeRef());
+                        if (result == SLANG_OK)
+                        {
+                            ComPtr<slang::IBlob> diagnosticsBlob;
+                            auto result = precompileService->getPrecompiledTargetCode(
+                                compileTarget,
+                                downstreamIR.writeRef(),
+                                diagnosticsBlob.writeRef());
+                            if (result == SLANG_OK)
+                            {
+                                kernelCodes.add(downstreamIR);
+                            }
+                            if (diagnosticsBlob)
+                            {
+                                DebugMessageType msgType = DebugMessageType::Warning;
+                                if (result != SLANG_OK)
+                                    msgType = DebugMessageType::Error;
+                                getDebugCallback()->handleMessage(
+                                    msgType,
+                                    DebugMessageSource::Slang,
+                                    (char*)diagnosticsBlob->getBufferPointer());
+                            }
+                        }
+                        SLANG_RETURN_ON_FAIL(result);
+                    }
+                }
+            }
+        }
+
+        SLANG_RETURN_ON_FAIL(createShaderModule(entryPointInfo, kernelCodes));
         return SLANG_OK;
     };
 
@@ -1160,10 +1234,10 @@ Result ShaderProgramBase::compileShaders(RendererBase* device)
 
 Result ShaderProgramBase::createShaderModule(
     slang::EntryPointReflection* entryPointInfo,
-    ComPtr<ISlangBlob> kernelCode)
+    List<ComPtr<ISlangBlob>>& kernelCodes)
 {
     SLANG_UNUSED(entryPointInfo);
-    SLANG_UNUSED(kernelCode);
+    SLANG_UNUSED(kernelCodes);
     return SLANG_OK;
 }
 
-- 
cgit v1.2.3