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

19 files changed, 289 insertions, 97 deletions

diff --git a/source/compiler-core/slang-artifact-associated-impl.cpp b/source/compiler-core/slang-artifact-associated-impl.cpp
index f6f64e294..f29c0f596 100644
--- a/source/compiler-core/slang-artifact-associated-impl.cpp
+++ b/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
diff --git a/source/compiler-core/slang-artifact-associated-impl.h b/source/compiler-core/slang-artifact-associated-impl.h
index 5144fe7a7..a6e323b0a 100644
--- a/source/compiler-core/slang-artifact-associated-impl.h
+++ b/source/compiler-core/slang-artifact-associated-impl.h
@@ -155,13 +155,15 @@ public:
     
     // 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
diff --git a/source/compiler-core/slang-artifact-associated.h b/source/compiler-core/slang-artifact-associated.h
index f63b8a673..766494271 100644
--- a/source/compiler-core/slang-artifact-associated.h
+++ b/source/compiler-core/slang-artifact-associated.h
@@ -124,6 +124,9 @@ public:
 
         /// 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
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())
diff --git a/source/slang/slang-compiler.cpp b/source/slang/slang-compiler.cpp
index 428532658..a5a09204b 100644
--- a/source/slang/slang-compiler.cpp
+++ b/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);
diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index 9d796f48d..54f43b382 100755
--- a/source/slang/slang-compiler.h
+++ b/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;
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 103cd15ab..caa8ca8ea 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -58,6 +58,7 @@
 #include "slang-ir-metadata.h"
 #include "slang-ir-optix-entry-point-uniforms.h"
 #include "slang-ir-pytorch-cpp-binding.h"
+#include "slang-ir-redundancy-removal.h"
 #include "slang-ir-restructure.h"
 #include "slang-ir-restructure-scoping.h"
 #include "slang-ir-sccp.h"
@@ -415,6 +416,37 @@ bool checkStaticAssert(IRInst* inst, DiagnosticSink* sink)
     return false;
 }
 
