"""
BC-250 (gfx1010) Softmax Fix Test
Tests if manual softmax works where fused kernel crashes with VGPR overflow.
"""
import torch
import time
import os
import gc

print(f'Device: {torch.cuda.get_device_name(0)}')
print()

def safe_softmax(x, dim=-1):
    """Manual softmax using elementwise ops - avoids fused kernel VGPR overflow."""
    x_max = x.max(dim=dim, keepdim=True).values
    exp_x = torch.exp(x - x_max)
    return exp_x / exp_x.sum(dim=dim, keepdim=True)

# Test 1: Does safe_softmax work where torch.softmax crashes?
print('=== Test 1: safe_softmax on [1, 24, 512, 4096] ===', flush=True)
try:
    x = torch.randn(1, 24, 512, 4096, device='cuda', dtype=torch.float32)
    torch.cuda.synchronize()
    t = time.time()
    y = safe_softmax(x, dim=-1)
    torch.cuda.synchronize()
    print(f'  OK: {time.time()-t:.3f}s', flush=True)
    # Verify correctness
    assert y.min() >= 0, "Negative values in softmax"
    assert abs(y.sum(dim=-1).mean().item() - 1.0) < 0.001, "Softmax doesn't sum to 1"
    print(f'  Correctness: OK (sum={y.sum(dim=-1).mean().item():.6f})', flush=True)
    del x, y
    gc.collect()
    torch.cuda.empty_cache()
except Exception as e:
    print(f'  FAIL: {e}', flush=True)

# Test 2: Full attention pipeline with safe_softmax
print('=== Test 2: Full attention with safe_softmax (24h, 4096 seq) ===', flush=True)
try:
    heads, seq, hd = 24, 4096, 128
    chunk = 512
    
    q = torch.randn(1, heads, seq, hd, device='cuda', dtype=torch.float32)
    k = torch.randn(1, heads, seq, hd, device='cuda', dtype=torch.float32)
    v = torch.randn(1, heads, seq, hd, device='cuda', dtype=torch.float32)
    out = torch.zeros(1, heads, seq, hd, device='cuda', dtype=torch.float32)
    
    scale = hd ** -0.5
    
    torch.cuda.synchronize()
    t = time.time()
    for i in range(0, seq, chunk):
        q_chunk = q[:, :, i:i+chunk, :] * scale
        scores = torch.matmul(q_chunk, k.transpose(-2, -1))
        attn = safe_softmax(scores, dim=-1)
        out[:, :, i:i+chunk, :] = torch.matmul(attn, v)
        del scores, attn
    torch.cuda.synchronize()
    elapsed = time.time() - t
    print(f'  OK: {elapsed:.3f}s', flush=True)
    del q, k, v, out
    gc.collect()
    torch.cuda.empty_cache()
except Exception as e:
    print(f'  FAIL: {e}', flush=True)

# Test 3: Also verify torch.softmax on SMALL tensors still works
print('=== Test 3: torch.softmax on small tensor [1, 8, 256, 256] ===', flush=True)
try:
    x = torch.randn(1, 8, 256, 256, device='cuda', dtype=torch.float32)
    torch.cuda.synchronize()
    y = torch.softmax(x, dim=-1)
    torch.cuda.synchronize()
    print(f'  OK (native softmax works on small tensors)', flush=True)
    del x, y
except Exception as e:
    print(f'  FAIL: {e}', flush=True)

# Test 4: Find the threshold where torch.softmax breaks
print('=== Test 4: Find native softmax failure threshold ===', flush=True)
for last_dim in [256, 512, 1024, 2048, 4096]:
    try:
        x = torch.randn(1, 24, 512, last_dim, device='cuda', dtype=torch.float32)
        torch.cuda.synchronize()
        y = torch.softmax(x, dim=-1)
        torch.cuda.synchronize()
        print(f'  dim={last_dim}: OK', flush=True)
        del x, y
        gc.collect()
        torch.cuda.empty_cache()
    except Exception as e:
        print(f'  dim={last_dim}: FAIL ({e})', flush=True)
        gc.collect()
        torch.cuda.empty_cache()
        break

# Test 5: Test SDPA (scaled_dot_product_attention) - what ComfyUI actually uses
print('=== Test 5: F.scaled_dot_product_attention ===', flush=True)
try:
    q = torch.randn(1, 24, 4096, 128, device='cuda', dtype=torch.float32)
    k = torch.randn(1, 24, 4096, 128, device='cuda', dtype=torch.float32)
    v = torch.randn(1, 24, 4096, 128, device='cuda', dtype=torch.float32)
    torch.cuda.synchronize()
    t = time.time()
    out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
    torch.cuda.synchronize()
    print(f'  SDPA OK: {time.time()-t:.3f}s', flush=True)
    del q, k, v, out
except Exception as e:
    print(f'  SDPA FAIL: {e}', flush=True)

# Test 5b: SDPA with smaller dims
print('=== Test 5b: SDPA smaller (8h, 1024 seq) ===', flush=True)
try:
    q = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
    k = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
    v = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
    torch.cuda.synchronize()
    t = time.time()
    out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
    torch.cuda.synchronize()
    print(f'  SDPA OK: {time.time()-t:.3f}s', flush=True)
    del q, k, v, out
except Exception as e:
    print(f'  SDPA FAIL: {e}', flush=True)

print()
print('=== ALL DONE ===', flush=True)
torch.cuda.synchronize()
os._exit(0)