Skip to content

Commit

Permalink
WGSL: Fix issue where global calls are generated (#5768)
Browse files Browse the repository at this point in the history
* Split out SPIR-V -specific legalization of global functions

This is a refactoring and should not affect generated code.

* Move global inst inlining code into separate function

This is a refactoring and should not affect generated code.

* Take SPIR-V -specific parts out of GlobalInstInliningContext

This is a refactoring and should not affect generated code.

* Move 'inlineGlobalValues' to generic inlining context

This is a refactoring and should not affect generated code.

* Move 'setInsertBeforeOutsideASM' to generic inlining context

This is a refactoring and should not affect generated code.

* Move generic inlining context into own file

This is a refactoring and should not affect generated code.

* Run global inlining for WGSL as well

* Make the 'getOutsideASM' function generic as well

* Enable language-feature/constants/static-const-in-generic-interface.slang for WebGPU

* Clarify when it's safe to remove and deallocate an IRInst

* Remove globals if they're left unused after inlining

This closes #5607.

* Handle IRGlobalValueRef in C-like emitter

* format code

---------

Co-authored-by: slangbot <186143334+slangbot@users.noreply.github.com>
  • Loading branch information
aleino-nv and slangbot authored Dec 12, 2024
1 parent 79dc7ef commit 9c82ed3
Show file tree
Hide file tree
Showing 8 changed files with 358 additions and 212 deletions.
5 changes: 5 additions & 0 deletions source/slang/slang-emit-c-like.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
247 changes: 247 additions & 0 deletions source/slang/slang-ir-legalize-global-values.cpp
Original file line number Diff line number Diff line change
@@ -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
46 changes: 46 additions & 0 deletions source/slang/slang-ir-legalize-global-values.h
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit 9c82ed3

Please sign in to comment.