DeviceInterface and unit tests (Data Manager 3/9)

src/data_manager/data_manager.ts

@@ -1,12 +1,4 @@
-import type {
-  ChunkId,
-  ChunkPriority,
-  ChunkEntry,
-  DataManagerLimits,
-  LocalChunkId,
-  Chunk,
-  DeviceInterface,
-} from "./types.js";
+import type { ChunkId, ChunkPriority, ChunkEntry, DataManagerLimits, LocalChunkId, Chunk } from "./types.js";
 import {
   chunkIdToString,
   ChunkState,
@@ -23,6 +15,8 @@ import {
 import type { NumberType, TypedArray } from "../types.js";
 import PriorityQueue from "./priority_queue.js";
 import { VolumeDims } from "../VolumeDims.js";
+import type { DeviceInterface } from "./device_interface.js";
+import { dataTypeToSize } from "./utils.js";
 
 const SUBSCRIBER_ID = Symbol("DataManager.subscriberId");
 
@@ -60,17 +54,6 @@ const swapRemove = <T>(arr: T[], index: number) => {
   }
 };
 
-const dataTypeToByteLength: { [T in NumberType]: number } = {
-  int8: 1,
-  int16: 2,
-  int32: 4,
-  uint8: 1,
-  uint16: 2,
-  uint32: 4,
-  float32: 4,
-  float64: 8,
-};
-
 type ChunkQueue = PriorityQueue<ChunkPriority, string>;
 
 class ChunkQueues {
@@ -208,7 +191,7 @@ export default class DataManager<Dev, Tex> {
       return 0;
     }
     const { x, y, z, dataType } = dims;
-    return x * y * z * dataTypeToByteLength[dataType];
+    return x * y * z * dataTypeToSize[dataType];
   }
 
   /**

src/data_manager/device_interface.ts

@@ -0,0 +1,141 @@
+import {
+  ByteType,
+  Data3DTexture,
+  FloatType,
+  IntType,
+  type PixelFormat,
+  PixelFormatGPU,
+  RedFormat,
+  RedIntegerFormat,
+  ShortType,
+  type TextureDataType,
+  UnsignedByteType,
+  UnsignedIntType,
+  UnsignedShortType,
+  type WebGLRenderer,
+} from "three";
+import type { Renderer } from "three/webgpu";
+
+import type { NumberType, TypedArray } from "../types.js";
+
+/** Abstraction over the graphics API or library used by `DataManager` to move chunks to/from the GPU. */
+export interface DeviceInterface<Dev, Tex> {
+  /**
+   * Type predicate for this interface's *device handle*, for APIs (like WebGPU) that require explicit access to a
+   * shared resource to perform GPU operations.
+   */
+  isDeviceHandle(val: unknown): val is Dev;
+  /** Create a texture with the given `data`, `dtype`, and `size` using `deviceHandle`. */
+  createTexture(data: TypedArray, dtype: NumberType, size: [number, number, number], deviceHandle: Dev): Tex;
+  /** Destroy a texture and clean up its GPU resources. */
+  destroyTexture(tex: Tex): void;
+  /** Called when `DataManager` is done with an update cycle, in case the interface has state to clean up. */
+  finishUpdate(): void;
+}
+
+// MARK: Three
+
+const dataTypeToThreeTextureProperties: {
+  [T in Exclude<NumberType, "float64">]: [TextureDataType, PixelFormat, PixelFormatGPU];
+} = {
+  int8: [ByteType, RedIntegerFormat, "R8I"],
+  int16: [ShortType, RedIntegerFormat, "R16I"],
+  int32: [IntType, RedIntegerFormat, "R32I"],
+  uint8: [UnsignedByteType, RedIntegerFormat, "R8UI"],
+  uint16: [UnsignedShortType, RedIntegerFormat, "R16UI"],
+  uint32: [UnsignedIntType, RedIntegerFormat, "R32UI"],
+  float32: [FloatType, RedFormat, "R32F"],
+};
+
+export class ThreeInterface implements DeviceInterface<Renderer | WebGLRenderer, Data3DTexture> {
+  isDeviceHandle(val: unknown): val is Renderer | WebGLRenderer {
+    return (
+      typeof val === "object" && val !== null && typeof (val as Renderer | WebGLRenderer).initTexture === "function"
+    );
+  }
+
+  createTexture(
+    data: TypedArray,
+    dtype: NumberType,
+    [x, y, z]: [number, number, number],
+    renderer: Renderer | WebGLRenderer
+  ) {
+    const [texType, texFormat, texInternalFormat] = dataTypeToThreeTextureProperties[dtype];
+    const texture = new Data3DTexture(data, x, y, z);
+    texture.type = texType;
+    texture.format = texFormat;
+    texture.internalFormat = texInternalFormat;
+    texture.needsUpdate = true;
+    renderer.initTexture(texture);
+    return texture;
+  }
+
+  destroyTexture(tex: Data3DTexture): void {
+    tex.dispose();
+  }
+
+  finishUpdate() {
+    // noop
+  }
+}
+
+// MARK: WebGPU
+
+export type GPUDeviceHandle = {
+  device: GPUDevice;
+  queue: GPUQueue;
+};
+
+const dataTypeToFormat: { [T in Exclude<NumberType, "int32" | "uint32" | "float64">]: GPUTextureFormat } = {
+  int8: "r8snorm",
+  int16: "r16snorm",
+  uint8: "r8unorm",
+  uint16: "r16unorm",
+  float32: "r32float",
+};
+
+export class WebGPUInterface implements DeviceInterface<GPUDeviceHandle, GPUTexture> {
+  isDeviceHandle(val: unknown): val is GPUDeviceHandle {
+    return (val as GPUDeviceHandle).device instanceof GPUDevice && (val as GPUDeviceHandle).queue instanceof GPUQueue;
+  }
+
+  createTexture(
+    data: TypedArray,
+    dtype: NumberType,
+    size: [number, number, number],
+    { device, queue }: GPUDeviceHandle
+  ): GPUTexture {
+    const format = dataTypeToFormat[dtype];
+    const texture = device.createTexture({
+      format,
+      size,
+      mipLevelCount: 1,
+      sampleCount: 1,
+      dimension: "3d",
+      // 2 === GPUTextureUsage.COPY_DST; 4 === GPUTextureUsage.TEXTURE_BINDING
+      usage: 2 | 4,
+    });
+    const copyInfo: GPUTexelCopyTextureInfo = {
+      texture,
+      mipLevel: 0,
+      origin: [0, 0, 0],
+      aspect: "all",
+    };
+    const layout: GPUTexelCopyBufferLayout = {
+      offset: 0,
+      bytesPerRow: texture.width * data.BYTES_PER_ELEMENT,
+      rowsPerImage: texture.height,
+    };
+    queue.writeTexture(copyInfo, data, layout, size);
+    return texture;
+  }
+
+  destroyTexture(texture: GPUTexture) {
+    texture.destroy();
+  }
+
+  finishUpdate(): void {
+    // noop for now, but this class could be made stateful to e.g. reuse same-sized textures on eviction and use
+    // this method to definitively destroy unused textures
+  }
+}

src/data_manager/types.ts

@@ -46,21 +46,6 @@ export const validateDataManagerLimits = (limits: DataManagerLimits): DataManage
   concurrentPrefetches: Math.min(limits.concurrentRequests, limits.concurrentPrefetches),
 });
 
