Support mixture of precompiled and non-precompiled modules (#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 #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

1 files changed, 135 insertions, 3 deletions

diff --git a/source/slang/slang-compiler-tu.cpp b/source/slang/slang-compiler-tu.cpp
index d4f3976a6..fe778148b 100644
--- a/source/slang/slang-compiler-tu.cpp
+++ b/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())