Fix & tests

Author: christiangnrd

src/linalg.jl

@@ -32,6 +32,7 @@ end
 
 # Supported values are :auto, :MPS, :MPSGraph, and :GPUArrays
 const matmul_alg = ScopedValue(:auto)
+matmul_alg_error(alg, inT, outT) = error("Matrix multiplication algorithm `:$alg` is not supported for input eltype $inT and output eltype $outT.")
 
 LinearAlgebra.generic_matmatmul!(C::MtlMatrix, tA, tB, A::MtlMatrix, B::MtlMatrix, _add::MulAddMul) =
     LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, _add.alpha, _add.beta)
@@ -59,15 +60,19 @@ LinearAlgebra.generic_matmatmul!(C::MtlMatrix, tA, tB, A::MtlMatrix, B::MtlMatri
     transB = tB == 'T' || tB == 'C'
 
     alg = matmul_alg[]
+    mps_supported = supports_mps_matmul(A, B, C, MPS_VALID_MATMUL_TYPES)
+    mpsgraph_supported = supports_mpsgraph_matmul(A, B, C, MPSGRAPH_VALID_MATMUL_TYPES)
     # If possible, dispatch to MPSGraphs, then performance shaders
-    if supports_mpsgraph_matmul(A, B, C, MPSGRAPH_VALID_MATMUL_TYPES) && (alg === :MPSGraph || (alg === :auto && !should_use_MPS(A, B, C)))
+    if alg === :MPSGraph || (alg === :auto && mpsgraph_supported && !should_use_MPS(A, B, C))
+        mpsgraph_supported || matmul_alg_error(alg, eltype(A), eltype(C))
         graph_matmul!(C, A, B, alpha, beta, transA, transB)
-    elseif supports_mps_matmul(A, B, C, MPS_VALID_MATMUL_TYPES) && (alg === :MPS || alg === :auto)
+    elseif alg === :MPS || (alg === :auto && mps_supported)
+        mps_supported || matmul_alg_error(alg, eltype(A), eltype(C))
         matmul!(C, A, B, alpha, beta, transA, transB)
     elseif alg === :GPUArrays || alg === :auto
         GPUArrays.generic_matmatmul!(C, wrap(A, tA), wrap(B, tB), alpha, beta)
     else
-        error("Invalid matmul algorithm and input combination.")
+        error(":$alg is not a valid matmul algorithm. Options are: `:auto`, `:MPS`, `:MPSGraph`, `:GPUArrays`")
     end
 end
 
@@ -97,15 +102,19 @@ LinearAlgebra.generic_matvecmul!(C::MtlVector, tA::AbstractChar, A::MtlMatrix, B
     transA = tA == 'T' || tA == 'C'
 
     alg = matmul_alg[]
+    mps_supported = supports_mps_matmul(A, B, C, MPS_VALID_MATVECMUL_TYPES)
+    mpsgraph_supported = supports_mpsgraph_matmul(A, B, C, MPSGRAPH_VALID_MATVECMUL_TYPES)
     # If possible, dispatch to MPSGraphs, then performance shaders
-    if supports_mpsgraph_matmul(A, B, C, MPSGRAPH_VALID_MATVECMUL_TYPES) && (alg === :MPSGraph || alg === :auto)
+    if alg === :MPSGraph || (alg === :auto && mpsgraph_supported)
+        mpsgraph_supported || matmul_alg_error(alg, eltype(A), eltype(C))
         graph_matvecmul!(C, A, B, alpha, beta, transA)
-    elseif supports_mps_matmul(A, B, C, MPS_VALID_MATVECMUL_TYPES) && (alg === :MPS || alg === :auto)
+    elseif alg === :MPS || (alg === :auto && mps_supported)
+        mps_supported || matmul_alg_error(alg, eltype(A), eltype(C))
         matvecmul!(C, A, B, alpha, beta, transA)
     elseif alg === :GPUArrays || alg === :auto
         GPUArrays.generic_matmatmul!(C, wrap(A, tA), B, alpha, beta)
     else
-        error("Invalid matmul algorithm and input combination.")
+        error(":$alg is not a valid matmul algorithm. Options are: `:auto`, `:MPS`, `:MPSGraph`, `:GPUArrays`")
     end
 end
 

test/Project.toml

@@ -14,6 +14,7 @@ ObjectiveC = "e86c9b32-1129-44ac-8ea0-90d5bb39ded9"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

test/linalg.jl

@@ -1,7 +1,28 @@
-using LinearAlgebra
+using LinearAlgebra, ScopedValues
 
 if MPS.is_supported(device())
 
+@testset "matmul algorithm selection" begin
+    # test that unsupported configurations error properly
+    N = 20
+    function test_matmul(inT, outT; vec_b=false, alg=:auto)
+        a = MtlArray(rand(inT, N, N))
+        b = MtlArray(rand(inT, vec_b ? (N,) : (N, N)))
+        c = fill!(similar(b, outT), zero(outT))
+
+        @with (Metal.matmul_alg => alg) mul!(c,a,b)
+    end
+
+    # Unsupported for MPS and MPSGraph
+    for vec_b in (true, false)
+        @test_throws "Matrix multiplication algorithm `:MPS`" test_matmul(Int8, Int16; vec_b, alg=:MPS)
+        @test_throws "Matrix multiplication algorithm `:MPSGraph`" test_matmul(Int8, Int16; vec_b, alg=:MPSGraph)
+
+        # Invalid algorithm Symbol
+        @test_throws ":bad is not a valid matmul algorithm." test_matmul(Int8, Int16; vec_b, alg=:bad)
+        @test_throws ":bad is not a valid matmul algorithm." test_matmul(Float16, Float16; vec_b, alg=:bad)
+    end
+end
 
 @testset "test matrix vector multiplication of views" begin
     N = 20