Optimize Dawn Wire WriteBuffer

Reduces the number of memcpys that are done when uploading buffer data
via WriteBuffer using the wire client. Does so by directly managing the
allocation of shared memory and writing to that rather than copying the
data to the wire command buffer first, which may involve chunking and
additional copies.

In tests so far this performs up to 2X better than the previous
WriteBuffer implementation, depending on data size. (Larger uploads
typically see a bigger improvement.)

Derived from Loko's prototype:
https://dawn-review.googlesource.com/c/dawn/+/269134
https://dawn-review.googlesource.com/c/dawn/+/269814

Alternative to my previous effort to move the WriteBuffer implementation
client-side: https://dawn-review.googlesource.com/c/dawn/+/270895

Bug: 441900745
Change-Id: I1d7aa1b7f64e2ac6582512d64a4b981f2e208cd6
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/271535
Reviewed-by: Kai Ninomiya <kainino@chromium.org>
Commit-Queue: Brandon Jones <bajones@chromium.org>
Reviewed-by: Loko Kung <lokokung@google.com>
diff --git a/include/dawn/wire/WireServer.h b/include/dawn/wire/WireServer.h
index 0d53226..d64fc57 100644
--- a/include/dawn/wire/WireServer.h
+++ b/include/dawn/wire/WireServer.h
@@ -140,6 +140,11 @@
                                            size_t size) = 0;
         std::span<uint8_t> GetTarget() const;
 
+        // Returns a direct pointer to the source data that will
+        // be copied into Target in DeserializeDataUpdate if accessible, nullptr
+        // otherwise.
+        virtual uint8_t* GetSourceData() const { return nullptr; }
+
       private:
         WriteHandle(const WriteHandle&) = delete;
         WriteHandle& operator=(const WriteHandle&) = delete;
diff --git a/src/dawn/dawn_wire.json b/src/dawn/dawn_wire.json
index 0fd2c01..c5d7387 100644
--- a/src/dawn/dawn_wire.json
+++ b/src/dawn/dawn_wire.json
@@ -87,8 +87,11 @@
             {"name": "queue id", "type": "ObjectId", "id_type": "queue" },
             {"name": "buffer id", "type": "ObjectId", "id_type": "buffer" },
             {"name": "buffer offset", "type": "uint64_t"},
