Support mixture of precompiled and non-precompiled modules (shader-slang#4860)

* Support mixture of precompiled and non-precompiled modules

This changes the implementation of precompile DXIL modules to
accept combinations of modules with precompiled DXIL, ones without,
and ones with a mixture of precompiled DXIL and Slang IR.

During precompilation, module IR is analyzed to find public functions
which appear to be capable of being compiled as HLSL, and those
functions are given a HLSLExport decoration, ensuring they are emitted
as HLSL and preserved in the precompiled DXIL blob. The IR for those
functions is then tagged with a new decoration AvailableInDXIL, which
marks that their implementation is present in the embedded DXIL blob.
The DXIL blob is attached to the IR as before, inside a EmbeddedDXIL
BlobLit instruction.

The logic that determines whether or not functions should be
precompiled to DXIL is a placeholder at this point, returning true
always. A subsequent change will add selection criteria.

During module linking, the full module IR is available, as well
as the optional EmbeddedDXIL blob. The IR for functions implemented
by the blob are tagged with AvailableInDXIL in the module IR.

After linking the IR for all modules to program level IR, the IR for
the functions marked AvailableInDXIL are deleted from the linked IR,
prior to emitting HLSL and compiling linking the result.

This change also changes the point of time when the module IR is
checked for EmbeddedDXIL blobs. Instead of happening at load time
as before, it happens during immediately before final linking, meaning
that the blob does not need to be independently stored with the module
separate from the IR as was done previously.

Work on shader-slang#4792

* Clean up debug prints

* Call isSimpleHLSLDataType stub

* Address feedback on precompiled dxil support

Allow for IR filtering both before and after linking.
Only mark AvailableInDXIL those functions which pass
both filtering stages. Functions are corrlated using
mangled function names.

Rather than delete functions entirely when linking with
libraries that include precompiled DXIL, instead convert
the IR function definitions to declarations by gutting
them, removing child blocks.

* Use artifact metadata and name list instead of linkedir hack

* Use String instead of UnownedStringSlice

* Update tests

* Renaming

* Minor edits

* Don't fully remove functions post-link

* Unexport before collecting metadata

Author: cheneym2

include/slang.h

@@ -4929,9 +4929,7 @@ namespace slang
             int                     targetIndex,
             bool                    value) = 0;
 
-        virtual SLANG_NO_THROW void SLANG_MCALL setTargetEmbedDXIL(
-            int                     targetIndex,
-            bool                    value) = 0;
+        virtual SLANG_NO_THROW void SLANG_MCALL setEmbedDXIL(bool value) = 0;
 
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetForceDXLayout(int targetIndex, bool value) = 0;
     };

source/compiler-core/slang-artifact-associated-impl.cpp

@@ -308,4 +308,9 @@ Slice<ShaderBindingRange> ArtifactPostEmitMetadata::getUsedBindingRanges()
     return Slice<ShaderBindingRange>(m_usedBindings.getBuffer(), m_usedBindings.getCount()); 
 }
 
+Slice<String> ArtifactPostEmitMetadata::getExportedFunctionMangledNames()
+{
+    return Slice<String>(m_exportedFunctionMangledNames.getBuffer(), m_exportedFunctionMangledNames.getCount());
+}
+
 } // namespace Slang

source/compiler-core/slang-artifact-associated-impl.h

@@ -155,13 +155,15 @@ class ArtifactPostEmitMetadata : public ComBaseObject, public IArtifactPostEmitM
 
     // IArtifactPostEmitMetadata
     SLANG_NO_THROW virtual Slice<ShaderBindingRange> SLANG_MCALL getUsedBindingRanges() SLANG_OVERRIDE;
-    
+    SLANG_NO_THROW virtual Slice<String> SLANG_MCALL getExportedFunctionMangledNames() SLANG_OVERRIDE;
+
     void* getInterface(const Guid& uuid);
     void* getObject(const Guid& uuid);
 
     static ComPtr<IArtifactPostEmitMetadata> create() { return ComPtr<IArtifactPostEmitMetadata>(new ThisType); }
 
     List<ShaderBindingRange> m_usedBindings;
+    List<String> m_exportedFunctionMangledNames;
 };
 
 } // namespace Slang

source/compiler-core/slang-artifact-associated.h

@@ -124,6 +124,9 @@ class IArtifactPostEmitMetadata : public ICastable
 
         /// Get the binding ranges
     SLANG_NO_THROW virtual Slice<ShaderBindingRange> SLANG_MCALL getUsedBindingRanges() = 0;
+
+        /// Get the list of functions that were exported in the linked IR
+    SLANG_NO_THROW virtual Slice<String> SLANG_MCALL getExportedFunctionMangledNames() = 0;
 };
 
 } // namespace Slang

