Skip to content

Commit 36416d4

Browse files
kevchengcodeskevchengcodes
committed
remove some redundant array copying
1 parent f151066 commit 36416d4

3 files changed

Lines changed: 49 additions & 60 deletions

File tree

‎swift/Sources/CoreAISpeech/ParakeetTDTDecoder.swift‎

Lines changed: 47 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,11 @@ public struct ParakeetTDTDecoder: SpeechDecoder {
2727
}
2828

2929
/// The `joint` graph plus the descriptors its buffers are built from.
30+
///
31+
/// No `decoderIn`: the step graph's `decoder_output` buffer is handed to the joint
32+
/// as `decoder_hidden_states` directly, so there's no second buffer to allocate.
3033
private struct JointGraph {
3134
let fn: InferenceFunction
32-
let decoderIn: NDArrayDescriptor
3335
let encoderIn: NDArrayDescriptor
3436
let logits: NDArrayDescriptor
3537
}
@@ -38,12 +40,15 @@ public struct ParakeetTDTDecoder: SpeechDecoder {
3840
/// place on each step rather than reallocated per emission.
3941
private struct Buffers {
4042
var inputIds: NDArray
43+
/// LSTM state, double-buffered: a step reads `hIn`/`cIn` and writes `hOut`/`cOut`,
44+
/// and the call site swaps each pair to adopt the new state.
4145
var hIn: NDArray
4246
var cIn: NDArray
43-
var decOut: NDArray
4447
var hOut: NDArray
4548
var cOut: NDArray
46-
var jointDecIn: NDArray
49+
/// Doubles as the joint's `decoder_hidden_states` input: same `[1, 1, hidden]` shape,
50+
/// same precision, so the step's output needs no copy to become the joint's input.
51+
var decOut: NDArray
4752
var jointEncIn: NDArray
4853
var logits: NDArray
4954

@@ -54,14 +59,18 @@ public struct ParakeetTDTDecoder: SpeechDecoder {
5459
inputIds = NDArray(descriptor: step.inputIds.resolvingDynamicDimensions([1, 1]))
5560
hIn = NDArray(descriptor: step.hiddenIn.resolvingDynamicDimensions(lstmShape))
5661
cIn = NDArray(descriptor: step.cellIn.resolvingDynamicDimensions(lstmShape))
57-
decOut = NDArray(descriptor: step.decoderOut.resolvingDynamicDimensions([1, 1, hidden]))
5862
hOut = NDArray(descriptor: step.newHidden.resolvingDynamicDimensions(lstmShape))
5963
cOut = NDArray(descriptor: step.newCell.resolvingDynamicDimensions(lstmShape))
60-
jointDecIn = NDArray(
61-
descriptor: joint.decoderIn.resolvingDynamicDimensions([1, 1, hidden]))
64+
decOut = NDArray(descriptor: step.decoderOut.resolvingDynamicDimensions([1, 1, hidden]))
6265
jointEncIn = NDArray(
6366
descriptor: joint.encoderIn.resolvingDynamicDimensions([1, 1, hidden]))
6467
logits = NDArray(descriptor: joint.logits.resolvingDynamicDimensions([1, 1, logitsSize]))
68+
69+
// The first step reads `hIn`/`cIn` before anything has written them, so seed the
70+
// zero state here rather than relying on fresh-allocation contents.
71+
let zeros = [Float](repeating: 0, count: lstmShape.reduce(1, *))
72+
fillFloatNDArray(&hIn, with: zeros)
73+
fillFloatNDArray(&cIn, with: zeros)
6574
}
6675
}
6776

@@ -93,16 +102,15 @@ public struct ParakeetTDTDecoder: SpeechDecoder {
93102
case .ndArray(let newCellDesc) = stepDesc.outputDescriptor(of: "new_cell_state")
94103
else { throw SpeechError.missingModel("Unexpected decoder_step descriptors") }
95104

96-
guard case .ndArray(let jointDecDesc) = jointDesc.inputDescriptor(of: "decoder_hidden_states"),
105+
guard case .ndArray(_) = jointDesc.inputDescriptor(of: "decoder_hidden_states"),
97106
case .ndArray(let jointEncDesc) = jointDesc.inputDescriptor(of: "encoder_hidden_states"),
98107
case .ndArray(let logitsDesc) = jointDesc.outputDescriptor(of: "logits")
99108
else { throw SpeechError.missingModel("Unexpected joint descriptors") }
100109

101110
self.stepGraph = StepGraph(
102111
fn: stepFn, inputIds: inputIdsDesc, hiddenIn: hiddenInDesc, cellIn: cellInDesc,
103112
decoderOut: decoderOutDesc, newHidden: newHiddenDesc, newCell: newCellDesc)
104-
self.jointGraph = JointGraph(
105-
fn: jointFn, decoderIn: jointDecDesc, encoderIn: jointEncDesc, logits: logitsDesc)
113+
self.jointGraph = JointGraph(fn: jointFn, encoderIn: jointEncDesc, logits: logitsDesc)
106114
}
107115

108116
public func decode(
@@ -130,7 +138,6 @@ public struct ParakeetTDTDecoder: SpeechDecoder {
130138
// of frames that carry real audio. Dynamic exports pass tEnc, so cap == tEnc.
131139
let cap = max(1, min(tEnc, validEncoderFrames))
132140
let lstmShape = [cfg.numDecoderLayers, 1, hidden]
133-
let lstmCount = cfg.numDecoderLayers * hidden
134141

135142
// Pull the full encoder output once; slice frame-by-frame in pure Swift.
136143
// flattenAsFloat inspects the array's own scalar type, so this reads an
@@ -140,17 +147,12 @@ public struct ParakeetTDTDecoder: SpeechDecoder {
140147
step: stepGraph, joint: jointGraph,
141148
lstmShape: lstmShape, hidden: hidden, logitsSize: logitsSize)
142149

143-
// Swift-side LSTM state — zero-seeded, advanced only on non-blank *input*.
144-
var hState = [Float](repeating: 0, count: lstmCount)
145-
var cState = [Float](repeating: 0, count: lstmCount)
146-
147150
// Previous iteration's symbol, blanks included — fed back as the next `input_ids`.
148151
// Distinct from `emitted`, which keeps only the non-blank symbols.
149152
var previousSymbol: Int32 = cfg.blankTokenId
150153
var emitted: [Int32] = []
151154
var frame = 0
152155
var firstStep = true
153-
var cachedDecBuf: [Float]? = nil // last decoder output; reused on blank-input iterations
154156
let emitCap = cap * cfg.maxSymbolsPerStep
155157

156158
var stepTimesMs: [Double] = []
@@ -169,26 +171,32 @@ public struct ParakeetTDTDecoder: SpeechDecoder {
169171
let inputIsBlank = (previousSymbol == cfg.blankTokenId)
170172

171173
// Blank-skip (optimization): a blank input reproduces the last decoder
172-
// output, since the state was held too.
173-
let decBuf: [Float]
174-
if !firstStep, inputIsBlank, let cached = cachedDecBuf {
175-
decBuf = cached
174+
// output, since the state was held too — and `buffers.decOut` still holds
175+
// it, because this branch doesn't run the step graph that would overwrite it.
176+
if !firstStep, inputIsBlank {
176177
coverage.blankSkipReuses += 1
177178
} else {
178-
let stepped = try await runDecoderStep(
179-
token: previousSymbol, hState: hState, cState: cState,
180-
buffers: &buffers)
181-
decBuf = stepped.decoderOutput
182-
cachedDecBuf = decBuf
183-
if capturedStep == nil { capturedStep = stepped }
179+
try await runDecoderStep(token: previousSymbol, buffers: &buffers)
180+
181+
// Parity capture only: gating on nil keeps the flattens to once per decode
182+
// instead of once per step, and reads `hOut`/`cOut` before the swap below.
183+
if capturedStep == nil {
184+
capturedStep = (
185+
decoderOutput: flattenAsFloat(buffers.decOut),
186+
newHidden: flattenAsFloat(buffers.hOut),
187+
newCell: flattenAsFloat(buffers.cOut)
188+
)
189+
}
184190

185191
// Load-bearing: a blank carries no label, so the state must not absorb
186192
// it. Mirrors HF `cache.update(..., mask=~blank_mask)`. Blanks normally
187193
// never get here at all — they take the reuse branch above — but the
188-
// guard still has to hold if the cache is ever invalidated.
194+
// guard still has to hold if that branch is ever changed. Not adopting
195+
// means not swapping: `hIn`/`cIn` keep the state the step ran from, and
196+
// the next step overwrites `hOut`/`cOut`.
189197
if firstStep || !inputIsBlank {
190-
hState = stepped.newHidden
191-
cState = stepped.newCell
198+
swap(&buffers.hIn, &buffers.hOut)
199+
swap(&buffers.cIn, &buffers.cOut)
192200
coverage.lstmStateAdvances += 1
193201
}
194202
firstStep = false
@@ -197,7 +205,6 @@ public struct ParakeetTDTDecoder: SpeechDecoder {
197205
// Joint(decoder_output, encoder[:, frame:frame+1, :]).
198206
let encOffset = frame * hidden
199207
let logitsFlat = try await runJoint(
200-
decoderOutput: decBuf,
201208
encoderFrame: encFlat[encOffset..<encOffset + hidden],
202209
buffers: &buffers)
203210
if capturedLogits == nil { capturedLogits = logitsFlat }
@@ -296,18 +303,14 @@ public struct ParakeetTDTDecoder: SpeechDecoder {
296303

297304
// MARK: - Graph invocations
298305

299-
/// One LSTM step: feed `token` at (`hState`, `cState`) and return the decoder output
300-
/// alongside the resulting state.
306+
/// One LSTM step: feed `token` at the state in `buffers.hIn`/`cIn`, leaving the decoder
307+
/// output in `buffers.decOut` and the resulting state in `buffers.hOut`/`cOut`.
301308
///
302-
/// Returning the new state rather than storing it keeps the decision of *whether* to
303-
/// adopt it — the HF `mask=~blank_mask` rule — at the call site with the blank
304-
/// bookkeeping it depends on.
305-
private func runDecoderStep(
306-
token: Int32, hState: [Float], cState: [Float], buffers: inout Buffers
307-
) async throws -> (decoderOutput: [Float], newHidden: [Float], newCell: [Float]) {
309+
/// Writing the new state to the output buffers rather than adopting it keeps the decision
310+
/// of *whether* to adopt — the HF `mask=~blank_mask` rule — at the call site with the
311+
/// blank bookkeeping it depends on. The call site adopts by swapping the buffer pairs.
312+
private func runDecoderStep(token: Int32, buffers: inout Buffers) async throws {
308313
fillNDArray(&buffers.inputIds, as: Int32.self, with: [token])
309-
fillFloatNDArray(&buffers.hIn, with: hState)
310-
fillFloatNDArray(&buffers.cIn, with: cState)
311314

312315
var stepOut = InferenceFunction.MutableViews()
313316
stepOut.insert(&buffers.decOut, for: "decoder_output")
@@ -321,26 +324,22 @@ public struct ParakeetTDTDecoder: SpeechDecoder {
321324
],
322325
states: InferenceFunction.MutableViews(),
323326
outputViews: consume stepOut)
324-
325-
return (
326-
decoderOutput: flattenAsFloat(buffers.decOut),
327-
newHidden: flattenAsFloat(buffers.hOut),
328-
newCell: flattenAsFloat(buffers.cOut)
329-
)
330327
}
331328

332329
/// `joint(decoder_output, encoder_frame)` → the flattened `[vocab | durations]` row.
330+
///
331+
/// The decoder side comes straight from `buffers.decOut`, whether the last step wrote it
332+
/// or the blank-skip branch left it in place.
333333
private func runJoint(
334-
decoderOutput: [Float], encoderFrame: ArraySlice<Float>, buffers: inout Buffers
334+
encoderFrame: ArraySlice<Float>, buffers: inout Buffers
335335
) async throws -> [Float] {
336-
fillFloatNDArray(&buffers.jointDecIn, with: decoderOutput)
337336
fillFloatNDArray(&buffers.jointEncIn, with: encoderFrame)
338337

339338
var jointOut = InferenceFunction.MutableViews()
340339
jointOut.insert(&buffers.logits, for: "logits")
341340
_ = try await jointGraph.fn.run(
342341
inputs: [
343-
"decoder_hidden_states": buffers.jointDecIn,
342+
"decoder_hidden_states": buffers.decOut,
344343
"encoder_hidden_states": buffers.jointEncIn,
345344
],
346345
states: InferenceFunction.MutableViews(),

‎swift/Sources/CoreAISpeech/SpeechDecoder.swift‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@ public struct WhisperDecoder: SpeechDecoder {
120120

121121
guard case .ndArray(let inputIdsNDDesc) = decDesc.inputDescriptor(of: "input_ids"),
122122
case .ndArray(let posIdsNDDesc) = decDesc.inputDescriptor(of: "position_ids"),
123-
case .ndArray(let encHSNDDesc) = decDesc.inputDescriptor(of: "encoder_hidden_states"),
123+
case .ndArray(_) = decDesc.inputDescriptor(of: "encoder_hidden_states"),
124124
case .ndArray(let keyCacheNDDesc) = decDesc.stateDescriptor(of: "keyCache"),
125125
case .ndArray(let valCacheNDDesc) = decDesc.stateDescriptor(of: "valueCache"),
126126
case .ndArray(let logitsNDDesc) = decDesc.outputDescriptor(of: "logits")
@@ -132,14 +132,6 @@ public struct WhisperDecoder: SpeechDecoder {
132132
let vcShape = valCacheNDDesc.shape.map { $0 < 0 ? maxTargetPos : $0 }
133133
var keyCache = NDArray(descriptor: keyCacheNDDesc.resolvingDynamicDimensions(kcShape))
134134
var valueCache = NDArray(descriptor: valCacheNDDesc.resolvingDynamicDimensions(vcShape))
135-
136-
var encHSArray = NDArray(descriptor: encHSNDDesc.resolvingDynamicDimensions(encoderOutputShape))
137-
// Both sides dispatch on the array's own scalar type: an f16 export hands back an
138-
// f16 encoder output and expects an f16 `encoder_hidden_states`, so reading or
139-
// filling these as `Float` would trap on the scalar-type check.
140-
let encFlat = flattenAsFloat(encoderOutput)
141-
fillFloatNDArray(&encHSArray, with: encFlat)
142-
143135
var logitsArray = NDArray(descriptor: logitsNDDesc.resolvingDynamicDimensions([1, 1, vocabSize]))
144136

145137
func step(_ tok: Int32, pos: Int) async throws {
@@ -155,7 +147,7 @@ public struct WhisperDecoder: SpeechDecoder {
155147
_ = try await decFn.run(
156148
inputs: [
157149
"input_ids": ids, "position_ids": posIds,
158-
"encoder_hidden_states": encHSArray,
150+
"encoder_hidden_states": encoderOutput,
159151
],
160152
states: consume st, outputViews: consume out)
161153
}

‎swift/Sources/Tools/speech-recognizer/SpeechRecognizerMain.swift‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -253,8 +253,6 @@ func runLegacy(model: String, audioPath: String?, warmup: Bool) async throws {
253253
URL(fileURLWithPath: path), targetSampleRate: 16_000)
254254
let floats = MelSpectrogram.fromPCM(pcm)
255255
melArray = NDArray(descriptor: melNDDesc.resolvingDynamicDimensions([1, 128, 3000]))
256-
// `fillFloatNDArray`, not `fillNDArray(as: Float.self,…)`: an f16 export's
257-
// `input_features` is an f16 tensor, and the strict variant traps on it.
258256
fillFloatNDArray(&melArray, with: floats)
259257
} else {
260258
melArray = NDArray(descriptor: melNDDesc.resolvingDynamicDimensions([1, 128, 3000]))

0 commit comments

Comments
 (0)