Codex
415 lines
13 KiB
JavaScript
415 lines
13 KiB
JavaScript
/**
|
|
* 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;
|