I have two n-by-3 blocks contiguous in memory ("n vectors of length 3") and I'd like to compute the dot product between each of the rows as fast as possible. In numpy, using einsum is the fastest variant to the best of my knowledge (einsum("ij,ij->i", a, b)).

I've compared this to the transposed variant (also contiguous in memory), i.e., two 3-by-n blocks. Here, the each of the components (x, y, z) form a big block in memory rather than n contiguous 3-blocks. I found that computation here can be about twice as fast:

enter image description here

I'm out of memory at $2^{25}$ vectors so I'd be interested in results beyond that point.

Any idea on why the transposed variant is so much faster at that size range?

Code to reproduce the plot:

import numpy
import perfplot

def setup(n):
    a = numpy.random.rand(n, 3)
    b = numpy.random.rand(n, 3)
    aT = numpy.ascontiguousarray(a.T)
    bT = numpy.ascontiguousarray(b.T)
    return (a, b), (aT, bT)

    n_range=[2 ** k for k in range(1, 25)],
        lambda data: numpy.einsum("ij, ij->i", data[0][0], data[0][1]),
        lambda data: numpy.einsum("ij, ij->j", data[1][0], data[1][1]),
    labels=["einsum", "einsum.T"],
    xlabel="len(a), len(b)",
    flops=lambda n: n,
  • 3
    $\begingroup$ yesterday spent some time experimenting and went under the hood of the einsum implementation inside numpy on github. Man, that function is complicated. $\endgroup$
    – Anton Menshov
    Nov 27, 2019 at 15:59
  • $\begingroup$ Yes, einsum is a monster. If one can reproduce this with plain C, that'd also be great. $\endgroup$ Nov 27, 2019 at 16:01
  • $\begingroup$ BTW, there's a GPU-specific library for this kind of operation -- docs.nvidia.com/cuda/cutensor/api/… $\endgroup$ Nov 2, 2020 at 18:53


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.