import os
import socket
import torch
from torch.utils.cpp_extension import load_inline

# 提前获取 hostname 和 pid，方便打印出触发 coredump 的准确命令
hostname = socket.gethostname()
pid = os.getpid()
pipe_path = f"/tmp/cuda_pipe_{hostname}_{pid}"
dump_path = f"/tmp/cuda_coredump_full_{hostname}_{pid}"

cuda_source = """
__global__ void hang_kernel(int* dummy) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx == 0) {
        volatile int counter = 0;
        // 死循环，让 GPU 彻底卡在这里
        while(true) {
            counter++; // 使用 volatile 变量防止循环被 NVCC 优化掉
        }
    }
}

void run_hang_kernel(torch::Tensor dummy) {
    hang_kernel<<<1, 1>>>(dummy.data_ptr<int>());
}
"""

cpp_source = """
void run_hang_kernel(torch::Tensor dummy);
"""

print("=> 正在编译 CUDA Kernel (可能需要几秒钟)...")
module = load_inline(
    name='cuda_hang_ext',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['run_hang_kernel'],
    verbose=False,
)

print("\n=> 编译完成！")
print(f"=> 当前 Hostname: {hostname}")
print(f"=> 当前 PID: {pid}")
print("\n" + "="*60)
print("请保留当前终端卡住，打开另一个新终端，执行以下命令触发 coredump:")
print(f"echo 1 > {pipe_path}")
print("="*60 + "\n")

# 分配一个 dummy tensor，触发 CUDA 上下文初始化
dummy_tensor = torch.zeros(1, dtype=torch.int32, device='cuda')

# 启动 kernel
module.run_hang_kernel(dummy_tensor)

print("=> Kernel 已下发！正在等待 torch.cuda.synchronize() ... (预期会永远卡在此处)")
torch.cuda.synchronize()