Compute high-dimensional matrix multiplication using the matmul operator

peiyingxin: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/6/7615e09396250859f99027df53e422151b086de2.png" data-download-href="https://discuss.pytorch.org/uploads/default/7615e09396250859f99027df53e422151b086de2" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/6/7615e09396250859f99027df53e422151b086de2_2_690x404.png" alt="image" data-base62-sha1="gQDddXyxVN7pKJCdk8oZMz1LPuq" width="690" height="404" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/6/7615e09396250859f99027df53e422151b086de2_2_690x404.png, https://discuss.pytorch.org/uploads/default/original/3X/7/6/7615e09396250859f99027df53e422151b086de2.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/6/7615e09396250859f99027df53e422151b086de2.png 2x" data-dominant-color="EEEFF0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">949×556 60.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/0/b086a3584f9cd9c7b499b188244bd5caf8b5965d.png" data-download-href="https://discuss.pytorch.org/uploads/default/b086a3584f9cd9c7b499b188244bd5caf8b5965d" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/0/b086a3584f9cd9c7b499b188244bd5caf8b5965d.png" alt="image" data-base62-sha1="pbCuINDsJ7ljhUObhjxZIEeqXg1" width="690" height="319" data-dominant-color="F5F5F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">837×388 16.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>compute high-dimensional matrix multiplication very slow, why?</p>

vince62s: <p>not sure why but with those dimensions:</p>
<pre><code class="lang-auto">import torch
#torch.backends.cudnn.benchmark = True
repeat = 1000

import time
with torch.no_grad():
    a = torch.rand((1, 16, 4096, 128), device=torch.device("cuda"), dtype=torch.float32)
    b = torch.rand((1, 16, 128, 4096), device=torch.device("cuda"), dtype=torch.float32)
    torch.cuda.synchronize()
    start_time = time.time()
    for i in range(repeat):
        c = torch.matmul(a, b)
    torch.cuda.synchronize()
    print(f"matmul time is: {time.time() - start_time}")
    print(c.shape)
    
    torch.cuda.synchronize()
    start_time = time.time()
    for i in range(repeat):
        d = torch.einsum('abij, abjl -&gt; abil', a, b)
    torch.cuda.synchronize()
    print(f"einsum time is: {time.time() - start_time}")
    print(c.shape)
</code></pre>
<p>matmul time is: 1.6723120212554932<br>
torch.Size([1, 16, 4096, 4096])<br>
einsum time is: 1.6882061958312988<br>
torch.Size([1, 16, 4096, 4096])</p>