feat: eagerly upload chunk textures to the GPU (#698)

Summary:

this pr changes when chunk textures are sent to the gpu. instead of uploading a
chunk lazily the first time it is drawn, each chunk's texture is now uploaded as
soon as its data finishes loading and that gpu texture is owned centrally for
the chunk's whole lifetime rather than by individual layers. once a chunk is on
the gpu its cpu copy is released to save memory, and its texture is freed when
the chunk is no longer needed. the image, label, and volume layers now share
these textures instead of each building and tracking their own.

there's a minor temporary regression: because a chunk's texture is freed the
moment it is no longer needed scrubbing through time can briefly show gaps
where the previous frame used to stay on screen while the next one loaded.
that smoothness returns once eviction is deferred under a budget, keeping
recently used textures resident. this does not affect volume rendering.

Author: shlomnissan

src/core/renderable_object.ts

@@ -20,13 +20,19 @@ export abstract class RenderableObject extends Node {
   private cullFaceMode_: CullingMode = "none";
 
   public setTexture(index: number, texture: Texture) {
-    const oldTexture = this.textures_[index];
-    if (oldTexture !== undefined) {
-      this.staleTextures_.push(oldTexture);
-    }
     this.textures_[index] = texture;
   }
 
+  protected clearTextures() {
+    this.textures_.length = 0;
+  }
+
+  protected markStaleTexture(texture: Texture | undefined) {
+    if (texture !== undefined) {
+      this.staleTextures_.push(texture);
+    }
+  }
+
   public popStaleTextures() {
     const stale = this.staleTextures_;
     this.staleTextures_ = [];

src/data/chunk.ts

@@ -1,4 +1,7 @@
-import { TextureUnpackRowAlignment } from "../objects/textures/texture";
+import {
+  type Texture,
+  TextureUnpackRowAlignment,
+} from "../objects/textures/texture";
 import { Logger } from "../utilities/logger";
 
 const chunkDataTypes = [
@@ -34,6 +37,7 @@ export type ChunkViewState = {
 
 export type Chunk = {
   data?: ChunkData;
+  texture?: Texture;
   state: "unloaded" | "queued" | "loading" | "loaded";
   lod: number;
   shape: {

src/data/chunk_manager.ts

@@ -1,19 +1,34 @@
-import { ChunkSource } from "./chunk";
-import { ChunkQueue } from "./chunk_queue";
-import { chunkMemoryStats } from "./chunk_memory";
+import { Chunk, ChunkSource } from "./chunk";
+import { ChunkQueue, comparePriority } from "./chunk_queue";
+import { chunkMemoryStats, clearChunkData } from "./chunk_memory";
+import { ChunkStore } from "./chunk_store";
+import { ChunkStoreView } from "./chunk_store_view";
+import { ImageSourcePolicy } from "../core/image_source_policy";
+import { Texture } from "../objects/textures/texture";
+import { Texture3D } from "../objects/textures/texture_3d";
 
 export type QueueStats = {
   pending: number;
   running: number;
 };
-import { ChunkStore } from "./chunk_store";
-import { ChunkStoreView } from "./chunk_store_view";
-import { ImageSourcePolicy } from "../core/image_source_policy";
+
+const MAX_UPLOADS_PER_STORE_PER_UPDATE = 4;
 
 export class ChunkManager {
   private readonly stores_: { source: ChunkSource; store: ChunkStore }[] = [];
   private readonly queue_ = new ChunkQueue();
 
+  private readonly uploadTexture_?: (texture: Texture) => void;
+  private readonly disposeTexture_?: (texture: Texture) => void;
+
+  constructor(
+    uploadTexture?: (texture: Texture) => void,
+    disposeTexture?: (texture: Texture) => void
+  ) {
+    this.uploadTexture_ = uploadTexture;
+    this.disposeTexture_ = disposeTexture;
+  }
+
   public get queueStats(): QueueStats {
     return {
       pending: this.queue_.pendingCount,
@@ -44,12 +59,16 @@ export class ChunkManager {
       for (const chunk of updatedChunks) {
         if (chunk.priority === null) {
           this.queue_.cancel(chunk);
+          this.disposeChunkTexture(chunk);
+          clearChunkData(chunk);
         } else if (chunk.state === "queued") {
           this.queue_.enqueue(chunk, (signal) =>
             source.loader.loadChunkData(chunk, signal)
           );
         }
       }
+
+      this.uploadLoadedChunks(updatedChunks);
     }
 
     this.queue_.flush();
@@ -60,4 +79,38 @@ export class ChunkManager {
       }
     }
   }
+
+  private uploadLoadedChunks(chunks: Set<Chunk>) {
+    if (!this.uploadTexture_) return;
+
+    const pending: Chunk[] = [];
+
+    for (const chunk of chunks) {
+      if (chunk.state === "loaded" && chunk.texture === undefined) {
+        pending.push(chunk);
+      }
+    }
+
+    if (pending.length === 0) return;
+
+    pending.sort(comparePriority);
+
+    const limit = Math.min(pending.length, MAX_UPLOADS_PER_STORE_PER_UPDATE);
+
+    for (let i = 0; i < limit; i++) {
+      const chunk = pending[i];
+      const texture = Texture3D.createWithChunk(chunk);
+      this.uploadTexture_(texture);
+      chunk.texture = texture;
+      clearChunkData(chunk);
+    }
+  }
+
+  private disposeChunkTexture(chunk: Chunk) {
+    if (!this.disposeTexture_) return;
+
+    if (chunk.texture === undefined) return;
+    this.disposeTexture_(chunk.texture);
+    chunk.texture = undefined;
+  }
 }

src/data/chunk_memory.ts

@@ -26,6 +26,7 @@ export function clearChunkData(chunk: Chunk): void {
     cpuChunkCount -= 1;
   }
   chunk.data = undefined;
+  chunk.texture?.releaseCpuData();
 }
 
 export function chunkMemoryStats() {

src/data/chunk_queue.ts

@@ -7,6 +7,20 @@ type LoaderFn = (signal: AbortSignal) => Promise<void>;
 
 type PendingItem = { chunk: Chunk; fn: LoaderFn };
 
+export function comparePriority(a: Chunk, b: Chunk) {
+  const priorityA = a.priority ?? Number.MAX_SAFE_INTEGER;
+  const priorityB = b.priority ?? Number.MAX_SAFE_INTEGER;
+
+  if (priorityA === priorityB) {
+    return (
+      (a.orderKey ?? Number.MAX_SAFE_INTEGER) -
+      (b.orderKey ?? Number.MAX_SAFE_INTEGER)
+    );
+  }
+
+  return priorityA - priorityB;
+}
+
 export class ChunkQueue {
   private readonly maxConcurrent_: number;
   private readonly pending_: PendingItem[] = [];
@@ -61,19 +75,7 @@ export class ChunkQueue {
       return;
     }
 
-    this.pending_.sort((a, b) => {
-      const priorityA = a.chunk.priority ?? Number.MAX_SAFE_INTEGER;
-      const priorityB = b.chunk.priority ?? Number.MAX_SAFE_INTEGER;
-
-      if (priorityA === priorityB) {
-        return (
-          (a.chunk.orderKey ?? Number.MAX_SAFE_INTEGER) -
-          (b.chunk.orderKey ?? Number.MAX_SAFE_INTEGER)
-        );
-      }
-
-      return priorityA - priorityB;
-    });
+    this.pending_.sort((a, b) => comparePriority(a.chunk, b.chunk));
 
     while (
       this.running_.size < this.maxConcurrent_ &&

src/data/chunk_store.ts

@@ -1,5 +1,4 @@
 import { Chunk, SourceDimensionMap } from "./chunk";
-import { clearChunkData } from "./chunk_memory";
 import { almostEqual } from "../utilities/almost_equal";
 import { Logger } from "../utilities/logger";
 import { ChunkStoreView } from "./chunk_store_view";
@@ -218,7 +217,6 @@ export class ChunkStore {
         );
       }
 
-      clearChunkData(chunk);
       chunk.state = "unloaded";
       chunk.orderKey = null;
       Logger.debug(

src/data/chunk_store_view.ts

@@ -75,7 +75,8 @@ export class ChunkStoreView {
     const lowResChunks: Chunk[] = [];
 
     for (const [chunk, state] of this.chunkViewStates_) {
-      if (!state.visible || chunk.state !== "loaded") continue;
+      if (!state.visible || chunk.state !== "loaded" || !chunk.texture)
+        continue;
       if (chunk.lod === currentLOD) {
         currentLODChunks.push(chunk);
       } else if (chunk.lod === fallbackLOD && currentLOD !== fallbackLOD) {
@@ -113,7 +114,7 @@ export class ChunkStoreView {
         "ChunkStoreView",
         "updateChunkViewStates called with no chunks initialized"
       );
-      this.chunkViewStates_.clear();
+      this.chunkViewStates_.forEach(resetChunkViewState);
       return;
     }
 
@@ -202,7 +203,7 @@ export class ChunkStoreView {
         "ChunkStoreView",
         "updateChunksForVolume called with no chunks initialized"
       );
-      this.chunkViewStates_.clear();
+      this.chunkViewStates_.forEach(resetChunkViewState);
       return;
     }
 

src/idetik.ts

@@ -113,7 +113,10 @@ export class Idetik {
     this.canvas = params.canvas;
 
     this.renderer_ = renderer ?? new WebGLRenderer(this.canvas);
-    this.chunkManager_ = new ChunkManager();
+    this.chunkManager_ = new ChunkManager(
+      (texture) => this.renderer_.uploadTexture(texture),
+      (texture) => this.renderer_.disposeTexture(texture)
+    );
     this.context_ = {
       chunkManager: this.chunkManager_,
     };

src/layers/image_layer.ts

@@ -11,13 +11,13 @@ import {
   validateChannelPropsCount,
 } from "../core/channel";
 import { ImageRenderable } from "../objects/renderable/image_renderable";
-import { Texture3D } from "../objects/textures/texture_3d";
 import { Color } from "../math/color";
 import { EventContext } from "../core/event_dispatcher";
 import { vec2, vec3 } from "gl-matrix";
 import { handlePointPickingEvent, PointPickingResult } from "./point_picking";
 import { clamp } from "../utilities/clamp";
 import { RenderablePool } from "../utilities/renderable_pool";
+import { Texture } from "../objects/textures/texture";
 
 export type ImageLayerProps = LayerOptions & {
   source: ChunkSource;
@@ -128,8 +128,13 @@ export class ImageLayer extends Layer implements ChannelsEnabled {
     if (!this.chunkStoreView_) return;
     if (this.state !== "ready") this.setState("ready");
 
+    const visibleChunksResident = Array.from(this.visibleChunks_.keys()).every(
+      (chunk) => chunk.texture !== undefined
+    );
+
     if (
       this.visibleChunks_.size > 0 &&
+      visibleChunksResident &&
       !this.chunkStoreView_.allVisibleFallbackLODLoaded() &&
       !this.isPresentationStale()
     ) {
@@ -147,8 +152,7 @@ export class ImageLayer extends Layer implements ChannelsEnabled {
 
     this.clearObjects();
     for (const chunk of orderedByLOD) {
-      if (chunk.state !== "loaded") continue;
-      const image = this.getImageForChunk(chunk);
+      const image = this.getImageForChunk(chunk, chunk.texture!);
       this.visibleChunks_.set(chunk, image);
       this.addObject(image);
     }
@@ -209,35 +213,32 @@ export class ImageLayer extends Layer implements ChannelsEnabled {
     }
   }
 
-  private getImageForChunk(chunk: Chunk) {
+  private getImageForChunk(chunk: Chunk, texture: Texture) {
     const existing = this.visibleChunks_.get(chunk);
     if (existing) return existing;
 
     const pooled = this.pool_.acquire(poolKeyForImageRenderable(chunk));
     if (pooled) {
-      const texture = pooled.textures[0] as Texture3D;
-      texture.updateWithChunk(chunk);
-
+      pooled.setTexture(0, texture);
       pooled.zTexCoord = this.zTexCoordForChunk(chunk);
       pooled.setChannelProps(this.getChannelPropsForChunk(chunk));
       this.updateImageChunk(pooled, chunk);
-
       return pooled;
     }
 
-    return this.createImage(chunk);
+    return this.createImage(chunk, texture);
   }
 
   private getChannelPropsForChunk(chunk: Chunk): ChannelProps[] {
     if (!this.channelProps_) return [{}];
     return [this.channelProps_[chunk.chunkIndex.c] ?? {}];
   }
 
-  private createImage(chunk: Chunk) {
+  private createImage(chunk: Chunk, texture: Texture) {
     const image = new ImageRenderable(
       chunk.shape.x,
       chunk.shape.y,
-      Texture3D.createWithChunk(chunk),
+      texture,
       this.getChannelPropsForChunk(chunk)
     );
     image.zTexCoord = this.zTexCoordForChunk(chunk);

src/layers/label_layer.ts

@@ -10,7 +10,7 @@ import {
   LabelColorMap,
   LabelColorMapProps,
 } from "../objects/renderable/label_color_map";
-import { Texture3D } from "../objects/textures/texture_3d";
+import { Texture } from "../objects/textures/texture";
 import { EventContext } from "../core/event_dispatcher";
 import { vec2, vec3 } from "gl-matrix";
 import { handlePointPickingEvent, PointPickingResult } from "./point_picking";
@@ -114,8 +114,13 @@ export class LabelLayer extends Layer {
     if (!this.chunkStoreView_) return;
     if (this.state !== "ready") this.setState("ready");
 
+    const visibleChunksResident = Array.from(this.visibleChunks_.keys()).every(
+      (chunk) => chunk.texture !== undefined
+    );
+
     if (
       this.visibleChunks_.size > 0 &&
+      visibleChunksResident &&
       !this.chunkStoreView_.allVisibleFallbackLODLoaded() &&
       !this.isPresentationStale()
     ) {
@@ -133,8 +138,7 @@ export class LabelLayer extends Layer {
 
     this.clearObjects();
     for (const chunk of orderedByLOD) {
-      if (chunk.state !== "loaded") continue;
-      const label = this.getLabelForChunk(chunk);
+      const label = this.getLabelForChunk(chunk, chunk.texture!);
       this.visibleChunks_.set(chunk, label);
       this.addObject(label);
     }
@@ -248,15 +252,13 @@ export class LabelLayer extends Layer {
     return (await label.textures[0].readTexel?.(x, y, z)) ?? null;
   }
 
-  private getLabelForChunk(chunk: Chunk) {
+  private getLabelForChunk(chunk: Chunk, texture: Texture) {
     const existing = this.visibleChunks_.get(chunk);
     if (existing) return existing;
 
     const pooled = this.pool_.acquire(poolKeyForImageRenderable(chunk));
     if (pooled) {
-      const texture = pooled.textures[0] as Texture3D;
-      texture.updateWithChunk(chunk);
-
+      pooled.setTexture(0, texture);
       pooled.zTexCoord = this.zTexCoordForChunk(chunk);
       pooled.setColorMap(this.colorMap_);
       pooled.setSelectedValue(this.selectedValue_);
@@ -265,14 +267,14 @@ export class LabelLayer extends Layer {
       return pooled;
     }
 
-    return this.createLabel(chunk);
+    return this.createLabel(chunk, texture);
   }
 
-  private createLabel(chunk: Chunk) {
+  private createLabel(chunk: Chunk, texture: Texture) {
     const label = new LabelImageRenderable({
       width: chunk.shape.x,
       height: chunk.shape.y,
-      imageData: Texture3D.createWithChunk(chunk),
+      imageData: texture,
       colorMap: this.colorMap_,
       outlineSelected: this.outlineSelected_,
       selectedValue: this.selectedValue_,