source/slang/slang-compiler-tu.cpp

@@ -4,10 +4,91 @@
 #include "../core/slang-basic.h"
 #include "slang-compiler.h"
 #include "slang-ir-insts.h"
+#include "slang-ir-util.h"
 #include "slang-capability.h"
 
 namespace Slang
 {
+    // Only attempt to precompile functions:
+    // 1) With function bodies (not just empty decls)
+    // 2) Not marked with unsafeForceInlineDecoration
+    // 3) Have a simple HLSL data type as the return or parameter type
+    static bool attemptPrecompiledExport(IRInst* inst)
+    {
+        if (inst->getOp() != kIROp_Func)
+        {
+            return false;
+        }
+
+        // Skip functions with no body
+        bool hasBody = false;
+        for (auto child : inst->getChildren())
+        {
+            if (child->getOp() == kIROp_Block)
+            {
+                hasBody = true;
+                break;
+            }
+        }
+        if (!hasBody)
+        {
+            return false;
+        }
+
+        // Skip functions marked with unsafeForceInlineDecoration
+        if (inst->findDecoration<IRUnsafeForceInlineEarlyDecoration>())
+        {
+            return false;
+        }
+
+        // Skip non-simple HLSL data types, filters out generics
+        if (!isSimpleHLSLDataType(inst))
+        {
+            return false;
+        }
+
+        return true;
+    }
+
+    /*
+    * Precompile the module for the given target.
+    *
+    * This function creates a target program and emits the precompiled blob as
+    * an embedded blob in the module IR, e.g. DXIL.
+    * Because the IR for the Slang Module may violate the restrictions of the
+    * target language, the emitted target blob may not be able to include the
+    * full module, but rather only the subset that can be precompiled. For
+    * example, DXIL libraries do not allow resources like structured buffers
+    * to appear in the library interface. Also, no target languages allow
+    * generics to be precompiled.
+    *
+    * Some restrictions can be enforced up front before linking, but some are
+    * done during target generation in between IR linking+legalization and
+    * target source emission.
+    *
+    * Functions which can be rejected up front:
+    * - Functions with no body
+    * - Functions marked with unsafeForceInlineDecoration
+    * - Functions that define or use generics
+    *
+    * The functions not rejected up front are marked with
+    * DownstreamModuleExportDecoration which indicates functions we're trying to
+    * export for precompilation, and this also helps to identify the functions
+    * in the linked IR which survived the additional pruning.
+    *
+    * Functions that are rejected after linking+legalization (inside
+    * emitPrecompiled*):
+    * - (DXIL) Functions that return or take a HLSLStructuredBufferType
+    * - (DXIL) Functions that return or take a Matrix type
+    *
+    * emitPrecompiled* produces the output artifact containing target language
+    * blob, and as metadata, the list of functions which survived the second
+    * phase of filtering.
+    *
+    * The original module IR functions matching those are then marked with
+    * "AvailableIn*" (e.g. AvailableInDXILDecoration) to indicate to future
+    * module users which functions are present in the precompiled blob.
+    */
     SLANG_NO_THROW SlangResult SLANG_MCALL Module::precompileForTarget(
         SlangCompileTarget target,
         slang::IBlob** outDiagnostics)
@@ -20,6 +101,7 @@ namespace Slang
 
         auto module = getIRModule();
         auto linkage = getLinkage();
+        auto builder = IRBuilder(module);
 
         DiagnosticSink sink(linkage->getSourceManager(), Lexer::sourceLocationLexer);
         applySettingsToDiagnosticSink(&sink, &sink, linkage->m_optionSet);
@@ -48,6 +130,7 @@ namespace Slang
         {
         case CodeGenTarget::DXIL:
             tp.getOptionSet().add(CompilerOptionName::Profile, Profile::RawEnum::DX_Lib_6_6);
+            tp.getOptionSet().add(CompilerOptionName::EmbedDXIL, true);
             break;
         }
 
@@ -59,20 +142,69 @@ namespace Slang
         CodeGenContext::Shared sharedCodeGenContext(&tp, entryPointIndices, &sink, nullptr);
         CodeGenContext codeGenContext(&sharedCodeGenContext);
 
+        // Mark all public functions as exported, ensure there's at least one. Store a mapping
+        // of function name to IRInst* for later reference. After linking is done, we'll scan
+        // the linked result to see which functions survived the pruning and are included in the
+        // precompiled blob.
+        Dictionary<String, IRInst*> nameToFunction;
+        bool hasAtLeastOneFunction = false;
+        for (auto inst : module->getGlobalInsts())
+        {
+            if (attemptPrecompiledExport(inst))
+            {
+                hasAtLeastOneFunction = true;
+                builder.addDecoration(inst, kIROp_DownstreamModuleExportDecoration);
+                nameToFunction[inst->findDecoration<IRExportDecoration>()->getMangledName()] = inst;
+            }
+        }
+
+        // Bail if there are no functions to export. That's not treated as an error
+        // because it's possible that the module just doesn't have any simple HLSL.
+        if (!hasAtLeastOneFunction)
+        {
+            return SLANG_OK;
+        }
+
         ComPtr<IArtifact> outArtifact;
-        SlangResult res = codeGenContext.emitTranslationUnit(outArtifact);
+        SlangResult res = codeGenContext.emitPrecompiledDXIL(outArtifact);
 
         sink.getBlobIfNeeded(outDiagnostics);
-
         if (res != SLANG_OK)
         {
             return res;
         }
 
+        auto metadata = findAssociatedRepresentation<IArtifactPostEmitMetadata>(outArtifact);
+        if (!metadata)
+        {
+            return SLANG_E_NOT_AVAILABLE;
+        }
+
+        for (const auto& mangledName : metadata->getExportedFunctionMangledNames())
+        {
+            auto moduleInst = nameToFunction[mangledName];
+            builder.addDecoration(moduleInst, kIROp_AvailableInDXILDecoration);
+            auto moduleDec = moduleInst->findDecoration<IRDownstreamModuleExportDecoration>();
+            moduleDec->removeAndDeallocate();
+        }
+
+        // Finally, clean up the transient export decorations left over in the module. These are
+        // represent functions that were pruned from the IR after linking, before target generation.
+        for (auto moduleInst : module->getGlobalInsts())
+        {
+            if (moduleInst->getOp() == kIROp_Func)
+            {
+                if (auto dec = moduleInst->findDecoration<IRDownstreamModuleExportDecoration>())
+                {
+                    dec->removeAndDeallocate();
+                }
+            }
+        }
+
         ISlangBlob* blob;
         outArtifact->loadBlob(ArtifactKeep::Yes, &blob);
 
-        auto builder = IRBuilder(module);
+        // Add the precompiled blob to the module
         builder.setInsertInto(module);
 
         switch (targetReq->getTarget())

source/slang/slang-compiler.cpp

@@ -350,7 +350,7 @@ namespace Slang
     Profile Profile::lookUp(UnownedStringSlice const& name)
     {
         #define PROFILE(TAG, NAME, STAGE, VERSION)	if(name == UnownedTerminatedStringSlice(#NAME)) return Profile::TAG;
-        #define PROFILE_ALIAS(TAG, DEF, NAME)		if(name == UnownedTerminatedStringSlice(#NAME)) return Profile::TAG;
+        #define PROFILE_ALIAS(TAG, DEF, NAME)     	if(name == UnownedTerminatedStringSlice(#NAME)) return Profile::TAG;
         #include "slang-profile-defs.h"
 
         return Profile::Unknown;
@@ -767,7 +767,7 @@ namespace Slang
 #   pragma warning(pop)
 #endif
 
-    SlangResult CodeGenContext::emitTranslationUnit(ComPtr<IArtifact>& outArtifact)
+    SlangResult CodeGenContext::emitPrecompiledDXIL(ComPtr<IArtifact>& outArtifact)
     {
         return emitWithDownstreamForEntryPoints(outArtifact);
     }
@@ -1094,23 +1094,6 @@ namespace Slang
         return SLANG_OK;
     }
 
-    bool CodeGenContext::isPrecompiled()
-    {
-        auto program = getProgram();
-
-        bool allPrecompiled = true;
-        program->enumerateIRModules([&](IRModule* irModule)
-            {
-                // TODO: Conditionalize this on target
-                if (!irModule->precompiledDXIL)
-                {
-                    allPrecompiled = false;
-                }
-            });
-
-        return allPrecompiled;
-    }
-
     SlangResult CodeGenContext::emitWithDownstreamForEntryPoints(ComPtr<IArtifact>& outArtifact)
     {
         outArtifact.setNull();
@@ -1272,15 +1255,17 @@ namespace Slang
         }
         else
         {
-            if (!isPrecompiled())
+            CodeGenContext sourceCodeGenContext(this, sourceTarget, extensionTracker);
+
+            if (target == CodeGenTarget::DXILAssembly || target == CodeGenTarget::DXIL)
             {
-                CodeGenContext sourceCodeGenContext(this, sourceTarget, extensionTracker);
+                sourceCodeGenContext.removeAvailableInDXIL = true;
+            }
 
-                SLANG_RETURN_ON_FAIL(sourceCodeGenContext.emitEntryPointsSource(sourceArtifact));
-                sourceCodeGenContext.maybeDumpIntermediate(sourceArtifact);
+            SLANG_RETURN_ON_FAIL(sourceCodeGenContext.emitEntryPointsSource(sourceArtifact));
+            sourceCodeGenContext.maybeDumpIntermediate(sourceArtifact);
 
-                sourceLanguage = (SourceLanguage)TypeConvertUtil::getSourceLanguageFromTarget((SlangCompileTarget)sourceTarget);
-            }
+            sourceLanguage = (SourceLanguage)TypeConvertUtil::getSourceLanguageFromTarget((SlangCompileTarget)sourceTarget);
         }
 
         if (sourceArtifact)
