Add MetalPerformanceShadersGraph Bindings (#14303)

I'm very pleased to present full bindings to the MetalPerformanceShadersGraph framework!

I'm happy with how everything turned out with the exception of a few notes and questions below.

I re-implemented Apple's MNIST sample (from https://developer.apple.com/documentation/metalperformanceshadersgraph/training_a_neural_network_using_mps_graph) here:

https://gist.github.com/praeclarum/b8077771fb341a1f9c28240113e00425

It's also added as a unit test.

Fixes #14286

### Notes

* Although the API says it works on macOS 11, it has bugs and crashes with errors even with Apple’s Swift examples. It’s better on macOS 12. iOS 14 and on is fine.

* `MPSGraphSparseStorageType` has terrible names. They match Apple's but I wish they were better.

* I added convenience methods to `MPSNDArray` and `MPSGrapTensorData` and the `Variable` and `Constant` operations to decrease the amount of unsafe code users have to write. I currently do this for 32-bit floats, the most common data type.

Co-authored-by: Alex Soto <[email protected]>
Co-authored-by: Rolf Bjarne Kvinge <[email protected]>
Co-authored-by: Manuel de la Pena <[email protected]>

CODEOWNERS

@@ -126,6 +126,8 @@
 /src/metalkit.cs
 /src/MetalPerformanceShaders
 /src/metalperformanceshaders.cs
+/src/MetalPerformanceShadersGraph @praeclarum
+/src/metalperformanceshadersgraph.cs @praeclarum
 /src/MobileCoreServices
 /src/mobilecoreservices.cs
 /src/ModelIO

src/Foundation/NSObject.mac.cs

@@ -104,6 +104,7 @@ public partial class NSObject : INativeObject
 		static IntPtr ios = Dlfcn.dlopen (Constants.IOSurfaceLibrary, 1);
 		static IntPtr ex = Dlfcn.dlopen (Constants.ExternalAccessoryLibrary, 1);
 		static IntPtr ms = Dlfcn.dlopen (Constants.MetalPerformanceShadersLibrary, 1);
+		static IntPtr msg = Dlfcn.dlopen (Constants.MetalPerformanceShadersGraphLibrary, 1);
 		static IntPtr bc = Dlfcn.dlopen (Constants.BusinessChatLibrary, 1);
 		static IntPtr ad = Dlfcn.dlopen (Constants.AdSupportLibrary, 1);
 		static IntPtr nl = Dlfcn.dlopen (Constants.NaturalLanguageLibrary, 1);

src/MetalPerformanceShaders/MPSNDArray.cs

@@ -5,6 +5,18 @@
 namespace MetalPerformanceShaders {
 	public partial class MPSNDArray {
 
+		public static MPSNDArray Create (IMTLDevice device, ReadOnlySpan<float> values, params int[] shape)
+		{
+			var ushape = new nuint [shape.Length];
+			for (var i = 0; i < shape.Length; i++) {
+				ushape [i] = (nuint) shape [i];
+			}
+			var desc = MPSNDArrayDescriptor.Create (MPSDataType.Float32, ushape);
+			var ndarray = new MPSNDArray (device, desc);
+			ndarray.Write (values);
+			return ndarray;
+		}
+
 		public void ExportData (IMTLCommandBuffer cmdBuf, IMTLBuffer buffer, MPSDataType sourceDataType, nuint offset)
 		{
 			ExportData (cmdBuf, buffer, sourceDataType, offset, IntPtr.Zero);
@@ -37,6 +49,21 @@ public unsafe void WriteBytes (IntPtr buffer, nint[] strideBytesPerDimension)
 				WriteBytes (buffer, (IntPtr)p);
 			}
 		}
+		public unsafe void Write (ReadOnlySpan<float> values)
+		{
+			if (DataType != MPSDataType.Float32)
+				throw new InvalidOperationException($"Attempted to write array data of type {DataType} to span of Float32s.");
+			nuint length = 1;
+			var ndims = NumberOfDimensions;
+			for (nuint i = 0; i < ndims; i++) {
+				length *= GetLength (i);
+			}
+			if (length != (nuint) values.Length)
+				throw new ArgumentException ($"The number of values ({values.Length}) does not match the shape length ({length}).");
+			fixed (float* p = values) {
+				WriteBytes ((IntPtr) p, strideBytesPerDimension: IntPtr.Zero);
+			}
+		}
 
 		public void ReadBytes (IntPtr buffer)
 		{
@@ -48,5 +75,20 @@ public unsafe void ReadBytes (IntPtr buffer, nint[] strideBytesPerDimension)
 				ReadBytes (buffer, (IntPtr)p);
 			}
 		}
+		public unsafe void Read (Span<float> values)
+		{
+			if (DataType != MPSDataType.Float32)
+				throw new InvalidOperationException ($"Attempted to read array data of type {DataType} to span of Float32s.");
+			nuint length = 1;
+			var ndims = NumberOfDimensions;
+			for (nuint i = 0; i < ndims; i++) {
+				length *= GetLength (i);
+			}
+			if (length != (nuint) values.Length)
+				throw new ArgumentException ($"The number of values ({values.Length}) does not match the shape length ({length}).");
+			fixed (float* p = values) {
+				ReadBytes ((IntPtr) p, strideBytesPerDimension: IntPtr.Zero);
+			}
+		}
 	}
 }

