Support more type combinations

Author: christiangnrd

lib/mpsgraphs/MPSGraphs.jl

@@ -16,21 +16,24 @@ using CEnum
 using ObjectiveC, .Foundation, .Dispatch
 
 # Valid combination of input (A and B matrices) and output (C) types
-# TODO: support the commented type combinations
+#   The commented type combinations work but are slower than with MPSMatrixMultiplicatiom
 const MPSGRAPH_VALID_MATMUL_TYPES =
     [
     #  (Int8, Float16),
     #  (Int8, Float32),
     #  (Int16, Float32),
      (Float16, Float16),
-    #  (Float16, Float32),
+     (Float16, Float32),
      (Float32, Float32),
     ]
 
 const MPSGRAPH_VALID_MATVECMUL_TYPES =
     [
+     (Int8, Float16),
+     (Int8, Float32),
+     (Int16, Float32),
      (Float16, Float16),
-    #  (Float16, Float32),
+     (Float16, Float32),
      (Float32, Float32),
     ]
 

lib/mpsgraphs/matmul.jl

@@ -1,27 +1,34 @@
-function _matmul!(c::MPSMatrix, ::Type{T1}, a::MPSMatrix, ::Type{T2}, b::MPSMatrix, ::Type{T3}, alpha::Number, beta::Number, transpose_a, transpose_b) where {T1, T2, T3}
+function _matmul!(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), T2)
-    placeB = placeholderTensor(graph, size(b), T3)
+    placeA = placeholderTensor(graph, size(a), Tab)
+    placeB = placeholderTensor(graph, size(b), Tab)
+
+    castA, castB = if Tc != Tab
+        castTensor(graph, placeA, Tc, "castA"),
+        castTensor(graph, placeB, Tc, "castB")
+    else
+        placeA, placeB
+    end
 
     transA = if transpose_a
-        transposeTensor(graph, placeA, 0, 1, "transpose_a")
+        transposeTensor(graph, castA, 0, 1, "transpose_a")
     else
-        placeA
+        castA
     end
 
     transB = if transpose_b
-        transposeTensor(graph, placeB, 0, 1, "transpose_b")
+        transposeTensor(graph, castB, 0, 1, "transpose_b")
     else
-        placeB
+        castB
     end
 
     matmul = matrixMultiplicationWithPrimaryTensor(graph, transB, transA)
 
     afteralpha = if alpha == 1
         matmul
     else
-        alphatensor = constantWithScalar(graph, alpha, T1)
+        alphatensor = constantWithScalar(graph, alpha, Tc)
         multiplicationWithPrimaryTensor(graph, alphatensor, matmul)
     end
 
@@ -33,9 +40,9 @@ function _matmul!(c::MPSMatrix, ::Type{T1}, a::MPSMatrix, ::Type{T2}, b::MPSMatr
     afterbeta = if beta == 0
         afteralpha
     else
-        placeC = placeholderTensor(graph, UInt.(size(c)), T1)
+        placeC = placeholderTensor(graph, UInt.(size(c)), Tc)
         feed[placeC] = MPSGraphTensorData(c)
-        betatensor = constantWithScalar(graph, beta, T1)
+        betatensor = constantWithScalar(graph, beta, Tc)
         betaC = multiplicationWithPrimaryTensor(graph, betatensor, placeC)
         additionWithPrimaryTensor(graph, afteralpha, betaC)
     end
@@ -46,12 +53,12 @@ function _matmul!(c::MPSMatrix, ::Type{T1}, a::MPSMatrix, ::Type{T2}, b::MPSMatr
     return MPSNDArray(resultdata)
 end
 
-function graph_matmul!(c::MtlArray{T1, N}, a::MtlArray{T2, N}, b::MtlArray{T3, N}, alpha::Number = true, beta::Number = false, transpose_a = false, transpose_b = false) where {T1, T2, T3, N}
-    resultndarr = _matmul!(MPSMatrix(c), T1, MPSMatrix(a), T2, MPSMatrix(b), T3, alpha, beta, transpose_a, transpose_b)
+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)
 end
 
-function graph_matvecmul!(c::MtlVector{T1}, a::MtlMatrix{T2}, b::MtlVector{T3}, alpha::Number = true, beta::Number = false, transpose = false) where {T1, T2, T3}
-    resultndarr = _matmul!(MPSMatrix(c), T1, MPSMatrix(a), T2, MPSMatrix(b), T3, alpha, beta, transpose, false)
+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)
 end

lib/mpsgraphs/operations.jl

@@ -1,4 +1,11 @@
 
+function castTensor(graph::MPSGraph, tensor::MPSGraphTensor, toType, name="cast")
+    obj = @objc [graph::id{MPSGraph} castTensor:tensor::id{MPSGraphTensor}
+                                toType:toType::MPSDataType
+                                name:name::id{NSString}]::id{MPSGraphTensor}
+    MPSGraphTensor(obj)
+end
+
 function constantWithScalar(graph::MPSGraph, scalar::Number, dataType)
     obj = @objc [graph::id{MPSGraph} constantWithScalar:scalar::Float64
                                 dataType:dataType::MPSDataType]::id{MPSGraphTensor}