Unwrap triangular matrices in broadcast (#1332)

In broadcasting over triangular matrices, we loop only over the stored
indices. For these indices, indexing into a triangular matrix is
equivalent to indexing into its parent. We may therefore replace an
`UpperOrLowerTriangular` matrix by its parent, which removes the branch
in `getindex`. This improves performance:
```julia
julia> L = LowerTriangular(zeros(600,600));

julia> L2 = copy(L);

julia> @Btime broadcast!(+, $L2, $L, $L);
  161.176 μs (0 allocations: 0 bytes) # master
  80.894 μs (0 allocations: 0 bytes) # this PR
```
This replacement is performed recursively on a `Broadcasted` object by
looping over its `args`, and non-triangular elements are left untouched.
Only `UpperOrLowerTriangular` matrices will be replaced by their
`parent`s.

(cherry picked from commit 5165fd3)

Author: jishnub

src/structuredbroadcast.jl

@@ -264,13 +264,24 @@ function copyto!(dest::Tridiagonal, bc::Broadcasted{<:StructuredMatrixStyle})
     return dest
 end
 
+# Recursively replace wrapped matrices by their parents to improve broadcasting performance
+# We may do this because the indexing within `copyto!` is restricted to the stored indices
+preprocess_broadcasted(::Type{T}, A) where {T} = _preprocess_broadcasted(T, A)
+function preprocess_broadcasted(::Type{T}, bc::Broadcasted) where {T}
+    args = map(x -> preprocess_broadcasted(T, x), bc.args)
+    Broadcast.Broadcasted(bc.f, args, bc.axes)
+end
+_preprocess_broadcasted(::Type{LowerTriangular}, A) = lowertridata(A)
+_preprocess_broadcasted(::Type{UpperTriangular}, A) = uppertridata(A)
+
 function copyto!(dest::LowerTriangular, bc::Broadcasted{<:StructuredMatrixStyle})
     isvalidstructbc(dest, bc) || return copyto!(dest, convert(Broadcasted{Nothing}, bc))
     axs = axes(dest)
     axes(bc) == axs || Broadcast.throwdm(axes(bc), axs)
+    bc_unwrapped = preprocess_broadcasted(LowerTriangular, bc)
     for j in axs[2]
         for i in j:axs[1][end]
-            @inbounds dest.data[i,j] = bc[CartesianIndex(i, j)]
+            @inbounds dest.data[i,j] = bc_unwrapped[CartesianIndex(i, j)]
         end
     end
     return dest
@@ -280,9 +291,10 @@ function copyto!(dest::UpperTriangular, bc::Broadcasted{<:StructuredMatrixStyle}
     isvalidstructbc(dest, bc) || return copyto!(dest, convert(Broadcasted{Nothing}, bc))
     axs = axes(dest)
     axes(bc) == axs || Broadcast.throwdm(axes(bc), axs)
+    bc_unwrapped = preprocess_broadcasted(UpperTriangular, bc)
     for j in axs[2]
         for i in 1:j
-            @inbounds dest.data[i,j] = bc[CartesianIndex(i, j)]
+            @inbounds dest.data[i,j] = bc_unwrapped[CartesianIndex(i, j)]
         end
     end
     return dest

src/triangular.jl

@@ -658,10 +658,8 @@ end
 
 uppertridata(A) = A
 lowertridata(A) = A
-# we restrict these specializations only to strided matrices to avoid cases where an UpperTriangular type
-# doesn't share its indexing with the parent
-uppertridata(A::UpperTriangular{<:Any, <:StridedMatrix}) = parent(A)
-lowertridata(A::LowerTriangular{<:Any, <:StridedMatrix}) = parent(A)
+uppertridata(A::UpperTriangular) = parent(A)
+lowertridata(A::LowerTriangular) = parent(A)
 
 @inline _rscale_add!(A::AbstractTriangular, B::AbstractTriangular, C::Number, alpha::Number, beta::Number) =
     @stable_muladdmul _triscale!(A, B, C, MulAddMul(alpha, beta))

test/structuredbroadcast.jl

@@ -385,4 +385,12 @@ end
     @test ind == CartesianIndex(1,1)
 end
 
+@testset "nested triangular broadcast" begin
+    for T in (LowerTriangular, UpperTriangular)
+        L = T(rand(Int,4,4))
+        M = Matrix(L)
+        @test L .+ L .+ 0 .+ L .+ 0 .- L == 2M
+    end
+end
+
 end