Support real-time ML inference on the CPU
System Frameworks 进阶 20m

BNNS Graph：在 CPU 上跑实时机器学习推理

Support real-time ML inference on the CPU

2024年6月10日

在 Apple 官方观看视频

一句话判断

如果你需要做音频或其他实时场景的 ML 推理且不能依赖 GPU，BNNS Graph 提供了比传统 BNNS primitives 平均快 2 倍的 CPU 推理能力——自动优化、零代码改动就能受益。

这场 Session 讲了什么

BNNS（Basic Neural Network Subroutines）是 Apple Accelerate 框架中的机器学习库，长期以来提供 CPU 上的 ML 推理和训练能力。今年 Apple 在此基础上推出了 BNNS Graph——一个全新的图级别 API，能够消费整个计算图而非单个操作原语。

演讲者 Simon Gladman 首先回顾了传统 BNNS 的使用方式：你需要为每一层手动创建 n 维数组描述符（输入、输出、权重、偏置），然后用这些描述符构建参数结构体，再创建层对象，最后逐层执行推理。如果你要实现一个完整的模型，每一层都要重复这个流程。

BNNS Graph 的使用方式完全不同：你提供一个编译好的 Core ML 模型文件（.mlmodelc），BNNS Graph 自动将整个模型编译成优化后的图对象，包含推理需要执行的 kernel 列表和中间张量的内存布局。你只需要构建一次图对象，然后用一个 context 执行推理。

Session 还重点讲解了 BNNS Graph 在实时场景（如音频处理 Audio Unit）中的应用——关键约束是推理阶段不能有任何内存分配和多线程操作，BNNS Graph 提供了精细的控制来满足这些要求。

值得深挖的点

四种自动优化：为什么图级别 API 比逐层 API 快

BNNS Graph 相比传统 BNNS 的性能提升来自四种自动优化，不需要你写任何额外代码。

第一种是数学变换（mathematical transformation）。如果你的模型最后有一个 slice 操作只取张量的一部分，BNNS Graph 会把这个 slice 移到模型开头——只计算需要的那部分元素，而不是先算完整张量再截取。

第二种是层融合（layer fusion）。BNNS Graph 会把连续的操作合并成一个。比如卷积后面紧跟激活函数，这两个操作在传统 API 中是分别执行的（中间结果需要写入内存再读出），BNNS Graph 把它们合成一个操作，省去了中间的内存读写。

第三种是拷贝消除（copy elision）。slice 操作通常需要把数据拷贝到新的张量，BNNS Graph 优化为传递原数据的窗口视图（view），避免不必要的拷贝。

第四种是权重重排（weight repacking）。BNNS Graph 会根据 CPU 缓存特性重新排列权重数据的内存布局——从行优先重排为分块迭代顺序，提高缓存命中率。

这四种优化叠加后，BNNS Graph 平均比传统 BNNS primitives 快 2 倍以上。

实时音频场景的严格约束

音频处理场景对 ML 推理有特殊要求：执行阶段（execute phase）绝对不能有内存分配或多线程调度。原因是这些操作可能触发操作系统上下文切换到内核态，导致实时截止时间被违反——用户听到的就是音频卡顿或爆音。

BNNS Graph 针对实时场景提供了精细控制。图编译阶段（build graph）可以正常使用内存分配和多线程——这个阶段只在初始化时执行一次。推理阶段（execute context）则使用预分配的内存和单线程模式。你还可以在创建图时指定编译选项来控制内存策略和线程行为。

代码片段

BNNS Graph 的基本使用流程

import Accelerate

// 第一步：从编译好的 Core ML 模型创建图对象
// mlModelcPath 是 Xcode 自动编译 mlpackage 生成的路径
let graphURL = Bundle.main.url(forResource: "Bitcrusher", 
                                withExtension: "mlmodelc")!

// 编译图对象（只需执行一次，可以放在后台线程）
let graph = try BNNSGraph.compile(
    modelAt: graphURL,
    options: .default
)

// 第二步：创建推理上下文
let context = try BNNSGraphContext(graph: graph)

