Support SPIR-V deferred linking option (#6500)

The new option "SkipDownstreamLinking" will defer final downstream IR
linking to the user application. This option only has an effect if
there are modules that were precompiled to the target IR using
precompileForTarget().

Until now, the default behavior for SPIR-V was to use deferred linking, and
the default behavior for DXIL was to use immediate/internal linking in Slang.

This change only affects the SPIR-V behavior such that both deferred and
non-deferred linking is supported based on the new option.

To support the non-deferred option, Slang will internally call into
SPIRV-Tools-link to reconstitute a complete SPIR-V shader program when
necessary (due to modules having been precompiled to target IR).
Otherwise, if SkipDownstreamLinking is enabled, the shader returned by
e.g. getTargetCode() or getEntryPointCode() may have import linkage to
the SPIR-V embedded in the constituent modules.

Closes #4994

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

3 files changed, 76 insertions, 2 deletions

diff --git a/source/slang/slang-compiler.cpp b/source/slang/slang-compiler.cpp
index 58cc55e71..3839e0722 100644
--- a/source/slang/slang-compiler.cpp
+++ b/source/slang/slang-compiler.cpp
@@ -2669,6 +2669,12 @@ bool CodeGenContext::shouldDumpIR()
     return getTargetProgram()->getOptionSet().getBoolOption(CompilerOptionName::DumpIr);
 }
 
+bool CodeGenContext::shouldSkipDownstreamLinking()
+{
+    return getTargetProgram()->getOptionSet().getBoolOption(
+        CompilerOptionName::SkipDownstreamLinking);
+}
+
 bool CodeGenContext::shouldReportCheckpointIntermediates()
 {
     return getTargetProgram()->getOptionSet().getBoolOption(
diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index 10da32400..cfcbe816f 100644
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -1384,7 +1384,8 @@ enum class PassThroughMode : SlangPassThroughIntegral
     LLVM = SLANG_PASS_THROUGH_LLVM,                  ///< LLVM 'compiler'
     SpirvOpt = SLANG_PASS_THROUGH_SPIRV_OPT,         ///< pass thorugh spirv to spirv-opt
     MetalC = SLANG_PASS_THROUGH_METAL,
-    Tint = SLANG_PASS_THROUGH_TINT, ///< pass through spirv to Tint API
+    Tint = SLANG_PASS_THROUGH_TINT,            ///< pass through spirv to Tint API
+    SpirvLink = SLANG_PASS_THROUGH_SPIRV_LINK, ///< pass through spirv to spirv-link
     CountOf = SLANG_PASS_THROUGH_COUNT_OF,
 };
 void printDiagnosticArg(StringBuilder& sb, PassThroughMode val);
@@ -2886,6 +2887,12 @@ public:
     // removed between IR linking and target source generation.
     bool removeAvailableInDownstreamIR = false;
 
+    // Determines if program level compilation like getTargetCode() or getEntryPointCode()
+    // should return a fully linked downstream program or just the glue SPIR-V/DXIL that
+    // imports and uses the precompiled SPIR-V/DXIL from constituent modules.
+    // This is a no-op if modules are not precompiled.
+    bool shouldSkipDownstreamLinking();
+
 protected:
     CodeGenTarget m_targetFormat = CodeGenTarget::Unknown;
     ExtensionTracker* m_extensionTracker = nullptr;
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index ddb4ea67a..94ea66d71 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -2093,10 +2093,71 @@ SlangResult emitSPIRVForEntryPointsDirectly(
     if (compiler)
     {
 #if 0
-        // Dump the unoptimized SPIRV after lowering from slang IR -> SPIRV
+        // Dump the unoptimized/unlinked SPIRV after lowering from slang IR -> SPIRV
         compiler->disassemble((uint32_t*)spirv.getBuffer(), int(spirv.getCount() / 4));
 #endif
 
+        bool isPrecompilation = codeGenContext->getTargetProgram()->getOptionSet().getBoolOption(
+            CompilerOptionName::EmbedDownstreamIR);
+
+        if (!isPrecompilation && !codeGenContext->shouldSkipDownstreamLinking())
+        {
+            ComPtr<IArtifact> linkedArtifact;
+
+            // collect spirv files
+            List<uint32_t*> spirvFiles;
+            List<uint32_t> spirvSizes;
+
+            // Start with the SPIR-V we just generated.
+            // SPIRV-Tools-link expects the size in 32-bit words
+            // whereas the spirv blob size is in bytes.
+            spirvFiles.add((uint32_t*)spirv.getBuffer());
+            spirvSizes.add(int(spirv.getCount()) / 4);
+
+            // Iterate over all modules in the linkedIR. For each module, if it
+            // contains an embedded downstream ir instruction, add it to the list
+            // of spirv files.
+            auto program = codeGenContext->getProgram();
+
+            program->enumerateIRModules(
+                [&](IRModule* irModule)
+                {
+                    for (auto globalInst : irModule->getModuleInst()->getChildren())
+                    {
+                        if (auto inst = as<IREmbeddedDownstreamIR>(globalInst))
+                        {
+                            if (inst->getTarget() == CodeGenTarget::SPIRV)
+                            {
+                                auto slice = inst->getBlob()->getStringSlice();
+                                spirvFiles.add((uint32_t*)slice.begin());
+                                spirvSizes.add(int(slice.getLength()) / 4);
+                            }
+                        }
+                    }
+                });
+
+            SLANG_ASSERT(int(spirv.getCount()) % 4 == 0);
+            SLANG_ASSERT(spirvFiles.getCount() == spirvSizes.getCount());
+
+            if (spirvFiles.getCount() > 1)
+            {
+                SlangResult linkresult = compiler->link(
+                    (const uint32_t**)spirvFiles.getBuffer(),
+                    (const uint32_t*)spirvSizes.getBuffer(),
+                    (uint32_t)spirvFiles.getCount(),
+                    linkedArtifact.writeRef());
+
+                if (linkresult != SLANG_OK)
+                {
+                    return SLANG_FAIL;
+                }
+
+                ComPtr<ISlangBlob> blob;
+                linkedArtifact->loadBlob(ArtifactKeep::No, blob.writeRef());
+                artifact = _Move(linkedArtifact);
+            }
+        }
+
         if (!codeGenContext->shouldSkipSPIRVValidation())
         {
             StringBuilder runSpirvValEnvVar;