What Shapes Do Matrix Multiplications Like?

Q1.py

import torch
from triton.testing import do_bench
torch.set_default_device('cuda')

for M, K, N in [(2047, 2048, 2048), (2048, 2047, 2048), (2048, 2048, 2047)]:
    A = torch.randn(M, K, dtype=torch.bfloat16)
    B = torch.randn(K, N, dtype=torch.bfloat16)
    print(f"M={M}, K={K}, N={N}")
    print(do_bench(lambda: torch.mm(A, B)))

Q2.py

import torch
from triton.testing import do_bench
torch.set_default_device('cuda')

M=108 * 256
N=3486
K=4679
A = torch.zeros(M, K, dtype=torch.bfloat16)
B = torch.zeros(N, K, dtype=torch.bfloat16).t()
print(f"M={M}, K={K}, N={N}")
torch.mm(A, B)

# ncu --metrics launch__waves_per_multiprocessor python Q2.py 

Q4.py

import torch
from triton.testing import do_bench
torch.set_default_device('cuda')

for M, K, N in [(2047, 2048, 2048), (2048, 2047, 2048), (2048, 2048, 2047)]:
    A = torch.randn(K, M, dtype=torch.bfloat16).t()
    B = torch.randn(K, N, dtype=torch.bfloat16)
    print(f"M={M}, K={K}, N={N}")
    print(do_bench(lambda: torch.mm(A, B)))

Companion to https://www.thonking.ai/p/what-shapes-do-matrix-multiplications