From 7826afcaad78cc33c976bb3db3cdc9eada4c77e8 Mon Sep 17 00:00:00 2001
From: Ellie Hermaszewska <ellieh@nvidia.com>
Date: Wed, 18 Oct 2023 06:26:00 +0800
Subject: 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>
---
 source/slang/slang-type-layout.cpp | 83 ++++++++++++++++++++++++++++++++++----
 1 file changed, 75 insertions(+), 8 deletions(-)

(limited to 'source/slang/slang-type-layout.cpp')

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`.
-- 
cgit v1.2.3