Skip to content

performance-optimizations.md

Latest commit

 

History

History
262 lines (182 loc) · 13.8 KB

performance-optimizations.md

File metadata and controls

262 lines (182 loc) · 13.8 KB
Raw

DeepFilterNet-mlx Performance Optimizations

Summary

This document describes the performance optimizations implemented in DeepFilterNet-mlx that bring offline processing of a 52-second audio file down to 0.47s (112x real-time, Hybrid configuration) and streaming to ~2ms/hop (4.9x real-time) on Apple Silicon.

Baseline

The starting point was a straightforward MLX port of the PyTorch DeepFilterNet3 model: standard NCHW tensor layout, no weight caching, no kernel fusion, and a naive per-timestep GPU GRU. Offline processing ran at ~33x real-time (1.55s for 52s audio) and streaming at ~3ms/hop.

Optimization 1: Pre-computed Weight Caches

Problem: Every conv2d call transposed weights from PyTorch's OIHW to MLX's OHWI format. Every batch norm recomputed scale = gamma / sqrt(var + eps) and shift = beta - mean * scale. Every GRU step transposed weight matrices. These are constant — they never change after model load.

Solution: At model init time, pre-compute and cache all derived weight forms:

  • Conv2d weights: conv2dWeightsOHWI — one-time OIHW→OHWI transpose for all conv layers (DeepFilterNetModel.swift:103)
  • Batch norm affine: bnScale / bnBias — pre-fused gamma/sqrt(var+eps) and beta - mean*scale stored as [1,C,1,1] tensors (DeepFilterNetModel.swift:100-101)
  • GRU transposed weights: gruTransposedWeights — W.T computed once for all GRU layers (DeepFilterNetModel.swift:109)
  • ConvTranspose dense weights: convTransposeDenseWeights — pre-transposed for the decoder upsampling layers (DeepFilterNetModel.swift:104-107)

Impact: Eliminates redundant per-forward-pass work. Every layer lookup is now a dictionary hit returning a ready-to-use tensor.

Optimization 2: NHWC Memory Layout for Streaming Encoder

Problem: MLX's conv2d operates on NHWC layout. The standard BCHW encoder path requires transpose(0,2,3,1) before each conv and transpose(0,3,1,2) after — two layout shuffles per layer, adding GPU dispatches for data that never needs to be in BCHW.

Solution: The streaming encoder operates entirely in NHWC. Each StreamConvLayer stores both BCHW and NHWC batch norm parameters:

struct StreamConvLayer {
    let bnScale: MLXArray      // [1,C,1,1] for BCHW path
    let bnScaleNHWC: MLXArray  // [1,1,1,C] for NHWC path
    ...
}

The streaming hot path (applyConvLastNHWC) passes NHWC tensors directly through the 6-layer encoder chain (4 ERB conv + 2 DF conv) without any intermediate transposes. The single BCHW→NHWC conversion happens once at the entry point.

Impact: Eliminates 12+ transpose operations per hop in the encoder, saving ~6 GPU dispatches.

Optimization 3: Metal Kernel Fusions

Problem: DeepFilterNet's inference involves many small element-wise operations (GRU gates, mask multiply, ERB inverse, deep filtering). Each operation is a separate GPU dispatch with ~10-15µs overhead. For a ~200-operation-per-hop model, dispatch overhead alone costs ~2-3ms.

Solution: 8 custom Metal kernels via MLXFast.metalKernel() fuse multi-step computations into single GPU dispatches:

Kernel Operations Fused Dispatches Saved
dfn_gru_gate_step sigmoid(xr+hr), sigmoid(xz+hz), tanh(xn+r*hn), h update 3→1 per GRU step
dfn_mask_multiply spec * gains with broadcasting over real/imag 2→1
dfn_mask_erbinv_apply ERB inverse filterbank matmul + mask multiply ~3→1
dfn_stream_df_frame Complex multiply-accumulate over dfOrder=5 5→1 per streaming hop
dfn_stream_df_frame_packed Same with packed coefficient layout 5→1
dfn_offline_df Batch deep filter with padding logic ~T*5→1 for offline

Each kernel has FP32 and FP16 variants. The ensureRowContiguous flag controls whether MLX should force contiguous layout before dispatch.

GRU Gate Fusion (most impactful for streaming): Each GRU step normally requires 6 slices, 3 additions, 2 sigmoids, 1 tanh, 2 multiplies, and 1 subtraction — ~15 separate GPU dispatches. The fused kernel computes the entire gate update in one dispatch:

float r = 1.0f / (1.0f + exp(-(xr + hr)));
float z = 1.0f / (1.0f + exp(-(xz + hz)));
float n = tanh(xn + r * hn);
out[idx] = (1.0f - z) * n + z * prevh;

With 3 GRU layers × 1 step per streaming hop, this saves ~42 dispatches per hop.

