Efficient way of doing matmul along 3rd dimension

General Usage
1

Hi,

I’ve been looking for a way to do matmul operations along slices of the 3rd dims of a 3D tensor.
I’ve had a look at Tullio.jl and TensorCast.jl but I’m not sure they are designed to do this. Basically I want to do this:

A = rand(5,5,10)
B = rand(5,5,10) 
C = zeros(5,5,10)
for k in 1:10
    C[:,:,k] = A[:,:,k] * B[:,:,k]
end

It’s like a broadcast but only along the 3rd dims. Ideally, i could just use mapslices but it only accepts one collection at a time.
This MWE works fine but I’m guessing there exists a package with functions or macros that do this more efficiently.

2

Some alternatives. Maybe I am not using Tullio quite properly here. I have 4 threads, and mul! is generally multi-threaded because of BLAS:

julia> using Tullio, LinearAlgebra

julia> function f!(C,A,B)
           for k in 1:10
               C[:,:,k] = A[:,:,k] * B[:,:,k]
           end
       end
f! (generic function with 1 method)

julia> function f2!(C,A,B)
           for k in 1:10
               @views mul!(C[:,:,k],A[:,:,k],B[:,:,k])
           end
       end
f2! (generic function with 1 method)

julia> f3!(C,A,B) = @views @tullio C[:,:,k] = A[:,:,k] * B[:,:,k]
f3! (generic function with 1 method)

julia> A = rand(5,5,10); B = rand(5,5,10); C = zeros(5,5,10);

julia> @btime f!($C,$A,$B);
  3.387 μs (30 allocations: 7.50 KiB)

julia> @btime f2!($C,$A,$B);
  1.750 μs (0 allocations: 0 bytes)

julia> @btime f3!($C,$A,$B);
  2.240 μs (10 allocations: 2.50 KiB)

(edit: added @views to Tullio).

Octavian gives me this:

julia> using Octavian

julia> function f5!(C,A,B)
           for k in 1:10
               @views Octavian.matmul!(C[:,:,k],A[:,:,k],B[:,:,k])
           end
       end
f5! (generic function with 1 method)

julia> @btime f5!($C,$A,$B);
  196.513 ns (0 allocations: 0 bytes)

to good to be true?

(the timings are probably very dependent on the matrix sizes, so you probably need to test with the actual data).

1 Like
3

eachslice(A; dims=3) .* eachslice(B; dims=3) gives almost exactly what you need, but unfortunately assigning to eachslice(C; dims=3) throws a weird error: ERROR: MethodError: no method matching ndims(::Type{Base.Generator{Base.OneTo{Int64}, Base.var"#230#231"{Array{Float64, 3}, Tuple{}, Tuple{Colon, Colon}}}}).

4

Okay thanks I’ll take a look at Octavian too.
Though I like Tullio also because I (think I) can do other stuff than matmul along the dimension.
Like

foo(A::Matrix, B::Matrix) = A*max.(B,0.0)
@views @tullio C[:,:,k] = foo(A[:,:,k], B[:,:,k])

Should work ?