diff --git a/src/data_manager/data_manager.ts b/src/data_manager/data_manager.ts index 5ce886aa..c8badbca 100644 --- a/src/data_manager/data_manager.ts +++ b/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 = (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; class ChunkQueues { @@ -208,7 +191,7 @@ export default class DataManager { return 0; } const { x, y, z, dataType } = dims; - return x * y * z * dataTypeToByteLength[dataType]; + return x * y * z * dataTypeToSize[dataType]; } /** diff --git a/src/data_manager/device_interface.ts b/src/data_manager/device_interface.ts new file mode 100644 index 00000000..901df9ac --- /dev/null +++ b/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 { + /** + * 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]: [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 { + 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]: GPUTextureFormat } = { + int8: "r8snorm", + int16: "r16snorm", + uint8: "r8unorm", + uint16: "r16unorm", + float32: "r32float", +}; + +export class WebGPUInterface implements DeviceInterface { + 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 + } +} diff --git a/src/data_manager/types.ts b/src/data_manager/types.ts index 39125fcd..a20dfb7e 100644 --- a/src/data_manager/types.ts +++ b/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 { - /** - * 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 { /** diff --git a/src/data_manager/utils.ts b/src/data_manager/utils.ts new file mode 100644 index 00000000..122931ae --- /dev/null +++ b/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, +}; diff --git a/src/test/data_manager.test.ts b/src/test/data_manager.test.ts new file mode 100644 index 00000000..47b4fb52 --- /dev/null +++ b/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(() => [ + { + 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(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 => ({ + 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 => ({ + 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 + }); +});