src/MetalPerformanceShadersGraph/MPSGraphEnums.cs

@@ -0,0 +1,117 @@
+using System;
+using System.Runtime.InteropServices;
+
+using Foundation;
+using ObjCRuntime;
+using Metal;
+
+namespace MetalPerformanceShadersGraph
+{
+	[Flags]
+	public enum MPSGraphOptions : ulong
+	{
+		None = 0,
+		SynchronizeResults = 1,
+		Verbose = 2,
+		Default = SynchronizeResults,
+	}
+
+	[Native]
+	public enum MPSGraphTensorNamedDataLayout : ulong
+	{
+		Nchw = 0,
+		Nhwc = 1,
+		Oihw = 2,
+		Hwio = 3,
+		Chw = 4,
+		Hwc = 5,
+		Hw = 6,
+	}
+
+	[Native]
+	public enum MPSGraphPaddingStyle : ulong
+	{
+		Explicit = 0,
+		Valid = 1,
+		Same = 2,
+		ExplicitOffset = 3,
+	}
+
+	[Native]
+	public enum MPSGraphPaddingMode : long
+	{
+		Constant = 0,
+		Reflect = 1,
+		Symmetric = 2,
+		ClampToEdge = 3,
+		Zero = 4,
+		Periodic = 5,
+		AntiPeriodic = 6,
+	}
+
+	[Native]
+	public enum MPSGraphReductionMode : ulong
+	{
+		Min = 0,
+		Max = 1,
+		Sum = 2,
+		Product = 3,
+		ArgumentMin = 4,
+		ArgumentMax = 5,
+	}
+
+	[Native]
+	public enum MPSGraphResizeMode : ulong
+	{
+		Nearest = 0,
+		Bilinear = 1,
+	}
+
+	[Native]
+	public enum MPSGraphScatterMode : long
+	{
+		Add = 0,
+		Sub = 1,
+		Mul = 2,
+		Div = 3,
+		Min = 4,
+		Max = 5,
+		Set = 6,
+	}
+
+	public enum MPSGraphDeviceType : uint
+	{
+		Metal = 0,
+	}
+
+	public enum MPSGraphLossReductionType : ulong
+	{
+		Axis = 0,
+		Sum = 1,
+		Mean = 2,
+	}
+
+	// For COO, indexTensor0 is x index and indexTensor1 is y index
+	// For CSC, indexTensor0 and indexTensor1 correspond to rowIndex and colStarts respectively.
+	// For CSR, indexTensor0 and indexTensor1 correspond to colIndex and rowStarts respectively.
+	public enum MPSGraphSparseStorageType : ulong
+	{
+		Coo = 0,
+		Csc = 1,
+		Csr = 2,
+	}
+
+	public enum MPSGraphRandomDistribution : ulong
+	{
+		Uniform = 0,
+		Normal = 1,
+		TruncatedNormal = 2,
+	}
+
+	public enum MPSGraphRandomNormalSamplingMethod : ulong
+	{
+		InvCdf = 0,
+		BoxMuller = 1,
+	}
+
+}

src/MetalPerformanceShadersGraph/MPSGraphExtensions.cs

@@ -0,0 +1,60 @@
+#nullable enable
+
+using System;
+using System.Buffers;
+using System.Runtime.InteropServices;
+
+using Foundation;
+using ObjCRuntime;
+using Metal;
+using MetalPerformanceShaders;
+
+namespace MetalPerformanceShadersGraph
+{
+	public static partial class MPSGraphMemoryOps_Extensions
+	{
+		public static unsafe MPSGraphTensor Constant (this MPSGraph graph, float scalar)
+		{
+			return graph.Constant ((double) scalar, new [] { 1 }, MPSDataType.Float32);
+		}
+
+		public static unsafe MPSGraphTensor Constant (this MPSGraph graph, ReadOnlySpan<float> values, int[] shape)
+		{
+			var length = 1;
+			for (var i = 0; i < shape.Length; i++)
+				length *= shape [i];
+			if (length != values.Length)
+				throw new ArgumentException ($"The number of values ({values.Length}) does not match the shape length ({length}).");
+			fixed (float* p = values) {
+				using var data = NSData.FromBytesNoCopy ((IntPtr) p, (nuint) (values.Length * 4), freeWhenDone: false);
+				return graph.Constant (data, shape, MPSDataType.Float32);
+			}
+		}
+
+		public static MPSGraphTensor Variable (this MPSGraph graph, float initialValue, int[] shape, string? name = null)
+		{
+			var length = 1;
+			for (var i = 0; i < shape.Length; i++)
+				length *= shape [i];
+			var pool = ArrayPool<float>.Shared;
+			var a = pool.Rent (length);
+			Array.Fill (a, initialValue);
+			var v = Variable (graph, a, shape, name);
+			pool.Return (a);
+			return v;
+		}
+
+		public static unsafe MPSGraphTensor Variable (this MPSGraph graph, ReadOnlySpan<float> initialValues, int[] shape, string? name = null)
+		{
+			var length = 1;
+			for (var i = 0; i < shape.Length; i++)
+				length *= shape [i];
+			if (length != initialValues.Length)
+				throw new ArgumentException ($"The number of initial values ({initialValues.Length}) does not match the shape length ({length}).");
+			fixed (float* p = initialValues) {
+				using var data = NSData.FromBytesNoCopy ((IntPtr) p, (nuint) (initialValues.Length * 4), freeWhenDone: false);
+				return graph.Variable (data, shape, MPSDataType.Float32, name);
+			}
+		}
+	}
+}