@@ -1571,24 +1556,29 @@ namespace Slang
             libraries.addRange(linkage->m_libModules.getBuffer(), linkage->m_libModules.getCount());
         }
 
-        if (isPrecompiled())
+        auto program = getProgram();
+
+        // Load embedded precompiled libraries from IR into library artifacts
+        program->enumerateIRModules([&](IRModule* irModule)
         {
-            auto program = getProgram();
-            program->enumerateIRModules([&](IRModule* irModule)
+            for (auto inst : irModule->getModuleInst()->getChildren())
+            {
+                if (target == CodeGenTarget::DXILAssembly || target == CodeGenTarget::DXIL)
                 {
-                    // TODO: conditionalize on target
-                    if (irModule->precompiledDXIL)
+                    if (inst->getOp() == kIROp_EmbeddedDXIL)
                     {
+                        auto slice = static_cast<IRBlobLit*>(inst->getOperand(0))->getStringSlice();
                         ArtifactDesc desc = ArtifactDescUtil::makeDescForCompileTarget(SLANG_DXIL);
                         desc.kind = ArtifactKind::Library;
 
                         auto library = ArtifactUtil::createArtifact(desc);
 
-                        library->addRepresentationUnknown(irModule->precompiledDXIL);
+                        library->addRepresentationUnknown(StringBlob::create(slice));
                         libraries.add(library);
                     }
-                });
-        }
+                }
+            }
+        });
 
         options.compilerSpecificArguments = allocator.allocate(compilerSpecificArguments);
         options.requiredCapabilityVersions = SliceUtil::asSlice(requiredCapabilityVersions);

