34 lines
1.0 KiB
Python
34 lines
1.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Check BI-V100 warp size."""
|
|
import torch
|
|
print(f"torch.cuda.get_device_properties(0).warp_size: "
|
|
f"{getattr(torch.cuda.get_device_properties(0), 'warp_size', 'N/A')}")
|
|
|
|
# Also check via CUDA kernel
|
|
from torch.utils.cpp_extension import load
|
|
import tempfile, os
|
|
cu_code = r'''
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
__global__ void check_warp(int* out) {
|
|
if (threadIdx.x == 0 && threadIdx.y == 0) {
|
|
out[0] = warpSize;
|
|
}
|
|
}
|
|
torch::Tensor get_warp_size() {
|
|
auto out = torch::zeros({1}, torch::dtype(torch::kInt32).device(torch::kCUDA));
|
|
check_warp<<<1, 32>>>(out.data_ptr<int>());
|
|
return out;
|
|
}
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("get_warp_size", &get_warp_size);
|
|
}
|
|
'''
|
|
with tempfile.NamedTemporaryFile(suffix='.cu', mode='w', delete=False) as f:
|
|
f.write(cu_code)
|
|
cu_path = f.name
|
|
ext = load(name="warpcheck", sources=[cu_path], verbose=False)
|
|
ws = ext.get_warp_size().item()
|
|
print(f"CUDA kernel warpSize: {ws}")
|
|
os.unlink(cu_path)
|