react-native-webgpu 0.5.11

package/android/cpp/AndroidPlatformContext.h

@@ -0,0 +1,207 @@
+#pragma once
+
+#include <android/bitmap.h>
+#include <jni.h>
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "webgpu/webgpu_cpp.h"
+
+#include "PlatformContext.h"
+#include "RNWebGPUManager.h"
+
+namespace rnwgpu {
+
+namespace jsi = facebook::jsi;
+namespace jni = facebook::jni;
+
+class AndroidPlatformContext : public PlatformContext {
+private:
+  jobject _blobModule;
+
+  std::vector<uint8_t> resolveBlob(JNIEnv *env, const std::string &blobId,
+                                   double offset, double size) {
+    if (!_blobModule) {
+      throw std::runtime_error("BlobModule instance is null");
+    }
+
+    jclass blobModuleClass = env->GetObjectClass(_blobModule);
+    if (!blobModuleClass) {
+      throw std::runtime_error("Couldn't find BlobModule class");
+    }
+
+    jmethodID resolveMethod = env->GetMethodID(blobModuleClass, "resolve",
+                                               "(Ljava/lang/String;II)[B");
+    env->DeleteLocalRef(blobModuleClass);
+
+    if (!resolveMethod) {
+      throw std::runtime_error("Couldn't find resolve method in BlobModule");
+    }
+
+    jstring jBlobId = env->NewStringUTF(blobId.c_str());
+    jbyteArray blobData = (jbyteArray)env->CallObjectMethod(
+        _blobModule, resolveMethod, jBlobId, static_cast<jint>(offset),
+        static_cast<jint>(size));
+    env->DeleteLocalRef(jBlobId);
+
+    if (!blobData) {
+      throw std::runtime_error("Couldn't retrieve blob data");
+    }
+
+    jsize len = env->GetArrayLength(blobData);
+    std::vector<uint8_t> data(len);
+    env->GetByteArrayRegion(blobData, 0, len,
+                            reinterpret_cast<jbyte *>(data.data()));
+    env->DeleteLocalRef(blobData);
+    return data;
+  }
+
+public:
+  explicit AndroidPlatformContext(jobject blobModule)
+      : _blobModule(blobModule) {}
+  ~AndroidPlatformContext() {
+    if (_blobModule) {
+      JNIEnv *env = facebook::jni::Environment::current();
+      env->DeleteGlobalRef(_blobModule);
+      _blobModule = nullptr;
+    }
+  }
+
+  wgpu::Surface makeSurface(wgpu::Instance instance, void *window, int width,
+                            int height) override {
+    wgpu::SurfaceSourceAndroidNativeWindow androidSurfaceDesc;
+    androidSurfaceDesc.window = reinterpret_cast<ANativeWindow *>(window);
+    wgpu::SurfaceDescriptor surfaceDescriptor;
+    surfaceDescriptor.nextInChain = &androidSurfaceDesc;
+    return instance.CreateSurface(&surfaceDescriptor);
+  }
+
+  ImageData createImageBitmap(std::string blobId, double offset,
+                              double size) override {
+    jni::Environment::ensureCurrentThreadIsAttached();
+
+    JNIEnv *env = facebook::jni::Environment::current();
+    if (!env) {
+      throw std::runtime_error("Couldn't get JNI environment");
+    }
+
+    auto data = resolveBlob(env, blobId, offset, size);
+    return createImageBitmapFromData(data);
+  }
+
+  void
+  createImageBitmapAsync(std::string blobId, double offset, double size,
+                         std::function<void(ImageData)> onSuccess,
+                         std::function<void(std::string)> onError) override {
+    std::thread([this, blobId = std::move(blobId), offset, size,
+                 onSuccess = std::move(onSuccess),
+                 onError = std::move(onError)]() {
+      jni::Environment::ensureCurrentThreadIsAttached();
+      try {
+        JNIEnv *env = facebook::jni::Environment::current();
+        if (!env) {
+          throw std::runtime_error("Couldn't get JNI environment");
+        }
+        auto data = resolveBlob(env, blobId, offset, size);
+        auto result = createImageBitmapFromData(data);
+        onSuccess(std::move(result));
+      } catch (const std::exception &e) {
+        onError(e.what());
+      }
+    }).detach();
+  }
+
+  ImageData createImageBitmapFromData(std::span<const uint8_t> data) override {
+    jni::Environment::ensureCurrentThreadIsAttached();
+
+    JNIEnv *env = facebook::jni::Environment::current();
+    if (!env) {
+      throw std::runtime_error("Couldn't get JNI environment");
+    }
+
+    // Create jbyteArray from the raw bytes
+    jbyteArray byteArray = env->NewByteArray(static_cast<jsize>(data.size()));
+    if (!byteArray) {
+      throw std::runtime_error("Couldn't allocate byte array");
+    }
+    env->SetByteArrayRegion(byteArray, 0, static_cast<jsize>(data.size()),
+                            reinterpret_cast<const jbyte *>(data.data()));
+
+    // Decode via BitmapFactory
+    jclass bitmapFactoryClass =
+        env->FindClass("android/graphics/BitmapFactory");
+    if (!bitmapFactoryClass) {
+      env->DeleteLocalRef(byteArray);
+      throw std::runtime_error("Couldn't find BitmapFactory class");
+    }
+    jmethodID decodeByteArrayMethod =
+        env->GetStaticMethodID(bitmapFactoryClass, "decodeByteArray",
+                               "([BII)Landroid/graphics/Bitmap;");
+    if (!decodeByteArrayMethod) {
+      env->DeleteLocalRef(byteArray);
+      env->DeleteLocalRef(bitmapFactoryClass);
+      throw std::runtime_error("Couldn't find decodeByteArray method");
+    }
+    jint length = static_cast<jint>(data.size());
+    jobject bitmap = env->CallStaticObjectMethod(
+        bitmapFactoryClass, decodeByteArrayMethod, byteArray, 0, length);
+    env->DeleteLocalRef(bitmapFactoryClass);
+
+    if (!bitmap) {
+      env->DeleteLocalRef(byteArray);
+      throw std::runtime_error("Couldn't decode image");
+    }
+
+    AndroidBitmapInfo bitmapInfo;
+    if (AndroidBitmap_getInfo(env, bitmap, &bitmapInfo) !=
+        ANDROID_BITMAP_RESULT_SUCCESS) {
+      env->DeleteLocalRef(byteArray);
+      env->DeleteLocalRef(bitmap);
+      throw std::runtime_error("Couldn't get bitmap info");
+    }
+
+    void *bitmapPixels;
+    if (AndroidBitmap_lockPixels(env, bitmap, &bitmapPixels) !=
+        ANDROID_BITMAP_RESULT_SUCCESS) {
+      env->DeleteLocalRef(byteArray);
+      env->DeleteLocalRef(bitmap);
+      throw std::runtime_error("Couldn't lock bitmap pixels");
+    }
+
+    ImageData result;
+    result.width = static_cast<int>(bitmapInfo.width);
+    result.height = static_cast<int>(bitmapInfo.height);
+    result.data.resize(bitmapInfo.height * bitmapInfo.stride);
+    memcpy(result.data.data(), bitmapPixels, result.data.size());
+
+    AndroidBitmap_unlockPixels(env, bitmap);
+
+    env->DeleteLocalRef(byteArray);
+    env->DeleteLocalRef(bitmap);
+
+    return result;
+  }
+
+  void createImageBitmapFromDataAsync(
+      std::span<const uint8_t> data, std::function<void(ImageData)> onSuccess,
+      std::function<void(std::string)> onError) override {
+    std::thread([this,
+                 ownedData = std::vector<uint8_t>(data.begin(), data.end()),
+                 onSuccess = std::move(onSuccess),
+                 onError = std::move(onError)]() mutable {
+      jni::Environment::ensureCurrentThreadIsAttached();
+      try {
+        auto result = createImageBitmapFromData(ownedData);
+        onSuccess(std::move(result));
+      } catch (const std::exception &e) {
+        onError(e.what());
+      }
+    }).detach();
+  }
+};
+
+} // namespace rnwgpu

