Adding inplace multiplication for (unit)triangular matrices

[mcognetta]

This PR closes #36828.

Adding in place mul! for (ε|Unit)(Upper|Lower)Triangular matrices.

Before, something like

julia> x = UpperTriangular(rand(3, 3))
3×3 UpperTriangular{Float64,Matrix{Float64}}:
 0.896774  0.0032823  0.133023
  ⋅        0.842607   0.414514
  ⋅         ⋅         0.98005

julia> A = UpperTriangular(rand(3, 3))
3×3 UpperTriangular{Float64,Matrix{Float64}}:
 0.42575  0.567843  0.19379
  ⋅       0.137596  0.636028
  ⋅        ⋅        0.959794

julia> mul!(x, A, A)

would fail with a method error, while

julia> mul!(rand(3, 3), A, A)

would work but destroy the underlying triangular structure.

Now,

julia> mul!(x, A, A)

works and maintains the triangular structure.

[mcognetta]

Gentle bump

[mcognetta]

I was inspired by @bjack205 's JuliaCon talk (sorry for the random mention), so I changed a slice to a view to reduce allocations. It now allocates half has much as using a slice. Making the code

        view(C, :, i) = A*view(B, :, i)

seems to still allocate as much as the slicing. Hopefully someone with more experience with views will see this and explain.

[mcognetta]

I was able to reduce it to only 1 allocation using a preallocated vector (a la https://discourse.julialang.org/t/inplace-multiplication-by-a-square-matrix/1702/4), but in benchmarking I found the original version + @views to be the most performant.

For the original change:

julia> @btime mul!(x, A, A)
  10.779 ms (1000 allocations: 3.97 MiB)

For the first view change:

julia> @btime mul!(x, A, A)
  11.684 ms (1000 allocations: 3.97 MiB)

For the preallocated view change:

julia> @btime mul!(x, A, A)
  49.991 ms (1 allocation: 4.06 KiB)

For the @views change:

julia> @btime mul!(x, A, A)
  10.594 ms (500 allocations: 1.98 MiB)

[mcognetta]

Sorry, I am using this thread as my personal notebook. I can get it slightly faster (seems to scale at about 10% faster from my benchmarks) and with zero allocations using:

julia> function mul!(C::UpperTriangular, A::UpperTriangular, B::UpperTriangular)
           m, n = size(B, 1), size(B, 2)
           if m != size(A, 1)
               throw(DimensionMismatch("right hand side B needs first dimension of size $(size(A,1)), has size $m"))
           end
           if C === A || C === B
               throw(ArgumentError("output matrix must not be aliased with input matrix"))
           end
           @views for i in 1:n
               mul!(parent(C)[:, i], A, B[:, i])
           end
           return C
       end

Current version (slicing):

julia> @btime mul!(x, A, A)
  11.551 ms (1000 allocations: 3.97 MiB)

The above version:

julia> @btime mul!(x, A, A)
  9.076 ms (0 allocations: 0 bytes)

But this strikes me as hacky and highly dependant on triangular matrices being backed by a dense matrix and parent being an O(1) time operation. Without parent(C) I get a method error:

ERROR: MethodError: no method matching lmul!(::UpperTriangular{Float64,Matrix{Float64}}, ::SubArray{Float64,1,UpperTriangular{Float64,Matrix{Float64}},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},false})

which stems from

julia/stdlib/LinearAlgebra/src/triangular.jl

Line 699 in 41e603e

mul!(C::AbstractVector , A::AbstractTriangular, B::AbstractVector) = lmul!(A, copyto!(C, B))

This can probably be fixed, but seems like more work than what should be done here.

Pinging @dkarrasch since you seem to be active in this area. Any thoughts?

[oxinabox]

The docstring of mul! says that normally one should extend 5 arg mul!, and 3 arg mul! will redispatch to that.

It seems like it would not be too hard to bring this PR inline with that.

[mcognetta]

Not ready for review anymore. Adding the 5-argument method causes a big slowdown due to allocations. I am working on it now.

[dkarrasch]

What about

Suggested change

	@views for i in 1:n
	C[:, i] = alpha*A*B[:, i] + beta*C[:, i]
	end
	for (Bi, Ci) in zip(eachcol(B), eachcol(C))
	mul!(Ci, A, Bi, alpha, beta)
	end

eachcol, zip, and 5-arg mul! should all be essentially allocation free.

[dkarrasch]

rand(Int,...) may yield very large integers, which may overflow upon multiplication and addition.

[dkarrasch]

For 3-arg mul!, there is another option to fix this:

function lmul!(A::UpperTriangular, B::UpperTriangular)
    lmul!(A, parent(B))
    return B
end

etc. This overwrites also the hidden part of parent(B), but hey, it's a mutating function.

EDIT: I forgot to erase the hidden terms, so one would have to run tril! or triu! on parent(B) first.

[dkarrasch]

Ok, so I did some more testing.

1. So my eachcol proposal did improve the performance significantly, but implementing the 5-arg mul! columnwise is not a good idea. For both small and large matrices, this is outperformed by the currently called generic_matmat_mul!. I'd suggest to simply remove those methods.

2. For the 3-arg mul!, I strongly suggest to do (in the big loop)

@eval mul!(C::$cty, A::$aty, B::$bty) = (lmul!(A, copyto!(parent(C), B)); C)

because this has then the chance to call BLAS triangular multiplication. That (a) solves the issue, and (b) is faster than A*B, because you avoid the allocation of the array behind C.

[dkarrasch]

Sorry, @mcognetta, I couldn't withstand pushing some changes and simplifying the tests. They were pretty heavy. Now, the tests also check for the correct output type.

[dkarrasch]

Bump?

[mcognetta]

@dkarrasch sorry for the delay. Thanks for your changes. It seems they are substantially faster than what I tried earlier. I don't think there is much else to do?

[mcognetta]

Added a few tests for the 5-arg version.

[dkarrasch]

There is the gc-failure, which I'm afraid of ignoring. I thought rebasing onto current master should help.

[oxinabox]

Since this is changing 3 arg mul! and not 5 arg mul! do we need to add a condition into the 5-arg mul! to redispatch mul!(C, A, B, false, true) to use it?

[dkarrasch]

There are two strange errors for which tests pass, though, and two clearly unrelated test failures that are known and tracked. I suggest to finally merge this.

[mcognetta]

Shall we close this since it's merged?

[dkarrasch]

Wait, no, it's not merged yet. I was seeing if anyone raised objections.