Impact: Reduces streaming per-hop time by ~0.5-1ms. The fused GRU gates alone eliminate the largest dispatch concentration in the inference path.

Optimization 4: Compiled Graph Caching

Problem: MLX's lazy evaluation builds a computation graph each hop. For fixed-shape operations (FFT, synthesis, decoder conv chains), the graph structure is identical every time but gets rebuilt from scratch.

Solution: Pre-compile fixed-shape subgraphs using MLX.compile(shapeless: false) at streamer init time:

  • compiledAnalysisFeatureStep — FFT + ERB features + DF features (fused analysis+features)
  • compiledSynthesisStep — iFFT + overlap-add synthesis
  • compiledStreamErbDecoderStep — Full ERB decoder (8 conv layers + BN + skip connections)
  • compiledStreamDfConvpStep — DF pathway conv layers
  • compiledStreamInferAssignStep — Deep filter application + spectral assignment
  • Per-GRU-layer step caches via compiledStepCache

These compiled functions capture all weight references at compile time and execute as single graph evaluations.

Impact: Reduces graph construction overhead and enables MLX's internal optimizations (operation fusion, memory planning). The ERB decoder compilation is especially valuable — it fuses 8 conv + BN + relu operations into one call.

Optimization 5: Tensor Ring Buffers

Problem: Streaming requires maintaining history of past frames (spectra, encoder features, DF coefficients). Naive approaches use array concatenation + slicing each hop, which allocates new arrays and copies data every step.

Solution: TensorRingBuffer — a fixed-capacity circular buffer that overwrites the oldest entry on each push:

final class TensorRingBuffer {
    private(set) var values: [MLXArray]
    private(set) var totalWritten: Int = 0

    func push(_ value: MLXArray) {
        values[totalWritten % capacity] = value
        totalWritten += 1
    }

    func get(absoluteIndex: Int) -> MLXArray? { ... }
}

Five ring buffers track streaming state: spec (full spectrum), specLow (low-frequency bins), encErb (ERB encoder features), encDf (DF encoder features), dfConvp (DF convolution pathway output).

Impact: Zero allocation per hop for frame history management. Lookups by absolute index are O(1).

Optimization 6: Accelerate-Optimized GRU (Offline Only)

Problem: The GRU is the offline bottleneck. With 5,212 frames across 3 GRU layers, the sequential recurrence requires ~15,000 tiny GPU matmuls (h @ W_hh) that cannot be parallelized — each timestep depends on the previous hidden state. GPU dispatch overhead per matmul (~10-15µs) dominates the actual compute.

Solution: Ported from mlx-audio-swift (commit f2c6ff2). Split the GRU into two phases:

  1. Batch input projection on GPU: gxAll = x @ W_ih + b_ih — a single large matmul across all timesteps
  2. Per-step hidden projection on CPU: gh = h @ W_hh via vDSP_mmul with scalar GRU gate computation
// Single GPU matmul for all timesteps
let gxAllMLX = (MLX.matmul(x2D, wihT) + bih).reshaped([batchSize, t, h3])
eval(gxAllMLX)

// Extract to CPU, run per-step loop with vDSP
let gxAll = gxAllMLX.asType(.float32).reshaped([-1]).asArray(Float.self)
for ti in 0..<t {
    vDSP_mmul(state + stOff, 1, wHH, 1, ghBuf, 1, 1, h3, hiddenSize)
    // Scalar gate computation: r, z, n, h_new
}

Why not for streaming: Each streaming hop has only 1 GRU timestep. The eval() + asArray() CPU↔GPU transfer overhead (~2ms) exceeds the GPU dispatch savings. Streaming keeps the Metal fused GRU gates (Optimization 3).

Impact: 3.3x offline speedup — from 1.55s (33.5x RT) to 0.47s (112x RT) for 52s audio.

Optimization 7: Materialization Interval Tuning

Problem: MLX's lazy evaluation accumulates computation graphs across hops. Without periodic eval(), the graph grows unboundedly, eventually causing memory pressure and slow graph traversal. Too-frequent evaluation wastes time on synchronization.

Solution: Periodic eval() every N hops, configured via materializeEveryHops (default: 512, tuned from initial 96):

hopsSinceMaterialize += 1
if config.materializeEveryHops > 0,
   hopsSinceMaterialize >= config.materializeEveryHops {
    materializeStreamingState(output: out)
    hopsSinceMaterialize = 0
}

materializeStreamingState() evaluates the current output plus all ring buffer contents and recurrent states, ensuring no lazy graph spans more than 512 hops.

Impact: Reduces amortized overhead from ~0.5ms/hop (at 96) to ~0.1ms/hop (at 512). For short audio (10s), this avoids the materialization spike entirely.

Optimization 8: Streaming Encoder Sequence State

Problem: The encoder conv layers have temporal kernel size 3, requiring the last 3 frames of features. Naively stacking ring buffer entries each hop creates fresh tensors.

