Still working

Author: christiangnrd

lib/mpsgraphs/matmul.jl

@@ -1,4 +1,4 @@
-function _matmul!(c::MPSMatrix, ::Type{Tc}, a::MPSMatrix, b::MPSMatrix, ::Type{Tab}, alpha::Number, beta::Number, transpose_a, transpose_b) where {Tc, Tab}
+function _matmul!(cmdbuf::MPSCommandBuffer, c::MPSMatrix, ::Type{Tc}, a::MPSMatrix, b::MPSMatrix, ::Type{Tab}, alpha::Number, beta::Number, transpose_a, transpose_b) where {Tc, Tab}
     graph = MPSGraph()
 
     placeA = placeholderTensor(graph, size(a), Tab)
@@ -7,6 +7,10 @@ function _matmul!(c::MPSMatrix, ::Type{Tc}, a::MPSMatrix, b::MPSMatrix, ::Type{T
     castA, castB = if Tc != Tab
         castTensor(graph, placeA, Tc, "castA"),
             castTensor(graph, placeB, Tc, "castB")
+
+    # castA, castB = if Tab != Float32
+    #     castTensor(graph, placeA, Float32, "castA"),
+    #         castTensor(graph, placeB, Float32, "castB")
     else
         placeA, placeB
     end
@@ -47,18 +51,28 @@ function _matmul!(c::MPSMatrix, ::Type{Tc}, a::MPSMatrix, b::MPSMatrix, ::Type{T
         additionWithPrimaryTensor(graph, afteralpha, betaC)
     end
 
+    castC = if Tc != Float32
+        afterbeta
+        # castTensor(graph, afterbeta, Tc, "castC")
+    else
+        afterbeta
+    end
+
     # Encode and commit matmul kernel
-    cmdbuf = MPSCommandBuffer(Metal.global_queue(device()))
-    resultdict = encode!(cmdbuf, graph, NSDictionary(feeds), NSArray([afterbeta]))
+    # resultdict = encode!(cmdbuf, graph, NSDictionary(feeds), NSArray([afterbeta]))
+    resultdict = encode!(cmdbuf, graph, NSDictionary(feeds), NSArray([castC]))
     commitAndContinue!(cmdbuf)
 
-    resultdata = MPSGraphTensorData(id{MPSGraphTensorData}(resultdict[afterbeta]))
+    # resultdata = MPSGraphTensorData(id{MPSGraphTensorData}(resultdict[afterbeta]))
+    resultdata = MPSGraphTensorData(id{MPSGraphTensorData}(resultdict[castC]))
 
-    return cmdbuf, MPSNDArray(resultdata)
+    return 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}
-    cmdbuf, resultndarr = _matmul!(MPSMatrix(c), Tc, MPSMatrix(a), MPSMatrix(b), Tab, alpha, beta, transpose_a, transpose_b)
+    cmdbuf = MPSCommandBuffer(Metal.global_queue(device()))
+
+    resultndarr = _matmul!(cmdbuf, 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)
@@ -70,7 +84,9 @@ function graph_matmul!(c::MtlArray{Tc, N}, a::MtlArray{Tab, N}, b::MtlArray{Tab,
 end
 
 function graph_matvecmul!(c::MtlVector{Tc}, a::MtlMatrix{Tab}, b::MtlVector{Tab}, alpha::Number = true, beta::Number = false, transpose = false) where {Tc, Tab}
-    cmdbuf, resultndarr = _matmul!(MPSMatrix(c), Tc, MPSMatrix(a), MPSMatrix(b), Tab, alpha, beta, transpose, false)
+    cmdbuf = MPSCommandBuffer(Metal.global_queue(device()))
+
+    resultndarr = _matmul!(cmdbuf, MPSMatrix(c), Tc, MPSMatrix(a), MPSMatrix(b), Tab, alpha, beta, transpose, false)
 
     commit!(cmdbuf) do cmdBuf
         exportDataWithCommandBuffer(resultndarr, cmdBuf, c.data[], Tc, c.offset)