package/android/cpp/cpp-adapter.cpp

@@ -0,0 +1,71 @@
+#include <memory>
+#include <unordered_map>
+
+#include <fbjni/fbjni.h>
+#include <jni.h>
+#include <jsi/jsi.h>
+
+#include <ReactCommon/CallInvokerHolder.h>
+#include <android/native_window_jni.h>
+#include <webgpu/webgpu_cpp.h>
+
+#include "AndroidPlatformContext.h"
+#include "GPUCanvasContext.h"
+#include "RNWebGPUManager.h"
+
+#define LOG_TAG "WebGPUModule"
+
+std::shared_ptr<rnwgpu::RNWebGPUManager> manager;
+
+extern "C" JNIEXPORT void JNICALL Java_com_webgpu_WebGPUModule_initializeNative(
+    JNIEnv *env, jobject /* this */, jlong jsRuntime,
+    jobject jsCallInvokerHolder, jobject blobModule) {
+  auto runtime = reinterpret_cast<facebook::jsi::Runtime *>(jsRuntime);
+  jobject globalBlobModule = env->NewGlobalRef(blobModule);
+  auto jsCallInvoker{
+      facebook::jni::alias_ref<facebook::react::CallInvokerHolder::javaobject>{
+          reinterpret_cast<facebook::react::CallInvokerHolder::javaobject>(
+              jsCallInvokerHolder)} -> cthis()
+          ->getCallInvoker()};
+  auto platformContext =
+      std::make_shared<rnwgpu::AndroidPlatformContext>(globalBlobModule);
+  manager = std::make_shared<rnwgpu::RNWebGPUManager>(runtime, jsCallInvoker,
+                                                      platformContext);
+}
+
+extern "C" JNIEXPORT void JNICALL Java_com_webgpu_WebGPUView_onSurfaceChanged(
+    JNIEnv *env, jobject thiz, jobject surface, jint contextId, jfloat width,
+    jfloat height) {
+  auto &registry = rnwgpu::SurfaceRegistry::getInstance();
+  registry.getSurfaceInfo(contextId)->resize(static_cast<int>(width),
+                                             static_cast<int>(height));
+}
+
+extern "C" JNIEXPORT void JNICALL Java_com_webgpu_WebGPUView_onSurfaceCreate(
+    JNIEnv *env, jobject thiz, jobject jSurface, jint contextId, jfloat width,
+    jfloat height) {
+  auto window = ANativeWindow_fromSurface(env, jSurface);
+  // ANativeWindow_acquire(window);
+  auto &registry = rnwgpu::SurfaceRegistry::getInstance();
+  auto gpu = manager->_gpu;
+  auto surface = manager->_platformContext->makeSurface(
+      gpu, window, static_cast<int>(width), static_cast<int>(height));
+  registry
+      .getSurfaceInfoOrCreate(contextId, gpu, static_cast<int>(width),
+                              static_cast<int>(height))
+      ->switchToOnscreen(window, surface);
+}
+
+extern "C" JNIEXPORT void JNICALL
+Java_com_webgpu_WebGPUView_switchToOffscreenSurface(JNIEnv *env, jobject thiz,
+                                                    jint contextId) {
+  auto &registry = rnwgpu::SurfaceRegistry::getInstance();
+  auto nativeSurface = registry.getSurfaceInfo(contextId)->switchToOffscreen();
+  // ANativeWindow_release(reinterpret_cast<ANativeWindow *>(nativeSurface));
+}
+
+extern "C" JNIEXPORT void JNICALL Java_com_webgpu_WebGPUView_onSurfaceDestroy(
+    JNIEnv *env, jobject thiz, jint contextId) {
+  auto &registry = rnwgpu::SurfaceRegistry::getInstance();
+  registry.removeSurfaceInfo(contextId);
+}

