react-native-webgpu 0.5.11

package/cpp/jsi/RuntimeAwareCache.h

@@ -0,0 +1,100 @@
+#pragma once
+
+#include <jsi/jsi.h>
+
+#include <cassert>
+#include <memory>
+#include <unordered_map>
+#include <utility>
+
+#include "RuntimeLifecycleMonitor.h"
+
+namespace rnwgpu {
+
+namespace jsi = facebook::jsi;
+
+class BaseRuntimeAwareCache {
+public:
+  static void setMainJsRuntime(jsi::Runtime *rt) { _mainRuntime = rt; }
+  static jsi::Runtime *getMainJsRuntime() {
+    assert(_mainRuntime != nullptr &&
+           "Expected main Javascript runtime to be set in the "
+           "BaseRuntimeAwareCache class.");
+
+    return _mainRuntime;
+  }
+
+private:
+  static jsi::Runtime *_mainRuntime;
+};
+
+/**
+ * Provides a way to keep data specific to a jsi::Runtime instance that gets
+ * cleaned up when that runtime is destroyed. This is necessary because JSI does
+ * not allow for its associated objects to be retained past the runtime
+ * lifetime. If an object (e.g. jsi::Values or jsi::Function instances) is kept
+ * after the runtime is torn down, its destructor (once it is destroyed
+ * eventually) will result in a crash (JSI objects keep a pointer to memory
+ * managed by the runtime, accessing that portion of the memory after runtime is
+ * deleted is the root cause of that crash).
+ *
+ * In order to provide an efficient implementation that does not add an overhead
+ * for the cases when only a single runtime is used, which is the primary
+ * usecase, the following assumption has been made: Only for secondary runtimes
+ * we track destruction and clean up the store associated with that runtime. For
+ * the first runtime we assume that the object holding the store is destroyed
+ * prior to the destruction of that runtime.
+ *
+ * The above assumption makes it work without any overhead when only single
+ * runtime is in use. Specifically, we don't perform any additional operations
+ * related to tracking runtime lifecycle when only a single runtime is used.
+ */
+template <typename T>
+class RuntimeAwareCache : public BaseRuntimeAwareCache,
+                          public RuntimeLifecycleListener {
+
+public:
+  void onRuntimeDestroyed(jsi::Runtime *rt) override {
+    if (getMainJsRuntime() != rt) {
+      // We are removing a secondary runtime
+      _secondaryRuntimeCaches.erase(rt);
+    }
+  }
+
+  ~RuntimeAwareCache() {
+    for (auto &cache : _secondaryRuntimeCaches) {
+      RuntimeLifecycleMonitor::removeListener(
+          *static_cast<jsi::Runtime *>(cache.first), this);
+    }
+  }
+
+  T &get(jsi::Runtime &rt) {
+    // We check if we're accessing the main runtime - this is the happy path
+    // to avoid us having to lookup by runtime for caches that only has a single
+    // runtime
+    if (getMainJsRuntime() == &rt) {
+      return _primaryCache;
+    } else {
+      if (_secondaryRuntimeCaches.count(&rt) == 0) {
+        // we only add listener when the secondary runtime is used, this assumes
+        // that the secondary runtime is terminated first. This lets us avoid
+        // additional complexity for the majority of cases when objects are not
+        // shared between runtimes. Otherwise we'd have to register all objecrts
+        // with the RuntimeMonitor as opposed to only registering ones that are
+        // used in secondary runtime. Note that we can't register listener here
+        // with the primary runtime as it may run on a separate thread.
+        RuntimeLifecycleMonitor::addListener(rt, this);
+
+        T cache;
+        _secondaryRuntimeCaches.emplace(&rt, std::move(cache));
+      }
+    }
+    return _secondaryRuntimeCaches.at(&rt);
+  }
+
+private:
+  std::unordered_map<void *, T> _secondaryRuntimeCaches;
+  T _primaryCache;
+};
+
+} // namespace rnwgpu

package/cpp/jsi/RuntimeLifecycleMonitor.cpp

