#!/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

print('=== BC-250 Transfer Test (no empty_cache) ===')
print(f'Device: {torch.cuda.get_device_name(0)}')

# Test: Keep all tensors on GPU, no empty_cache(), do compute
print()
print('Test: 50x transfer+matmul [2560x2560 f16], NO empty_cache')
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
    if (i+1) % 10 == 0:
        mem = torch.cuda.memory_allocated() / 1024**2
        print(f'  {i+1}/50 done, GPU mem: {mem:.0f}MB ({time.time()-t0:.1f}s)')
print(f'  Total: {time.time()-t0:.1f}s')

# Test 2: Simulate lowvram layer loading pattern
print()
print('Test2: Simulated UNET forward (453 layers)')
t0 = time.time()
for i in range(453):
    # Simulate dequant on CPU -> transfer to GPU
    w = torch.randn(1024, 1024, dtype=torch.float16).to('cuda:0', non_blocking=False)
    # Simulate compute
    x = torch.randn(1, 1024, dtype=torch.float16, device='cuda:0')
    y = torch.matmul(x, w.T)
    torch.cuda.synchronize()
    del w, x, y
    if (i+1) % 100 == 0:
        mem = torch.cuda.memory_allocated() / 1024**2
        print(f'  Layer {i+1}/453, GPU mem: {mem:.0f}MB ({time.time()-t0:.1f}s)')
print(f'  Total: {time.time()-t0:.1f}s')

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