Reuse the encoder analysis scratch buffers

internal/celt/celt.go

@@ -15,4 +15,26 @@ const (
 	maxFrameSampleCount = shortBlockSampleCount << maxLM
 	maxBands            = 21
 	hybridStartBand     = 17
+	// maxBandSampleCount is the widest band — bands 20 and 21 span 22 edge
+	// units — at the longest frame, which bounds any per-band scratch.
+	maxBandSampleCount = 22 << maxLM
 )
+
+// encoderScratch holds the analysis path's working buffers. EncodeFrame runs 50
+// times a second per stream, so allocating these per frame is pure garbage —
+// they are sized for the worst case once and reused, the same way
+// decoderScratch works.
+type encoderScratch struct {
+	autocorr [pitchLPCOrder + 1]float32
+	lpc      [pitchLPCOrder]float32
+	// pitchSearch decimates by a further 2, so its buffers are a quarter of the
+	// window it is handed.
+	pitchX   [maxFrameSampleCount >> 2]float32
+	pitchY   [(maxFrameSampleCount + combFilterMaxPeriod) >> 2]float32
+	pitchXC  [combFilterMaxPeriod >> 1]float32
+	yyLookup [(combFilterMaxPeriod >> 1) + 1]float32
+	// tfAnalysis works one band at a time, so the widest band at the longest
+	// frame bounds both of its copies.
+	tfTmp    [maxBandSampleCount]float32
+	tfTmpOne [maxBandSampleCount]float32
+}

internal/celt/encoder.go

