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, 120 insertions, 5 deletions

diff --git a/tools/gfx-util/shader-cursor.cpp b/tools/gfx-util/shader-cursor.cpp
index 5a2cdfd33..c849673bf 100644
--- a/tools/gfx-util/shader-cursor.cpp
+++ b/tools/gfx-util/shader-cursor.cpp
@@ -20,6 +20,58 @@ Result gfx::ShaderCursor::getDereferenced(ShaderCursor& outCursor) const
     }
 }
 
+ShaderCursor ShaderCursor::getExplicitCounter() const
+{
+    // Similar to getField below
+
+    // The alternative to handling this here would be to augment IResourceView
+    // with a `getCounterResourceView()`, and set that also in `setResource`
+    if(const auto counterVarLayout = m_typeLayout->getExplicitCounter())
+    {
+        ShaderCursor counterCursor;
+
+        // The counter cursor will point into the same parent object.
+        counterCursor.m_baseObject = m_baseObject;
+
+        // The type being pointed to is the type of the field.
+        counterCursor.m_typeLayout = counterVarLayout->getTypeLayout();
+
+        // The byte offset is the current offset plus the relative offset of the counter.
+        // The offset in binding ranges is computed similarly.
+        counterCursor.m_offset.uniformOffset
+            = m_offset.uniformOffset + SlangInt(counterVarLayout->getOffset());
+        counterCursor.m_offset.bindingRangeIndex
+            = m_offset.bindingRangeIndex + GfxIndex(m_typeLayout->getExplicitCounterBindingRangeOffset());
+
+        // The index of the counter within any binding ranges will be the same
+        // as the index computed for the parent structure.
+        //
+        // Note: this case would arise for an array of structured buffers
+        //
+        //      AppendStructuredBuffer g[4];
+        //
+        // In this scenario, `g` holds two binding ranges:
+        //
+        // * Range #0 comprises 4 element buffers, representing `g[...].elements`
+        // * Range #1 comprises 4 counter buffers, representing `g[...].counter`
+        //
+        // A cursor for `g[2]` would have a `bindingRangeIndex` of zero but
+        // a `bindingArrayIndex` of 2, indicating that we could end up
+        // referencing either range, but no matter what we know the index
+        // is 2. Thus when we form a cursor for `g[2].counter` we want to
+        // apply the binding range offset to get a `bindingRangeIndex` of
+        // 1, while the `bindingArrayIndex` is unmodified.
+        //
+        // The result is that `g[2].counter` is stored in range #1 at array index 2.
+        //
+        counterCursor.m_offset.bindingArrayIndex = m_offset.bindingArrayIndex;
+
+        return counterCursor;
+    }
+    // Otherwise, return an invalid cursor
+    return ShaderCursor{};
+}
+
 Result ShaderCursor::getField(const char* name, const char* nameEnd, ShaderCursor& outCursor) const
 {
     // If this cursor is invalid, then can't possible fetch a field.
diff --git a/tools/gfx-util/shader-cursor.h b/tools/gfx-util/shader-cursor.h
index 5dfc57f6d..3d35bb2ed 100644
--- a/tools/gfx-util/shader-cursor.h
+++ b/tools/gfx-util/shader-cursor.h
@@ -63,6 +63,13 @@ struct ShaderCursor
         return cursor;
     }
 
+    /// Some resources such as RWStructuredBuffer, AppendStructuredBuffer and
+    /// ConsumeStructuredBuffer need to have their counter explicitly bound on
+    /// APIs other than DirectX, this will return a valid ShaderCursor pointing
+    /// to that resource if that is the case.
+    /// Otherwise, this returns an invalid cursor.
+    ShaderCursor getExplicitCounter() const;
+
     ShaderCursor getElement(GfxIndex index) const;
 
     static Result followPath(const char* path, ShaderCursor& ioCursor);