@@ -0,0 +1,72 @@
+#include "RuntimeLifecycleMonitor.h"
+
+#include <mutex>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+
+namespace rnwgpu {
+
+static std::unordered_map<jsi::Runtime *,
+                          std::unordered_set<RuntimeLifecycleListener *>>
+    listeners;
+static std::mutex listenersMutex;
+
+struct RuntimeLifecycleMonitorObject : public jsi::HostObject {
+  jsi::Runtime *_rt;
+  explicit RuntimeLifecycleMonitorObject(jsi::Runtime *rt) : _rt(rt) {}
+  ~RuntimeLifecycleMonitorObject() {
+    std::unordered_set<RuntimeLifecycleListener *> listenersCopy;
+    {
+      std::lock_guard<std::mutex> lock(listenersMutex);
+      auto listenersSet = listeners.find(_rt);
+      if (listenersSet != listeners.end()) {
+        listenersCopy = listenersSet->second;
+        listeners.erase(listenersSet);
+      }
+    }
+    for (auto listener : listenersCopy) {
+      listener->onRuntimeDestroyed(_rt);
+    }
+  }
+};
+
+void RuntimeLifecycleMonitor::addListener(jsi::Runtime &rt,
+                                          RuntimeLifecycleListener *listener) {
+  bool shouldInstallMonitor = false;
+  {
+    std::lock_guard<std::mutex> lock(listenersMutex);
+    auto listenersSet = listeners.find(&rt);
+    if (listenersSet == listeners.end()) {
+      std::unordered_set<RuntimeLifecycleListener *> newSet;
+      newSet.insert(listener);
+      listeners.emplace(&rt, std::move(newSet));
+      shouldInstallMonitor = true;
+    } else {
+      listenersSet->second.insert(listener);
+    }
+  }
+  if (shouldInstallMonitor) {
+    // We install a global host object in the provided runtime, this way we can
+    // use that host object destructor to get notified when the runtime is being
+    // terminated. We use a unique name for the object as it gets saved with the
+    // runtime's global object.
+    rt.global().setProperty(
+        rt, "__rnwgpu_rt_lifecycle_monitor",
+        jsi::Object::createFromHostObject(
+            rt, std::make_shared<RuntimeLifecycleMonitorObject>(&rt)));
+  }
+}
+
+void RuntimeLifecycleMonitor::removeListener(
+    jsi::Runtime &rt, RuntimeLifecycleListener *listener) {
+  std::lock_guard<std::mutex> lock(listenersMutex);
+  auto listenersSet = listeners.find(&rt);
+  if (listenersSet == listeners.end()) {
+    // nothing to do here
+  } else {
+    listenersSet->second.erase(listener);
+  }
+}
+
+} // namespace rnwgpu

package/cpp/jsi/RuntimeLifecycleMonitor.h

@@ -0,0 +1,32 @@
+#pragma once
+
+#include <jsi/jsi.h>
+
+#include <memory>
+
+namespace rnwgpu {
+
+namespace jsi = facebook::jsi;
+
+/**
+ * Listener interface that allows for getting notified when a jsi::Runtime
+ * instance is destroyed.
+ */
+struct RuntimeLifecycleListener {
+  virtual ~RuntimeLifecycleListener() {}
+  virtual void onRuntimeDestroyed(jsi::Runtime *) = 0;
+};
+
+/**
+ * This class provides an API via static methods for registering and
+ * unregistering runtime lifecycle listeners. The listeners can be used to
+ * cleanup any data that references a given jsi::Runtime instance before it gets
+ * destroyed.
+ */
+struct RuntimeLifecycleMonitor {
+  static void addListener(jsi::Runtime &rt, RuntimeLifecycleListener *listener);
+  static void removeListener(jsi::Runtime &rt,
+                             RuntimeLifecycleListener *listener);
+};
+
+} // namespace rnwgpu

package/cpp/rnwgpu/ArrayBuffer.h

