Hoist resource variables using a global initializer function correctly

source/slang/slang-emit.cpp

@@ -61,6 +61,7 @@
 #include "slang-ir-restructure.h"
 #include "slang-ir-restructure-scoping.h"
 #include "slang-ir-sccp.h"
+#include "slang-ir-simplify-global-vars.h"
 #include "slang-ir-specialize.h"
 #include "slang-ir-specialize-arrays.h"
 #include "slang-ir-specialize-buffer-load-arg.h"
@@ -1216,10 +1217,18 @@ Result linkAndOptimizeIR(
         break;
     }
 
+    // Sometimes we will have a GlobalVar with a 'init' that aliases another
+    // concrete value.
+    // These sort of GlobalVar's should be removed in-favor of directly
+    // using the known value since this may be an indirect use of a resource
+    // variable (which is illegal for almost targets).
+    simplifyGlobalVars(irModule);
+
     switch( target )
     {
     default:
         break;
+    case CodeGenTarget::HLSL:
     case CodeGenTarget::GLSL:
         moveGlobalVarInitializationToEntryPoints(irModule);
         break;

source/slang/slang-ir-legalize-types.cpp

@@ -343,7 +343,6 @@ struct LegalCallBuilder
                 // result type of the function, so we know that
                 // the legalization funciton/call will use a `void`
                 // result type.
-                //
                 _emitCall(m_context->builder->getVoidType());
                 return resultVal;
             }
