[WebGPU] Bind group layout with multiple external textures does not work

https://bugs.webkit.org/show_bug.cgi?id=262550
<radar://116403353>

Reviewed by Dan Glastonbury.

When external textures are used, they can expand to more than one bind slot
per source, so binding.index != (not necessarily) [id(k)].

Build a mapping and pass it to the compiler. The reverse we don't yet need,
since the bind group layout returned from the compiler breaks down the external
texture into its subcomponents.

* Source/WebGPU/WGSL/WGSL.h:
* Source/WebGPU/WebGPU/BindGroup.mm:
(WebGPU::Device::createBindGroup):
* Source/WebGPU/WebGPU/BindGroupLayout.h:
* Source/WebGPU/WebGPU/BindGroupLayout.mm:
(WebGPU::Device::createBindGroupLayout):
(WebGPU::BindGroupLayout::argumentBufferIndexForEntryIndex const):
* Source/WebGPU/WebGPU/RenderPipeline.mm:
(WebGPU::Device::addPipelineLayouts):
* Source/WebGPU/WebGPU/ShaderModule.mm:
(WebGPU::ShaderModule::convertPipelineLayout):

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

Source/WebCore/Modules/WebGPU/Implementation/WebGPUDeviceImpl.cpp

@@ -176,23 +176,12 @@ Ref<BindGroupLayout> DeviceImpl::createBindGroupLayout(const BindGroupLayoutDesc
 {
     auto label = descriptor.label.utf8();
 
-    auto backingExternalTextureEntries = descriptor.entries.map([](const auto&) {
-        return WGPUExternalTextureBindGroupLayoutEntry {
-            {
-                nullptr,
-                static_cast<WGPUSType>(WGPUSTypeExtended_BindGroupLayoutEntryExternalTexture),
-            }, {
-                nullptr,
-            },
-        };
-    });
-
     Vector<WGPUBindGroupLayoutEntry> backingEntries;
     backingEntries.reserveInitialCapacity(descriptor.entries.size());
     for (size_t i = 0; i < descriptor.entries.size(); ++i) {
         const auto& entry = descriptor.entries[i];
         backingEntries.uncheckedAppend({
-            entry.externalTexture ? &backingExternalTextureEntries[i].chain : nullptr,
+            nullptr,
             entry.binding,
             m_convertToBackingContext->convertShaderStageFlagsToBacking(entry.visibility), {
                 nullptr,
@@ -204,9 +193,9 @@ Ref<BindGroupLayout> DeviceImpl::createBindGroupLayout(const BindGroupLayoutDesc
                 entry.sampler ? m_convertToBackingContext->convertToBacking(entry.sampler->type) : WGPUSamplerBindingType_Undefined,
             }, {
                 nullptr,
-                entry.texture ? m_convertToBackingContext->convertToBacking(entry.texture->sampleType) : WGPUTextureSampleType_Undefined,
-                entry.texture ? m_convertToBackingContext->convertToBacking(entry.texture->viewDimension) : WGPUTextureViewDimension_Undefined,
-                entry.texture ? entry.texture->multisampled : false,
+                entry.externalTexture ? static_cast<WGPUTextureSampleType>(WGPUTextureSampleType_ExternalTexture) : (entry.texture ? m_convertToBackingContext->convertToBacking(entry.texture->sampleType) : WGPUTextureSampleType_Undefined),
+                (!entry.externalTexture && entry.texture) ? m_convertToBackingContext->convertToBacking(entry.texture->viewDimension) : WGPUTextureViewDimension_Undefined,
+                (!entry.externalTexture && entry.texture) ? entry.texture->multisampled : false,
             }, {
                 nullptr,
                 entry.storageTexture ? m_convertToBackingContext->convertToBacking(entry.storageTexture->access) : WGPUStorageTextureAccess_Undefined,

Source/WebGPU/WGSL/WGSL.h

@@ -150,8 +150,11 @@ struct BindGroupLayoutEntry {
     OptionSet<ShaderStage> visibility;
     using BindingMember = std::variant<BufferBindingLayout, SamplerBindingLayout, TextureBindingLayout, StorageTextureBindingLayout, ExternalTextureBindingLayout>;
     BindingMember bindingMember;
+    std::optional<uint32_t> vertexArgumentBufferIndex;
     std::optional<uint32_t> vertexBufferDynamicOffset;
+    std::optional<uint32_t> fragmentArgumentBufferIndex;
     std::optional<uint32_t> fragmentBufferDynamicOffset;
+    std::optional<uint32_t> computeArgumentBufferIndex;
     std::optional<uint32_t> computeBufferDynamicOffset;
 };
 

Source/WebGPU/WebGPU/BindGroup.mm

@@ -477,7 +477,7 @@ static MTLResourceUsage resourceUsageForBindingAcccess(BindGroupLayout::BindingA
             if (!optionalAccess)
                 continue;
 
-            auto index = entry.binding;
+            auto index = bindGroupLayout.argumentBufferIndexForEntryIndex(entry.binding, stage);
             MTLResourceUsage resourceUsage = resourceUsageForBindingAcccess(*optionalAccess);
 
             if (bufferIsPresent) {

Source/WebGPU/WebGPU/BindGroupLayout.h

@@ -25,6 +25,7 @@
 
 #pragma once
 
+#import <wtf/EnumeratedArray.h>
 #import <wtf/FastMalloc.h>
 #import <wtf/HashMap.h>
 #import <wtf/HashTraits.h>
@@ -49,11 +50,15 @@ class Device;
 class BindGroupLayout : public WGPUBindGroupLayoutImpl, public RefCounted<BindGroupLayout> {
     WTF_MAKE_FAST_ALLOCATED;
 public:
+    template <typename T>
+    using ShaderStageArray = EnumeratedArray<ShaderStage, T, ShaderStage::Compute>;
+    using ArgumentBufferIndices = ShaderStageArray<std::optional<uint32_t>>;
     struct Entry {
         uint32_t binding;
         WGPUShaderStageFlags visibility;
         using BindingLayout = std::variant<WGPUBufferBindingLayout, WGPUSamplerBindingLayout, WGPUTextureBindingLayout, WGPUStorageTextureBindingLayout, WGPUExternalTextureBindingLayout>;
         BindingLayout bindingLayout;
+        ArgumentBufferIndices argumentBufferIndices;
         std::optional<uint32_t> vertexDynamicOffset;
         std::optional<uint32_t> fragmentDynamicOffset;
         std::optional<uint32_t> computeDynamicOffset;
@@ -96,12 +101,12 @@ class BindGroupLayout : public WGPUBindGroupLayoutImpl, public RefCounted<BindGr
     id<MTLArgumentEncoder> computeArgumentEncoder() const { return m_computeArgumentEncoder; }
 
     std::optional<StageMapValue> bindingAccessForBindingIndex(uint32_t bindingIndex, ShaderStage renderStage) const;
+    NSUInteger argumentBufferIndexForEntryIndex(uint32_t bindingIndex, ShaderStage renderStage) const;
 
     static bool isPresent(const WGPUBufferBindingLayout&);
     static bool isPresent(const WGPUSamplerBindingLayout&);
     static bool isPresent(const WGPUTextureBindingLayout&);
     static bool isPresent(const WGPUStorageTextureBindingLayout&);
-    static bool isPresent(const WGPUExternalTextureBindingLayout&);
 
     const EntriesContainer& entries() const { return m_bindGroupLayoutEntries; }
     uint32_t sizeOfVertexDynamicOffsets() const;

Source/WebGPU/WebGPU/BindGroupLayout.mm

@@ -126,24 +126,6 @@ static uint64_t makeKey(uint32_t firstInteger, auto secondValue)
     return descriptor;
 }
 
-bool BindGroupLayout::isPresent(const WGPUExternalTextureBindingLayout&)
-{
-    return true;
-}
-
-static MTLArgumentDescriptor *createArgumentDescriptor(const WGPUExternalTextureBindingLayout& externalTexture)
-{
-    if (externalTexture.nextInChain)
-        return nil;
-
-    // FIXME: Implement this properly.
-    // External textures have a bunch of information included in them, not just a single texture.
-    auto descriptor = [MTLArgumentDescriptor new];
-    descriptor.dataType = MTLDataTypeTexture;
-    descriptor.access = BindGroupLayout::BindingAccessReadOnly;
-    return descriptor;
-}
-
 static void addDescriptor(NSMutableArray<MTLArgumentDescriptor *> *arguments, MTLArgumentDescriptor *descriptor, NSUInteger index)
 {
     MTLArgumentDescriptor *stageDescriptor = [descriptor copy];
@@ -169,65 +151,95 @@ static bool containsStage(WGPUShaderStageFlags stageBitfield, auto stage)
     for (uint32_t i = 0; i < stageCount; ++i)
         arguments[i] = [NSMutableArray arrayWithCapacity:descriptor.entryCount];
 
+    Vector<WGPUBindGroupLayoutEntry> descriptorEntries(descriptor.entries, descriptor.entryCount);
+    std::sort(descriptorEntries.begin(), descriptorEntries.end(), [](const WGPUBindGroupLayoutEntry& a, const WGPUBindGroupLayoutEntry& b) {
+        return a.binding < b.binding;
+    });
+
+    constexpr ShaderStage stages[] = { ShaderStage::Vertex, ShaderStage::Fragment, ShaderStage::Compute };
     BindGroupLayout::EntriesContainer bindGroupLayoutEntries;
     size_t sizeOfDynamicOffsets[stageCount] = { 0, 0, 0 };
-    for (uint32_t i = 0; i < descriptor.entryCount; ++i) {
-        const WGPUBindGroupLayoutEntry& entry = descriptor.entries[i];
+    uint32_t bindingOffset[stageCount] = { 0, 0, 0 };
+    for (auto& entry : descriptorEntries) {
         if (entry.nextInChain) {
             if (entry.nextInChain->sType != static_cast<WGPUSType>(WGPUSTypeExtended_BindGroupLayoutEntryExternalTexture))
                 return BindGroupLayout::createInvalid(*this);
             if (entry.nextInChain->next)
                 return BindGroupLayout::createInvalid(*this);
         }
 
-        MTLArgumentDescriptor *descriptor = nil;
+        bool isExternalTexture = false;
+        constexpr int maxGeneratedDescriptors = 4;
+        MTLArgumentDescriptor *descriptors[maxGeneratedDescriptors] = { nil, nil, nil, nil };
         BindGroupLayout::Entry::BindingLayout bindingLayout;
         auto processBindingLayout = [&](const auto& type) {
             if (!BindGroupLayout::isPresent(type))
                 return true;
-            if (descriptor)
+            if (descriptors[0])
                 return false;
-            descriptor = createArgumentDescriptor(type);
-            if (!descriptor)
+            descriptors[0] = createArgumentDescriptor(type);
+            if (!descriptors[0])
                 return false;
             bindingLayout = type;
             return true;
         };
-        if (!processBindingLayout(entry.buffer))
-            return BindGroupLayout::createInvalid(*this);
-        if (!processBindingLayout(entry.sampler))
-            return BindGroupLayout::createInvalid(*this);
-        if (!processBindingLayout(entry.texture))
-            return BindGroupLayout::createInvalid(*this);
-        if (!processBindingLayout(entry.storageTexture))
-            return BindGroupLayout::createInvalid(*this);
-        if (entry.nextInChain) {
-            const WGPUExternalTextureBindGroupLayoutEntry& externalTextureEntry = *reinterpret_cast<const WGPUExternalTextureBindGroupLayoutEntry*>(entry.nextInChain);
-            ASSERT(externalTextureEntry.chain.sType == static_cast<WGPUSType>(WGPUSTypeExtended_BindGroupLayoutEntryExternalTexture));
-            processBindingLayout(entry.storageTexture);
-            if (!processBindingLayout(externalTextureEntry.externalTexture))
+        if (entry.texture.sampleType == WGPUTextureSampleType_ExternalTexture) {
+            isExternalTexture = true;
+            descriptors[0] = createArgumentDescriptor(entry.texture);
+            descriptors[1] = createArgumentDescriptor(entry.texture);
+            WGPUBufferBindingLayout bufferLayout {
+                .nextInChain = nullptr,
+                .type = static_cast<WGPUBufferBindingType>(WGPUBufferBindingType_Float3x2),
+                .hasDynamicOffset = static_cast<WGPUBool>(false),
+                .minBindingSize = 0
+            };
+            descriptors[2] = createArgumentDescriptor(bufferLayout);
+            bufferLayout.type = static_cast<WGPUBufferBindingType>(WGPUBufferBindingType_Float4x3);
+            descriptors[3] = createArgumentDescriptor(bufferLayout);
+            bindingLayout = WGPUExternalTextureBindingLayout();
+        } else {
+            if (!processBindingLayout(entry.buffer))
+                return BindGroupLayout::createInvalid(*this);
+            if (!processBindingLayout(entry.sampler))
+                return BindGroupLayout::createInvalid(*this);
+            if (!processBindingLayout(entry.texture))
+                return BindGroupLayout::createInvalid(*this);
+            if (!processBindingLayout(entry.storageTexture))
                 return BindGroupLayout::createInvalid(*this);
         }
 
-        if (!descriptor)
+        if (!descriptors[0])
             return BindGroupLayout::createInvalid(*this);
 
-        std::optional<uint32_t> dynamicOffsets[stageCount] = { std::nullopt, std::nullopt, std::nullopt };
+        std::optional<uint32_t> dynamicOffsets[stageCount];
+        BindGroupLayout::ArgumentBufferIndices argumentBufferIndices;
         for (uint32_t stage = 0; stage < stageCount; ++stage) {
             if (containsStage(entry.visibility, stage)) {
-                indicesForBinding.add(makeKey(entry.binding, stage), descriptor.access);
-                addDescriptor(arguments[stage], descriptor, entry.binding);
+                indicesForBinding.add(makeKey(entry.binding, stage), descriptors[0].access);
+                auto renderStage = stages[stage];
+                argumentBufferIndices[renderStage] = entry.binding + bindingOffset[stage];
                 if (entry.buffer.hasDynamicOffset) {
                     dynamicOffsets[stage] = sizeOfDynamicOffsets[stage];
                     sizeOfDynamicOffsets[stage] += sizeof(uint32_t);
                 }
+
+                for (int descriptorIndex = 0; descriptorIndex < maxGeneratedDescriptors; ++descriptorIndex) {
+                    if (MTLArgumentDescriptor *descriptor = descriptors[descriptorIndex])
+                        addDescriptor(arguments[stage], descriptor, *argumentBufferIndices[renderStage] + descriptorIndex);
+                    else
+                        break;
+                }
+
+                if (isExternalTexture)
+                    bindingOffset[stage] += maxGeneratedDescriptors;
             }
         }
 
         bindGroupLayoutEntries.add(entry.binding, BindGroupLayout::Entry {
             entry.binding,
             entry.visibility,
             WTFMove(bindingLayout),
+            WTFMove(argumentBufferIndices),
             WTFMove(dynamicOffsets[0]),
             WTFMove(dynamicOffsets[1]),
             WTFMove(dynamicOffsets[2])
@@ -237,15 +249,7 @@ static bool containsStage(WGPUShaderStageFlags stageBitfield, auto stage)
     auto label = fromAPI(descriptor.label);
     id<MTLArgumentEncoder> argumentEncoders[stageCount];
     for (size_t stage = 0; stage < stageCount; ++stage) {
-        NSArray<MTLArgumentDescriptor *> *sortedArray = [arguments[stage] sortedArrayUsingComparator:^NSComparisonResult(MTLArgumentDescriptor *a, MTLArgumentDescriptor *b) {
-            if (a.index < b.index)
-                return NSOrderedAscending;
-            if (a.index == b.index)
-                return NSOrderedSame;
-
-            return NSOrderedDescending;
-        }];
-        argumentEncoders[stage] = arguments[stage].count ? [m_device newArgumentEncoderWithArguments:sortedArray] : nil;
+        argumentEncoders[stage] = arguments[stage].count ? [m_device newArgumentEncoderWithArguments:arguments[stage]] : nil;
         argumentEncoders[stage].label = label;
         if (arguments[stage].count && !argumentEncoders[stage])
             return BindGroupLayout::createInvalid(*this);
@@ -317,6 +321,16 @@ static bool containsStage(WGPUShaderStageFlags stageBitfield, auto stage)
     return m_sizeOfComputeDynamicOffsets;
 }
 
+NSUInteger BindGroupLayout::argumentBufferIndexForEntryIndex(uint32_t bindingIndex, ShaderStage renderStage) const
+{
+    if (auto it = m_bindGroupLayoutEntries.find(bindingIndex); it != m_bindGroupLayoutEntries.end()) {
+        auto result = it->value.argumentBufferIndices[renderStage];
+        return result ? *result : NSNotFound;
+    }
+
+    return NSNotFound;
+}
+
 } // namespace WebGPU
 
 #pragma mark WGPU Stubs

Source/WebGPU/WebGPU/RenderPipeline.mm

@@ -442,10 +442,16 @@ static WGPUTextureViewDimension convertViewDimension(WGSL::TextureViewDimension
         auto& bindGroupLayout = pipelineLayout.bindGroupLayouts[pipelineLayoutIndex];
         auto& entries = pipelineEntries[pipelineLayoutIndex];
         for (auto& entry : bindGroupLayout.entries) {
+            if (auto existingIndex = entries.findIf([&](auto& existingEntry) {
+                return existingEntry.binding == entry.binding;
+            }); existingIndex != WTF::notFound) {
+                entries[existingIndex].visibility |= convertVisibility(entry.visibility);
+                continue;
+            }
+
             WGPUBindGroupLayoutEntry newEntry = { };
             newEntry.binding = entry.binding;
             newEntry.visibility = convertVisibility(entry.visibility);
-            bool isExternalTexture = false;
             WTF::switchOn(entry.bindingMember, [&](const WGSL::BufferBindingLayout& bufferBinding) {
                 newEntry.buffer = WGPUBufferBindingLayout {
                     .nextInChain = nullptr,
@@ -473,42 +479,15 @@ static WGPUTextureViewDimension convertViewDimension(WGSL::TextureViewDimension
                     .viewDimension = convertViewDimension(storageTexture.viewDimension)
                 };
             }, [&](const WGSL::ExternalTextureBindingLayout&) {
-                isExternalTexture = true;
                 newEntry.texture = WGPUTextureBindingLayout {
                     .nextInChain = nullptr,
-                    .sampleType = WGPUTextureSampleType_Float,
+                    .sampleType = static_cast<WGPUTextureSampleType>(WGPUTextureSampleType_ExternalTexture),
                     .viewDimension = WGPUTextureViewDimension_2D,
                     .multisampled = false
                 };
             });
 
             entries.append(newEntry);
-            // FIXME: - https://bugs.webkit.org/show_bug.cgi?id=257978
-            // perform breakdown in BindGroupLayout.mm
-            if (isExternalTexture) {
-                ++newEntry.binding;
-                entries.append(newEntry);
-
-                WGPUBindGroupLayoutEntry bufferEntry = { };
-                bufferEntry.binding = newEntry.binding + 1;
-                bufferEntry.visibility = newEntry.binding;
-                bufferEntry.buffer = WGPUBufferBindingLayout {
-                    .nextInChain = nullptr,
-                    .type = static_cast<WGPUBufferBindingType>(WGPUBufferBindingType_Float3x2),
-                    .hasDynamicOffset = false,
-                    .minBindingSize = 0,
-                };
-                entries.append(bufferEntry);
-
-                bufferEntry.binding = newEntry.binding + 2;
-                bufferEntry.buffer = WGPUBufferBindingLayout {
-                    .nextInChain = nullptr,
-                    .type = static_cast<WGPUBufferBindingType>(WGPUBufferBindingType_Float4x3),
-                    .hasDynamicOffset = false,
-                    .minBindingSize = 0,
-                };
-                entries.append(bufferEntry);
-            }
         }
     }
 }

Source/WebGPU/WebGPU/ShaderModule.mm

@@ -367,8 +367,11 @@ static auto wgslViewDimension(WGPUTextureViewDimension viewDimension)
             wgslEntry.binding = entry.value.binding;
             wgslEntry.visibility.fromRaw(entry.value.visibility);
             wgslEntry.bindingMember = convertBindingLayout(entry.value.bindingLayout);
+            wgslEntry.vertexArgumentBufferIndex = entry.value.argumentBufferIndices[WebGPU::ShaderStage::Vertex];
             wgslEntry.vertexBufferDynamicOffset = entry.value.vertexDynamicOffset;
+            wgslEntry.fragmentArgumentBufferIndex = entry.value.argumentBufferIndices[WebGPU::ShaderStage::Fragment];
             wgslEntry.fragmentBufferDynamicOffset = entry.value.fragmentDynamicOffset;
+            wgslEntry.computeArgumentBufferIndex = entry.value.argumentBufferIndices[WebGPU::ShaderStage::Compute];
             wgslEntry.computeBufferDynamicOffset = entry.value.computeDynamicOffset;
             wgslBindGroupLayout.entries.append(wgslEntry);
         }

Source/WebGPU/WebGPU/WebGPUExt.h

@@ -68,6 +68,8 @@ typedef struct WGPUInstanceCocoaDescriptor {
     __unsafe_unretained WGPUScheduleWorkBlock scheduleWorkBlock;
 } WGPUInstanceCocoaDescriptor;
 
+const int WGPUTextureSampleType_ExternalTexture = WGPUTextureSampleType_Force32 - 1;
+
 typedef void (^WGPURenderBuffersWereRecreatedBlockCallback)(CFArrayRef ioSurfaces);
 typedef void (^WGPUOnSubmittedWorkScheduledCallback)(WGPUWorkItem);
 typedef void (^WGPUCompositorIntegrationRegisterBlockCallback)(WGPURenderBuffersWereRecreatedBlockCallback renderBuffersWereRecreated, WGPUOnSubmittedWorkScheduledCallback onSubmittedWorkScheduledCallback);
@@ -80,11 +82,6 @@ typedef struct WGPUExternalTextureBindingLayout {
     WGPUChainedStruct const * nextInChain;
 } WGPUExternalTextureBindingLayout;
 
-typedef struct WGPUExternalTextureBindGroupLayoutEntry {
-    WGPUChainedStruct chain;
-    WGPUExternalTextureBindingLayout externalTexture;
-} WGPUExternalTextureBindGroupLayoutEntry;
-
 typedef struct WGPUBindGroupExternalTextureEntry {
     WGPUChainedStruct chain;
     WGPUExternalTexture externalTexture;