WWDC Quick Look 💓 By SwiftGGTeam
中文 EN 日本語
What’s new in BNNS Graph

What’s new in BNNS Graph

Machine Learning
Watch original video

Highlight

BNNSGraph debuted last year with a file-based API that loaded a graph from a CoreML package and ran inference. This year’s BNNSGraphBuilder is a brand-new Swift DSL. You define the compute graph directly in Swift code — argument declarations, operator calls, output returns — all inside a single makeContext closure.

Core content

Last year’s Bitcrusher demo showed how BNNSGraph replaced the chore of writing discrete layers by hand, but you still had to walk the PyTorch → CoreML package → mlmodelc pipeline. For an audio effect with only a dozen tensor ops, maintaining a separate Python file and building a CoreML package costs more than it saves. For pre- and post-processing graphs, developers kept bouncing between writing Swift code that called BNNS primitives one by one and detouring through PyTorch to compile a model.

This year Apple flattens the whole flow. BNNSGraphBuilder lets you define the graph inside a Swift closure: declare an argument, call operators, return the output tensors. One BNNSGraph.makeContext call gives you an executable context. The code lives in the same source file as the rest of your Swift code. Type checks happen at compile time. Once a tensor shape is known at Swift runtime, you can pass it into the graph and get the speed of static shapes. Combined with BNNSGraph’s existing layer fusion, copy elision, and weight repacking, small models and pre/post-processing graphs no longer need PyTorch as a middleman.

Detailed content

The entry point of the API is the new type method BNNSGraph.makeContext. The builder closure parameter is where you declare inputs, assemble operators, and return outputs. The snippet below is the smallest runnable example from the session at 8:31 (8:31). It computes the element-wise product of two 8-element vectors and takes the mean:

let context = try BNNSGraph.makeContext { builder in
    let x = builder.argument(name: "x",
                             dataType: Float.self,
                             shape: [8])
    let y = builder.argument(name: "y",
                             dataType: Float.self,
                             shape: [8])

    let product = x * y
    let mean = product.mean(axes: [0], keepDimensions: true)

    // Prints "shape: [1] | stride: [1]".
    print("mean", mean)

    return [product, mean]
}

var args = context.argumentNames().map { name in
    context.tensor(argument: name, fillKnownDynamicShapes: false)!
}
args[0].allocate(as: Float.self, count: 8)   // product
args[1].allocate(as: Float.self, count: 1)   // mean
args[2].allocate(initializingFrom: [1, 2, 3, 4, 5, 6, 7, 8] as [Float])
args[3].allocate(initializingFrom: [8, 7, 6, 5, 4, 3, 2, 1] as [Float])

try context.executeFunction(arguments: &args)

Key points:

  • BNNSGraph.makeContext { builder in ... } compiles a Swift closure into a reusable context. You usually build it once at app launch.
  • builder.argument(name:dataType:shape:) declares an input tensor. shape: [8] is a static shape, so no runtime inference is needed.
  • x * y uses a Swift operator to trigger element-wise multiplication, equivalent to BNNS’s element-wise multiply primitive.
  • product.mean(axes:keepDimensions:) calls a reduction operator. keepDimensions: true keeps the reduced axis and yields a [1] shape.
  • print("mean", mean) prints the shape and stride of an intermediate tensor at compile/build time, which helps debugging.
  • return [product, mean] decides the output tensors. context.argumentNames() lists outputs first, inputs second.
  • executeFunction(arguments:) actually runs the graph and writes results into the output tensors you allocated.

The type system rejects implicit data-type mixing at compile time. The strongly typed example at 12:04 (12:04) shows the form bases.pow(y: exponents.cast(to: Float16.self)) — a Float16 base and an Int32 exponent must be cast explicitly, or makeContext fails to compile. The boolean tensor returned by the comparison mask0 .< mask1 likewise needs cast(to: Float16.self) before it can multiply with an FP16 tensor.

Slicing uses Swift subscript syntax (14:15). BNNSGraph.Builder.SliceRange(startIndex:endIndex:) supports negative indices that count from the tail. SliceRange.fillAll keeps the whole dimension. Underneath, BNNSGraph treats the slice as a reference to the original data — zero copy, zero extra allocation.

The post-processing example (19:04) combines softmax(axis:) with topK(_:axis:findLargest:). The latter returns a tuple of values and indices — the typical decode flow for a classification model.

Migrating Bitcrusher is the clearest before/after (21:03). The old PyTorch version inherits nn.Module, defines forward, and runs through coremltools. The Swift version calls source * saturationGain, destination.tanh(), destination.round() directly on the tensors. The mapping to PyTorch is nearly one-to-one, but everything lives in a single Swift file. Switch precision by changing BITCRUSHER_PRECISION from Float32 to Float16 (22:34). The session measured FP16 as faster than FP32.

Core takeaways

  • What to do: Move the pre/post-processing of an ML model from Python to Swift and rewrite it with BNNSGraphBuilder. Why it pays off: Pre-processing (cropping, normalization, thresholding) and post-processing (softmax, topK, NMS) usually involve only a dozen operators, yet they sit in a separate Python project. Once moved into Swift, they live alongside your business code. Compile-time type checks catch a lot of shape bugs. How to start: Pick a tiny pre-processing flow (say, center-crop plus normalize), rewrite it with BNNSGraph.makeContext, and verify the numbers against the original Python output.

  • What to do: When you build a custom Audio Unit for Logic Pro or GarageBand, write the DSP graph directly in Swift. Why it pays off: BNNS targets the CPU, low latency, and predictable memory and threading — exactly what real-time audio needs. Once an effect like Bitcrusher is expressed entirely in BNNSGraphBuilder, you no longer have to maintain a PyTorch project. How to start: Translate your existing saturation, quantization, and dry-wet mix formulas into builder ops like *, tanh(), round(). Default BITCRUSHER_PRECISION to Float16.

  • What to do: Use BNNSGraphBuilder instead of the CoreML package path for small-model scenarios. Why it pays off: The CoreML package path is mature and solid, but it requires the Python toolchain. When the model has only a few layers and parameters depend on runtime variables, the file-based API becomes a burden. The Swift DSL can build a static-shape graph at runtime from values you already know, which beats flexible-shape performance. How to start: Debug shapes with print(mean) and similar prints on intermediate tensors. Once they look right, run end-to-end inference with context.executeFunction(arguments:).

Comments

GitHub Issues · utterances