import torch
import os

print(f"HSA_OVERRIDE_GFX_VERSION={os.environ.get('HSA_OVERRIDE_GFX_VERSION','NOT SET')}")
print(f"PyTorch arch list: {torch.cuda.get_arch_list()}")
print(f"Device: {torch.cuda.get_device_name(0)}")
print(f"GCN Arch: {torch.cuda.get_device_properties(0).gcnArchName}")

try:
    t = torch.randn(256, 256, device="cuda")
    r = torch.mm(t, t)
    val = r[0,0].item()
    print(f"GPU matmul OK! result[0,0]={val:.4f}")
    
    # Bigger test
    a = torch.randn(1024, 1024, device="cuda")
    b = torch.randn(1024, 1024, device="cuda")
    c = torch.mm(a, b)
    print(f"Large matmul OK! shape={c.shape}")
    
    # Test fp32 conv
    x = torch.randn(1, 3, 64, 64, device="cuda")
    conv = torch.nn.Conv2d(3, 16, 3, padding=1).cuda()
    y = conv(x)
    print(f"Conv2d OK! output shape={y.shape}")
    
    print("ALL GPU TESTS PASSED!")
except Exception as e:
    print(f"FAILED: {e}")