norm(a) could be a lot faster

[stevengj]

Based on #697 (comment):

Here is a simple implementation of the 2-norm that correctly scales in case of underflow, overflow, and subnormal quantities (updated, see below), but which seems to be way faster than our built-in norm function (which calls generic_norm2 for lengths < 32 and BLAS.nrm2 otherwise):

function mynorm2(x::AbstractArray) # could also work on iterators
    isempty(x) && return float(sum(x))
    norm² = sum(abs2 ∘ float, x)
    if isfinite(norm²) & (norm² >= floatmin(norm²))
        return sqrt(norm²)  # fast path: no overflow/underflow or subnormals
    else
        scale = isinf(norm²) ? floatmin(sqrt(inv(norm²)))*4 : inv(floatmin(sqrt(norm²))*4)
        return sqrt(sum(x -> abs2(x * scale), x)) / scale
    end
end

Here is a benchmark, for both rand(n) arrays (no overflow or underflow) and randn(n) * 1e200 (overflow, which requires an extra pass over the array in mynorm2 but takes the same time for BLAS.nrm2 which apparently has no "fast path"). On my machine (Apple M4), it makes no difference whether I do BLAS.set_num_threads(1) or BLAS.set_num_threads(4). mynorm2 is nearly 10 times faster for many lengths!

Benchmark code

using BenchmarkTools, LinearAlgebra

n = sort!([5; 15; 25; 30; 10 .^ (0:7)])

t_mynorm2 = map(n) do n
    @show n
    a = rand(n)
    @belapsed mynorm2($a)
end

t_mynorm2_overflow = map(n) do n
    @show n
    a = rand(n) * 1e200
    @belapsed mynorm2($a)
end

BLAS.set_num_threads(1)

t_norm = map(n) do n
    @show n
    a = rand(n)
    @belapsed norm($a)
end

t_norm_overflow = map(n) do n
    @show n
    a = rand(n)
    @belapsed norm($a)
end

using Plots

plot(n, 1e-9 * n ./ t_mynorm2, label="mynorm2", marker=:star, color=:blue)
plot!(n, 1e-9 * n ./ t_mynorm2_overflow, label="mynorm2 (overflow)", marker=:star, color=:cyan)
plot!(n, 1e-9 * n ./ t_norm, label="norm", marker=:circle, color=:red)
plot!(n, 1e-9 * n ./ t_norm_overflow, label="norm (overflow)", marker=:circle, color=:magenta)
plot!(xaxis=:log10, xlabel="array length n", ylabel="speed (elements/ns)", legend=:outertopright,
      size=(800,400), left_margin=2mm, bottom_margin=3mm)

[stevengj]

I also noticed that our current generic_norm2 function (used for $n < 32$) only checks for over/underflow, and not subnormal values, so it is significantly less accurate than mynorm2 in the case where the squared entries are subnormal, e.g.:

julia> n = 10; a = fill(1e-161, n);

julia> norm(a), mynorm2(a), sqrt(n) * a[1]
(3.1434555694052575e-161, 3.1622776601683794e-161, 3.1622776601683798e-161)

(the third value is the correct result, which nearly matches mynorm2 but only matches norm to two digits).

[araujoms]

Nice! But why is it faster? Because of the pairwise summation algorithm in sum? And which BLAS are you using? OpenBLAS or AppleAccelerate?

[stevengj]

I'm using OpenBLAS. The LAPACK reference dnrm2 implementation seems crazily over-complicated to me, with a loop full of branches, but I'm not sure what OpenBLAS does. And for BLAS1 functions like this, there often seems to be little advantage nowadays to the "optimized" library routines compared to what LLVM can spit out from a simple loop.

Pairwise summation doesn't usually improve speed (it is slightly slower than a straight loop, which is its base case), it is to improve accuracy.

[araujoms]

You still get overflow in edge cases, because inv(inv(floatmax(Float64))) == Inf. For example for floatmax(Float64)/sqrt(3)*ones(3). A possible solution is to do instead

function mynorm3(x::AbstractArray) # could also work on iterators
    isempty(x) && return float(sum(x))
    norm² = sum(abs2 ∘ float, x)
    if isfinite(norm²) & (norm² >= floatmin(norm²))
        return sqrt(norm²)  # fast path: no overflow/underflow or subnormals
    else
        scale = isinf(norm²) ? floatmax(sqrt(norm²)) : floatmin(sqrt(norm²))
        return sqrt(sum(x -> abs2(x / scale), x)) * scale
    end
