perf(asr): bulk-fill MLMultiArray resets and copies - #941
vakharwalad23 wants to merge 7 commits into
Conversation
MLArrayCache.returnArray reset the 240000-sample preprocessor input one NSNumber at a time, about 20 ms per transcription on the path that returns the transcript, and TdtDecoderState(from:) copied the LSTM state the same way before every recoverable decode. Reset through memset over the backing extent, fill other values through a typed pointer, and memcpy between identical layouts; the element loops remain as fallbacks.
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Replace the overlapping-region memcpy path with safe overlap handling, and strengthen padded-byte coverage.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 1
Open (1)
What changed in this PR
Optimizes MLMultiArray reset and copy operations for ASR performance.
Changes:
- Adds bulk zeroing, typed fills, and compatible bulk copies.
- Preserves stride-aware fallback behavior.
- Adds correctness and performance regression tests.
| File | Description |
|---|---|
Tests/FluidAudioTests/Shared/MLArrayCacheTests.swift |
Tests cache resets and padded strides; padding-byte coverage remains a nit. |
Tests/FluidAudioTests/ASR/Parakeet/SlidingWindow/TDT/Decoder/TdtDecoderStateV3Tests.swift |
Tests reset, copy, stride, and performance behavior. |
Sources/FluidAudio/ASR/Parakeet/SlidingWindow/TDT/Decoder/TdtDecoderState.swift |
Implements optimized reset/copy paths; the memcpy fast path must handle overlapping regions safely. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
memcpy is undefined when source and destination overlap, which zero-copy views of one allocation and a self-copy can produce. memmove keeps the bulk path correct there, a self-copy returns early, and the tests cover overlapping views, self-copy, and the padding bytes of an aligned array.
|
Thanks for the speedup. One correctness concern that should be addressed before merge, two test fixes that go with it, and a few optional cleanups. Blocking
Tests to fix alongside
Non-blocking
|
withUnsafeMutableBytes reports the padded byte span, so a view built with padded strides over a tighter allocation would be overrun by a span-wide memset or copy. The bulk paths now require the span to equal count times the element size; padded layouts keep the element loop, so nothing past the last element is written. The copy goes through copyMemory, which is overlap-safe, and the doc comments state the overlap and -0.0 behavior.
resetData and copyData move next to reset(to:) in Shared, and reset(to:) delegates to them, so MLArrayCache no longer reaches into the decoder for its helper. reset(to:) used to walk count contiguous slots, which skipped the last rows of a padded array; the stride-aware fill covers them. The warm-up's private vDSP fill goes the same way.
The only getArray consumer overwrites the full extent of the preprocessor input with memcpy, and every caller pads the audio before that, so the reset on return was dead work: 240000 boxed stores, about 20 ms, on the path that returns the transcript.
createZeroCopyView checked the logical element count against the source but built the view with padded strides, so a view whose innermost dimension is not a multiple of 16 could extend past the source storage. The check now uses the span the strides imply, as ANEMemoryUtils does.
|
Working through these now on the same branch: the fast path will apply only to contiguous storage (span == count * element size) with the element loop kept for padded strides, the padded-stride test poisons the full backing before asserting, the budget tests are gated off in CI, and the non-blocking items (helpers consolidated under Shared with reset(to:) delegating, copyMemory-based copy, overlap and -0.0 notes, full-array test checks, a float64 case) go in alongside. I will also drop the zero-on-return reset as suggested, since the only getArray consumer overwrites the full extent. Push coming shortly. |
|
Pushed five commits covering the review. The bulk paths now run only on contiguous storage (span == count * element size), with the element loop kept for padded strides and sentinel tests for the overrun; createZeroCopyView bounds by its padded span, which was the root cause. The padded test poisons the backing first, the timing budgets skip under CI, the reset on return is gone, resetData/copyData live in Shared with reset(to:) delegating (that also fixed its padded-array gap), copyData is the copyMemory shape, the overlap and -0.0 notes are in, and the tests use the file helpers with a float64 case added. Full suite green locally, download-gated classes included. |
Why is this change needed?
MLArrayCache.returnArrayclears every returned array withMLMultiArray.resetData(to: 0), which stores oneNSNumberper element through the subscript. The preprocessor input is[1, 240000]float32, so everyAsrManager.transcribecall spends about 20 ms zero-filling it (release build, M3 Pro) on the path that returns the transcript, after the models are already done.TdtDecoderState(from:)copies the LSTM state the same way, one element at a time, and since #910 it runs before every recoverable inference.This PR fills and copies through the backing storage instead:
resetData(to: 0)is onememsetover the array's byte extent for every data type.copyData(from:)is onememcpywhen shape, data type and strides match; otherwise it keeps the element loop.withUnsafeMutableBytesreports the full backing extent, so the ANE-aligned arrays with padded strides ([10, 10]with strides[16, 1]) are covered too.Measurements
Apple M3 Pro, macOS 26.6, Xcode 26.6.
Unit level, debug test build:
resetData(to: 0),[1, 240000]float32copyData(from:),[1, 240000]float32Release build, the same 240000-element reset in isolation: 21.5 ms as an
NSNumberloop, 0.006 ms asmemset.End to end, release build,
AsrManager.transcribeon an 11 s clip with a Parakeet TDT 0.6B v3 bundle (Orukeet), median of 20 warm runs,mainvs this branch:transcribewall timeTranscripts are identical before and after.
Tests
MLArrayCacheTests.testReturnArrayResetsPreprocessorBufferWithinBudget(fails onmainat 42.9 ms against a 5 ms budget, best of 5)MLArrayCacheTests.testReturnArrayResetsPaddedStrideArrayTdtDecoderStateV3Tests.testMLMultiArrayCopyDataLargeArrayWithinBudget(fails onmainat 62.9 ms)TdtDecoderStateV3Tests.testMLMultiArrayCopyDataAcrossStrideLayoutsTdtDecoderStateV3Tests.testMLMultiArrayResetDataInt32Value,testMLMultiArrayResetDataFloat16ValueThe full suite passed locally with model downloads enabled (2534 tests).