[WebGPU] mapAsync shouldn't be calling getMappedRange on the entire b…

…uffer

https://bugs.webkit.org/show_bug.cgi?id=273028
<radar://126800249>

Reviewed by Tadeu Zagallo.

Simplify buffer mapping so we don't pass a copy of the entire buffer from the
GPU process to the web process when it is not needed.

* LayoutTests/platform/mac-wk2/TestExpectations:
Add passing expectations for api,operation,buffers,* and
api,validation,buffer,* tests.

Passing expectations already exist, tests were still marked Skipped.

* Source/WebCore/Modules/WebGPU/GPUBuffer.cpp:
(WebCore::makeArrayBuffer):
(WebCore::GPUBuffer::getMappedRange):
(WebCore::GPUBuffer::internalUnmap):
* Source/WebCore/Modules/WebGPU/GPUBuffer.h:
* Source/WebCore/Modules/WebGPU/Implementation/WebGPUBufferImpl.cpp:
(WebCore::WebGPU::BufferImpl::getBufferContents):
(WebCore::WebGPU::BufferImpl::copy):
* Source/WebCore/Modules/WebGPU/Implementation/WebGPUBufferImpl.h:
* Source/WebCore/Modules/WebGPU/InternalAPI/WebGPUBuffer.h:
* Source/WebGPU/WebGPU/Buffer.h:
* Source/WebGPU/WebGPU/Buffer.mm:
(WebGPU::Buffer::getBufferContents):
(wgpuBufferGetBufferContents):
* Source/WebGPU/WebGPU/WebGPU.h:
* Source/WebKit/GPUProcess/graphics/WebGPU/RemoteBuffer.cpp:
(WebKit::RemoteBuffer::mapAsync):
(WebKit::RemoteBuffer::getMappedRange):
(WebKit::RemoteBuffer::unmap):
(WebKit::RemoteBuffer::copy):
(WebKit::RemoteBuffer::destroy):
* Source/WebKit/GPUProcess/graphics/WebGPU/RemoteBuffer.h:
* Source/WebKit/GPUProcess/graphics/WebGPU/RemoteBuffer.messages.in:
* Source/WebKit/WebProcess/GPU/graphics/WebGPU/RemoteBufferProxy.cpp:
(WebKit::WebGPU::RemoteBufferProxy::mapAsync):
(WebKit::WebGPU::RemoteBufferProxy::getMappedRange):
(WebKit::WebGPU::RemoteBufferProxy::getBufferContents):
(WebKit::WebGPU::RemoteBufferProxy::copy):
(WebKit::WebGPU::RemoteBufferProxy::unmap):
* Source/WebKit/WebProcess/GPU/graphics/WebGPU/RemoteBufferProxy.h:
Simplify buffer mapping by storing fewer state variables refering to the same
data and remove copies.

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

LayoutTests/platform/mac-wk2/TestExpectations

@@ -1782,6 +1782,18 @@ webkit.org/b/261356 [ Debug ] fast/mediastream/device-change-event-2.html [ Pass
 webkit.org/b/263396 css3/color/text.html [ Skip ]
 
 # WebGPU
+
+[ Release ] http/tests/webgpu/webgpu/api/operation/buffers/map.html [ Pass ]
+[ Release ] http/tests/webgpu/webgpu/api/operation/buffers/map_ArrayBuffer.html [ Pass ]
+[ Release ] http/tests/webgpu/webgpu/api/operation/buffers/map_detach.html [ Pass ]
+[ Release ] http/tests/webgpu/webgpu/api/operation/buffers/map_oom.html [ Pass ]
+[ Release ] http/tests/webgpu/webgpu/api/operation/buffers/threading.html [ Pass ]
+
+[ Release ] http/tests/webgpu/webgpu/api/validation/buffer/create.html [ Pass ]
+[ Release ] http/tests/webgpu/webgpu/api/validation/buffer/destroy.html [ Pass ]
+[ Release ] http/tests/webgpu/webgpu/api/validation/buffer/mapping.html [ Pass ]
+
+# Skip tests until they can be validated to consistently pass in EWS
 http/tests/webgpu/webgpu/idl/constructable.html [ Skip ]
 http/tests/webgpu/webgpu/idl/constants/flags.html [ Skip ]
 

Source/WebCore/Modules/WebGPU/GPUBuffer.cpp

@@ -97,10 +97,10 @@ void GPUBuffer::mapAsync(GPUMapModeFlags mode, std::optional<GPUSize64> offset,
     });
 }
 
