Precompiled SPIR-V import support (#5048)

* Precompiled SPIR-V import support

Adds appropriate linkage and function declaration syntax
for SPIR-V functions that are declared, to be imported
from another SPIR-V module.

Unlike DXIL, stripping the Slang IR for a function down
to a declaration requires retaining a block of parameters,
as the function declaration must be emitted to SPIR-V
with the same parameters as a definition. Because that
thwarts the logic in Slang to tell the difference between
a declaration and definition, and explicit decoration is
introduced to explicitly mark functions which need to be
treated as declarations during emit phase.

Fixes #4992

Co-authored-by: Yong He <yonghe@outlook.com>

4 files changed, 101 insertions, 17 deletions

diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index e134c579e..954caf24b 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -6,6 +6,7 @@
 #include "slang-ir-call-graph.h"
 #include "slang-ir-insts.h"
 #include "slang-ir-layout.h"
+#include "slang-ir-redundancy-removal.h"
 #include "slang-ir-spirv-legalize.h"
 #include "slang-ir-spirv-snippet.h"
 #include "slang-ir-util.h"
@@ -2628,14 +2629,75 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
     /// Emit a declaration for the given `irFunc`
     SpvInst* emitFuncDeclaration(IRFunc* irFunc)
     {
-        if (irFunc->findDecorationImpl(kIROp_SPIRVOpDecoration))
-            return nullptr;
-        // For now we aren't handling function declarations;
-        // we expect to deal only with fully linked modules.
+        // [2.4: Logical Layout of a Module]
+        //
+        // > All function declarations("declarations" are functions without a
+        //   body; there is no forward declaration to a function with a body).
+        //
+        auto section = getSection(SpvLogicalSectionID::FunctionDeclarations);
+
+        // > A function declaration is as follows.
+        // > * Function declaration, using OpFunction.
+        // > * Function parameter declarations, using OpFunctionParameter.
+        // > * Function end, using OpFunctionEnd.
+        //
+
+        // [3.24. Function Control]
+        //
+        // TODO: We should eventually support emitting the "function control"
+        // mask to include inline and other hint bits based on decorations
+        // set on `irFunc`.
+        //
+        SpvFunctionControlMask spvFunctionControl = SpvFunctionControlMaskNone;
+
+        // [3.32.9. Function Instructions]
+        //
+        // > OpFunction
+        //
+        // Note that the type <id> of a SPIR-V function uses the
+        // *result* type of the function, while the actual function
+        // type is given as a later operand. Slan IR instead uses
+        // the type of a function instruction store, you know, its *type*.
+        //
+        SpvInst* spvFunc = emitOpFunction(
+            section,
+            irFunc,
+            irFunc->getDataType()->getResultType(),
+            spvFunctionControl,
+            irFunc->getDataType());
+
+        // > OpFunctionParameter
+        //
+        // Though parameters always belong to blocks in Slang, there are no
+        // blocks in a function declaration, so we will emit the parameters
+        // as derived from the function's type.
+        //
+        auto funcType = irFunc->getDataType();
+        auto paramCount = funcType->getParamCount();
+        for (UInt pp = 0; pp < paramCount; ++pp)
+        {
+            auto paramType = funcType->getParamType(pp);
+            SpvInst* spvParam = emitOpFunctionParameter(spvFunc, nullptr, paramType);
+            maybeEmitPointerDecoration(spvParam, paramType, false, kIROp_Param);
+        }
+
+        // [3.32.9. Function Instructions]
+        //
+        // > OpFunctionEnd
+        //
+        // In the SPIR-V encoding a function is logically the parent of any
+        // instructions up to a matching `OpFunctionEnd`. In our intermediate
+        // structure we will make the `OpFunctionEnd` be the last child of
+        // the `OpFunction`.
         //
-        m_sink->diagnose(irFunc, Diagnostics::internalCompilerError);
-        SLANG_UNEXPECTED("function declaration in SPIR-V emit");
-        UNREACHABLE_RETURN(nullptr);
+        emitOpFunctionEnd(spvFunc, nullptr);
+
+        // We will emit any decorations pertinent to the function to the
+        // appropriate section of the module.
+        //
+        emitDecorations(irFunc, getID(spvFunc));
+
+        return spvFunc;
     }
 
     /// Emit a SPIR-V function definition for the Slang IR function `irFunc`.
@@ -4358,6 +4420,21 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
                     SpvLinkageTypeExport);
                 break;
             }
+        case kIROp_DownstreamModuleImportDecoration:
+            {
+                requireSPIRVCapability(SpvCapabilityLinkage);
+                auto name =
+                    decoration->getParent()->findDecoration<IRExportDecoration>()->getMangledName();
+                emitInst(
+                    getSection(SpvLogicalSectionID::Annotations),
+                    decoration,
+                    SpvOpDecorate,
+                    dstID,
+                    SpvDecorationLinkageAttributes,
+                    name,
+                    SpvLinkageTypeImport);
+                break;
+            }
             // ...
         }
 
