Torch.einsum is around ~4x faster than broadcasting torch.matmul for my use case

[KFrank]

Hi Santosh!

@Santosh_Gupta

torch.matmul(phase2.unsqueeze(2), weight.unsqueeze(2)).squeeze(2).squeeze(2)

and

torch.einsum('abc, bc -> ab', phase2, weight)

I don’t know why einsum() is faster than matmul(), but we have seen things
like this before, for example, as in this post:

Why are batched operations slower than a for loop?

Hi Howard! Good question – I don’t know. I can reproduce your result. including on a recent nightly build. It’s interesting to note that a non-loop einsum() version is competitive with your “non-batched” loop version, and, at least on my system, turns out to be somewhat faster. (I see slower, but similar relative timings using just the cpu, with einsum() again being the fastest version.) Running a tweaked version of your test code, with einsum() added as “Batch3,” I get: 1.10.0 10.2 G…

(We’ve also seen cases where einsum() is unexpectedly and unreasonably
slow.)

What could be causing the 4x speedup in einsum?

As an aside, I might guess that it would be better to describe this as an
(unexpected) slowdown in matmul(), rather than as a speedup in einsum().
Have you considered comparing the matmul() timings with a loop version?

Best.

K. Frank