Type layouts for structured buffers with counters (#3269)

* More tests for append structured buffer

* Append and Consume structured buffer tests for DX12

* neaten

* test wobble

* Add counter layout information to append/consume structured buffers

* add getRWStructuredBufferType

* Correct definition of get size for append/consume structured buffers

* tweak append structured buffer test

* Allow initializing counter buffer in render test

* vulkan test for consume structured buffer

* Handle null counterVarLayout in getExplicitCounterBindingRangeOffset

* remove dead code

* Implement atomic counter increment/decrement for spirv

* explicit spirv test

* Add missing check on result

* Hold on to counter resources

---------

Co-authored-by: Yong He <yonghe@outlook.com>

9 files changed, 290 insertions, 41 deletions

diff --git a/source/slang/slang-ast-builder.cpp b/source/slang/slang-ast-builder.cpp
index 6a9aad257..ce13ab650 100644
--- a/source/slang/slang-ast-builder.cpp
+++ b/source/slang/slang-ast-builder.cpp
@@ -358,6 +358,11 @@ HLSLStructuredBufferType* ASTBuilder::getStructuredBufferType(Type* elementType)
     return as<HLSLStructuredBufferType>(getSpecializedBuiltinType(elementType, "HLSLStructuredBufferType"));
 }
 
+HLSLRWStructuredBufferType* ASTBuilder::getRWStructuredBufferType(Type* elementType)
+{
+    return as<HLSLRWStructuredBufferType>(getSpecializedBuiltinType(elementType, "HLSLRWStructuredBufferType"));
+}
+
 SamplerStateType* ASTBuilder::getSamplerStateType()
 {
     return as<SamplerStateType>(getSpecializedBuiltinType(nullptr, "HLSLStructuredBufferType"));
diff --git a/source/slang/slang-ast-builder.h b/source/slang/slang-ast-builder.h
index a5e1cd40c..aff3088ab 100644
--- a/source/slang/slang-ast-builder.h
+++ b/source/slang/slang-ast-builder.h
@@ -473,6 +473,8 @@ public:
 
     HLSLStructuredBufferType* getStructuredBufferType(Type* elementType);
 
+    HLSLRWStructuredBufferType* getRWStructuredBufferType(Type* elementType);
+
     SamplerStateType* getSamplerStateType();
 
     DifferentialPairType* getDifferentialPairType(
diff --git a/source/slang/slang-emit-spirv-ops.h b/source/slang/slang-emit-spirv-ops.h
index f64fc7f9a..f1a9ff3e4 100644
--- a/source/slang/slang-emit-spirv-ops.h
+++ b/source/slang/slang-emit-spirv-ops.h
@@ -2313,4 +2313,56 @@ SpvInst* emitOpExecutionModeIdLocalSizeId(
         zSize
     );
 }
+
+template<typename T1, typename T2, typename T3, typename T4>
+SpvInst* emitOpAtomicIIncrement(
+    SpvInstParent* parent,
+    IRInst* inst,
+    const T1& idResultType,
+    const T2& pointer,
+    const T3& memory,
+    const T4& semantics
+)
+{
+    static_assert(isSingular<T1>);
+    static_assert(isSingular<T2>);
+    static_assert(isSingular<T3>);
+    static_assert(isSingular<T4>);
+    return emitInst(
+        parent,
+        inst,
+        SpvOpAtomicIIncrement,
+        idResultType,
+        kResultID,
+        pointer,
+        memory,
+        semantics
+    );
+}
+
+template<typename T1, typename T2, typename T3, typename T4>
+SpvInst* emitOpAtomicIDecrement(
+    SpvInstParent* parent,
+    IRInst* inst,
+    const T1& idResultType,
+    const T2& pointer,
+    const T3& memory,
+    const T4& semantics
+)
+{
+    static_assert(isSingular<T1>);
+    static_assert(isSingular<T2>);
+    static_assert(isSingular<T3>);
+    static_assert(isSingular<T4>);
+    return emitInst(
+        parent,
+        inst,
+        SpvOpAtomicIDecrement,
+        idResultType,
+        kResultID,
+        pointer,
+        memory,
+        semantics
+    );
+}
 #endif        // SLANG_IN_SPIRV_EMIT_CONTEXT
diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 16a88b8af..aa21127e9 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -2390,6 +2390,20 @@ struct SPIRVEmitContext
             return emitImageStore(parent, as<IRImageStore>(inst));
         case kIROp_ImageSubscript:
             return emitImageSubscript(parent, as<IRImageSubscript>(inst));
+        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);
+            }
+        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);
+            }
         }
     }
 