package/android/src/main/AndroidManifest.xml

@@ -0,0 +1,3 @@
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+          package="com.webgpu">
+</manifest>

package/android/src/main/java/com/webgpu/WebGPUAHBView.java

@@ -0,0 +1,298 @@
+package com.webgpu;
+
+import android.annotation.SuppressLint;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.graphics.Canvas;
+import android.graphics.Paint;
+import android.graphics.PixelFormat;
+import android.graphics.Rect;
+import android.hardware.HardwareBuffer;
+import android.media.Image;
+import android.media.ImageReader;
+import android.os.Build;
+import android.os.Handler;
+import android.os.Looper;
+import android.view.Surface;
+import android.view.View;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.RequiresApi;
+
+import java.nio.ByteBuffer;
+
+@RequiresApi(api = Build.VERSION_CODES.Q)
+@SuppressLint("ViewConstructor")
+public class WebGPUAHBView extends View implements ImageReader.OnImageAvailableListener {
+
+  private static final int MAX_IMAGES = 3;
+  private static final String TAG = "WebGPUAHBView";
+
+  private final WebGPUAPI mApi;
+  private ImageReader mImageReader;
+  private Surface mSurface;
+  private Image mCurrentImage;
+  private Bitmap mCachedBitmap;
+  private final Paint mPaint;
+  private final Handler mHandler;
+  private final Object mImageLock = new Object();
+  private boolean mSurfaceCreated = false;
+  private int mConfiguredWidth = 0;
+  private int mConfiguredHeight = 0;
+
+  public WebGPUAHBView(Context context, WebGPUAPI api) {
+    super(context);
+    mApi = api;
+    mPaint = new Paint(Paint.FILTER_BITMAP_FLAG);
+    mHandler = new Handler(Looper.getMainLooper());
+
+    // Enable hardware acceleration for this view
+    setLayerType(LAYER_TYPE_HARDWARE, null);
+
+    // Make sure we get drawn
+    setWillNotDraw(false);
+  }
+
+  @Override
+  protected void onSizeChanged(int w, int h, int oldw, int oldh) {
+    super.onSizeChanged(w, h, oldw, oldh);
+
+    if (w > 0 && h > 0) {
+      // Recreate ImageReader with new dimensions
+      setupImageReader(w, h);
+    }
+  }
+
+  private void setupImageReader(int width, int height) {
+    // Don't recreate if dimensions haven't changed
+    if (width == mConfiguredWidth && height == mConfiguredHeight && mImageReader != null) {
+      return;
+    }
+
+    // Clean up previous ImageReader
+    cleanupImageReader();
+
+    try {
+      // Create ImageReader with HardwareBuffer support
+      mImageReader = ImageReader.newInstance(
+        width,
+        height,
+        PixelFormat.RGBA_8888,
+        MAX_IMAGES,
+        HardwareBuffer.USAGE_GPU_SAMPLED_IMAGE |
+          HardwareBuffer.USAGE_GPU_COLOR_OUTPUT |
+          HardwareBuffer.USAGE_COMPOSER_OVERLAY |
+          HardwareBuffer.USAGE_CPU_READ_RARELY // For fallback to Bitmap if needed
+      );
+
+      mImageReader.setOnImageAvailableListener(this, mHandler);
+
+      // Get the Surface for WebGPU to render to
+      mSurface = mImageReader.getSurface();
+
+      mConfiguredWidth = width;
+      mConfiguredHeight = height;
+
+      // Notify WebGPU about the new surface
+      if (!mSurfaceCreated) {
+        mApi.surfaceCreated(mSurface);
+        mSurfaceCreated = true;
+      } else {
+        mApi.surfaceChanged(mSurface);
+      }
+
+    } catch (Exception e) {
+      e.printStackTrace();
+      // Fallback to offscreen if ImageReader creation fails
+      mApi.surfaceOffscreen();
+    }
+  }
+
+  @Override
+  public void onImageAvailable(ImageReader reader) {
+    synchronized (mImageLock) {
+      // Close previous image if exists
+      if (mCurrentImage != null) {
+        mCurrentImage.close();
+        mCurrentImage = null;
+      }
+
+      try {
+        // Get the latest image
+        mCurrentImage = reader.acquireLatestImage();
+
+        if (mCurrentImage != null) {
+          // Request a redraw on the UI thread
+          postInvalidateOnAnimation();
+        }
+      } catch (Exception e) {
+        e.printStackTrace();
+      }
+    }
+  }
+
+  @Override
+  protected void onDraw(Canvas canvas) {
+    super.onDraw(canvas);
+
+    synchronized (mImageLock) {
+      if (mCurrentImage == null) {
+        return;
+      }
+
+      try {
+        // Try to use HardwareBuffer directly (most efficient path)
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
+          HardwareBuffer hardwareBuffer = mCurrentImage.getHardwareBuffer();
+
+          if (hardwareBuffer != null) {
+            // Draw using HardwareBuffer
+            drawHardwareBuffer(canvas, hardwareBuffer);
+            return;
+          }
+        }
+
+        // Fallback: Convert Image to Bitmap
+        drawImageAsBitmap(canvas, mCurrentImage);
+
+      } catch (Exception e) {
+        e.printStackTrace();
+      }
+    }
+  }
+
+  private void drawHardwareBuffer(Canvas canvas, HardwareBuffer hardwareBuffer) {
+    // On Android Q+, we can create a Bitmap from HardwareBuffer
+    try {
+      // Create or reuse bitmap with matching dimensions
+      if (mCachedBitmap == null ||
+        mCachedBitmap.getWidth() != hardwareBuffer.getWidth() ||
+        mCachedBitmap.getHeight() != hardwareBuffer.getHeight()) {
+
+        if (mCachedBitmap != null) {
+          mCachedBitmap.recycle();
+        }
+
+        // Create a hardware-backed Bitmap from the HardwareBuffer
+        mCachedBitmap = Bitmap.wrapHardwareBuffer(hardwareBuffer, null);
+      } else {
+        // Reuse existing bitmap and update with new HardwareBuffer content
+        mCachedBitmap = Bitmap.wrapHardwareBuffer(hardwareBuffer, null);
+      }
+
+      if (mCachedBitmap != null) {
+        // Draw the bitmap to canvas
+        canvas.drawBitmap(mCachedBitmap, 0, 0, mPaint);
+      }
+    } catch (Exception e) {
+      e.printStackTrace();
+      // Fallback to Image-based rendering
+      drawImageAsBitmap(canvas, mCurrentImage);
+    }
+  }
+
+  private void drawImageAsBitmap(Canvas canvas, Image image) {
+    // Fallback method: manually convert Image to Bitmap
+    if (image.getFormat() != PixelFormat.RGBA_8888) {
+      return;
+    }
+
+    Image.Plane[] planes = image.getPlanes();
+    if (planes.length == 0) {
+      return;
+    }
+
+    ByteBuffer buffer = planes[0].getBuffer();
+    int pixelStride = planes[0].getPixelStride();
+    int rowStride = planes[0].getRowStride();
+    int rowPadding = rowStride - pixelStride * image.getWidth();
+
+    // Create or reuse bitmap
+    int bitmapWidth = image.getWidth() + rowPadding / pixelStride;
+    if (mCachedBitmap == null ||
+      mCachedBitmap.getWidth() != bitmapWidth ||
+      mCachedBitmap.getHeight() != image.getHeight()) {
+
+      if (mCachedBitmap != null) {
+        mCachedBitmap.recycle();
+      }
+      mCachedBitmap = Bitmap.createBitmap(
+        bitmapWidth,
+        image.getHeight(),
+        Bitmap.Config.ARGB_8888
+      );
+    }
+
+    mCachedBitmap.copyPixelsFromBuffer(buffer);
+
+    // Draw only the valid portion (without padding)
+    Rect src = new Rect(0, 0, image.getWidth(), image.getHeight());
+    Rect dst = new Rect(0, 0, getWidth(), getHeight());
+    canvas.drawBitmap(mCachedBitmap, src, dst, mPaint);
+  }
+
+  private void cleanupImageReader() {
+    synchronized (mImageLock) {
+      if (mCurrentImage != null) {
+        mCurrentImage.close();
+        mCurrentImage = null;
+      }
+
+      if (mImageReader != null) {
+        mImageReader.close();
+        mImageReader = null;
+      }
+
+      if (mSurface != null) {
+        mSurface.release();
+        mSurface = null;
+      }
+    }
+  }
+
+  @Override
+  protected void onDetachedFromWindow() {
+    super.onDetachedFromWindow();
+
+    // Notify WebGPU that surface is being destroyed
+    if (mSurfaceCreated) {
+      mApi.surfaceDestroyed();
+      mSurfaceCreated = false;
+    }
+
+    // Clean up resources
+    cleanupImageReader();
+
+    if (mCachedBitmap != null) {
+      mCachedBitmap.recycle();
+      mCachedBitmap = null;
+    }
+  }
+
+  @Override
+  protected void onAttachedToWindow() {
+    super.onAttachedToWindow();
+
+    // Re-setup if we have valid dimensions
+    if (getWidth() > 0 && getHeight() > 0) {
+      setupImageReader(getWidth(), getHeight());
+    }
+  }
+
+  @Override
+  public void setAlpha(float alpha) {
+    super.setAlpha(alpha);
+    mPaint.setAlpha((int) (alpha * 255));
+    invalidate();
+  }
+
+  @Override
+  protected void onVisibilityChanged(@NonNull View changedView, int visibility) {
+    super.onVisibilityChanged(changedView, visibility);
+
+    if (visibility == VISIBLE && mSurface == null && getWidth() > 0 && getHeight() > 0) {
+      // Re-create surface if needed when becoming visible
+      setupImageReader(getWidth(), getHeight());
+    }
+  }
+}

