Size check in 2-argument mul (#1315)

This PR adds a size check in the 2-argument `mul`, so that now the
destination array is allocated only if the sizes of the arguments are
compatible with matrix multiplication. This means that we don't allocate
in case of an error anymore.

The performance for small-matrix multiplication seems largely similar
(it's comparable to
#1310, and seems
identical within the noise limit):
```julia
julia> A = [1 2; 3 4];

julia> @Btime $A * $A;
  42.304 ns (2 allocations: 112 bytes) # before this PR
  44.203 ns (2 allocations: 112 bytes) # this PR
```

We also redirect the generic `mul` to `_mul` now, which is the function
that defines the multiplication code. This allows us to reuse the `_mul`
definition elsewhere without having to repeat code. Currently, this is
mainly necessary in the `Bidiagonal`-triangular multiplications.

Author: ViralBShah

src/bidiag.jl

@@ -1261,26 +1261,22 @@ function _dibimul_nonzeroalpha!(C::Bidiagonal, A::Diagonal, B::Bidiagonal, _add)
 end
 
 function mul(A::UpperOrUnitUpperTriangular, B::Bidiagonal)
-    TS = promote_op(matprod, eltype(A), eltype(B))
-    C = mul!(similar(A, TS, size(A)), A, B)
+    C = _mul(A, B)
     return B.uplo == 'U' ? UpperTriangular(C) : C
 end
 
 function mul(A::LowerOrUnitLowerTriangular, B::Bidiagonal)
-    TS = promote_op(matprod, eltype(A), eltype(B))
-    C = mul!(similar(A, TS, size(A)), A, B)
+    C = _mul(A, B)
     return B.uplo == 'L' ? LowerTriangular(C) : C
 end
 
 function mul(A::Bidiagonal, B::UpperOrUnitUpperTriangular)
-    TS = promote_op(matprod, eltype(A), eltype(B))
-    C = mul!(similar(B, TS, size(B)), A, B)
+    C = _mul(A, B)
     return A.uplo == 'U' ? UpperTriangular(C) : C
 end
 
 function mul(A::Bidiagonal, B::LowerOrUnitLowerTriangular)
-    TS = promote_op(matprod, eltype(A), eltype(B))
-    C = mul!(similar(B, TS, size(B)), A, B)
+    C = _mul(A, B)
     return A.uplo == 'L' ? LowerTriangular(C) : C
 end
 

src/matmul.jl

@@ -51,11 +51,13 @@ end
 
 # Matrix-vector multiplication
 function (*)(A::StridedMaybeAdjOrTransMat{T}, x::StridedVector{S}) where {T<:BlasFloat,S<:Real}
+    matmul_size_check(size(A), size(x))
     TS = promote_op(matprod, T, S)
     y = isconcretetype(TS) ? convert(AbstractVector{TS}, x) : x
     mul!(similar(x, TS, size(A,1)), A, y)
 end
 function (*)(A::AbstractMatrix{T}, x::AbstractVector{S}) where {T,S}
+    matmul_size_check(size(A), size(x))
     TS = promote_op(matprod, T, S)
     mul!(similar(x, TS, axes(A,1)), A, x)
 end
@@ -113,7 +115,10 @@ julia> [1 1; 0 1] * [1 0; 1 1]
 """
 (*)(A::AbstractMatrix, B::AbstractMatrix) = mul(A, B)
 # we add an extra level of indirection to avoid ambiguities in *
-function mul(A::AbstractMatrix, B::AbstractMatrix)
+# We also define the core functionality within _mul to reuse the code elsewhere
+mul(A::AbstractMatrix, B::AbstractMatrix) = _mul(A, B)
+function _mul(A::AbstractMatrix, B::AbstractMatrix)
+    matmul_size_check(size(A), size(B))
     TS = promote_op(matprod, eltype(A), eltype(B))
     mul!(matprod_dest(A, B, TS), A, B)
 end

src/symmetric.jl

@@ -710,12 +710,14 @@ end
 mul(A::HermOrSym, B::HermOrSym) = A * copyto!(similar(parent(B)), B)
 # catch a few potential BLAS-cases
 function mul(A::HermOrSym{<:BlasFloat,<:StridedMatrix}, B::AdjOrTrans{<:BlasFloat,<:StridedMatrix})
+    matmul_size_check(size(A), size(B))
     T = promote_type(eltype(A), eltype(B))
     mul!(similar(B, T, (size(A, 1), size(B, 2))),
             convert(AbstractMatrix{T}, A),
             copy_oftype(B, T)) # make sure the AdjOrTrans wrapper is resolved
 end
 function mul(A::AdjOrTrans{<:BlasFloat,<:StridedMatrix}, B::HermOrSym{<:BlasFloat,<:StridedMatrix})
+    matmul_size_check(size(A), size(B))
     T = promote_type(eltype(A), eltype(B))
     mul!(similar(B, T, (size(A, 1), size(B, 2))),
             copy_oftype(A, T), # make sure the AdjOrTrans wrapper is resolved