-static auto makeArrayBuffer(auto source, auto byteLength, auto& cachedArrayBuffers, auto& device, auto& buffer)
+static auto makeArrayBuffer(auto source, size_t offset, auto byteLength, auto& cachedArrayBuffers, auto& device, auto& buffer)
 {
     auto arrayBuffer = ArrayBuffer::create(source, byteLength);
-    cachedArrayBuffers.append({ arrayBuffer.ptr(), reinterpret_cast<uint8_t*>(source) });
+    cachedArrayBuffers.append({ arrayBuffer.ptr(), offset });
     cachedArrayBuffers.last().buffer->pin();
     if (device)
         device->addBufferToUnmap(buffer);
@@ -175,12 +175,12 @@ ExceptionOr<Ref<JSC::ArrayBuffer>> GPUBuffer::getMappedRange(std::optional<GPUSi
     if (!mappedRange.source) {
         m_arrayBuffers.clear();
         if (m_mappedAtCreation || !size)
-            return makeArrayBuffer(static_cast<size_t>(0U), 1, m_arrayBuffers, m_device, *this);
+            return makeArrayBuffer(static_cast<size_t>(0U), 0, 1, m_arrayBuffers, m_device, *this);
 
         return Exception { ExceptionCode::OperationError, "getMappedRange failed"_s };
     }
 
-    return makeArrayBuffer(mappedRange.source, size, m_arrayBuffers, m_device, *this);
+    return makeArrayBuffer(mappedRange.source, offset, size, m_arrayBuffers, m_device, *this);
 }
 
 void GPUBuffer::unmap(ScriptExecutionContext& scriptExecutionContext)
@@ -207,15 +207,15 @@ void GPUBuffer::internalUnmap(ScriptExecutionContext& scriptExecutionContext)
     for (auto& arrayBufferAndOffset : m_arrayBuffers) {
         auto& arrayBuffer = arrayBufferAndOffset.buffer;
         if (arrayBuffer && arrayBuffer->data() && arrayBuffer->byteLength()) {
-            memcpy(arrayBufferAndOffset.source, arrayBuffer->data(), arrayBuffer->byteLength());
+            m_backing->copy(arrayBuffer->toVector(), arrayBufferAndOffset.offset);
             JSC::ArrayBufferContents emptyBuffer;
             arrayBuffer->unpin();
             arrayBuffer->transferTo(scriptExecutionContext.vm(), emptyBuffer);
         }
     }
 
-    m_arrayBuffers.clear();
     m_backing->unmap();
+    m_arrayBuffers.clear();
 }
 
 void GPUBuffer::destroy(ScriptExecutionContext& scriptExecutionContext)

Source/WebCore/Modules/WebGPU/GPUBuffer.h

@@ -80,7 +80,7 @@ class GPUBuffer : public RefCounted<GPUBuffer>, public CanMakeWeakPtr<GPUBuffer>
     Ref<WebGPU::Buffer> m_backing;
     struct ArrayBufferWithOffset {
         RefPtr<JSC::ArrayBuffer> buffer;
-        uint8_t* source;
+        size_t offset { 0 };
     };
     Vector<ArrayBufferWithOffset> m_arrayBuffers;
     size_t m_bufferSize { 0 };

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

