#!/bin/bash
export HSA_OVERRIDE_GFX_VERSION=10.1.0
export HSA_ENABLE_SDMA=0
export HIP_VISIBLE_DEVICES=0
export TORCHDYNAMO_DISABLE=1
export PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.8,max_split_size_mb:128
export TORCH_BLAS_PREFER_HIPBLASLT=0

/home/fabian/ComfyUI/venv/bin/python -c "
import torch
import time
import gc

print('=== BC-250 GPU Transfer Stress Test ===')
print(f'Device: {torch.cuda.get_device_name(0)}')
print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB')

# Test 1: Many small .to(cuda) transfers
print()
print('Test 1: 100x small .to(cuda) [1024x1024 f16]')
t0 = time.time()
for i in range(100):
    cpu_t = torch.randn(1024, 1024, dtype=torch.float16)
    gpu_t = cpu_t.to('cuda:0', non_blocking=False)
    del gpu_t, cpu_t
    if (i+1) % 10 == 0:
        print(f'  {i+1}/100 done ({time.time()-t0:.1f}s)')
torch.cuda.synchronize()
print(f'  Total: {time.time()-t0:.1f}s')
gc.collect()
torch.cuda.empty_cache()

# Test 2: Larger tensors like attention weights
print()
print('Test 2: 50x medium .to(cuda) [3840x2560 f16]')
t0 = time.time()
for i in range(50):
    cpu_t = torch.randn(3840, 2560, dtype=torch.float16)
    gpu_t = cpu_t.to('cuda:0', non_blocking=False)
    del gpu_t, cpu_t
    if (i+1) % 10 == 0:
        print(f'  {i+1}/50 done ({time.time()-t0:.1f}s)')
torch.cuda.synchronize()
print(f'  Total: {time.time()-t0:.1f}s')
gc.collect()
torch.cuda.empty_cache()

# Test 3: .to(cuda) + matmul (actual compute)
print()
print('Test 3: 20x transfer + matmul [2560x2560 f16]')
t0 = time.time()
for i in range(20):
    a = torch.randn(2560, 2560, dtype=torch.float16).to('cuda:0', non_blocking=False)
    b = torch.randn(2560, 2560, dtype=torch.float16).to('cuda:0', non_blocking=False)
    c = torch.matmul(a, b)
    torch.cuda.synchronize()
    del a, b, c
    if (i+1) % 5 == 0:
        print(f'  {i+1}/20 done ({time.time()-t0:.1f}s)')
torch.cuda.empty_cache()
print(f'  Total: {time.time()-t0:.1f}s')

# Test 4: Keep tensors on GPU (like lowvram loads many layers)
print()
print('Test 4: Load 30 layers to GPU simultaneously [1024x2560 f16]')
t0 = time.time()
layers = []
for i in range(30):
    cpu_t = torch.randn(1024, 2560, dtype=torch.float16)
    gpu_t = cpu_t.to('cuda:0', non_blocking=False)
    layers.append(gpu_t)
    del cpu_t
    if (i+1) % 10 == 0:
        mem = torch.cuda.memory_allocated() / 1024**2
        print(f'  {i+1}/30 done, GPU mem: {mem:.0f}MB ({time.time()-t0:.1f}s)')
torch.cuda.synchronize()
print(f'  Total: {time.time()-t0:.1f}s')

# Cleanup
del layers
torch.cuda.empty_cache()
gc.collect()

# Test 5: Rapid alloc/free cycle (simulating lowvram)
print()
print('Test 5: 50x rapid alloc-compute-free cycle [2560x2560 f16]')
t0 = time.time()
for i in range(50):
    a = torch.randn(2560, 2560, dtype=torch.float16).to('cuda:0', non_blocking=False)
    b = torch.randn(2560, 2560, dtype=torch.float16).to('cuda:0', non_blocking=False)
    c = torch.matmul(a, b)
    torch.cuda.synchronize()
    del a, b, c
    torch.cuda.empty_cache()
    if (i+1) % 10 == 0:
        print(f'  {i+1}/50 done ({time.time()-t0:.1f}s)')
print(f'  Total: {time.time()-t0:.1f}s')

print()
print('ALL TESTS PASSED')
"
