Function benchmark_baselines

prepare.py:276–321

Benchmark torch.matmul at all sizes and dtypes. Returns results dict.

()

# ---------------------------------------------------------------------------

def benchmark_baselines() -> dict:
    """Benchmark torch.matmul at all sizes and dtypes. Returns results dict."""

    print("Benchmarking PyTorch baselines...")
    results = {}

    for size_name, dims in MATMUL_SIZES:
        M, N, K = dims["M"], dims["N"], dims["K"]
        flops = _matmul_flops(M, N, K)

        for dtype in TEST_DTYPES:
            tag = _dtype_tag(dtype)

            # Load cached test data if available, else generate on the fly
            save_path = os.path.join(TEST_DATA_DIR, "matmul", size_name, f"{tag}.pt")
            if os.path.exists(save_path):
                data = torch.load(save_path, weights_only=True)
                A = data["A"].cuda()
                B = data["B"].cuda()
            else:
                gen = torch.Generator(device="cpu")
                gen.manual_seed(_SEED)
                A = torch.randn(M, K, generator=gen, dtype=dtype).cuda()
                B = torch.randn(K, N, generator=gen, dtype=dtype).cuda()

            latency_us = _benchmark_fn(torch.matmul, A, B)
            tflops = flops / (latency_us * 1e-6) / 1e12

            key = f"matmul_{size_name}_{tag}"
            results[key] = {
                "kernel_type": "matmul",
                "size": size_name,
                "dtype": tag,
                "M": M, "N": N, "K": K,
                "latency_us": round(latency_us, 2),
                "throughput_tflops": round(tflops, 3),
            }

            print(f"  matmul {size_name} {tag}: {tflops:.1f} TFLOPS ({latency_us:.2f} us)")

            # Free GPU memory
            del A, B
            torch.cuda.empty_cache()

    print()
    return results


# ---------------------------------------------------------------------------