#!/bin/bash
# GPU benchmark - test various operations
export HSA_OVERRIDE_GFX_VERSION=10.1.0
export HSA_ENABLE_SDMA=0
export HIP_VISIBLE_DEVICES=0

cd /home/fabian/ComfyUI
source venv/bin/activate

python3 -c "
import torch
import time

print(f'PyTorch: {torch.__version__}')
print(f'CUDA: {torch.cuda.is_available()}')
print(f'Device: {torch.cuda.get_device_name(0)}')
print()

# Test 1: Small matmul (should be fast)
print('=== Test 1: Small matmul 256x256 ===')
a = torch.randn(256, 256, device='cuda', dtype=torch.float32)
b = torch.randn(256, 256, device='cuda', dtype=torch.float32)
torch.cuda.synchronize()
t = time.time()
for _ in range(10):
    c = torch.mm(a, b)
torch.cuda.synchronize()
print(f'  10x matmul 256x256 fp32: {time.time()-t:.3f}s')

# Test 2: Large matmul (typical for model inference)
print('=== Test 2: Large matmul 2048x2048 ===')
a = torch.randn(2048, 2048, device='cuda', dtype=torch.float32)
b = torch.randn(2048, 2048, device='cuda', dtype=torch.float32)
torch.cuda.synchronize()
t = time.time()
c = torch.mm(a, b)
torch.cuda.synchronize()
print(f'  1x matmul 2048x2048 fp32: {time.time()-t:.3f}s')

# Test 3: fp16 matmul
print('=== Test 3: fp16 matmul 2048x2048 ===')
a = torch.randn(2048, 2048, device='cuda', dtype=torch.float16)
b = torch.randn(2048, 2048, device='cuda', dtype=torch.float16)
torch.cuda.synchronize()
t = time.time()
c = torch.mm(a, b)
torch.cuda.synchronize()
print(f'  1x matmul 2048x2048 fp16: {time.time()-t:.3f}s')

# Test 4: Conv2d (typical for VAE)
print('=== Test 4: Conv2d 64ch ===')
conv = torch.nn.Conv2d(64, 64, 3, padding=1).cuda().float()
x = torch.randn(1, 64, 64, 64, device='cuda', dtype=torch.float32)
torch.cuda.synchronize()
t = time.time()
y = conv(x)
torch.cuda.synchronize()
print(f'  Conv2d 64x64x64 fp32: {time.time()-t:.3f}s')

# Test 5: Attention-like operation
print('=== Test 5: Attention (batch matmul) ===')
q = torch.randn(1, 8, 1024, 64, device='cuda', dtype=torch.float32)
k = torch.randn(1, 8, 1024, 64, device='cuda', dtype=torch.float32)
torch.cuda.synchronize()
t = time.time()
attn = torch.matmul(q, k.transpose(-2, -1))
torch.cuda.synchronize()
print(f'  BMM 8x1024x64 fp32: {time.time()-t:.3f}s')

# Test 6: Larger attention (closer to model size)
print('=== Test 6: Large attention ===')
q = torch.randn(1, 16, 4096, 128, device='cuda', dtype=torch.float32)
k = torch.randn(1, 16, 4096, 128, device='cuda', dtype=torch.float32)
torch.cuda.synchronize()
t = time.time()
attn = torch.matmul(q, k.transpose(-2, -1))
torch.cuda.synchronize()
print(f'  BMM 16x4096x128 fp32: {time.time()-t:.3f}s')

# Test 7: bf16
print('=== Test 7: bf16 matmul ===')
try:
    a = torch.randn(2048, 2048, device='cuda', dtype=torch.bfloat16)
    b = torch.randn(2048, 2048, device='cuda', dtype=torch.bfloat16)
    torch.cuda.synchronize()
    t = time.time()
    c = torch.mm(a, b)
    torch.cuda.synchronize()
    print(f'  bf16 2048x2048: {time.time()-t:.3f}s')
except Exception as e:
    print(f'  bf16 FAILED: {e}')

print()
print('=== All tests done ===')
" 2>&1