diff --git a/src/linalg.jl b/src/linalg.jl index 02aa920..a1c87db 100644 --- a/src/linalg.jl +++ b/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)) diff --git a/test/blasmultests.jl b/test/blasmultests.jl index af62bf0..69553ec 100644 --- a/test/blasmultests.jl +++ b/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 diff --git a/test/gpu.jl b/test/gpu.jl index ab38519..35a108b 100644 --- a/test/gpu.jl +++ b/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