@@ -364,41 +363,39 @@ struct LegalCallBuilder
         {
         case LegalType::Flavor::simple:
             {
-                // In the leaf case we have a simple type, and
-                // we just want to declare a local variable based on it.
-                //
+                // Here we emit a local/global var which will be assigned as an out-parameter argument, then we 'Load' this var
+                // and use that 'Load' as the result of simplification.
+
                 auto simpleType = resultType.getSimple();
                 auto builder = m_context->builder;
 
-                // Recall that a local variable in our IR represents a *pointer*
-                // to storage of the appropriate type.
-                //
-                auto varPtr = builder->emitVar(simpleType);
+                // If we were going to emit an IRVar in global scope we have 1 important case to resolve: when IRVar is a resource
+                // and is assigned something in global scope. This is important since this means we need to make an IRGlobalVar 
+                // instead of IRVar.
+                // A later pass will handle any out-parameter hoisting into IRGlobalVar init block logic
 
+                IRInst* varPtr = nullptr;
+                if (!builder->getInsertLoc().getParent() || builder->getInsertLoc().getParent()->getOp() == kIROp_Module)
+                    varPtr = builder->createGlobalVar(simpleType);
+                else
+                    varPtr = builder->emitVar(simpleType);
                 // We need to pass that pointer as an argument to our new
                 // `call` instruction, so that it can receive the value
                 // written by the callee.
                 //
                 m_args.add(varPtr);
 
                 // Note: Because `varPtr` is a pointer to the value we want,
-                // we have the small problem of needing to return a `LegalVal`
-                // that has dereferenced the value after the call.
-                //
-                // We solve this problem by inserting as `load` from our
-                // new variable immediately after the call, before going
-                // and resetting the insertion point to continue inserting
-                // stuff before the call (which is where we wnat the local
-                // variable declarations to go).
-                //
-                // TODO: Confirm that this logic can't go awry if (somehow)
-                // there is no instruction after `m_call`. That should not
-                // be possible inside of a function body, but it could in
-                // theory be a problem if we ever have top-level module-scope
-                // code representing initialization of constants and/or globals.
-                //
+                // but we want a value, we need to 'Load' the 'varPtr' result
+                // after our m_call.
+                // 
+                // Since we may have more than 1 parameter we 'setInsertBefore(m_call);'
                 builder->setInsertBefore(m_call->getNextInst());
-                auto val = builder->emitLoad(simpleType, varPtr);
+                IRInst* val = nullptr;
+                if (as<IRGlobalVar>(varPtr))
+                    val = varPtr;
+                else
+                    val = builder->emitLoad(simpleType, varPtr);
                 builder->setInsertBefore(m_call);
 
                 return LegalVal::simple(val);

source/slang/slang-ir-simplify-global-vars.cpp

@@ -0,0 +1,108 @@
+// slang-ir-explicit-global-init.cpp
+
+#include "slang-ir-simplify-global-vars.h"
+#include "slang-ir-insts.h"
+#include "slang-ir.h"
+#include "slang-ir-util.h"
+
+namespace Slang
+{
+    void tryToSimplifyUseOfGlobalVar(IRInst* user, List<IRInst*>& toDestroy)
+    {
+        switch (user->getOp())
+        {
+        case kIROp_Store:
+            {
+                auto store = as<IRStore>(user);
+                if (store->getParent() && store->getParent()->getOp() == kIROp_Module)
+                {
+                    // Store in Global scope into a GlobalVar is equal to 'insertReturn(global->GetFirstBlock(), Store->GetVal())'
+                    if (auto globalVal = as<IRGlobalValueWithCode>(store->getPtr()))
+                    {
+                        IRBuilder builder(store);
+                        IRBlock* block = globalVal->getFirstBlock();
+                        if (!block)
+                        {
+                            builder.setInsertInto(globalVal);
+                            block = builder.emitBlock();
+                        }
+                        builder.setInsertInto(block);
+                        builder.emitReturn(store->getVal());
+                        toDestroy.add(user);
+                    }
+                }
+                break;
+            }
+        default:
+            break;
+        }
+    }
+
+    void simplifyGlobalVars(IRModule* module)
+    {
+        IRBuilder builder(module);
+        for (auto inst : module->getGlobalInsts())
+        {
+            auto globalVar = as<IRGlobalVar>(inst);
+            if (!globalVar)
+                continue;
+
+            // Only simplify resource types since otherwise we need additional logic
+            // not yet implemented
+            auto globalVarType = globalVar->getDataType();
+            auto globalPtrType = as<IRPtrTypeBase>(globalVarType);
+            if (!globalPtrType
+                || !isResourceType(globalPtrType->getValueType()))
+                continue;
+
+            // Find any trivial use which is supposed-to-be part of the
+            // init block, hoist these.
+            // Do some extra auixillary checks.
+            List<IRInst*> toDestroy;
+            bool onlyGlobalSideEffects = true;
+            traverseUses(globalVar, [&](IRUse* use)
+                {
+                    auto user = use->getUser();
+                    tryToSimplifyUseOfGlobalVar(user, toDestroy);
+                    if (user->getParent()->getOp() != kIROp_Module
+                        && user->mightHaveSideEffects())
+                        onlyGlobalSideEffects = false;
+                });
+            for (auto i : toDestroy)
+                i->removeAndDeallocate();
+
+            // Check if our IRGlobalVar is an alias for another variable, if so, simplify
+            if (!onlyGlobalSideEffects)
+                continue;
+
+            auto firstBlock = globalVar->getFirstBlock();
+            if (!firstBlock)
+                continue;
+
+            // For an alias we must encounter a return op first
+            bool fail = false;
+            IRReturn* firstReturn = nullptr;
+
+            for (auto localInst : firstBlock->getChildren())
+            {
+                if (localInst->getOp() != kIROp_Return)
+                {
+                    fail = true;
+                    break;
+                }
+                else if (localInst->getOp() == kIROp_Return)
+                {
+                    firstReturn = as<IRReturn>(localInst);
+                    break;
+                }
+            }
+
+            if (fail || !firstReturn)
+                continue;
+
+            globalVar->replaceUsesWith(firstReturn->getVal());
+            firstBlock->removeAndDeallocateAllDecorationsAndChildren();
+            firstBlock->removeAndDeallocate();
+        }
+    }
+}

source/slang/slang-ir-simplify-global-vars.h

@@ -0,0 +1,11 @@
+// slang-ir-explicit-global-init.h
+#pragma once
+
+namespace Slang
+{
+struct IRModule;
+
+/// Simplify IRGlobalVars and possibley remove if an alias.
+void simplifyGlobalVars(IRModule* module);
+
+}

source/slang/slang-ir-specialize-resources.cpp

@@ -1202,7 +1202,9 @@ bool specializeResourceOutputs(
     HashSet<IRFunc*>&          unspecializableFuncs)
 {
     auto targetRequest = codeGenContext->getTargetReq();
-    if(isD3DTarget(targetRequest) || isKhronosTarget(targetRequest))
+    // D3D & SPIRV/GLSL cannot easily use resources as parameters
+    // Metal cannot assign globals to resource-outputting functions (in global scope)
+    if(isD3DTarget(targetRequest) || isKhronosTarget(targetRequest) || isMetalTarget(targetRequest))
     {}
     else
     {

tests/compute/type-legalize-global-with-init-function.slang

@@ -0,0 +1,49 @@
+// type-legalize-global-with-init.slang
+//
+// Confirm that type legalization can handle a global constant
+// with a resource type or a type that recursively contains
+// resources.
+//
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=BUF): -shaderobj
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=BUF): -vk -shaderobj
+//
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<uint> outputBuffer;
+
+//TEST_INPUT:ubuffer(data=[1 2 3 4 5 6 7 8], stride=4):name=inputBuffer
+RWStructuredBuffer<uint> inputBuffer;
+
+static const RWStructuredBuffer<uint> gBuffer = inputBuffer;
+
+struct Stuff
+{
+    __init(RWStructuredBuffer<uint> inA, RWStructuredBuffer<uint> inB)
+    {
+        a = inA;
+        b = inB;
+    }
+    RWStructuredBuffer<uint> a;
+    RWStructuredBuffer<uint> b;
+}
+
+static const Stuff gStuff = Stuff( inputBuffer, inputBuffer );
+
+uint test(uint x)
+{
+    return gBuffer[x]
+        + gStuff.a[x + 1] * 16
+        + gStuff.b[x + 2] * 256;
+}
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    let tid = dispatchThreadID.x;
+    let inVal = tid;
+    let outVal = test(inVal);
+    outputBuffer[tid] = outVal;
+}
+//BUF: 321
+//BUF: 432
+//BUF: 543
+//BUF: 654

tests/compute/type-legalize-global-with-init.slang

@@ -4,7 +4,8 @@
 // with a resource type or a type that recursively contains
 // resources.
 //
-//TEST(compute):COMPARE_COMPUTE: -shaderobj
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=BUF): -shaderobj
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=BUF): -vk -shaderobj
 //
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<uint> outputBuffer;
@@ -37,3 +38,8 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     let outVal = test(inVal);
     outputBuffer[tid] = outVal;
 }
+//BUF: 321
+//BUF: 432
+//BUF: 543
+//BUF: 654
+