-/** Abstraction over the graphics API or library used by `DataManager` to move chunks to/from the GPU. */
-export interface DeviceInterface<Dev, Tex> {
-  /**
-   * Type predicate for this interface's *device handle*, for APIs (like WebGPU) that require explicit access to a
-   * shared resource to perform GPU operations.
-   */
-  isDeviceHandle(val: unknown): val is Dev;
-  /** Create a texture with the given `data`, `dtype`, and `size` using `deviceHandle`. */
-  createTexture(data: TypedArray, dtype: NumberType, size: [number, number, number], deviceHandle: Dev): Tex;
-  /** Destroy a texture and clean up its GPU resources. */
-  destroyTexture(tex: Tex): void;
-  /** Called when `DataManager` is done with an update cycle, in case the interface has state to clean up. */
-  finishUpdate(): void;
-}
-
 /** Reasons that a chunk is tracked by `DataManager`, in increasing order of priority. */
 export const enum ChunkPriorityLevel {
   /**

src/data_manager/utils.ts

@@ -0,0 +1,12 @@
+import type { NumberType } from "../types.js";
+
+export const dataTypeToSize: { [T in NumberType]: number } = {
+  int8: 1,
+  int16: 2,
+  int32: 4,
+  uint8: 1,
+  uint16: 2,
+  uint32: 4,
+  float32: 4,
+  float64: 8,
+};

src/test/data_manager.test.ts

@@ -0,0 +1,102 @@
+import type { IChunkSource, IDataSubscriber } from "../data_manager/data_manager.js";
+import DataManager from "../data_manager/data_manager.js";
+import type { DeviceInterface } from "../data_manager/device_interface.js";
+import { type ChunkPriority, ChunkPriorityLevel } from "../data_manager/types.js";
+import type { TypedArray, NumberType } from "../types.js";
+
+const mockSource = () => {
+  return {
+    getDims: vi.fn<IChunkSource["getDims"]>(() => [
+      {
+        shape: [7, 7, 7, 7, 7],
+        spacing: [1, 1, 1, 1, 1],
+        chunkShape: [1, 1, 2, 2, 2],
+        dataType: "uint8",
+        spaceUnit: "micrometer",
+        timeUnit: "second",
+      },
+    ]),
+    getChunk: vi.fn<IChunkSource["getChunk"]>(async ({ tczyx }) => {
+      const length = tczyx.slice(2).reduce((len, coord) => (coord >= 3 ? len : len * 2), 1);
+      const data = new Uint8Array(length).fill(0);
+      data.set(tczyx);
+      return { data, dtype: "uint8" };
+    }),
+  } satisfies IChunkSource;
+};
+
+type MockTex = {
+  data: TypedArray;
+  dtype: NumberType;
+  size: [number, number, number];
+  id: number;
+};
+
+let mockTexCount = 0;
+
+const mockInterface = (): DeviceInterface<void, MockTex> => ({
+  isDeviceHandle: (_val): _val is void => true,
+  createTexture: vi.fn((data, dtype, size) => {
+    const result = { data, dtype, size, id: mockTexCount };
+    mockTexCount += 1;
+    return result;
+  }),
+  destroyTexture: vi.fn(),
+  finishUpdate: vi.fn(),
+});
+
+const mockSubscriber = (): IDataSubscriber<MockTex> => ({
+  onChunkLoaded: vi.fn(),
+  onChunkOnGpu: vi.fn(),
+});
+
+describe("request priority", () => {
+  it("orders requests by priority", () => {
+    const manager = new DataManager(mockInterface());
+    const sourceMock = mockSource();
+    const source = manager.addSource(sourceMock);
+    const subscriber = mockSubscriber();
+
+    const queue = (x: number, score: number) =>
+      manager.addChunkRequest(
+        subscriber,
+        { source, multiscale: 0, tczyx: [0, 0, 0, 0, x] },
+        { level: ChunkPriorityLevel.VISIBLE, score },
+        false
+      );
+    queue(0, 1);
+    queue(1, 0);
+    queue(2, 2);
+    manager.update();
+
+    const getCalls = sourceMock.getChunk.mock.calls;
+    expect(getCalls[0][0].tczyx[4]).toEqual(2);
+    expect(getCalls[1][0].tczyx[4]).toEqual(0);
+    expect(getCalls[2][0].tczyx[4]).toEqual(1);
+  });
+
+  it("orders `VISIBLE` requests ahead of `PREFETCH` ones", () => {
+    const manager = new DataManager(mockInterface());
+    const sourceMock = mockSource();
+    const source = manager.addSource(sourceMock);
+    const subscriber = mockSubscriber();
+
+    const queue = (x: number, priority: ChunkPriority) =>
+      manager.addChunkRequest(subscriber, { source, multiscale: 0, tczyx: [0, 0, 0, 0, x] }, priority, false);
+    queue(0, { level: ChunkPriorityLevel.PREFETCH, score: 1000 });
+    queue(1, { level: ChunkPriorityLevel.VISIBLE, score: 0 });
+    manager.update();
+
+    const getCalls = sourceMock.getChunk.mock.calls;
+    expect(getCalls[0][0].tczyx[4]).toEqual(1);
+    expect(getCalls[1][0].tczyx[4]).toEqual(0);
+  });
+
+  it("resolves chunk priorities across all subscribers", () => {
+    // TODO
+  });
+
+  it("handles memory and device priorities separately", () => {
+    // TODO
+  });
+});