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,48 @@
 
 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)
     {
         return false;
     }
 
-    // Skip functions with no body
-    bool hasBody = false;
-    for (auto child : inst->getChildren())
+    if (inst->getOp() == kIROp_GlobalVar)
     {
-        if (child->getOp() == kIROp_Block)
+        inst->dump();
+        if (inst->findDecoration<IRExportDecoration>())
         {
-            hasBody = true;
-            break;
+            return true;
         }
-    }
-    if (!hasBody)
-    {
         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 +63,18 @@ static bool attemptPrecompiledExport(IRInst* inst)
     return true;
 }
 
+static bool needsImport(IRInst* inst)
+{
+    if (inst->getOp() == kIROp_GlobalVar)
+    {
+        if (inst->findDecoration<IRUserExternDecoration>())
+        {
+            return true;
+        }
+    }
+    return false;
+}
+
 /*
  * Precompile the module for the given target.
  *
@@ -160,24 +184,31 @@ 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;
     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))
+        {
+            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 export. That's not treated as an error
+    // because it's possible that the module just doesn't have any simple code.
+    if (!hasAtLeastOneExport)
     {
         return SLANG_OK;
     }
 
+    // dump module
+    module->getModuleInst()->dump();
+
     ComPtr<IArtifact> outArtifact;
     SlangResult res = codeGenContext.emitPrecompiledDownstreamIR(outArtifact);
 
@@ -208,7 +239,7 @@ Module::precompileForTarget(SlangCompileTarget target, slang::IBlob** outDiagnos
     // represent functions that were pruned from the IR after linking, before target generation.
     for (auto moduleInst : module->getGlobalInsts())
     {
-        if (moduleInst->getOp() == kIROp_Func)
+        //if (moduleInst->getOp() == kIROp_Func)
         {
             if (auto dec = moduleInst->findDecoration<IRDownstreamModuleExportDecoration>())
             {

source/slang/slang-emit-spirv.cpp

@@ -2556,7 +2556,7 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
                 varInst,
                 SpvLiteralInteger::from32(int32_t(0)));
         }
-
+        
         bool anyModifiers = (var->findDecoration<IRInterpolationModeDecoration>() != nullptr);
 
         // If the user didn't explicitly qualify a varying
@@ -2727,15 +2727,46 @@ 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
+        IRInst* initializer = nullptr;
+        for (auto child : globalVar->getChildren())
+        {
+            if (auto blockInst = as<IRBlock>(child))
+            {
+                for (auto inst : blockInst->getChildren())
+                {
+                    inst->dump();
+                    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;
     }
 
@@ -4860,8 +4891,19 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
         case kIROp_DownstreamModuleImportDecoration:
             {
                 requireSPIRVCapability(SpvCapabilityLinkage);
-                auto name =
-                    decoration->getParent()->findDecoration<IRExportDecoration>()->getMangledName();
+                printf("DownstreamModuleImportDecoration\n");
+                decoration->dump();
+                printf("\nParent to decoration\n");
+                decoration->getParent()->dump();
+                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,10 +8371,12 @@ SlangResult emitSPIRVFromIR(
         }
         if (generateWholeProgram)
         {
-            if (auto func = as<IRFunc>(inst))
+            //if (auto func = as<IRFunc>(inst))
             {
-                if (func->findDecoration<IRDownstreamModuleExportDecoration>())
+                if (inst->findDecoration<IRDownstreamModuleExportDecoration>())
                 {
+                    printf("exporting...............\n");
+                    inst->dump();
                     context.ensureInst(inst);
                     symbolsEmitted = true;
                 }

source/slang/slang-emit.cpp

@@ -2191,8 +2191,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.cpp

@@ -2367,6 +2367,11 @@ IRStringLit* IRBuilder::getStringValue(const UnownedStringSlice& inSlice)
 
 IRBlobLit* IRBuilder::getBlobValue(ISlangBlob* blob)
 {
+    if (!blob || blob->getBufferSize() == 0)
+    {
+        return nullptr;
+    }
+
     IRConstant keyInst;
     memset(&keyInst, 0, sizeof(keyInst));
 

tests/library/export-library-global.slang

@@ -1,7 +1,7 @@
-// TEST_IGNORE_FILE
-
-// export-library-global.slang
-
-export static int myGlobalVar = 42;
+//TEST:SIMPLE(filecheck=CHECK): -o tests/library/export-library-global.slang-module -target spirv -embed-downstream-ir -dump-ir
 
+// CHECK: EmbeddedDownstreamIR(6 : Int,
+// CHECK: OpCapability Linkage
+// CHECK: OpDecorate %myGlobalVar LinkageAttributes "_SV11myGlobalVar" Export
 
+export static int myGlobalVar = 42;

tests/library/precompiled-spirv-global.slang

@@ -1,8 +1,12 @@
-//TEST:COMPILE: tests/library/precompiled-spirv-global.slang -o tests/library/precompiled-spirv-global.slang-module -target spirv -embed-downstream-ir
+//TEST:COMPILE: tests/library/export-library-global.slang -o tests/library/export-library-global.slang-module -target spirv -embed-downstream-ir
+//TEST:SIMPLE(filecheck=CHECK): -o tests/library/precompiled-spirv-global.slang-module -target spirv -embed-downstream-ir -dump-ir
+//TEST:COMPILE: tests/library/export-library-global.slang-module tests/library/precompiled-spirv-global.slang-module -target spirv
 
-//import "export-library-global.slang";  //This would make the test pass
+// CHECK: EmbeddedDownstreamIR(6 : Int,
+// CHECK: OpCapability Linkage
+// CHECK: OpDecorate %myGlobalVar LinkageAttributes "_SV11myGlobalVar" Import
 
-extern int myGlobalVar; //This makes the test fail
+extern static int myGlobalVar;
 
 layout(location = 0) out float4 fragColor;