diff --git a/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp b/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp
index 9df5f4e7e..ed04541ef 100644
--- a/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp
+++ b/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp
@@ -154,14 +154,15 @@ namespace Slang
             builder.setInsertInto(getDimensionsFunc);
             builder.emitBlock();
             auto bufferParam = builder.emitParam(structType);
-            auto counterBuffer = builder.emitFieldExtract(counterBufferType, bufferParam, counterBufferKey);
-            IRInst* getCounterPtrArgs[] = { counterBuffer, builder.getIntValue(builder.getIntType(), 0) };
-            auto counterBufferPtr = builder.emitIntrinsicInst(builder.getPtrType(builder.getIntType()), kIROp_RWStructuredBufferGetElementPtr, 2, getCounterPtrArgs);
-            auto counter = builder.emitLoad(counterBufferPtr);
-            counter = builder.emitCast(builder.getUIntType(), counter);
-            auto stride = builder.getIntValue(builder.getUIntType(), elementSize.getStride());
-            IRInst* vecArgs[] = { counter, stride };
-            builder.emitReturn(builder.emitMakeVector(uint2Type, 2, vecArgs));
+            auto elementBuffer = builder.emitFieldExtract(elementBufferType, bufferParam, elementBufferKey);
+
+            const auto dim = builder.emitIntrinsicInst(
+                uint2Type,
+                kIROp_StructuredBufferGetDimensions,
+                1,
+                &elementBuffer
+            );
+            builder.emitReturn(dim);
         }
 
         // Replace all insts with synthesized functions.
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index c90230c1f..2813918b6 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -8633,6 +8633,23 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         return v;
     }
 
