Benchmark for https://github.com/pytorch/pytorch/pull/45358

kron_benchmark.py

import numpy, torch, cupy
import itertools
import time
from torch.utils._benchmark import Timer
from torch.utils._benchmark import Compare
import sys

print('Using pytorch %s' % (torch.__version__))
print('Using cupy %s' % (cupy.__version__))
print('Using numpy %s' % (numpy.__version__))

numpy.random.seed(1234)

shapes = [(32,), (32, 32), (2, 16, 32), (2, 16, 32, 32)]
results = []
num_threads = 8
dtype = torch.float32
repeats = 2

for mat1_shape, mat2_shape in itertools.product(shapes, reversed(shapes)):
    mat1 = torch.rand(*mat1_shape, dtype=dtype, device='cuda')
    mat2 = torch.rand(*mat2_shape, dtype=dtype, device='cuda')

    mat1_cpu = mat1.to('cpu')
    mat2_cpu = mat2.to('cpu')

    mat1_numpu = mat1_cpu.numpy()
    mat2_numpy = mat2_cpu.numpy()

    mat1_cupy = cupy.array(mat1_cpu)
    mat2_cupy = cupy.array(mat2_cpu)

    tasks = [("torch.kron(mat1, mat2)", "torch.kron CUDA"), ("cupy.kron(mat1_cupy, mat2_cupy)", "cupy.kron CUDA"), ("torch.kron(mat1_cpu, mat2_cpu)", "torch.kron CPU"), ("numpy.kron(mat1_cpu, mat2_cpu)", "numpy.kron CPU")]
    timers = [Timer(stmt=stmt, num_threads=num_threads, label=f"kron {dtype}", sub_label=f"{(mat1_shape, mat2_shape)}", description=label, globals=globals()) for stmt, label in tasks]

    for i, timer in enumerate(timers * repeats):
        results.append(
            timer.blocked_autorange()
        )
        print(f"\r{i + 1} / {len(timers) * repeats}", end="")
        sys.stdout.flush()
comparison = Compare(results)
comparison.print()