Branch on Bool alpha in scaling mul! (#1286)

By dealing with the `alpha == 0` case separately, we ensure that if
`alpha::Bool`, it must be `true`. This reduces the branches in
`_lscale_add` from 4 to 2 in the common case of 3-argument `mul!`. This
leads to a latency reduction, as each branch has to compile a different
broadcast expression, and we currently compile four but use only one.
Primarily, this PR leads to a reduction in allocations.

```julia
julia> using LinearAlgebra

julia> v = 1:4; w = similar(v);

julia> @time mul!(w, 1, v);
  0.171120 seconds (1.04 M allocations: 52.799 MiB, 99.98% compilation time) # nightly
  0.163178 seconds (702.63 k allocations: 35.533 MiB, 99.98% compilation time) # this PR
```

Something similar usually doesn't lead to a big gain in the
`_rscale_add` method, as `s * alpha` often has the same type as `s`, and
therefore the branches on `alpha` compile the same code.

Author: jishnub

src/generic.jl

@@ -202,24 +202,35 @@ _lscale_add!(C::StridedArray, s::Number, X::StridedArray, alpha::Number, beta::N
     generic_mul!(C, s, X, alpha, beta)
 @inline function _lscale_add!(C::AbstractArray, s::Number, X::AbstractArray, alpha::Number, beta::Number)
     if axes(C) == axes(X)
-        if isone(alpha)
-            if iszero(beta)
-                @. C = s * X
-            else
-                @. C = s * X + C * beta
-            end
-        else
-            if iszero(beta)
-                @. C = s * X * alpha
-            else
-                @. C = s * X * alpha + C * beta
-            end
-        end
+        iszero(alpha) && return _rmul_or_fill!(C, beta)
+        _lscale_add_nonzeroalpha!(C, s, X, alpha, beta)
     else
         generic_mul!(C, s, X, alpha, beta)
     end
     return C
 end
+function _lscale_add_nonzeroalpha!(C::AbstractArray, s::Number, X::AbstractArray, alpha::Number, beta::Number)
+    if isone(alpha)
+        # since alpha is unused, we might as well set to `true` to avoid recompiling
+        # the branch if an `alpha` of a different type is used
+        _lscale_add_nonzeroalpha!(C, s, X, true, beta)
+    else
+        if iszero(beta)
+            @. C = s * X * alpha
+        else
+            @. C = s * X * alpha + C * beta
+        end
+    end
+    C
+end
+function _lscale_add_nonzeroalpha!(C::AbstractArray, s::Number, X::AbstractArray, alpha::Bool, beta::Number)
+    if iszero(beta)
+        @. C = s * X
+    else
+        @. C = s * X + C * beta
+    end
+    C
+end
 @inline mul!(C::AbstractArray, X::AbstractArray, s::Number, alpha::Number, beta::Number) =
     _rscale_add!(C, X, s, alpha, beta)
 
@@ -228,24 +239,26 @@ _rscale_add!(C::StridedArray, X::StridedArray, s::Number, alpha::Number, beta::N
 @inline function _rscale_add!(C::AbstractArray, X::AbstractArray, s::Number, alpha::Number, beta::Number)
     if axes(C) == axes(X)
         if isone(alpha)
-            if iszero(beta)
-                @. C = X * s
-            else
-                @. C = X * s + C * beta
-            end
+            # since alpha is unused, we might as well ignore it in this branch.
+            # This avoids recompiling the branch if an `alpha` of a different type is used
+            _rscale_add_alphaisone!(C, X, s, beta)
         else
             s_alpha = s * alpha
-            if iszero(beta)
-                @. C = X * s_alpha
-            else
-                @. C = X * s_alpha + C * beta
-            end
+            _rscale_add_alphaisone!(C, X, s_alpha, beta)
         end
     else
         generic_mul!(C, X, s, alpha, beta)
     end
     return C
 end
+function _rscale_add_alphaisone!(C::AbstractArray, X::AbstractArray, s::Number, beta::Number)
+    if iszero(beta)
+        @. C = X * s
+    else
+        @. C = X * s + C * beta
+    end
+    C
+end
 
 # For better performance when input and output are the same array
 # See https://github.com/JuliaLang/julia/issues/8415#issuecomment-56608729

test/generic.jl

@@ -857,4 +857,26 @@ end
     end
 end
 
+@testset "scaling mul" begin
+    v = 1:4
+    w = similar(v)
+    @test mul!(w, 2, v) == 2v
+    @test mul!(w, v, 2) == 2v
+    # 5-arg equivalent to the 3-arg method, but with non-Bool alpha
+    @test mul!(copy!(similar(v), v), 2, v, 1, 0) == 2v
+    @test mul!(copy!(similar(v), v), v, 2, 1, 0) == 2v
+    # 5-arg tests with alpha::Bool
+    @test mul!(copy!(similar(v), v), 2, v, true, 1) == 3v
+    @test mul!(copy!(similar(v), v), v, 2, true, 1) == 3v
+    @test mul!(copy!(similar(v), v), 2, v, false, 2) == 2v
+    @test mul!(copy!(similar(v), v), v, 2, false, 2) == 2v
+    # 5-arg tests
+    @test mul!(copy!(similar(v), v), 2, v, 1, 3) == 5v
+    @test mul!(copy!(similar(v), v), v, 2, 1, 3) == 5v
+    @test mul!(copy!(similar(v), v), 2, v, 2, 3) == 7v
+    @test mul!(copy!(similar(v), v), v, 2, 2, 3) == 7v
+    @test mul!(copy!(similar(v), v), 2, v, 2, 0) == 4v
+    @test mul!(copy!(similar(v), v), v, 2, 2, 0) == 4v
+end
+
 end # module TestGeneric