More MPSGraphs API (#624)

* MPSGraphShapedType

* MPSGraphDevice

* MPSGraph compilation/serialization

* Fix cenversion to MPSShape

* Oops

* Add test

Author: christiangnrd

lib/mps/MPS.jl

@@ -21,7 +21,7 @@ using BFloat16s
 const MtlFloat = Union{Float32, Float16}
 
 const MPSShape = NSArray#{NSNumber}
-Base.convert(::Type{MPSShape}, tuple::Union{Vector{T},NTuple{T, <:Integer}}) where T = NSArray(NSNumber.(collect(tuple)))
+Base.convert(::Type{MPSShape}, tuple::Union{Vector{<:Integer},NTuple{N, <:Integer}}) where N = NSArray(NSNumber.(collect(tuple)))
 
 # Valid combination of input (A and B matrices) and output (C) types
 const MPS_VALID_MATMUL_TYPES =

lib/mpsgraphs/core.jl

@@ -11,15 +11,15 @@ end
 
 # @objcwrapper immutable=true MPSGraphShapedType <: MPSGraphType
 
-# XXX: Not used yet and needs fixing
-# function MPSGraphShapedType(shape::MPSShape, dataType)
-#     tmp = @objc [MPSGraphShapedType alloc]::id{MPSGraphShapedType}
-#     obj = MPSGraphShapedType(tmp)
-#     finalizer(release, obj)
-#     @objc [obj::id{MPSGraphShapedType} initWithShape:shape::id{MPSShape}
-#                                        dataType:dataType::MPSDataType]::id{MPSGraphShapedType}
-#     return obj
-# end
+MPSGraphShapedType(shape, dataType) = MPSGraphShapedType(convert(MPSShape, shape), dataType)
+function MPSGraphShapedType(shape::MPSShape, dataType)
+    tmp = @objc [MPSGraphShapedType alloc]::id{MPSGraphShapedType}
+    obj = MPSGraphShapedType(tmp)
+    finalizer(release, obj)
+    @objc [obj::id{MPSGraphShapedType} initWithShape:shape::id{MPSShape}
+                                       dataType:dataType::MPSDataType]::id{MPSGraphShapedType}
+    return obj
+end
 
 ## MPSGraphDevice.h
 # @objcwrapper MPSGraphDevice <: MPSGraphType

lib/mpsgraphs/execution.jl

@@ -32,3 +32,27 @@ function run(graph::MPSGraph, commandQueue::MTLCommandQueue, feeds::MPSGraphTens
                                          targetOperations:nil::id{Object}]::id{MPSGraphTensorDataDictionary}
     MPSGraphTensorDataDictionary(obj)
 end
+
+const MPSGraphTensorShapedTypeDictionary = NSDictionary#{MPSGraphTensor, MPSGraphTensorShapedType}
+
+compile(graph::MPSGraph, dev::MTLDevice, feeds::MPSGraphTensorShapedTypeDictionary, targetTensors::NSArray, targetOperations=nil, compilationDescriptor=nil) = compile(graph, MPSGraphDevice(dev), feeds, targetTensors, targetOperations, compilationDescriptor)
+function compile(graph::MPSGraph, dev::MPSGraphDevice, feeds::MPSGraphTensorShapedTypeDictionary, targetTensors::NSArray, targetOperations=nil, compilationDescriptor=nil)
+    exec = @objc [graph::id{MPSGraph} compileWithDevice:dev::id{MPSGraphDevice}
+                                     feeds:feeds::id{MPSGraphTensorShapedTypeDictionary}
+                             targetTensors:targetTensors::id{NSArray}
+                          targetOperations:targetOperations::id{Object}
+                     compilationDescriptor:compilationDescriptor::id{Object}]::id{MPSGraphExecutable}
+    return MPSGraphExecutable(exec)
+end
+
+function MPSGraphExecutableSerializationDescriptor()
+    tmp = @objc [MPSGraphExecutableSerializationDescriptor alloc]::id{MPSGraphExecutableSerializationDescriptor}
+    obj = MPSGraphExecutableSerializationDescriptor(tmp)
+    return obj
+end
+
+serialize(graphExe::MPSGraphExecutable, url, descriptor=MPSGraphExecutableSerializationDescriptor()) = serialize(graphExe, NSFileURL(url), descriptor)
+function serialize(graphExe::MPSGraphExecutable, url::NSURL, descriptor=MPSGraphExecutableSerializationDescriptor())
+    @objc [graphExe::id{MPSGraphExecutable} serializeToMPSGraphPackageAtURL:url::id{NSURL}
+                              descriptor:descriptor::id{MPSGraphExecutableSerializationDescriptor}]::Nothing
+end

lib/mpsgraphs/libmpsgraph.jl

@@ -12,7 +12,7 @@ const MPSGraphTensorDataDictionary = NSDictionary
 
 @objcwrapper immutable = true MPSGraphType <: MPSGraphObject
 
-@objcwrapper immutable = true MPSGraphShapedType <: MPSGraphType
+@objcwrapper immutable = false MPSGraphShapedType <: MPSGraphType
 
 @objcproperties MPSGraphShapedType begin
     @autoproperty shape::id{MPSShape} setter = setShape

res/wrap/libmpsgraph.toml

@@ -27,3 +27,6 @@ macro_mode = "disable"
 
 [api.MPSGraphTensorData]
 immutable=false
+
+[api.MPSGraphShapedType]
+immutable=false

test/mpsgraphs/core.jl

@@ -1,7 +1,7 @@
 if MPS.is_supported(device())
 
 using .MPS: MPSShape
-using .MPSGraphs: MPSGraph, MPSGraphDevice
+using .MPSGraphs: MPSGraph, MPSGraphDevice, MPSGraphShapedType
 @testset "Core" begin
 
 graph = MPSGraph()
@@ -13,6 +13,11 @@ graphdev = MPSGraphDevice(dev)
 @test graphdev.type == MPSGraphs.MPSGraphDeviceTypeMetal
 @test graphdev.metalDevice == dev
 
+mpsh = convert(MPS.MPSShape, (2,3,4))
+shtyp = MPSGraphShapedType(mpsh, Float32)
+@test shtyp.shape == convert(MPS.MPSShape,(2,3,4))
+@test shtyp.dataType == MPS.MPSDataTypeFloat32
+
 end # @testset "Core"
 
 end # MPS.is_supported(device())