@@ -47,6 +47,8 @@ type Encoder struct {
 	normalizedBands   [2][]float32
 	pitchBuf          []float32
 	pitchChannels     [2][]float32
+	bandTmpScratch    []float32
+	scratch           encoderScratch
 	// rdoScratch and rdoState back the stereo-angle search. They live here so
 	// the per-frame band state can borrow them instead of allocating.
 	rdoScratch [4][]float32
@@ -138,6 +140,7 @@ func (e *Encoder) Reset() {
 		e.normalizedBands[ch] = make([]float32, 0, maxFrameSampleCount)
 	}
 	e.cwrsScratch = make([]uint32, 0, cwrsMaxPulseCount+2)
+	e.bandTmpScratch = make([]float32, 0, maxFrameSampleCount)
 	e.pitchBuf = make([]float32, 0, (combFilterMaxPeriod+maxFrameSampleCount)>>1)
 
 	// libopus seeds tonal_average at 256 (celt_encoder.c:3091), the midpoint
@@ -446,19 +449,19 @@ func (e *Encoder) choosePrefilter(
 
 	pitchLen := (combFilterMaxPeriod + frameSampleCount) >> 1
 	buf := slicetools.Resize(&e.pitchBuf, pitchLen)
-	pitchDownsample(pitchInput, buf, pitchLen, 2)
+	pitchDownsample(pitchInput, buf, pitchLen, 2, &e.scratch)
 
 	// The top 1.5 octave of the range is skipped: short-term correlation there
 	// produces too many false positives.
 	pitchPeriod := pitchSearch(
 		buf[combFilterMaxPeriod>>1:], buf,
-		frameSampleCount, combFilterMaxPeriod-3*combFilterMinPeriod,
+		frameSampleCount, combFilterMaxPeriod-3*combFilterMinPeriod, &e.scratch,
 	)
 	pitchPeriod = combFilterMaxPeriod - pitchPeriod
 
 	pitchGain := removeDoubling(
 		buf, combFilterMaxPeriod, combFilterMinPeriod, frameSampleCount,
-		&pitchPeriod, e.analysis.prefilter.period, e.analysis.prefilter.gain,
+		&pitchPeriod, e.analysis.prefilter.period, e.analysis.prefilter.gain, &e.scratch,
 	)
 	pitchPeriod = min(pitchPeriod, combFilterMaxPeriod-2)
 
@@ -812,6 +815,7 @@ func (e *Encoder) newBandState(
 		norm:           e.bandNorm[:0],
 		lowbandScratch: e.bandLowScratch[:0],
 		collapseMasks:  e.bandCollapseMasks[:0],
+		tmpScratch:     e.bandTmpScratch[:0],
 		// libopus only searches the stereo angle at complexity 8 and up: it
 		// encodes every stereo band twice, so the cost is real.
 		thetaRDO:   e.complexity >= thetaRDOComplexity && info.channelCount == 2,
@@ -936,6 +940,7 @@ func (e *Encoder) EncodeFrame(pcm [][]float32, dst []byte, frameBytes, startBand
 		info.tfSelect = tfAnalysis(
 			normalized[tfChan], info.lm, info.startBand, info.endBand, info.transient,
 			lambda, tfEstimate, 0, len(analysis.mdct[tfChan]), &dr.importance, &info.tfChange,
+			&e.scratch,
 		)
 	} else {
 		for band := info.startBand; band < info.endBand; band++ {

internal/celt/pitch_search.go

@@ -31,9 +31,9 @@ func celtFir5(x []float32, num [5]float32) {
 	}
 }
 
-// celtAutocorr returns lag+1 autocorrelation values of x.
-func celtAutocorr(x []float32, lag int) []float32 {
-	ac := make([]float32, lag+1)
+// celtAutocorr returns lag+1 autocorrelation values of x, written into ac.
+func celtAutocorr(x []float32, lag int, ac []float32) []float32 {
+	ac = ac[:lag+1]
 	for k := 0; k <= lag; k++ {
 		var d float64
 		for i := k; i < len(x); i++ {
@@ -47,8 +47,9 @@ func celtAutocorr(x []float32, lag int) []float32 {
 
 // celtLPC solves for p LPC coefficients by Levinson-Durbin recursion.
 // Port of libopus _celt_lpc (celt/celt_lpc.c).
-func celtLPC(ac []float32, p int) []float32 {
-	lpc := make([]float32, p)
+func celtLPC(ac []float32, p int, lpc []float32) []float32 {
+	lpc = lpc[:p]
+	clear(lpc)
 	if ac[0] <= 1e-10 {
 		return lpc
 	}
@@ -79,7 +80,7 @@ func celtLPC(ac []float32, p int) []float32 {
 
 // pitchDownsample decimates the channels by factor into xLP, sums them, then
 // whitens the result with a 4th-order LPC filter plus a fixed zero.
-func pitchDownsample(x [][]float32, xLP []float32, length, factor int) {
+func pitchDownsample(x [][]float32, xLP []float32, length, factor int, scratch *encoderScratch) {
 	offset := factor / 2
 	for i := 1; i < length; i++ {
 		xLP[i] = 0.25*x[0][factor*i-offset] + 0.25*x[0][factor*i+offset] + 0.5*x[0][factor*i]
@@ -92,15 +93,15 @@ func pitchDownsample(x [][]float32, xLP []float32, length, factor int) {
 		xLP[0] += 0.25*right[offset] + 0.5*right[0]
 	}
 
-	ac := celtAutocorr(xLP[:length], pitchLPCOrder)
+	ac := celtAutocorr(xLP[:length], pitchLPCOrder, scratch.autocorr[:])
 	// Noise floor at -40 dB, then lag windowing: both keep the LPC solve from
 	// chasing a near-singular autocorrelation on quiet or very tonal frames.
 	ac[0] *= 1.0001
 	for i := 1; i <= pitchLPCOrder; i++ {
 		ac[i] -= ac[i] * (0.008 * float32(i)) * (0.008 * float32(i))
 	}
 
-	lpc := celtLPC(ac, pitchLPCOrder)
+	lpc := celtLPC(ac, pitchLPCOrder, scratch.lpc[:])
 	tmp := float32(1.0)
 	for i := range pitchLPCOrder {
 		tmp *= 0.9
@@ -183,12 +184,13 @@ func refineXcorr(xLP, y, xcorr []float32, length, maxPitch int, coarse [2]int) {
 // pitchSearch finds the lag of the strongest correlation between xLP and y.
 // Port of libopus pitch_search: a coarse pass on a further 2x decimation, then
 // a finer pass restricted to the neighborhood of the two best coarse lags.
-func pitchSearch(xLP, y []float32, length, maxPitch int) int {
+func pitchSearch(xLP, y []float32, length, maxPitch int, scratch *encoderScratch) int {
 	lag := length + maxPitch
 
-	xLP4 := make([]float32, length>>2)
-	yLP4 := make([]float32, lag>>2)
-	xcorr := make([]float32, maxPitch>>1)
+	xLP4 := scratch.pitchX[:length>>2]
+	yLP4 := scratch.pitchY[:lag>>2]
+	xcorr := scratch.pitchXC[:maxPitch>>1]
+	clear(xcorr)
 
 	// Downsample by 2 again.
 	for j := range xLP4 {
@@ -252,6 +254,7 @@ func dualInnerProdLag(x []float32, base, n, lagA, lagB int) (a, b float32) {
 //nolint:cyclop,gocognit // Mirrors the reference sub-multiple scan.
 func removeDoubling(
 	x []float32, maxPeriod, minPeriod, n int, period *int, prevPeriod int, prevGain float32,
+	scratch *encoderScratch,
 ) float32 {
 	minPeriod0 := minPeriod
 	maxPeriod /= 2
@@ -274,7 +277,7 @@ func removeDoubling(
 
 	// yyLookup[i] is the energy of the window ending i samples back, updated
 	// incrementally so each candidate lag costs one inner product, not two.
-	yyLookup := make([]float32, maxPeriod+1)
+	yyLookup := scratch.yyLookup[:maxPeriod+1]
 	yyLookup[0] = xx
 	yy := xx
 	for i := 1; i <= maxPeriod; i++ {

internal/celt/pitch_test.go

@@ -28,15 +28,16 @@ func runPitchPipeline(pcm []float32) (int, float32) {
 	const frameSampleCount = 960
 	pitchLen := (combFilterMaxPeriod + frameSampleCount) >> 1
 	buf := make([]float32, pitchLen)
-	pitchDownsample([][]float32{pcm}, buf, pitchLen, 2)
+	var scratch encoderScratch
+	pitchDownsample([][]float32{pcm}, buf, pitchLen, 2, &scratch)
 
 	period := pitchSearch(
 		buf[combFilterMaxPeriod>>1:], buf,
-		frameSampleCount, combFilterMaxPeriod-3*combFilterMinPeriod,
+		frameSampleCount, combFilterMaxPeriod-3*combFilterMinPeriod, &scratch,
 	)
 	period = combFilterMaxPeriod - period
 	gain := removeDoubling(
-		buf, combFilterMaxPeriod, combFilterMinPeriod, frameSampleCount, &period, 0, 0)
+		buf, combFilterMaxPeriod, combFilterMinPeriod, frameSampleCount, &period, 0, 0, &scratch)
 
 	return period, gain
 }
@@ -83,7 +84,7 @@ func TestCeltLPCFlatSpectrum(t *testing.T) {
 	// White-noise autocorrelation (only ac[0] non-zero) has no prediction gain,
 	// so every coefficient must come out at zero.
 	ac := []float32{1, 0, 0, 0, 0}
-	lpc := celtLPC(ac, 4)
+	lpc := celtLPC(ac, 4, make([]float32, 4))
 
 	for i, v := range lpc {
 		assert.InDelta(t, 0, v, 1e-6, "coefficient %d", i)

internal/celt/tf_analysis.go

@@ -37,12 +37,13 @@ func tfAnalysis(
 	tfChan, frameSize int,
 	importance *[maxBands]int,
 	tfRes *[maxBands]int,
+	scratch *encoderScratch,
 ) int {
 	bias := 0.04 * maxFloat32(-0.25, 0.5-tfEstimate)
 
 	widest := int(bandEdges[endBand]-bandEdges[endBand-1]) << lm
-	tmp := make([]float32, widest)
-	tmpOne := make([]float32, widest)
+	tmp := scratch.tfTmp[:widest]
+	tmpOne := scratch.tfTmpOne[:widest]
 
 	var metric [maxBands]int
 	for band := startBand; band < endBand; band++ {