# Speed comparison matrix multiplication in Julia

**URL:** <https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656>\
**Category:** Performance\
**Tags:** question, linearalgebra, optimization, tullio\
**Created:** [August 19, 2021, 6:20am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656 "2021-08-19T06:20:00Z")\
**Posts on this page:** 20\
**Page:** 2

<div class="post-metadata">

**Author:** ![jling](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jling/32/212909_2.png) [@jling](https://discourse.julialang.org/u/jling)\
**Post date:** [August 19, 2021, 9:18am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/21 "2021-08-19T09:18:29Z")

</div>

you know Jax is not single threaded right?:

```julia
In [3]: for x in range(200):
   ...: jnp.matmul(a,b)
   ...: 

```

![image](https://global.discourse-cdn.com/julialang/original/3X/2/6/266f8be2482214822635abb9d412a2ed047ddaff.png)

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [August 19, 2021, 9:25am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/22 "2021-08-19T09:25:20Z")

</div>

> [@sivakon](#):
>
> I gain no speed up in single threaded

Relative to what? Are you running BLAS single-threaded? You can check with

```julia
BLAS._get_num_threads()

```

and set it to one thread with

```julia
BLAS.set_num_threads(1)

```

As for orientation, there is a _huge_ difference between row and column major here:

```julia
jl> a = Array(reshape(Int32.(1:2*2000*400), 2,2000,400));

jl> b = Array(reshape(Int32.(1:2*2000*400), 2,400,2000));

jl> @btime @tullio c[i, j, k] := $a[i, j, q] * $b[i, q, k];
  170.893 ms (116 allocations: 30.52 MiB)

```

vs

```julia
l> a = Array(reshape(Int32.(1:2*2000*400), 2000,400,2));

jl> b = Array(reshape(Int32.(1:2*2000*400), 400,2000,2));

jl> @btime @tullio c[j, k, i] := $a[j, q, i] * $b[q, k, i];
  24.013 ms (116 allocations: 30.52 MiB)

```

---

<div class="post-metadata">

**Author:** ![sivakon](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sivakon/32/5691_2.png) [@sivakon](https://discourse.julialang.org/u/sivakon)\
**Post date:** [August 19, 2021, 9:27am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/23 "2021-08-19T09:27:40Z")

</div>

```julia
import jax.numpy as jnp
from jax.config import config
config.update("jax_enable_x64", False)

import os
os.environ["XLA_FLAGS"] = ("--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=0")

a = jnp.arange(2 * 2000 * 400).reshape((2, 2000, 400)) + 1
b = jnp.arange(2 * 400 * 2000).reshape((2, 400, 2000)) + 1

print(a.dtype) # int32

%%timeit
c = jnp.matmul(a,b) # 291ms

```

---

<div class="post-metadata">

**Author:** ![jling](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jling/32/212909_2.png) [@jling](https://discourse.julialang.org/u/jling)\
**Post date:** [August 19, 2021, 9:32am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/24 "2021-08-19T09:32:11Z")

</div>

```julia
In [1]: import jax.numpy as jnp
   ...: from jax.config import config
   ...: config.update("jax_enable_x64", False)
   ...: 
   ...: import os
   ...: os.environ["XLA_FLAGS"] = ("--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=0")
   ...: 
   ...: a = jnp.arange(2 * 2000 * 400).reshape((2, 2000, 400)) + 1
   ...: b = jnp.arange(2 * 400 * 2000).reshape((2, 400, 2000)) + 1

In [3]: %timeit jnp.matmul(a,b)
1 s ± 1.98 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

```

```julia
julia> function mm(A,B)
           C = Array{Int32, 3}(undef, size(A,1),size(A,2), size(B,3))
           for i in axes(A,1)
               @views C[i,:,:] = A[i,:,:] * B[i,:,:]
           end
           C
       end
mm (generic function with 1 method)

julia> a = reshape(Int32.(1:2*2000*400), 2,2000,400);

julia> b = reshape(Int32.(1:2*2000*400), 2,400,2000);

julia> using BenchmarkTools

julia> import LinearAlgebra

julia> LinearAlgebra.BLAS.set_num_threads(1)

julia> @btime mm($a,$b);
  827.518 ms (18 allocations: 61.04 MiB)

```

Julia already faster

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [August 19, 2021, 9:32am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/25 "2021-08-19T09:32:25Z")

</div>

If you don’t correct for major dimension orientation, these comparisons won’t be very meaningful.

---

<div class="post-metadata">

**Author:** ![sivakon](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sivakon/32/5691_2.png) [@sivakon](https://discourse.julialang.org/u/sivakon)\
**Post date:** [August 19, 2021, 9:32am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/26 "2021-08-19T09:32:31Z")

</div>

```julia
➜ julia -q ~
julia> using BenchmarkTools, Tullio, LinearAlgebra

julia> BLAS.set_num_threads(1)

julia> a = Array(reshape(Int32.(1:2*2000*400), 2,2000,400));

julia> b = Array(reshape(Int32.(1:2*2000*400), 2,400,2000));

julia> @btime @tullio c[i, j, k] := $a[i, j, q] * $b[i, q, k];
  2.304 s (2 allocations: 30.52 MiB)

```

---

<div class="post-metadata">

**Author:** ![jling](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jling/32/212909_2.png) [@jling](https://discourse.julialang.org/u/jling)\
**Post date:** [August 19, 2021, 9:33am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/27 "2021-08-19T09:33:44Z")

</div>

> [@sivakon](#):
>
> `2,2000,400`

You want: `2000, 400, 2`, also, Tullio doesn’t use BLAS

---

<div class="post-metadata">

**Author:** ![sivakon](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sivakon/32/5691_2.png) [@sivakon](https://discourse.julialang.org/u/sivakon)\
**Post date:** [August 19, 2021, 9:35am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/28 "2021-08-19T09:35:04Z")

</div>

```julia
➜ jupyter console ~
Jupyter console 6.4.0

Python 3.7.10 (default, Apr 27 2021, 08:49:44)
Type 'copyright', 'credits' or 'license' for more information
IPython 7.25.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import jax.numpy as jnp
   ...: from jax.config import config
   ...: config.update("jax_enable_x64", False)
   ...:
   ...: import os
   ...: os.environ["XLA_FLAGS"] = ("--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=0")
   ...:
   ...: a = jnp.arange(2 * 2000 * 400).reshape((2, 2000, 400)) + 1
   ...: b = jnp.arange(2 * 400 * 2000).reshape((2, 400, 2000)) + 1
   ...:
   ...: print(a.dtype) # int32
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
int32

In [2]: %%timeit
   ...: c = jnp.matmul(a,b)
   ...:
   ...:
395 ms ± 8.41 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

```

I am not sure why your Jax is slower

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [August 19, 2021, 9:37am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/29 "2021-08-19T09:37:18Z")

</div>

This is getting confusing.

Changing the number of BLAS threads should change the performance of the `mm` function, if you are comparing that to Tullio, it does not affect Tullio.

In order to compare single-threaded performance, you should

- Start Julia with a single thread
- set BLAS threads to 1
- set jax to use one thread
- Then compare all three

For multi-threading, do the same, but set up each to have the same number of threads as you have physical cores.

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [August 19, 2021, 9:38am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/30 "2021-08-19T09:38:21Z")

</div>

And, very importantly: change orientation of your arrays for Julia, since Julia is column major.

---

<div class="post-metadata">

**Author:** ![jling](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jling/32/212909_2.png) [@jling](https://discourse.julialang.org/u/jling)\
**Post date:** [August 19, 2021, 9:39am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/31 "2021-08-19T09:39:06Z")

</div>

it maybe due to CPU, can you check `top` and run it in a loop (like what I did) to make sure Jax is indeed using single thread? the fact that your single thread vs. multi only changed 5% is weird, on my end, I see 1s → 40ms difference with 48 cores.

And you really should use the correct orientation (i.e. memory layout) to be fair.

---

<div class="post-metadata">

**Author:** ![sivakon](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sivakon/32/5691_2.png) [@sivakon](https://discourse.julialang.org/u/sivakon)\
**Post date:** [August 19, 2021, 10:10am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/32 "2021-08-19T10:10:47Z")

</div>

I see your point. Seems like Jax is using multi-threading on my Mac. Same code and flags give the result you shared on Linux. This might be a bug in Jax on Mac.

---

<div class="post-metadata">

**Author:** ![sivakon](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sivakon/32/5691_2.png) [@sivakon](https://discourse.julialang.org/u/sivakon)\
**Post date:** [August 19, 2021, 10:47am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/33 "2021-08-19T10:47:26Z")

</div>

Updated the post, take a look. Thanks for the help

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [August 19, 2021, 10:51am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/34 "2021-08-19T10:51:18Z")

</div>

@sivakon I apologize for repeating myself, but why do you not fix the major axis orientation issue? Isn’t it correct that that jax (like numpy) uses row-major arrays? In that case these results are not comparable. Is there something I’m missing, since you’re not addressing this concern?

In the benchmark I posted further up, changing orientation leads to a 7x speedup on my computer.

---

<div class="post-metadata">

**Author:** ![jling](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jling/32/212909_2.png) [@jling](https://discourse.julialang.org/u/jling)\
**Post date:** [August 19, 2021, 11:03am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/35 "2021-08-19T11:03:19Z")

</div>

so the latest result suggested that Julia is not slower than Jax even when we’re using the awkward memory layout because col-major vs. row-major? good to know, because Jax is Numpy re-write with some TF runtime and LLVM backend that compiles kernel on its own

also you really should fix col-major vs. row-major to make it apple to apple

---

<div class="post-metadata">

**Author:** ![sivakon](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sivakon/32/5691_2.png) [@sivakon](https://discourse.julialang.org/u/sivakon)\
**Post date:** [August 19, 2021, 11:13am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/36 "2021-08-19T11:13:03Z")

</div>

```julia
using Tullio, BenchmarkTools
a = Array(reshape(Int32.(1:2*2000*400), 2,2000,400));
b = Array(reshape(Int32.(1:2*2000*400), 2,400,2000));
@btime @tullio c[i, j, k] := $a[i, j, q] * $b[i, q, k]; 
# 2.212 s (2 allocations: 30.52 MiB)

using Tullio, BenchmarkTools
a = Array(reshape(Int32.(1:2*2000*400), 2000,400,2));
b = Array(reshape(Int32.(1:2*2000*400), 400,2000,2));
@btime @tullio c[j, k, i] := $a[j, q, i] * $b[q, k, i]; 
# 1.854 s (2 allocations: 30.52 MiB)

```

Is this what you mean? I don’t see any speed-up here.

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [August 19, 2021, 11:28am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/37 "2021-08-19T11:28:13Z")

</div>

That is strange. Here are my results:

BTW, did you use

```julia
using Tullio, LoopVectorization

```

?

Single-threaded:

```julia
jl> Threads.nthreads()
1

jl> using Tullio, LoopVectorization

jl> a = Array(reshape(Int32.(1:2*2000*400), 2,2000,400));

jl> b = Array(reshape(Int32.(1:2*2000*400), 2,400,2000));

jl> @btime @tullio c[i, j, k] := $a[i, j, q] * $b[i, q, k];
  713.923 ms (2 allocations: 30.52 MiB)

jl> a = Array(reshape(Int32.(1:2*2000*400), 2000,400,2));

jl> b = Array(reshape(Int32.(1:2*2000*400), 400,2000,2));

jl> @btime @tullio c[j, k, i] := $a[j, q, i] * $b[q, k, i];
  133.608 ms (2 allocations: 30.52 MiB)

```

8 threads:

```julia
jl> Threads.nthreads()
8

jl> using Tullio, LoopVectorization

jl> a = Array(reshape(Int32.(1:2*2000*400), 2,2000,400));

jl> b = Array(reshape(Int32.(1:2*2000*400), 2,400,2000));

jl> @btime @tullio c[i, j, k] := $a[i, j, q] * $b[i, q, k];
  154.384 ms (117 allocations: 30.52 MiB)

jl> a = Array(reshape(Int32.(1:2*2000*400), 2000,400,2));

jl> b = Array(reshape(Int32.(1:2*2000*400), 400,2000,2));

jl> @btime @tullio c[j, k, i] := $a[j, q, i] * $b[q, k, i];
  23.660 ms (117 allocations: 30.52 MiB)

```

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [August 19, 2021, 11:32am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/38 "2021-08-19T11:32:53Z")

</div>

Mystery solved. Using _just_ Tullio _without_ LoopVectorization yields times around 1.8s. if I have

```julia
using Tullio, LoopVectorization

```

it gives me a 2.5x speedup with row-major orientation, but a _13.5_x speedup for colum major.

But anyway, I was just getting a bit exasperated that numerous requests to fix the orientation issue were not acknowledged or noticed.

---

<div class="post-metadata">

**Author:** ![carstenbauer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/carstenbauer/32/4981_2.png) [@carstenbauer](https://discourse.julialang.org/u/carstenbauer)\
**Post date:** [August 19, 2021, 11:38am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/39 "2021-08-19T11:38:00Z")

</div>

Can confirm this:

```julia
julia> using Tullio, LoopVectorization, BenchmarkTools

julia> a = Array(reshape(Int32.(1:2*2000*400), 2,2000,400));

julia> b = Array(reshape(Int32.(1:2*2000*400), 2,400,2000));

julia> @btime @tullio c[i, j, k] := $a[i, j, q] * $b[i, q, k];
  1.415 s (2 allocations: 30.52 MiB)

julia> a = Array(reshape(Int32.(1:2*2000*400), 2000,400,2));

julia> b = Array(reshape(Int32.(1:2*2000*400), 400,2000,2));

julia> @btime @tullio c[j, k, i] := $a[j, q, i] * $b[q, k, i];
  251.543 ms (2 allocations: 30.52 MiB)

```

Without LoopVectorization it takes 3 - 4 seconds with the latter case being only about 20% faster.

---

<div class="post-metadata">

**Author:** ![sivakon](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sivakon/32/5691_2.png) [@sivakon](https://discourse.julialang.org/u/sivakon)\
**Post date:** [August 19, 2021, 11:52am UTC](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656/40 "2021-08-19T11:52:13Z")

</div>

> [@DNF](#):
>
> But anyway, I was just getting a bit exasperated that numerous requests to fix the orientation issue were not acknowledged or noticed.

I wanted to fix my issue first, have apples to apples comparison when comparing speed between Julia and Jax, and later optimize my Julia code.

```julia
julia> using Tullio, LoopVectorization, BenchmarkTools

julia> a = Array(reshape(Int32.(1:2*2000*400), 2,2000,400));

julia> b = Array(reshape(Int32.(1:2*2000*400), 2,400,2000));

julia> @btime @tullio c[i, j, k] := $a[i, j, q] * $b[i, q, k];
  682.080 ms (2 allocations: 30.52 MiB)

julia> a = Array(reshape(Int32.(1:2*2000*400), 2000,400,2));

julia> b = Array(reshape(Int32.(1:2*2000*400), 400,2000,2));

julia> @btime @tullio c[j, k, i] := $a[j, q, i] * $b[q, k, i];
  126.235 ms (2 allocations: 30.52 MiB)

```

My code is hella fast now, thanks.

[Previous page](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656.md?page=1)

[Next page](https://discourse.julialang.org/t/speed-comparison-matrix-multiplication-in-julia/66656.md?page=3)