Implementation more similar to MPS implementation

Author: christiangnrd

lib/mpsgraphs/MPSGraphs.jl

@@ -10,7 +10,8 @@ module MPSGraphs
 
 using ..Metal
 using .MTL
-using .MPS: MPSDataType, MPSMatrix, MPSVector, MPSShape, MPSNDArray, exportToMtlArray!
+using .MPS
+using .MPS: MPSDataType, MPSShape, exportDataWithCommandBuffer
 
 using CEnum
 using ObjectiveC, .Foundation, .Dispatch

lib/mpsgraphs/matmul.jl

@@ -32,33 +32,51 @@ function _matmul!(c::MPSMatrix, ::Type{Tc}, a::MPSMatrix, b::MPSMatrix, ::Type{T
         multiplicationWithPrimaryTensor(graph, alphatensor, matmul)
     end
 
-    feed = Dict(
+    feeds = Dict{MPSGraphTensor, MPSGraphTensorData}(
         placeA => MPSGraphTensorData(a),
         placeB => MPSGraphTensorData(b)
     )
 
     afterbeta = if beta == 0
         afteralpha
     else
-        placeC = placeholderTensor(graph, UInt.(size(c)), Tc)
-        feed[placeC] = MPSGraphTensorData(c)
+        placeC = placeholderTensor(graph, size(c), Tc)
+        feeds[placeC] = MPSGraphTensorData(c)
         betatensor = constantWithScalar(graph, beta, Tc)
         betaC = multiplicationWithPrimaryTensor(graph, betatensor, placeC)
         additionWithPrimaryTensor(graph, afteralpha, betaC)
     end
 
-    res = run(graph, feed, [afterbeta])
-    resultdata = only(Dict{MPSGraphTensor, MPSGraphTensorData}(res)).second
+    # Encode and commit matmul kernel
+    cmdbuf = MPSCommandBuffer(Metal.global_queue(device()))
+    resultdict = encode!(cmdbuf, graph, NSDictionary(feeds), NSArray([afterbeta]))
+    commitAndContinue!(cmdbuf)
 
-    return MPSNDArray(resultdata)
+    resultdata = MPSGraphTensorData(id{MPSGraphTensorData}(resultdict[afterbeta]))
+
+    return cmdbuf, MPSNDArray(resultdata)
 end
 
 function graph_matmul!(c::MtlArray{Tc, N}, a::MtlArray{Tab, N}, b::MtlArray{Tab, N}, alpha::Number = true, beta::Number = false, transpose_a = false, transpose_b = false) where {Tc, Tab, N}
-    resultndarr = _matmul!(MPSMatrix(c), Tc, MPSMatrix(a), MPSMatrix(b), Tab, alpha, beta, transpose_a, transpose_b)
-    return exportToMtlArray!(c, resultndarr)
+    cmdbuf, resultndarr = _matmul!(MPSMatrix(c), Tc, MPSMatrix(a), MPSMatrix(b), Tab, alpha, beta, transpose_a, transpose_b)
+
+    commit!(cmdbuf) do cmdBuf
+        exportDataWithCommandBuffer(resultndarr, cmdBuf, c.data[], Tc, c.offset)
+    end
+
+    wait_completed(cmdbuf)
+
+    return c
 end
 
 function graph_matvecmul!(c::MtlVector{Tc}, a::MtlMatrix{Tab}, b::MtlVector{Tab}, alpha::Number = true, beta::Number = false, transpose = false) where {Tc, Tab}
-    resultndarr = _matmul!(MPSMatrix(c), Tc, MPSMatrix(a), MPSMatrix(b), Tab, alpha, beta, transpose, false)
-    return exportToMtlArray!(c, resultndarr)
+    cmdbuf, resultndarr = _matmul!(MPSMatrix(c), Tc, MPSMatrix(a), MPSMatrix(b), Tab, alpha, beta, transpose, false)
+
+    commit!(cmdbuf) do cmdBuf
+        exportDataWithCommandBuffer(resultndarr, cmdBuf, c.data[], Tc, c.offset)
+    end
+
+    wait_completed(cmdbuf)
+
+    return c
 end