PyTorch 2.12 稀疏矩阵深度解析：COO、CSR、CSC、BSR 与 BSC

from __future__ import annotations

import argparse
import platform
import time
from dataclasses import dataclass
from typing import Callable

import torch


@dataclass
class Case:
    name: str
    dense: torch.Tensor
    rhs: torch.Tensor


@dataclass
class Result:
    case: str
    layout: str
    nnz: int
    storage_bytes: int | None
    storage_ratio: float | None
    time_ms: float | None
    time_ratio: float | None
    note: str


def tensor_bytes(tensor: torch.Tensor) -> int:
    return tensor.untyped_storage().nbytes()


def sparse_storage_bytes(tensor: torch.Tensor) -> int:
    layout = tensor.layout
    if layout == torch.sparse_coo:
        return (
            tensor.indices().untyped_storage().nbytes()
            + tensor.values().untyped_storage().nbytes()
        )
    if layout in {torch.sparse_csr, torch.sparse_bsr}:
        return (
            tensor.crow_indices().untyped_storage().nbytes()
            + tensor.col_indices().untyped_storage().nbytes()
            + tensor.values().untyped_storage().nbytes()
        )
    if layout in {torch.sparse_csc, torch.sparse_bsc}:
        return (
            tensor.ccol_indices().untyped_storage().nbytes()
            + tensor.row_indices().untyped_storage().nbytes()
            + tensor.values().untyped_storage().nbytes()
        )
    raise ValueError(f"unsupported layout for storage accounting: {layout}")