+    void addSpecializedForTargetDecorations(IRInst* inst, Decl* decl)
+    {
+        // If this declaration was marked as being an intrinsic for a particular
+        // target, then we should reflect that here.
+        for (auto targetMod : decl->getModifiersOfType<SpecializedForTargetModifier>())
+        {
+            // `targetMod` indicates that this particular declaration represents
+            // a specialized definition of the particular function for the given
+            // target, and we need to reflect that at the IR level.
+
+            auto targetName = targetMod->targetToken.getContent();
+            auto targetCap = findCapabilityAtom(targetName);
+
+            getBuilder()->addTargetDecoration(inst, CapabilitySet(targetCap));
+        }
+    }
+
     // Attach target-intrinsic decorations to an instruction,
     // based on modifiers on an AST declaration.
     void addTargetIntrinsicDecorations(
@@ -9170,19 +9187,7 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
 
         getBuilder()->addHighLevelDeclDecoration(irFunc, decl);
 
-        // If this declaration was marked as being an intrinsic for a particular
-        // target, then we should reflect that here.
-        for (auto targetMod : decl->getModifiersOfType<SpecializedForTargetModifier>())
-        {
-            // `targetMod` indicates that this particular declaration represents
-            // a specialized definition of the particular function for the given
-            // target, and we need to reflect that at the IR level.
-
-            auto targetName = targetMod->targetToken.getContent();
-            auto targetCap = findCapabilityAtom(targetName);
-
-            getBuilder()->addTargetDecoration(irFunc, CapabilitySet(targetCap));
-        }
+        addSpecializedForTargetDecorations(irFunc, decl);
 
         // If this declaration was marked as having a target-specific lowering
         // for a particular target, then handle that here.
diff --git a/source/slang/slang-reflection-api.cpp b/source/slang/slang-reflection-api.cpp
index a52b7b886..9d687245f 100644
--- a/source/slang/slang-reflection-api.cpp
+++ b/source/slang/slang-reflection-api.cpp
@@ -935,6 +935,14 @@ SLANG_API SlangInt spReflectionTypeLayout_findFieldIndexByName(SlangReflectionTy
     return -1;
 }
 
+SLANG_API SlangReflectionVariableLayout* spReflectionTypeLayout_GetExplicitCounter(SlangReflectionTypeLayout* inTypeLayout)
+{
+    const auto typeLayout = convert(inTypeLayout);
+    if(const auto structuredBufferTypeLayout = as<StructuredBufferTypeLayout>(typeLayout))
+        return (SlangReflectionVariableLayout*) structuredBufferTypeLayout->counterVarLayout.Ptr();
+    return nullptr;
+}
+
 SLANG_API size_t spReflectionTypeLayout_GetElementStride(SlangReflectionTypeLayout* inTypeLayout, SlangParameterCategory category)
 {
     auto typeLayout = convert(inTypeLayout);
@@ -1784,6 +1792,75 @@ namespace Slang
                 m_extendedInfo->m_bindingRanges.add(bindingRange);
                 m_extendedInfo->m_subObjectRanges.add(subObjectRange);
             }
+            else if(const auto structuredBufferTypeLayout = as<StructuredBufferTypeLayout>(typeLayout))
+            {
+                // For structured buffers we expect them to consume a single
+                // resource descriptor slot (not counting the possible counter
+                // buffer)
+                SLANG_ASSERT(typeLayout->resourceInfos.getCount() == 1);
+                const auto& resInfo = typeLayout->resourceInfos[0];
+
+                const auto bindingType = as<HLSLStructuredBufferType>(typeLayout->getType())
+                    ? SLANG_BINDING_TYPE_RAW_BUFFER
+                    : SLANG_BINDING_TYPE_MUTABLE_RAW_BUFFER;
+
+                // We now allocate a descriptor range for this buffer
+                TypeLayout::ExtendedInfo::DescriptorRangeInfo descriptorRange;
+                descriptorRange.kind = resInfo.kind;
+                descriptorRange.bindingType = bindingType;
+                // Note that we don't use resInfo.count here, as each
+                // structuredBufferType is essentially a struct of 2 fields
+                // (elements, counter) and not an array of length 2.
+                SLANG_ASSERT(resInfo.count != 2 || structuredBufferTypeLayout->counterVarLayout);
+                SLANG_ASSERT(resInfo.count != 1 || !structuredBufferTypeLayout->counterVarLayout);
+                descriptorRange.count = multiplier;
+                descriptorRange.indexOffset = _calcIndexOffset(path.primary, resInfo.kind);
+
+                Int descriptorSetIndex = _findOrAddDescriptorSet(_calcSpaceOffset(path.primary, resInfo.kind));
+                const RefPtr<TypeLayout::ExtendedInfo::DescriptorSetInfo> descriptorSet
+                    = m_extendedInfo->m_descriptorSets[descriptorSetIndex];
+                auto descriptorRangeIndex = descriptorSet->descriptorRanges.getCount();
+                descriptorSet->descriptorRanges.add(descriptorRange);
+
+                // We will map the elements buffer to a single binding range
+                TypeLayout::ExtendedInfo::BindingRangeInfo bindingRange;
+                bindingRange.leafTypeLayout = typeLayout;
+                bindingRange.leafVariable = path.primary ? path.primary->var->getVariable() : nullptr;
+                bindingRange.bindingType = bindingType;
+                bindingRange.count = multiplier;
+                bindingRange.descriptorSetIndex = descriptorSetIndex;
+                bindingRange.firstDescriptorRangeIndex = descriptorRangeIndex;
+                bindingRange.descriptorRangeCount = 1;
+
+                auto bindingRangeIndex = m_extendedInfo->m_bindingRanges.getCount();
+                m_extendedInfo->m_bindingRanges.add(bindingRange);
+
+                // We also make sure to report it as a sub-object range.
+                TypeLayout::ExtendedInfo::SubObjectRangeInfo subObjectRange;
+                subObjectRange.bindingRangeIndex = bindingRangeIndex;
+                subObjectRange.offsetVarLayout = createOffsetVarLayout(typeLayout, path);
+                subObjectRange.spaceOffset = 0;
+                m_extendedInfo->m_subObjectRanges.add(subObjectRange);
+
+                // If we have an associated counter for this structured buffer,
+                // add its ranges
+                if(structuredBufferTypeLayout->counterVarLayout)
+                {
+                    ExtendedBindingRangePath counterPath(
+                        path,
+                        structuredBufferTypeLayout->counterVarLayout
+                    );
+                    // This should always be 1, because it comes after the
+                    // single binding range we just added
+                    structuredBufferTypeLayout->counterVarLayout->bindingRangeOffset =
+                        m_extendedInfo->m_bindingRanges.getCount() - bindingRangeIndex;
+                    addRangesRec(
+                        structuredBufferTypeLayout->counterVarLayout->typeLayout,
+                        counterPath,
+                        multiplier
+                    );
+                }
+            }
             else
             {
                 // Here we have the catch-all case that handles "leaf" fields
@@ -1940,19 +2017,7 @@ namespace Slang
                     bindingRange.descriptorRangeCount++;
                 }
 
-                auto bindingRangeIndex = m_extendedInfo->m_bindingRanges.getCount();
-
                 m_extendedInfo->m_bindingRanges.add(bindingRange);
-
-                // For `StructuredBuffer` fields, we also make sure to report it as a sub-object range.
-                if (const auto structuredBufferTypeLayout = as<StructuredBufferTypeLayout>(typeLayout))
-                {
-                    TypeLayout::ExtendedInfo::SubObjectRangeInfo subObjectRange;
-                    subObjectRange.bindingRangeIndex = bindingRangeIndex;
-                    subObjectRange.offsetVarLayout = createOffsetVarLayout(typeLayout, path);
-                    subObjectRange.spaceOffset = 0;
-                    m_extendedInfo->m_subObjectRanges.add(subObjectRange);
-                }
             }
         }
     };