-            {"name": "data", "type": "uint8_t", "annotation": "const*", "length": "size", "wire_is_data_only": true},
-            {"name": "size", "type": "uint64_t"}
+            { "name": "size", "type": "uint64_t"},
+            { "name": "write handle create info length", "type": "uint64_t" },
+            { "name": "write handle create info", "type": "uint8_t", "annotation": "const*", "length": "write handle create info length", "skip_serialize": true},
+            { "name": "write data update info length", "type": "uint64_t" },
+            { "name": "write data update info", "type": "uint8_t", "annotation": "const*", "length": "write data update info length", "skip_serialize": true}
         ],
         "queue write texture": [
             {"name": "queue id", "type": "ObjectId", "id_type": "queue" },
diff --git a/src/dawn/wire/client/Queue.cpp b/src/dawn/wire/client/Queue.cpp
index a724441..490d1ec 100644
--- a/src/dawn/wire/client/Queue.cpp
+++ b/src/dawn/wire/client/Queue.cpp
@@ -25,6 +25,11 @@
 // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+#ifdef UNSAFE_BUFFERS_BUILD
+// TODO(crbug.com/439062058): Remove this and convert code to safer constructs.
+#pragma allow_unsafe_buffers
+#endif
+
 #include "dawn/wire/client/Queue.h"
 
 #include <memory>
@@ -32,6 +37,7 @@
 #include <utility>
 
 #include "dawn/common/StringViewUtils.h"
+#include "dawn/wire/BufferConsumer_impl.h"
 #include "dawn/wire/client/Client.h"
 #include "dawn/wire/client/EventManager.h"
 #include "partition_alloc/pointers/raw_ptr.h"
@@ -121,15 +127,43 @@
                            const void* data,
                            size_t size) {
     Buffer* buffer = FromAPI(cBuffer);
+    Client* client = GetClient();
+
+    // Create write handle and prepare to serialize command.
+    size_t writeHandleCreateInfoLength = 0;
+    std::unique_ptr<MemoryTransferService::WriteHandle> writeHandle(
+        client->GetMemoryTransferService()->CreateWriteHandle(size));
+    if (writeHandle == nullptr) {
+        // Trigger a device loss.
+        client->Disconnect();
+        return;
+    }
+    writeHandleCreateInfoLength = writeHandle->SerializeCreateSize();
+
+    // Write the data to the allocated memory.
+    memcpy(writeHandle->GetData(), data, size);
+
+    // Prepare to serialize data update command.
+    size_t writeDataUpdateInfoLength = writeHandle->SizeOfSerializeDataUpdate(0u, size);
 
     QueueWriteBufferCmd cmd;
-    cmd.queueId = GetWireHandle(GetClient()).id;
-    cmd.bufferId = buffer->GetWireHandle(GetClient()).id;
+    cmd.queueId = GetWireHandle(client).id;
+    cmd.bufferId = buffer->GetWireHandle(client).id;
     cmd.bufferOffset = bufferOffset;
-    cmd.data = static_cast<const uint8_t*>(data);
     cmd.size = size;
+    cmd.writeHandleCreateInfoLength = writeHandleCreateInfoLength;
+    cmd.writeHandleCreateInfo = nullptr;
+    cmd.writeDataUpdateInfoLength = writeDataUpdateInfoLength;
+    cmd.writeDataUpdateInfo = nullptr;
 
-    GetClient()->SerializeCommand(cmd);
+    client->SerializeCommand(
+        cmd,
+        CommandExtension{
+            writeHandleCreateInfoLength,
+            [&](char* writeHandleBuffer) { writeHandle->SerializeCreate(writeHandleBuffer); }},
+        CommandExtension{writeDataUpdateInfoLength, [&](char* writeHandleBuffer) {
+                             writeHandle->SerializeDataUpdate(writeHandleBuffer, 0u, cmd.size);
+                         }});
 }
 
 void Queue::APIWriteTexture(const WGPUTexelCopyTextureInfo* destination,
diff --git a/src/dawn/wire/server/ServerQueue.cpp b/src/dawn/wire/server/ServerQueue.cpp
index 8ce8973..d35355f 100644
--- a/src/dawn/wire/server/ServerQueue.cpp
+++ b/src/dawn/wire/server/ServerQueue.cpp
@@ -26,6 +26,7 @@
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 #include <limits>
+#include <memory>
 
 #include "dawn/common/Assert.h"
 #include "dawn/wire/server/Server.h"
@@ -61,13 +62,46 @@
 WireResult Server::DoQueueWriteBuffer(Known<WGPUQueue> queue,
                                       Known<WGPUBuffer> buffer,
                                       uint64_t bufferOffset,
-                                      const uint8_t* data,
-                                      uint64_t size) {
+                                      uint64_t size,
+                                      uint64_t writeHandleCreateInfoLength,
+                                      const uint8_t* writeHandleCreateInfo,
+                                      uint64_t writeDataUpdateInfoLength,
+                                      const uint8_t* writeDataUpdateInfo) {
     if (size > std::numeric_limits<size_t>::max()) {
         return WireResult::FatalError;
     }
 
-    mProcs.queueWriteBuffer(queue->handle, buffer->handle, bufferOffset, data,
+    MemoryTransferService::WriteHandle* writeHandle = nullptr;
+    // Deserialize metadata produced from the client to create a companion server handle.
+    if (!mMemoryTransferService->DeserializeWriteHandle(
+            writeHandleCreateInfo, static_cast<size_t>(writeHandleCreateInfoLength),
+            &writeHandle)) {
+        return WireResult::FatalError;
+    }
+
+    // Try first to use GetSourceData if the memory transfer service implements
+    // it. If so, we can avoid a copy.
+    uint8_t* sourceData = writeHandle->GetSourceData();
+    if (sourceData) {
+        mProcs.queueWriteBuffer(queue->handle, buffer->handle, bufferOffset, sourceData,
+                                static_cast<size_t>(size));
+        return WireResult::Success;
+    }
+
+    // Otherwise, fall back to DeserializeDataUpdate.
+    auto backingData = std::make_unique<char[]>(size);
+    writeHandle->SetTarget(backingData.get());
+    writeHandle->SetDataLength(size);
+
+    // Deserialize the flush info and flush updated data from the handle into the target
+    // of the handle that's just a temporary allocation from above right now.
+    if (!writeHandle->DeserializeDataUpdate(writeDataUpdateInfo,
+                                            static_cast<size_t>(writeDataUpdateInfoLength), 0u,
+                                            static_cast<size_t>(size))) {
+        return WireResult::FatalError;
+    }
+
+    mProcs.queueWriteBuffer(queue->handle, buffer->handle, bufferOffset, backingData.get(),
                             static_cast<size_t>(size));
     return WireResult::Success;
 }