Batched dot(x, A, y)

[3f6a]

Motivation and description

Say I have the arrays x[i,b], y[j,b] and A[i,j,b]. Is there an efficient way to do the following "batched dot" operation:

[sum(x[i,b] * A[i,j,b] * y[j,b] for i = axes(A,1) for j = axes(A,2)) for b = ...]

where b traverses the batch dimension. As usual, we could have size(x,2) == 1, size(A,3)==1, ..., which would mean the corresponding missing dimension is broadcasted.

Apologies if there is already a way to do this (efficiently) with existing functions in NNlib, I could not figure it out.

Possible Implementation

No response

[mcabbott]

The code you sketch looks like it'll return a 3-array z[i,j,b], but I think you intend to sum i and j? If so, sorry, I misread.

There is nothing built into this package. But some ways you could write it are:

using NNlib, Tullio, TensorCast, LinearAlgebra

x = rand(10,7); A = rand(10,10,7); y = rand(10,7);

bdot1(x,A,y) = @views [dot(x[:,b], A[:,:,b], y[:,b]) for b in axes(x,2)]
bdot2(x,A,y) = @tullio z[b] := x[i,b] * A[i,j,b] * y[j,b] 
bdot3(x,A,y) = diag(transpose(x) * NNlib.batched_vec(A, y))  # NNlib computes Ay[i,b]
function bdot4(x,A,y)
    Ay = NNlib.batched_vec(A, y) 
    @tullio z[b] := x[i,b] * Ay[i,b]
end
bdot5(x,A,y) = @reduce z[b] := sum(i,j) x[i,b] * A[i,j,b] * y[j,b]

bdot1(x,A,y) ≈ bdot2(x,A,y) ≈ bdot3(x,A,y) ≈ bdot4(x,A,y) ≈ bdot5(x,A,y)  # true

let N = 50
    x = rand(N,N); A = rand(N,N,N); y = rand(N,N);
    @show N
    @btime bdot1($x,$A,$y);  # fastest, for me, 2x
    @btime bdot2($x,$A,$y);
    @btime bdot3($x,$A,$y);  # 10x worse
    @btime bdot4($x,$A,$y);
    @btime bdot5($x,$A,$y);
end;

The relative times are likely to depend a lot on how big the arrays are:

N = 50
  22.666 μs (102 allocations: 23.91 KiB)
  39.375 μs (2 allocations: 480 bytes)
  319.708 μs (39 allocations: 43.58 KiB)
  312.375 μs (36 allocations: 23.50 KiB)
  62.834 μs (9 allocations: 992.70 KiB)

N = 500
  23.158 ms (1503 allocations: 1.99 MiB)
  65.800 ms (50 allocations: 6.72 KiB)
  19.679 ms (40 allocations: 3.88 MiB)
  18.665 ms (37 allocations: 1.94 MiB)
  204.979 ms (10 allocations: 953.69 MiB)

[3f6a]

Nice that the simplest version

bdot1(x,A,y) = @views [dot(x[:,b], A[:,:,b], y[:,b]) for b in axes(x,2)]

is performant.

Could we add this as a function in NNlib ?

[mcabbott]

I'm somewhat surprised by these performance numbers, and recommend checking them on other systems, other sizes.

To add something like this to NNlib, the minimum standard is that it:

• Needs to work for GPU array -- at present bdot1 will still make a Vector, and will probably be slow. Maybe bdot3 is better.
• Needs to work with Zygote & Enzyme -- which probably means writing gradient rules, although maybe bdot3 works & is reasonably fast?

[3f6a]

Do you have any idea why @tullio seems to underperform here?