feat: add WebGPU DAS backend
This commit is contained in:
Codex
@@ -0,0 +1,414 @@
|
||||
/**
|
||||
* Two-pass frequency-domain delay-and-sum beamforming kernel.
|
||||
*
|
||||
* Dispatch shapes:
|
||||
* compute_steering = (ceil(microphoneCount / WORKGROUP_SIZE), pointCount, binCount)
|
||||
* compute_das = (ceil(pointCount / WORKGROUP_SIZE), binCount, 1)
|
||||
*
|
||||
* The first pass computes trigonometry once per (bin, point, microphone).
|
||||
* The second pass reuses those steering vectors for every STFT frame. Callers
|
||||
* can bound the intermediate buffer by submitting point/bin batches.
|
||||
*
|
||||
* Flat buffer layouts:
|
||||
* spectrum: [bin][frame][microphone][real, imaginary]
|
||||
* distances: [point][microphone]
|
||||
* calibration: [microphone][inverse sensitivity gain, phase offset radians]
|
||||
* steering: [bin][point][microphone][real, imaginary]
|
||||
* output: [bin][point]
|
||||
*/
|
||||
|
||||
export const WORKGROUP_SIZE = 64;
|
||||
export const DAS_WORKGROUP_SIZE = WORKGROUP_SIZE;
|
||||
export const MAX_MICROPHONES = 1024;
|
||||
|
||||
// WGSL uniform structs are rounded to a 16-byte boundary. The reserved word
|
||||
// keeps the JavaScript and shader representations explicitly 32 bytes.
|
||||
export const DAS_PARAMS_BYTE_LENGTH = 32;
|
||||
|
||||
export const DAS_PARAMS_LAYOUT = Object.freeze({
|
||||
byteLength: DAS_PARAMS_BYTE_LENGTH,
|
||||
pointCount: Object.freeze({ offset: 0, type: "u32" }),
|
||||
microphoneCount: Object.freeze({ offset: 4, type: "u32" }),
|
||||
frameCount: Object.freeze({ offset: 8, type: "u32" }),
|
||||
binCount: Object.freeze({ offset: 12, type: "u32" }),
|
||||
sampleRate: Object.freeze({ offset: 16, type: "f32" }),
|
||||
fftSize: Object.freeze({ offset: 20, type: "u32" }),
|
||||
speedOfSound: Object.freeze({ offset: 24, type: "f32" }),
|
||||
reserved: Object.freeze({ offset: 28, type: "u32" }),
|
||||
});
|
||||
|
||||
export const DAS_BUFFER_LAYOUT = Object.freeze({
|
||||
params: Object.freeze({ group: 0, binding: 0, access: "uniform", layout: DAS_PARAMS_LAYOUT }),
|
||||
binIndices: Object.freeze({ group: 0, binding: 1, access: "read", element: "u32" }),
|
||||
spectrum: Object.freeze({
|
||||
group: 0,
|
||||
binding: 2,
|
||||
access: "read",
|
||||
element: "f32",
|
||||
order: "bin-frame-microphone-complex",
|
||||
}),
|
||||
distances: Object.freeze({
|
||||
group: 0,
|
||||
binding: 3,
|
||||
access: "read",
|
||||
element: "f32",
|
||||
order: "point-microphone",
|
||||
}),
|
||||
calibration: Object.freeze({
|
||||
group: 0,
|
||||
binding: 4,
|
||||
access: "read",
|
||||
element: "vec2<f32>",
|
||||
order: "microphone-(inverseGain-phaseRadians)",
|
||||
}),
|
||||
steering: Object.freeze({
|
||||
group: 0,
|
||||
binding: 5,
|
||||
access: "read_write",
|
||||
element: "vec2<f32>",
|
||||
order: "bin-point-microphone-complex",
|
||||
}),
|
||||
output: Object.freeze({
|
||||
group: 0,
|
||||
binding: 6,
|
||||
access: "read_write",
|
||||
element: "f32",
|
||||
order: "bin-point",
|
||||
}),
|
||||
});
|
||||
|
||||
export const DAS_ENTRY_POINTS = Object.freeze({
|
||||
steering: "compute_steering",
|
||||
beamforming: "compute_das",
|
||||
});
|
||||
|
||||
function requireUint32(name, value) {
|
||||
if (!Number.isInteger(value) || value < 0 || value > 0xffff_ffff) {
|
||||
throw new RangeError(`${name} must be an unsigned 32-bit integer.`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function requireUint32Product(name, values) {
|
||||
let product = 1;
|
||||
for (const value of values) {
|
||||
if (value !== 0 && product > 0xffff_ffff / value) {
|
||||
throw new RangeError(`${name} exceeds the shader's u32 address space.`);
|
||||
}
|
||||
product *= value;
|
||||
}
|
||||
return product;
|
||||
}
|
||||
|
||||
function requirePositiveF32(name, value) {
|
||||
const rounded = Math.fround(value);
|
||||
if (!(rounded > 0) || !Number.isFinite(rounded)) {
|
||||
throw new RangeError(`${name} must be a finite positive f32 value.`);
|
||||
}
|
||||
return rounded;
|
||||
}
|
||||
|
||||
function requireFiniteF32(name, value) {
|
||||
const rounded = Math.fround(value);
|
||||
if (!Number.isFinite(rounded)) {
|
||||
throw new RangeError(`${name} must be a finite f32 value.`);
|
||||
}
|
||||
return rounded;
|
||||
}
|
||||
|
||||
function requireFiniteNumber(name, value) {
|
||||
if (!Number.isFinite(value)) throw new RangeError(`${name} must be finite.`);
|
||||
return value;
|
||||
}
|
||||
|
||||
/** Pack the binding-0 uniform using the offsets in DAS_PARAMS_LAYOUT. */
|
||||
export function packDasParams(options) {
|
||||
const {
|
||||
pointCount,
|
||||
microphoneCount,
|
||||
frameCount,
|
||||
binCount,
|
||||
sampleRate,
|
||||
fftSize,
|
||||
speedOfSound,
|
||||
} = options;
|
||||
const counts = {
|
||||
pointCount: requireUint32("pointCount", pointCount),
|
||||
microphoneCount: requireUint32("microphoneCount", microphoneCount),
|
||||
frameCount: requireUint32("frameCount", frameCount),
|
||||
binCount: requireUint32("binCount", binCount),
|
||||
fftSize: requireUint32("fftSize", fftSize),
|
||||
};
|
||||
if (counts.microphoneCount > MAX_MICROPHONES) {
|
||||
throw new RangeError(`microphoneCount must not exceed ${MAX_MICROPHONES}.`);
|
||||
}
|
||||
if (counts.fftSize === 0) throw new RangeError("fftSize must be greater than zero.");
|
||||
|
||||
requireUint32Product("distance buffer element count", [
|
||||
counts.pointCount,
|
||||
counts.microphoneCount,
|
||||
]);
|
||||
requireUint32Product("output buffer element count", [counts.binCount, counts.pointCount]);
|
||||
requireUint32Product("steering buffer element count", [
|
||||
counts.binCount,
|
||||
counts.pointCount,
|
||||
counts.microphoneCount,
|
||||
]);
|
||||
requireUint32Product("spectrum buffer element count", [
|
||||
counts.binCount,
|
||||
counts.frameCount,
|
||||
counts.microphoneCount,
|
||||
2,
|
||||
]);
|
||||
|
||||
const packedSampleRate = requirePositiveF32("sampleRate", sampleRate);
|
||||
const packedSpeedOfSound = requirePositiveF32("speedOfSound", speedOfSound);
|
||||
const buffer = new ArrayBuffer(DAS_PARAMS_BYTE_LENGTH);
|
||||
const view = new DataView(buffer);
|
||||
view.setUint32(0, counts.pointCount, true);
|
||||
view.setUint32(4, counts.microphoneCount, true);
|
||||
view.setUint32(8, counts.frameCount, true);
|
||||
view.setUint32(12, counts.binCount, true);
|
||||
view.setFloat32(16, packedSampleRate, true);
|
||||
view.setUint32(20, counts.fftSize, true);
|
||||
view.setFloat32(24, packedSpeedOfSound, true);
|
||||
view.setUint32(28, 0, true);
|
||||
return buffer;
|
||||
}
|
||||
|
||||
/**
|
||||
* Pack already-linear calibration as tightly packed vec2<f32> entries.
|
||||
* phase is the microphone phase offset in radians. The shader subtracts it
|
||||
* from the propagation phase, exactly as the CPU reference does.
|
||||
*/
|
||||
export function packDasCalibration(calibration) {
|
||||
if (!Array.isArray(calibration)) throw new TypeError("calibration must be an array.");
|
||||
const packed = new Float32Array(calibration.length * 2);
|
||||
for (let index = 0; index < calibration.length; index += 1) {
|
||||
const entry = calibration[index];
|
||||
const gain = requireFiniteF32(`calibration[${index}].gain`, entry?.gain);
|
||||
if (gain < 0) throw new RangeError(`calibration[${index}].gain must not be negative.`);
|
||||
packed[index * 2] = gain;
|
||||
packed[index * 2 + 1] = requireFiniteF32(
|
||||
`calibration[${index}].phase`,
|
||||
entry?.phase,
|
||||
);
|
||||
}
|
||||
return packed;
|
||||
}
|
||||
|
||||
/** Convert array-definition microphone metadata to the shader calibration layout. */
|
||||
export function packMicrophoneCalibration(microphones) {
|
||||
if (!Array.isArray(microphones)) throw new TypeError("microphones must be an array.");
|
||||
const packed = new Float32Array(microphones.length * 2);
|
||||
for (let index = 0; index < microphones.length; index += 1) {
|
||||
const microphone = microphones[index];
|
||||
const sensitivityDb = requireFiniteNumber(
|
||||
`microphones[${index}].sensitivityDb`,
|
||||
microphone?.sensitivityDb ?? 0,
|
||||
);
|
||||
const phaseOffsetDeg = requireFiniteNumber(
|
||||
`microphones[${index}].phaseOffsetDeg`,
|
||||
microphone?.phaseOffsetDeg ?? 0,
|
||||
);
|
||||
packed[index * 2] = requireFiniteF32(
|
||||
`microphones[${index}] inverse sensitivity gain`,
|
||||
10 ** (-sensitivityDb / 20),
|
||||
);
|
||||
packed[index * 2 + 1] = requireFiniteF32(
|
||||
`microphones[${index}] phase offset`,
|
||||
(phaseOffsetDeg * Math.PI) / 180,
|
||||
);
|
||||
}
|
||||
return packed;
|
||||
}
|
||||
|
||||
export const DAS_WGSL = /* wgsl */ `
|
||||
struct DasParams {
|
||||
point_count: u32,
|
||||
microphone_count: u32,
|
||||
frame_count: u32,
|
||||
bin_count: u32,
|
||||
sample_rate: f32,
|
||||
fft_size: u32,
|
||||
speed_of_sound: f32,
|
||||
reserved: u32,
|
||||
}
|
||||
|
||||
struct U32Buffer {
|
||||
values: array<u32>,
|
||||
}
|
||||
|
||||
struct F32Buffer {
|
||||
values: array<f32>,
|
||||
}
|
||||
|
||||
struct Vec2Buffer {
|
||||
values: array<vec2<f32>>,
|
||||
}
|
||||
|
||||
@group(0) @binding(0) var<uniform> params: DasParams;
|
||||
@group(0) @binding(1) var<storage, read> bin_indices: U32Buffer;
|
||||
@group(0) @binding(2) var<storage, read> spectrum: F32Buffer;
|
||||
@group(0) @binding(3) var<storage, read> distances: F32Buffer;
|
||||
@group(0) @binding(4) var<storage, read> calibration: Vec2Buffer;
|
||||
@group(0) @binding(5) var<storage, read_write> steering_vectors: Vec2Buffer;
|
||||
@group(0) @binding(6) var<storage, read_write> output_maps: F32Buffer;
|
||||
|
||||
const TAU: f32 = 6.283185307179586;
|
||||
const MAX_MICROPHONES: u32 = ${MAX_MICROPHONES}u;
|
||||
const MAX_U32: u32 = 0xffffffffu;
|
||||
|
||||
fn multiplication_overflows(left: u32, right: u32) -> bool {
|
||||
return left != 0u && right > MAX_U32 / left;
|
||||
}
|
||||
|
||||
fn wrap_phase(phase: f32) -> f32 {
|
||||
return phase - round(phase / TAU) * TAU;
|
||||
}
|
||||
|
||||
@compute @workgroup_size(${WORKGROUP_SIZE}, 1, 1)
|
||||
fn compute_steering(@builtin(global_invocation_id) invocation: vec3<u32>) {
|
||||
let microphone_index = invocation.x;
|
||||
let point_index = invocation.y;
|
||||
let bin_index = invocation.z;
|
||||
|
||||
if (
|
||||
microphone_index >= params.microphone_count ||
|
||||
point_index >= params.point_count ||
|
||||
bin_index >= params.bin_count
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (
|
||||
params.microphone_count == 0u ||
|
||||
params.microphone_count > MAX_MICROPHONES ||
|
||||
params.fft_size == 0u ||
|
||||
params.sample_rate <= 0.0 ||
|
||||
params.speed_of_sound <= 0.0 ||
|
||||
multiplication_overflows(params.point_count, params.microphone_count) ||
|
||||
multiplication_overflows(params.bin_count, params.point_count)
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
let bin_point_count = params.bin_count * params.point_count;
|
||||
if (multiplication_overflows(bin_point_count, params.microphone_count)) {
|
||||
return;
|
||||
}
|
||||
|
||||
let required_distance_values = params.point_count * params.microphone_count;
|
||||
let required_steering_values = bin_point_count * params.microphone_count;
|
||||
if (
|
||||
arrayLength(&bin_indices.values) < params.bin_count ||
|
||||
arrayLength(&calibration.values) < params.microphone_count ||
|
||||
arrayLength(&distances.values) < required_distance_values ||
|
||||
arrayLength(&steering_vectors.values) < required_steering_values
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
let bin = f32(bin_indices.values[bin_index]);
|
||||
let angular_wave_number =
|
||||
(TAU * bin * params.sample_rate) /
|
||||
(f32(params.fft_size) * params.speed_of_sound);
|
||||
let distance_index = point_index * params.microphone_count + microphone_index;
|
||||
let steering_index =
|
||||
(bin_index * params.point_count + point_index) * params.microphone_count +
|
||||
microphone_index;
|
||||
let microphone_calibration = calibration.values[microphone_index];
|
||||
let phase = wrap_phase(
|
||||
angular_wave_number * distances.values[distance_index] -
|
||||
microphone_calibration.y
|
||||
);
|
||||
steering_vectors.values[steering_index] =
|
||||
vec2<f32>(cos(phase), sin(phase)) * microphone_calibration.x;
|
||||
}
|
||||
|
||||
@compute @workgroup_size(${WORKGROUP_SIZE}, 1, 1)
|
||||
fn compute_das(@builtin(global_invocation_id) invocation: vec3<u32>) {
|
||||
let point_index = invocation.x;
|
||||
let bin_index = invocation.y;
|
||||
|
||||
if (point_index >= params.point_count || bin_index >= params.bin_count) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (
|
||||
params.microphone_count == 0u ||
|
||||
params.microphone_count > MAX_MICROPHONES ||
|
||||
params.frame_count == 0u ||
|
||||
multiplication_overflows(params.bin_count, params.point_count) ||
|
||||
multiplication_overflows(params.bin_count, params.frame_count)
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
let bin_point_count = params.bin_count * params.point_count;
|
||||
let bin_frame_count = params.bin_count * params.frame_count;
|
||||
if (
|
||||
multiplication_overflows(bin_point_count, params.microphone_count) ||
|
||||
multiplication_overflows(bin_frame_count, params.microphone_count)
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
let spectrum_complex_count = bin_frame_count * params.microphone_count;
|
||||
if (multiplication_overflows(spectrum_complex_count, 2u)) {
|
||||
return;
|
||||
}
|
||||
|
||||
let required_output_values = bin_point_count;
|
||||
let required_steering_values = bin_point_count * params.microphone_count;
|
||||
let required_spectrum_values = spectrum_complex_count * 2u;
|
||||
if (
|
||||
arrayLength(&output_maps.values) < required_output_values ||
|
||||
arrayLength(&steering_vectors.values) < required_steering_values ||
|
||||
arrayLength(&spectrum.values) < required_spectrum_values
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
let microphone_scale = 1.0 / f32(params.microphone_count);
|
||||
let steering_base =
|
||||
(bin_index * params.point_count + point_index) * params.microphone_count;
|
||||
var accumulated_power = 0.0;
|
||||
|
||||
for (var frame_index = 0u; frame_index < params.frame_count; frame_index += 1u) {
|
||||
var sum = vec2<f32>(0.0, 0.0);
|
||||
for (
|
||||
var microphone_index = 0u;
|
||||
microphone_index < params.microphone_count;
|
||||
microphone_index += 1u
|
||||
) {
|
||||
let complex_index = (
|
||||
(
|
||||
(bin_index * params.frame_count + frame_index) * params.microphone_count +
|
||||
microphone_index
|
||||
) * 2u
|
||||
);
|
||||
let input = vec2<f32>(
|
||||
spectrum.values[complex_index],
|
||||
spectrum.values[complex_index + 1u]
|
||||
);
|
||||
let steering = steering_vectors.values[steering_base + microphone_index];
|
||||
|
||||
// (input.real + i*input.imaginary) *
|
||||
// (steering.real + i*steering.imaginary)
|
||||
sum += vec2<f32>(
|
||||
input.x * steering.x - input.y * steering.y,
|
||||
input.x * steering.y + input.y * steering.x
|
||||
);
|
||||
}
|
||||
|
||||
accumulated_power +=
|
||||
dot(sum, sum) * microphone_scale * microphone_scale;
|
||||
}
|
||||
|
||||
let output_index = bin_index * params.point_count + point_index;
|
||||
output_maps.values[output_index] = accumulated_power / f32(params.frame_count);
|
||||
}
|
||||
`;
|
||||
|
||||
export const DAS_SHADER_SOURCE = DAS_WGSL;
|
||||
Reference in New Issue
Block a user