feat: add WebGPU DAS backend
This commit is contained in:
Codex
@@ -0,0 +1,252 @@
|
||||
import {
|
||||
DAS_PARAMS_BYTE_LENGTH,
|
||||
packDasParams,
|
||||
packMicrophoneCalibration,
|
||||
} from "/packages/webgpu-kernels/das.wgsl.js";
|
||||
import {
|
||||
WebGpuLimitError,
|
||||
WebGpuUnavailableError,
|
||||
estimateWebGpuBatchPlan,
|
||||
estimateWebGpuBatchSize,
|
||||
packStftBatch,
|
||||
} from "/packages/webgpu-kernels/das.js";
|
||||
import { computeDasWithBackend } from "/packages/webgpu-kernels/compute.js";
|
||||
|
||||
function assertNear(actual, expected, tolerance, label) {
|
||||
assert(
|
||||
Math.abs(actual - expected) <= tolerance,
|
||||
`${label}: expected ${expected} ± ${tolerance}, received ${actual}`,
|
||||
);
|
||||
}
|
||||
|
||||
test("WebGPU DAS uniform과 microphone calibration을 WGSL layout으로 packing한다", () => {
|
||||
const params = packDasParams({
|
||||
pointCount: 81,
|
||||
microphoneCount: 16,
|
||||
frameCount: 7,
|
||||
binCount: 3,
|
||||
sampleRate: 8_000,
|
||||
fftSize: 256,
|
||||
speedOfSound: 343.25,
|
||||
});
|
||||
const view = new DataView(params);
|
||||
|
||||
assertEqual(params.byteLength, DAS_PARAMS_BYTE_LENGTH, "uniform byte length");
|
||||
assertEqual(view.getUint32(0, true), 81, "point count offset");
|
||||
assertEqual(view.getUint32(4, true), 16, "microphone count offset");
|
||||
assertEqual(view.getUint32(8, true), 7, "frame count offset");
|
||||
assertEqual(view.getUint32(12, true), 3, "bin count offset");
|
||||
assertEqual(view.getFloat32(16, true), Math.fround(8_000), "sample rate offset");
|
||||
assertEqual(view.getUint32(20, true), 256, "FFT size offset");
|
||||
assertEqual(view.getFloat32(24, true), Math.fround(343.25), "speed of sound offset");
|
||||
assertEqual(view.getUint32(28, true), 0, "reserved word");
|
||||
|
||||
const calibration = packMicrophoneCalibration([
|
||||
{ sensitivityDb: 0, phaseOffsetDeg: 0 },
|
||||
{ sensitivityDb: 6.020599913, phaseOffsetDeg: 90 },
|
||||
{ sensitivityDb: -6.020599913, phaseOffsetDeg: -45 },
|
||||
]);
|
||||
assertEqual(calibration.length, 6, "two f32 values per microphone");
|
||||
assertNear(calibration[0], 1, 1e-7, "unity inverse gain");
|
||||
assertNear(calibration[1], 0, 1e-7, "zero phase");
|
||||
assertNear(calibration[2], 0.5, 1e-6, "+6 dB inverse gain");
|
||||
assertNear(calibration[3], Math.PI / 2, 1e-6, "+90 degree phase");
|
||||
assertNear(calibration[4], 2, 1e-6, "-6 dB inverse gain");
|
||||
assertNear(calibration[5], -Math.PI / 4, 1e-6, "-45 degree phase");
|
||||
});
|
||||
|
||||
test("STFT batch를 bin-frame-microphone-complex 순서로 packing한다", () => {
|
||||
const stft = {
|
||||
frameCount: 2,
|
||||
microphoneCount: 2,
|
||||
spectra: new Map([
|
||||
[3, {
|
||||
real: Float32Array.of(1, 2, 3, 4),
|
||||
imaginary: Float32Array.of(11, 12, 13, 14),
|
||||
}],
|
||||
[7, {
|
||||
real: Float32Array.of(5, 6, 7, 8),
|
||||
imaginary: Float32Array.of(15, 16, 17, 18),
|
||||
}],
|
||||
]),
|
||||
};
|
||||
|
||||
const packed = packStftBatch(stft, [7, 3]);
|
||||
assert(packed instanceof Float32Array, "packed STFT type");
|
||||
const expected = [
|
||||
5, 15, 6, 16, 7, 17, 8, 18,
|
||||
1, 11, 2, 12, 3, 13, 4, 14,
|
||||
];
|
||||
assertEqual(packed.length, expected.length, "packed STFT length");
|
||||
expected.forEach((value, index) => assertEqual(packed[index], value, `packed STFT value ${index}`));
|
||||
});
|
||||
|
||||
test("WebGPU batch size가 storage, dispatch, memory 한도를 모두 지킨다", () => {
|
||||
const workload = {
|
||||
selectedBins: [4, 5, 6, 7, 8],
|
||||
bandBins: [[4, 5, 6, 7, 8]],
|
||||
frameCount: 2,
|
||||
microphoneCount: 4,
|
||||
pointCount: 8,
|
||||
distanceBytes: 128,
|
||||
outputBytes: 64,
|
||||
spectrumBytesPerBin: 64,
|
||||
baseWorkingBytes: 256,
|
||||
};
|
||||
const limits = {
|
||||
maxStorageBufferBindingSize: 192,
|
||||
maxBufferSize: 192,
|
||||
maxComputeWorkgroupsPerDimension: 2,
|
||||
maxComputeInvocationsPerWorkgroup: 64,
|
||||
maxComputeWorkgroupSizeX: 64,
|
||||
};
|
||||
|
||||
const batchSize = estimateWebGpuBatchSize({
|
||||
workload,
|
||||
limits,
|
||||
maximumGpuWorkingBytes: 4_096,
|
||||
maximumHostWorkingBytes: 4_096,
|
||||
});
|
||||
assertEqual(batchSize, 2, "dispatch y dimension is the tightest batch limit");
|
||||
const plan = estimateWebGpuBatchPlan({
|
||||
workload,
|
||||
limits,
|
||||
maximumGpuWorkingBytes: 4_096,
|
||||
maximumHostWorkingBytes: 4_096,
|
||||
});
|
||||
assertEqual(plan.pointBatchSize, 2, "steering point dispatch batch");
|
||||
assertEqual(plan.steeringBytes, 128, "two-pass steering buffer bytes");
|
||||
|
||||
let limitError;
|
||||
try {
|
||||
estimateWebGpuBatchSize({
|
||||
workload,
|
||||
limits: { ...limits, maxStorageBufferBindingSize: 31 },
|
||||
maximumGpuWorkingBytes: 4_096,
|
||||
maximumHostWorkingBytes: 4_096,
|
||||
});
|
||||
} catch (error) {
|
||||
limitError = error;
|
||||
}
|
||||
assert(limitError instanceof WebGpuLimitError, "static calibration buffer limit error type");
|
||||
|
||||
let invocationLimitError;
|
||||
try {
|
||||
estimateWebGpuBatchSize({
|
||||
workload,
|
||||
limits,
|
||||
maximumGpuWorkingBytes: 4_096,
|
||||
maximumHostWorkingBytes: 4_096,
|
||||
maximumMacsPerInvocation: 7,
|
||||
});
|
||||
} catch (error) {
|
||||
invocationLimitError = error;
|
||||
}
|
||||
assert(invocationLimitError instanceof WebGpuLimitError, "per-invocation watchdog guard error type");
|
||||
|
||||
const scalingWorkload = {
|
||||
selectedBins: [10, 11],
|
||||
bandBins: [[10, 11]],
|
||||
frameCount: 46,
|
||||
microphoneCount: 192,
|
||||
pointCount: 40_000,
|
||||
distanceBytes: 40_000 * 192 * 4,
|
||||
outputBytes: 40_000 * 4 * 2,
|
||||
spectrumBytesPerBin: 46 * 192 * 8,
|
||||
baseWorkingBytes: 40_000 * 192 * 4 + 40_000 * 4 * 3 + 512 * 8 * 2,
|
||||
};
|
||||
const scalingPlan = estimateWebGpuBatchPlan({
|
||||
workload: scalingWorkload,
|
||||
limits: {
|
||||
maxStorageBufferBindingSize: 128 * 1024 ** 2,
|
||||
maxBufferSize: 256 * 1024 ** 2,
|
||||
maxComputeWorkgroupsPerDimension: 65_535,
|
||||
maxComputeInvocationsPerWorkgroup: 256,
|
||||
maxComputeWorkgroupSizeX: 256,
|
||||
maxStorageBuffersPerShaderStage: 8,
|
||||
maxUniformBufferBindingSize: 65_536,
|
||||
},
|
||||
});
|
||||
assertEqual(scalingPlan.binBatchSize, 2, "200x200 preserves one CPU STFT bin batch");
|
||||
assert(scalingPlan.pointBatchSize < 40_000, "200x200 steering uses point chunks");
|
||||
assert(scalingPlan.steeringBytes <= 64 * 1024 ** 2, "steering target is respected");
|
||||
});
|
||||
|
||||
test("backend auto는 typed WebGPU 운용 오류만 CPU로 fallback한다", async () => {
|
||||
const calls = [];
|
||||
const options = { experiment: { beamforming: { backend: "auto" } } };
|
||||
const cpuResult = { backend: "cpu", bands: [], marker: "fallback-result" };
|
||||
const result = await computeDasWithBackend(options, {
|
||||
computeGpu: async () => {
|
||||
calls.push("gpu");
|
||||
throw new WebGpuUnavailableError("adapter unavailable");
|
||||
},
|
||||
computeCpu: async () => {
|
||||
calls.push("cpu");
|
||||
return cpuResult;
|
||||
},
|
||||
});
|
||||
|
||||
assertEqual(calls.join(","), "gpu,cpu", "auto backend order");
|
||||
assertEqual(result.backend, "cpu", "fallback backend");
|
||||
assertEqual(result.marker, "fallback-result", "CPU result is preserved");
|
||||
|
||||
const limitFallback = await computeDasWithBackend(options, {
|
||||
computeGpu: async () => {
|
||||
throw new WebGpuLimitError("device limit");
|
||||
},
|
||||
computeCpu: async () => cpuResult,
|
||||
});
|
||||
assertEqual(limitFallback.backend, "cpu", "device limit also uses CPU fallback in auto mode");
|
||||
|
||||
let programmingError;
|
||||
try {
|
||||
await computeDasWithBackend(options, {
|
||||
computeGpu: async () => {
|
||||
throw new TypeError("programming error");
|
||||
},
|
||||
computeCpu: async () => {
|
||||
throw new Error("CPU must not run");
|
||||
},
|
||||
});
|
||||
} catch (error) {
|
||||
programmingError = error;
|
||||
}
|
||||
assert(programmingError instanceof TypeError, "unexpected GPU errors are not hidden by fallback");
|
||||
|
||||
const controller = new AbortController();
|
||||
controller.abort();
|
||||
let abortError;
|
||||
try {
|
||||
await computeDasWithBackend(
|
||||
{ ...options, signal: controller.signal },
|
||||
{
|
||||
computeGpu: async () => {
|
||||
throw new WebGpuUnavailableError("adapter unavailable after cancellation");
|
||||
},
|
||||
computeCpu: async () => {
|
||||
throw new Error("cancelled auto mode must not start CPU fallback");
|
||||
},
|
||||
},
|
||||
);
|
||||
} catch (error) {
|
||||
abortError = error;
|
||||
}
|
||||
assertEqual(abortError?.name, "AbortError", "cancellation wins over GPU fallback");
|
||||
|
||||
const explicitOptions = { experiment: { beamforming: { backend: "webgpu" } } };
|
||||
let explicitError;
|
||||
try {
|
||||
await computeDasWithBackend(explicitOptions, {
|
||||
computeGpu: async () => {
|
||||
throw new WebGpuLimitError("forced device limit");
|
||||
},
|
||||
computeCpu: async () => {
|
||||
throw new Error("explicit webgpu must not run CPU fallback");
|
||||
},
|
||||
});
|
||||
} catch (error) {
|
||||
explicitError = error;
|
||||
}
|
||||
assert(explicitError instanceof WebGpuLimitError, "explicit WebGPU preserves operational error");
|
||||
});
|
||||
Reference in New Issue
Block a user