Support link time constants in embedded SPIR-V

When modules contain extern references to globals they don't define, or when the contain exports of globals they may or may not use, utilize SPIR-V linkage attributes to define and use the variables in embedded SPIR-V IR.

Fixes shader-slang#6524

Author: cheneym2

source/slang/slang-compiler-tu.cpp

@@ -10,36 +10,51 @@
 
 namespace Slang
 {
-// Only attempt to precompile functions:
+// Only attempt to precompile functions and global variables:
 // 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)
+    if (inst->getOp() != kIROp_Func && inst->getOp() != kIROp_GlobalVar && inst->getOp() != kIROp_GlobalConstant)
     {
         return false;
     }
 
-    // Skip functions with no body
-    bool hasBody = false;
-    for (auto child : inst->getChildren())
+    if (inst->getOp() == kIROp_GlobalVar ||
+        inst->getOp() == kIROp_GlobalConstant)
     {
-        if (child->getOp() == kIROp_Block)
+        printf("GlobalVar or GlobalConstant\n");
+        if (inst->findDecoration<IRExportDecoration>())
         {
-            hasBody = true;
-            break;
+            printf("Exporting global variable\n");
+            return true;
         }
-    }
-    if (!hasBody)
-    {
+        printf("Not exporting global variable\n");
         return false;
     }
-
-    // Skip functions marked with unsafeForceInlineDecoration
-    if (inst->findDecoration<IRUnsafeForceInlineEarlyDecoration>())
+    else 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
@@ -51,6 +66,18 @@ static bool attemptPrecompiledExport(IRInst* inst)
     return true;
 }
 
+static bool needsImport(IRInst* inst)
+{
+    if (inst->getOp() == kIROp_GlobalVar || inst->getOp() == kIROp_GlobalConstant)
+    {
+        if (inst->findDecoration<IRUserExternDecoration>())
+        {
+            return true;
+        }
+    }
+    return false;
+}
+
 /*
  * Precompile the module for the given target.
  *
@@ -67,28 +94,28 @@ static bool attemptPrecompiledExport(IRInst* inst)
  * done during target generation in between IR linking+legalization and
  * target source emission.
  *
- * Functions which can be rejected up front:
+ * Language features 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
+ * The instructions not rejected up front are marked with
+ * DownstreamModuleExportDecoration which indicates what we're trying to
+ * export for precompilation, and this also helps to identify the instructions
  * in the linked IR which survived the additional pruning.
  *
- * Functions that are rejected after linking+legalization (inside
+ * Instructions that are rejected after linking+legalization (inside
  * emitPrecompiledDownstreamIR):
  * - (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
+ * blob, and as metadata, the list of instructions which survived the second
  * phase of filtering.
  *
- * The original module IR functions matching those are then marked with
+ * The original module IR instructions matching those are then marked with
  * "AvailableInDownstreamIRDecoration" to indicate to future
- * module users which functions are present in the precompiled blob.
+ * module users which instructions are present in the precompiled blob.
  */
 SLANG_NO_THROW SlangResult SLANG_MCALL
 Module::precompileForTarget(SlangCompileTarget target, slang::IBlob** outDiagnostics)
@@ -107,6 +134,9 @@ Module::precompileForTarget(SlangCompileTarget target, slang::IBlob** outDiagnos
         }
     }
 
+    if (getEntryPoints().getCount() != 0)
+        return SLANG_OK;
+
     auto module = getIRModule();
     auto linkage = getLinkage();
     auto builder = IRBuilder(module);
@@ -160,20 +190,26 @@ Module::precompileForTarget(SlangCompileTarget target, slang::IBlob** outDiagnos
     // the linked result to see which functions survived the pruning and are included in the
     // precompiled blob.
     Dictionary<String, IRInst*> nameToFunction;
-    bool hasAtLeastOneFunction = false;
+    bool hasAtLeastOneExport = false;
+    bool hasAtLeastOneImport = false;
     for (auto inst : module->getGlobalInsts())
     {
         if (attemptPrecompiledExport(inst))
         {
-            hasAtLeastOneFunction = true;
+            hasAtLeastOneExport = true;
             builder.addDecoration(inst, kIROp_DownstreamModuleExportDecoration);
             nameToFunction[inst->findDecoration<IRExportDecoration>()->getMangledName()] = inst;
         }
+        if (needsImport(inst))
+        {
+            hasAtLeastOneImport = true;
+            builder.addDecoration(inst, kIROp_DownstreamModuleImportDecoration);
+        }
     }
 
-    // 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)
+    // Bail if there is nothing to import/export. That's not treated as an error
+    // because it's possible that the module just doesn't have any simple code.
+    if (!hasAtLeastOneExport && !hasAtLeastOneImport)
     {
         return SLANG_OK;
     }
@@ -201,19 +237,17 @@ Module::precompileForTarget(SlangCompileTarget target, slang::IBlob** outDiagnos
             kIROp_AvailableInDownstreamIRDecoration,
             builder.getIntValue(builder.getIntType(), (int)targetReq->getTarget()));
         auto moduleDec = moduleInst->findDecoration<IRDownstreamModuleExportDecoration>();
-        moduleDec->removeAndDeallocate();
+        if (moduleDec)
+            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>())
         {
-            if (auto dec = moduleInst->findDecoration<IRDownstreamModuleExportDecoration>())
-            {
-                dec->removeAndDeallocate();
-            }
+            dec->removeAndDeallocate();
         }
     }
 