@@ -0,0 +1,73 @@
+#pragma once
+#include <jsi/jsi.h>
+
+#include <memory>
+
+#include "JSIConverter.h"
+
+namespace rnwgpu {
+
+namespace jsi = facebook::jsi;
+
+struct ArrayBuffer : jsi::MutableBuffer {
+  ArrayBuffer(void *data, size_t size, size_t bytesPerElement)
+      : _data(data), _size(size), _bytesPerElement(bytesPerElement) {}
+
+  ~ArrayBuffer() override {}
+
+  size_t size() const override { return _size; }
+
+  uint8_t *data() override { return static_cast<uint8_t *>(_data); }
+
+  void *_data;
+  size_t _size;
+  size_t _bytesPerElement;
+};
+
+static std::shared_ptr<ArrayBuffer>
+createArrayBufferFromJSI(jsi::Runtime &runtime,
+                         const jsi::ArrayBuffer &arrayBuffer,
+                         size_t bytesPerElement) {
+  auto size = arrayBuffer.size(runtime);
+  return std::make_shared<ArrayBuffer>(arrayBuffer.data(runtime), size,
+                                       bytesPerElement);
+}
+
+template <> struct JSIConverter<std::shared_ptr<ArrayBuffer>> {
+  static std::shared_ptr<ArrayBuffer>
+  fromJSI(jsi::Runtime &runtime, const jsi::Value &arg, bool outOfBound) {
+    if (arg.isObject()) {
+      auto obj = arg.getObject(runtime);
+      if (obj.isArrayBuffer(runtime)) {
+        return createArrayBufferFromJSI(runtime, obj.getArrayBuffer(runtime),
+                                        1);
+      }
+      if (obj.hasProperty(runtime, "buffer")) {
+        auto bufferProp = obj.getProperty(runtime, "buffer");
+        if (bufferProp.isObject() &&
+            bufferProp.getObject(runtime).isArrayBuffer(runtime)) {
+          auto buff = bufferProp.getObject(runtime);
+          auto bytesPerElements =
+              obj.getProperty(runtime, "BYTES_PER_ELEMENT").asNumber();
+          auto arrayBuffer = buff.getArrayBuffer(runtime);
+          auto byteOffset = static_cast<size_t>(
+              obj.getProperty(runtime, "byteOffset").asNumber());
+          auto byteLength = static_cast<size_t>(
+              obj.getProperty(runtime, "byteLength").asNumber());
+          return std::make_shared<ArrayBuffer>(
+              arrayBuffer.data(runtime) + byteOffset, byteLength,
+              static_cast<size_t>(bytesPerElements));
+        }
+      }
+    }
+    throw std::runtime_error("ArrayBuffer::fromJSI: argument is not an object "
+                             "with an ArrayBuffer 'buffer' property");
+  }
+
+  static jsi::Value toJSI(jsi::Runtime &runtime,
+                          std::shared_ptr<ArrayBuffer> arg) {
+    return jsi::ArrayBuffer(runtime, arg);
+  }
+};
+
+} // namespace rnwgpu

package/cpp/rnwgpu/PlatformContext.h

@@ -0,0 +1,46 @@
+#pragma once
+
+#include <functional>
+#include <memory>
+#include <span>
+#include <string>
+#include <vector>
+
+#include "webgpu/webgpu_cpp.h"
+
+namespace rnwgpu {
+
+struct ImageData {
+  std::vector<uint8_t> data;
+  size_t width;
+  size_t height;
+  wgpu::TextureFormat format;
+};
+
+class PlatformContext {
+public:
+  PlatformContext() = default;
+  virtual ~PlatformContext() = default;
+
+  virtual wgpu::Surface makeSurface(wgpu::Instance instance, void *surface,
+                                    int width, int height) = 0;
+  virtual ImageData createImageBitmap(std::string blobId, double offset,
+                                      double size) = 0;
+
+  // Async version that performs image decoding on a background thread
+  virtual void
+  createImageBitmapAsync(std::string blobId, double offset, double size,
+                         std::function<void(ImageData)> onSuccess,
+                         std::function<void(std::string)> onError) = 0;
+
+  // Create ImageBitmap from raw encoded image bytes (PNG/JPEG/etc.)
+  virtual ImageData
+  createImageBitmapFromData(std::span<const uint8_t> data) = 0;
+
+  virtual void
+  createImageBitmapFromDataAsync(std::span<const uint8_t> data,
+                                 std::function<void(ImageData)> onSuccess,
+                                 std::function<void(std::string)> onError) = 0;
+};
+
+} // namespace rnwgpu

