perf(asr): bulk-fill MLMultiArray resets and copies

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..<count {
-            self[i] = value
-        }
-    }
-
-    func copyData(from source: MLMultiArray) {
-        for i in 0..<count {
-            self[i] = source[i]
-        }
-    }
-}

Sources/FluidAudio/Shared/ANEMemoryOptimizer.swift

@@ -88,11 +88,14 @@ public final class ANEMemoryOptimizer {
         shape: [NSNumber],
         offset: Int = 0
     ) throws -> 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
         )
     }

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
         }

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..<count {
+            self[i] = value
+        }
+    }
+
+    /// Copies every element from `source`.
+    ///
+    /// Identical contiguous layouts copy the storage in bulk; that copy is overlap-safe, so two
+    /// views of one allocation may overlap. Any other pair reads the whole source before writing,
+    /// so overlapping views with different layouts stay safe too.
+    func copyData(from source: MLMultiArray) {
+        let elementSize = ANEMemoryUtils.getElementSize(for: dataType)
+        if dataType == source.dataType, shape == source.shape, strides == source.strides {
+            let copied = withUnsafeMutableBytes { destination, _ -> Bool in
+                guard destination.count == count * elementSize else {
+                    return false
+                }
+                source.withUnsafeBytes { destination.copyMemory(from: $0) }
+                return true
+            }
+            if copied {
+                return
+            }
+        }
+        let values = (0..<count).map { source[$0] }
+        for i in 0..<count {
+            self[i] = values[i]
         }
     }
 }

Sources/FluidAudio/Shared/ModelWarmup.swift

@@ -1,4 +1,3 @@
-import Accelerate
 import CoreML
 import Foundation
 
@@ -27,7 +26,7 @@ enum ModelWarmup {
             shape: inputShape.map { NSNumber(value: $0) },
             dataType: .float32
         )
-        array.resetToZeros()
+        array.resetData(to: 0)
 
         let features = try MLDictionaryFeatureProvider(dictionary: [
             inputName: MLFeatureValue(multiArray: array)
@@ -85,13 +84,13 @@ enum ModelWarmup {
                 shape: featureShape.map { NSNumber(value: $0) },
                 dataType: .float32
             )
-            featureArray.resetToZeros()
+            featureArray.resetData(to: 0)
 
             let weightArray = try MLMultiArray(
                 shape: weightsShape.map { NSNumber(value: $0) },
                 dataType: .float32
             )
-            weightArray.resetToZeros()
+            weightArray.resetData(to: 0)
 
             let provider = try MLDictionaryFeatureProvider(dictionary: [
                 "fbank_features": MLFeatureValue(multiArray: featureArray),
@@ -110,7 +109,7 @@ enum ModelWarmup {
                 shape: [1, 1, 1, NSNumber(value: totalElements)],
                 dataType: .float32
             )
-            combinedArray.resetToZeros()
+            combinedArray.resetData(to: 0)
 
             let provider = try MLDictionaryFeatureProvider(dictionary: [
                 "audio_and_weights": MLFeatureValue(multiArray: combinedArray)
@@ -126,13 +125,13 @@ enum ModelWarmup {
             shape: [1, 1, NSNumber(value: audioSamples)],
             dataType: .float32
         )
-        audioArray.resetToZeros()
+        audioArray.resetData(to: 0)
 
         let weightArray = try MLMultiArray(
             shape: [1, NSNumber(value: weightFrames)],
             dataType: .float32
         )
-        weightArray.resetToZeros()
+        weightArray.resetData(to: 0)
 
         let provider = try MLDictionaryFeatureProvider(dictionary: [
             "audio": MLFeatureValue(multiArray: audioArray),
@@ -142,12 +141,3 @@ enum ModelWarmup {
         _ = try model.prediction(from: provider)
     }
 }
-
-extension MLMultiArray {
-    fileprivate func resetToZeros() {
-        let pointer = dataPointer.assumingMemoryBound(to: Float.self)
-        let count = self.count
-        var zero: Float = 0
-        vDSP_vfill(&zero, pointer, 1, vDSP_Length(count))
-    }
-}