WGSL: Fix issue where global calls are generated

source/slang/slang-emit-c-like.cpp

@@ -2963,6 +2963,11 @@ void CLikeSourceEmitter::defaultEmitInstExpr(IRInst* inst, const EmitOpInfo& inO
         {
             break; // should already have been parsed and used.
         }
+    case kIROp_GlobalValueRef:
+        {
+            emitOperand(as<IRGlobalValueRef>(inst)->getOperand(0), getInfo(EmitOp::General));
+            break;
+        }
     default:
         diagnoseUnhandledInst(inst);
         break;

source/slang/slang-ir-legalize-global-values.cpp

@@ -0,0 +1,247 @@
+#include "slang-ir-legalize-global-values.h"
+
+#include "slang-ir-clone.h"
+#include "slang-ir-util.h"
+
+namespace Slang
+{
+
+void GlobalInstInliningContextGeneric::inlineGlobalValuesAndRemoveIfUnused(IRModule* module)
+{
+    List<IRUse*> globalInstUsesToInline;
+
+    for (auto globalInst : module->getGlobalInsts())
+    {
+        if (isInlinableGlobalInst(globalInst))
+        {
+            for (auto use = globalInst->firstUse; use; use = use->nextUse)
+            {
+                if (getParentFunc(use->getUser()) != nullptr)
+                    globalInstUsesToInline.add(use);
+            }
+        }
+    }
+
+    HashSet<IRInst*> globalInstsToConsiderDeleting;
+    for (auto use : globalInstUsesToInline)
+    {
+        auto user = use->getUser();
+        IRBuilder builder(user);
+        builder.setInsertBefore(getOutsideASM(user));
+        IRCloneEnv cloneEnv;
+        auto val = maybeInlineGlobalValue(builder, use->getUser(), use->get(), cloneEnv);
+        if (val != use->get())
+        {
+            // Since certain globals that appear in the IR are considered illegal for all targets,
+            // e.g. calls to functions, we delete the globals we've inlined.
+            // Note that the inlining is done such that none of the descendants of the global will
+            // have any uses either.
+            globalInstsToConsiderDeleting.add(use->usedValue);
+
+            builder.replaceOperand(use, val);
+        }
+    }
+
+    for (auto globalInst : globalInstsToConsiderDeleting)
+    {
+        if (!globalInst->hasUses())
+            globalInst->removeAndDeallocate();
+    }
+}
+
+bool GlobalInstInliningContextGeneric::isLegalGlobalInst(IRInst* inst)
+{
+    switch (inst->getOp())
+    {
+    case kIROp_MakeStruct:
+    case kIROp_MakeArray:
+    case kIROp_MakeArrayFromElement:
+    case kIROp_MakeVector:
+    case kIROp_MakeMatrix:
+    case kIROp_MakeMatrixFromScalar:
+    case kIROp_MakeVectorFromScalar:
+        return true;
+    default:
+        if (as<IRConstant>(inst))
+            return true;
+        if (isLegalGlobalInstForTarget(inst))
+            return true;
+        return false;
+    }
+}
+
+bool GlobalInstInliningContextGeneric::isInlinableGlobalInst(IRInst* inst)
+{
+    switch (inst->getOp())
+    {
+    case kIROp_Add:
+    case kIROp_Sub:
+    case kIROp_Mul:
+    case kIROp_FRem:
+    case kIROp_IRem:
+    case kIROp_Lsh:
+    case kIROp_Rsh:
+    case kIROp_And:
+    case kIROp_Or:
+    case kIROp_Not:
+    case kIROp_Neg:
+    case kIROp_Div:
+    case kIROp_FieldExtract:
+    case kIROp_FieldAddress:
+    case kIROp_GetElement:
+    case kIROp_GetElementPtr:
+    case kIROp_GetOffsetPtr:
+    case kIROp_UpdateElement:
+    case kIROp_MakeTuple:
+    case kIROp_GetTupleElement:
+    case kIROp_MakeStruct:
+    case kIROp_MakeArray:
+    case kIROp_MakeArrayFromElement:
+    case kIROp_MakeVector:
+    case kIROp_MakeMatrix:
+    case kIROp_MakeMatrixFromScalar:
+    case kIROp_MakeVectorFromScalar:
+    case kIROp_swizzle:
+    case kIROp_swizzleSet:
+    case kIROp_MatrixReshape:
+    case kIROp_MakeString:
+    case kIROp_MakeResultError:
+    case kIROp_MakeResultValue:
+    case kIROp_GetResultError:
+    case kIROp_GetResultValue:
+    case kIROp_CastFloatToInt:
+    case kIROp_CastIntToFloat:
+    case kIROp_CastIntToPtr:
+    case kIROp_PtrCast:
+    case kIROp_CastPtrToBool:
+    case kIROp_CastPtrToInt:
+    case kIROp_BitAnd:
+    case kIROp_BitNot:
+    case kIROp_BitOr:
+    case kIROp_BitXor:
+    case kIROp_BitCast:
+    case kIROp_IntCast:
+    case kIROp_FloatCast:
+    case kIROp_Greater:
+    case kIROp_Less:
+    case kIROp_Geq:
+    case kIROp_Leq:
+    case kIROp_Neq:
+    case kIROp_Eql:
+    case kIROp_Call:
+        return true;
+    default:
+        if (isInlinableGlobalInstForTarget(inst))
+            return true;
+        return false;
+    }
+}
+
+bool GlobalInstInliningContextGeneric::shouldInlineInstImpl(IRInst* inst)
+{
+    // If 'inst' has an ancestor that is currently being inlined, then we
+    // better inline it since we'll be removing the ancestor.
+    bool ancestorShouldBeInlined = false;
+    for (IRInst* ancestor = inst->parent; ancestor != nullptr; ancestor = ancestor->parent)
+        if (m_mapGlobalInstToShouldInline.tryGetValue(inst, ancestorShouldBeInlined) &&
+            ancestorShouldBeInlined)
+            return true;
+
+    if (!isInlinableGlobalInst(inst))
+        return false;
+    if (isLegalGlobalInst(inst))
+    {
+        for (UInt i = 0; i < inst->getOperandCount(); i++)
+            if (shouldInlineInst(inst->getOperand(i)))
+                return true;
+        return false;
+    }
+    return true;
+}
+
+bool GlobalInstInliningContextGeneric::shouldInlineInst(IRInst* inst)
+{
+    bool result = false;
+    if (m_mapGlobalInstToShouldInline.tryGetValue(inst, result))
+        return result;
+    result = shouldInlineInstImpl(inst);
+    m_mapGlobalInstToShouldInline[inst] = result;
+    return result;
+}
+
+IRInst* GlobalInstInliningContextGeneric::inlineInst(
+    IRBuilder& builder,
+    IRCloneEnv& cloneEnv,
+    IRInst* inst)
+{
+    // We rely on this dictionary in order to force inlining of any nodes with that should be
+    // inlined
+    SLANG_ASSERT(m_mapGlobalInstToShouldInline[inst]);
+
+    IRInst* result;
+    if (cloneEnv.mapOldValToNew.tryGetValue(inst, result))
+        return result;
+
+    for (UInt i = 0; i < inst->getOperandCount(); i++)
+    {
+        auto operand = inst->getOperand(i);
+        IRBuilder operandBuilder(builder);
+        operandBuilder.setInsertBefore(getOutsideASM(builder.getInsertLoc().getInst()));
+        maybeInlineGlobalValue(operandBuilder, inst, operand, cloneEnv);
+    }
+    result = cloneInstAndOperands(&cloneEnv, &builder, inst);
+    cloneEnv.mapOldValToNew[inst] = result;
+    IRBuilder subBuilder(builder);
+    subBuilder.setInsertInto(result);
+    for (auto child : inst->getDecorations())
+    {
+        cloneInst(&cloneEnv, &subBuilder, child);
+    }
+    for (auto child : inst->getChildren())
+    {
+        m_mapGlobalInstToShouldInline[child] = true;
+        inlineInst(subBuilder, cloneEnv, child);
+    }
+    return result;
+}
+
+IRInst* GlobalInstInliningContextGeneric::maybeInlineGlobalValue(
+    IRBuilder& builder,
+    IRInst* user,
+    IRInst* inst,
+    IRCloneEnv& cloneEnv)
+{
+    if (!shouldInlineInst(inst))
+    {
+        switch (inst->getOp())
+        {
+        case kIROp_Func:
+        case kIROp_Specialize:
+        case kIROp_Generic:
+        case kIROp_LookupWitness:
+            return inst;
+        }
+        if (as<IRType>(inst))
+            return inst;
+
+        // If we encounter a global value that shouldn't be inlined, e.g. a const literal,
+        // we should insert a GlobalValueRef() inst to wrap around it, so all the dependent
+        // uses can be pinned to the function body.
+        auto result = inst;
+        bool shouldWrapGlobalRef = true;
+        if (!isLegalGlobalInst(user) && !getIROpInfo(user->getOp()).isHoistable())
+            shouldWrapGlobalRef = false;
+        else if (shouldBeInlinedForTarget(user))
+            shouldWrapGlobalRef = false;
+        if (shouldWrapGlobalRef)
+            result = builder.emitGlobalValueRef(inst);
+        cloneEnv.mapOldValToNew[inst] = result;
+        return result;
+    }
+
+    // If the global value is inlinable, we make all its operands avaialble locally, and
+    // then copy it to the local scope.
+    return inlineInst(builder, cloneEnv, inst);
+}
+
+} // namespace Slang

source/slang/slang-ir-legalize-global-values.h

@@ -0,0 +1,46 @@
+#pragma once
+
+#include "core/slang-dictionary.h"
+
+namespace Slang
+{
+struct IRBuilder;
+struct IRCloneEnv;
+struct IRInst;
+struct IRModule;
+
+struct GlobalInstInliningContextGeneric
+{
+    Dictionary<IRInst*, bool> m_mapGlobalInstToShouldInline;
+
+    // Target-specific control over how inlining happens
+    virtual bool isLegalGlobalInstForTarget(IRInst* inst) = 0;
+    virtual bool isInlinableGlobalInstForTarget(IRInst* inst) = 0;
+    virtual bool shouldBeInlinedForTarget(IRInst* user) = 0;
+    virtual IRInst* getOutsideASM(IRInst* beforeInst) = 0;
+
+    // Inline global values that can't represented by the target to their use sites.
+    // If this leaves any global unused, then remove it.
+    void inlineGlobalValuesAndRemoveIfUnused(IRModule* module);
+
+    // Opcodes that can exist in global scope, as long as the operands are.
+    bool isLegalGlobalInst(IRInst* inst);
+
+    // Opcodes that can be inlined into function bodies.
+    bool isInlinableGlobalInst(IRInst* inst);
+
+    bool shouldInlineInstImpl(IRInst* inst);
+
+    bool shouldInlineInst(IRInst* inst);
+
+    IRInst* inlineInst(IRBuilder& builder, IRCloneEnv& cloneEnv, IRInst* inst);
+
+    /// Inline `inst` in the local function body so they can be emitted as a local inst.
+    ///
+    IRInst* maybeInlineGlobalValue(
+        IRBuilder& builder,
+        IRInst* user,
+        IRInst* inst,
+        IRCloneEnv& cloneEnv);
+};
+} // namespace Slang