package/cpp/rnwgpu/RNWebGPUManager.cpp

@@ -0,0 +1,225 @@
+#include "RNWebGPUManager.h"
+
+#include "GPU.h"
+#include "NativeObject.h"
+#include "RNWebGPU.h"
+
+// GPU API classes (for instanceof support)
+#include "GPUAdapter.h"
+#include "GPUAdapterInfo.h"
+#include "GPUBindGroup.h"
+#include "GPUBindGroupLayout.h"
+#include "GPUBuffer.h"
+#include "GPUCanvasContext.h"
+#include "GPUCommandBuffer.h"
+#include "GPUCommandEncoder.h"
+#include "GPUCompilationInfo.h"
+#include "GPUCompilationMessage.h"
+#include "GPUComputePassEncoder.h"
+#include "GPUComputePipeline.h"
+#include "GPUDevice.h"
+#include "GPUDeviceLostInfo.h"
+#include "GPUError.h"
+#include "GPUExternalTexture.h"
+#include "GPUInternalError.h"
+#include "GPUOutOfMemoryError.h"
+#include "GPUPipelineLayout.h"
+#include "GPUQuerySet.h"
+#include "GPUQueue.h"
+#include "GPURenderBundle.h"
+#include "GPURenderBundleEncoder.h"
+#include "GPURenderPassEncoder.h"
+#include "GPURenderPipeline.h"
+#include "GPUSampler.h"
+#include "GPUShaderModule.h"
+#include "GPUSupportedLimits.h"
+#include "GPUTexture.h"
+#include "GPUTextureView.h"
+#include "GPUUncapturedErrorEvent.h"
+#include "GPUValidationError.h"
+
+// Enums
+#include "GPUBufferUsage.h"
+#include "GPUColorWrite.h"
+#include "GPUMapMode.h"
+#include "GPUShaderStage.h"
+#include "GPUTextureUsage.h"
+
+#include <memory>
+#include <utility>
+
+namespace rnwgpu {
+
+RNWebGPUManager::RNWebGPUManager(
+    jsi::Runtime *jsRuntime,
+    std::shared_ptr<facebook::react::CallInvoker> jsCallInvoker,
+    std::shared_ptr<PlatformContext> platformContext)
+    : _jsRuntime(jsRuntime), _jsCallInvoker(jsCallInvoker),
+      _platformContext(platformContext) {
+
+  // Register main runtime for RuntimeAwareCache
+  BaseRuntimeAwareCache::setMainJsRuntime(_jsRuntime);
+
+  auto gpu = std::make_shared<GPU>(*_jsRuntime);
+  auto rnWebGPU =
+      std::make_shared<RNWebGPU>(gpu, _platformContext, _jsCallInvoker);
+  _gpu = gpu->get();
+  _jsRuntime->global().setProperty(*_jsRuntime, "RNWebGPU",
+                                   RNWebGPU::create(*_jsRuntime, rnWebGPU));
+
+  // Install constructors for instanceof support
+  GPU::installConstructor(*_jsRuntime);
+  GPUAdapter::installConstructor(*_jsRuntime);
+  GPUAdapterInfo::installConstructor(*_jsRuntime);
+  GPUBindGroup::installConstructor(*_jsRuntime);
+  GPUBindGroupLayout::installConstructor(*_jsRuntime);
+  GPUBuffer::installConstructor(*_jsRuntime);
+  GPUCanvasContext::installConstructor(*_jsRuntime);
+  GPUCommandBuffer::installConstructor(*_jsRuntime);
+  GPUCommandEncoder::installConstructor(*_jsRuntime);
+  GPUCompilationInfo::installConstructor(*_jsRuntime);
+  GPUCompilationMessage::installConstructor(*_jsRuntime);
+  GPUComputePassEncoder::installConstructor(*_jsRuntime);
+  GPUComputePipeline::installConstructor(*_jsRuntime);
+  GPUDevice::installConstructor(*_jsRuntime);
+  GPUDeviceLostInfo::installConstructor(*_jsRuntime);
+  GPUError::installConstructor(*_jsRuntime);
+  GPUExternalTexture::installConstructor(*_jsRuntime);
+  GPUInternalError::installConstructor(*_jsRuntime);
+  GPUOutOfMemoryError::installConstructor(*_jsRuntime);
+  GPUValidationError::installConstructor(*_jsRuntime);
+  GPUUncapturedErrorEvent::installConstructor(*_jsRuntime);
+  GPUPipelineLayout::installConstructor(*_jsRuntime);
+  GPUQuerySet::installConstructor(*_jsRuntime);
+  GPUQueue::installConstructor(*_jsRuntime);
+  GPURenderBundle::installConstructor(*_jsRuntime);
+  GPURenderBundleEncoder::installConstructor(*_jsRuntime);
+  GPURenderPassEncoder::installConstructor(*_jsRuntime);
+  GPURenderPipeline::installConstructor(*_jsRuntime);
+  GPUSampler::installConstructor(*_jsRuntime);
+  GPUShaderModule::installConstructor(*_jsRuntime);
+  GPUSupportedLimits::installConstructor(*_jsRuntime);
+  GPUTexture::installConstructor(*_jsRuntime);
+  GPUTextureView::installConstructor(*_jsRuntime);
+
+  // Install constant objects as plain JS objects with own properties
+  _jsRuntime->global().setProperty(*_jsRuntime, "GPUBufferUsage",
+                                   GPUBufferUsage::create(*_jsRuntime));
+  _jsRuntime->global().setProperty(*_jsRuntime, "GPUColorWrite",
+                                   GPUColorWrite::create(*_jsRuntime));
+  _jsRuntime->global().setProperty(*_jsRuntime, "GPUMapMode",
+                                   GPUMapMode::create(*_jsRuntime));
+  _jsRuntime->global().setProperty(*_jsRuntime, "GPUShaderStage",
+                                   GPUShaderStage::create(*_jsRuntime));
+  _jsRuntime->global().setProperty(*_jsRuntime, "GPUTextureUsage",
+                                   GPUTextureUsage::create(*_jsRuntime));
+
+  // Install global helper functions for Worklets serialization
+  // These are standalone functions that don't require RNWebGPU instance
+  installWebGPUWorkletHelpers(*_jsRuntime);
+}
+
+void RNWebGPUManager::installWebGPUWorkletHelpers(jsi::Runtime &runtime) {
+  // __webgpuIsWebGPUObject - checks if a value is a WebGPU NativeObject
+  auto isWebGPUObjectFunc = jsi::Function::createFromHostFunction(
+      runtime, jsi::PropNameID::forUtf8(runtime, "__webgpuIsWebGPUObject"), 1,
+      [](jsi::Runtime &rt, const jsi::Value & /*thisVal*/,
+         const jsi::Value *args, size_t count) -> jsi::Value {
+        if (count < 1 || !args[0].isObject()) {
+          return jsi::Value(false);
+        }
+        auto obj = args[0].getObject(rt);
+
+        // Check if it has native state
+        if (!obj.hasNativeState(rt)) {
+          return jsi::Value(false);
+        }
+
+        // Check if it has Symbol.toStringTag on its prototype (WebGPU objects
+        // do)
+        auto objectCtor = rt.global().getPropertyAsObject(rt, "Object");
+        auto getPrototypeOf =
+            objectCtor.getPropertyAsFunction(rt, "getPrototypeOf");
+        auto proto = getPrototypeOf.call(rt, obj);
+
+        if (!proto.isObject()) {
+          return jsi::Value(false);
+        }
+
+        auto protoObj = proto.getObject(rt);
+        auto symbolCtor = rt.global().getPropertyAsObject(rt, "Symbol");
+        auto toStringTag = symbolCtor.getProperty(rt, "toStringTag");
+        if (toStringTag.isUndefined()) {
+          return jsi::Value(false);
+        }
+
+        auto getOwnPropertyDescriptor =
+            objectCtor.getPropertyAsFunction(rt, "getOwnPropertyDescriptor");
+        auto desc = getOwnPropertyDescriptor.call(rt, protoObj, toStringTag);
+        return jsi::Value(desc.isObject());
+      });
+  runtime.global().setProperty(runtime, "__webgpuIsWebGPUObject",
+                               std::move(isWebGPUObjectFunc));
+
+  // __webgpuBox - boxes a WebGPU object for Worklets serialization
+  auto boxFunc = jsi::Function::createFromHostFunction(
+      runtime, jsi::PropNameID::forUtf8(runtime, "__webgpuBox"), 1,
+      [](jsi::Runtime &rt, const jsi::Value & /*thisVal*/,
+         const jsi::Value *args, size_t count) -> jsi::Value {
+        if (count < 1 || !args[0].isObject()) {
+          throw jsi::JSError(rt,
+                             "__webgpuBox() requires a WebGPU object argument");
+        }
+
+        auto obj = args[0].getObject(rt);
+
+        // Check if it has native state
+        if (!obj.hasNativeState(rt)) {
+          throw jsi::JSError(
+              rt, "Object has no native state - not a WebGPU object");
+        }
+
+        // Get the brand name from Symbol.toStringTag on the prototype
+        auto objectCtor = rt.global().getPropertyAsObject(rt, "Object");
+        auto getPrototypeOf =
+            objectCtor.getPropertyAsFunction(rt, "getPrototypeOf");
+        auto proto = getPrototypeOf.call(rt, obj);
+
+        std::string brand;
+        if (proto.isObject()) {
+          auto protoObj = proto.getObject(rt);
+          auto symbolCtor = rt.global().getPropertyAsObject(rt, "Symbol");
+          auto toStringTag = symbolCtor.getProperty(rt, "toStringTag");
+          if (!toStringTag.isUndefined()) {
+            auto getOwnPropertyDescriptor = objectCtor.getPropertyAsFunction(
+                rt, "getOwnPropertyDescriptor");
+            auto desc =
+                getOwnPropertyDescriptor.call(rt, protoObj, toStringTag);
+            if (desc.isObject()) {
+              auto descObj = desc.getObject(rt);
+              auto value = descObj.getProperty(rt, "value");
+              if (value.isString()) {
+                brand = value.getString(rt).utf8(rt);
+              }
+            }
+          }
+        }
+
+        if (brand.empty()) {
+          throw jsi::JSError(rt, "Cannot determine WebGPU object type - no "
+                                 "Symbol.toStringTag found");
+        }
+
+        auto nativeState = obj.getNativeState(rt);
+        auto boxed = std::make_shared<BoxedWebGPUObject>(nativeState, brand);
+        return jsi::Object::createFromHostObject(rt, boxed);
+      });
+  runtime.global().setProperty(runtime, "__webgpuBox", std::move(boxFunc));
+}
+
+RNWebGPUManager::~RNWebGPUManager() {
+  _jsRuntime = nullptr;
+  _jsCallInvoker = nullptr;
+}
+
+} // namespace rnwgpu

