rmul! does not respect "strong zero": NaN*false=0.0

[tam724]

See the following MWE:

julia> using CUDA, LinearAlgebra

julia> rmul!([NaN, NaN, 1.0, -1.0], false)
4-element Vector{Float64}:
  0.0
  0.0
  0.0
 -0.0

julia> rmul!([NaN, NaN, 1.0, -1.0] |> cu, false)
4-element CuArray{Float32, 1, CUDA.DeviceMemory}:
 NaN
 NaN
   0.0
  -0.0

In julia base, false is defined as a "strong zero", see https://github.com/JuliaLang/julia/blob/5e9a32e7af2837e677e60543d4a15faa8d3a7297/base/bool.jl#L178. Hence, NaN*false = 0 and NaN*true = NaN. For consistency, the following dispatch could be defined for Bool.

function LinearAlgebra.rmul!(x::CUDA.DenseCuArray{<:CUDA.CUBLAS.CublasFloat}, k::Bool)
    k && return x
    return x .= copysign.(zero(eltype(x)), x)
end

This would bypass the fallback from rmul! to scal! defined here:

CUDA.jl/lib/cublas/linalg.jl

Lines 9 to 14 in 792aec5

	LinearAlgebra.rmul!(x::StridedCuArray{<:CublasFloat}, k::Number) =
	scal!(length(x), k, x)
	# Work around ambiguity with GPUArrays wrapper
	LinearAlgebra.rmul!(x::DenseCuArray{<:CublasFloat}, k::Real) =
	invoke(rmul!, Tuple{typeof(x), Number}, x, k)

I'd be happy to create a PR + tests.

[maleadt]

I'd be happy to create a PR + tests.

That would be great, thanks! I take it a similar fix needs to happen in GPUArrays? https://github.com/JuliaGPU/GPUArrays.jl/blob/87b95a9562fe59479a44f84098a006fb8c46ba7d/src/host/linalg.jl#L543-L552

[tam724]

I think the GPUArrays.jl implementation just uses the scalar multiplication Base.:*(::AbstractFloat, ::Bool). So it should comply with the rmul! implementation in LinearAlgebra.jl.

In https://github.com/tam724/CUDA.jl/tree/fix_rmul! I went through the "libraries/cublas/level 1" testset and included random CuArrays with NaN entries. If the tested function accepts a scalar factor it is also additionally tested with NaN, true and false.
This broke the tests of rmul! (fixed), axpby!, rotate! and reflect!.

I think it's hard to say how this should be treated generally, because even NaN*false = 0.0 is undocumented behavior in julia (documented in a comment).
EDIT: it is documented: https://docs.julialang.org/en/v1/manual/mathematical-operations/#Arithmetic-Operators
Maybe the "strong zero" should not be relied on, because e.g. even in LinearAlgebra.jl something like this happpens:

julia> using LinearAlgebra

julia> rotate!([NaN], [NaN], false, false) # rotate!(x, y, c, s) means: x = c*x + s*y and y = -conj(s)*x + c*y
([0.0], [NaN])

But a use case for rotate! with false arguments is hard to imagine..