def cuda_median_ms(
    fn: Callable[[], torch.Tensor],
    *,
    warmup: int,
    repeats: int,
) -> float:
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()

    times: list[float] = []
    for _ in range(repeats):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        fn()
        end.record()
        torch.cuda.synchronize()
        times.append(start.elapsed_time(end))
    return sorted(times)[len(times) // 2]


def cpu_median_ms(fn: Callable[[], torch.Tensor], *, min_run_time: float) -> float:
    for _ in range(3):
        fn()
    start = time.perf_counter()
    iters = 0
    samples: list[float] = []
    while time.perf_counter() - start < min_run_time or iters < 5:
        t0 = time.perf_counter()
        fn()
        samples.append((time.perf_counter() - t0) * 1000.0)
        iters += 1
    return sorted(samples)[len(samples) // 2]


def median_ms(
    fn: Callable[[], torch.Tensor],
    *,
    device: torch.device,
    warmup: int,
    repeats: int,
    min_run_time: float,
) -> float:
    if device.type == "cuda":
        return cuda_median_ms(fn, warmup=warmup, repeats=repeats)
    return cpu_median_ms(fn, min_run_time=min_run_time)


def make_unstructured_case(
    name: str,
    rows: int,
    cols: int,
    rhs_cols: int,
    density: float,
    *,
    seed: int,
    dtype: torch.dtype,
    device: torch.device,
) -> Case:
    generator = torch.Generator(device=device).manual_seed(seed)
    nnz = max(1, int(rows * cols * density))
    flat = torch.randperm(rows * cols, generator=generator, device=device)[:nnz]
    row = torch.div(flat, cols, rounding_mode="floor")
    col = flat.remainder(cols)
    values = torch.randn(nnz, generator=generator, dtype=dtype, device=device)
    dense = torch.zeros((rows, cols), dtype=dtype, device=device)
    dense[row, col] = values
    rhs = torch.randn((cols, rhs_cols), generator=generator, dtype=dtype, device=device)
    return Case(name=name, dense=dense, rhs=rhs)


def make_block_case(
    name: str,
    rows: int,
    cols: int,
    rhs_cols: int,
    block_size: tuple[int, int],
    block_density: float,
    *,
    seed: int,
    dtype: torch.dtype,
    device: torch.device,
) -> Case:
    brow, bcol = block_size
    if rows % brow or cols % bcol:
        raise ValueError("rows and cols must be divisible by block size")
    generator = torch.Generator(device=device).manual_seed(seed)
    block_rows = rows // brow
    block_cols = cols // bcol
    n_blocks = max(1, int(block_rows * block_cols * block_density))
    chosen = torch.randperm(block_rows * block_cols, generator=generator, device=device)[:n_blocks]
    dense = torch.zeros((rows, cols), dtype=dtype, device=device)
    for block in chosen.tolist():
        br = block // block_cols
        bc = block % block_cols
        dense[br * brow : (br + 1) * brow, bc * bcol : (bc + 1) * bcol] = torch.randn(
            (brow, bcol), generator=generator, dtype=dtype, device=device
        )
    rhs = torch.randn((cols, rhs_cols), generator=generator, dtype=dtype, device=device)
    return Case(name=name, dense=dense, rhs=rhs)


def convert_layouts(dense: torch.Tensor, block_size: tuple[int, int]) -> dict[str, torch.Tensor]:
    coo = dense.to_sparse_coo().coalesce()
    return {
        "dense": dense,
        "coo": coo,
        "csr": dense.to_sparse_csr(),
        "csc": dense.to_sparse_csc(),
        "bsr": dense.to_sparse_bsr(blocksize=block_size),
        "bsc": dense.to_sparse_bsc(blocksize=block_size),
    }


def run_case(
    case: Case,
    block_size: tuple[int, int],
    *,
    device: torch.device,
    warmup: int,
    repeats: int,
    min_run_time: float,
) -> list[Result]:
    layouts = convert_layouts(case.dense, block_size)
    dense_bytes = tensor_bytes(case.dense)
    dense_time_ms = median_ms(
        lambda: torch.mm(case.dense, case.rhs),
        device=device,
        warmup=warmup,
        repeats=repeats,
        min_run_time=min_run_time,
    )
    results = [
        Result(
            case=case.name,
            layout="dense",
            nnz=int(torch.count_nonzero(case.dense).item()),
            storage_bytes=dense_bytes,
            storage_ratio=1.0,
            time_ms=dense_time_ms,
            time_ratio=1.0,
            note="baseline",
        )
    ]
    for layout_name in ["coo", "csr", "csc", "bsr", "bsc"]:
        sparse = layouts[layout_name]
        storage_bytes = sparse_storage_bytes(sparse)
        try:
            torch.sparse.mm(sparse, case.rhs)
            if device.type == "cuda":
                torch.cuda.synchronize()
            time_ms = median_ms(
                lambda sparse=sparse: torch.sparse.mm(sparse, case.rhs),
                device=device,
                warmup=warmup,
                repeats=repeats,
                min_run_time=min_run_time,
            )
            note = "ok"
        except Exception as exc:
            time_ms = None
            note = f"unsupported: {type(exc).__name__}: {str(exc).splitlines()[0]}"
        results.append(
            Result(
                case=case.name,
                layout=layout_name,
                nnz=int(sparse._nnz()),
                storage_bytes=storage_bytes,
                storage_ratio=storage_bytes / dense_bytes,
                time_ms=time_ms,
                time_ratio=None if time_ms is None else time_ms / dense_time_ms,
                note=note,
            )
        )
    return results


def fmt_bytes(value: int | None) -> str:
    if value is None:
        return "-"
    units = ["B", "KiB", "MiB", "GiB"]
    size = float(value)
    for unit in units:
        if size < 1024 or unit == units[-1]:
            return f"{size:.1f} {unit}" if unit != "B" else f"{int(size)} B"
        size /= 1024
    return f"{value} B"


def fmt_ratio(value: float | None) -> str:
    return "-" if value is None else f"{value:.3f}x"


def fmt_ms(value: float | None) -> str:
    return "-" if value is None else f"{value:.3f}"


def print_markdown(results: list[Result]) -> None:
    print(
        "| case | layout | nnz / blocks | storage | storage vs dense | "
        "median ms | time vs dense | note |"
    )
    print("| --- | ---: | ---: | ---: | ---: | ---: | ---: | --- |")
    for row in results:
        print(
            f"| {row.case} | {row.layout} | {row.nnz} | {fmt_bytes(row.storage_bytes)} | "
            f"{fmt_ratio(row.storage_ratio)} | {fmt_ms(row.time_ms)} | "
            f"{fmt_ratio(row.time_ratio)} | {row.note} |"
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
    parser.add_argument("--rows", type=int, default=2048)
    parser.add_argument("--cols", type=int, default=2048)
    parser.add_argument("--rhs-cols", type=int, default=64)
    parser.add_argument("--block-size", type=int, default=16)
    parser.add_argument("--warmup", type=int, default=10)
    parser.add_argument("--repeats", type=int, default=50)
    parser.add_argument("--min-run-time", type=float, default=0.2)
    parser.add_argument("--require-torch-prefix", default="2.12.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.require_torch_prefix and not torch.__version__.startswith(args.require_torch_prefix):
        raise SystemExit(
            f"Expected torch {args.require_torch_prefix}*, got torch {torch.__version__}"
        )
    if args.device == "cuda" and not torch.cuda.is_available():
        raise SystemExit("CUDA device requested, but torch.cuda.is_available() is false")

    device = torch.device(args.device)
    dtype = torch.float32
    block_size = (args.block_size, args.block_size)
    cases = [
        make_unstructured_case("random 0.1%", args.rows, args.cols, args.rhs_cols, 0.001, seed=11, dtype=dtype, device=device),
        make_unstructured_case("random 1%", args.rows, args.cols, args.rhs_cols, 0.01, seed=12, dtype=dtype, device=device),
        make_unstructured_case("random 5%", args.rows, args.cols, args.rhs_cols, 0.05, seed=13, dtype=dtype, device=device),
        make_block_case("16x16 blocks 1%", args.rows, args.cols, args.rhs_cols, block_size, 0.01, seed=21, dtype=dtype, device=device),
    ]

    print(f"torch: {torch.__version__}")
    print(f"python: {platform.python_version()} ({platform.machine()})")
    print(f"device: {device}")
    if device.type == "cuda":
        props = torch.cuda.get_device_properties(device)
        print(f"cuda: {torch.version.cuda}")
        print(f"gpu: {props.name}")
        print(f"compute capability: {props.major}.{props.minor}")
        print(f"total memory: {fmt_bytes(props.total_memory)}")
    else:
        print(f"threads: {torch.get_num_threads()}")
    print(f"shape: {args.rows}x{args.cols}, rhs: {args.cols}x{args.rhs_cols}, dtype: float32")
    print(f"block size: {block_size[0]}x{block_size[1]}")
    print()

    all_results: list[Result] = []
    for case in cases:
        all_results.extend(
            run_case(
                case,
                block_size,
                device=device,
                warmup=args.warmup,
                repeats=args.repeats,
                min_run_time=args.min_run_time,
            )
        )
    print_markdown(all_results)


if __name__ == "__main__":
    main()