diff --git a/tools/gfx/d3d12/d3d12-device.cpp b/tools/gfx/d3d12/d3d12-device.cpp
index 983f93bce..eb9e597bf 100644
--- a/tools/gfx/d3d12/d3d12-device.cpp
+++ b/tools/gfx/d3d12/d3d12-device.cpp
@@ -1665,9 +1665,11 @@ Result DeviceImpl::createBufferView(
 {
     auto resourceImpl = (BufferResourceImpl*)buffer;
     auto resourceDesc = *resourceImpl->getDesc();
+    const auto counterResourceImpl = static_cast<BufferResourceImpl*>(counterBuffer);
 
     RefPtr<ResourceViewImpl> viewImpl = new ResourceViewImpl();
     viewImpl->m_resource = resourceImpl;
+    viewImpl->m_counterResource = counterResourceImpl;
     viewImpl->m_desc = desc;
 
     switch (desc.type)
@@ -1722,7 +1724,6 @@ Result DeviceImpl::createBufferView(
         }
         else
         {
-            auto counterResourceImpl = static_cast<BufferResourceImpl*>(counterBuffer);
             SLANG_RETURN_ON_FAIL(m_cpuViewHeap->allocate(&viewImpl->m_descriptor));
             viewImpl->m_allocator = m_cpuViewHeap;
             m_device->CreateUnorderedAccessView(
diff --git a/tools/gfx/d3d12/d3d12-resource-views.h b/tools/gfx/d3d12/d3d12-resource-views.h
index 12e3d0714..fd3f44116 100644
--- a/tools/gfx/d3d12/d3d12-resource-views.h
+++ b/tools/gfx/d3d12/d3d12-resource-views.h
@@ -26,6 +26,8 @@ class ResourceViewImpl
 {
 public:
     Slang::RefPtr<Resource> m_resource;
+    // null, unless this is a structuredbuffer with a separate counter buffer
+    Slang::RefPtr<Resource> m_counterResource;
     virtual SLANG_NO_THROW Result SLANG_MCALL getNativeHandle(InteropHandle* outHandle) override;
 };
 
diff --git a/tools/gfx/d3d12/d3d12-shader-object.cpp b/tools/gfx/d3d12/d3d12-shader-object.cpp
index c5389f6ea..74760eabd 100644
--- a/tools/gfx/d3d12/d3d12-shader-object.cpp
+++ b/tools/gfx/d3d12/d3d12-shader-object.cpp
@@ -181,7 +181,11 @@ Result ShaderObjectImpl::init(
         // referenced by descriptors in this object does not get
         // freed while the object is still live.
         //
+        // The doubling here is because any buffer resource could
+        // have a counter buffer associated with it, which we 
+        // also need to ensure isn't destroyed prematurely.
         m_boundResources.setCount(resourceCount);
+        m_boundCounterResources.setCount(resourceCount);
     }
     if (auto samplerCount = layout->getSamplerSlotCount())
     {
@@ -950,8 +954,9 @@ Result ShaderObjectImpl::setResource(ShaderOffset const& offset, IResourceView*
     {
         auto resourceViewImpl = static_cast<ResourceViewImpl*>(resourceView);
         // Hold a reference to the resource to prevent its destruction.
-        m_boundResources[bindingRange.baseIndex + offset.bindingArrayIndex] =
-            resourceViewImpl->m_resource;
+        const auto resourceOffset = bindingRange.baseIndex + offset.bindingArrayIndex;
+        m_boundResources[resourceOffset] = resourceViewImpl->m_resource;
+        m_boundCounterResources[resourceOffset] = resourceViewImpl->m_counterResource;
         internalResourceView = resourceViewImpl;
     }
     break;
diff --git a/tools/gfx/d3d12/d3d12-shader-object.h b/tools/gfx/d3d12/d3d12-shader-object.h
index d6560628c..84d305ae7 100644
--- a/tools/gfx/d3d12/d3d12-shader-object.h
+++ b/tools/gfx/d3d12/d3d12-shader-object.h
@@ -197,6 +197,7 @@ public:
     DescriptorSet m_cachedGPUDescriptorSet;
 
     ShortList<RefPtr<Resource>, 8> m_boundResources;
+    ShortList<RefPtr<Resource>, 8> m_boundCounterResources;
     List<D3D12_GPU_VIRTUAL_ADDRESS> m_rootArguments;
     /// A constant buffer used to stored ordinary data for this object
     /// and existential-type sub-objects.
diff --git a/tools/render-test/render-test-main.cpp b/tools/render-test/render-test-main.cpp
index 3bdcdba07..1f2b3a5ae 100644
--- a/tools/render-test/render-test-main.cpp
+++ b/tools/render-test/render-test-main.cpp
@@ -209,11 +209,49 @@ struct AssignValsFromLayoutContext
         ComPtr<IBufferResource> bufferResource;
         SLANG_RETURN_ON_FAIL(ShaderRendererUtil::createBufferResource(srcBuffer, /*entry.isOutput,*/ bufferSize, bufferData.getBuffer(), device, bufferResource));
 
+        ComPtr<IBufferResource> counterResource;
+        const auto explicitCounterCursor = dstCursor.getExplicitCounter();
+        if(srcBuffer.counter != ~0u)
+        {
+            if(explicitCounterCursor.isValid())
+            {
+                // If this cursor has a full buffer object associated with the
+                // resource, then assign to that.
+                ShaderInputLayout::BufferVal counterVal;
+                counterVal.bufferData.add(srcBuffer.counter);
+                assignBuffer(explicitCounterCursor, &counterVal);
+            }
+            else
+            {
+                // Otherwise, this API (D3D) must be handling the buffer object
+                // specially, in which case create the buffer resource to pass
+                // into `createBufferView`
+                const InputBufferDesc& counterBufferDesc{
+                    InputBufferType::StorageBuffer,
+                    sizeof(uint32_t),
+                    Format::Unknown,
+                };
+                SLANG_RETURN_ON_FAIL(ShaderRendererUtil::createBufferResource(
+                    counterBufferDesc,
+                    sizeof(srcBuffer.counter),
+                    &srcBuffer.counter,
+                    device,
+                    counterResource
+                ));
+            }
+        }
+        else if(explicitCounterCursor.isValid())
+        {
+            // If we know we require a counter for this resource but haven't
+            // been given one, error
+            return SLANG_E_INVALID_ARG;
+        }
+
         IResourceView::Desc viewDesc = {};
         viewDesc.type = IResourceView::Type::UnorderedAccess;
         viewDesc.format = srcBuffer.format;
         viewDesc.bufferElementSize = srcVal->bufferDesc.stride;
-        auto bufferView = device->createBufferView(bufferResource, nullptr, viewDesc);
+        auto bufferView = device->createBufferView(bufferResource, counterResource, viewDesc);
         dstCursor.setResource(bufferView);
         maybeAddOutput(dstCursor, srcVal, bufferResource);
 
@@ -977,7 +1015,7 @@ Result RenderTestApp::writeBindingOutput(const String& fileName)
                 m_transientHeap->finish();
                 m_transientHeap->synchronizeAndReset();
 
-                m_device->readBufferResource(stagingBuffer, 0, bufferSize, blob.writeRef());
+                SLANG_RETURN_ON_FAIL(m_device->readBufferResource(stagingBuffer, 0, bufferSize, blob.writeRef()));
             }
 
             if (!blob)
diff --git a/tools/render-test/shader-input-layout.cpp b/tools/render-test/shader-input-layout.cpp
index b8c505607..80ce4e316 100644
--- a/tools/render-test/shader-input-layout.cpp
+++ b/tools/render-test/shader-input-layout.cpp
@@ -203,6 +203,11 @@ namespace renderer_test
                 parser.Read("=");
                 val->bufferDesc.stride = parser.ReadInt();
             }
+            else if (word == "counter")
+            {
+                parser.Read("=");
+                val->bufferDesc.counter = parser.ReadInt();
+            }
             else if (word == "random")
             {
                 parser.Read("(");
diff --git a/tools/render-test/shader-input-layout.h b/tools/render-test/shader-input-layout.h
index 78d545114..2803d1915 100644
--- a/tools/render-test/shader-input-layout.h
+++ b/tools/render-test/shader-input-layout.h
@@ -58,6 +58,10 @@ struct InputBufferDesc
     InputBufferType type = InputBufferType::StorageBuffer;
     int stride = 0; // stride == 0 indicates an unstructured buffer.
     Format format = Format::Unknown;
+    // For RWStructuredBuffer, AppendStructuredBuffer, ConsumeStructuredBuffer
+    // the default value of 0xffffffff indicates that a counter buffer should
+    // not be assigned
+    uint32_t counter = ~0u;
 };
 
 struct InputSamplerDesc