+static void unexportNonEmbeddableDXIL(IRModule* irModule)
+{
+    for (auto inst : irModule->getGlobalInsts())
+    {
+        if (inst->getOp() == kIROp_Func)
+        {
+            // DXIL does not permit HLSLStructureBufferType in exported functions
+            // or sadly Matrices (https://github.com/shader-slang/slang/issues/4880)
+            auto type = as<IRFuncType>(inst->getFullType());
+            auto argCount = type->getOperandCount();
+            for (UInt aa = 0; aa < argCount; ++aa)
+            {
+                auto operand = type->getOperand(aa);
+                if (operand->getOp() == kIROp_HLSLStructuredBufferType ||
+                    operand->getOp() == kIROp_MatrixType)
+                {
+                    if (auto dec = inst->findDecoration<IRPublicDecoration>())
+                    {
+                        dec->removeAndDeallocate();
+                    }
+                    if (auto dec = inst->findDecoration<IRDownstreamModuleExportDecoration>())
+                    {
+                        dec->removeAndDeallocate();
+                    }
+                    break;
+                }
+            }
+        }
+    }
+}
+
 Result linkAndOptimizeIR(
     CodeGenContext*                         codeGenContext,
     LinkingAndOptimizationOptions const&    options,
@@ -746,6 +778,11 @@ Result linkAndOptimizeIR(
         break;
     }
 
+    if (codeGenContext->removeAvailableInDXIL)
+    {
+        removeAvailableInDownstreamModuleDecorations(irModule);
+    }
+
     if (targetProgram->getOptionSet().shouldRunNonEssentialValidation())
     {
         checkForRecursiveTypes(irModule, sink);
@@ -1453,6 +1490,13 @@ Result linkAndOptimizeIR(
     auto metadata = new ArtifactPostEmitMetadata;
     outLinkedIR.metadata = metadata;
 
+    if (targetProgram->getOptionSet().getBoolOption(CompilerOptionName::EmbedDXIL))
+    {
+        // We need to make sure that we don't try to export any functions that can't
+        // be part of a DXIL library interface, eg. with resources.
+        unexportNonEmbeddableDXIL(irModule);
+    }
+
     collectMetadata(irModule, *metadata);
 
     outLinkedIR.metadata = metadata;
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index eb4e88c41..09cb7952c 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -790,6 +790,7 @@ INST_RANGE(BindingQuery, GetRegisterIndex, GetRegisterSpace)
     INST(PreciseDecoration,                 precise,                0, 0)
     INST(PublicDecoration,                  public,                 0, 0)
     INST(HLSLExportDecoration,              hlslExport,             0, 0)
+    INST(DownstreamModuleExportDecoration,  downstreamModuleExport, 0, 0)
     INST(PatchConstantFuncDecoration,       patchConstantFunc,      1, 0)
     INST(OutputControlPointsDecoration,     outputControlPoints,    1, 0)
     INST(OutputTopologyDecoration,          outputTopology,         1, 0)
@@ -800,6 +801,8 @@ INST_RANGE(BindingQuery, GetRegisterIndex, GetRegisterSpace)
     INST(NumThreadsDecoration,              numThreads,             3, 0)
     INST(WaveSizeDecoration,                waveSize,               1, 0)
 
+    INST(AvailableInDXILDecoration,         availableInDXIL,        0, 0)
+
         // Added to IRParam parameters to an entry point
     /* GeometryInputPrimitiveTypeDecoration */
         INST(PointInputPrimitiveTypeDecoration,  pointPrimitiveType,     0, 0)
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index f8836219e..9b94005d9 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -430,6 +430,8 @@ IR_SIMPLE_DECORATION(NonCopyableTypeDecoration)
 IR_SIMPLE_DECORATION(HLSLMeshPayloadDecoration)
 IR_SIMPLE_DECORATION(GlobalInputDecoration)
 IR_SIMPLE_DECORATION(GlobalOutputDecoration)
+IR_SIMPLE_DECORATION(AvailableInDXILDecoration)
+IR_SIMPLE_DECORATION(DownstreamModuleExportDecoration)
 
 struct IRGLSLLocationDecoration : IRDecoration
 {
diff --git a/source/slang/slang-ir-link.cpp b/source/slang/slang-ir-link.cpp
index 8b08b9045..94ba75a44 100644
--- a/source/slang/slang-ir-link.cpp
+++ b/source/slang/slang-ir-link.cpp
@@ -1461,7 +1461,7 @@ static bool _isHLSLExported(IRInst* inst)
     for (auto decoration : inst->getDecorations())
     {
         const auto op = decoration->getOp();
-        if (op == kIROp_HLSLExportDecoration)
+        if (op == kIROp_HLSLExportDecoration || op == kIROp_DownstreamModuleExportDecoration)
         {
             return true;
         }
diff --git a/source/slang/slang-ir-metadata.cpp b/source/slang/slang-ir-metadata.cpp
index 641c272a2..030fde157 100644
--- a/source/slang/slang-ir-metadata.cpp
+++ b/source/slang/slang-ir-metadata.cpp
@@ -42,8 +42,18 @@ static void _insertBinding(List<ShaderBindingRange>& ranges, LayoutResourceKind
 void collectMetadata(const IRModule* irModule, ArtifactPostEmitMetadata& outMetadata)
 {
     // Scan the instructions looking for global resource declarations
+    // and exported functions.
     for (const auto& inst : irModule->getGlobalInsts())
     {
+        if (auto func = as<IRFunc>(inst))
+        {
+            if (func->findDecoration<IRDownstreamModuleExportDecoration>())
+            {
+                auto name = func->findDecoration<IRExportDecoration>()->getMangledName();
+                outMetadata.m_exportedFunctionMangledNames.add(name);
+            }
+        }
+
         auto param = as<IRGlobalParam>(inst);
         if (!param) continue;
         
diff --git a/source/slang/slang-ir-redundancy-removal.cpp b/source/slang/slang-ir-redundancy-removal.cpp
index fa3cd444a..038412fbf 100644
--- a/source/slang/slang-ir-redundancy-removal.cpp
+++ b/source/slang/slang-ir-redundancy-removal.cpp
@@ -159,6 +159,36 @@ bool removeRedundancyInFunc(IRGlobalValueWithCode* func)
     return result;
 }
 
+// Remove IR definitions from all [AvailableInDXIL] functions when compiling DXIL,
+// as these functions are already defined in the embedded precompiled DXIL library.
+void removeAvailableInDownstreamModuleDecorations(IRModule* module)
+{
+    List<IRInst*> toRemove;
+    for (auto globalInst : module->getGlobalInsts())
+    {
+        auto funcInst = as<IRFunc>(globalInst);
+        if (!funcInst)
+        {
+            continue;
+        }
+        if (globalInst->findDecoration<IRAvailableInDXILDecoration>())
+        {
+            // Gut the function definition, turning it into a declaration
+            for (auto inst : funcInst->getChildren())
+            {
+                if (inst->getOp() == kIROp_Block)
+                {
+                    toRemove.add(inst);
+                }
+            }
+        }
+    }
+    for (auto inst : toRemove)
+    {
+        inst->removeAndDeallocate();
+    }
+}
+
 static IRInst* _getRootVar(IRInst* inst)
 {
     while (inst)
diff --git a/source/slang/slang-ir-redundancy-removal.h b/source/slang/slang-ir-redundancy-removal.h
index c2df7853e..117f52708 100644
--- a/source/slang/slang-ir-redundancy-removal.h
+++ b/source/slang/slang-ir-redundancy-removal.h
@@ -10,4 +10,6 @@ namespace Slang
     bool removeRedundancyInFunc(IRGlobalValueWithCode* func);
 
     bool eliminateRedundantLoadStore(IRGlobalValueWithCode* func);
+
+    void removeAvailableInDownstreamModuleDecorations(IRModule* module);
 }
diff --git a/source/slang/slang-ir-util.cpp b/source/slang/slang-ir-util.cpp
index 817c10ec2..c4fc60bd2 100644
--- a/source/slang/slang-ir-util.cpp
+++ b/source/slang/slang-ir-util.cpp
@@ -244,6 +244,14 @@ bool isSimpleDataType(IRType* type)
     }
 }
 
+bool isSimpleHLSLDataType(IRInst* inst)
+{
+    // TODO: Add criteria
+    // https://github.com/shader-slang/slang/issues/4792
+    SLANG_UNUSED(inst);
+    return true;
+}
+
 SourceLoc findFirstUseLoc(IRInst* inst)
 {
     for (auto use = inst->firstUse; use; use = use->nextUse)
diff --git a/source/slang/slang-ir-util.h b/source/slang/slang-ir-util.h
index c7d6a1544..78d24d441 100644
--- a/source/slang/slang-ir-util.h
+++ b/source/slang/slang-ir-util.h
@@ -99,6 +99,8 @@ bool isValueType(IRInst* type);
 
 bool isSimpleDataType(IRType* type);
 
+bool isSimpleHLSLDataType(IRInst* inst);
+
 SourceLoc findFirstUseLoc(IRInst* inst);
 
 inline bool isChildInstOf(IRInst* inst, IRInst* parent)
diff --git a/source/slang/slang-ir.h b/source/slang/slang-ir.h
index 99c62b214..375107d1d 100644
--- a/source/slang/slang-ir.h
+++ b/source/slang/slang-ir.h
@@ -2389,9 +2389,6 @@ public:
     {
         return m_containerPool;
     }
-
-    // TODO: make a map with lookup by target?
-    ComPtr<ISlangBlob> precompiledDXIL;
 private:
     IRModule() = delete;
 
diff --git a/source/slang/slang-module-library.cpp b/source/slang/slang-module-library.cpp
index 42b9ff77a..02ace07d3 100644
--- a/source/slang/slang-module-library.cpp
+++ b/source/slang/slang-module-library.cpp
@@ -83,15 +83,6 @@ SlangResult loadModuleLibrary(const Byte* inBytes, size_t bytesCount, String pat
                     module, &sink);
                 if (!loadedModule)
                     return SLANG_FAIL;
-
-                for (auto inst : module.irModule->getModuleInst()->getChildren())
-                {
-                    if (inst->getOp() == kIROp_EmbeddedDXIL)
-                    {
-                        auto slice = static_cast<IRBlobLit*>(inst->getOperand(0))->getStringSlice();
-                        module.irModule->precompiledDXIL = StringBlob::create(slice);
-                    }
-                }
                 library->m_modules.add(loadedModule);
             }
         }
diff --git a/source/slang/slang-options.cpp b/source/slang/slang-options.cpp
index 12b32998f..cf0dd2f20 100644
--- a/source/slang/slang-options.cpp
+++ b/source/slang/slang-options.cpp
@@ -702,8 +702,6 @@ struct OptionsParser
     void setProfile(RawTarget* rawTarget, Profile profile);
     void addCapabilityAtom(RawTarget* rawTarget, CapabilityName atom);
 
-    SlangResult addEmbeddedLibrary(const CodeGenTarget format, CompilerOptionName option);
-    
     void setFloatingPointMode(RawTarget* rawTarget, FloatingPointMode mode);
     
     SlangResult parse(
@@ -1640,23 +1638,6 @@ SlangResult OptionsParser::_parseProfile(const CommandLineArg& arg)
     return SLANG_OK;
 }
 
-// Creates a target of the specified type whose output will be embedded as IR metadata
-SlangResult OptionsParser::addEmbeddedLibrary(const CodeGenTarget format, CompilerOptionName option)
-{
-    RawTarget rawTarget;
-    rawTarget.format = format;
-    // Silently allow redundant targets if it is the same as the last specified target.
-    if (m_rawTargets.getCount() == 0 || m_rawTargets.getLast().format != rawTarget.format)
-    {
-        m_rawTargets.add(rawTarget);
-    }
-
-    getCurrentTarget()->optionSet.add(option, true);
-    getCurrentTarget()->optionSet.add(CompilerOptionName::GenerateWholeProgram, true);
-
-    return SLANG_OK;
-}
-
 SlangResult OptionsParser::_parse(
     int             argc,
     char const* const* argv)
@@ -1948,7 +1929,7 @@ SlangResult OptionsParser::_parse(
                 linkage->m_optionSet.set(optionKind, compressionType);
                 break;
             }
-            case OptionKind::EmbedDXIL: SLANG_RETURN_ON_FAIL(addEmbeddedLibrary(CodeGenTarget::DXIL, CompilerOptionName::EmbedDXIL)); break;
+            case OptionKind::EmbedDXIL: m_compileRequest->setEmbedDXIL(true); break;
             case OptionKind::Target:
             {
                 CommandLineArg name;
@@ -2795,11 +2776,6 @@ SlangResult OptionsParser::_parse(
             {
                 m_compileRequest->setTargetGenerateWholeProgram(targetID, true);
             }
-
-            if (rawTarget.optionSet.getBoolOption(CompilerOptionName::EmbedDXIL))
-            {
-                m_compileRequest->setTargetEmbedDXIL(targetID, true);
-            }
         }
 
         // Next we need to sort out the output files specified with `-o`, and
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 23e25249d..4a6d33363 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -3300,19 +3300,16 @@ SlangResult EndToEndCompileRequest::executeActionsInner()
 
     // If requested, attempt to compile the translation unit all the way down to the target language
     // and stash the result blob in an IR op.
-    for (auto targetReq : getLinkage()->targets)
+    if (getOptionSet().getBoolOption(CompilerOptionName::EmbedDXIL))
     {
-        if (targetReq->getOptionSet().getBoolOption(CompilerOptionName::EmbedDXIL))
-        {
-            auto frontEndReq = getFrontEndReq();
+        auto frontEndReq = getFrontEndReq();
 
-            for (auto translationUnit : frontEndReq->translationUnits)
-            {
-                SlangCompileTarget target = SlangCompileTarget(targetReq->getTarget());
-                translationUnit->getModule()->precompileForTarget(
-                    target,
-                    nullptr);
-            }
+        for (auto translationUnit : frontEndReq->translationUnits)
+        {
+            SlangCompileTarget target = SlangCompileTarget(SlangCompileTarget::SLANG_DXIL);
+            SLANG_RETURN_ON_FAIL(translationUnit->getModule()->precompileForTarget(
+                target,
+                nullptr));
         }
     }
 
@@ -5935,9 +5932,9 @@ void EndToEndCompileRequest::setTargetGenerateWholeProgram(int targetIndex, bool
     getTargetOptionSet(targetIndex).set(CompilerOptionName::GenerateWholeProgram, value);
 }
 
-void EndToEndCompileRequest::setTargetEmbedDXIL(int targetIndex, bool value)
+void EndToEndCompileRequest::setEmbedDXIL(bool value)
 {
-    getTargetOptionSet(targetIndex).set(CompilerOptionName::EmbedDXIL, value);
+    getOptionSet().set(CompilerOptionName::EmbedDXIL, value);
 }
 
 void EndToEndCompileRequest::setTargetLineDirectiveMode(
@@ -6861,4 +6858,4 @@ SlangResult EndToEndCompileRequest::isParameterLocationUsed(Int entryPointIndex,
     return SLANG_OK;
 }
 
-} // namespace Slang
\ No newline at end of file
+} // namespace Slang