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>
---
 tools/gfx/d3d12/d3d12-device.cpp        | 3 ++-
 tools/gfx/d3d12/d3d12-resource-views.h  | 2 ++
 tools/gfx/d3d12/d3d12-shader-object.cpp | 9 +++++++--
 tools/gfx/d3d12/d3d12-shader-object.h   | 1 +
 4 files changed, 12 insertions(+), 3 deletions(-)

(limited to 'tools/gfx')

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