[GFX] Fix d3d12 buffer view creation logic for StructuredBuffers. (#3954)

6 files changed, 267 insertions, 123 deletions

diff --git a/tools/gfx/d3d12/d3d12-device.cpp b/tools/gfx/d3d12/d3d12-device.cpp
index 33375497b..03f9997f4 100644
--- a/tools/gfx/d3d12/d3d12-device.cpp
+++ b/tools/gfx/d3d12/d3d12-device.cpp
@@ -1674,126 +1674,16 @@ Result DeviceImpl::createBufferView(
     viewImpl->m_counterResource = counterResourceImpl;
     viewImpl->m_desc = desc;
 
-    switch (desc.type)
-    {
-    default:
-        return SLANG_FAIL;
-
-    case IResourceView::Type::UnorderedAccess:
-    {
-        D3D12_UNORDERED_ACCESS_VIEW_DESC uavDesc = {};
-        uavDesc.ViewDimension = D3D12_UAV_DIMENSION_BUFFER;
-        uavDesc.Format = D3DUtil::getMapFormat(desc.format);
-        uavDesc.Buffer.FirstElement = desc.bufferRange.firstElement;
-        uint64_t viewSize = 0;
-        if (desc.bufferElementSize)
-        {
-            uavDesc.Buffer.StructureByteStride = (UINT)desc.bufferElementSize;
-            uavDesc.Buffer.NumElements =
-                desc.bufferRange.elementCount == 0
-                ? UINT(resourceDesc.sizeInBytes / desc.bufferElementSize)
-                : (UINT)desc.bufferRange.elementCount;
-            viewSize = (uint64_t)desc.bufferElementSize * uavDesc.Buffer.NumElements;
-        }
-        else if (desc.format == Format::Unknown)
-        {
-            uavDesc.Format = DXGI_FORMAT_R32_TYPELESS;
-            uavDesc.Buffer.NumElements = desc.bufferRange.elementCount == 0
-                ? UINT(resourceDesc.sizeInBytes / 4)
-                : UINT(desc.bufferRange.elementCount / 4);
-            uavDesc.Buffer.Flags |= D3D12_BUFFER_UAV_FLAG_RAW;
-            viewSize = 4ull * uavDesc.Buffer.NumElements;
-        }
-        else
-        {
-            FormatInfo sizeInfo;
-            gfxGetFormatInfo(desc.format, &sizeInfo);
-            assert(sizeInfo.pixelsPerBlock == 1);
-            uavDesc.Buffer.NumElements =
-                desc.bufferRange.elementCount == 0
-                ? UINT(resourceDesc.sizeInBytes / sizeInfo.blockSizeInBytes)
-                : (UINT)desc.bufferRange.elementCount;
-            viewSize = (uint64_t)uavDesc.Buffer.NumElements * sizeInfo.blockSizeInBytes;
-        }
-
-        if (viewSize >= (1ull << 32) - 8)
-        {
-            // D3D12 does not support view descriptors that has size near 4GB.
-            // We will not create actual SRV/UAVs for such large buffers.
-            // However, a buffer this large can still be bound as root parameter.
-            // So instead of failing, we quietly ignore descriptor creation.
-            viewImpl->m_descriptor.cpuHandle.ptr = 0;
-        }
-        else
-        {
-            SLANG_RETURN_ON_FAIL(m_cpuViewHeap->allocate(&viewImpl->m_descriptor));
-            viewImpl->m_allocator = m_cpuViewHeap;
-            m_device->CreateUnorderedAccessView(
-                resourceImpl->m_resource,
-                counterResourceImpl ? counterResourceImpl->m_resource.getResource() : nullptr,
-                &uavDesc,
-                viewImpl->m_descriptor.cpuHandle);
-        }
-    }
-    break;
-
-    case IResourceView::Type::ShaderResource:
-    {
-        D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
-        srvDesc.ViewDimension = D3D12_SRV_DIMENSION_BUFFER;
-        srvDesc.Format = D3DUtil::getMapFormat(desc.format);
-        srvDesc.Buffer.StructureByteStride = 0;
-        srvDesc.Buffer.FirstElement = desc.bufferRange.firstElement;
-        srvDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
-        uint64_t viewSize = 0;
-        if (desc.bufferElementSize)
-        {
-            srvDesc.Buffer.StructureByteStride = (UINT)desc.bufferElementSize;
-            srvDesc.Buffer.NumElements =
-                desc.bufferRange.elementCount == 0
-                ? UINT(resourceDesc.sizeInBytes / desc.bufferElementSize)
-                : (UINT)desc.bufferRange.elementCount;
-            viewSize = (uint64_t)desc.bufferElementSize * srvDesc.Buffer.NumElements;
-        }
-        else if (desc.format == Format::Unknown)
-        {
-            srvDesc.Format = DXGI_FORMAT_R32_TYPELESS;
-            srvDesc.Buffer.NumElements = desc.bufferRange.elementCount == 0
-                ? UINT(resourceDesc.sizeInBytes / 4)
-                : UINT(desc.bufferRange.elementCount / 4);
-            srvDesc.Buffer.Flags |= D3D12_BUFFER_SRV_FLAG_RAW;
-            viewSize = 4ull * srvDesc.Buffer.NumElements;
-        }
-        else
-        {
-            FormatInfo sizeInfo;
-            gfxGetFormatInfo(desc.format, &sizeInfo);
-            assert(sizeInfo.pixelsPerBlock == 1);
-            srvDesc.Buffer.NumElements =
-                desc.bufferRange.elementCount == 0
-                ? UINT(resourceDesc.sizeInBytes / sizeInfo.blockSizeInBytes)
-                : (UINT)desc.bufferRange.elementCount;
-            viewSize = (uint64_t)srvDesc.Buffer.NumElements * sizeInfo.blockSizeInBytes;
-        }
-        if (viewSize >= (1ull << 32) - 8)
-        {
-            // D3D12 does not support view descriptors that has size near 4GB.
-            // We will not create actual SRV/UAVs for such large buffers.
-            // However, a buffer this large can still be bound as root parameter.
-            // So instead of failing, we quietly ignore descriptor creation.
-            viewImpl->m_descriptor.cpuHandle.ptr = 0;
-        }
-        else
-        {
-            SLANG_RETURN_ON_FAIL(m_cpuViewHeap->allocate(&viewImpl->m_descriptor));
-            viewImpl->m_allocator = m_cpuViewHeap;
-            m_device->CreateShaderResourceView(
-                resourceImpl->m_resource, &srvDesc, viewImpl->m_descriptor.cpuHandle);
-        }
-    }
-    break;
-    }
-
+    SLANG_RETURN_ON_FAIL(createD3D12BufferDescriptor(
+        resourceImpl,
+        counterResourceImpl,
+        desc,
+        this,
+        m_cpuViewHeap.get(),
+        &viewImpl->m_descriptor));
+    if (viewImpl->m_descriptor.cpuHandle.ptr != 0)
+        viewImpl->m_allocator = m_cpuViewHeap.get();
+ 
     returnComPtr(outView, viewImpl);
     return SLANG_OK;
 }
diff --git a/tools/gfx/d3d12/d3d12-resource-views.cpp b/tools/gfx/d3d12/d3d12-resource-views.cpp
index 5a044dd3b..b0b441f87 100644
--- a/tools/gfx/d3d12/d3d12-resource-views.cpp
+++ b/tools/gfx/d3d12/d3d12-resource-views.cpp
@@ -1,5 +1,6 @@
 // d3d12-resource-views.cpp
 #include "d3d12-resource-views.h"
+#include "d3d12-device.h"
 
 namespace gfx
 {
@@ -12,6 +13,208 @@ ResourceViewInternalImpl::~ResourceViewInternalImpl()
 {
     if (m_descriptor.cpuHandle.ptr)
         m_allocator->free(m_descriptor);
+    for (auto desc : m_mapBufferStrideToDescriptor)
+    {
+        m_allocator->free(desc.second);
+    }
+}
+
+SlangResult createD3D12BufferDescriptor(
+    BufferResourceImpl* buffer,
+    BufferResourceImpl* counterBuffer,
+    IResourceView::Desc const& desc,
+    DeviceImpl* device,
+    D3D12GeneralExpandingDescriptorHeap* descriptorHeap,
+    D3D12Descriptor* outDescriptor)
+{
+
+    auto resourceImpl = (BufferResourceImpl*)buffer;
+    auto resourceDesc = *resourceImpl->getDesc();
+    const auto counterResourceImpl = static_cast<BufferResourceImpl*>(counterBuffer);
+
+    switch (desc.type)
+    {
+    default:
+        return SLANG_FAIL;
+
+    case IResourceView::Type::UnorderedAccess:
+    {
+        D3D12_UNORDERED_ACCESS_VIEW_DESC uavDesc = {};
+        uavDesc.ViewDimension = D3D12_UAV_DIMENSION_BUFFER;
+        uavDesc.Format = D3DUtil::getMapFormat(desc.format);
+        uavDesc.Buffer.FirstElement = desc.bufferRange.firstElement;
+        uint64_t viewSize = 0;
+        if (desc.bufferElementSize)
+        {
+            uavDesc.Buffer.StructureByteStride = (UINT)desc.bufferElementSize;
+            uavDesc.Buffer.NumElements =
+                desc.bufferRange.elementCount == 0
+                ? UINT(resourceDesc.sizeInBytes / desc.bufferElementSize)
+                : (UINT)desc.bufferRange.elementCount;
+            viewSize = (uint64_t)desc.bufferElementSize * uavDesc.Buffer.NumElements;
+        }
+        else if (desc.format == Format::Unknown)
+        {
+            uavDesc.Format = DXGI_FORMAT_R32_TYPELESS;
+            uavDesc.Buffer.NumElements = desc.bufferRange.elementCount == 0
+                ? UINT(resourceDesc.sizeInBytes / 4)
+                : UINT(desc.bufferRange.elementCount / 4);
+            uavDesc.Buffer.Flags |= D3D12_BUFFER_UAV_FLAG_RAW;
+            viewSize = 4ull * uavDesc.Buffer.NumElements;
+        }
+        else
+        {
+            FormatInfo sizeInfo;
+            gfxGetFormatInfo(desc.format, &sizeInfo);
+            assert(sizeInfo.pixelsPerBlock == 1);
+            uavDesc.Buffer.NumElements =
+                desc.bufferRange.elementCount == 0
+                ? UINT(resourceDesc.sizeInBytes / sizeInfo.blockSizeInBytes)
+                : (UINT)desc.bufferRange.elementCount;
+            viewSize = (uint64_t)uavDesc.Buffer.NumElements * sizeInfo.blockSizeInBytes;
+        }
+
+        if (viewSize >= (1ull << 32) - 8)
+        {
+            // D3D12 does not support view descriptors that has size near 4GB.
+            // We will not create actual SRV/UAVs for such large buffers.
+            // However, a buffer this large can still be bound as root parameter.
+            // So instead of failing, we quietly ignore descriptor creation.
+            outDescriptor->cpuHandle.ptr = 0;
+        }
+        else
+        {
+            SLANG_RETURN_ON_FAIL(descriptorHeap->allocate(outDescriptor));
+            device->m_device->CreateUnorderedAccessView(
+                resourceImpl->m_resource,
+                counterResourceImpl ? counterResourceImpl->m_resource.getResource() : nullptr,
+                &uavDesc,
+                outDescriptor->cpuHandle);
+        }
+    }
+    break;
+
+    case IResourceView::Type::ShaderResource:
+    {
+        D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
+        srvDesc.ViewDimension = D3D12_SRV_DIMENSION_BUFFER;
+        srvDesc.Format = D3DUtil::getMapFormat(desc.format);
+        srvDesc.Buffer.StructureByteStride = 0;
+        srvDesc.Buffer.FirstElement = desc.bufferRange.firstElement;
+        srvDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
+        uint64_t viewSize = 0;
+        if (desc.bufferElementSize)
+        {
+            srvDesc.Buffer.StructureByteStride = (UINT)desc.bufferElementSize;
+            srvDesc.Buffer.NumElements =
+                desc.bufferRange.elementCount == 0
+                ? UINT(resourceDesc.sizeInBytes / desc.bufferElementSize)
+                : (UINT)desc.bufferRange.elementCount;
+            viewSize = (uint64_t)desc.bufferElementSize * srvDesc.Buffer.NumElements;
+        }
+        else if (desc.format == Format::Unknown)
+        {
+            srvDesc.Format = DXGI_FORMAT_R32_TYPELESS;
+            srvDesc.Buffer.NumElements = desc.bufferRange.elementCount == 0
+                ? UINT(resourceDesc.sizeInBytes / 4)
+                : UINT(desc.bufferRange.elementCount / 4);
+            srvDesc.Buffer.Flags |= D3D12_BUFFER_SRV_FLAG_RAW;
+            viewSize = 4ull * srvDesc.Buffer.NumElements;
+        }
+        else
+        {
+            FormatInfo sizeInfo;
+            gfxGetFormatInfo(desc.format, &sizeInfo);
+            assert(sizeInfo.pixelsPerBlock == 1);
+            srvDesc.Buffer.NumElements =
+                desc.bufferRange.elementCount == 0
+                ? UINT(resourceDesc.sizeInBytes / sizeInfo.blockSizeInBytes)
+                : (UINT)desc.bufferRange.elementCount;
+            viewSize = (uint64_t)srvDesc.Buffer.NumElements * sizeInfo.blockSizeInBytes;
+        }
+        if (viewSize >= (1ull << 32) - 8)
+        {
+            // D3D12 does not support view descriptors that has size near 4GB.
+            // We will not create actual SRV/UAVs for such large buffers.
+            // However, a buffer this large can still be bound as root parameter.
+            // So instead of failing, we quietly ignore descriptor creation.
+            outDescriptor->cpuHandle.ptr = 0;
+        }
+        else
+        {
+            SLANG_RETURN_ON_FAIL(descriptorHeap->allocate(outDescriptor));
+            device->m_device->CreateShaderResourceView(
+                resourceImpl->m_resource, &srvDesc, outDescriptor->cpuHandle);
+        }
+    }
+    break;
+    }
+    return SLANG_OK;
+}
+
+SlangResult ResourceViewInternalImpl::getBufferDescriptorForBinding(
+    DeviceImpl* device,
+    ResourceViewImpl* view,
+    uint32_t bufferStride,
+    D3D12Descriptor& outDescriptor)
+{
+    // If stride is 0, just use the default descriptor.
+    if (bufferStride == 0)
+    {
+        outDescriptor = m_descriptor;
+        return SLANG_OK;
+    }
+
+    // Otherwise, look for an existing descriptor from the cache if it exists.
+    if (auto descriptor = m_mapBufferStrideToDescriptor.tryGetValue(bufferStride))
+    {
+        outDescriptor = *descriptor;
+        return SLANG_OK;
+    }
+
+    // We need to create and cache a d3d12 descriptor for the resource view that encodes
+    // the given buffer stride.
+    auto bufferResImpl = static_cast<BufferResourceImpl*>(view->m_resource.get());
+    auto desc = view->m_desc;
+    uint64_t bufferSize = 0;
+    if (desc.bufferElementSize == 0)
+    {
+        // If buffer element size is 0, we assume the buffer range from original desc is in bytes.
+        bufferSize = desc.bufferRange.elementCount;
+        if (bufferSize == 0)
+        {
+            bufferSize = bufferResImpl->getDesc()->sizeInBytes - desc.bufferRange.firstElement;
+        }
+        desc.bufferElementSize = bufferStride;
+        desc.bufferRange.firstElement /= bufferStride;
+        desc.bufferRange.elementCount = bufferSize / bufferStride;
+    }
+    else
+    {
+        // If buffer element size is not 0, we assume the buffer range from original desc is in elements
+        // of original stride.
+        if (desc.bufferRange.elementCount == 0)
+        {
+            bufferSize = bufferResImpl->getDesc()->sizeInBytes - desc.bufferRange.firstElement * desc.bufferElementSize;
+        }
+        else
+        {
+            bufferSize = desc.bufferRange.elementCount * desc.bufferElementSize;
+        }
+        desc.bufferElementSize = bufferStride;
+        desc.bufferRange.firstElement = desc.bufferRange.firstElement * desc.bufferElementSize / bufferStride;
+        desc.bufferRange.elementCount = bufferSize / bufferStride;
+    }
+    SLANG_RETURN_ON_FAIL(createD3D12BufferDescriptor(
+        bufferResImpl,
+        static_cast<BufferResourceImpl*>(view->m_counterResource.get()),
+        desc,
+        device,
+        m_allocator,
+        &outDescriptor));
+    m_mapBufferStrideToDescriptor[bufferStride] = outDescriptor;
+
+    return SLANG_OK;
 }
 
 Result ResourceViewImpl::getNativeHandle(InteropHandle* outHandle)
diff --git a/tools/gfx/d3d12/d3d12-resource-views.h b/tools/gfx/d3d12/d3d12-resource-views.h
index fd3f44116..f8c6654f2 100644
--- a/tools/gfx/d3d12/d3d12-resource-views.h
+++ b/tools/gfx/d3d12/d3d12-resource-views.h
@@ -12,14 +12,37 @@ namespace d3d12
 
 using namespace Slang;
 
+class ResourceViewImpl;
+
 class ResourceViewInternalImpl
 {
 public:
+    // The default descriptor for the view.
     D3D12Descriptor m_descriptor;
+
+    // StructuredBuffer descriptors for different strides.
+    Dictionary<uint32_t, D3D12Descriptor> m_mapBufferStrideToDescriptor;
+
     RefPtr<D3D12GeneralExpandingDescriptorHeap> m_allocator;
+
     ~ResourceViewInternalImpl();
+
+    // Get a d3d12 descriptor from the buffer view with the given buffer element stride.
+    SlangResult getBufferDescriptorForBinding(
+        DeviceImpl* device,
+        ResourceViewImpl* view,
+        uint32_t bufferStride,
+        D3D12Descriptor& outDescriptor);
 };
 
+SlangResult createD3D12BufferDescriptor(
+    BufferResourceImpl* buffer,
+    BufferResourceImpl* counterBuffer,
+    IResourceView::Desc const& desc,
+    DeviceImpl* device,
+    D3D12GeneralExpandingDescriptorHeap* descriptorHeap,
+    D3D12Descriptor* outDescriptor);
+
 class ResourceViewImpl
     : public ResourceViewBase
     , public ResourceViewInternalImpl
diff --git a/tools/gfx/d3d12/d3d12-shader-object-layout.cpp b/tools/gfx/d3d12/d3d12-shader-object-layout.cpp
index 0c0095621..ef59c3672 100644
--- a/tools/gfx/d3d12/d3d12-shader-object-layout.cpp
+++ b/tools/gfx/d3d12/d3d12-shader-object-layout.cpp
@@ -116,6 +116,7 @@ Result ShaderObjectLayoutImpl::Builder::setElementTypeLayout(
         uint32_t count = (uint32_t)typeLayout->getBindingRangeBindingCount(r);
         slang::TypeLayoutReflection* slangLeafTypeLayout =
             typeLayout->getBindingRangeLeafTypeLayout(r);
+        
         BindingRangeInfo bindingRangeInfo = {};
         bindingRangeInfo.bindingType = slangBindingType;
         bindingRangeInfo.resourceShape = slangLeafTypeLayout->getResourceShape();
@@ -126,6 +127,21 @@ Result ShaderObjectLayoutImpl::Builder::setElementTypeLayout(
             typeLayout,
             r);
         bindingRangeInfo.isSpecializable = typeLayout->isBindingRangeSpecializable(r);
+        switch (slangBindingType)
+        {
+        case slang::BindingType::RawBuffer:
+        case slang::BindingType::TypedBuffer:
+        case slang::BindingType::MutableRawBuffer:
+        case slang::BindingType::MutableTypedBuffer:
+            {
+                auto bufferElementType = slangLeafTypeLayout->getElementTypeLayout();
+                if (bufferElementType)
+                {
+                    bindingRangeInfo.bufferElementStride = (uint32_t)bufferElementType->getStride();
+                }
+            }
+            break;
+        }
         if (bindingRangeInfo.isRootParameter)
         {
             RootParameterInfo rootInfo = {};
diff --git a/tools/gfx/d3d12/d3d12-shader-object-layout.h b/tools/gfx/d3d12/d3d12-shader-object-layout.h
index 4b46df0b0..b8b3082f6 100644
--- a/tools/gfx/d3d12/d3d12-shader-object-layout.h
+++ b/tools/gfx/d3d12/d3d12-shader-object-layout.h
@@ -66,6 +66,9 @@ public:
         /// as a sub-object.
         uint32_t subObjectIndex;
 
+        /// The stride of a structured buffer.
+        uint32_t bufferElementStride;
+
         bool isRootParameter;
 
         /// Is this binding range represent a specialization point, such as an existential value, or a `ParameterBlock<IFoo>`.
diff --git a/tools/gfx/d3d12/d3d12-shader-object.cpp b/tools/gfx/d3d12/d3d12-shader-object.cpp
index 1b7b51106..3d38df5ab 100644
--- a/tools/gfx/d3d12/d3d12-shader-object.cpp
+++ b/tools/gfx/d3d12/d3d12-shader-object.cpp
@@ -939,6 +939,8 @@ Result ShaderObjectImpl::setResource(ShaderOffset const& offset, IResourceView*
     }
 
     ResourceViewInternalImpl* internalResourceView = nullptr;
+    auto resourceViewImpl = static_cast<ResourceViewImpl*>(resourceView);
+
     switch (resourceView->getViewDesc()->type)
     {
 #if SLANG_GFX_HAS_DXR_SUPPORT
@@ -953,7 +955,6 @@ Result ShaderObjectImpl::setResource(ShaderOffset const& offset, IResourceView*
 #endif
     default:
     {
-        auto resourceViewImpl = static_cast<ResourceViewImpl*>(resourceView);
         // Hold a reference to the resource to prevent its destruction.
         const auto resourceOffset = bindingRange.baseIndex + offset.bindingArrayIndex;
         m_boundResources[resourceOffset] = resourceViewImpl->m_resource;
@@ -964,13 +965,21 @@ Result ShaderObjectImpl::setResource(ShaderOffset const& offset, IResourceView*
     }
 
     auto descriptorSlotIndex = bindingRange.baseIndex + (int32_t)offset.bindingArrayIndex;
-    if (internalResourceView->m_descriptor.cpuHandle.ptr)
+    D3D12Descriptor srcDescriptor = {};
+
+    SLANG_RETURN_ON_FAIL(internalResourceView->getBufferDescriptorForBinding(
+        static_cast<DeviceImpl*>(m_device.get()),
+        resourceViewImpl,
+        bindingRange.bufferElementStride,
+        srcDescriptor));
+
+    if (srcDescriptor.cpuHandle.ptr)
     {
         d3dDevice->CopyDescriptorsSimple(
             1,
             m_descriptorSet.resourceTable.getCpuHandle(
                 bindingRange.baseIndex + (int32_t)offset.bindingArrayIndex),
-            internalResourceView->m_descriptor.cpuHandle,
+            srcDescriptor.cpuHandle,
             D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
     }
     else