diff --git a/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/TDT/Decoder/TdtDecoderState.swift b/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/TDT/Decoder/TdtDecoderState.swift index b89de500..b25a65e1 100644 --- a/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/TDT/Decoder/TdtDecoderState.swift +++ b/Sources/FluidAudio/ASR/Parakeet/SlidingWindow/TDT/Decoder/TdtDecoderState.swift @@ -94,17 +94,3 @@ public struct TdtDecoderState: Sendable { // Keep LSTM states as they represent the final linguistic context } } - -extension MLMultiArray { - func resetData(to value: NSNumber) { - for i in 0.. MLMultiArray { - // Ensure we have enough data - let sourceElements = sourceArray.shape.map { $0.intValue }.reduce(1, *) - let viewElements = shape.map { $0.intValue }.reduce(1, *) - - guard offset + viewElements <= sourceElements else { + // The view pads its innermost stride, so its storage span can exceed its element count; + // bound the span, not the count, or a padded view would run past a tighter source. + let strides = calculateOptimalStrides(for: shape, dataType: sourceArray.dataType) + let viewSpan = shape.isEmpty ? 0 : strides[0].intValue * shape[0].intValue + let sourceSpan = + sourceArray.shape.isEmpty ? 0 : sourceArray.strides[0].intValue * sourceArray.shape[0].intValue + + guard offset + viewSpan <= sourceSpan else { throw DiarizerError.invalidArrayBounds } @@ -107,7 +110,7 @@ public final class ANEMemoryOptimizer { dataPointer: offsetPointer, shape: shape, dataType: sourceArray.dataType, - strides: calculateOptimalStrides(for: shape, dataType: sourceArray.dataType), + strides: strides, deallocator: nil // No deallocation since it's a view ) } diff --git a/Sources/FluidAudio/Shared/MLArrayCache.swift b/Sources/FluidAudio/Shared/MLArrayCache.swift index b62ba4fb..29ea5995 100644 --- a/Sources/FluidAudio/Shared/MLArrayCache.swift +++ b/Sources/FluidAudio/Shared/MLArrayCache.swift @@ -31,7 +31,8 @@ actor MLArrayCache { return try ANEMemoryUtils.createAlignedArray(shape: shape, dataType: dataType) } - /// Return an array to the cache for reuse + /// Return an array to the cache for reuse. Its contents are kept: every consumer of `getArray` + /// overwrites the full extent before use, so clearing here would be wasted work. func returnArray(_ array: MLMultiArray) { let key = CacheKey( shape: array.shape.map { $0.intValue }, @@ -42,7 +43,6 @@ actor MLArrayCache { // Limit cache size per key if arrays.count < maxCacheSize / max(cache.count, 1) { - array.resetData(to: 0) arrays.append(array) cache[key] = arrays } diff --git a/Sources/FluidAudio/Shared/MLMultiArray+Extensions.swift b/Sources/FluidAudio/Shared/MLMultiArray+Extensions.swift index 779ece06..a6d6c399 100644 --- a/Sources/FluidAudio/Shared/MLMultiArray+Extensions.swift +++ b/Sources/FluidAudio/Shared/MLMultiArray+Extensions.swift @@ -1,15 +1,70 @@ import CoreML +import Foundation extension MLMultiArray { /// Reset all elements in the array to the given value. func reset(to value: NSNumber) { - let count = self.count - if self.dataType == .float32 { - let ptr = self.dataPointer.bindMemory(to: Float.self, capacity: count) - ptr.update(repeating: value.floatValue, count: count) - } else if self.dataType == .int32 { - let intPtr = self.dataPointer.bindMemory(to: Int32.self, capacity: count) - intPtr.update(repeating: value.int32Value, count: count) + resetData(to: value) + } + + /// Fills every element with `value`. + /// + /// Contiguous storage fills in bulk: zero is one `memset` for every data type, other values + /// fill through a typed pointer for float32, float64 and int32. Padded strides and other data + /// types fill element by element, so nothing past the last element is ever written. `value` is + /// compared as an `NSNumber`, so `-0.0` takes the zero path and lands as `+0.0`. + func resetData(to value: NSNumber) { + let elementSize = ANEMemoryUtils.getElementSize(for: dataType) + let filled = withUnsafeMutableBytes { bytes, _ -> Bool in + guard bytes.count == count * elementSize, let base = bytes.baseAddress else { + return false + } + if value == 0 { + memset(base, 0, bytes.count) + return true + } + switch dataType { + case .float32: + bytes.bindMemory(to: Float.self).update(repeating: value.floatValue) + case .float64: + bytes.bindMemory(to: Double.self).update(repeating: value.doubleValue) + case .int32: + bytes.bindMemory(to: Int32.self).update(repeating: value.int32Value) + default: + return false + } + return true + } + if filled { + return + } + for i in 0.. Bool in + guard destination.count == count * elementSize else { + return false + } + source.withUnsafeBytes { destination.copyMemory(from: $0) } + return true + } + if copied { + return + } + } + let values = (0...allocate(capacity: elements) + defer { storage.deallocate() } + storage.initialize(repeating: 1, count: elements) + let sentinel: Float = 12345 + for i in outside { + storage[i] = sentinel + } + let view = try MLMultiArray( + dataPointer: UnsafeMutableRawPointer(storage), shape: [2, 10], dataType: .float32, + strides: [16, 1], deallocator: nil) + + view.resetData(to: 0) + + verifyArrayIsZero(view) + for i in outside { + XCTAssertEqual(storage[i], sentinel, "Storage outside the elements was written at \(i)") + } + } + func testMLMultiArrayCopyData() throws { let sourceArray = try MLMultiArray(shape: [3, 4], dataType: .float32) let destArray = try MLMultiArray(shape: [3, 4], dataType: .float32) @@ -214,6 +315,124 @@ final class TdtDecoderStateV3Tests: XCTestCase { verifyArraysEqual(destArray, sourceArray) } + func testMLMultiArrayCopyDataLargeArrayWithinBudget() throws { + // The decoder state is snapshotted before every inference, so the copy must be a bulk + // transfer. A per-element copy of 240000 samples costs tens of milliseconds. Timed locally + // only: the parallel CI job shares its machine. + try XCTSkipIf(ProcessInfo.processInfo.environment["CI"] != nil, "Timing budgets run locally only") + let shape: [NSNumber] = [1, NSNumber(value: ASRConstants.maxModelSamples)] + let sourceArray = try MLMultiArray(shape: shape, dataType: .float32) + let destArray = try MLMultiArray(shape: shape, dataType: .float32) + sourceArray[sourceArray.count - 1] = NSNumber(value: Float(3)) + let clock = ContinuousClock() + var best = Double.infinity + + for _ in 0..<5 { + let elapsed = clock.measure { + destArray.copyData(from: sourceArray) + } + best = min(best, elapsed / .milliseconds(1)) + } + + XCTAssertEqual(destArray[destArray.count - 1].floatValue, 3) + XCTAssertLessThan(best, 5, "copyData took \(best) ms for \(ASRConstants.maxModelSamples) elements") + } + + func testMLMultiArrayCopyDataAcrossStrideLayouts() throws { + // A plain array and an ANE-aligned array of the same shape have different strides; + // the copy must still land every element. + let shape: [NSNumber] = [10, 10] + let sourceArray = try ANEMemoryUtils.createAlignedArray(shape: shape, dataType: .float32) + let destArray = try MLMultiArray(shape: shape, dataType: .float32) + XCTAssertNotEqual(sourceArray.strides, destArray.strides) + + fillArrayWithTestData(sourceArray, multiplier: 1.5) + + destArray.copyData(from: sourceArray) + + verifyArraysEqual(destArray, sourceArray) + } + + func testMLMultiArrayCopyDataBetweenOverlappingViews() throws { + // Two zero-copy views of one allocation, offset by 16 elements, overlap on 48 of their 64 + // elements; the copy must behave as if the source were read completely first. + let backing = try ANEMemoryUtils.createAlignedArray(shape: [1, 96], dataType: .float32) + for i in 0...allocate(capacity: elements) + let destinationStorage = UnsafeMutablePointer.allocate(capacity: elements) + defer { + sourceStorage.deallocate() + destinationStorage.deallocate() + } + sourceStorage.initialize(repeating: 1, count: elements) + let sentinel: Float = 12345 + destinationStorage.initialize(repeating: sentinel, count: elements) + let sourceView = try MLMultiArray( + dataPointer: UnsafeMutableRawPointer(sourceStorage), shape: [2, 10], dataType: .float32, + strides: [16, 1], deallocator: nil) + let destinationView = try MLMultiArray( + dataPointer: UnsafeMutableRawPointer(destinationStorage), shape: [2, 10], dataType: .float32, + strides: [16, 1], deallocator: nil) + fillArrayWithTestData(sourceView, multiplier: 3) + + destinationView.copyData(from: sourceView) + + verifyArraysEqual(destinationView, sourceView) + for i in outside { + XCTAssertEqual(destinationStorage[i], sentinel, "Storage outside the elements was written at \(i)") + } + } + func testMLMultiArrayCopyDataNonFloat() throws { let sourceArray = try MLMultiArray(shape: [2, 3], dataType: .int32) let destArray = try MLMultiArray(shape: [2, 3], dataType: .int32) diff --git a/Tests/FluidAudioTests/Shared/ANEMemoryOptimizerTests.swift b/Tests/FluidAudioTests/Shared/ANEMemoryOptimizerTests.swift index 0bf57ba4..af6d5f2d 100644 --- a/Tests/FluidAudioTests/Shared/ANEMemoryOptimizerTests.swift +++ b/Tests/FluidAudioTests/Shared/ANEMemoryOptimizerTests.swift @@ -193,6 +193,25 @@ final class ANEMemoryOptimizerTests: XCTestCase { // MARK: - Memory Pressure Tests + // MARK: - Zero-Copy View Bounds + + func testZeroCopyViewRejectsPaddedSpanBeyondSource() throws { + // A [10, 10] view pads its rows to 16 elements, so its storage span is 160 elements; a + // 100-element source cannot back it even though it holds 100 logical elements. + let source = try optimizer.createAlignedArray(shape: [100], dataType: .float32) + + XCTAssertThrowsError(try optimizer.createZeroCopyView(from: source, shape: [10, 10])) + } + + func testZeroCopyViewAcceptsSpanWithinSource() throws { + let source = try optimizer.createAlignedArray(shape: [1024], dataType: .float32) + + let view = try optimizer.createZeroCopyView(from: source, shape: [256], offset: 768) + + XCTAssertEqual(view.shape, [256]) + XCTAssertEqual(view.dataPointer, source.dataPointer.advanced(by: 768 * MemoryLayout.stride)) + } + func testMemoryPressureHandling() throws { // Create many buffers to simulate memory pressure var buffers: [MLMultiArray] = [] diff --git a/Tests/FluidAudioTests/Shared/MLArrayCacheTests.swift b/Tests/FluidAudioTests/Shared/MLArrayCacheTests.swift index 09673423..16728596 100644 --- a/Tests/FluidAudioTests/Shared/MLArrayCacheTests.swift +++ b/Tests/FluidAudioTests/Shared/MLArrayCacheTests.swift @@ -47,27 +47,6 @@ final class MLArrayCacheTests: XCTestCase { XCTAssertEqual(array2.dataType, .float32) } - func testReturnArrayResetsData() async throws { - let shape: [NSNumber] = [10] - let array = try await cache.getArray(shape: shape, dataType: .float32) - - // Set some values - for i in 0..