Solution: Maintain rolling sequence states (encErbSeqState, encDfSeqState, dfConvpSeqState) that are updated incrementally:

static func appendTimeFrame(_ seq: MLXArray, frame: MLXArray) -> MLXArray {
    let t = seq.shape[2]
    return MLX.concatenated([seq[0..., 0..., 1..<t, 0...], frame], axis: 2)
}

The conv layers receive these pre-built sequences directly and extract only the last output frame, avoiding redundant computation over already-processed frames.

Optimization 9: Deep Filter Vectorization (Offline)

Problem: The offline deep filter applies complex multiply-accumulate over dfOrder=5 taps across all frequency bins and timesteps. A naive loop over time creates thousands of small operations.

Solution: The Metal kernel dfn_offline_df parallelizes across all (batch, time, freq) dimensions in a single dispatch. Each thread handles one (b, t, f) position, looping only over the 5 filter taps:

for (uint k = 0; k < K; ++k) {
    int srcT = int(t) + int(k) - pad_left;
    if (srcT < 0 || srcT >= int(T)) continue;
    // Complex multiply-accumulate
    outR += sr * cr - si * ci;
    outI += sr * ci + si * cr;
}

Impact: Replaces T×F×K individual operations with a single GPU dispatch for the entire batch.

Optimization 10: Pre-extracted CPU Weight Arrays

Problem: If the Accelerate GRU path called asArray(Float.self) on weight tensors every forward pass, it would add ~1ms of GPU→CPU transfer per GRU layer.

Solution: StreamGRULayer pre-extracts all weight arrays to CPU [Float] buffers at init time:

init(wihT: MLXArray, whhT: MLXArray, bih: MLXArray, bhh: MLXArray) {
    eval(wihT, whhT, bih, bhh)
    self.cpuWIH = wihT.asType(.float32).reshaped([-1]).asArray(Float.self)
    self.cpuWHH = whhT.asType(.float32).reshaped([-1]).asArray(Float.self)
    self.cpuBiasIH = bih.asType(.float32).reshaped([-1]).asArray(Float.self)
    self.cpuBiasHH = bhh.asType(.float32).reshaped([-1]).asArray(Float.self)
}

Impact: Zero per-forward-pass weight transfer cost. The only runtime CPU↔GPU transfer is the input projection result and the final output.

Performance Configuration

All optimizations are controlled via DeepFilterNetPerformanceConfig:

public struct DeepFilterNetPerformanceConfig {
    var enableMetalFusedMaskMultiply: Bool      // Optimization 3
    var enableMetalFusedErbInvMaskApply: Bool    // Optimization 3
    var enableMetalFusedStreamingDeepFilter: Bool // Optimization 3
    var enableMetalFusedOfflineDeepFilter: Bool  // Optimization 3
    var enableMetalFusedGRUGates: Bool           // Optimization 3
    var enableAccelerateGRU: Bool                // Optimization 6
    var preferCompiledGraphs: Bool               // Optimization 4
    var ensureKernelContiguousInputs: Bool       // Metal kernel layout
    var kernelThreadGroupSize: Int               // Metal dispatch tuning
}

Two presets are available:

  • throughput: All optimizations enabled (production default)
  • safe: All fusions disabled, for debugging/validation

Results Summary

Metric Before Optimizations After Optimizations Speedup
Offline (52s audio) 1.55s (33.5x RT) 0.47s / 0.55s (112x / 95x RT) 3.3x / 2.8x
Streaming (ms/hop) ~3.1ms ~1.98ms 1.6x
Streaming RT factor ~3.2x 4.9x 1.5x

Offline: Hybrid achieves 0.47s (112x RT), MLX GPU throughput achieves 0.55s (95x RT). Both use the Accelerate GRU; the Hybrid configuration gains an additional edge from compiled conv graphs.

Optimization Contribution Breakdown (Offline)

Optimization Approximate Impact
Accelerate GRU (Opt 6) 3.3x (dominant — eliminates sequential GPU dispatches)
Weight caching (Opt 1) ~5-10% (eliminates per-pass transpose/compute)
Metal kernel fusions (Opt 3) ~5% offline, ~20% streaming
Compiled graphs (Opt 4) ~5% (graph construction overhead)
NHWC layout (Opt 2) ~3% streaming (fewer transposes)
Ring buffers (Opt 5) ~2% streaming (allocation elimination)
Materialization tuning (Opt 7) Variable (prevents graph explosion)

Output Quality

All optimizations maintain near-perfect output quality:

Comparison Correlation
Optimized MLX vs PyTorch reference 0.9999984
Accelerate GRU vs standard GRU 1.0000000 (bit-identical)

The Accelerate GRU produces bit-identical output to the standard MLX GRU path because both perform the exact same mathematical operations — the optimization only changes where (CPU vs GPU) and how (batch vs sequential) the computation runs.