Workaround for WebGL crash for datasets with many segmentation layers (#6995)

* reduce shader complexity when having multiple segmentation layers (shader needs to be compiled more often then, though)

* fix tests

* clean up

* update changelog

Author: philippotto

CHANGELOG.unreleased.md

@@ -19,6 +19,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 - Fixed that changing a segment color could lead to a crash. [#7000](https://github.com/scalableminds/webknossos/pull/7000)
 - Fixed rendering issues on some affected systems that led to "black holes". [#7018](https://github.com/scalableminds/webknossos/pull/7018)
 - Fixed a bug that made downloads of public annotations fail occasionally. [#7025](https://github.com/scalableminds/webknossos/pull/7025)
+- Added a workaround for a WebGL crash which could appear when a dataset contained many segmentation layers. [#6995](https://github.com/scalableminds/webknossos/pull/6995)
 
 ### Removed
 

frontend/javascripts/oxalis/model/bucket_data_handling/data_rendering_logic.tsx

@@ -200,6 +200,7 @@ export function calculateTextureSizeAndCountForLayer(
 function buildTextureInformationMap<
   Layer extends {
     elementClass: ElementClass;
+    category: "color" | "segmentation";
   },
 >(
   layers: Array<Layer>,
@@ -238,6 +239,7 @@ function getSmallestCommonBucketCapacity<
 function getRenderSupportedLayerCount<
   Layer extends {
     elementClass: ElementClass;
+    category: "color" | "segmentation";
   },
 >(
   specs: GpuSpecs,
@@ -262,15 +264,23 @@ function getRenderSupportedLayerCount<
     (specs.maxTextureCount - textureCountForSegmentation - lookupTextureCount) /
       maximumTextureCountForLayer,
   );
+
+  // Without any GPU restrictions, WK would be able to render all color layers
+  // plus one segmentation layer. Use that as the upper layer count limit to avoid
+  // compiling too complex shaders.
+  const maximumLayerCount =
+    Array.from(textureInformationPerLayer.keys()).filter((l) => l.category === "color").length +
+    (hasSegmentation ? 1 : 0);
   return {
-    maximumLayerCountToRender,
+    maximumLayerCountToRender: Math.min(maximumLayerCountToRender, maximumLayerCount),
     maximumTextureCountForLayer,
   };
 }
 
 export function computeDataTexturesSetup<
   Layer extends {
     elementClass: ElementClass;
+    category: "color" | "segmentation";
   },
 >(
   specs: GpuSpecs,

frontend/javascripts/test/model/binary/data_rendering_logic.spec.ts

@@ -1,3 +1,4 @@
+import _ from "lodash";
 import {
   calculateTextureSizeAndCountForLayer,
   computeDataTexturesSetup,
@@ -25,6 +26,34 @@ const grayscaleElementClass = "uint8";
 const volumeByteCount = 4;
 const volumeElementClass = "uint32";
 
+/*
+ * The current rendering logic in WK only allows
+ * as many layers as necessary. This is done to avoid
+ * that the shaders are compiled for N layers even though
+ * N layers will never be rendered at the same time (because
+ * only one segmentation layer can be rendered at a time).
+ * For that reason, testing the specs has to be done with a
+ * sufficiently large amount of layers. To achieve this,
+ * the helper function createLayers is used.
+ */
+
+const createGrayscaleLayer = () => ({
+  byteCount: grayscaleByteCount,
+  elementClass: grayscaleElementClass,
+  category: "color",
+});
+const createVolumeLayer = () => ({
+  byteCount: volumeByteCount,
+  elementClass: volumeElementClass,
+  category: "segmentation",
+});
+
+function createLayers(grayscaleCount: number, volumeCount: number) {
+  const grayscaleLayers = _.range(0, grayscaleCount).map(() => createGrayscaleLayer());
+  const volumeLayers = _.range(0, volumeCount).map(() => createVolumeLayer());
+  return grayscaleLayers.concat(volumeLayers);
+}
+
 test("calculateTextureSizeAndCountForLayer: grayscale data + minSpecs", (t) => {
   const { textureSize, textureCount } = calculateTextureSizeAndCountForLayer(
     minSpecs,
@@ -90,24 +119,8 @@ test("calculateTextureSizeAndCountForLayer: color data + betterSpecs", (t) => {
   t.is(textureSize, midSpecs.supportedTextureSize);
   t.is(textureCount, 1);
 });
-const grayscaleLayer1 = {
-  byteCount: grayscaleByteCount,
-  elementClass: grayscaleElementClass,
-};
-const grayscaleLayer2 = {
-  byteCount: grayscaleByteCount,
-  elementClass: grayscaleElementClass,
-};
-const grayscaleLayer3 = {
-  byteCount: grayscaleByteCount,
-  elementClass: grayscaleElementClass,
-};
-const volumeLayer1 = {
-  byteCount: volumeByteCount,
-  elementClass: volumeElementClass,
-};
 
-type Layer = typeof grayscaleLayer1;
+type Layer = ReturnType<typeof createGrayscaleLayer>;
 
 const getByteCount = (layer: Layer) => layer.byteCount;
 
@@ -123,7 +136,7 @@ function computeDataTexturesSetupCurried(spec: typeof minSpecs, hasSegmentation:
   return (layers: Layer[]) =>
     computeDataTexturesSetup(
       spec,
-      layers as { elementClass: ElementClass }[],
+      layers as { elementClass: ElementClass; category: "color" | "segmentation" }[],
       getByteCount as any,
       hasSegmentation,
       DEFAULT_REQUIRED_BUCKET_CAPACITY,
@@ -137,59 +150,37 @@ test("Basic support (no segmentation): all specs", (t) => {
     [midSpecs, 15],
     [betterSpecs, 31],
   ];
+  const hundredGrayscaleLayers = createLayers(100, 0);
   for (const [spec, expectedLayerCount] of specs) {
     const computeDataTexturesSetupPartial = computeDataTexturesSetupCurried(spec, false);
-    testSupportFlags(t, computeDataTexturesSetupPartial([grayscaleLayer1]), expectedLayerCount);
     testSupportFlags(
       t,
-      computeDataTexturesSetupPartial([grayscaleLayer1, grayscaleLayer2]),
+      computeDataTexturesSetupPartial(hundredGrayscaleLayers),
       expectedLayerCount,
     );
     testSupportFlags(
       t,
-      computeDataTexturesSetupPartial([grayscaleLayer1, grayscaleLayer2, grayscaleLayer3]),
+      computeDataTexturesSetupPartial(hundredGrayscaleLayers),
+      expectedLayerCount,
+    );
+    testSupportFlags(
+      t,
+      computeDataTexturesSetupPartial(hundredGrayscaleLayers),
       expectedLayerCount,
     );
   }
 });
 
 test("Basic support + volume: min specs", (t) => {
   const computeDataTexturesSetupPartial = computeDataTexturesSetupCurried(minSpecs, true);
-  testSupportFlags(t, computeDataTexturesSetupPartial([grayscaleLayer1, grayscaleLayer2]), 4);
-  testSupportFlags(t, computeDataTexturesSetupPartial([grayscaleLayer1, volumeLayer1]), 1);
-  testSupportFlags(
-    t,
-    computeDataTexturesSetupPartial([grayscaleLayer1, grayscaleLayer2, volumeLayer1]),
-    1,
-  );
-  testSupportFlags(
-    t,
-    computeDataTexturesSetupPartial([
-      grayscaleLayer1,
-      grayscaleLayer2,
-      grayscaleLayer3,
-      volumeLayer1,
-    ]),
-    1,
-  );
+  testSupportFlags(t, computeDataTexturesSetupPartial(createLayers(10, 0)), 4);
+  testSupportFlags(t, computeDataTexturesSetupPartial(createLayers(10, 1)), 1);
+  testSupportFlags(t, computeDataTexturesSetupPartial(createLayers(10, 1)), 1);
+  testSupportFlags(t, computeDataTexturesSetupPartial(createLayers(10, 1)), 1);
 });
 
 test("Basic support + volume: mid specs", (t) => {
   const computeDataTexturesSetupPartial = computeDataTexturesSetupCurried(midSpecs, true);
-  testSupportFlags(t, computeDataTexturesSetupPartial([grayscaleLayer1, volumeLayer1]), 12);
-  testSupportFlags(
-    t,
-    computeDataTexturesSetupPartial([grayscaleLayer1, grayscaleLayer2, volumeLayer1]),
-    12,
-  );
-  testSupportFlags(
-    t,
-    computeDataTexturesSetupPartial([
-      grayscaleLayer1,
-      grayscaleLayer2,
-      grayscaleLayer3,
-      volumeLayer1,
-    ]),
-    12,
-  );
+  testSupportFlags(t, computeDataTexturesSetupPartial(createLayers(20, 1)), 12);
+  testSupportFlags(t, computeDataTexturesSetupPartial(createLayers(5, 1)), 6);
 });