Example code for MPS graph not working

I'm trying out MPS Graph today and faced the exact same issue as yours. Crashes on commit, wrapping it in checks and it will stuck at encoding for Iteration 64, losses are basically 0.0, empty documentation, and the forum post left unreplied for a month. I really wonder if this feature is production ready for developers since it looks nice on the surface and works really terrible on the example project.

My setup: Version 13.4.1 (13F100) macOS 12.5 MacBook Air (M1) 2020 8GB

Maybe it's the asynchronous part that is broken. I later changed the program to use the blocking graph.run, signal for completion by hand, and signal the semaphore by hand. With a blocking MPS Graph demo it seems to run properly. But the documentation issue is still very nasty.

Hello there! I tried my best again and here what I got:

First

If you replace MTLCommandBuffer with MPSCommandBuffer

    func encodeTrainingBatch(commandBuffer: MPSCommandBuffer,
                             sourceTensorData: MPSGraphTensorData,
                             labelsTensorData: MPSGraphTensorData,
                             completion: ((Float) -> Void)?) -> MPSGraphTensorData {

and call it like this:

        let commandBuffer = gCommandQueue.makeCommandBuffer()!
        let mpsCommandBuffer = MPSCommandBuffer(commandBuffer: commandBuffer)

        var yLabels: MPSNDArray? = nil
        let xInput = dataset.getRandomTrainingBatch(device: gDevice, batchSize: batchSize, labels: &yLabels)

        _ = classiferGraph.encodeTrainingBatch(commandBuffer: mpsCommandBuffer,
                                               sourceTensorData: MPSGraphTensorData(xInput),
                                               labelsTensorData: MPSGraphTensorData(yLabels!),
                                               completion: updateProgressCubeAndLoss)
        mpsCommandBuffer.commitAndContinue()

it starts working... But only on iPad version on Mac :D It won't do anything on physical iPhone device (more on this down the line)

Second one

If you replace graph.encode(to: to graph.runAsync( you no longer need to control commandBuffer and it's also solves the issue

Third

As I mentioned previously, you can replace MTLCommandBuffer with MPSCommandBuffer it will solve the issue on iPad build running on macOS. BUT it not working on physical iPhone. I mean, it doesn't crash, but descriptor not calling completionHandler.

BUT

If you add completion handler with addCompletedHandler on commandBuffer, it WILL run!

Appendix

I don't know why it works the way I explained before, but it seems that for now best way is to use runAsync because it does not need MPSCommandBuffer which cause the bug #1 and it calls completionHandler of MPSGraphExecutionDescriptor in every configuration

Hope it will help someone and I hope for Apple developers answers >_>

Hi, we already have a fix on the way to the sample code which had a minor bug, attaching the diff below for anyone to be unblocked, please apply it with:

git apply fix.diff

diff --git a/MPSGraphClassifier/MNISTClassifierGraph.swift b/MPSGraphClassifier/MNISTClassifierGraph.swift
index 0f93b6d..2d6e6ef 100644
--- a/MPSGraphClassifier/MNISTClassifierGraph.swift
+++ b/MPSGraphClassifier/MNISTClassifierGraph.swift
@@ -209,7 +209,7 @@ class MNISTClassifierGraph: NSObject {
     let doubleBufferingSemaphore = DispatchSemaphore(value: 2)
 
     // Encode training batch to command buffer using double buffering
-    func encodeTrainingBatch(commandBuffer: MTLCommandBuffer,
+    func encodeTrainingBatch(commandBuffer: MPSCommandBuffer,
                              sourceTensorData: MPSGraphTensorData,
                              labelsTensorData: MPSGraphTensorData,
                              completion: ((Float) -> Void)?) -> MPSGraphTensorData {
@@ -237,7 +237,7 @@ class MNISTClassifierGraph: NSObject {
         let feed = [sourcePlaceholderTensor: sourceTensorData,
                     labelsPlaceholderTensor: labelsTensorData]
         
-        let fetch = graph.encode(to: MPSCommandBuffer(commandBuffer: commandBuffer),
+        let fetch = graph.encode(to: commandBuffer,
                                  feeds: feed,
                                  targetTensors: targetTrainingTensors,
                                  targetOperations: targetTrainingOps,
@@ -247,7 +247,7 @@ class MNISTClassifierGraph: NSObject {
     }
 
     // Encode inference batch to command buffer using double buffering
-    func encodeInferenceBatch(commandBuffer: MTLCommandBuffer,
+    func encodeInferenceBatch(commandBuffer: MPSCommandBuffer,
                               sourceTensorData: MPSGraphTensorData,
                               labelsTensorData: MPSGraphTensorData) -> MPSGraphTensorData {
         doubleBufferingSemaphore.wait()
@@ -286,7 +286,7 @@ class MNISTClassifierGraph: NSObject {
             self.doubleBufferingSemaphore.signal()
         }
         
-        let fetch = graph.encode(to: MPSCommandBuffer(commandBuffer: commandBuffer),
+        let fetch = graph.encode(to: commandBuffer,
                                   feeds: [sourcePlaceholderTensor: sourceTensorData,
                                           labelsPlaceholderTensor: labelsTensorData],
                                   targetTensors: targetInferenceTensors,
diff --git a/MPSGraphClassifier/ViewController.swift b/MPSGraphClassifier/ViewController.swift
index 6efc787..2aafe88 100644
--- a/MPSGraphClassifier/ViewController.swift
+++ b/MPSGraphClassifier/ViewController.swift
@@ -247,7 +247,7 @@ class ViewController: UIViewController, CanvasDelegate {
     
     // Run a training iteration batch
     func runTrainingIterationBatch() -> MTLCommandBuffer {
-        let commandBuffer = gCommandQueue.makeCommandBuffer()!
+        let commandBuffer = MPSCommandBuffer(commandBuffer: gCommandQueue.makeCommandBuffer()!)
         
         var yLabels: MPSNDArray? = nil
         let xInput = dataset.getRandomTrainingBatch(device: gDevice, batchSize: batchSize, labels: &yLabels)
@@ -297,7 +297,7 @@ class ViewController: UIViewController, CanvasDelegate {
         
         // encoding each image
         for currImageIdx in stride(from: 0, to: dataset.totalNumberOfTestImages, by: Int(batchSize)) {
-            let commandBuffer = gCommandQueue.makeCommandBuffer()!
+            let commandBuffer = MPSCommandBuffer(commandBuffer: gCommandQueue.makeCommandBuffer()!)
             
             xInput = dataset.getTrainingBatchWithDevice(device: gDevice,
                                                         batchIndex: Int(currImageIdx) / Int(batchSize),