Enable NonUniformResourceIndex support for glsl, hlsl and spirv (#3899)

Fixes #387676* ForceInline SampleLevel to allow decorations to apply

* explictly add all the SPIRVAsmOperand Insts in non-differentiable list, which might get inadvertently processed when these functions are inlined into the main shader

* Support NonUniformResourceIndex for SPIR-V target

Fixes #3876

* add a new IR instruction for NonUniformResourceIndex

* slang ir emitter for nonuniform resource index

* update the hlsl meta slang

* Add test cases for NonUniformResourceIndex access for buffers and textures, with/without cast, nested access etc.

* add default c-like emitter for nonuniformresourceinfo

* added hlsl emitter

* added glsl emitter

* requisites for spirv enabling

- new decorator for nonuniformresourceindex
- emitter for nonuniformresourceindex signature change

* add hasResourceType checker

* add rwStructBuffType in resourcetype checker

* add a case for nonuniformres in emitDecorations

* DO NOT COMMIT: This change adds special handling for RWStructBuf within the isResourceType function, if it is a pointer to this resource, return true to make it work with nonuniformres test

* spirv emitter for decorations - update the emitLocalInst to perform decorations at the end

* added main spirv emitter code

* slang emit spirv bugfix

* hacky way of supporting Call Inst

* move code to cleanup nonuniform inst into helper function

* remove stale codefrom test

* add spirv decoration for nonuniform

* update test to remove global variables

* update coherent-2 test

* update comment for special handling

* update the spirv legalize to handle nested nonuniforms

improved logic that handles call ops, rwstructbuf, nested nonuniforms
etc.

* update nonuniform-array-of-tex test

* missed removing nonuniform inst causing duplicate decorations

* add glsl and hlsl variants of nonuniform tests

* repurpose the hasResource function into something specific for nonuniform inst decoration helper

* clean up comments and code around spirv-legalization to emit nonuniform inst by recursively looking into the inst

* use the helper canDecorateNonUniformInst to convert `nonUniformResourceInfo` inst to decoration

* converted compute/unbounded-array-of-array cross compile test into a simple check test

* update contains Resource helper function to be more generic

* clean up the case for opcall handling with nonuniform resource inst

* update ptr to struct buffer check to be more explicit and rename the function to check for ptr to resource type

* update comments and fix the test for coherent

* fix typos

* update logic on spirv legalize to delete dead instructions - for some reason this doesn't automatically happen

* add comments to declarations

* add NonuniformResourceIndex to the non-differential inst list

1 files changed, 121 insertions, 54 deletions

diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index b5a94cf0d..3a551f301 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -2588,6 +2588,7 @@ struct SPIRVEmitContext
         /// Emit an instruction that is local to the body of the given `parent`.
     SpvInst* emitLocalInst(SpvInstParent* parent, IRInst* inst)
     {
+        SpvInst* result = nullptr;
         switch( inst->getOp() )
         {
         default:
@@ -2600,65 +2601,91 @@ struct SPIRVEmitContext
             }
         case kIROp_Specialize:
         case kIROp_MissingReturn:
-            return nullptr;
+            break;
         case kIROp_Var:
-            return emitVar(parent, inst);
+            result = emitVar(parent, inst);
+            break;
         case kIROp_Call:
-            return emitCall(parent, static_cast<IRCall*>(inst));
+            result = emitCall(parent, static_cast<IRCall*>(inst));
+            break;
         case kIROp_FieldAddress:
-            return emitFieldAddress(parent, as<IRFieldAddress>(inst));
+            result = emitFieldAddress(parent, as<IRFieldAddress>(inst));
+            break;
         case kIROp_FieldExtract:
-            return emitFieldExtract(parent, as<IRFieldExtract>(inst));
+            result = emitFieldExtract(parent, as<IRFieldExtract>(inst));
+            break;
         case kIROp_GetElementPtr:
-            return emitGetElementPtr(parent, as<IRGetElementPtr>(inst));
+            result = emitGetElementPtr(parent, as<IRGetElementPtr>(inst));
+            break;
         case kIROp_GetOffsetPtr:
-            return emitGetOffsetPtr(parent, inst);
+            result = emitGetOffsetPtr(parent, inst);
+            break;
         case kIROp_GetElement:
-            return emitGetElement(parent, as<IRGetElement>(inst));
+            result = emitGetElement(parent, as<IRGetElement>(inst));
+            break;
         case kIROp_MakeStruct:
-            return emitCompositeConstruct(parent, inst);
+            result = emitCompositeConstruct(parent, inst);
+            break;
         case kIROp_MakeArrayFromElement:
-            return emitMakeArrayFromElement(parent, inst);
+            result = emitMakeArrayFromElement(parent, inst);
+            break;
         case kIROp_MakeMatrixFromScalar:
-            return emitMakeMatrixFromScalar(parent, inst);
+            result = emitMakeMatrixFromScalar(parent, inst);
+            break;
         case kIROp_MakeMatrix:
-            return emitMakeMatrix(parent, inst);
+            result = emitMakeMatrix(parent, inst);
+            break;
         case kIROp_Load:
-            return emitLoad(parent, as<IRLoad>(inst));
+            result = emitLoad(parent, as<IRLoad>(inst));
+            break;
         case kIROp_Store:
-            return emitStore(parent, as<IRStore>(inst));
+            result = emitStore(parent, as<IRStore>(inst));
+            break;
         case kIROp_SwizzledStore:
-            return emitSwizzledStore(parent, as<IRSwizzledStore>(inst));
+            result = emitSwizzledStore(parent, as<IRSwizzledStore>(inst));
+            break;
         case kIROp_swizzleSet:
-            return emitSwizzleSet(parent, as<IRSwizzleSet>(inst));
+            result = emitSwizzleSet(parent, as<IRSwizzleSet>(inst));
+            break;
         case kIROp_RWStructuredBufferGetElementPtr:
-            return emitStructuredBufferGetElementPtr(parent, inst);
+            result = emitStructuredBufferGetElementPtr(parent, inst);
+            break;
         case kIROp_StructuredBufferGetDimensions:
-            return emitStructuredBufferGetDimensions(parent, inst);
+            result = emitStructuredBufferGetDimensions(parent, inst);
+            break;
         case kIROp_swizzle:
-            return emitSwizzle(parent, as<IRSwizzle>(inst));
+            result = emitSwizzle(parent, as<IRSwizzle>(inst));
+            break;
         case kIROp_IntCast:
-            return emitIntCast(parent, as<IRIntCast>(inst));
+            result = emitIntCast(parent, as<IRIntCast>(inst));
+            break;
         case kIROp_FloatCast:
-            return emitFloatCast(parent, as<IRFloatCast>(inst));
+            result = emitFloatCast(parent, as<IRFloatCast>(inst));
+            break;
         case kIROp_CastIntToFloat:
-            return emitIntToFloatCast(parent, as<IRCastIntToFloat>(inst));
+            result = emitIntToFloatCast(parent, as<IRCastIntToFloat>(inst));
+            break;
         case kIROp_CastFloatToInt:
-            return emitFloatToIntCast(parent, as<IRCastFloatToInt>(inst));
+            result = emitFloatToIntCast(parent, as<IRCastFloatToInt>(inst));
+            break;
         case kIROp_CastPtrToInt:
-            return emitCastPtrToInt(parent, inst);
+            result = emitCastPtrToInt(parent, inst);
+            break;
         case kIROp_CastPtrToBool:
-            return emitCastPtrToBool(parent, inst);
+            result = emitCastPtrToBool(parent, inst);
+            break;
         case kIROp_CastIntToPtr:
-            return emitCastIntToPtr(parent, inst);
+            result = emitCastIntToPtr(parent, inst);
+            break;
         case kIROp_PtrCast:
         case kIROp_BitCast:
-            return emitOpBitcast(
+            result = emitOpBitcast(
                 parent,
                 inst,
                 inst->getDataType(),
                 inst->getOperand(0)
             );
+            break;
         case kIROp_Add:
         case kIROp_Sub:
         case kIROp_Mul:
@@ -2681,12 +2708,14 @@ struct SPIRVEmitContext
         case kIROp_Geq:
         case kIROp_Rsh:
         case kIROp_Lsh:
-            return emitArithmetic(parent, inst);
+            result = emitArithmetic(parent, inst);
+            break;
         case kIROp_GlobalValueRef:
             {
                 auto inner = ensureInst(inst->getOperand(0));
                 registerInst(inst, inner);
-                return inner;
+                result = inner;
+                break;
             }
         case kIROp_GetVulkanRayTracingPayloadLocation:
             {
@@ -2700,15 +2729,18 @@ struct SPIRVEmitContext
                 }
                 auto inner = ensureInst(location);
                 registerInst(inst, inner);
-                return inner;
+                result = inner;
+                break;
             }
         case kIROp_Return:
             if (as<IRReturn>(inst)->getVal()->getOp() == kIROp_VoidLit)
-                return emitOpReturn(parent, inst);
+                result = emitOpReturn(parent, inst);
             else
-                return emitOpReturnValue(parent, inst, as<IRReturn>(inst)->getVal());
+                result = emitOpReturnValue(parent, inst, as<IRReturn>(inst)->getVal());
+            break;
         case kIROp_discard:
-            return emitOpKill(parent, inst);
+            result = emitOpKill(parent, inst);
+            break;
         case kIROp_unconditionalBranch:
             {
                 // If we are jumping to the main block of a loop,
@@ -2719,7 +2751,8 @@ struct SPIRVEmitContext
                 if (isLoopTargetBlock(targetBlock, loopInst))
                     return emitOpBranch(parent, inst, getIRInstSpvID(loopInst));
                 // Otherwise, emit a normal branch inst into the target block.
-                return emitOpBranch(parent, inst, getIRInstSpvID(targetBlock));
+                result = emitOpBranch(parent, inst, getIRInstSpvID(targetBlock));
+                break;
             }
         case kIROp_loop:
             {
@@ -2735,7 +2768,8 @@ struct SPIRVEmitContext
                 // from the actual loop target block) are emitted first.
                 emitOpBranch(parent, nullptr, blockId);
         
-                return block;
+                result = block;
+                break;
             }
         case kIROp_ifElse:
             {
@@ -2743,7 +2777,7 @@ struct SPIRVEmitContext
                 auto afterBlockID = getIRInstSpvID(ifelseInst->getAfterBlock());
                 emitOpSelectionMerge(parent, nullptr, afterBlockID, SpvSelectionControlMaskNone);
                 auto falseLabel = ifelseInst->getFalseBlock();
-                return emitOpBranchConditional(
+                result = emitOpBranchConditional(
                     parent,
                     inst,
                     ifelseInst->getCondition(),
@@ -2751,13 +2785,14 @@ struct SPIRVEmitContext
                     falseLabel ? getID(ensureInst(falseLabel)) : afterBlockID,
                     makeArray<SpvLiteralInteger>()
                 );
+                break;
             }
         case kIROp_Switch:
             {
                 auto switchInst = as<IRSwitch>(inst);
                 auto mergeBlockID = getIRInstSpvID(switchInst->getBreakLabel());
                 emitOpSelectionMerge(parent, nullptr, mergeBlockID, SpvSelectionControlMaskNone);
-                return emitInstCustomOperandFunc(parent, inst, SpvOpSwitch, [&]() {
+                result = emitInstCustomOperandFunc(parent, inst, SpvOpSwitch, [&]() {
                     emitOperand(switchInst->getCondition());
                     auto defaultLabel = switchInst->getDefaultLabel();
                     emitOperand(defaultLabel ? getID(ensureInst(defaultLabel)) : mergeBlockID);
@@ -2771,13 +2806,17 @@ struct SPIRVEmitContext
                         emitOperand(caseLabel ? getID(ensureInst(caseLabel)) : mergeBlockID);
                     }
                 });
+                break;
             }
         case kIROp_Unreachable:
-            return emitOpUnreachable(parent, inst);
+            result = emitOpUnreachable(parent, inst);
+            break;
         case kIROp_conditionalBranch:
             SLANG_UNEXPECTED("Unstructured branching is not supported by SPIRV.");
+            break;
         case kIROp_MakeVector:
-            return emitConstruct(parent, inst);
+            result = emitConstruct(parent, inst);
+            break;
         case kIROp_MakeVectorFromScalar:
             {
                 const auto scalar = inst->getOperand(0);
@@ -2785,45 +2824,62 @@ struct SPIRVEmitContext
                 SLANG_ASSERT(vecTy);
                 const auto numElems = as<IRIntLit>(vecTy->getElementCount());
                 SLANG_ASSERT(numElems);
-                return emitSplat(parent, inst, scalar, numElems->getValue());
+                result = emitSplat(parent, inst, scalar, numElems->getValue());
             }
+            break;
         case kIROp_MakeArray:
-            return emitConstruct(parent, inst);
+            result = emitConstruct(parent, inst);
+            break;
         case kIROp_Select:
-            return emitInst(parent, inst, SpvOpSelect, inst->getFullType(), kResultID, OperandsOf(inst));
+            result = emitInst(parent, inst, SpvOpSelect, inst->getFullType(), kResultID, OperandsOf(inst));
+            break;
         case kIROp_DebugLine:
-            return emitDebugLine(parent, as<IRDebugLine>(inst));
+            result = emitDebugLine(parent, as<IRDebugLine>(inst));
+            break;
         case kIROp_DebugVar:
-            return emitDebugVar(parent, as<IRDebugVar>(inst));
+            result = emitDebugVar(parent, as<IRDebugVar>(inst));
+            break;
         case kIROp_DebugValue:
-            return emitDebugValue(parent, as<IRDebugValue>(inst));
+            result = emitDebugValue(parent, as<IRDebugValue>(inst));
+            break;
         case kIROp_GetStringHash:
-            return emitGetStringHash(inst);
+            result = emitGetStringHash(inst);
+            break;
         case kIROp_undefined:
-            return emitOpUndef(parent, inst, inst->getDataType());
+            result = emitOpUndef(parent, inst, inst->getDataType());
+            break;
         case kIROp_SPIRVAsm:
-            return emitSPIRVAsm(parent, as<IRSPIRVAsm>(inst));
+            result = emitSPIRVAsm(parent, as<IRSPIRVAsm>(inst));
+            break;
         case kIROp_ImageLoad:
-            return emitImageLoad(parent, as<IRImageLoad>(inst));
+            result = emitImageLoad(parent, as<IRImageLoad>(inst));
+            break;
         case kIROp_ImageStore:
-            return emitImageStore(parent, as<IRImageStore>(inst));
+            result = emitImageStore(parent, as<IRImageStore>(inst));
+            break;
         case kIROp_ImageSubscript:
-            return emitImageSubscript(parent, as<IRImageSubscript>(inst));
+            result = emitImageSubscript(parent, as<IRImageSubscript>(inst));
+            break;
         case kIROp_AtomicCounterIncrement:
             {
                 IRBuilder builder{inst};
                 const auto memoryScope =  emitIntConstant(IRIntegerValue{SpvScopeDevice}, builder.getUIntType());
                 const auto memorySemantics =  emitIntConstant(IRIntegerValue{SpvMemorySemanticsMaskNone}, builder.getUIntType());
-                return emitOpAtomicIIncrement(parent, inst, inst->getFullType(), inst->getOperand(0), memoryScope, memorySemantics);
+                result = emitOpAtomicIIncrement(parent, inst, inst->getFullType(), inst->getOperand(0), memoryScope, memorySemantics);
             }
+            break;
         case kIROp_AtomicCounterDecrement:
             {
                 IRBuilder builder{inst};
                 const auto memoryScope =  emitIntConstant(IRIntegerValue{SpvScopeDevice}, builder.getUIntType());
                 const auto memorySemantics =  emitIntConstant(IRIntegerValue{SpvMemorySemanticsMaskNone}, builder.getUIntType());
-                return emitOpAtomicIDecrement(parent, inst, inst->getFullType(), inst->getOperand(0), memoryScope, memorySemantics);
+                result = emitOpAtomicIDecrement(parent, inst, inst->getFullType(), inst->getOperand(0), memoryScope, memorySemantics);
             }
+            break;
         }
+        if (result)
+            emitDecorations(inst, getID(result));
+        return result;
     }
 
     SpvInst* emitImageLoad(SpvInstParent* parent, IRImageLoad* load)
@@ -2925,7 +2981,6 @@ struct SPIRVEmitContext
         }
     }
 
-
     SpvExecutionMode getDepthOutputExecutionMode(IRInst* builtinVar)
     {
         SpvExecutionMode result = SpvExecutionModeMax;
@@ -3301,6 +3356,18 @@ struct SPIRVEmitContext
             }
             break;
 
+        case kIROp_SPIRVNonUniformResourceDecoration:
+            {
+                requireSPIRVCapability(SpvCapabilityShaderNonUniform);
+                emitOpDecorate(
+                    getSection(SpvLogicalSectionID::Annotations),
+                    decoration,
+                    dstID,
+                    SpvDecorationNonUniform
+                );
+            }
+            break;
+
         case kIROp_OutputTopologyDecoration:
             {
                 const auto o = cast<IROutputTopologyDecoration>(decoration);