@@ -2298,6 +2363,22 @@ SLANG_API SlangInt spReflectionTypeLayout_getFieldBindingRangeOffset(SlangReflec
     return 0;
 }
 
+SLANG_API SlangInt spReflectionTypeLayout_getExplicitCounterBindingRangeOffset(SlangReflectionTypeLayout* inTypeLayout)
+{
+    auto typeLayout = convert(inTypeLayout);
+    if(!typeLayout) return 0;
+
+    if(const auto structuredBufferTypeLayout = as<StructuredBufferTypeLayout>(typeLayout))
+    {
+        getExtendedTypeLayout(structuredBufferTypeLayout);
+        return structuredBufferTypeLayout->counterVarLayout
+            ? structuredBufferTypeLayout->counterVarLayout->bindingRangeOffset
+            : 0;
+    }
+
+    return 0;
+}
+
 #if 0
 SLANG_API SlangInt spReflectionTypeLayout_getSubObjectRangeCount(SlangReflectionTypeLayout* inTypeLayout)
 {
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index 978fa6fbb..1534b6cc5 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -228,6 +228,11 @@ struct DefaultLayoutRulesImpl : SimpleLayoutRulesImpl
         // and instead it will be injected in the concrete implementations
         // that require it.
     }
+
+    bool DoStructuredBuffersNeedSeparateCounterBuffer() override
+    {
+        return true;
+    }
 };
 
     /// Common behavior for GLSL-family layout.
@@ -600,6 +605,14 @@ struct HLSLStructuredBufferLayoutRulesImpl : DefaultLayoutRulesImpl
     // but retain the rules around not adjusting the size of an array or
     // structure to its alignment. In this way they should match our
     // default layout rules.
+    //
+    // DirectX does however allow transparently managing the counter buffer
+    // resource for StructuredBuffers.
+
+    bool DoStructuredBuffersNeedSeparateCounterBuffer() override
+    {
+        return false;
+    }
 };
 
 struct DefaultVaryingLayoutRulesImpl : DefaultLayoutRulesImpl
@@ -733,8 +746,7 @@ struct GLSLObjectLayoutRulesImpl : ObjectLayoutRulesImpl
 
         // In Vulkan GLSL, pretty much every object is just a descriptor-table slot.
         // Except for AppendConsumeStructuredBuffer, which takes two slots.
-        if (kind == ShaderParameterKind::AppendConsumeStructuredBuffer)
-            slotCount = 2;
+        // This, however, is added in 'createStructuredBufferWithCounterTypeLayout'
 
         if (options.hlslToVulkanKindFlags)
         {
@@ -2785,6 +2797,49 @@ RefPtr<TypeLayout> createParameterGroupTypeLayout(
         elementTypeRules);
 }
 
