# Efficient way of doing matmul along 3rd dimension

**URL:** <https://discourse.julialang.org/t/efficient-way-of-doing-matmul-along-3rd-dimension/77075>\
**Category:** General Usage\
**Created:** [February 25, 2022, 2:39pm UTC](https://discourse.julialang.org/t/efficient-way-of-doing-matmul-along-3rd-dimension/77075 "2022-02-25T14:39:17Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![HenriDeh](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/henrideh/32/8316_2.png) [@HenriDeh](https://discourse.julialang.org/u/HenriDeh)\
**Post date:** [February 25, 2022, 2:39pm UTC](https://discourse.julialang.org/t/efficient-way-of-doing-matmul-along-3rd-dimension/77075/1 "2022-02-25T14:39:17Z")

</div>

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:

```julia
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.

---

<div class="post-metadata">

**Author:** ![lmiq](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lmiq/32/18314_2.png) [@lmiq](https://discourse.julialang.org/u/lmiq)\
**Post date:** [February 25, 2022, 2:52pm UTC](https://discourse.julialang.org/t/efficient-way-of-doing-matmul-along-3rd-dimension/77075/2 "2022-02-25T14:52:10Z")

</div>

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
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
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).

---

<div class="post-metadata">

**Author:** ![aplavin](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/aplavin/32/222056_2.png) [@aplavin](https://discourse.julialang.org/u/aplavin)\
**Post date:** [February 25, 2022, 3:04pm UTC](https://discourse.julialang.org/t/efficient-way-of-doing-matmul-along-3rd-dimension/77075/3 "2022-02-25T15:04:57Z")

</div>

`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}}}})`.

---

<div class="post-metadata">

**Author:** ![HenriDeh](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/henrideh/32/8316_2.png) [@HenriDeh](https://discourse.julialang.org/u/HenriDeh)\
**Post date:** [February 25, 2022, 3:05pm UTC](https://discourse.julialang.org/t/efficient-way-of-doing-matmul-along-3rd-dimension/77075/4 "2022-02-25T15:05:35Z")

</div>

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

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

```

Should work ?