# Tensor multiplication

**URL:** <https://discourse.julialang.org/t/tensor-multiplication/21450>\
**Category:** Performance\
**Tags:** arrays, tensors\
**Created:** [March 4, 2019, 11:47am UTC](https://discourse.julialang.org/t/tensor-multiplication/21450 "2019-03-04T11:47:54Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![Michiel](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/michiel/32/222428_2.png) [@Michiel](https://discourse.julialang.org/u/Michiel)\
**Post date:** [March 4, 2019, 11:47am UTC](https://discourse.julialang.org/t/tensor-multiplication/21450/1 "2019-03-04T11:47:54Z")

</div>

Hi all,

I am implementing a function to perform a generalization of matrix multiplication to a general N-dimensional array or tensor. This product is denoted as \times\_m to multiply a conformable matrix A with a tensor \mathcal{X} according to dimension n.

A working example is given below (note, I already tried several things to make it more performant: @inbounds, efficient looping over array elements etc).

```julia
function tensormult4!(X, A, Y, T, dim)
    @assert size(X, dim) == size(A, 2)
    Tmax = last(CartesianIndices(T))
    m, n = size(A)
    @simd for k in 1:n
        fill!(T, zero(eltype(T)))
        @simd for I in CartesianIndices(X)
            @inbounds T[min(Tmax, I)] += A[k, I[dim]] * X[I]
        end
        @simd for I in CartesianIndices(Y)
            if I[dim] == k # remove this?
                @inbounds Y[I] = T[min(Tmax, I)]
            end
        end
    end
    return Y
end

```

a working example:

```julia
A = rand(4, 4)
X = rand(3, 4, 2)
Y = similar(X)
T = sum(X, dims=2) # T is a temporary placeholder

tensormult4!(X, A, Y, T, 2)

```

Code seems to run not too slowly, but I suspect there is quite a bit of low hanging fruit.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [March 4, 2019, 12:41pm UTC](https://discourse.julialang.org/t/tensor-multiplication/21450/2 "2019-03-04T12:41:36Z")

</div>

This doesn’t run for me at the moment. Can you say more precisely what you want to achieve?

Many such operations can be written with Einsum.jl or TensorOperations.jl, here is one reading of how your Y is related to X and A:

```julia
using Einsum
@einsum Y[i,a,k] := A[a,b] * X[i,b,k] # sum over b=1:4

```

---

<div class="post-metadata">

**Author:** ![chakravala](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chakravala/32/6832_2.png) [@chakravala](https://discourse.julialang.org/u/chakravala)\
**Post date:** [March 4, 2019, 3:10pm UTC](https://discourse.julialang.org/t/tensor-multiplication/21450/3 "2019-03-04T15:10:16Z")

</div>

My package [Grassmann.jl](https://github.com/chakravala/Grassmann.jl) is designed to work with high dimensional conformal tensor algebras, but is still a work in progress. Due to sparse multivector representation and design, it outperforms “higher dimensional arrays” in multiple applications, as more sparse methods get implemented.

---

<div class="post-metadata">

**Author:** ![Michiel](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/michiel/32/222428_2.png) [@Michiel](https://discourse.julialang.org/u/Michiel)\
**Post date:** [March 4, 2019, 3:13pm UTC](https://discourse.julialang.org/t/tensor-multiplication/21450/4 "2019-03-04T15:13:01Z")

</div>

There was a mistake in initializing `T`, which is now fixed in the example.

I want to build indeed similar code as Einsum.jl and TensorOperations.jl. I not that my code seems to be faster than `@einsum` but substantially slower than `@tensor`.

At the moment, the reason is mainly autodidactic. It is not clear from the manual how TensorOperations.jl is so fast and how I can improve my own code.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [March 4, 2019, 4:27pm UTC](https://discourse.julialang.org/t/tensor-multiplication/21450/5 "2019-03-04T16:27:19Z")

</div>

Good, now it works. But I think your loops can be much simpler, there is no need for temporary storage `T`.

Einsum just generates a set of loops, which you can examine by writing `@macroexpand @einsum Y[i,a,k] = A[a,b] * X[i,b,k]`. (They may have a few `let` blocks which IIRC were needed for Julia 0.6.).

TensorOperations is much more complicated, but its goal here is essentially to call `mul!(A,B,C)` on the arrays. You could simulate this by writing `reshape(permutedims(X, (2, 1, 3)),size(X,2),:)` to make a matrix which you can multiply by `A`, and then un-doing that at the end. This should be faster on large arrays, BLAS is optimised in lots of clever ways. But for these small ones it seems to be beaten by Einsum:

```julia
julia> using TensorOperations, Einsum
julia> ein(X,A,Y) = @einsum Y[i,a,k] = A[a,b] * X[i,b,k]
julia> ten(X,A,Y) = @tensor Y[i,a,k] = A[a,b] * X[i,b,k]

julia> @btime tensormult4!($X, $A, $Y, $T, 2);
  300.207 ns (0 allocations: 0 bytes)

julia> @btime ein($X, $A, $Y);
  72.467 ns (0 allocations: 0 bytes)

julia> @btime ten($X, $A, $Y);
  1.760 μs (42 allocations: 2.66 KiB)

```

If you want to treat N-dimensional arrays at once, then things will get more complicated… you may be able to write things like `reshape(X,size(X,1),:)`, but you may have to get into `ntuple` indices and all that, especially if you want an argument `dims` like this (which may want to be `Val(dims)` etc.)

---

<div class="post-metadata">

**Author:** ![Michiel](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/michiel/32/222428_2.png) [@Michiel](https://discourse.julialang.org/u/Michiel)\
**Post date:** [March 20, 2019, 1:23pm UTC](https://discourse.julialang.org/t/tensor-multiplication/21450/6 "2019-03-20T13:23:45Z")

</div>

I rewrote the function using a reshape rather than for-loops, this gives a better running time, but of course much more memory use. I think it works faster than Einsum, but still much slower than TensorOperations.

```julia
function tensormult5!(X, A, Y, dim)

    @assert size(X, dim) == size(A, 2)
    m, n = size(A)
    s = collect(1:ndims(X))
    d = [size(Y)...]
    for i in 1:(2-1)
       s[i], s[i+1] = s[i+1], s[i]
    end
    T = A * reshape(permutedims(X, s), size(X, dim), :)
    Y[:] = permutedims(reshape(T, d[s]...), s)
    return Y
end

```

So benchmarking gives:

```julia
julia> using BenchmarkTools

julia> using TensorOperations, Einsum

julia> ein(X,A,Y) = @einsum Y[i,a,k] = A[a,b] * X[i,b,k]
ein (generic function with 1 method)

julia> ten(X,A,Y) = @tensor Y[i,a,k] = A[a,b] * X[i,b,k]
ten (generic function with 1 method)

julia> X = rand(100, 200, 300);

julia> A = rand(200, 200);

julia> Y = similar(X);

julia> dim = 2;

julia> @btime tensormult5!($X, $A, $Y, $dim);
  63.654 ms (24 allocations: 137.33 MiB)

julia> @btime ein($X, $A, $Y);
  1.277 s (0 allocations: 0 bytes)

julia> @btime ten($X, $A, $Y);
  39.209 ms (328 allocations: 11.88 KiB)

```

Is reshaping arrays for such applications the most efficient? Because in general I have to copy the complete array so that it corresponds to the right multiplication. I would expect that for-loops might be more efficient because no data is copied.

---

<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:** [March 20, 2019, 1:39pm UTC](https://discourse.julialang.org/t/tensor-multiplication/21450/7 "2019-03-20T13:39:53Z")

</div>

The trick for performance is having the correct loops.

For example, optimized matrix multiplication is actually done with at least five loops, instead of the naive three:  
The three expected loops on a small computational kernel. This kernel is small enough so that the entire output fits on the CPU registers. Eg, create/update 8x6 block blocks at a time on a computer with avx2 (or a 16x14 block at a time with avx512).  
Then you loop over this kernel.

You’d need a similar approach for higher dimensional tensors.

Also, I think you can avoid copying data with the BLAS version.

You could take a look at this:

> **[TACO: The Tensor Algebra Compiler](http://tensor-compiler.org)**
>
> Website for the TACO project

IIRC, someone was working on something like that for Julia.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [March 20, 2019, 2:01pm UTC](https://discourse.julialang.org/t/tensor-multiplication/21450/8 "2019-03-20T14:01:15Z")

</div>

Reshaping is quite efficient, but `permutedims` makes a copy. You can experiment with `PermutedDimsArray` which is its lazy cousin, however if you use that in here:

> [@Michiel](#):
>
> A \* reshape(permutedims(X, s), …

then it seems like `*` falls back to some slow-but-generic code, giving times similar to the `@einsum` version. But using that for copying into `Y`, or else `permutedims!(Y,...`, will save you a copy.

As @Elrod says, the efficient BLAS `*` is much more complicated than 3 loops, for reasons of cache-efficiency. TensorOperations is doing something similar, and you can also use its back-end Strided.jl alone. Adding `@strided` I think replaces `permutedims` with a lazy version specific to that package, and likewise replaces `*` with something a bit more flexible. Both this and BLAS should be multi-threaded; I believe that Strided uses Julia’s threads (e.g. `Threads.nthreads()`) but don’t swear to it.

---

<div class="post-metadata">

**Author:** ![Michiel](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/michiel/32/222428_2.png) [@Michiel](https://discourse.julialang.org/u/Michiel)\
**Post date:** [March 20, 2019, 2:22pm UTC](https://discourse.julialang.org/t/tensor-multiplication/21450/9 "2019-03-20T14:22:51Z")

</div>

Thanks for the hints @mcabbott and @Elrod!

I gather that my code cannot improved much implementation wise, but should be rather optimized at the level of the algorithm.

I tried `@strided`, but this does not seem to do anything for time and memory use. As far as I understand, this is what makes `TensorOperations.jl` so fast.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [March 20, 2019, 2:47pm UTC](https://discourse.julialang.org/t/tensor-multiplication/21450/10 "2019-03-20T14:47:34Z")

</div>

> [@Elrod](#):
>
> You could take a look at this:  
> [http://tensor-compiler.org](http://tensor-compiler.org)
> 
> IIRC, someone was working on something like that for Julia.

There was a nice talk on this by Peter Ahrens, including the relationship between Tortilla (which doesn’t seem to be public yet) and Taco: [https://youtu.be/Rp7sTl9oPNI](https://youtu.be/Rp7sTl9oPNI)