Handle empty products in mul! without BLAS, and treat β = 0 as a strong zero

src/linalg.jl

@@ -54,6 +54,13 @@ function LinearAlgebra.mul!(
         A::StridedView{<:Any, 2}, B::StridedView{<:Any, 2},
         α::Number = true, β::Number = false
     ) where {T}
+    # an empty product only scales `C`, which does not need BLAS; GPU BLAS libraries can not
+    # handle the empty arrays anyway
+    if iszero(size(A, 2))
+        (size(C, 1) == size(A, 1) && size(C, 2) == size(B, 2) && iszero(size(B, 1))) ||
+            throw(DimensionMismatch("A has size $(size(A)), B has size $(size(B)), C has size $(size(C))"))
+        return _scale!(C, β)
+    end
     if !(eltype(C) <: LinearAlgebra.BlasFloat && eltype(A) == eltype(B) == eltype(C))
         return __mul!(C, A, B, α, β)
     end
@@ -72,6 +79,9 @@ function LinearAlgebra.mul!(
     return C
 end
 
+# `C = β * C`, where `β == 0` overwrites `C`, which may contain `NaN`s, e.g. if it is uninitialized
+_scale!(C::StridedView, β::Number) = iszero(β) ? fill!(C, zero(eltype(C))) : rmul!(C, β)
+
 function isblasmatrix(A::StridedView{T, 2}) where {T <: LinearAlgebra.BlasFloat}
     if A.op == identity
         return stride(A, 1) == 1 || stride(A, 2) == 1
@@ -169,7 +179,7 @@ function __mul!(
     C2 = sreshape(C, (m, n, 1))
 
     if α == 0 || k == 0
-        rmul!(C, β)
+        _scale!(C, β)
     elseif α == 1
         if β == 0
             _mapreducedim!(*, +, zero, (m, n, k), (C2, A2, B2))

test/blasmultests.jl

@@ -95,6 +95,20 @@ end
     B1 = StridedView(copy(B))
     C1 = StridedView(copy(C))
     @test mul!(C, A, B, α, β) ≈ mul!(C1, A1, B1, α, β)
+
+    # `β = 0` overwrites `C` and ignores `NaN`s, also when the product can not be computed by BLAS
+    for T in (Float64, ComplexF64), op in (identity, conj), α in (zero(T), one(T))
+        A2 = StridedView(rand(T, 2, 0))
+        B2 = StridedView(rand(T, 0, 2))
+        C2 = StridedView(fill(T(NaN), 2, 2))
+        @test iszero(mul!(C2, op(A2), B2, α, zero(T)))
+    end
+    for T in (Float64, ComplexF64), op in (identity, conj)
+        A2 = StridedView(rand(T, 2, 3))
+        B2 = StridedView(rand(T, 3, 2))
+        C2 = StridedView(fill(T(NaN), 2, 2))
+        @test iszero(mul!(C2, op(A2), B2, zero(T), zero(T)))
+    end
 end
 
 # small enough to never be split over threads, such that no tasks need to be spawned

test/gpu.jl

@@ -128,6 +128,10 @@ end
     vA2 = view(StridedView(rand(T, (32, 64))), 1:32, 1:2:64)
     vA3 = view(StridedView(rand(T, (64, 32))), 1:2:64, 1:32)
     @test compare((C, A, B) -> mul!(C, A, B, α, β), AT, vA1, vA2, vA3)
+    # empty inner dimension, which only scales `C`
+    A1, A2, A3 = StridedView(rand(T, (4, 4))), StridedView(rand(T, (4, 0))), StridedView(rand(T, (0, 4)))
+    @test compare((C, A, B) -> mul!(C, A, B, α, β), AT, A1, A2, A3)
+    @test compare((C, A, B) -> mul!(C, A, B, α, zero(T)), AT, A1, A2, A3)
 end
 
 @testset "map, scale!, axpy!, axpby! ($AT)" for AT in ATs