import tensorflow as tf
import time

print("TF version:", tf.__version__)
print("Built with CUDA:", tf.test.is_built_with_cuda())
print("GPUs:", tf.config.list_physical_devices('GPU'))

# ✅ GPU가 잡혔으면 어떤 장치인지 출력
if tf.config.list_physical_devices('GPU'):
    print("✅ GPU is AVAILABLE")
    print("GPU device:", tf.test.gpu_device_name())
else:
    print("❌ GPU NOT FOUND")
    exit()

# ✅ 실제 연산이 GPU에서 돌아가는지 확인 (matrix multiply)
with tf.device('/GPU:0'):
    a = tf.random.normal([4096, 4096])
    b = tf.random.normal([4096, 4096])

    start = time.time()
    c = tf.matmul(a, b)

    # ✅ 반드시 실행 완료까지 기다림 (TF 2.x 안정적인 방식)
    _ = c.numpy()

    end = time.time()

print("✅ MatMul done on GPU")
print("Time:", end - start, "sec")
print("Result shape:", c.shape)

# ✅ 디바이스 확인 (진짜 GPU에서 수행됐는지)
print("Executed on:", c.device)