diff --git a/apps/docs/content/api/gpu-device-extensions.mdx b/apps/docs/content/api/gpu-device-extensions.mdx index 6bab25583..bd751894c 100644 --- a/apps/docs/content/api/gpu-device-extensions.mdx +++ b/apps/docs/content/api/gpu-device-extensions.mdx @@ -1,10 +1,44 @@ --- title: Native Extensions -description: Native interop extensions to GPUDevice and the WebGPU globals. +description: React Native extensions to WebGPU objects and globals. --- React Native WebGPU adds a handful of native extensions on top of the WebGPU spec. +## GPUBuffer.readbackSync + +`readbackSync()` synchronously copies a small `MAP_READ` buffer from the GPU into an owned CPU-accessible `ArrayBuffer`. It blocks the calling JavaScript or worklet thread until previously submitted GPU work completes, copies the mapped bytes, and unmaps the GPU buffer before returning. + + + Synchronous readback is not faster than `mapAsync()`. It stalls the calling + thread and removes CPU/GPU parallelism. Keep dependent work on the GPU when + possible, and use `mapAsync()` when the result can be consumed in a later + JavaScript turn. Use `readbackSync()` only when a small CPU result is genuinely + required in the current turn. + + +Create a `MAP_READ | COPY_DST` staging buffer, copy the compute result into it, submit the commands, and read it: + +```tsx +const staging = device.createBuffer({ + size: 16, + usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, +}); + +encoder.copyBufferToBuffer(storage, 0, staging, 0, 16); +device.queue.submit([encoder.finish()]); + +const values = new Uint32Array(staging.readbackSync(0, 16, 500)); +``` + +| Parameter | Default | Description | +| ----------- | --------------------- | ---------------------------------------- | +| `offset` | `0` | Byte offset where the read begins | +| `size` | Remaining buffer size | Number of bytes to copy, capped at 1 MiB | +| `timeoutMs` | `2000` | Maximum time to wait for the GPU | + +`readbackSync()` throws if the range is invalid, the read exceeds 1 MiB, mapping fails, the timeout elapses, or the Dawn instance does not support timed waits. The last case can occur with an externally provided Dawn instance. + ## importSharedTextureMemory Imports a native buffer (IOSurface / AHardwareBuffer) for zero-copy sampling. diff --git a/packages/webgpu/cpp/rnwgpu/ArrayBuffer.h b/packages/webgpu/cpp/rnwgpu/ArrayBuffer.h index 3e6d38cb7..f177b9216 100644 --- a/packages/webgpu/cpp/rnwgpu/ArrayBuffer.h +++ b/packages/webgpu/cpp/rnwgpu/ArrayBuffer.h @@ -13,17 +13,27 @@ namespace jsi = facebook::jsi; struct ArrayBuffer : jsi::MutableBuffer { ArrayBuffer(void *data, size_t size, size_t bytesPerElement) - : _data(data), _size(size), _bytesPerElement(bytesPerElement) {} + : _data(data), _size(size), _bytesPerElement(bytesPerElement), + _ownsData(false) {} - ~ArrayBuffer() override {} + ArrayBuffer(size_t size, size_t bytesPerElement) + : _ownedData(size == 0 ? nullptr : new uint8_t[size]), + _data(_ownedData.get()), _size(size), _bytesPerElement(bytesPerElement), + _ownsData(true) {} + + ~ArrayBuffer() override = default; size_t size() const override { return _size; } uint8_t *data() override { return static_cast(_data); } + bool ownsData() const { return _ownsData; } + + std::unique_ptr _ownedData; void *_data; size_t _size; size_t _bytesPerElement; + bool _ownsData; }; static std::shared_ptr @@ -110,7 +120,11 @@ template <> struct JSIConverter> { static jsi::Value toJSI(jsi::Runtime &runtime, std::shared_ptr arg) { - return jsi::ArrayBuffer(runtime, arg); + jsi::ArrayBuffer arrayBuffer(runtime, arg); + if (arg->ownsData()) { + arrayBuffer.setExternalMemoryPressure(runtime, arg->size()); + } + return jsi::Value(runtime, arrayBuffer); } }; diff --git a/packages/webgpu/cpp/rnwgpu/api/GPUBuffer.cpp b/packages/webgpu/cpp/rnwgpu/api/GPUBuffer.cpp index db698fff7..4e14973b8 100644 --- a/packages/webgpu/cpp/rnwgpu/api/GPUBuffer.cpp +++ b/packages/webgpu/cpp/rnwgpu/api/GPUBuffer.cpp @@ -1,6 +1,10 @@ #include "GPUBuffer.h" +#include +#include +#include #include +#include #include #include "Convertors.h" @@ -35,6 +39,105 @@ GPUBuffer::getMappedRange(std::optional o, std::optional size) { return array_buffer; } +std::shared_ptr +GPUBuffer::readbackSync(std::optional offsetIn, + std::optional sizeIn, + std::optional timeoutMsIn) { + // Synchronous small-buffer readback: blocks the calling thread until all + // previously submitted GPU work using this buffer completes, then returns + // a copy of the mapped bytes. Built for tiny compute results (landmarks, + // ranges, counters) that must be consumed in the SAME frame - the async + // mapAsync path forces at least one frame of staleness in render loops + // that cannot await. Requires MAP_READ usage (pair with COPY_DST and copy + // into this buffer from your storage buffer). The wait uses + // Instance::WaitAny, which this library's instance enables via the + // TimedWaitAny feature at creation; external (Skia-provided) instances + // without it fail the wait and throw rather than hang. + auto toByteSize = [](const char *name, double value) -> size_t { + constexpr double kMaxSafeInteger = 9'007'199'254'740'991.0; + if (!std::isfinite(value) || value < 0 || std::floor(value) != value || + value > kMaxSafeInteger || + value > static_cast(std::numeric_limits::max())) { + throw std::runtime_error(std::string("GPUBuffer.readbackSync ") + name + + " must be a non-negative safe integer"); + } + return static_cast(value); + }; + + const size_t offset = + offsetIn.has_value() ? toByteSize("offset", *offsetIn) : 0; + const uint64_t bufferSize = _instance.GetSize(); + if (offset > bufferSize) { + throw std::runtime_error( + "GPUBuffer.readbackSync offset exceeds the buffer size"); + } + const size_t size = sizeIn.has_value() + ? toByteSize("size", *sizeIn) + : static_cast(bufferSize - offset); + if (size > bufferSize - offset) { + throw std::runtime_error( + "GPUBuffer.readbackSync range exceeds the buffer size"); + } + + constexpr size_t kMaxReadbackSyncBytes = 1 << 20; + if (size > kMaxReadbackSyncBytes) { + throw std::runtime_error( + "GPUBuffer.readbackSync is limited to 1 MiB; use mapAsync for larger " + "readbacks"); + } + + constexpr double kDefaultTimeoutMs = 2'000.0; + constexpr double kNanosecondsPerMillisecond = 1'000'000.0; + const double timeoutMs = timeoutMsIn.value_or(kDefaultTimeoutMs); + const double maxTimeoutMs = + static_cast(std::numeric_limits::max()) / + kNanosecondsPerMillisecond; + if (!std::isfinite(timeoutMs) || timeoutMs < 0 || timeoutMs > maxTimeoutMs) { + throw std::runtime_error( + "GPUBuffer.readbackSync timeoutMs must be a finite, non-negative " + "number"); + } + const uint64_t timeoutNs = + static_cast(timeoutMs * kNanosecondsPerMillisecond); + + struct MapResult { + wgpu::MapAsyncStatus status = wgpu::MapAsyncStatus::Error; + std::string message = "callback never ran"; + }; + auto mapResult = std::make_shared(); + auto future = _instance.MapAsync( + wgpu::MapMode::Read, offset, size, wgpu::CallbackMode::WaitAnyOnly, + [mapResult](wgpu::MapAsyncStatus status, wgpu::StringView message) { + mapResult->status = status; + mapResult->message = std::string(message); + }); + auto waitStatus = _async->instance().WaitAny(future, timeoutNs); + if (waitStatus != wgpu::WaitStatus::Success) { + // Cancels the pending map request. The callback owns MapResult so a late + // completion cannot access stack memory after this method returns. + _instance.Unmap(); + throw std::runtime_error( + "GPUBuffer.readbackSync did not complete before timeoutMs, or the Dawn " + "instance does not support timed waits"); + } + if (mapResult->status != wgpu::MapAsyncStatus::Success) { + throw std::runtime_error("GPUBuffer.readbackSync mapping failed: " + + mapResult->message); + } + const void *ptr = _instance.GetConstMappedRange(offset, size); + if (ptr == nullptr) { + _instance.Unmap(); + throw std::runtime_error( + "GPUBuffer.readbackSync could not access the mapped range"); + } + // Allocate owned native storage. The JSI converter wraps this memory without + // copying it again and reports its external memory pressure to the runtime. + auto result = std::make_shared(size, 1); + memcpy(result->data(), ptr, size); + _instance.Unmap(); + return result; +} + void GPUBuffer::destroy() { _instance.Destroy(); } async::AsyncTaskHandle GPUBuffer::mapAsync(jsi::Runtime &runtime, diff --git a/packages/webgpu/cpp/rnwgpu/api/GPUBuffer.h b/packages/webgpu/cpp/rnwgpu/api/GPUBuffer.h index 1b5cf123e..ede646201 100644 --- a/packages/webgpu/cpp/rnwgpu/api/GPUBuffer.h +++ b/packages/webgpu/cpp/rnwgpu/api/GPUBuffer.h @@ -38,6 +38,9 @@ class GPUBuffer : public NativeObject { std::optional size); std::shared_ptr getMappedRange(std::optional offset, std::optional size); + std::shared_ptr readbackSync(std::optional offset, + std::optional size, + std::optional timeoutMs); void unmap(); void destroy(); @@ -57,6 +60,7 @@ class GPUBuffer : public NativeObject { &GPUBuffer::mapAsync); installMethod(runtime, prototype, "getMappedRange", &GPUBuffer::getMappedRange); + installMethod(runtime, prototype, "readbackSync", &GPUBuffer::readbackSync); installMethod(runtime, prototype, "unmap", &GPUBuffer::unmap); installMethod(runtime, prototype, "destroy", &GPUBuffer::destroy); installGetter(runtime, prototype, "size", &GPUBuffer::getSize); diff --git a/packages/webgpu/src/__tests__/ReadbackSync.spec.ts b/packages/webgpu/src/__tests__/ReadbackSync.spec.ts new file mode 100644 index 000000000..fbded86dc --- /dev/null +++ b/packages/webgpu/src/__tests__/ReadbackSync.spec.ts @@ -0,0 +1,123 @@ +import { client } from "./setup"; + +describe("readbackSync", () => { + it("reads back compute results synchronously, same tick", async () => { + const result = await client.eval(({ device }) => { + const storage = device.createBuffer({ + size: 16, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC, + }); + const staging = device.createBuffer({ + size: 16, + usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, + }); + const module = device.createShaderModule({ + code: ` +@group(0) @binding(0) var out: array; +@compute @workgroup_size(1) +fn main() { + out[0] = 1u; + out[1] = 2u; + out[2] = 3u; + out[3] = 42u; +}`, + }); + const pipeline = device.createComputePipeline({ + layout: "auto", + compute: { module, entryPoint: "main" }, + }); + const bindGroup = device.createBindGroup({ + layout: pipeline.getBindGroupLayout(0), + entries: [{ binding: 0, resource: { buffer: storage } }], + }); + const encoder = device.createCommandEncoder(); + const pass = encoder.beginComputePass(); + pass.setPipeline(pipeline); + pass.setBindGroup(0, bindGroup); + pass.dispatchWorkgroups(1); + pass.end(); + encoder.copyBufferToBuffer(storage, 0, staging, 0, 16); + device.queue.submit([encoder.finish()]); + // No await between submit and read: the readback is synchronous. + return Array.from(new Uint32Array(staging.readbackSync())); + }); + expect(result).toEqual([1, 2, 3, 42]); + }); + it("respects offset and size", async () => { + const result = await client.eval(({ device }) => { + const staging = device.createBuffer({ + size: 16, + usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, + }); + device.queue.writeBuffer( + staging, + 0, + new Uint32Array([10, 20, 30, 40]).buffer, + ); + return Array.from(new Uint32Array(staging.readbackSync(8, 8))); + }); + expect(result).toEqual([30, 40]); + }); + it("is repeatable on the same buffer", async () => { + const result = await client.eval(({ device }) => { + const staging = device.createBuffer({ + size: 4, + usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, + }); + const reads: number[] = []; + for (let i = 0; i < 3; i++) { + device.queue.writeBuffer(staging, 0, new Uint32Array([i]).buffer); + reads.push(new Uint32Array(staging.readbackSync())[0]!); + } + return reads; + }); + expect(result).toEqual([0, 1, 2]); + }); + it("rejects reads above the size cap", async () => { + const result = await client.eval(({ device }) => { + const staging = device.createBuffer({ + size: 2 * 1024 * 1024, + usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, + }); + try { + staging.readbackSync(); + return "no error"; + } catch (e) { + return e instanceof Error && e.message.includes("limited to 1 MiB") + ? "capped" + : "wrong error"; + } + }); + expect(result).toBe("capped"); + }); + it("accepts a configurable timeout", async () => { + const result = await client.eval(({ device }) => { + const staging = device.createBuffer({ + size: 4, + usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, + }); + device.queue.writeBuffer(staging, 0, new Uint32Array([42]).buffer); + return new Uint32Array( + staging.readbackSync(undefined, undefined, 5_000), + )[0]; + }); + expect(result).toBe(42); + }); + it("rejects an invalid timeout", async () => { + const result = await client.eval(({ device }) => { + const staging = device.createBuffer({ + size: 4, + usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, + }); + try { + staging.readbackSync(undefined, undefined, -1); + return "no error"; + } catch (e) { + return e instanceof Error && e.message.includes("timeoutMs") + ? "invalid timeout" + : "wrong error"; + } + }); + expect(result).toBe("invalid timeout"); + }); +}); diff --git a/packages/webgpu/src/index.tsx b/packages/webgpu/src/index.tsx index cd895e20d..b1f5d705b 100644 --- a/packages/webgpu/src/index.tsx +++ b/packages/webgpu/src/index.tsx @@ -80,6 +80,30 @@ declare global { readonly nativePointer: bigint; } + interface GPUBuffer { + /** + * Blocks the calling thread until pending GPU work completes, then returns + * a copy of this buffer's mapped bytes. + * + * This React Native extension is intended for results up to 1 MiB that + * must be consumed in the current JavaScript turn. Prefer `mapAsync()` + * whenever a later turn is acceptable because synchronous readback stalls + * the calling thread and reduces CPU/GPU parallelism. + * + * The buffer must have `GPUBufferUsage.MAP_READ` usage. The returned + * `ArrayBuffer` owns its storage and remains valid after this method + * unmaps the GPU buffer. + * + * @throws If the range is invalid, exceeds 1 MiB, mapping fails, the wait + * times out, or the Dawn instance does not support timed waits. + */ + readbackSync( + offset?: number, + size?: number, + timeoutMs?: number, + ): ArrayBuffer; + } + interface GPUDevice { importSharedTextureMemory( descriptor: GPUSharedTextureMemoryDescriptor,