import os

import torch
import torch.utils.benchmark as benchmark

from itertools import product

results = []

batches = [(), (16,), (32,)]
sizes = [16, 128, 512,  2048]
dtypes = [torch.float32, torch.float64, torch.complex64, torch.complex128]

ENVVAR = "TORCH_LINALG_EIGH_BACKEND"
BACKENDS = {
    1: "syevd",
    2: "syevj",
    3: "syevj_batched",
}


for b, n, dtype in product(batches, sizes, dtypes):
    shape = b + (n, n)
    print(f"Testing shape={shape}, dtype={dtype}")
    label = "torch.linalg.eigh"
    sub_label = f"{shape}, {dtype}"
    X = torch.rand(*shape, dtype=dtype, device="cuda")
    X = X + X.mT.conj()
    stmt = "torch.linalg.eigh(X)"
    for mode, name in BACKENDS.items():
        os.environ[ENVVAR] = str(mode)
        # warm-up
        for _ in range(5):
            exec(stmt)

        results.append(benchmark.Timer(
            stmt=stmt,
            globals={'X': X},
            label=label,
            sub_label=sub_label,
            description=name,
        ).blocked_autorange(min_run_time=1))

compare = benchmark.Compare(results)
compare.print()

[--------------------------------- torch.linalg.eigh ----------------------------------]
                                          |    syevd    |     syevj     |  syevj_batched
1 threads: -----------------------------------------------------------------------------
      (16, 16), torch.float32             |      130.4  |        261.6  |        128.8  
      (16, 16), torch.float64             |      361.6  |        872.8  |        361.0  
      (16, 16), torch.complex64           |      131.6  |        316.7  |        130.6  
      (16, 16), torch.complex128          |      482.9  |       2787.3  |        473.5  
      (128, 128), torch.float32           |     1373.5  |       1495.0  |       1370.6  
      (128, 128), torch.float64           |     4038.9  |      12203.9  |       4032.1  
      (128, 128), torch.complex64         |     1439.5  |       2010.9  |       1439.0  
      (128, 128), torch.complex128        |     4889.0  |      34441.1  |       4898.3  
      (512, 512), torch.float32           |     3375.4  |      10695.3  |       3374.4  
      (512, 512), torch.float64           |    11973.0  |     115825.4  |      11764.8  
      (512, 512), torch.complex64         |     3766.1  |      17226.0  |       3759.7  
      (512, 512), torch.complex128        |    20258.5  |     373294.1  |      20233.5  
      (2048, 2048), torch.float32         |    26396.7  |     141881.3  |      26466.9  
      (2048, 2048), torch.float64         |   116787.9  |    1964422.3  |     116736.7  
      (2048, 2048), torch.complex64       |    36661.3  |     317082.7  |      36678.7  
      (2048, 2048), torch.complex128      |   310149.7  |    7228237.5  |     309390.3  
      (16, 16, 16), torch.float32         |     1989.5  |       3983.7  |        147.1  
      (16, 16, 16), torch.float64         |     5907.9  |      15396.8  |        380.9  
      (16, 16, 16), torch.complex64       |     2044.8  |       4902.2  |        147.3  
      (16, 16, 16), torch.complex128      |     8041.2  |      48572.2  |        442.0  
      (16, 128, 128), torch.float32       |    22692.3  |      25441.5  |        969.3  
      (16, 128, 128), torch.float64       |    70677.4  |     194437.7  |       4669.9  
      (16, 128, 128), torch.complex64     |    26051.2  |      34769.5  |       1096.1  
      (16, 128, 128), torch.complex128    |    78945.3  |     544737.4  |       9740.3  
      (16, 512, 512), torch.float32       |    54176.4  |     199518.5  |      20227.1  
      (16, 512, 512), torch.float64       |   188504.8  |    1867793.8  |      62640.4  
      (16, 512, 512), torch.complex64     |    60293.3  |     274358.2  |      26721.6  
      (16, 512, 512), torch.complex128    |   324978.8  |    6001170.3  |     131001.7  
      (16, 2048, 2048), torch.float32     |   422230.6  |    2199166.9  |     352754.5  
      (16, 2048, 2048), torch.float64     |  1872350.6  |   31005349.9  |    1679366.3  
      (16, 2048, 2048), torch.complex64   |   591507.1  |    5106029.0  |     555647.0  
      (16, 2048, 2048), torch.complex128  |  4951373.7  |  114891808.9  |    4610288.0  
      (32, 16, 16), torch.float32         |     4175.4  |       8940.8  |        168.7  
      (32, 16, 16), torch.float64         |    12171.5  |      32038.5  |        423.1  
      (32, 16, 16), torch.complex64       |     4290.7  |      10810.8  |        167.1  
      (32, 16, 16), torch.complex128      |    16349.0  |     106911.7  |        515.3  
      (32, 128, 128), torch.float32       |    45842.8  |      53155.2  |       1278.7  
      (32, 128, 128), torch.float64       |   143017.6  |     392286.8  |       6659.3  
      (32, 128, 128), torch.complex64     |    52562.7  |      70499.9  |       1510.7  
      (32, 128, 128), torch.complex128    |   161991.5  |    1104503.2  |      13099.4  
      (32, 512, 512), torch.float32       |   107676.2  |     403025.7  |      30682.5  
      (32, 512, 512), torch.float64       |   378800.5  |    3755004.9  |      96316.8  
      (32, 512, 512), torch.complex64     |   121609.9  |     617220.1  |      40746.4  
      (32, 512, 512), torch.complex128    |   655783.5  |   12012476.8  |     233843.4  
      (32, 2048, 2048), torch.float32     |   855990.9  |    4453487.0  |     658549.1  
      (32, 2048, 2048), torch.float64     |  3744886.3  |   62298724.0  |    3243343.8  
      (32, 2048, 2048), torch.complex64   |  1189833.9  |   10124234.7  |    1062783.0  
      (32, 2048, 2048), torch.complex128  |  9905053.1  |  230245338.3  |    9101891.5  

Times are in microseconds (us).