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 '19 at 15:59
  • $\begingroup$ Yes, einsum is a monster. If one can reproduce this with plain C, that'd also be great. $\endgroup$ – Nico Schlömer Nov 27 '19 at 16:01
  • $\begingroup$ BTW, there's a GPU-specific library for this kind of operation -- docs.nvidia.com/cuda/cutensor/api/… $\endgroup$ – Yaroslav Bulatov Nov 2 '20 at 18:53

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.