source/slang/slang-compiler.h

@@ -555,7 +555,6 @@ namespace Slang
             /// and parsing via Slang reflection, and is not recommended for future APIs to use.
             ///
         Scope* _getOrCreateScopeForLegacyLookup(ASTBuilder* astBuilder);
-
     protected:
         ComponentType(Linkage* linkage);
 
@@ -2732,10 +2731,14 @@ namespace Slang
 
         SlangResult emitEntryPoints(ComPtr<IArtifact>& outArtifact);
 
-        SlangResult emitTranslationUnit(ComPtr<IArtifact>& outArtifact);
+        SlangResult emitPrecompiledDXIL(ComPtr<IArtifact>& outArtifact);
 
         void maybeDumpIntermediate(IArtifact* artifact);
 
+        // Used to cause instructions available in precompiled DXIL to be
+        // removed between IR linking and target source generation.
+        bool removeAvailableInDXIL = false;
+
     protected:
         CodeGenTarget m_targetFormat = CodeGenTarget::Unknown;
         ExtensionTracker* m_extensionTracker = nullptr;
@@ -2772,11 +2775,6 @@ namespace Slang
 
 
         SlangResult _emitEntryPoints(ComPtr<IArtifact>& outArtifact);
-
-	/* Checks if all modules in the target program are already compiled to the
-        target language, indicating that a pass-through linking using the
-        downstream compiler is viable.*/
-        bool isPrecompiled();
     private:
         Shared* m_shared = nullptr;
     };
@@ -2816,7 +2814,7 @@ namespace Slang
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetForceGLSLScalarBufferLayout(int targetIndex, bool value) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetForceDXLayout(int targetIndex, bool value) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetGenerateWholeProgram(int targetIndex, bool value) SLANG_OVERRIDE;
-        virtual SLANG_NO_THROW void SLANG_MCALL setTargetEmbedDXIL(int targetIndex, bool value) SLANG_OVERRIDE;
+        virtual SLANG_NO_THROW void SLANG_MCALL setEmbedDXIL(bool value) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setMatrixLayoutMode(SlangMatrixLayoutMode mode) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setDebugInfoLevel(SlangDebugInfoLevel level) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setOptimizationLevel(SlangOptimizationLevel level) SLANG_OVERRIDE;