[WebGPU] bind group dynamic offsets are not implemented correctly, th…

…ey don't correspond to the proper stage

https://bugs.webkit.org/show_bug.cgi?id=266815>
<radar://120049732>

Reviewed by Tadeu Zagallo.

Dynamic offsets assumed buffers had linearly increasing bind
group indices, which does not necessarily hold and is not
the case in Unity's demos.

Break this assumption by computing the corresponding index
from the source dynamic offsets buffer during bind group
layout generation time.

Then use this precomputed index instead of a linearly increasing
value in the PipelineLayout.

* Source/WebGPU/WebGPU/BindGroupLayout.h:
* Source/WebGPU/WebGPU/BindGroupLayout.mm:
(WebGPU::Device::createBindGroupLayout):
* Source/WebGPU/WebGPU/PipelineLayout.mm:
(WebGPU::PipelineLayout::offsetVectorForBindGroup):

Canonical link: https://commits.webkit.org/272522@main

Source/WebGPU/WebGPU/BindGroupLayout.h

@@ -63,6 +63,7 @@ class BindGroupLayout : public WGPUBindGroupLayoutImpl, public RefCounted<BindGr
         std::optional<uint32_t> vertexDynamicOffset;
         std::optional<uint32_t> fragmentDynamicOffset;
         std::optional<uint32_t> computeDynamicOffset;
+        uint32_t dynamicOffsetsIndex;
     };
     using EntriesContainer = HashMap<uint32_t, Entry, DefaultHash<uint32_t>, WTF::UnsignedWithZeroKeyHashTraits<uint32_t>>;
 

Source/WebGPU/WebGPU/BindGroupLayout.mm

@@ -342,6 +342,27 @@ static bool containsStage(WGPUShaderStageFlags stageBitfield, auto stage)
             }
         }
 
+        auto hasDynamicOffsets = ^(size_t* sizeOfDynamicOffsets, size_t stageCount) {
+            for (size_t i = 0; i < stageCount; ++i) {
+                if (sizeOfDynamicOffsets[i])
+                    return true;
+            }
+
+            return false;
+        };
+        HashMap<uint32_t, uint32_t, DefaultHash<uint32_t>, WTF::UnsignedWithZeroKeyHashTraits<uint32_t>> dynamicOffsetIndices;
+        if (hasDynamicOffsets(sizeOfDynamicOffsets, stageCount)) {
+            std::sort(descriptorEntries.begin(), descriptorEntries.end(), [](const WGPUBindGroupLayoutEntry& a, const WGPUBindGroupLayoutEntry& b) {
+                return a.binding < b.binding;
+            });
+            uint32_t nextDynamicOffsetIndex = 0;
+            for (auto& entry : descriptorEntries) {
+                if (entry.buffer.hasDynamicOffset)
+                    dynamicOffsetIndices.add(entry.binding, nextDynamicOffsetIndex++);
+            }
+        }
+
+        auto dynamicOffsetsIt = dynamicOffsetIndices.find(entry.binding);
         bindGroupLayoutEntries.add(entry.binding, BindGroupLayout::Entry {
             .binding = entry.binding,
             .visibility = entry.visibility,
@@ -350,7 +371,8 @@ static bool containsStage(WGPUShaderStageFlags stageBitfield, auto stage)
             .bufferSizeArgumentBufferIndices = WTFMove(bufferSizeArgumentBufferIndices),
             .vertexDynamicOffset = WTFMove(dynamicOffsets[0]),
             .fragmentDynamicOffset = WTFMove(dynamicOffsets[1]),
-            .computeDynamicOffset = WTFMove(dynamicOffsets[2])
+            .computeDynamicOffset = WTFMove(dynamicOffsets[2]),
+            .dynamicOffsetsIndex = dynamicOffsetsIt == dynamicOffsetIndices.end() ? std::numeric_limits<uint32_t>::max() : dynamicOffsetsIt->value
         });
     }
 

Source/WebGPU/WebGPU/PipelineLayout.mm

@@ -209,7 +209,7 @@ static size_t returnOffsetOfGroup0(auto& v, auto groupIndex)
     auto& bindGroupLayouts = *m_bindGroupLayouts;
     if (auto it = stageOffsets.find(bindGroupIndex); it != stageOffsets.end()) {
         auto& container = it->value;
-        uint32_t dynamicOffsetIndex = 0, stageOffsetIndex = 0;
+        uint32_t stageOffsetIndex = 0;
         auto& bindGroupLayout = bindGroupLayouts[bindGroupIndex];
         for (auto& entryKvp : bindGroupLayout->entries()) {
             auto& entry = entryKvp.value;
@@ -218,12 +218,10 @@ static size_t returnOffsetOfGroup0(auto& v, auto groupIndex)
                 continue;
 
             if (entry.visibility & stage) {
-                ASSERT(container.size() > stageOffsetIndex && dynamicOffsets.size() > dynamicOffsetIndex);
-                container[stageOffsetIndex] = dynamicOffsets[dynamicOffsetIndex];
+                RELEASE_ASSERT(container.size() > stageOffsetIndex && dynamicOffsets.size() > entry.dynamicOffsetsIndex);
+                container[stageOffsetIndex] = dynamicOffsets[entry.dynamicOffsetsIndex];
                 ++stageOffsetIndex;
             }
-
-            ++dynamicOffsetIndex;
         }
 
         return &container;