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())