Complete size checks in BLAS.[sy/he]mm! (JuliaLang#45605)

ffucci · Aug 11, 2022 · 7d55f9c · 7d55f9c
1 parent 53ba43e
commit 7d55f9c
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 10 deletions.
diff --git a/stdlib/LinearAlgebra/src/blas.jl b/stdlib/LinearAlgebra/src/blas.jl
@@ -1540,11 +1540,27 @@ for (mfname, elty) in ((:dsymm_,:Float64),
             require_one_based_indexing(A, B, C)
             m, n = size(C)
             j = checksquare(A)
-            if j != (side == 'L' ? m : n)
-                throw(DimensionMismatch(lazy"A has size $(size(A)), C has size ($m,$n)"))
-            end
-            if size(B,2) != n
-                throw(DimensionMismatch(lazy"B has second dimension $(size(B,2)) but needs to match second dimension of C, $n"))
+            M, N = size(B)
+            if side == 'L'
+                if j != m
+                    throw(DimensionMismatch(lazy"A has first dimension $j but needs to match first dimension of C, $m"))
+                end
+                if N != n
+                    throw(DimensionMismatch(lazy"B has second dimension $N but needs to match second dimension of C, $n"))
+                end
+                if j != M
+                    throw(DimensionMismatch(lazy"A has second dimension $j but needs to match first dimension of B, $M"))
+                end
+            else
+                if j != n
+                    throw(DimensionMismatch(lazy"B has second dimension $j but needs to match second dimension of C, $n"))
+                end
+                if N != j
+                    throw(DimensionMismatch(lazy"A has second dimension $N but needs to match first dimension of B, $j"))
+                end
+                if M != m
+                    throw(DimensionMismatch(lazy"A has first dimension $M but needs to match first dimension of C, $m"))
+                end
             end
             chkstride1(A)
             chkstride1(B)
@@ -1614,11 +1630,27 @@ for (mfname, elty) in ((:zhemm_,:ComplexF64),
             require_one_based_indexing(A, B, C)
             m, n = size(C)
             j = checksquare(A)
-            if j != (side == 'L' ? m : n)
-                throw(DimensionMismatch(lazy"A has size $(size(A)), C has size ($m,$n)"))
-            end
-            if size(B,2) != n
-                throw(DimensionMismatch(lazy"B has second dimension $(size(B,2)) but needs to match second dimension of C, $n"))
+            M, N = size(B)
+            if side == 'L'
+                if j != m
+                    throw(DimensionMismatch(lazy"A has first dimension $j but needs to match first dimension of C, $m"))
+                end
+                if N != n
+                    throw(DimensionMismatch(lazy"B has second dimension $N but needs to match second dimension of C, $n"))
+                end
+                if j != M
+                    throw(DimensionMismatch(lazy"A has second dimension $j but needs to match first dimension of B, $M"))
+                end
+            else
+                if j != n
+                    throw(DimensionMismatch(lazy"B has second dimension $j but needs to match second dimension of C, $n"))
+                end
+                if N != j
+                    throw(DimensionMismatch(lazy"A has second dimension $N but needs to match first dimension of B, $j"))
+                end
+                if M != m
+                    throw(DimensionMismatch(lazy"A has first dimension $M but needs to match first dimension of C, $m"))
+                end
             end
             chkstride1(A)
             chkstride1(B)

diff --git a/stdlib/LinearAlgebra/test/blas.jl b/stdlib/LinearAlgebra/test/blas.jl
@@ -227,11 +227,19 @@ Random.seed!(100)
                 @test_throws DimensionMismatch BLAS.symm('R','U',Cmn,Cnn)
                 @test_throws DimensionMismatch BLAS.symm!('L','U',one(elty),Asymm,Cnn,one(elty),Cmn)
                 @test_throws DimensionMismatch BLAS.symm!('L','U',one(elty),Asymm,Cnn,one(elty),Cnm)
+                @test_throws DimensionMismatch BLAS.symm!('L','U',one(elty),Asymm,Cmn,one(elty),Cnn)
+                @test_throws DimensionMismatch BLAS.symm!('R','U',one(elty),Asymm,Cnm,one(elty),Cmn)
+                @test_throws DimensionMismatch BLAS.symm!('R','U',one(elty),Asymm,Cnn,one(elty),Cnm)
+                @test_throws DimensionMismatch BLAS.symm!('R','U',one(elty),Asymm,Cmn,one(elty),Cnn)
                 if elty <: BlasComplex
                     @test_throws DimensionMismatch BLAS.hemm('L','U',Cnm,Cnn)
                     @test_throws DimensionMismatch BLAS.hemm('R','U',Cmn,Cnn)
                     @test_throws DimensionMismatch BLAS.hemm!('L','U',one(elty),Aherm,Cnn,one(elty),Cmn)
                     @test_throws DimensionMismatch BLAS.hemm!('L','U',one(elty),Aherm,Cnn,one(elty),Cnm)
+                    @test_throws DimensionMismatch BLAS.hemm!('L','U',one(elty),Aherm,Cmn,one(elty),Cnn)
+                    @test_throws DimensionMismatch BLAS.hemm!('R','U',one(elty),Aherm,Cnm,one(elty),Cmn)
+                    @test_throws DimensionMismatch BLAS.hemm!('R','U',one(elty),Aherm,Cnn,one(elty),Cnm)
+                    @test_throws DimensionMismatch BLAS.hemm!('R','U',one(elty),Aherm,Cmn,one(elty),Cnn)
                 end
             end
         end

diff --git a/stdlib/LinearAlgebra/test/symmetric.jl b/stdlib/LinearAlgebra/test/symmetric.jl
@@ -352,6 +352,9 @@ end
                 C = zeros(eltya,n,n)
                 @test Hermitian(aherm) * a ≈ aherm * a
                 @test a * Hermitian(aherm) ≈ a * aherm
+                # rectangular multiplication
+                @test [a; a] * Hermitian(aherm) ≈ [a; a] * aherm
+                @test Hermitian(aherm) * [a a] ≈ aherm * [a a]
                 @test Hermitian(aherm) * Hermitian(aherm) ≈ aherm*aherm
                 @test_throws DimensionMismatch Hermitian(aherm) * Vector{eltya}(undef, n+1)
                 LinearAlgebra.mul!(C,a,Hermitian(aherm))
@@ -360,6 +363,9 @@ end
                 @test Symmetric(asym) * Symmetric(asym) ≈ asym*asym
                 @test Symmetric(asym) * a ≈ asym * a
                 @test a * Symmetric(asym) ≈ a * asym
+                # rectangular multiplication
+                @test Symmetric(asym) * [a a] ≈ asym * [a a]
+                @test [a; a] * Symmetric(asym) ≈ [a; a] * asym
                 @test_throws DimensionMismatch Symmetric(asym) * Vector{eltya}(undef, n+1)
                 LinearAlgebra.mul!(C,a,Symmetric(asym))
                 @test C ≈ a*asym