"""Test if torch.cuda.synchronize() crashes after F.linear on gfx1010"""
import torch
import torch.nn.functional as F
import time, os, signal, sys

os.environ['HSA_OVERRIDE_GFX_VERSION'] = '10.1.0'
os.environ['HSA_ENABLE_SDMA'] = '0'
os.environ['TORCHDYNAMO_DISABLE'] = '1'

# Catch SIGSEGV
def segfault_handler(signum, frame):
    print(f"\n*** SIGNAL {signum} (SIGSEGV) CAUGHT ***", flush=True)
    sys.exit(1)
signal.signal(signal.SIGSEGV, segfault_handler)

print(f"Device: {torch.cuda.get_device_name(0)}", flush=True)

# Simulate exactly what the model does:
# For each layer: transfer from CPU → F.linear → sync → next layer
shapes = [
    ("CBW505", (1024, 256)),
    ("CBW506", (256, 1024)),
    ("CBW507", (3840, 2560)),
    ("CBW508", (3840, 64)),
    ("CBW509", (11520, 3840)),   # This is where ComfyUI crashes
    ("CBW510", (3840, 3840)),
    ("CBW511", (3840, 3840)),
    ("CBW512", (11520, 3840)),
    ("CBW513", (3840, 11520)),
]

# Start with latent-like input
x = torch.randn(1, 1024, 256, dtype=torch.float16, device='cuda:0')
torch.cuda.synchronize()
print("Initial input on GPU", flush=True)

for i, (name, (out_feat, in_feat)) in enumerate(shapes):
    print(f"\n--- {name}: Linear({in_feat} -> {out_feat}) ---", flush=True)
    
    # Step 1: Pre-sync (like v20 patch)
    print(f"  pre-sync...", flush=True) 
    torch.cuda.synchronize()
    print(f"  pre-sync OK", flush=True)
    
    # Step 2: CPU dequant simulation → transfer to GPU
    w_cpu = torch.randn(out_feat, in_feat, dtype=torch.float16)
    b_cpu = torch.randn(out_feat, dtype=torch.float16)
    print(f"  CPU weights created", flush=True)
    
    w_gpu = w_cpu.to('cuda:0', non_blocking=False)
    b_gpu = b_cpu.to('cuda:0', non_blocking=False)
    print(f"  transferred to GPU", flush=True)
    
    # Step 3: Adjust input dimension
    if x.shape[-1] != in_feat:
        x = torch.randn(1, 1024, in_feat, dtype=torch.float16, device='cuda:0')
        torch.cuda.synchronize()
        print(f"  input resized to [{1}, {1024}, {in_feat}]", flush=True)
    
    # Step 4: F.linear (dispatches GPU kernel)
    print(f"  F.linear starting...", flush=True)
    x = F.linear(x, w_gpu, b_gpu)
    print(f"  F.linear dispatched, out={list(x.shape)}", flush=True)
    
    # Step 5: Post-sync (forces GPU to finish and report errors)
    print(f"  post-sync...", flush=True)
    torch.cuda.synchronize()
    print(f"  post-sync OK", flush=True)
    
    del w_cpu, b_cpu, w_gpu, b_gpu
    print(f"  {name} DONE", flush=True)

print("\n=== ALL PASSED ===")