Scheduler performance optimizations in Swift

Author: alejandro-isaza

swift/StableDiffusion/pipeline/DPMSolverMultistepScheduler.swift

@@ -83,14 +83,20 @@ public final class DPMSolverMultistepScheduler: Scheduler {
     /// Convert the model output to the corresponding type the algorithm needs.
     /// This implementation is for second-order DPM-Solver++ assuming epsilon prediction.
     func convertModelOutput(modelOutput: MLShapedArray<Float32>, timestep: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
-        assert(modelOutput.scalars.count == sample.scalars.count)
+        assert(modelOutput.scalarCount == sample.scalarCount)
+        let scalarCount = modelOutput.scalarCount
         let (alpha_t, sigma_t) = (self.alpha_t[timestep], self.sigma_t[timestep])
-        
-        // This could be optimized with a Metal kernel if we find we need to
-        let x0_scalars = zip(modelOutput.scalars, sample.scalars).map { m, s in
-            (s - m * sigma_t) / alpha_t
+
+        return MLShapedArray(unsafeUninitializedShape: modelOutput.shape) { scalars, _ in
+            assert(scalars.count == scalarCount)
+            modelOutput.withUnsafeShapedBufferPointer { modelOutput, _, _ in
+                sample.withUnsafeShapedBufferPointer { sample, _, _ in
+                    for i in 0 ..< scalarCount {
+                        scalars.initializeElement(at: i, to: (sample[i] - modelOutput[i] * sigma_t) / alpha_t)
+                    }
+                }
+            }
         }
-        return MLShapedArray(scalars: x0_scalars, shape: modelOutput.shape)
     }
 
     /// One step for the first-order DPM-Solver (equivalent to DDIM).

swift/StableDiffusion/pipeline/Scheduler.swift

@@ -1,6 +1,7 @@
 // For licensing see accompanying LICENSE.md file.
 // Copyright (C) 2022 Apple Inc. All Rights Reserved.
 
+import Accelerate
 import CoreML
 
 @available(iOS 16.2, macOS 13.1, *)
@@ -59,18 +60,21 @@ public extension Scheduler {
     ///   - values: The arrays to be weighted and summed
     /// - Returns: sum_i weights[i]*values[i]
     func weightedSum(_ weights: [Double], _ values: [MLShapedArray<Float32>]) -> MLShapedArray<Float32> {
+        let scalarCount = values.first!.scalarCount
         assert(weights.count > 1 && values.count == weights.count)
-        assert(values.allSatisfy({ $0.scalarCount == values.first!.scalarCount }))
-        var w = Float(weights.first!)
-        var scalars = values.first!.scalars.map({ $0 * w })
-        for next in 1 ..< values.count {
-            w = Float(weights[next])
-            let nextScalars = values[next].scalars
-            for i in 0 ..< scalars.count {
-                scalars[i] += w * nextScalars[i]
+        assert(values.allSatisfy({ $0.scalarCount == scalarCount }))
+
+        return MLShapedArray(unsafeUninitializedShape: values.first!.shape) { scalars, _ in
+            scalars.initialize(repeating: 0.0)
+            for i in 0 ..< values.count {
+                let w = Float(weights[i])
+                values[i].withUnsafeShapedBufferPointer { buffer, _, _ in
+                    assert(buffer.count == scalarCount)
+                    // scalars[j] = w * values[i].scalars[j]
+                    cblas_saxpy(Int32(scalarCount), w, buffer.baseAddress, 1, scalars.baseAddress, 1)
+                }
             }
         }
-        return MLShapedArray(scalars: scalars, shape: values.first!.shape)
     }
 
     func addNoise(

swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift

@@ -286,20 +286,19 @@ public struct StableDiffusionPipeline: ResourceManaging {
     }
 
     func performGuidance(_ noise: MLShapedArray<Float32>, _ guidanceScale: Float) -> MLShapedArray<Float32> {
-
-        let blankNoiseScalars = noise[0].scalars
-        let textNoiseScalars = noise[1].scalars
-
-        var resultScalars =  blankNoiseScalars
-
-        for i in 0..<resultScalars.count {
-            // unconditioned + guidance*(text - unconditioned)
-            resultScalars[i] += guidanceScale*(textNoiseScalars[i]-blankNoiseScalars[i])
-        }
-
         var shape = noise.shape
         shape[0] = 1
-        return MLShapedArray<Float32>(scalars: resultScalars, shape: shape)
+        return MLShapedArray<Float>(unsafeUninitializedShape: shape) { result, _ in
+            noise.withUnsafeShapedBufferPointer { scalars, _, strides in
+                for i in 0 ..< result.count {
+                    // unconditioned + guidance*(text - unconditioned)
+                    result.initializeElement(
+                        at: i,
+                        to: scalars[i] + guidanceScale * (scalars[strides[0] + i] - scalars[i])
+                    )
+                }
+            }
+        }
     }
 
     func decodeToImages(_ latents: [MLShapedArray<Float32>], configuration config: Configuration) throws -> [CGImage?] {