src/MetalPerformanceShadersGraph/MPSGraphTensorData.cs

@@ -0,0 +1,33 @@
+#nullable enable
+
+using System;
+using System.Buffers;
+using System.Runtime.InteropServices;
+
+using Foundation;
+using ObjCRuntime;
+using Metal;
+using MetalPerformanceShaders;
+
+namespace MetalPerformanceShadersGraph
+{
+	public partial class MPSGraphTensorData
+	{
+		public static MPSGraphTensorData Create (IMTLDevice device, ReadOnlySpan<float> values, params int[] shape)
+		{
+			var ndarray = MPSNDArray.Create (device, values, shape);
+			return new MPSGraphTensorData (ndarray);
+		}
+
+		public static MPSGraphTensorData Create (params MPSImage[] imageBatch)
+		{
+			return new MPSGraphTensorData (NSArray<MPSImage>.FromNSObjects (imageBatch));
+		}
+
+		public void Read (Span<float> values)
+		{
+			using var ndarray = this.MPSNDArray;
+			ndarray.Read (values);
+		}
+	}
+}

src/frameworks.sources

@@ -1187,6 +1187,15 @@ METALPERFORMANCESHADERS_SOURCES = \
 	MetalPerformanceShaders/MPSStateBatch.cs \
 	MetalPerformanceShaders/MPSStateResourceList.cs \
 
+# MetalPerformanceShadersGraph
+
+METALPERFORMANCESHADERSGRAPH_API_SOURCES = \
+	MetalPerformanceShadersGraph/MPSGraphEnums.cs \
+
+METALPERFORMANCESHADERSGRAPH_SOURCES = \
+	MetalPerformanceShadersGraph/MPSGraphExtensions.cs \
+	MetalPerformanceShadersGraph/MPSGraphTensorData.cs \
+
 # MetricKit
 
 METRICKIT_SOURCES = \
@@ -2026,6 +2035,7 @@ MACOS_FRAMEWORKS =          \
 	Metal                   \
 	MetalKit                \
 	MetalPerformanceShaders \
+	MetalPerformanceShadersGraph \
 	MetricKit               \
 	MLCompute               \
 	MobileCoreServices      \
@@ -2135,6 +2145,7 @@ IOS_FRAMEWORKS =         \
 	Metal \
 	MetalKit \
 	MetalPerformanceShaders \
+	MetalPerformanceShadersGraph \
 	MetricKit \
 	MLCompute \
 	MobileCoreServices \
@@ -2239,6 +2250,7 @@ TVOS_FRAMEWORKS =           \
 	Metal                   \
 	MetalKit                \
 	MetalPerformanceShaders \
+	MetalPerformanceShadersGraph \
 	MLCompute               \
 	MobileCoreServices      \
 	ModelIO                 \
@@ -2321,6 +2333,7 @@ MACCATALYST_FRAMEWORKS = \
 	Metal \
 	MetalKit \
 	MetalPerformanceShaders \
+	MetalPerformanceShadersGraph \
 	MetricKit \
 	MLCompute \
 	MobileCoreServices \

src/generator.cs

@@ -791,6 +791,8 @@ public NamespaceManager (BindingTouch binding_touch, string customObjCRuntimeNS,
 			ImplicitNamespaces.Add ("ModelIO");
 		if (Frameworks.HaveMetal)
 			ImplicitNamespaces.Add ("Metal");
+		if (Frameworks.HaveMetalPerformanceShadersGraph)
+			ImplicitNamespaces.Add ("MetalPerformanceShadersGraph");
 
 		if (Frameworks.HaveCoreImage)
 			ImplicitNamespaces.Add ("CoreImage");