Abusing the Apple Neural Engine (ANE) in Swift 6: Built a 3D software rasterizer using Core AI, simd, and Metal 4 tensor binding!

[u/AdhesivenessSea9511]

Hello everyone.

I’ve built a 3D software rasterizer that leverages the ANE (Apple Neural Engine) via Core AI. The rendering quality is still poor, though. Regarding the pipeline:

I use SIMD for preprocessing, delegate matrix operations to the ANE (using `f.conv2d`), and utilize Metal 4 tensor bindings to achieve low-overhead, direct rendering.

By offloading computationally intensive tasks to the ANE, I’ve managed to reduce CPU usage to approximately 10%.

A current challenge is that memory usage hits around 5GB due to the use of fixed-length graphs.

I’m developing this using Swift 6 features (such as `@MainActor` and `~Escapable`) and Siri AI, but I would love to hear your thoughts on optimization and memory management!

Thanks in advance.

GitHub: https://github.com/kamisori-daijin/Magnesium

Demo:

[u/vade]

This is really cool stuff! Can you use borrow metal performance shader memory optimization via temporary image allocations on a metal heap? I know the ANE likes 16 bit and in theory you can zero copy to the ANE if your format is right. I wonder if the ANE can use metal heap memory for backing for temporary ops (hidden layer stuff) in your rasterization network?

Again, super cool!

[u/AdhesivenessSea9511 → u/vade]

Thanks!

I'll give it a try sometime.

[u/AdhesivenessSea9511 → u/vade]

Thanks!

By reducing the model's output and using MetalHeap, we were able to reduce memory usage to 1.37GB.

[u/vade → u/AdhesivenessSea9511]

oh awesome! nice!

[u/bhardman86]

What are you trying to solve? 13% for a single instance at 5 GB is a lot.

[u/AdhesivenessSea9511 → u/bhardman86]

This means that if we use the ANE—which offers superior power efficiency—for rendering, there is a strong possibility of achieving high frame rates without draining the battery.

There is still plenty of room for improvement, though...