# Fast \`diag(A' \* B \* A)\`

**URL:** <https://discourse.julialang.org/t/fast-diag-a-b-a/98216>\
**Category:** General Usage\
**Tags:** performance, linearalgebra\
**Created:** [May 2, 2023, 5:46pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216 "2023-05-02T17:46:25Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![e3c6](https://avatars.discourse-cdn.com/v4/letter/e/e79b87/32.png) [@e3c6](https://discourse.julialang.org/u/e3c6)\
**Post date:** [May 2, 2023, 5:46pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/1 "2023-05-02T17:46:25Z")

</div>

Let `A` be a k \times n matrix and `B` a k\times k matrix, with k \ll n.

I want to compute the diagonal

```julia
diag(A' * B * A)

```

in a faster way. Here `A' * B * A` is a large n \times n matrix, which I would like to avoid instantiating.

Is there a faster (BLAS?) way to do this operation?

---

<div class="post-metadata">

**Author:** ![e3c6](https://avatars.discourse-cdn.com/v4/letter/e/e79b87/32.png) [@e3c6](https://discourse.julialang.org/u/e3c6)\
**Post date:** [May 2, 2023, 5:48pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/2 "2023-05-02T17:48:15Z")

</div>

Here is a possibility:

```julia
"""
    diagprod(A, B)

Compute `diag(A' * B * A)` in an efficient way.
"""
function diagprod(A::AbstractMatrix, B::AbstractMatrix)
    @assert size(A, 1) == size(B, 1) == size(B, 2)
    return vec(sum(A .* (B * A); dims=1))
end

```

Can we make it faster?

---

<div class="post-metadata">

**Author:** ![tbeason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tbeason/32/15898_2.png) [@tbeason](https://discourse.julialang.org/u/tbeason)\
**Post date:** [May 2, 2023, 6:14pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/3 "2023-05-02T18:14:02Z")

</div>

Wouldn’t a loop be faster? Or maybe Tullio can provide a happy medium?

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [May 2, 2023, 6:24pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/4 "2023-05-02T18:24:25Z")

</div>

I think a loop with LoopVectorization would blow this out of the water by removing the need for allocations

---

<div class="post-metadata">

**Author:** ![mikmoore](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mikmoore/32/31109_2.png) [@mikmoore](https://discourse.julialang.org/u/mikmoore)\
**Post date:** [May 2, 2023, 6:34pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/5 "2023-05-02T18:34:03Z")

</div>

Your `diagprod` is _algorithmically_ about as fast as one could expect and uses BLAS for the multiplication. I wouldn’t expect you to do much better. The largest remaining opportunities likely lie in reducing allocations.

You could pass in workspace variables to avoid the allocation for `B * A`, but that’s a bit of a nuisance and likely won’t save a ton. But you can do the `A .* BA` calculation in-place to re-use an existing allocation:

```julia
BA = B * A
BA .= conj.(A) .* BA
return vec(sum(BA; dims=1))

```

(note that I added `conj` to the left `A` since it’s actually `A'`).

Further, you can avoid a memory pass through the data by simply using

```julia
map(LinearAlgebra.dot, eachcol(A), eachcol(B * A))

```

Something with `LoopVectorization` or `Tullio` could be a little faster, I expect. However, if you don’t use BLAS for the `B * A` calculation you’ll likely lose performance at large matrix sizes and the remaining `sum(conj.(A) .* BA; dims=1)` piece only has so much room for improvement.

If you’re dealing with small matrices, there’s much more room for non-BLAS solutions to shine.

EDIT: at least for some input sizes, posters below have demonstrated speedups notably bigger than what I might have anticipated by using `LoopVectorization` and related packages.

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [May 2, 2023, 6:51pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/6 "2023-05-02T18:51:09Z")

</div>

I assume you meant that you wanted `A` to be k \times n? Here’s a few options:

```julia
#+begin_src julia
using Tullio, LoopVectorization

f1(A, B) = diag(A' * B * A)
f2(A, B) = vec(sum(A .* (B * A); dims=1))
f3(A, B) = @tullio D[i] := A[j, i] * B[j, k] * A[k, i]

function f4(A::Matrix{T}, B::Matrix{U}) where {T, U}
    V = promote_type(T, U)
    @assert size(A, 1) == size(B, 1) == size(B, 2)
    D = zeros(V, size(A, 2))
    @tturbo for i ∈ eachindex(D)
        for j ∈ axes(B, 1)
            for k ∈ axes(B, 2)
                D[i] += conj(A[j, i]) * B[j, k] * A[k, i]
            end
        end
    end
    D
end
f5(A, B) = map(LinearAlgebra.dot, eachcol(A), eachcol(B * A))

function f6(A::Matrix{T}, B::Matrix{U}) where {T, U}
    V = promote_type(T, U)
    @assert size(A, 1) == size(B, 1) == size(B, 2)
    D = Vector{promote_type(T,U)}(undef, size(A, 2))
    @tturbo for i ∈ eachindex(D)
        di = zero(eltype(D))
        for j ∈ axes(B, 1)
            for k ∈ axes(B, 2)
                di += conj(A[j, i]) * B[j, k] * A[k, i]
            end
        end
        D[i] = di
    end
end

let n = 1000, k = 10
    A = randn(k, n)
    B = randn(k, k)
    for f ∈ (f1, f2, f3, f4, f5, f6)
        print(f, " ")
        @btime $f($A, $B)
    end
end;

#+end_src

```

```julia
#+RESULTS:
: f1 599.030 μs (5 allocations: 7.71 MiB)
: f2 21.550 μs (7 allocations: 164.36 KiB)
: f3 20.880 μs (1 allocation: 7.94 KiB)
: f4 4.636 μs (1 allocation: 7.94 KiB)
: f5 23.940 μs (3 allocations: 86.11 KiB)
: f6 4.481 μs (1 allocation: 7.94 KiB)

```

Looks like writing the manual loop in LoopVectorization.jl is the winner here.

* * *

Edit1: I missed on the of the suggestions above so I added it as `f5`.  
Edit2: I messed up the order of indices for `f4` so it was faster than it should have been. However, even with this fix it’s still the fastest option, just by a lesser margin.  
Edit3: Fixed the problem pointed out here: [Fast `diag(A' * B * A)` - #9 by mikmoore](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/9). This again has a negative performance impact on `f4`, but it’s still the fastest.  
Edit4: I found yet another problem with `f3` and `f4` where I accidrntally wrote `B[j, k] * A[j, i]` instead of `B[j, k] * A[k, i]`. Fixing this slows down `f3` and `f4` a bit but `f3` is hit harder than `f4` and `f4` remains the fastest.  
Edit5: added another LoopVectorization.jl example that reduces the number of array accesses required so speeds things up a little.

---

<div class="post-metadata">

**Author:** ![mikmoore](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mikmoore/32/31109_2.png) [@mikmoore](https://discourse.julialang.org/u/mikmoore)\
**Post date:** [May 2, 2023, 7:17pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/7 "2023-05-02T19:17:59Z")

</div>

> [@Mason](#):
>
> `D = Vector{V}(undef, size(A, 1))`

Shouldn’t this be `size(A,2)`? This is doing 100x less work than it should be.

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [May 2, 2023, 7:20pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/8 "2023-05-02T19:20:53Z")

</div>

Oops, yes, it should be `size(A, 2)`, but it’s not doing `100x` more work it’s `100x` less work than it should because the original poster had the wrong worder of indices for `A` (I think). A should be `k, n` not `n, k` if it should be multiplying `B` as `A' * B * A`.

I’ll edit my comment with the correct timings

---

<div class="post-metadata">

**Author:** ![mikmoore](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mikmoore/32/31109_2.png) [@mikmoore](https://discourse.julialang.org/u/mikmoore)\
**Post date:** [May 2, 2023, 7:58pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/9 "2023-05-02T19:58:56Z")

</div>

> [@Mason](#):
>
> `D[i] = conj(A[j, i]) * B[j, k] * A[j, i]`

Sorry to not catch this all at once.

Shouldn’t this be `+=` and initialized to `D = zeros(V, size(A,2))`? Or maybe use a temporary variable `Di = zero(V)` to do the accumulation before saving to `D[i]`, if the compiler won’t hoist that on its own…

As written, the loop only does “work” on the final pair of `j,k`. The compiler may not be smart enough to have caught on, so it may not have a performance impact, but I’m suspicious that this function manages to be so much faster than even `B*A` despite that it should be doing “more” calculations. The bulk of that is likely the small dimension `k=10` giving a meaningful opportunity to beat BLAS.

EDIT: I think there may also be an indexing bug in the expression above. It looks like something more like

```julia
D[i] += A[j,i]' * B[j,k] * A[k,i]

```

is correct but I derived this using different variables and might have messed up the translation.

---

<div class="post-metadata">

**Author:** ![e3c6](https://avatars.discourse-cdn.com/v4/letter/e/e79b87/32.png) [@e3c6](https://discourse.julialang.org/u/e3c6)\
**Post date:** [May 2, 2023, 8:12pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/10 "2023-05-02T20:12:14Z")

</div>

> [@Mason](#):
>
> I assume you meant that you wanted `A` to be k \times n?

Whoops, yes. Corrected.

---

<div class="post-metadata">

**Author:** ![e3c6](https://avatars.discourse-cdn.com/v4/letter/e/e79b87/32.png) [@e3c6](https://discourse.julialang.org/u/e3c6)\
**Post date:** [May 2, 2023, 8:33pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/11 "2023-05-02T20:33:28Z")

</div>

Thanks a lot for the detailed response.

`f4` is not working for me, some issue with LoopVectorization. What version are you on?

Here is what I tried:

```julia
function diagprod(A::Matrix{T}, B::Matrix{U}) where {T, U} # same as Mason's f4
    V = promote_type(T, U)
    @assert size(A, 1) == size(B, 1) == size(B, 2)
    D = Vector{V}(undef, size(A, 2))
    @tturbo for i ∈ eachindex(D)
        for j ∈ axes(B, 1)
            for k ∈ axes(B, 2)
                D[i] = conj(A[j, i]) * B[j, k] * A[j, i]
            end
        end
    end
end

A = randn(5, 100)
B = randn(5, 5)
diag(A' * B * A) ≈ diagprod(A, B)

```

This leads to a complicated error. Brieflly:

```julia
ERROR: MethodError: no method matching __vstore!(::Ptr{Float64}, ::VectorizationBase.VecUnroll{3, 2, Float64, VectorizationBase.Vec{2, Float64}}, ::Static.StaticInt{0}, ::Static.False, ::Static.False, ::Static.False, ::Static.StaticInt{16})

```

See [stacktrace.txt · GitHub](https://gist.github.com/cossio/be8f6771c73e64697ce6986042b665e7) for details and stack trace. Any ideas?

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [May 2, 2023, 8:38pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/12 "2023-05-02T20:38:12Z")

</div>

> [@mikmoore](#):
>
> Shouldn’t this be `+=` and initialized to `D = zeros(V, size(A,2))`? Or maybe use a temporary variable `Di = zero(V)` to do the accumulation before saving to `D[i]`, if the compiler won’t hoist that on its own…
> 
> As written, the loop only does “work” on the final pair of `j,k`. The compiler may not be smart enough to have caught on, so it may not have a performance impact, but I’m suspicious that this function manages to be so much faster than even `B*A` despite that it should be doing “more” calculations. The bulk of that is likely the small dimension `k=10` giving a meaningful opportunity to beat BLAS.

Crap, yeah you’re right that I should have done `+=` on `zeros`. The good news is that it does not result in a big performance loss. Updating my post yet again.

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [May 2, 2023, 8:47pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/13 "2023-05-02T20:47:43Z")

</div>

> [@e3c6](#):
>
> `f4` is not working for me, some issue with LoopVectorization. What version are you on?

Hm, I can’t reproduce the problem you encountered. I’ve got

```julia
LoopVectorization v0.12.158

```

Can you try and see if

```julia
function diagprod(A::Matrix{T}, B::Matrix{U}) where {T, U}
    V = promote_type(T, U)
    @assert size(A, 1) == size(B, 1) == size(B, 2)
    D = zeros(size(A, 2))
    @tturbo for i ∈ eachindex(D)
        for j ∈ axes(B, 1)
            for k ∈ axes(B, 2)
                D[i] += conj(A[j, i]) * B[j, k] * A[k, i]
            end
        end
    end
    D
end

A = randn(5, 100)
B = randn(5, 5)
diag(A' * B * A) ≈ diagprod(A, B)

```

works for you?

---

<div class="post-metadata">

**Author:** ![e3c6](https://avatars.discourse-cdn.com/v4/letter/e/e79b87/32.png) [@e3c6](https://discourse.julialang.org/u/e3c6)\
**Post date:** [May 2, 2023, 8:50pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/14 "2023-05-02T20:50:30Z")

</div>

Yes that works. I see what’s different is you’re not using the `promote_type` anymore. Do you understand the problem?

> [@Mason](#):
>
> ```julia
> LoopVectorization v0.12.158
> 
> ```

Yep, I am using the same version of LoopVectorization.  
I am on Julia v1.8.5, macOS.

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [May 2, 2023, 8:52pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/15 "2023-05-02T20:52:38Z")

</div>

I’m still using `promote_type`. The difference was just that I was traversing the correct loops. Before we accidentally did `B[j, k] * A[j, i]` instead of `B[j, k] * A[k, i]`, but I don’t know why that would result in the error you saw, I didn’t get an error just a wrong answer 😅

---

<div class="post-metadata">

**Author:** ![e3c6](https://avatars.discourse-cdn.com/v4/letter/e/e79b87/32.png) [@e3c6](https://discourse.julialang.org/u/e3c6)\
**Post date:** [May 2, 2023, 8:55pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/16 "2023-05-02T20:55:43Z")

</div>

> [@Mason](#):
>
> I’m still using `promote_type`.

You’re calling `V = promote_type(T, U)` but then `V` is not used anywhere. However I confirm this is not the issue. Yes I see the index difference, that must be it.

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [May 2, 2023, 8:56pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/17 "2023-05-02T20:56:50Z")

</div>

Oh, I see, yeah it should have been `D = zeros(V, size(A, 2))`, but that won’t effect this test.

I’ll update [Fast `diag(A' * B * A)` - #6 by Mason](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/6)

---

<div class="post-metadata">

**Author:** ![e3c6](https://avatars.discourse-cdn.com/v4/letter/e/e79b87/32.png) [@e3c6](https://discourse.julialang.org/u/e3c6)\
**Post date:** [May 2, 2023, 8:58pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/18 "2023-05-02T20:58:36Z")

</div>

> [@Mason](#):
>
> I’ll update [Fast `diag(A' * B * A)` - #6 by Mason](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/6)

Thanks!

---

<div class="post-metadata">

**Author:** ![Dan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dan/32/42581_2.png) [@Dan](https://discourse.julialang.org/u/Dan)\
**Post date:** [May 2, 2023, 10:27pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/19 "2023-05-02T22:27:48Z")

</div>

Just for the record,

```julia
f6(A, B) = map(LinearAlgebra.dot, eachcol(A), Iterators.repeated(B), eachcol(A))

```

also exists, and uses the `dot(v, B, v)` notation, which specifically computes this needed form. It is slower than suggested `f4` but uses the same amount of memory.

The benchmark here is for specific `n` and `k` and the results may depend on those, so @e3c6 , if there are specific values you need, it would be helpful to know (and if any of the matrices are sparse would also be interesting).

---

<div class="post-metadata">

**Author:** ![Elrod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elrod/32/22461_2.png) [@Elrod](https://discourse.julialang.org/u/Elrod)\
**Post date:** [May 2, 2023, 10:37pm UTC](https://discourse.julialang.org/t/fast-diag-a-b-a/98216/20 "2023-05-02T22:37:09Z")

</div>

> [@e3c6](#):
>
> Any ideas?

Looks like some part of codegen is failing to reduce along one of the dimensions.  
I’d be happy to walk someone through fixing this if they’re interested, otherwise I may took a look later, at least to see if fixing this would be easy (which it might be).

Also, FWIW, it may be faster to do

```julia
    D = Vector{promote_type(T,U)}(undef, size(A, 2))
    @tturbo for i ∈ eachindex(D)
        di = zero(eltype(D))
        for j ∈ axes(B, 1)
            for k ∈ axes(B, 2)
                di += conj(A[j, i]) * B[j, k] * A[k, i]
            end
        end
        D[i] = di
    end

```

[Next page](https://discourse.julialang.org/t/fast-diag-a-b-a/98216.md?page=2)