Fix Varying Variable Location Assignments With Hull Shaders (#4915)

* Fix Varying Variable Location Assignments With Hull Shaders

Fixes: #4913
Fixes: #4540

Changes:
1. Added `kIROp_ControlBarrier` to HLSL/GLSL emitting.
2. Added a method to track 'used' and 'unused' varyings for when legalizing GLSL. This allows us to assign correct offsets to automatically added varyings
    * Added a `ZeroLSB` check to UIntSet for this purpose

* add missing return

* code comment adjustment

* cleanup

* comment and HLSL controlBarrier mistake

* assume space for glsl/spriv varying is irrelevant

6 files changed, 72 insertions, 12 deletions

diff --git a/source/core/slang-uint-set.cpp b/source/core/slang-uint-set.cpp
index ba71254e1..457915449 100644
--- a/source/core/slang-uint-set.cpp
+++ b/source/core/slang-uint-set.cpp
@@ -3,6 +3,27 @@
 namespace Slang
 {
 
+Index UIntSet::getLSBZero()
+{
+    uint64_t offset = 0;
+    for (Element& element : this->m_buffer)
+    {
+        // Flip all bits so bitscanForward can find a 0 bit
+        Element flippedElement = ~element;
+
+        // continue if we don't have 0 bits
+        if (flippedElement == 0)
+        {
+            offset += sizeof(Element) * 8;
+            continue;
+        }
+
+        // Get LSBZero of current Block, add with offset
+        return bitscanForward(flippedElement) + offset;
+    }
+    return offset;
+}
+
 UIntSet& UIntSet::operator=(UIntSet&& other)
 {
     m_buffer = _Move(other.m_buffer);
diff --git a/source/core/slang-uint-set.h b/source/core/slang-uint-set.h
index 4ba067871..3531e5cfb 100644
--- a/source/core/slang-uint-set.h
+++ b/source/core/slang-uint-set.h
@@ -133,6 +133,9 @@ public:
         /// Returns true if set1 and set2 have a same value set (ie there is an intersection)
     static bool hasIntersection(const UIntSet& set1, const UIntSet& set2);
 
+    /// Get LSB Zero of UIntSet. LSB Zero is the smallest value missing from this UIntSet.
+    Index getLSBZero();
+
     struct Iterator
     {
         friend class UIntSet;
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index 50d71b6b6..902030791 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -1727,6 +1727,11 @@ bool GLSLSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOu
 {
     switch (inst->getOp())
     {
+        case kIROp_ControlBarrier:
+        {
+            m_writer->emit("barrier();\n");
+            return true;
+        }
         case kIROp_MakeVectorFromScalar:
         case kIROp_MatrixReshape:
         {
diff --git a/source/slang/slang-emit-hlsl.cpp b/source/slang/slang-emit-hlsl.cpp
index d3757f534..a867d4c06 100644
--- a/source/slang/slang-emit-hlsl.cpp
+++ b/source/slang/slang-emit-hlsl.cpp
@@ -500,6 +500,11 @@ bool HLSLSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOu
 {
     switch (inst->getOp())
     {
+        case kIROp_ControlBarrier:
+        {
+            m_writer->emit("GroupMemoryBatrierWithGroupSync();\n");
+            return true;
+        }
         case kIROp_MakeVector:
         case kIROp_MakeMatrix:
         {
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 044f79531..e91209108 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -1137,6 +1137,10 @@ Result linkAndOptimizeIR(
             ? as<GLSLExtensionTracker>(options.sourceEmitter->getExtensionTracker())
             : &glslExtensionTracker;
 
+#if 0
+            dumpIRIfEnabled(codeGenContext, irModule, "PRE GLSL LEGALIZED");
+#endif
+
         legalizeEntryPointsForGLSL(
             session,
             irModule,
diff --git a/source/slang/slang-ir-glsl-legalize.cpp b/source/slang/slang-ir-glsl-legalize.cpp
index f2e2fa7cb..0242076f5 100644
--- a/source/slang/slang-ir-glsl-legalize.cpp
+++ b/source/slang/slang-ir-glsl-legalize.cpp
@@ -259,11 +259,20 @@ List<IRInst*> ScalarizedVal::leafAddresses()
 
 struct GLSLLegalizationContext
 {
-    Session*                session;
-    GLSLExtensionTracker*   glslExtensionTracker;
-    DiagnosticSink*         sink;
-    Stage                   stage;
-    IRFunc*                 entryPointFunc;
+    Session*                   session;
+    GLSLExtensionTracker*      glslExtensionTracker;
+    DiagnosticSink*            sink;
+    Stage                      stage;
+    IRFunc*                    entryPointFunc;
+    
+    /// This dictionary stores all bindings of 'VaryingIn/VaryingOut'. We assume 'space' is 0.
+    Dictionary<LayoutResourceKind, UIntSet> usedBindingIndex;
+
+    GLSLLegalizationContext()
+    {
+        // Reserve for VaryingInput VaryingOutput
+        usedBindingIndex.reserve(2);
+    }
 
     struct SystemSemanticGlobal
     {
@@ -846,6 +855,7 @@ struct OuterParamInfoLink
 };
 
 void createVarLayoutForLegalizedGlobalParam(
+    GLSLLegalizationContext* context,
     IRInst* globalParam,
     IRBuilder* builder,
     IRVarLayout* inVarLayout,
@@ -857,6 +867,8 @@ void createVarLayoutForLegalizedGlobalParam(
     OuterParamInfoLink* outerParamInfo,
     GLSLSystemValueInfo* systemValueInfo)
 {
+    context->usedBindingIndex[kind].add(bindingIndex);
+
     // We need to construct a fresh layout for the variable, even
     // if the original had its own layout, because it might be
     // an `inout` parameter, and we only want to deal with the case
@@ -953,7 +965,7 @@ IRInst* getOrCreateBuiltinParamForHullShader(GLSLLegalizationContext* context, U
     return outputControlPointIdParam;
 }
 
-IRTypeLayout* createPatchConstantFuncResultTypeLayout(IRBuilder& irBuilder, IRType* type)
+IRTypeLayout* createPatchConstantFuncResultTypeLayout(GLSLLegalizationContext* context, IRBuilder& irBuilder, IRType* type)
 {
     if (auto structType = as<IRStructType>(type))
     {
@@ -961,8 +973,7 @@ IRTypeLayout* createPatchConstantFuncResultTypeLayout(IRBuilder& irBuilder, IRTy
         for (auto field : structType->getFields())
         {
             auto fieldType = field->getFieldType();
-
-            IRTypeLayout* fieldTypeLayout = createPatchConstantFuncResultTypeLayout(irBuilder, fieldType);
+            IRTypeLayout* fieldTypeLayout = createPatchConstantFuncResultTypeLayout(context, irBuilder, fieldType);
             IRVarLayout::Builder fieldVarLayoutBuilder(&irBuilder, fieldTypeLayout);
             auto decoration = field->getKey()->findDecoration<IRSemanticDecoration>();
             if (decoration)
@@ -970,6 +981,17 @@ IRTypeLayout* createPatchConstantFuncResultTypeLayout(IRBuilder& irBuilder, IRTy
                 if (decoration->getSemanticName().startsWithCaseInsensitive(toSlice("sv_")))
                     fieldVarLayoutBuilder.setSystemValueSemantic(decoration->getSemanticName(), 0);
             }
+            else
+            {
+                auto varLayoutForKind = fieldVarLayoutBuilder.findOrAddResourceInfo(LayoutResourceKind::VaryingOutput);
+                
+                UInt space = 0;
+                varLayoutForKind->space = space;
+
+                auto unusedBinding = context->usedBindingIndex[LayoutResourceKind::VaryingOutput].getLSBZero();
+                varLayoutForKind->offset = unusedBinding;
+                context->usedBindingIndex[LayoutResourceKind::VaryingOutput].add(unusedBinding);
+            }
             builder.addField(field->getKey(), fieldVarLayoutBuilder.build());
         }
         auto typeLayout = builder.build();
@@ -977,7 +999,7 @@ IRTypeLayout* createPatchConstantFuncResultTypeLayout(IRBuilder& irBuilder, IRTy
     }
     else if (auto arrayType = as<IRArrayTypeBase>(type))
     {
-        auto elementTypeLayout = createPatchConstantFuncResultTypeLayout(irBuilder, arrayType->getElementType());
+        auto elementTypeLayout = createPatchConstantFuncResultTypeLayout(context, irBuilder, arrayType->getElementType());
         IRArrayTypeLayout::Builder builder(&irBuilder, elementTypeLayout);
         return builder.build();
     }
@@ -1102,7 +1124,7 @@ void invokePathConstantFuncInHullShader(GLSLLegalizationContext* context, CodeGe
     builder.setInsertBefore(constantFunc->getFirstBlock()->getFirstOrdinaryInst());
     
     auto constantOutputType = constantFunc->getResultType();
-    IRTypeLayout* constantOutputLayout = createPatchConstantFuncResultTypeLayout(builder, constantOutputType);
+    IRTypeLayout* constantOutputLayout = createPatchConstantFuncResultTypeLayout(context, builder, constantOutputType);
     IRVarLayout::Builder resultVarLayoutBuilder(&builder, constantOutputLayout);
     if (auto semanticDecor = constantFunc->findDecoration<IRSemanticDecoration>())
         resultVarLayoutBuilder.setSystemValueSemantic(semanticDecor->getSemanticName(), 0);
@@ -1246,7 +1268,7 @@ ScalarizedVal createSimpleGLSLGlobalVarying(
             builder->addImportDecoration(globalParam, systemValueName);
 
             createVarLayoutForLegalizedGlobalParam(
-                globalParam, builder, inVarLayout, inTypeLayout, kind, bindingIndex, bindingSpace, declarator, outerParamInfo, systemValueInfo);
+                context, globalParam, builder, inVarLayout, inTypeLayout, kind, bindingIndex, bindingSpace, declarator, outerParamInfo, systemValueInfo);
 
             semanticGlobalTmp.globalParam = globalParam;
 
@@ -1382,7 +1404,7 @@ ScalarizedVal createSimpleGLSLGlobalVarying(
     }
 
     createVarLayoutForLegalizedGlobalParam(
-        globalParam, builder, inVarLayout, typeLayout, kind, bindingIndex, bindingSpace, declarator, outerParamInfo, systemValueInfo);
+        context, globalParam, builder, inVarLayout, typeLayout, kind, bindingIndex, bindingSpace, declarator, outerParamInfo, systemValueInfo);
     return val;
 }