source/slang/slang-emit-spirv.cpp

@@ -1793,6 +1793,8 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
             return emitGlobalParam(as<IRGlobalParam>(inst));
         case kIROp_GlobalVar:
             return emitGlobalVar(as<IRGlobalVar>(inst));
+        case kIROp_GlobalConstant:
+            return emitGlobalConstant(as<IRGlobalConstant>(inst));
         case kIROp_SPIRVAsmOperandBuiltinVar:
             return emitBuiltinVar(inst);
         case kIROp_Var:
@@ -2727,15 +2729,55 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
             if (ptrType->hasAddressSpace())
                 storageClass = addressSpaceToStorageClass(ptrType->getAddressSpace());
         }
-        auto varInst = emitOpVariable(
-            getSection(SpvLogicalSectionID::GlobalVariables),
-            globalVar,
-            globalVar->getDataType(),
-            storageClass);
+        
+        // iterate through block children looking for constant initializers
+        // TODO: This is a bit of a hack.
+        IRInst* initializer = nullptr;
+        for (auto child : globalVar->getChildren())
+        {
+            if (auto blockInst = as<IRBlock>(child))
+            {
+                for (auto inst : blockInst->getChildren())
+                {
+                    if (auto var = as<IRReturn>(inst))
+                    {
+                        initializer = var->getVal();
+                    }
+                }
+            }
+        }
+        SpvInst* varInst;
+        if (initializer)
+        {
+            varInst = emitOpVariable(
+                getSection(SpvLogicalSectionID::GlobalVariables),
+                globalVar,
+                globalVar->getDataType(),
+                storageClass,
+                initializer);
+        }
+        else
+        {
+            varInst = emitOpVariable(
+                getSection(SpvLogicalSectionID::GlobalVariables),
+                globalVar,
+                globalVar->getDataType(),
+                storageClass);
+        }        
         maybeEmitPointerDecoration(varInst, globalVar);
         if (layout)
             emitVarLayout(globalVar, varInst, layout);
         emitDecorations(globalVar, getID(varInst));
+
+        return varInst;
+    }
+
+    SpvInst* emitGlobalConstant(IRGlobalConstant* globalConstant)
+    {
+        SpvInst* varInst = emitOpConstant(
+            globalConstant,
+            globalConstant->getDataType(),
+            getLiteralBits(globalConstant->getDataType(), globalConstant));
         return varInst;
     }
 
@@ -4041,6 +4083,8 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
                     operands.getArrayView());
             }
             break;
+        case kIROp_GlobalConstant:
+            return emitGlobalConstant(as<IRGlobalConstant>(inst));
         }
         if (result)
             emitDecorations(inst, getID(result));