package/cpp/rnwgpu/RNWebGPUManager.h

@@ -0,0 +1,46 @@
+#pragma once
+
+#include <memory>
+
+#include "GPU.h"
+#include "PlatformContext.h"
+#include "SurfaceRegistry.h"
+
+namespace facebook {
+namespace jsi {
+class Runtime;
+} // namespace jsi
+namespace react {
+class CallInvoker;
+}
+} // namespace facebook
+
+namespace rnwgpu {
+
+namespace jsi = facebook::jsi;
+namespace react = facebook::react;
+
+class RNWebGPUManager {
+public:
+  RNWebGPUManager(jsi::Runtime *jsRuntime,
+                  std::shared_ptr<facebook::react::CallInvoker> jsCallInvoker,
+                  std::shared_ptr<PlatformContext> platformContext);
+  ~RNWebGPUManager();
+
+  /**
+   * Install global helper functions for Worklets serialization.
+   * This installs __webgpuIsWebGPUObject and __webgpuBox on the global object.
+   * Can be called on any runtime (main JS, UI, or custom worklet runtimes).
+   */
+  static void installWebGPUWorkletHelpers(jsi::Runtime &runtime);
+
+private:
+  jsi::Runtime *_jsRuntime;
+  std::shared_ptr<facebook::react::CallInvoker> _jsCallInvoker;
+
+public:
+  wgpu::Instance _gpu;
+  std::shared_ptr<PlatformContext> _platformContext;
+};
+
+} // namespace rnwgpu