+// Create a type layout for a structured buffer type with an associated counter
+RefPtr<StructuredBufferTypeLayout>
+createStructuredBufferWithCounterTypeLayout(
+    TypeLayoutContext const& context,
+    ShaderParameterKind kind,
+    Type* structuredBufferType,
+    RefPtr<TypeLayout> elementTypeLayout)
+{
+    auto typeLayout = createStructuredBufferTypeLayout(context, kind, structuredBufferType, elementTypeLayout);
+
+    const auto structuredBufferLayoutRules = context.getRulesFamily()->getStructuredBufferRules(context.targetReq);
+
+    const auto counterType = context.astBuilder->getIntType();
+    const auto counterBufferType = context.astBuilder->getRWStructuredBufferType(counterType);
+    const auto counterTypeLayout = createTypeLayoutWith(context, structuredBufferLayoutRules, counterBufferType);
+
+    const auto counterVarDecl = context.astBuilder->create<VarDecl>();
+    counterVarDecl->type.type = counterBufferType;
+    counterVarDecl->nameAndLoc.name =
+      context.astBuilder->getSharedASTBuilder()->getNamePool()->getName("counter");
+
+    RefPtr<VarLayout> counterVarLayout = new VarLayout();
+    counterVarLayout->varDecl = makeDeclRef(counterVarDecl);
+    counterVarLayout->typeLayout = counterTypeLayout;
+
+    for(auto& typeResourceInfo : typeLayout->resourceInfos)
+    {
+        const auto counterResourceInfo
+            = counterVarLayout->findOrAddResourceInfo(typeResourceInfo.kind);
+        const auto counterTypeResourceInfo
+            = counterVarLayout->getTypeLayout()->FindResourceInfo(typeResourceInfo.kind);
+        // We expect this index to be 1
+        counterResourceInfo->index = typeResourceInfo.count.getFiniteValue();
+        // likewise
+        typeResourceInfo.count += counterTypeResourceInfo->count;
+    }
+
+    typeLayout->counterVarLayout = counterVarLayout;
+    typeLayout->addResourceUsageFrom(counterTypeLayout);
+
+    return typeLayout;
+}
+
 // Create a type layout for a structured buffer type.
 RefPtr<StructuredBufferTypeLayout>
 createStructuredBufferTypeLayout(
@@ -2845,12 +2900,24 @@ createStructuredBufferTypeLayout(
         structuredBufferLayoutRules,
         elementType);
 
-    return createStructuredBufferTypeLayout(
-        context,
-        kind,
-        structuredBufferType,
-        elementTypeLayout);
-
+    if(kind == ShaderParameterKind::AppendConsumeStructuredBuffer
+        && structuredBufferLayoutRules->DoStructuredBuffersNeedSeparateCounterBuffer())
+    {
+        return createStructuredBufferWithCounterTypeLayout(
+            context,
+            kind,
+            structuredBufferType,
+            elementTypeLayout
+        );
+    }
+    else
+    {
+        return createStructuredBufferTypeLayout(
+            context,
+            kind,
+            structuredBufferType,
+            elementTypeLayout);
+    }
 }
 
     /// Create layout information for the given `type`.
diff --git a/source/slang/slang-type-layout.h b/source/slang/slang-type-layout.h
index e3dd719d6..1ebd26098 100644
--- a/source/slang/slang-type-layout.h
+++ b/source/slang/slang-type-layout.h
@@ -628,6 +628,10 @@ class StructuredBufferTypeLayout : public TypeLayout
 {
 public:
     RefPtr<TypeLayout> elementTypeLayout;
+    // If required, this is the VarLayout for a buffer which contains the
+    // counter associated with the buffer, most often used for
+    // AppendStructuredBuffer or ConsumeStructuredBuffer
+    RefPtr<VarLayout> counterVarLayout;
 };
 
     /// Type layout for a logical sequence type
@@ -962,6 +966,10 @@ struct SimpleLayoutRulesImpl
 
     // End layout for a struct, and finalize its size/alignment.
     virtual void EndStructLayout(UniformLayoutInfo* ioStructInfo) = 0;
+
+    // Do structured buffers need a separate binding for the counter buffer?
+    // (DirectX is the exception in managing these together)
+    virtual bool DoStructuredBuffersNeedSeparateCounterBuffer() = 0;
 };
 
 struct ObjectLayoutRulesImpl
@@ -1021,6 +1029,13 @@ struct LayoutRulesImpl
         return simpleRules->EndStructLayout(ioStructInfo);
     }
 
+    // Do structured buffers need a separate binding for the counter buffer?
+    // (DirectX is the exception in managing these together)
+    bool DoStructuredBuffersNeedSeparateCounterBuffer()
+    {
+        return simpleRules->DoStructuredBuffersNeedSeparateCounterBuffer();
+    }
+
     // Forward `ObjectLayoutRulesImpl` interface
 
     SimpleLayoutInfo GetObjectLayout(ShaderParameterKind kind, const ObjectLayoutRulesImpl::Options& options)
@@ -1306,6 +1321,13 @@ createStructuredBufferTypeLayout(
     Type*                structuredBufferType,
     Type*                elementType);
 
+RefPtr<StructuredBufferTypeLayout>
+createStructuredBufferTypeLayout(
+    TypeLayoutContext const& context,
+    ShaderParameterKind      kind,
+    Type*                    structuredBufferType,
+    RefPtr<TypeLayout>       elementTypeLayout);
+
     /// Create a type layout for an unspecialized `globalGenericParamDecl`.
 RefPtr<TypeLayout> createTypeLayoutForGlobalGenericTypeParam(
     TypeLayoutContext const&    context,