Make ANE utils concurrency safe

Author: BrandonWeng

Documentation/API.md

@@ -18,7 +18,7 @@ This page summarizes the primary public APIs across modules. See inline doc comm
 Main class for speaker diarization and "who spoke when" analysis.
 
 **Key Methods:**
-- `performCompleteDiarization(_:sampleRate:) throws -> DiarizerResult`
+- `performCompleteDiarization(_:sampleRate:) throws -> DiarizationResult`
   - Process complete audio file and return speaker segments
   - Parameters: `RandomAccessCollection<Float>` audio samples, sample rate (default: 16000)
   - Returns: `DiarizerResult` with speaker segments and timing

Sources/FluidAudio/Diarizer/Core/DiarizerManager.swift

@@ -1,5 +1,6 @@
 import Accelerate
 import CoreML
+import Dispatch
 import Foundation
 import OSLog
 
@@ -8,8 +9,6 @@ public final class DiarizerManager {
     internal let logger = AppLogger(category: "Diarizer")
     internal let config: DiarizerConfig
     private var models: DiarizerModels?
-    private var chunkBuffer: [Float] = []
-
     /// Public getter for segmentation model (for streaming)
     public var segmentationModel: MLModel? {
         return models?.segmentationModel
@@ -126,6 +125,7 @@ public final class DiarizerManager {
         let overlapDuration = Int(config.chunkOverlap.rounded())
         let chunkSize = sampleRate * chunkDuration
         let stepSize = chunkSize - (sampleRate * overlapDuration)
+        var chunkBuffer = [Float](repeating: 0.0, count: max(chunkSize, 1))
 
         var allSegments: [TimedSpeakerSegment] = []
 
@@ -145,7 +145,9 @@ public final class DiarizerManager {
                 chunk,
                 chunkOffset: chunkOffset,
                 models: models,
-                sampleRate: sampleRate
+                sampleRate: sampleRate,
+                chunkSize: chunkSize,
+                chunkBuffer: &chunkBuffer
             )
             allSegments.append(contentsOf: chunkSegments)
 
@@ -203,12 +205,13 @@ public final class DiarizerManager {
         _ chunk: C,
         chunkOffset: Double,
         models: DiarizerModels,
-        sampleRate: Int = 16000
+        sampleRate: Int = 16000,
+        chunkSize: Int,
+        chunkBuffer: inout [Float]
     ) throws -> ([TimedSpeakerSegment], ChunkTimings)
     where C: RandomAccessCollection, C.Element == Float, C.Index == Int {
         let segmentationStartTime = Date()
 
-        let chunkSize = sampleRate * 10
         let chunkCount = chunk.distance(from: chunk.startIndex, to: chunk.endIndex)
         let copyCount = min(chunkCount, chunkSize)
 
@@ -278,8 +281,9 @@ public final class DiarizerManager {
             masks.append(speakerMask)
         }
 
-        let embeddings = try embeddingExtractor.getEmbeddings(
-            audio: Array(paddedChunk),
+        let embeddings = try extractEmbeddingsSynchronously(
+            using: embeddingExtractor,
+            audio: paddedChunk,
             masks: masks,
             minActivityThreshold: config.minActiveFramesCount
         )
@@ -342,6 +346,40 @@ public final class DiarizerManager {
         return (segments, timings)
     }
 
+    private func extractEmbeddingsSynchronously<C>(
+        using extractor: EmbeddingExtractor,
+        audio: C,
+        masks: [[Float]],
+        minActivityThreshold: Float
+    ) throws -> [[Float]]
+    where C: RandomAccessCollection, C.Element == Float, C.Index == Int {
+        var embeddingsResult: [[Float]]?
+        var capturedError: Error?
+        let semaphore = DispatchSemaphore(value: 0)
+
+        Task {
+            do {
+                let result = try await extractor.getEmbeddings(
+                    audio: audio,
+                    masks: masks,
+                    minActivityThreshold: minActivityThreshold
+                )
+                embeddingsResult = result
+            } catch {
+                capturedError = error
+            }
+            semaphore.signal()
+        }
+
+        semaphore.wait()
+
+        if let error = capturedError {
+            throw error
+        }
+
+        return embeddingsResult ?? []
+    }
+
     /// Count activity frames per speaker.
     private func calculateSpeakerActivities(_ binarizedSegments: [[[Float]]]) -> [Float] {
         let numSpeakers = binarizedSegments[0][0].count

Sources/FluidAudio/Diarizer/Extraction/EmbeddingExtractor.swift

@@ -3,38 +3,17 @@ import CoreML
 import OSLog
 
 /// Embedding extractor with ANE-aligned memory and zero-copy operations
-public class EmbeddingExtractor {
+public actor EmbeddingExtractor {
     private let wespeakerModel: MLModel
     private let logger = AppLogger(category: "EmbeddingExtractor")
     private let memoryOptimizer = ANEMemoryOptimizer()
 
-    // Pre-allocated ANE-aligned buffers
-    private var waveformBuffer: MLMultiArray?
-    private var maskBuffer: MLMultiArray?
-
-    // Reusable feature providers
-    private var featureProviders: [ZeroCopyDiarizerFeatureProvider] = []
+    private let waveformShape: [NSNumber] = [3, 160_000]
+    private let waveformKey = "wespeaker_waveform_primary"
 
     public init(embeddingModel: MLModel) {
         self.wespeakerModel = embeddingModel
-
-        // Pre-allocate ANE-aligned buffers
-        do {
-            self.waveformBuffer = try memoryOptimizer.createAlignedArray(
-                shape: [3, 160000] as [NSNumber],
-                dataType: .float32
-            )
-
-            self.maskBuffer = try memoryOptimizer.createAlignedArray(
-                shape: [3, 1000] as [NSNumber],  // Typical mask size
-                dataType: .float32
-            )
-        } catch {
-            logger.error("Failed to allocate ANE-aligned buffers: \(error)")
-            // Buffers will remain nil, will be allocated on-demand in getEmbeddings
-        }
-
-        logger.info("EmbeddingExtractor initialized with ANE-aligned buffers")
+        logger.info("EmbeddingExtractor ready with ANE memory optimizer")
     }
 
     /// Extract speaker embeddings using the CoreML embedding model.
@@ -54,18 +33,32 @@ public class EmbeddingExtractor {
         minActivityThreshold: Float = 10.0
     ) throws -> [[Float]]
     where C: RandomAccessCollection, C.Element == Float, C.Index == Int {
-        // We need to return embeddings for ALL speakers, not just active ones
-        // to maintain compatibility with the rest of the pipeline
-        var embeddings: [[Float]] = []
+        guard let firstMask = masks.first else {
+            return []
+        }
+
+        let maskShape = [3, firstMask.count] as [NSNumber]
+        let maskKey = "wespeaker_mask_\(firstMask.count)"
+
+        let waveformLease = try memoryOptimizer.leaseBuffer(
+            key: waveformKey,
+            shape: waveformShape,
+            dataType: .float32
+        )
 
-        // Get or create appropriately sized mask buffer
-        let maskShape = [3, masks[0].count] as [NSNumber]
-        let currentMaskBuffer = try memoryOptimizer.getPooledBuffer(
-            key: "wespeaker_mask_\(masks[0].count)",
+        let maskLease = try memoryOptimizer.leaseBuffer(
+            key: maskKey,
             shape: maskShape,
             dataType: .float32
         )
 
+        let waveformBuffer = waveformLease.multiArray
+        let maskBuffer = maskLease.multiArray
+
+        // We need to return embeddings for ALL speakers, not just active ones
+        // to maintain compatibility with the rest of the pipeline
+        var embeddings: [[Float]] = []
+
         // Process all speakers but optimize for active ones
         for speakerIdx in 0..<masks.count {
             // Check if speaker is active
@@ -80,28 +73,28 @@ public class EmbeddingExtractor {
             // Use ANE-optimized copy for audio data
             memoryOptimizer.optimizedCopy(
                 from: audio,
-                to: waveformBuffer!,
+                to: waveformBuffer,
                 offset: 0  // First speaker slot
             )
 
             // Optimize mask creation with zero-copy view
             fillMaskBufferOptimized(
                 masks: masks,
                 speakerIndex: speakerIdx,
-                buffer: currentMaskBuffer
+                buffer: maskBuffer
             )
 
             // Create zero-copy feature provider
             let featureProvider = ZeroCopyDiarizerFeatureProvider(features: [
-                "waveform": MLFeatureValue(multiArray: waveformBuffer!),
-                "mask": MLFeatureValue(multiArray: currentMaskBuffer),
+                "waveform": MLFeatureValue(multiArray: waveformBuffer),
+                "mask": MLFeatureValue(multiArray: maskBuffer),
             ])
 
             // Run model with optimal prediction options
             let options = MLPredictionOptions()
             // Prefetch to Neural Engine for better performance
-            waveformBuffer!.prefetchToNeuralEngine()
-            currentMaskBuffer.prefetchToNeuralEngine()
+            waveformBuffer.prefetchToNeuralEngine()
+            maskBuffer.prefetchToNeuralEngine()
 
             let output = try wespeakerModel.prediction(from: featureProvider, options: options)
 

Sources/FluidAudio/Diarizer/Offline/Extraction/OfflineEmbeddingExtractor.swift

@@ -259,20 +259,21 @@ struct OfflineEmbeddingExtractor {
         var pldaBatchCallCount = 0
 
         func performEmbeddingWarmup() throws {
-            let warmupAudioArray = try memoryOptimizer.getPooledBuffer(
+            let warmupAudioLease = try memoryOptimizer.leaseBuffer(
                 key: "offline_embedding_warmup_audio",
                 shape: fbankInputShape,
                 dataType: .float32
             )
+            let warmupAudioArray = warmupAudioLease.multiArray
             let warmupAudioPointer = warmupAudioArray.dataPointer.assumingMemoryBound(to: Float.self)
             vDSP_vclr(warmupAudioPointer, 1, vDSP_Length(warmupAudioArray.count))
 
             let warmupFbankFeatures = try runFbankModel(audioArray: warmupAudioArray)
             let zeroWeights = [Float](repeating: 0, count: weightFrameCount)
-            let warmupWeightsArray = try prepareWeightsInput(weights: zeroWeights)
+            let warmupWeightsLease = try prepareWeightsInput(weights: zeroWeights)
             _ = try runEmbeddingModel(
                 fbankFeatures: warmupFbankFeatures,
-                weightsArray: warmupWeightsArray
+                weightsArray: warmupWeightsLease.multiArray
             )
         }
 
@@ -455,10 +456,10 @@ struct OfflineEmbeddingExtractor {
                 }
 
                 let embeddingStart = resampleEnd
-                let weightsArray = try prepareWeightsInput(weights: resampledMask)
+                let weightsLease = try prepareWeightsInput(weights: resampledMask)
                 let embedding256 = try runEmbeddingModel(
                     fbankFeatures: fbankFeatures,
-                    weightsArray: weightsArray
+                    weightsArray: weightsLease.multiArray
                 )
                 let embeddingEnd = clock.now
                 embeddingDuration += embeddingStart.duration(to: embeddingEnd)
@@ -702,12 +703,13 @@ struct OfflineEmbeddingExtractor {
 
     private func prepareWeightsInput(
         weights: [Float]
-    ) throws -> MLMultiArray {
-        let array = try memoryOptimizer.getPooledBuffer(
+    ) throws -> ANEMultiArrayLease {
+        let lease = try memoryOptimizer.leaseBuffer(
             key: "offline_embedding_weights_\(weightFrameCount)",
             shape: weightInputShape,
             dataType: .float32
         )
+        let array = lease.multiArray
 
         let pointer = array.dataPointer.assumingMemoryBound(to: Float.self)
         vDSP_vclr(pointer, 1, vDSP_Length(array.count))
@@ -726,7 +728,7 @@ struct OfflineEmbeddingExtractor {
             }
         }
 
-        return array
+        return lease
     }
 
     private func runEmbeddingModel(