oidn-webgpu 0.0.0-stage → 0.6.0

package/lib/nativeUNetCore.d.ts

@@ -0,0 +1,202 @@
+import { type OIDNGPUBindGroup, type OIDNGPUBuffer, type OIDNGPUComputePipeline, type OIDNGPUDevice, type OIDNGPUQuerySet } from './gpuTypes.js';
+import { type ExecutableModelNode, type ModelExecutionPlan } from './graphOptimizer.js';
+import type { UNetModelGraph, ValidatedUNetModel } from './modelSpec';
+import { OIDNResourceTracker, type OIDNResourceSnapshot } from './resourceTracker.js';
+export type NativeUNetPrecision = 'fp32' | 'fp16';
+export type NativeUNetPrecisionSetting = NativeUNetPrecision | 'auto';
+export type NativeUNetKernel = 'direct' | 'implicit-gemm' | 'spatial' | 'subgroup' | 'subgroup-matrix';
+export type NativeUNetKernelSetting = NativeUNetKernel | 'auto';
+export type NativeUNetGemmWorkgroup = readonly [4 | 8 | 16, 4 | 8];
+export interface NativeUNetGemmOptions {
+    /** Internal per-layer tile selection; explicit workgroup sizes remain fixed by default. */
+    tilePolicy?: 'fixed' | 'output-aligned';
+    /** Internal decoder source-branch scheduling experiment. */
+    decoderLoad?: 'per-load' | 'source-first';
+    /** Sub-pixel GEMM decoder experiment; defaults to FP16 with K-major weights. */
+    decoderSubPixel?: boolean;
+    /** Expand register-tile rows for native shader robustness; off by default in browsers. */
+    unrollRows?: boolean;
+    /** Internal pooling dispatch experiment for the GEMM execution path. */
+    poolLayout?: 'spatial' | 'channels';
+    /** Internal shared-memory padding experiment; does not change storage ABI. */
+    sharedLayout?: 'linear' | 'padded' | 'padded-input' | 'padded-weights';
+    /** Internal loop scheduling experiment; each output keeps the same K order. */
+    accumulationOrder?: 'k-major' | 'row-major';
+    /** Internal final-convolution experiments; retains the original FP16 partial order. */
+    finalLayer?: 'direct' | 'shared-input' | 'shared-input-weights' | 'shared-auto';
+    /** Internal load-layout experiments; packed loads preserve the FP16 storage bits. */
+    loadMode?: 'native' | 'packed-weights' | 'packed-all';
+    /** Incremental addressing (default), or the original analytic path for A/B comparisons. */
+    addressMode?: 'analytic' | 'incremental' | 'base-offset';
+    /** K-major (default) coalesces neighboring output-block loads; output-major preserves the old ABI. */
+    weightLayout?: 'output-major' | 'k-major';
+    /** Register tile height (default 8). The K reduction grouping stays fixed. */
+    rowsPerThread?: 2 | 4 | 8;
+    /** Workgroup width/height (default [8, 8]); width also selects output channel blocks. */
+    workgroupSize?: NativeUNetGemmWorkgroup;
+}
+/** Square multiply-add shape selected from `chromium-experimental-subgroup-matrix`. */
+export interface NativeSubgroupMatrixConfig {
+    componentType: 'f16' | 'f32';
+    resultType: 'f16' | 'f32';
+    /** M = N = K. */
+    size: 8 | 16;
+    subgroupSize: number;
+    /**
+     * WGSL spelling: Chrome's `chromium_experimental_subgroup_matrix`, or
+     * wgpu/Naga's `wgpu_cooperative_matrix` for native hosts.
+     */
+    dialect: 'tint' | 'naga';
+}
+/** Experimental subgroup-matrix tile tuning; not covered by semver. */
+export interface NativeSubgroupMatrixOptions {
+    /** Output channel fragments per workgroup (default 2). */
+    outputFragments?: 1 | 2 | 4;
+    /** Preferred accumulation type when the adapter offers several. */
+    accumulation?: 'f32' | 'f16';
+    /** Sub-pixel decomposition of the upsampling decoder layers (default true). */
+    subPixel?: boolean;
+    /**
+     * Double-buffered staging with register prefetch when the workgroup memory
+     * limit allows it (default false: 6-10% slower on an Apple M4 Max).
+     */
+    doubleBuffer?: boolean;
+}
+export interface NativeUNetOptions {
+    precision?: NativeUNetPrecisionSetting;
+    /** Experimental `subgroup-matrix` kernel tuning. */
+    subgroupMatrix?: NativeSubgroupMatrixOptions;
+    /**
+     * Convolution kernel selection. `auto` uses implicit GEMM for FP16/FP32
+     * convolutions and the direct kernel for the final output layer.
+     */
+    kernel?: NativeUNetKernelSetting;
+    /** Optional implicit-GEMM tuning; other convolution kernels ignore it. */
+    gemm?: NativeUNetGemmOptions;
+    /** Maximum number of shape-dependent activation plans retained. */
+    shapeCacheSize?: number;
+}
+export interface NativeUNetLayerTiming {
+    id: string;
+    durationMs: number;
+}
+export interface NativeUNetExecutionProfile {
+    totalMs: number;
+    layers: NativeUNetLayerTiming[];
+}
+export interface NativePipelineSpec {
+    key: string;
+    code: string;
+    kernel?: NativeUNetKernel;
+}
+interface ActivationSlot {
+    buffer: OIDNGPUBuffer;
+    capacity: number;
+}
+export interface CachedExecution {
+    plan: ModelExecutionPlan;
+    valueBuffers: Map<string, OIDNGPUBuffer>;
+    slots: ActivationSlot[];
+    nodeBindings: OIDNGPUBindGroup[];
+    nodePipelines: OIDNGPUComputePipeline[];
+    nodeKernels: NativeUNetKernel[];
+    /** Subgroup-matrix output fragments per node (0 for other kernels). */
+    nodeFragments: number[];
+    inputPipeline: OIDNGPUComputePipeline;
+    inputUniform: OIDNGPUBuffer;
+    ownedBuffers: OIDNGPUBuffer[];
+    cpuInputBuffers?: OIDNGPUBuffer[];
+    cpuReadbackBuffer?: OIDNGPUBuffer;
+    lastUsed: number;
+}
+export declare function createInputPackShader(precision: NativeUNetPrecision, sourceCount: number): string;
+/**
+ * Picks a square configuration, preferring f32 accumulation: mixed f16
+ * inputs first, then f32 inputs (whose f16-exact values give the same
+ * products), and f16 accumulation last.
+ */
+export declare function selectSubgroupMatrixConfig(device: OIDNGPUDevice, accumulation?: 'f32' | 'f16'): NativeSubgroupMatrixConfig | undefined;
+export declare function resolveNativeUNetPrecision(device: OIDNGPUDevice, requested?: NativeUNetPrecisionSetting): NativeUNetPrecision;
+/**
+ * Native, model-driven OIDN U-Net executor without asynchronous operations, so
+ * that scriptc can compile it for native hosts. `NativeUNetExecutor` adds the
+ * browser's asynchronous preparation, CPU readback and profiling.
+ */
+export declare class NativeUNetCore {
+    protected _device: OIDNGPUDevice;
+    readonly usedKernels: Set<NativeUNetKernel>;
+    /** Immutable packed weights and biases, for native resource sharing. */
+    get weightBuffers(): OIDNGPUBuffer[];
+    readonly precision: NativeUNetPrecision;
+    readonly kernelSetting: NativeUNetKernelSetting;
+    readonly gemm: Readonly<Required<NativeUNetGemmOptions>>;
+    readonly maxSpatialInputBlocks: number;
+    readonly subgroupsAvailable: boolean;
+    /** Selected configuration when the `subgroup-matrix` kernel can run. */
+    readonly subgroupMatrix?: NativeSubgroupMatrixConfig;
+    readonly subgroupMatrixOutputFragments: number;
+    readonly subgroupMatrixSubPixel: boolean;
+    readonly subgroupMatrixDoubleBuffer: boolean;
+    private _subgroupMatrixFixedFragments;
+    protected _model: UNetModelGraph;
+    private _gemmByOutputBlocks;
+    private _packedConvs;
+    protected _pipelineCache: Map<string, OIDNGPUComputePipeline>;
+    private _executionCache;
+    protected _retiredExecutions: Set<CachedExecution>;
+    private _clock;
+    private _shapeCacheSize;
+    private _profileNextExecution;
+    protected _profileOperations: number;
+    protected _resources: OIDNResourceTracker;
+    protected _disposed: boolean;
+    constructor(_device: OIDNGPUDevice, model: ValidatedUNetModel, options?: NativeUNetOptions);
+    private _pipeline;
+    protected _nodePipelineSpec(node: ExecutableModelNode, isFinal: boolean, fragments?: number): NativePipelineSpec;
+    /** The decoder's sources when it runs with the sub-pixel decomposition. */
+    private _subPixelDecoder;
+    private _subgroupMatrixSpec;
+    /**
+     * Output fragments per workgroup for a subgroup-matrix layer: the default
+     * unless that leaves fewer than SUBGROUP_MATRIX_MIN_WORKGROUPS workgroups,
+     * where one fragment keeps more of the GPU busy (low-resolution layers).
+     * Never more fragments than the output channels fill (the final RGB layer).
+     */
+    private _subgroupMatrixFragmentsFor;
+    private _gemmForOutput;
+    private _matrixPool;
+    private _coalescedPool;
+    private _selectConvKernel;
+    private _nodePipeline;
+    private _createExecution;
+    protected _execution(width: number, height: number): CachedExecution;
+    /**
+     * Releases an evicted plan. The synchronous core destroys it at once; the
+     * browser executor waits for the queue first.
+     */
+    protected _retireExecution(execution: CachedExecution): void;
+    /** Allocates shape-dependent execution resources before interactive use. */
+    prewarm(shapes: readonly {
+        width: number;
+        height: number;
+    }[]): void;
+    /** Captures per-pass GPU timestamps for the next execute call when supported. */
+    profileNextExecution(): boolean;
+    /** Shape-owned packed input; external preprocessing must write all lanes. */
+    packedInputBuffer(width: number, height: number): OIDNGPUBuffer;
+    /** Executes after a preprocessing pass has filled packedInputBuffer. */
+    executePacked(width: number, height: number): OIDNGPUBuffer;
+    execute(inputBuffers: readonly OIDNGPUBuffer[], width: number, height: number): OIDNGPUBuffer;
+    private _execute;
+    /**
+     * Receives the timestamp queries of a profiled execution. The synchronous
+     * core cannot read them back, so it releases them.
+     */
+    protected _readProfile(querySet: OIDNGPUQuerySet, resolveBuffer: OIDNGPUBuffer, readbackBuffer: OIDNGPUBuffer, labels: string[]): void;
+    protected _releaseBuffer(buffer: OIDNGPUBuffer | undefined): void;
+    protected _releaseQuerySet(querySet: OIDNGPUQuerySet | undefined): void;
+    protected _destroyExecution(execution: CachedExecution): void;
+    getResourceInfo(): OIDNResourceSnapshot;
+    dispose(): void;
+}
+export {};