end

but I'm not sure the slowdown is worth it.

[stevengj]

We could just do inv(floatmin) instead of floatmax as the scale factor:

scale = inv(isinf(norm²) ? inv(floatmin(sqrt(inv(norm²)))) : floatmin(sqrt(norm²)))

or simply

scale = isinf(norm²) ? floatmin(sqrt(inv(norm²))) : inv(floatmin(sqrt(norm²)))

(which is a good idea anyway, to ensure that the scale factor is not subnormal).

Updated my top post.

[araujoms]

That doesn't work either, because inv(floatmin(BigFloat)) == Inf. Why must every floating point format have such a shit design?

[stevengj]

I think we agreed in another issue (JuliaLang/julia#61935) not to worry about that kind of case for BigFloat, because no reasonable BigFloat application will operate near floatmax $\approx 10^{10^{18}}$?

But we can just use inv(floatmin*4) to avoid this.

[araujoms]

Sure, people playing with those numbers only have themselves to blame. The problem is that this path is taken when the norm is zero, with the result that

julia> mynorm2(BigFloat[0])
NaN

[stevengj]

Should be fixed with my latest version above, which uses floatmin * 4.

julia> mynorm2(BigFloat[0])
0.0

(And of course there are also other alternatives, as you mentioned above. The important thing is the fast path... if we lose a small factor in performance for cases where abs2 will over/underflow, or for an array of all zeros, I doubt that is a big problem.)

[stevengj]

I think we should be more conservative, since not all floating-point types even define floatmin (e.g. Decimals.jl: JuliaMath/Decimals.jl#114), and instead do something like:

scale = _norm2_scale(x, norm²)
invscale = inv(scale)
return isfinite(invscale) ? sqrt(sum(x -> abs2(x * invscale), x)) * scale : sqrt(sum(x -> abs2(x / scale), x)) * scale

(or rather norm2 instead of abs2), with

_norm2_scale(x, norm²::Union{Float16, Float32, Float64, BigFloat}) = isinf(norm²) ? inv(floatmin(sqrt(inv(norm²)))*4) : floatmin(sqrt(norm²))*4
function _norm2_scale(x, norm²) # slower fallback
    scale = normInf(x)
    return iszero(scale) ? oneunit(scale) : scale
end

[araujoms]

That's more like it! It even fixes an arcane overflow bug in norm:

julia> v = floatmax(BigFloat)/sqrt(big(3.0))*ones(3);

julia> norm(v)
Inf

julia> mynorm2(v)
5.875653789111587590936911998878442589938516392745498308333779606469323584389875e+1388255822130839282

I'm not sure the slower fallback is a good idea, though, it will unnecessarily slow down DoubleFloats and MultiFloats, and other packages that do a better job of subtyping AbstractFloat. IMHO this is a bug in Decimals and should be fixed there.

[stevengj]

floatmin also does not exist for IntervalArithmetic (JuliaIntervals/IntervalArithmetic.jl#760):

julia> using LinearAlgebra, IntervalArithmetic

julia> x = interval(1, 2)
[1.0, 2.0]_com

julia> norm([x])
[1.0, 2.0]_com

julia> floatmin(x)
ERROR: MethodError: no method matching floatmin(::Interval{Float64})

It also doesn't exist for Unitful.jl types (JuliaPhysics/Unitful.jl#850):

julia> using LinearAlgebra, Unitful

julia> x = 1.0u"m"
1 m

julia> norm([x])
1.0 m

julia> floatmin(x)
ERROR: MethodError: no method matching floatmin(::Quantity{Float64, 𝐋, Unitful.FreeUnits{(m,), 𝐋, nothing}})

(Though it does work for Measurements.jl.)

Basically, I'm worried that a lot of "exotic" numeric types could fail to implement this. Forcing them to pay a small performance penalty in an uncommon corner case seems like a small price to pay for generality. (And they could always implement _norm2_scale if they are really worried about it, which I doubt.)

[araujoms]

Neither IntervalArithmetic nor Unitful subtype AbstractFloat, so this could be solved simply by letting the fast fallback act only on AbstractFloat and leaving the slower fallback unrestricted.