chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use CUDA runtime API to get device prop instead of ATen

Browse Source
This commit is contained in:
Chenggang Zhao
2025-10-11 09:16:31 +08:00
parent 9f196058ae
commit f8f41145da
1 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/jit/device_runtime.hpp
+7 -2
View File
@@ -36,8 +36,13 @@ public:
} }
std::shared_ptr<cudaDeviceProp> get_prop() { std::shared_ptr<cudaDeviceProp> get_prop() {
if (cached_prop == nullptr) if (cached_prop == nullptr) {
cached_prop = std::make_shared<cudaDeviceProp>(*at::cuda::getCurrentDeviceProperties()); int device_idx;
cudaDeviceProp prop;
DG_CUDA_RUNTIME_CHECK(cudaGetDevice(&device_idx));
DG_CUDA_RUNTIME_CHECK(cudaGetDeviceProperties(&prop, device_idx));
cached_prop = std::make_shared<cudaDeviceProp>(prop);
}
return cached_prop; return cached_prop;
} }