Implement Gouraud-based shading using WebGPU.

The WebGPU feature hasn't been released yet but it's interesting to see how
we can use it in order to speed up the rendering of some objects.
This patch allows to render mesh patterns using WebGPU.

I didn't see any significant performance improvement on my machine (mac M2)
but it may be different on other platforms.

Author: calixteman

src/core/evaluator.js

@@ -105,6 +105,7 @@ const DefaultPartialEvaluatorOptions = Object.freeze({
   iccUrl: null,
   standardFontDataUrl: null,
   wasmUrl: null,
+  prepareWebGPU: null,
 });
 
 const PatternType = {
@@ -1513,7 +1514,8 @@ class PartialEvaluator {
         resources,
         this._pdfFunctionFactory,
         this.globalColorSpaceCache,
-        localColorSpaceCache
+        localColorSpaceCache,
+        this.options.prepareWebGPU
       );
       patternIR = shadingFill.getIR();
     } catch (reason) {

src/core/pattern.js

@@ -55,7 +55,8 @@ class Pattern {
     res,
     pdfFunctionFactory,
     globalColorSpaceCache,
-    localColorSpaceCache
+    localColorSpaceCache,
+    prepareWebGPU = null
   ) {
     const dict = shading instanceof BaseStream ? shading.dict : shading;
     const type = dict.get("ShadingType");
@@ -76,6 +77,7 @@ class Pattern {
         case ShadingType.LATTICE_FORM_MESH:
         case ShadingType.COONS_PATCH_MESH:
         case ShadingType.TENSOR_PATCH_MESH:
+          prepareWebGPU?.();
           return new MeshShading(
             shading,
             xref,
@@ -934,7 +936,7 @@ class MeshShading extends BaseShading {
   }
 
   _packData() {
-    let i, ii, j, jj;
+    let i, ii, j;
 
     const coords = this.coords;
     const coordsPacked = new Float32Array(coords.length * 2);
@@ -945,25 +947,27 @@ class MeshShading extends BaseShading {
     }
     this.coords = coordsPacked;
 
+    // Stride 4 (RGBA layout, alpha unused) so the buffer maps directly to
+    // array<u32> in the WebGPU vertex shader without any repacking.
     const colors = this.colors;
-    const colorsPacked = new Uint8Array(colors.length * 3);
+    const colorsPacked = new Uint8Array(colors.length * 4);
     for (i = 0, j = 0, ii = colors.length; i < ii; i++) {
       const c = colors[i];
       colorsPacked[j++] = c[0];
       colorsPacked[j++] = c[1];
       colorsPacked[j++] = c[2];
+      j++; // alpha — unused, stays 0
     }
     this.colors = colorsPacked;
 
+    // Store raw vertex indices (not byte offsets) so the GPU shader can
+    // address coords / colors without knowing their strides, and so the
+    // arrays are transferable Uint32Arrays.
     const figures = this.figures;
     for (i = 0, ii = figures.length; i < ii; i++) {
-      const figure = figures[i],
-        ps = figure.coords,
-        cs = figure.colors;
-      for (j = 0, jj = ps.length; j < jj; j++) {
-        ps[j] *= 2;
-        cs[j] *= 3;
-      }
+      const figure = figures[i];
+      figure.coords = new Uint32Array(figure.coords);
+      figure.colors = new Uint32Array(figure.colors);
     }
   }
 

src/core/pdf_manager.js

@@ -71,6 +71,20 @@ class BasePdfManager {
       FeatureTest.isOffscreenCanvasSupported;
     evaluatorOptions.isImageDecoderSupported &&=
       FeatureTest.isImageDecoderSupported;
+
+    // Set up a one-shot callback so evaluators can notify the main thread that
+    // WebGPU-acceleratable content was found. The flag ensures the message is
+    // sent at most once per document.
+    if (evaluatorOptions.enableWebGPU) {
+      let prepareWebGPUSent = false;
+      evaluatorOptions.prepareWebGPU = () => {
+        if (!prepareWebGPUSent) {
+          prepareWebGPUSent = true;
+          handler.send("PrepareWebGPU", null);
+        }
+      };
+    }
+    delete evaluatorOptions.enableWebGPU;
     this.evaluatorOptions = Object.freeze(evaluatorOptions);
 
     // Initialize image-options once per document.

src/display/api.js

@@ -79,6 +79,7 @@ import { DOMFilterFactory } from "./filter_factory.js";
 import { DOMStandardFontDataFactory } from "display-standard_fontdata_factory";
 import { DOMWasmFactory } from "display-wasm_factory";
 import { GlobalWorkerOptions } from "./worker_options.js";
+import { initWebGPUMesh } from "./webgpu_mesh.js";
 import { Metadata } from "./metadata.js";
 import { OptionalContentConfig } from "./optional_content_config.js";
 import { PagesMapper } from "./pages_mapper.js";
@@ -347,6 +348,7 @@ function getDocument(src = {}) {
       ? NodeFilterFactory
       : DOMFilterFactory);
   const enableHWA = src.enableHWA === true;
+  const enableWebGPU = src.enableWebGPU === true;
   const useWasm = src.useWasm !== false;
   const pagesMapper = src.pagesMapper || new PagesMapper();
 
@@ -440,6 +442,7 @@ function getDocument(src = {}) {
       iccUrl,
       standardFontDataUrl,
       wasmUrl,
+      enableWebGPU,
     },
   };
   const transportParams = {
@@ -2926,6 +2929,13 @@ class WorkerTransport {
       this.#onProgress(data);
     });
 
+    messageHandler.on("PrepareWebGPU", () => {
+      if (this.destroyed) {
+        return;
+      }
+      initWebGPUMesh();
+    });
+
     if (typeof PDFJSDev === "undefined" || !PDFJSDev.test("MOZCENTRAL")) {
       messageHandler.on("FetchBinaryData", async data => {
         if (this.destroyed) {

src/display/pattern_helper.js

@@ -13,6 +13,7 @@
  * limitations under the License.
  */
 
+import { drawMeshWithGPU, isWebGPUMeshReady } from "./webgpu_mesh.js";
 import {
   FormatError,
   info,
@@ -282,48 +283,48 @@ function drawTriangle(data, context, p1, p2, p3, c1, c2, c3) {
   const bytes = data.data,
     rowSize = data.width * 4;
   let tmp;
-  if (coords[p1 + 1] > coords[p2 + 1]) {
+  if (coords[p1 * 2 + 1] > coords[p2 * 2 + 1]) {
     tmp = p1;
     p1 = p2;
     p2 = tmp;
     tmp = c1;
     c1 = c2;
     c2 = tmp;
   }
-  if (coords[p2 + 1] > coords[p3 + 1]) {
+  if (coords[p2 * 2 + 1] > coords[p3 * 2 + 1]) {
     tmp = p2;
     p2 = p3;
     p3 = tmp;
     tmp = c2;
     c2 = c3;
     c3 = tmp;
   }
-  if (coords[p1 + 1] > coords[p2 + 1]) {
+  if (coords[p1 * 2 + 1] > coords[p2 * 2 + 1]) {
     tmp = p1;
     p1 = p2;
     p2 = tmp;
     tmp = c1;
     c1 = c2;
     c2 = tmp;
   }
-  const x1 = (coords[p1] + context.offsetX) * context.scaleX;
-  const y1 = (coords[p1 + 1] + context.offsetY) * context.scaleY;
-  const x2 = (coords[p2] + context.offsetX) * context.scaleX;
-  const y2 = (coords[p2 + 1] + context.offsetY) * context.scaleY;
-  const x3 = (coords[p3] + context.offsetX) * context.scaleX;
-  const y3 = (coords[p3 + 1] + context.offsetY) * context.scaleY;
+  const x1 = (coords[p1 * 2] + context.offsetX) * context.scaleX;
+  const y1 = (coords[p1 * 2 + 1] + context.offsetY) * context.scaleY;
+  const x2 = (coords[p2 * 2] + context.offsetX) * context.scaleX;
+  const y2 = (coords[p2 * 2 + 1] + context.offsetY) * context.scaleY;
+  const x3 = (coords[p3 * 2] + context.offsetX) * context.scaleX;
+  const y3 = (coords[p3 * 2 + 1] + context.offsetY) * context.scaleY;
   if (y1 >= y3) {
     return;
   }
-  const c1r = colors[c1],
-    c1g = colors[c1 + 1],
-    c1b = colors[c1 + 2];
-  const c2r = colors[c2],
-    c2g = colors[c2 + 1],
-    c2b = colors[c2 + 2];
-  const c3r = colors[c3],
-    c3g = colors[c3 + 1],
-    c3b = colors[c3 + 2];
+  const c1r = colors[c1 * 4],
+    c1g = colors[c1 * 4 + 1],
+    c1b = colors[c1 * 4 + 2];
+  const c2r = colors[c2 * 4],
+    c2g = colors[c2 * 4 + 1],
+    c2b = colors[c2 * 4 + 2];
+  const c3r = colors[c3 * 4],
+    c3g = colors[c3 * 4 + 1],
+    c3b = colors[c3 * 4 + 2];
 
   const minY = Math.round(y1),
     maxY = Math.round(y3);
@@ -494,26 +495,39 @@ class MeshShadingPattern extends BaseShadingPattern {
       paddedWidth,
       paddedHeight
     );
-    const tmpCtx = tmpCanvas.context;
 
-    const data = tmpCtx.createImageData(width, height);
-    if (backgroundColor) {
-      const bytes = data.data;
-      for (let i = 0, ii = bytes.length; i < ii; i += 4) {
-        bytes[i] = backgroundColor[0];
-        bytes[i + 1] = backgroundColor[1];
-        bytes[i + 2] = backgroundColor[2];
-        bytes[i + 3] = 255;
+    if (isWebGPUMeshReady()) {
+      tmpCanvas.context.drawImage(
+        drawMeshWithGPU(
+          this._figures,
+          context,
+          backgroundColor,
+          paddedWidth,
+          paddedHeight,
+          BORDER_SIZE
+        ),
+        0,
+        0
+      );
+    } else {
+      const data = tmpCanvas.context.createImageData(width, height);
+      if (backgroundColor) {
+        const bytes = data.data;
+        for (let i = 0, ii = bytes.length; i < ii; i += 4) {
+          bytes[i] = backgroundColor[0];
+          bytes[i + 1] = backgroundColor[1];
+          bytes[i + 2] = backgroundColor[2];
+          bytes[i + 3] = 255;
+        }
       }
+      for (const figure of this._figures) {
+        drawFigure(data, figure, context);
+      }
+      tmpCanvas.context.putImageData(data, BORDER_SIZE, BORDER_SIZE);
     }
-    for (const figure of this._figures) {
-      drawFigure(data, figure, context);
-    }
-    tmpCtx.putImageData(data, BORDER_SIZE, BORDER_SIZE);
-    const canvas = tmpCanvas.canvas;
 
     return {
-      canvas,
+      canvas: tmpCanvas.canvas,
       offsetX: offsetX - BORDER_SIZE * scaleX,
       offsetY: offsetY - BORDER_SIZE * scaleY,
       scaleX,