@@ -5019,9 +5096,9 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
         return nullptr;
     }
 
-    void maybeEmitPointerDecoration(SpvInst* varInst, IRInst* inst)
+    void maybeEmitPointerDecoration(SpvInst* varInst, IRType* type, bool isVar, IROp op)
     {
-        auto ptrType = as<IRPtrType>(unwrapArray(inst->getDataType()));
+        auto ptrType = as<IRPtrType>(unwrapArray(type));
         if (!ptrType)
             return;
         if (addressSpaceToStorageClass(ptrType->getAddressSpace()) ==
@@ -5033,7 +5110,7 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
                 getSection(SpvLogicalSectionID::Annotations),
                 nullptr,
                 varInst,
-                (as<IRVar>(inst) ? SpvDecorationAliasedPointer : SpvDecorationAliased));
+                (isVar ? SpvDecorationAliasedPointer : SpvDecorationAliased));
         }
         else
         {
@@ -5049,14 +5126,18 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
                     getSection(SpvLogicalSectionID::Annotations),
                     nullptr,
                     varInst,
-                    (inst->getOp() == kIROp_GlobalVar || inst->getOp() == kIROp_Var ||
-                             inst->getOp() == kIROp_DebugVar
+                    (op == kIROp_GlobalVar || op == kIROp_Var || op == kIROp_DebugVar
                          ? SpvDecorationAliasedPointer
                          : SpvDecorationAliased));
             }
         }
     }
 
+    void maybeEmitPointerDecoration(SpvInst* varInst, IRInst* inst)
+    {
+        maybeEmitPointerDecoration(varInst, inst->getDataType(), as<IRVar>(inst), inst->getOp());
+    }
+
     SpvInst* emitParam(SpvInstParent* parent, IRInst* inst)
     {
         auto paramSpvInst = emitOpFunctionParameter(parent, inst, inst->getFullType());
@@ -7534,6 +7615,8 @@ SlangResult emitSPIRVFromIR(
     }
 #endif
 
+    removeAvailableInDownstreamModuleDecorations(CodeGenTarget::SPIRV, irModule);
+
     auto shouldPreserveParams = codeGenContext->getTargetProgram()->getOptionSet().getBoolOption(
         CompilerOptionName::PreserveParameters);
     auto generateWholeProgram = codeGenContext->getTargetProgram()->getOptionSet().getBoolOption(
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index c98fae55e..33baefa51 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -822,6 +822,7 @@ INST_RANGE(BindingQuery, GetRegisterIndex, GetRegisterSpace)
     INST(PublicDecoration,                  public,                 0, 0)
     INST(HLSLExportDecoration,              hlslExport,             0, 0)
     INST(DownstreamModuleExportDecoration,  downstreamModuleExport, 0, 0)
+    INST(DownstreamModuleImportDecoration,  downstreamModuleImport, 0, 0)
     INST(PatchConstantFuncDecoration,       patchConstantFunc,      1, 0)
     INST(OutputControlPointsDecoration,     outputControlPoints,    1, 0)
     INST(OutputTopologyDecoration,          outputTopology,         1, 0)
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 7c87dd377..c44211c1c 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -452,6 +452,7 @@ IR_SIMPLE_DECORATION(HLSLMeshPayloadDecoration)
 IR_SIMPLE_DECORATION(GlobalInputDecoration)
 IR_SIMPLE_DECORATION(GlobalOutputDecoration)
 IR_SIMPLE_DECORATION(DownstreamModuleExportDecoration)
+IR_SIMPLE_DECORATION(DownstreamModuleImportDecoration)
 
 struct IRAvailableInDownstreamIRDecoration : IRDecoration
 {
diff --git a/source/slang/slang-ir-redundancy-removal.cpp b/source/slang/slang-ir-redundancy-removal.cpp
index 230f356a4..b90a8705d 100644
--- a/source/slang/slang-ir-redundancy-removal.cpp
+++ b/source/slang/slang-ir-redundancy-removal.cpp
@@ -171,6 +171,7 @@ bool removeRedundancyInFunc(IRGlobalValueWithCode* func)
 void removeAvailableInDownstreamModuleDecorations(CodeGenTarget target, IRModule* module)
 {
     List<IRInst*> toRemove;
+    auto builder = IRBuilder(module);
     for (auto globalInst : module->getGlobalInsts())
     {
         if (auto funcInst = as<IRFunc>(globalInst))
@@ -181,13 +182,11 @@ void removeAvailableInDownstreamModuleDecorations(CodeGenTarget target, IRModule
                     (dec->getTarget() == target))
                 {
                     // Gut the function definition, turning it into a declaration
-                    for (auto inst : funcInst->getChildren())
+                    for (auto block : funcInst->getBlocks())
                     {
-                        if (inst->getOp() == kIROp_Block)
-                        {
-                            toRemove.add(inst);
-                        }
+                        toRemove.add(block);
                     }
+                    builder.addDecoration(funcInst, kIROp_DownstreamModuleImportDecoration);
                 }
             }
         }