package/android/src/main/java/com/webgpu/WebGPUAPI.java

@@ -0,0 +1,20 @@
+package com.webgpu;
+
+import android.view.Surface;
+
+import com.facebook.proguard.annotations.DoNotStrip;
+
+public interface WebGPUAPI {
+
+  void surfaceCreated(
+    Surface surface
+  );
+
+  void surfaceChanged(
+    Surface surface
+  );
+
+  void surfaceDestroyed();
+
+  void surfaceOffscreen();
+}

package/android/src/main/java/com/webgpu/WebGPUBaseView.java

@@ -0,0 +1,61 @@
+package com.webgpu;
+
+import android.content.Context;
+import android.view.Surface;
+import android.view.View;
+import com.facebook.proguard.annotations.DoNotStrip;
+
+public abstract class WebGPUBaseView extends View {
+  protected Integer mContextId;
+
+  public WebGPUBaseView(Context context, Integer contextId) {
+    super(context);
+    mContextId = contextId;
+  }
+
+  protected void handleSurfaceCreate(Surface surface) {
+    float density = getResources().getDisplayMetrics().density;
+    float scaledWidth = getWidth() / density;
+    float scaledHeight = getHeight() / density;
+    onSurfaceCreate(surface, mContextId, scaledWidth, scaledHeight);
+  }
+
+  protected void handleSurfaceChanged(Surface surface) {
+    float density = getResources().getDisplayMetrics().density;
+    float scaledWidth = getWidth() / density;
+    float scaledHeight = getHeight() / density;
+    onSurfaceChanged(surface, mContextId, scaledWidth, scaledHeight);
+  }
+
+  protected void handleSurfaceDestroyed() {
+    onSurfaceDestroy(mContextId);
+  }
+
+  @Override
+  protected void onLayout(boolean changed, int left, int top, int right, int bottom) {
+    super.onLayout(changed, left, top, right, bottom);
+  }
+
+
+  @DoNotStrip
+  private native void onSurfaceCreate(
+    Surface surface,
+    int contextId,
+    float width,
+    float height
+  );
+
+  @DoNotStrip
+  private native void onSurfaceChanged(
+    Surface surface,
+    int contextId,
+    float width,
+    float height
+  );
+
+  @DoNotStrip
+  private native void onSurfaceDestroy(int contextId);
+
+  @DoNotStrip
+  protected native void switchToOffscreenSurface(int contextId);
+}