@@ -83,6 +83,21 @@ auto BufferImpl::getMappedRange(Size64 offset, std::optional<Size64> size) -> Ma
     return { static_cast<uint8_t*>(pointer) - actualOffset, actualSize };
 }
 
+auto BufferImpl::getBufferContents() -> MappedRange
+{
+    if (!m_backing.get())
+        return { nullptr, 0 };
+
+    auto* pointer = wgpuBufferGetBufferContents(m_backing.get());
+    auto bufferSize = wgpuBufferGetSize(m_backing.get());
+    return { static_cast<uint8_t*>(pointer), static_cast<size_t>(bufferSize) };
+}
+
+void BufferImpl::copy(Vector<uint8_t>&&, size_t)
+{
+    RELEASE_ASSERT_NOT_REACHED();
+}
+
 void BufferImpl::unmap()
 {
     wgpuBufferUnmap(m_backing.get());

Source/WebCore/Modules/WebGPU/Implementation/WebGPUBufferImpl.h

@@ -60,7 +60,9 @@ class BufferImpl final : public Buffer {
 
     void mapAsync(MapModeFlags, Size64 offset, std::optional<Size64> sizeForMap, CompletionHandler<void(bool)>&&) final;
     MappedRange getMappedRange(Size64 offset, std::optional<Size64>) final;
+    MappedRange getBufferContents() final;
     void unmap() final;
+    void copy(Vector<uint8_t>&&, size_t) final;
 
     void destroy() final;
 

Source/WebCore/Modules/WebGPU/InternalAPI/WebGPUBuffer.h

@@ -58,7 +58,8 @@ class Buffer : public RefCounted<Buffer>, public CanMakeWeakPtr<Buffer> {
     virtual void unmap() = 0;
 
     virtual void destroy() = 0;
-
+    virtual MappedRange getBufferContents() = 0;
+    virtual void copy(Vector<uint8_t>&&, size_t offset) = 0;
 protected:
     Buffer() = default;
 

Source/WebGPU/WebGPU/Buffer.h

@@ -91,6 +91,7 @@ class Buffer : public WGPUBufferImpl, public RefCounted<Buffer>, public CanMakeW
     bool isDestroyed() const;
     void setCommandEncoder(CommandEncoder&, bool mayModifyBuffer = false) const;
     uint32_t maxIndex(MTLIndexType) const;
+    uint8_t* getBufferContents();
 
 private:
     Buffer(id<MTLBuffer>, uint64_t initialSize, WGPUBufferUsageFlags, State initialState, MappingRange initialMappingRange, Device&);

Source/WebGPU/WebGPU/Buffer.mm

@@ -240,6 +240,11 @@ static size_t computeRangeSize(uint64_t size, size_t offset)
     return static_cast<char*>(m_buffer.contents) + offset;
 }
 
+uint8_t* Buffer::getBufferContents()
+{
+    return static_cast<uint8_t*>(m_buffer.contents);
+}
+
 NSString* Buffer::errorValidatingMapAsync(WGPUMapModeFlags mode, size_t offset, size_t rangeSize) const
 {
 #define ERROR_STRING(x) (@"GPUBuffer.mapAsync: " x)
@@ -445,6 +450,11 @@ WGPUBufferMapState wgpuBufferGetMapState(WGPUBuffer buffer)
     return WebGPU::fromAPI(buffer).getMappedRange(offset, size);
 }
 
+void* wgpuBufferGetBufferContents(WGPUBuffer buffer)
+{
+    return WebGPU::fromAPI(buffer).getBufferContents();
+}
+
 uint64_t wgpuBufferGetSize(WGPUBuffer buffer)
 {
     return WebGPU::fromAPI(buffer).initialSize();

Source/WebGPU/WebGPU/WebGPU.h

@@ -1597,6 +1597,7 @@ WGPU_EXPORT void wgpuBufferDestroy(WGPUBuffer buffer) WGPU_FUNCTION_ATTRIBUTE;
 WGPU_EXPORT void const * wgpuBufferGetConstMappedRange(WGPUBuffer buffer, size_t offset, size_t size) WGPU_FUNCTION_ATTRIBUTE;
 WGPU_EXPORT WGPUBufferMapState wgpuBufferGetMapState(WGPUBuffer buffer) WGPU_FUNCTION_ATTRIBUTE;
 WGPU_EXPORT void * wgpuBufferGetMappedRange(WGPUBuffer buffer, size_t offset, size_t size) WGPU_FUNCTION_ATTRIBUTE;
+WGPU_EXPORT void * wgpuBufferGetBufferContents(WGPUBuffer buffer) WGPU_FUNCTION_ATTRIBUTE;
 WGPU_EXPORT uint64_t wgpuBufferGetSize(WGPUBuffer buffer) WGPU_FUNCTION_ATTRIBUTE;
 WGPU_EXPORT WGPUBufferUsageFlags wgpuBufferGetUsage(WGPUBuffer buffer) WGPU_FUNCTION_ATTRIBUTE;
 WGPU_EXPORT void wgpuBufferMapAsync(WGPUBuffer buffer, WGPUMapModeFlags mode, size_t offset, size_t size, WGPUBufferMapCallback callback, void * userdata) WGPU_FUNCTION_ATTRIBUTE;

Source/WebKit/GPUProcess/graphics/WebGPU/RemoteBuffer.cpp

@@ -32,6 +32,8 @@
 #include "StreamServerConnection.h"
 #include "WebGPUObjectHeap.h"
 
+#include <wtf/CheckedArithmetic.h>
+
 namespace WebKit {
 
 RemoteBuffer::RemoteBuffer(WebCore::WebGPU::Buffer& buffer, WebGPU::ObjectHeap& objectHeap, Ref<IPC::StreamServerConnection>&& streamConnection, bool mappedAtCreation, WebGPUIdentifier identifier)
@@ -52,47 +54,57 @@ void RemoteBuffer::stopListeningForIPC()
     m_streamConnection->stopReceivingMessages(Messages::RemoteBuffer::messageReceiverName(), m_identifier.toUInt64());
 }
 
-void RemoteBuffer::mapAsync(WebCore::WebGPU::MapModeFlags mapModeFlags, WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> size, CompletionHandler<void(std::optional<Vector<uint8_t>>&&)>&& callback)
+void RemoteBuffer::mapAsync(WebCore::WebGPU::MapModeFlags mapModeFlags, WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> size, CompletionHandler<void(bool)>&& callback)
 {
     m_isMapped = true;
     m_mapModeFlags = mapModeFlags;
 
     m_backing->mapAsync(mapModeFlags, offset, size, [protectedThis = Ref<RemoteBuffer>(*this), callback = WTFMove(callback)] (bool success) mutable {
         if (!success) {
-            callback(std::nullopt);
+            callback(false);
             return;
         }
 
-        auto mappedRange = protectedThis->m_backing->getMappedRange(0, std::nullopt);
-        protectedThis->m_mappedRange = mappedRange;
-        callback(Vector(std::span { static_cast<const uint8_t*>(mappedRange.source), mappedRange.byteLength }));
+        callback(true);
     });
 }
 
 void RemoteBuffer::getMappedRange(WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> size, CompletionHandler<void(std::optional<Vector<uint8_t>>&&)>&& callback)
 {
     auto mappedRange = m_backing->getMappedRange(offset, size);
-    m_mappedRange = mappedRange;
     m_isMapped = true;
 
     callback(Vector(std::span { static_cast<const uint8_t*>(mappedRange.source), mappedRange.byteLength }));
 }
 
-void RemoteBuffer::unmap(Vector<uint8_t>&& data)
+void RemoteBuffer::unmap()
 {
-    if (m_isMapped && m_mappedRange && m_mappedRange->byteLength >= data.size() && m_mapModeFlags.contains(WebCore::WebGPU::MapMode::Write))
-        memcpy(m_mappedRange->source, data.data(), data.size());
-
     if (m_isMapped)
         m_backing->unmap();
     m_isMapped = false;
-    m_mappedRange = std::nullopt;
     m_mapModeFlags = { };
 }
 
+void RemoteBuffer::copy(Vector<uint8_t>&& data, size_t offset)
+{
+    if (!m_isMapped || !m_mapModeFlags.contains(WebCore::WebGPU::MapMode::Write))
+        return;
+
+    auto [buffer, bufferLength] = m_backing->getBufferContents();
+    if (!buffer || !bufferLength)
+        return;
+
+    auto dataSize = data.size();
+    auto endOffset = checkedSum<size_t>(offset, dataSize);
+    if (endOffset.hasOverflowed() || endOffset.value() > bufferLength)
+        return;
+
+    memcpy(buffer + offset, data.data(), data.size());
+}
+
 void RemoteBuffer::destroy()
 {
-    unmap(Vector<uint8_t>());
+    unmap();
     m_backing->destroy();
 }
 

Source/WebKit/GPUProcess/graphics/WebGPU/RemoteBuffer.h

@@ -79,8 +79,9 @@ class RemoteBuffer final : public IPC::StreamMessageReceiver {
     void didReceiveStreamMessage(IPC::StreamServerConnection&, IPC::Decoder&) final;
 
     void getMappedRange(WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> sizeForMap, CompletionHandler<void(std::optional<Vector<uint8_t>>&&)>&&);
-    void mapAsync(WebCore::WebGPU::MapModeFlags, WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> sizeForMap, CompletionHandler<void(std::optional<Vector<uint8_t>>&&)>&&);
-    void unmap(Vector<uint8_t>&&);
+    void mapAsync(WebCore::WebGPU::MapModeFlags, WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> sizeForMap, CompletionHandler<void(bool)>&&);
+    void copy(Vector<uint8_t>&&, size_t offset);
+    void unmap();
 
     void destroy();
     void destruct();
@@ -92,7 +93,6 @@ class RemoteBuffer final : public IPC::StreamMessageReceiver {
     Ref<IPC::StreamServerConnection> m_streamConnection;
     WebGPUIdentifier m_identifier;
     bool m_isMapped { false };
-    std::optional<WebCore::WebGPU::Buffer::MappedRange> m_mappedRange;
     WebCore::WebGPU::MapModeFlags m_mapModeFlags;
 };
 

Source/WebKit/GPUProcess/graphics/WebGPU/RemoteBuffer.messages.in

@@ -25,8 +25,9 @@
 
 messages -> RemoteBuffer NotRefCounted Stream {
     void GetMappedRange(WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> size) -> (std::optional<Vector<uint8_t>> data) Synchronous
-    void MapAsync(WebCore::WebGPU::MapModeFlags mapModeFlags, WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> size) -> (std::optional<Vector<uint8_t>> data)
-    void Unmap(Vector<uint8_t> data)
+    void MapAsync(WebCore::WebGPU::MapModeFlags mapModeFlags, WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> size) -> (bool success)
+    void Copy(Vector<uint8_t> data, size_t offset)
+    void Unmap()
     void Destroy()
     void Destruct()
     void SetLabel(String label)

Source/WebKit/WebProcess/GPU/graphics/WebGPU/RemoteBufferProxy.cpp

@@ -49,10 +49,9 @@ RemoteBufferProxy::~RemoteBufferProxy()
 
 void RemoteBufferProxy::mapAsync(WebCore::WebGPU::MapModeFlags mapModeFlags, WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> size, CompletionHandler<void(bool)>&& callback)
 {
-    auto sendResult = sendWithAsyncReply(Messages::RemoteBuffer::MapAsync(mapModeFlags, offset, size), [callback = WTFMove(callback), mapModeFlags, protectedThis = Ref { *this }](auto data) mutable {
-        if (!data)
+    auto sendResult = sendWithAsyncReply(Messages::RemoteBuffer::MapAsync(mapModeFlags, offset, size), [callback = WTFMove(callback), mapModeFlags, protectedThis = Ref { *this }](auto success) mutable {
+        if (!success)
             return callback(false);
-        protectedThis->m_data = WTFMove(data);
         protectedThis->m_mapModeFlags = mapModeFlags;
         callback(true);
     });
@@ -66,33 +65,35 @@ static bool offsetOrSizeExceedsBounds(size_t dataSize, WebCore::WebGPU::Size64 o
 
 auto RemoteBufferProxy::getMappedRange(WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> size) -> MappedRange
 {
-    if (m_data.has_value() && m_data->data()) {
-        if (offsetOrSizeExceedsBounds(m_data->size(), offset, size))
-            return { };
-
-        return { m_data->data() + offset, static_cast<size_t>(size.value_or(m_data->size() - offset)) };
-    }
-
     // FIXME: Implement error handling.
     auto sendResult = sendSync(Messages::RemoteBuffer::GetMappedRange(offset, size));
     auto [data] = sendResult.takeReplyOr(std::nullopt);
 
     if (!data || !data->data() || offsetOrSizeExceedsBounds(data->size(), offset, size))
         return { };
 
-    m_data = WTFMove(data);
-    return { m_data->data() + offset, static_cast<size_t>(size.value_or(m_data->size() - offset)) };
+    return { data->data() + offset, static_cast<size_t>(size.value_or(data->size() - offset)) };
+}
+
+auto RemoteBufferProxy::getBufferContents() -> MappedRange
+{
+    RELEASE_ASSERT_NOT_REACHED();
+}
+
+void RemoteBufferProxy::copy(Vector<uint8_t>&& data, size_t offset)
+{
+    if (!m_mapModeFlags.contains(WebCore::WebGPU::MapMode::Write))
+        return;
+
+    auto sendResult = send(Messages::RemoteBuffer::Copy(WTFMove(data), offset));
+    UNUSED_VARIABLE(sendResult);
 }
 
 void RemoteBufferProxy::unmap()
 {
-    Vector<uint8_t> data;
-    if (m_data && m_mapModeFlags.contains(WebCore::WebGPU::MapMode::Write))
-        data = WTFMove(*m_data);
-    auto sendResult = send(Messages::RemoteBuffer::Unmap(WTFMove(data)));
+    auto sendResult = send(Messages::RemoteBuffer::Unmap());
     UNUSED_VARIABLE(sendResult);
 
-    m_data = std::nullopt;
     m_mapModeFlags = { };
 }
 

Source/WebKit/WebProcess/GPU/graphics/WebGPU/RemoteBufferProxy.h

@@ -80,7 +80,9 @@ class RemoteBufferProxy final : public WebCore::WebGPU::Buffer {
 
     void mapAsync(WebCore::WebGPU::MapModeFlags, WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64> sizeForMap, CompletionHandler<void(bool)>&&) final;
     MappedRange getMappedRange(WebCore::WebGPU::Size64 offset, std::optional<WebCore::WebGPU::Size64>) final;
+    MappedRange getBufferContents() final;
     void unmap() final;
+    void copy(Vector<uint8_t>&&, size_t offset) final;
 
     void destroy() final;
 
@@ -91,7 +93,6 @@ class RemoteBufferProxy final : public WebCore::WebGPU::Buffer {
     WebGPUIdentifier m_backing;
     Ref<ConvertToBackingContext> m_convertToBackingContext;
     Ref<RemoteDeviceProxy> m_parent;
-    std::optional<Vector<uint8_t>> m_data;
     WebCore::WebGPU::MapModeFlags m_mapModeFlags;
 };