#!/usr/bin/env python3
"""
ROCm MIOpen Cache Warmer for MinerU

在 AMD RDNA 架构上，MIOpen 遇到新尺寸的卷积运算时需要搜索最优 kernel（冷启动）。
预热脚本提前跑一遍常用尺寸，将 kernel 缓存到 ~/.cache/miopen/，避免运行时等待。

缓存持久化到磁盘，重启不丢失；只有升级 ROCm 后才需要重新跑。

用法:
    python cache_warmer.py --device cuda --max_side 960 --step 32
"""

import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm


def get_args():
    p = argparse.ArgumentParser(description="ROCm MIOpen Cache Warmer")
    p.add_argument("--device", type=str, default="cuda")
    p.add_argument("--max_side", type=int, default=960)
    p.add_argument("--step", type=int, default=32)
    return p.parse_args()


class MockOCRModel(nn.Module):
    """模拟 MinerU OCR 模型的卷积结构，覆盖 MIOpen 常用 kernel 尺寸。"""

    def __init__(self, in_channels: int = 3):
        super().__init__()
        self.stem = nn.Conv2d(in_channels, 16, 3, stride=2, padding=1)
        self.dw_3x3 = nn.Conv2d(16, 16, 3, stride=1, padding=1, groups=16)
        self.pw_1 = nn.Conv2d(16, 64, 1)
        self.dw_5x5 = nn.Conv2d(64, 64, 5, stride=2, padding=2, groups=64)
        self.pw_2 = nn.Conv2d(64, 128, 1)
        self.dw_3x3_s2 = nn.Conv2d(128, 128, 3, stride=2, padding=1, groups=128)
        self.pw_3 = nn.Conv2d(128, 256, 1)
        self.out_conv = nn.Conv2d(256, 64, 1)
        self.binarize_conv = nn.Conv2d(64, 1, 3, stride=1, padding=1)
        self.act = nn.ReLU()

    def forward(self, x):
        x = self.stem(x)
        x = self.act(x)
        x = self.dw_3x3(x)
        x = self.pw_1(x)
        x = self.dw_5x5(x)
        x = self.act(x)
        x = self.pw_2(x)
        x = self.dw_3x3_s2(x)
        x = self.pw_3(x)
        x = self.out_conv(x)
        x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
        x = self.binarize_conv(x)
        return x


def main():
    args = get_args()
    assert torch.cuda.is_available(), "GPU not available"
    device = torch.device(args.device)

    print("=" * 50)
    print("ROCm MIOpen Cache Warmer")
    print(f"  GPU    : {torch.cuda.get_device_name(0)}")
    print(f"  ROCm   : {torch.version.hip}")
    print(f"  Cache  : ~/.cache/miopen/")
    print("=" * 50)

    model = MockOCRModel().to(device).eval()
    sizes = list(range(64, args.max_side + 1, args.step))
    combos = [(h, w) for h in sizes for w in sizes]

    print(f"Warming {len(combos)} shapes ({len(sizes)}x{len(sizes)} grid, "
          f"step={args.step})...")

    ok = 0
    with torch.no_grad():
        for h, w in tqdm(combos, desc="Warming"):
            try:
                model(torch.zeros((1, 3, h, w), device=device, dtype=torch.float32))
                ok += 1
            except RuntimeError as e:
                if "out of memory" in str(e):
                    torch.cuda.empty_cache()
                # 其他错误跳过，不影响后续

    print(f"\nDone! {ok}/{len(combos)} shapes cached (~3–4 min)")
    print("Kernels saved to ~/.cache/miopen/")


if __name__ == "__main__":
    main()