Repository: pytorch
Issue: einsum "jk,ijkl->il" is ~16x slower than numpy
Labels: module: performance,triaged,module: numpy,module: linear algebra

## 🐛 Bug

The high rank tensor contraction with einsum `"jk,ijkl->il"` is about 16x slower than numpy.

## To Reproduce

    import time
    import torch
    import numpy as np

    nb = 160
    dtype = torch.float64

    mat1 = torch.rand((nb, nb, nb, nb), dtype=dtype)
    mat2 = torch.rand((nb, nb), dtype=dtype)

    t0 = time.time()
    res1 = torch.einsum("ij,ijkl->kl", mat2, mat1)
    t1 = time.time()
    res2 = torch.einsum("kl,ijkl->ij", mat2, mat1)
    t2 = time.time()
    res3 = torch.einsum("jk,ijkl->il", mat2, mat1)
    t3 = time.time()
    res4 = torch.einsum("il,ijkl->jk", mat2, mat1)
    t4 = time.time()

    print("PyTorch")
    print("1:", t1 - t0)
    print("2:", t2 - t1)
    print("3:", t3 - t2)
    print("4:", t4 - t3)

    mat1 = np.asarray(mat1)
    mat2 = np.asarray(mat2)

    t0 = time.time()
    res1 = np.einsum("ij,ijkl->kl", mat2, mat1, optimize=True)
    t1 = time.time()
    res2 = np.einsum("kl,ijkl->ij", mat2, mat1, optimize=True)
    t2 = time.time()
    res3 = np.einsum("jk,ijkl->il", mat2, mat1, optimize=True)
    t3 = time.time()
    res4 = np.einsum("il,ijkl->jk", mat2, mat1, optimize=True)
    t4 = time.time()

    print("Numpy")
    print("1:", t1 - t0)
    print("2:", t2 - t1)
    print("3:", t3 - t2)
    print("4:", t4 - t3)

Prints:

    PyTorch
    1: 0.23056674003601074
    2: 0.23742151260375977
    3: 7.9162068367004395
    4: 8.927417755126953
    Numpy
    1: 0.27585268020629883
    2: 0.23679137229919434
    3: 0.5738704204559326
    4: 0.48691320419311523

## Expected behavior

Case 3 and 4 for PyTorch are about 40x slower than case 1 & 2.
On the other hand, case 3 & 4 for numpy are only about 2x slower (which is acceptable).
For PyTorch's case 3, I also tried several alternatives and all of them are as slow as the original case 3:

    # these cases are as slow as the original case 3
    res3 = torch.einsum("jk,iljk->il", mat2, mat1.movedim(-1, 1))
    res3 = (mat2.unsqueeze(-1) * mat1).reshape(mat1.shape[0], -1, mat1.shape[-1]).sum(dim=-2)
    res3 = torch.tensordot(mat2, mat1, dims=([0, 1], [1, 2]))

Seeing those cases are as slow as the original einsum, I'm not sure if this is the problem with `einsum`, or with matrix multiplication, or with something else.

## Environment

 - PyTorch Version (e.g., 1.0): 1.9.0.dev20210419
 - OS (e.g., Linux): WSL Ubuntu 20.04
 - How you installed PyTorch (`conda`, `pip`, source): conda
 - Build command you used (if compiling from source): N/A
 - Python version: 3.8.5
 - CUDA/cuDNN version: N/A
 - GPU models and configuration: N/A
 - Any other relevant information: N/A

## Additional context

EDIT: I have just discovered a faster form for case 4:

    res4 = torch.einsum("il,ijkl->jk", mat2, mat1)  # original case 4: slower, about 8s
    res4 = torch.einsum("il,ijkl->ijk", mat2, mat1).sum(dim=0)  # hack case 4: faster, about 0.2s

However, case 3 is still pretty slow...

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @jianyuh @nikitaved @pearu @mruberry @heitorschueroff @walterddr @IvanYashchuk @xwang233 @Lezcano @rgommers @VitalyFedyunin @ngimel