[AUDIO_WORKLET] Optimised output buffer copy (#24891)

A reworking of #22753, which "improves the copy back from the audio
worklet's heap to JS by 7-12x depending on the browser." From the
previous description:

Since we pass in the stack for the worklet from the caller's heap, its
address doesn't change. And since the render quantum size doesn't change
after the audio worklet creation, the stack positions for the audio
buffers do not change either. This optimisation adds one-time subarray
views and replaces the float-by-float copy with a simple `set()` per
channel (per output).

The existing interactive tests (written for the original PR) can be run
for comparison:
```
test/runner interactive.test_audio_worklet_stereo_io
test/runner interactive.test_audio_worklet_2x_stereo_io
test/runner interactive.test_audio_worklet_mono_io
test/runner interactive.test_audio_worklet_2x_hard_pan_io
test/runner interactive.test_audio_worklet_params_mixing
test/runner interactive.test_audio_worklet_memory_growth
test/runner interactive.test_audio_worklet_hard_pans
```
These test various input/output arrangements as well as parameters
(parameters are interesting because, depending on the browser, the sizes
change as the params move from static to varying).

The original benchmark of the extracted copy is still valid:

https://wip.numfum.com/cw/2024-10-29/index.html

This is tested with 32- and 64-bit wasm (which required a reordering of
how structs and data were stored to avoid alignment issues).

Some explanations:
- Fixed-position output buffer views are created once in
the`WasmAudioWorkletProcessor` constructor
- Stack allocations for the `process()` call are split into aligned
struct data (see the comments) and audio/param data
- The struct writes are simplified by this splitting of data
- `ASSERTIONS` are used to ensure everything fits and correctly aligns
- The tests account for size changes in the params, which can vary from
a single float to 128 floats (a single float nicely showing up any
8-byte alignment issues for wasm64)

~~Future improvements: the output views are sequential, so instead of of
being individual views covering each channel the views could cover one
to however-many-views needed, with a single `set()` being enough for all
outputs.~~

Author: cwoffenden

src/audio_worklet.js

@@ -32,51 +32,120 @@ function createWasmAudioWorkletProcessor(audioParams) {
       this.callback = {{{ makeDynCall('iipipipp', 'opts.callback') }}};
       this.userData = opts.userData;
       // Then the samples per channel to process, fixed for the lifetime of the
-      // context that created this processor. Note for when moving to Web Audio
-      // 1.1: the typed array passed to process() should be the same size as this
-      // 'render quantum size', and this exercise of passing in the value
-      // shouldn't be required (to be verified)
+      // context that created this processor. Even though this 'render quantum
+      // size' is fixed at 128 samples in the 1.0 spec, it will be variable in
+      // the 1.1 spec. It's passed in now, just to prove it's settable, but will
+      // eventually be a property of the  AudioWorkletGlobalScope (globalThis).
       this.samplesPerChannel = opts.samplesPerChannel;
+      this.bytesPerChannel = this.samplesPerChannel * {{{ getNativeTypeSize('float') }}};
+
+      // Prepare the output views; see createOutputViews(). The 'minimum alloc'
+      // firstly stops STACK_OVERFLOW_CHECK failing (since the stack will be
+      // full if we allocate all the available space, with 16 bytes being the
+      // minimum allo size due to alignments) leaving room for a single
+      // AudioSampleFrame as a minumum. There's an arbitrary maximum of 64, for
+      // the case where a multi-MB stack is passed.
+      this.outputViews = new Array(Math.min(((wwParams.stackSize - /*minimum alloc*/ 16) / this.bytesPerChannel) | 0, /*sensible limit*/ 64));
+#if ASSERTIONS
+      console.assert(this.outputViews.length > 0, `AudioWorklet needs more stack allocating (at least ${this.bytesPerChannel})`);
+#endif
+      this.createOutputViews();
+
+#if ASSERTIONS
+      // Explicitly verify this later in process(). Note to self, stackSave is a
+      // bit of a misnomer as it simply gets the stack address.
+      this.ctorOldStackPtr = stackSave();
+#endif
+    }
+
+    /**
+     * Create up-front as many typed views for marshalling the output data as
+     * may be required, allocated at the *top* of the worklet's stack (and whose
+     * addresses are fixed). 
+     */
+    createOutputViews() {
+      // These are still alloc'd to take advantage of the overflow checks, etc.
+      var oldStackPtr = stackSave();
+      var viewDataIdx = {{{ getHeapOffset('stackAlloc(this.outputViews.length * this.bytesPerChannel)', 'float') }}};
+#if WEBAUDIO_DEBUG
+      console.log(`AudioWorklet creating ${this.outputViews.length} buffer one-time views (for a stack size of ${wwParams.stackSize} at address ${ptrToString(viewDataIdx * 4)})`);
+#endif
+      // Inserted in reverse so the lowest indices are closest to the stack top
+      for (var n = this.outputViews.length - 1; n >= 0; n--) {
+        this.outputViews[n] = HEAPF32.subarray(viewDataIdx, viewDataIdx += this.samplesPerChannel);
+      }
+      stackRestore(oldStackPtr);
     }
 
     static get parameterDescriptors() {
       return audioParams;
     }
 
     /**
+     * Marshals all inputs and parameters to the Wasm memory on the thread's
+     * stack, then performs the wasm audio worklet call, and finally marshals
+     * audio output data back.
+     *
      * @param {Object} parameters
      */
     process(inputList, outputList, parameters) {
-      // Marshal all inputs and parameters to the Wasm memory on the thread stack,
-      // then perform the wasm audio worklet call,
-      // and finally marshal audio output data back.
+#if ALLOW_MEMORY_GROWTH
+      // Recreate the output views if the heap has changed
+      // TODO: add support for GROWABLE_ARRAYBUFFERS
+      if (HEAPF32.buffer != this.outputViews[0].buffer) {
+        this.createOutputViews();
+      }
+#endif
 
       var numInputs = inputList.length;
       var numOutputs = outputList.length;
 
       var entry; // reused list entry or index
       var subentry; // reused channel or other array in each list entry or index
 
-      // Calculate how much stack space is needed.
-      var bytesPerChannel = this.samplesPerChannel * {{{ getNativeTypeSize('float') }}};
-      var stackMemoryNeeded = (numInputs + numOutputs) * {{{ C_STRUCTS.AudioSampleFrame.__size__ }}};
+      // Calculate the required stack and output buffer views (stack is further
+      // split into aligned structs and the raw float data).
+      var stackMemoryStruct = (numInputs + numOutputs) * {{{ C_STRUCTS.AudioSampleFrame.__size__ }}};
+      var stackMemoryData = 0;
+      for (entry of inputList) {
+        stackMemoryData += entry.length;
+      }
+      stackMemoryData *= this.bytesPerChannel;
+      // Collect the total number of output channels (mapped to array views)
+      var outputViewsNeeded = 0;
+      for (entry of outputList) {
+        outputViewsNeeded += entry.length;
+      }
+      stackMemoryData += outputViewsNeeded * this.bytesPerChannel;
       var numParams = 0;
-      for (entry of inputList) stackMemoryNeeded += entry.length * bytesPerChannel;
-      for (entry of outputList) stackMemoryNeeded += entry.length * bytesPerChannel;
       for (entry in parameters) {
-        stackMemoryNeeded += parameters[entry].byteLength + {{{ C_STRUCTS.AudioParamFrame.__size__ }}};
         ++numParams;
+        stackMemoryStruct += {{{ C_STRUCTS.AudioParamFrame.__size__ }}};
+        stackMemoryData += parameters[entry].byteLength;
       }
-
-      // Allocate the necessary stack space.
       var oldStackPtr = stackSave();
-      var inputsPtr = stackAlloc(stackMemoryNeeded);
+#if ASSERTIONS
+      console.assert(oldStackPtr == this.ctorOldStackPtr, 'AudioWorklet stack address has unexpectedly moved');
+      console.assert(outputViewsNeeded <= this.outputViews.length, `Too many AudioWorklet outputs (need ${outputViewsNeeded} but have stack space for ${this.outputViews.length})`);
+#endif
+
+      // Allocate the necessary stack space. All pointer variables are in bytes;
+      // 'structPtr' starts at the first struct entry (all run sequentially)
+      // and is the working start to each record; 'dataPtr' is the same for the
+      // audio/params data, starting after *all* the structs.
+      // 'structPtr' begins 16-byte aligned, allocated from the internal
+      // _emscripten_stack_alloc(), as are the output views, and so to ensure
+      // the views fall on the correct addresses (and we finish at stacktop) we
+      // request additional bytes, taking this alignment into account, then
+      // offset `dataPtr` by the difference.
+      var stackMemoryAligned = (stackMemoryStruct + stackMemoryData + 15) & ~15;
+      var structPtr = stackAlloc(stackMemoryAligned);
+      var dataPtr = structPtr + (stackMemoryAligned - stackMemoryData);
 
-      // Copy input audio descriptor structs and data to Wasm ('structPtr' is
-      // reused as the working start to each struct record, 'dataPtr' start of
-      // the data section, usually after all structs).
-      var structPtr = inputsPtr;
-      var dataPtr = inputsPtr + numInputs * {{{ C_STRUCTS.AudioSampleFrame.__size__ }}};
+      // Copy input audio descriptor structs and data to Wasm (recall, structs
+      // first, audio data after). 'inputsPtr' is the start of the C callback's
+      // input AudioSampleFrame.
+      var /*const*/ inputsPtr = structPtr;
       for (entry of inputList) {
         // Write the AudioSampleFrame struct instance
         {{{ makeSetValue('structPtr', C_STRUCTS.AudioSampleFrame.numberOfChannels, 'entry.length', 'u32') }}};
@@ -86,28 +155,13 @@ function createWasmAudioWorkletProcessor(audioParams) {
         // Marshal the input audio sample data for each audio channel of this input
         for (subentry of entry) {
           HEAPF32.set(subentry, {{{ getHeapOffset('dataPtr', 'float') }}});
-          dataPtr += bytesPerChannel;
+          dataPtr += this.bytesPerChannel;
         }
       }
 
-      // Copy output audio descriptor structs to Wasm
-      var outputsPtr = dataPtr;
-      structPtr = outputsPtr;
-      var outputDataPtr = (dataPtr += numOutputs * {{{ C_STRUCTS.AudioSampleFrame.__size__ }}});
-      for (entry of outputList) {
-        // Write the AudioSampleFrame struct instance
-        {{{ makeSetValue('structPtr', C_STRUCTS.AudioSampleFrame.numberOfChannels, 'entry.length', 'u32') }}};
-        {{{ makeSetValue('structPtr', C_STRUCTS.AudioSampleFrame.samplesPerChannel, 'this.samplesPerChannel', 'u32') }}};
-        {{{ makeSetValue('structPtr', C_STRUCTS.AudioSampleFrame.data, 'dataPtr', '*') }}};
-        structPtr += {{{ C_STRUCTS.AudioSampleFrame.__size__ }}};
-        // Reserve space for the output data
-        dataPtr += bytesPerChannel * entry.length;
-      }
-
-      // Copy parameters descriptor structs and data to Wasm
-      var paramsPtr = dataPtr;
-      structPtr = paramsPtr;
-      dataPtr += numParams * {{{ C_STRUCTS.AudioParamFrame.__size__ }}};
+      // Copy parameters descriptor structs and data to Wasm. 'paramsPtr' is the
+      // start of the C callback's input AudioParamFrame.
+      var /*const*/ paramsPtr = structPtr;
       for (entry = 0; subentry = parameters[entry++];) {
         // Write the AudioParamFrame struct instance
         {{{ makeSetValue('structPtr', C_STRUCTS.AudioParamFrame.length, 'subentry.length', 'u32') }}};
@@ -118,20 +172,54 @@ function createWasmAudioWorkletProcessor(audioParams) {
         dataPtr += subentry.length * {{{ getNativeTypeSize('float') }}};
       }
 
+      // Copy output audio descriptor structs to Wasm. 'outputsPtr' is the start
+      // of the C callback's output AudioSampleFrame. 'dataPtr' will now be
+      // aligned with the output views, ending at stacktop (which is why this
+      // needs to be last).
+      var /*const*/ outputsPtr = structPtr;
+      for (entry of outputList) {
+        // Write the AudioSampleFrame struct instance
+        {{{ makeSetValue('structPtr', C_STRUCTS.AudioSampleFrame.numberOfChannels, 'entry.length', 'u32') }}};
+        {{{ makeSetValue('structPtr', C_STRUCTS.AudioSampleFrame.samplesPerChannel, 'this.samplesPerChannel', 'u32') }}};
+        {{{ makeSetValue('structPtr', C_STRUCTS.AudioSampleFrame.data, 'dataPtr', '*') }}};
+        structPtr += {{{ C_STRUCTS.AudioSampleFrame.__size__ }}};
+        // Advance the output pointer to the next output (matching the pre-allocated views)
+        dataPtr += this.bytesPerChannel * entry.length;
+      }
+
+#if ASSERTIONS
+      // If all the maths worked out, we arrived at the original stack address
+      console.assert(dataPtr == oldStackPtr, `AudioWorklet stack missmatch (audio data finishes at ${dataPtr} instead of ${oldStackPtr})`);
+
+      // Sanity checks. If these trip the most likely cause, beyond unforeseen
+      // stack shenanigans, is that the 'render quantum size' changed after
+      // construction (which shouldn't be possible).
+      if (numOutputs) {
+        // First that the output view addresses match the stack positions
+        dataPtr -= this.bytesPerChannel;
+        for (entry = 0; entry < outputViewsNeeded; entry++) {
+          console.assert(dataPtr == this.outputViews[entry].byteOffset, 'AudioWorklet internal error in addresses of the output array views');
+          dataPtr -= this.bytesPerChannel;
+        }
+        // And that the views' size match the passed in output buffers
+        for (entry of outputList) {
+          for (subentry of entry) {
+            console.assert(subentry.byteLength == this.bytesPerChannel, `AudioWorklet unexpected output buffer size (expected ${this.bytesPerChannel} got ${subentry.byteLength})`);
+          }
+        }
+      }
+#endif
+
       // Call out to Wasm callback to perform audio processing
       var didProduceAudio = this.callback(numInputs, inputsPtr, numOutputs, outputsPtr, numParams, paramsPtr, this.userData);
       if (didProduceAudio) {
         // Read back the produced audio data to all outputs and their channels.
-        // (A garbage-free function TypedArray.copy(dstTypedArray, dstOffset,
-        // srcTypedArray, srcOffset, count) would sure be handy..  but web does
-        // not have one, so manually copy all bytes in)
-        outputDataPtr = {{{ getHeapOffset('outputDataPtr', 'float') }}};
+        // The preallocated 'outputViews' already have the correct offsets and
+        // sizes into the stack (recall from createOutputViews() that they run
+        // backwards).
         for (entry of outputList) {
           for (subentry of entry) {
-            // repurposing structPtr for now
-            for (structPtr = 0; structPtr < this.samplesPerChannel; ++structPtr) {
-              subentry[structPtr] = HEAPF32[outputDataPtr++];
-            }
+            subentry.set(this.outputViews[--outputViewsNeeded]);
           }
         }
       }