@@ -4860,8 +4904,15 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
         case kIROp_DownstreamModuleImportDecoration:
             {
                 requireSPIRVCapability(SpvCapabilityLinkage);
-                auto name =
-                    decoration->getParent()->findDecoration<IRExportDecoration>()->getMangledName();
+                UnownedStringSlice name;
+                if (decoration->getParent()->findDecoration<IRImportDecoration>())
+                {
+                    name = decoration->getParent()->findDecoration<IRImportDecoration>()->getMangledName();
+                }
+                else
+                {
+                    name = decoration->getParent()->findDecoration<IRExportDecoration>()->getMangledName();
+                }
                 emitInst(
                     getSection(SpvLogicalSectionID::Annotations),
                     decoration,
@@ -8329,13 +8380,10 @@ SlangResult emitSPIRVFromIR(
         }
         if (generateWholeProgram)
         {
-            if (auto func = as<IRFunc>(inst))
+            if (inst->findDecoration<IRDownstreamModuleExportDecoration>())
             {
-                if (func->findDecoration<IRDownstreamModuleExportDecoration>())
-                {
-                    context.ensureInst(inst);
-                    symbolsEmitted = true;
-                }
+                context.ensureInst(inst);
+                symbolsEmitted = true;
             }
         }
     }

source/slang/slang-emit.cpp

@@ -2067,6 +2067,15 @@ SlangResult emitSPIRVForEntryPointsDirectly(
     // Outside because we want to keep IR in scope whilst we are processing emits
     LinkedIR linkedIR;
     LinkingAndOptimizationOptions linkingAndOptimizationOptions;
+
+    bool isPrecompilation = codeGenContext->getTargetProgram()->getOptionSet().getBoolOption(
+        CompilerOptionName::EmbedDownstreamIR);
+
+    if (!isPrecompilation)
+    {
+        codeGenContext->removeAvailableInDownstreamIR = true;
+    }
+
     SLANG_RETURN_ON_FAIL(
         linkAndOptimizeIR(codeGenContext, linkingAndOptimizationOptions, linkedIR));
 
@@ -2098,9 +2107,6 @@ SlangResult emitSPIRVForEntryPointsDirectly(
         compiler->disassemble((uint32_t*)spirv.getBuffer(), int(spirv.getCount() / 4));
 #endif
 
-        bool isPrecompilation = codeGenContext->getTargetProgram()->getOptionSet().getBoolOption(
-            CompilerOptionName::EmbedDownstreamIR);
-
         if (!isPrecompilation && !codeGenContext->shouldSkipDownstreamLinking())
         {
             ComPtr<IArtifact> linkedArtifact;
@@ -2191,8 +2197,8 @@ SlangResult emitSPIRVForEntryPointsDirectly(
             downstreamOptions.optimizationLevel = DownstreamCompileOptions::OptimizationLevel::None;
             break;
         case OptimizationLevel::Default:
-            downstreamOptions.optimizationLevel =
-                DownstreamCompileOptions::OptimizationLevel::Default;
+            downstreamOptions.optimizationLevel = DownstreamCompileOptions::OptimizationLevel::None;
+            //Default;
             break;
         case OptimizationLevel::High:
             downstreamOptions.optimizationLevel = DownstreamCompileOptions::OptimizationLevel::High;

source/slang/slang-ir-explicit-global-init.cpp

@@ -72,6 +72,22 @@ struct MoveGlobalVarInitializationToEntryPointsPass
         m_module = module;
         m_targetProgram = targetProgram;
 
+        // Hack. If there are no entrypoint functions, we don't need to do anything.
+        bool hasEntryPoints = false;
+        for (auto inst : m_module->getGlobalInsts())
+        {
+            auto func = as<IRFunc>(inst);
+            if (!func)
+                continue;
+
+            if (!func->findDecoration<IREntryPointDecoration>())
+                continue;
+
+            hasEntryPoints = true;
+        }
+        if (!hasEntryPoints)
+            return;
+
         // We start by looking for global variables with
         // initialization logic in the IR, and processing
         // each to produce a split variable (now without

source/slang/slang-ir-link.cpp

@@ -1720,7 +1720,7 @@ static void diagnoseUnresolvedSymbols(TargetRequest* req, DiagnosticSink* sink,
 {
     for (auto globalSym : module->getGlobalInsts())
     {
-        if (globalSym->findDecoration<IRImportDecoration>())
+        if (globalSym->findDecoration<IRImportDecoration>() && !globalSym->findDecoration<IRDownstreamModuleImportDecoration>())
         {
             for (;;)
             {

source/slang/slang-ir-metadata.cpp

@@ -118,13 +118,14 @@ void collectMetadata(const IRModule* irModule, ArtifactPostEmitMetadata& outMeta
     // and exported functions.
     for (const auto& inst : irModule->getGlobalInsts())
     {
+        if (inst->findDecoration<IRDownstreamModuleExportDecoration>())
+        {
+            auto name = inst->findDecoration<IRExportDecoration>()->getMangledName();
+            outMetadata.m_exportedFunctionMangledNames.add(name);
+        }
+
         if (auto func = as<IRFunc>(inst))
         {
-            if (func->findDecoration<IRDownstreamModuleExportDecoration>())
-            {
-                auto name = func->findDecoration<IRExportDecoration>()->getMangledName();
-                outMetadata.m_exportedFunctionMangledNames.add(name);
-            }
 
             // Collect metadata from entrypoint params.
             for (auto param : func->getParams())