import cupy as cp
import time

# Verify available GPU hardware allocation
try:
    device = cp.cuda.Device(0)
    print(f"Running on Device: {device}")
except Exception as e:
    print(f"CUDA Error: {e}. Make sure you are running on a GPU compute node!")
    exit(1)

# Initialize large matrices directly on the GPU memory space
print("Allocating multi-dimensional arrays on the GPU...")
size = 5000
x_gpu = cp.ones((size, size), dtype=cp.float32)
y_gpu = cp.ones((size, size), dtype=cp.float32)

# Perform matrix multiplication on the GPU cores
print(f"Performing matrix multiplication ({size}x{size})...")
start_time = time.time()
z_gpu = cp.dot(x_gpu, y_gpu)

# Force CUDA synchronization to get accurate benchmarks
cp.cuda.Stream.null.synchronize()
end_time = time.time()

print(f"Operation completed successfully in: {end_time - start_time:.4f} seconds")

# Pull calculation matrix value back to CPU host memory for confirmation
result_cpu = cp.asnumpy(z_gpu)
print(f"Verification check value: {result_cpu[0, 0]} (Expected: {float(size)})")