"""
使用PyTorch Profiler分析GPU性能
"""

import torch
import torch.nn as nn
from torch.profiler import profile, record_function, ProfilerActivity

class TransformerBlock(nn.Module):
    def __init__(self, hidden_size, num_heads):
        super().__init__()
        self.attention = nn.MultiheadAttention(hidden_size, num_heads)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.GELU(),
            nn.Linear(hidden_size * 4, hidden_size)
        )
        self.ln1 = nn.LayerNorm(hidden_size)
        self.ln2 = nn.LayerNorm(hidden_size)
    
    def forward(self, x):
        # 自注意力
        attn_out, _ = self.attention(x, x, x)
        x = self.ln1(x + attn_out)
        
        # MLP
        mlp_out = self.mlp(x)
        x = self.ln2(x + mlp_out)
        
        return x

def profile_transformer():
    """分析Transformer块的性能"""
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    model = TransformerBlock(768, 12).to(device)
    x = torch.randn(32, 128, 768).to(device)
    
    # 预热
    for _ in range(10):
        _ = model(x)
    
    # Profiling
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        profile_memory=True,
    ) as prof:
        with record_function("model_forward"):
            output = model(x)
    
    # 打印结果
    print(prof.key_averages().table(
        sort_by="cuda_time_total",
        row_limit=20
    ))
    
    # 导出Chrome trace
    prof.export_chrome_trace("transformer_trace.json")
    print("Trace saved to transformer_trace.json")

if __name__ == "__main__":
    profile_transformer()

"""
比较不同数据类型的性能
"""

import torch
import time

def benchmark_precision():
    """基准测试不同精度"""
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    sizes = [(512, 512), (1024, 1024), (2048, 2048)]
    dtypes = [torch.float32, torch.float16, torch.bfloat16]
    
    results = {}
    
    for size in sizes:
        print(f"\nMatrix size: {size}")
        print("-" * 50)
        
        M, K = size
        N = K
        
        for dtype in dtypes:
            A = torch.randn(M, K, dtype=dtype, device=device)
            B = torch.randn(K, N, dtype=dtype, device=device)
            
            # 预热
            for _ in range(10):
                _ = torch.matmul(A, B)
            
            torch.cuda.synchronize()
            start = time.time()
            
            iterations = 100
            for _ in range(iterations):
                C = torch.matmul(A, B)
            
            torch.cuda.synchronize()
            elapsed = (time.time() - start) / iterations * 1000
            
            # 计算GFLOPS
            flops = 2 * M * K * N
            gflops = flops / elapsed / 1e6 * 1000
            
            results[(size, dtype)] = {
                'time': elapsed,
                'gflops': gflops
            }
            
            print(f"{str(dtype):>20}: {elapsed:.3f} ms, {gflops:.1f} GFLOPS")
    
    return results

def check_tensor_core_support():
    """检查Tensor Core支持"""
    
    if not torch.cuda.is_available():
        print("CUDA not available")
        return
    
    device_props = torch.cuda.get_device_properties(0)
    print(f"GPU: {device_props['name']}")
    print(f"Compute Capability: {device_props['major']}.{device_props['minor']}")
    
    # Tensor Core要求
    if device_props['major'] >= 7:
        print("Tensor Core supported!")
        
        # 检查合适的维度
        print("\nFor optimal Tensor Core performance:")
        print("- Matrix dimensions should be multiples of 8 (FP16) or 16 (INT8)")
        print("- Use NHWC layout for convolutions")
    else:
        print("Tensor Core not available on this GPU")

if __name__ == "__main__":
    check_tensor_core_support()
    benchmark_precision()

"""
根据硬件特性优化模型
"""

import torch
import torch.nn as nn

class HardwareAwareModel(nn.Module):
    """
    根据硬件特性自动调整的模型
    """
    
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        
        # 检测硬件
        self.device_type = self._detect_hardware()
        
        # 根据硬件选择配置
        if self.device_type == 'gpu_with_tensor_core':
            # 确保维度是8的倍数以利用Tensor Core
            hidden_size = ((hidden_size + 7) // 8) * 8
            print("Optimizing for GPU with Tensor Core")
        elif self.device_type == 'tpu':
            # TPU偏好固定形状和bfloat16
            print("Optimizing for TPU")
        else:
            print("Using default configuration")
        
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()
    
    def _detect_hardware(self):
        """检测硬件类型"""
        if torch.cuda.is_available():
            props = torch.cuda.get_device_properties(0)
            if props['major'] >= 7:
                return 'gpu_with_tensor_core'
            return 'gpu'
        # 可以添加其他硬件检测
        return 'cpu'
    
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x


def optimize_for_hardware(model, device):
    """
    根据硬件进行运行时优化
    """
    
    if isinstance(device, torch.device) and device.type == 'cuda':
        # GPU优化
        
        # 启用TF32（Ampere及以后）
        if torch.cuda.get_device_capability()[0] >= 8:
            torch.backends.cuda.matmul.allow_tf32 = True
            print("Enabled TF32 for matmul")
        
        # 启用cuDNN benchmark
        torch.backends.cudnn.benchmark = True
        print("Enabled cuDNN benchmark")
    
    elif hasattr(torch, 'xla') and 'xla' in str(device):
        # TPU优化
        import torch_xla.core.xla_model as xm
        print("Optimizing for XLA/TPU")
    
    return model


def adaptive_batch_size(device, model_size_mb=100):
    """
    根据可用显存自适应选择批次大小
    """
    
    if torch.cuda.is_available():
        free_mem = torch.cuda.mem_get_info()[0]  # bytes
        free_mem_mb = free_mem / 1024 / 1024
        
        # 估算每个样本需要的显存
        sample_mem_mb = model_size_mb * 2 / 1024  # 粗略估计
        
        # 保留一些余量
        available_mb = free_mem_mb * 0.7
        
        batch_size = int(available_mb / max(sample_mem_mb, 1))
        batch_size = max(1, min(batch_size, 1024))  # 限制范围
        
        print(f"Free memory: {free_mem_mb:.0f} MB")
        print(f"Recommended batch size: {batch_size}")
        
        return batch_size
    
    return 32  # 默认值


if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    model = HardwareAwareModel(784, 256, 10)
    model = optimize_for_hardware(model, device)
    model.to(device)
    
    batch_size = adaptive_batch_size(device)
    x = torch.randn(batch_size, 784, device=device)
    output = model(x)
    
    print(f"Output shape: {output.shape}")