// 第三步：执行推理（可重复调用）
// 输入和输出张量需要预分配
let input = BNNSNDArrayDescriptor(
    shape: .shape(1, 256),
    dataType: .float
)
let output = BNNSNDArrayDescriptor(
    shape: .shape(1, 256),
    dataType: .float
)

// 执行推理
try context.execute(
    inputs: [input],
    outputs: [output]
)

坑点：BNNSGraph.compile 是耗时操作，不要在音频回调线程中调用。通常在 App 启动后或用户第一次使用 ML 功能时在后台线程编译。

在 Audio Unit 中使用 BNNS Graph

// Audio Unit 的实时处理回调中执行推理
class BitcrusherAudioUnit: AUAudioUnit {
    private var graph: BNNSGraph?
    private var context: BNNSGraphContext?
    private var inputBuffer: BNNSNDArrayDescriptor?
    private var outputBuffer: BNNSNDArrayDescriptor?
    
    // 初始化时编译图（非实时线程）
    func setupGraph() {
        DispatchQueue.global(qos: .userInitiated).async {
            self.graph = try? BNNSGraph.compile(
                modelAt: self.modelURL,
                // 实时场景的关键选项
                options: [.singleThreaded, .noMemoryAllocation]
            )
            if let graph = self.graph {
                self.context = try? BNNSGraphContext(graph: graph)
            }
            // 预分配输入输出缓冲区
            self.inputBuffer = BNNSNDArrayDescriptor(
                shape: .shape(1, 512),
                dataType: .float
            )
            self.outputBuffer = BNNSNDArrayDescriptor(
                shape: .shape(1, 512),
                dataType: .float
            )
        }
    }
    
    // 实时音频回调（不能有内存分配）
    func processAudio(input: UnsafePointer<Float>, 
                      output: UnsafeMutablePointer<Float>,
                      frameCount: UInt32) {
        guard let context = context,
              let inputBuf = inputBuffer,
              let outputBuf = outputBuffer else { return }
        
        // 将音频数据复制到预分配的缓冲区
        inputBuf.data(ofShape: .shape(1, frameCount))
            .initialize(from: input)
        
        // 执行推理——无内存分配、单线程
        try? context.execute(
            inputs: [inputBuf],
            outputs: [outputBuf]
        )
        
        // 将结果复制到输出
        outputBuf.data(ofShape: .shape(1, frameCount))
            .copyBytes(to: output)
    }
}

坑点：processAudio 在实时音频线程中运行，任何可能导致阻塞的操作都会造成音频故障。确保 context.execute 使用的是预编译的、无内存分配的模式。

用 SwiftUI 和 Swift Charts 可视化推理结果

// 将 BNNS Graph 的推理结果用 Swift Charts 展示
struct AudioVisualizationView: View {
    let samples: [Float]
    
    var body: some View {
        Chart(Array(samples.enumerated()), id: \.offset) { index, value in
            LineMark(
                x: .value("时间", index),
                y: .value("振幅", Double(value))
            )
            .foregroundStyle(.blue)
        }
        .chartYScale(domain: -1...1)
        .frame(height: 200)
    }
}

最佳实践

如果你的 App 已经在使用 Core ML 做 CPU 推理，BNNS Graph 是一个值得评估的替代方案。它不需要你改变模型训练流程——同样的 Core ML 模型文件（.mlpackage / .mlmodelc）可以直接被 BNNS Graph 消费。迁移的主要工作是用 BNNS Graph 的 API 替换 Core ML 的推理调用。

对于实时音频场景，关键的原则是”编译在后台、推理在前台”。图对象的编译是一次性耗时操作，放在 App 启动后的后台线程。推理上下文创建后就可以反复使用，不需要重新编译。音频回调中只做 context.execute 调用，其他所有准备工作都在初始化阶段完成。

还有什么值得关注

  • BNNS Graph 是 Accelerate 框架的一部分，不需要额外的框架依赖。
  • Core ML 在内部已经使用 BNNS 做 CPU 推理加速，BNNS Graph 给了你更直接的控制和更细的性能调优能力。
  • Xcode 的 Audio Unit Extension App 模板可以快速创建包含 BNNS Graph 的音频处理项目。
WWDC 2024