ggml : add ANE backend

[ggerganov]

According to this #336 (comment), there is a new CoreML API and an ANE backend might be possible to implement with latest Apple software/hardware.

[FSSRepo]

It seems that it is only available for matrix multiplications and not for custom operations. Apologies for my ignorance.

[qdrddr]

Would that be of any help the LM Studio has implemented MLX. And here is Anemll ANE library to work with MLX it is MIT Licensed. And there's FastMLX with an Apache 2.0 license.

[BB-fat]

I'm very interested in this direction and would like to share some findings from my experiments:

I created a demo implementation using MLTensor for matrix multiplication and compared it with a Metal implementation.

Example Code:

    func mulMat(a: [Float], aRows: Int, aCols: Int,
               b: [Float], bRows: Int, bCols: Int,
               c: inout [Float]) throws -> Double {
        
        print("CoreML Device: Performing matrix multiplication - Matrix dimensions [\(aRows)x\(aCols)] * [\(bRows)x\(bCols)] = [\(aRows)x\(bCols)]")

        // Validate input parameters
        if a.count < aRows * aCols || b.count < bRows * bCols || c.count < aRows * bCols {
            print("CoreML Device: Input matrix dimensions mismatch")
            throw DeviceError.invalidParameters("Input matrix dimensions mismatch")
        }
        
        // Validate dimension compatibility
        if aCols != bRows {
            print("CoreML Device: Matrix dimensions incompatible, A's columns (\(aCols)) must equal B's rows (\(bRows))")
            throw DeviceError.invalidParameters("Matrix dimensions incompatible")
        }

        if #available(macOS 15.0, *) {
            let startTime = CFAbsoluteTimeGetCurrent()

            let aTensor = MLTensor(shape: [aRows, aCols], scalars: a)
            let bTensor = MLTensor(shape: [bRows, bCols], scalars: b)
            let result = aTensor.matmul(bTensor)
            
            // Use semaphore to wait synchronously for async operation
            let semaphore = DispatchSemaphore(value: 0)
            var resultArray: [Float] = []
            
            Task {
                do {
                    let shapedArray = await result.shapedArray(of: Float.self)
                    resultArray = Array(shapedArray.scalars)
                }
                semaphore.signal()
            }
            
            // Wait for async operation to complete
            semaphore.wait()
            
            // Copy results to output parameter c
            for i in 0..<min(c.count, resultArray.count) {
                c[i] = resultArray[i]
            }

            let endTime = CFAbsoluteTimeGetCurrent()
            let duration = (endTime - startTime) * 1000 // Convert to milliseconds
            print("CoreML Device: Execution time - \(String(format: "%.2f", duration)) ms")
        
            return duration
        } else {
            throw DeviceError.executionFailed("System version lower than macOS 15.0, MLTensor unavailable")
        }
    }

Here's what I discovered:

1. Performance comparison between MLTensor and Metal showed no significant differences for matrix operations.

2. I used Apple's Instruments app to monitor hardware utilization during execution, and interestingly, MLTensor was actually utilizing the GPU rather than ANE for computations.

3. After thoroughly searching through Apple's official documentation (which is quite limited on this topic), I found the withMLTensorComputePolicy API. However, this API only supports two options: "CPU only" and "CPU and GPU".

Based on these observations, I suspect that the current CoreML framework still does not support ANE at the operator level. The API appears to lack explicit control for directing specific operations to the Neural Engine.

If my understanding is incorrect, I would appreciate any clarification.

[optlink]

3. After thoroughly searching through Apple's official documentation (which is quite limited on this topic), I found the withMLTensorComputePolicy API. However, this API only supports two options: "CPU only" and "CPU and GPU".

Based on these observations, I suspect that the current CoreML framework still does not support ANE at the operator level. The API appears to lack explicit control for directing specific operations to the Neural Engine.

If my understanding is incorrect, I would appreciate any clarification.

Forgive me if this is useless as I am not an expert on this topic but this comment and this link show otherwise.

[BB-fat]

@optlink Thank you for sharing the information. I've investigated #336 (comment) and found that the computation still occurs on the GPU even after setting MLComputeUnits.cpuAndNeuralEngine.

I tested with the following code:

resultArray = await withMLTensorComputePolicy(
    MLComputePolicy.init(.cpuAndNeuralEngine)
) {
    let aTensor = MLTensor(shape: [aRows, aCols], scalars: a)
    let bTensor = MLTensor(shape: [bRows, bCols], scalars: b)
    let result = aTensor.matmul(bTensor)
    let shapedArray = await result.shapedArray(of: Float.self)
    return Array(shapedArray.scalars)
}

However, I observed that the operations still executed on the GPU. Even after explicitly setting MLComputeUnits.cpuOnly, the computation didn't shift to the CPU as expected.

The documentation for withMLTensorComputePolicy is quite limited, so I'm not sure if my usage is incorrect or if there's an underlying issue with the API.

[ggerganov]

@BB-fat Did you try different shapes? Maybe the NPU supports only specific shapes (i.e. like multiples of 16/32 etc.).

[BB-fat]

Hi @ggerganov ,
I've tried shapes that are multiples of 16/32, but still can't use the ANE.

[BB-fat]

I found a reply from an Apple engineer on the Apple Developer Forum, which included a sample code snippet containing withMLTensorComputePolicy. It appears to be consistent with my usage.

[giladgd]

@BB-fat I've created a post on Apple Developer Forums to try and get some help from Apple engineers regarding this:
https://developer.apple.com/forums/thread/775589