You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
29 lines
786 B
29 lines
786 B
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
|
|
#pragma once
|
|
|
|
#include <ATen/cuda/Exceptions.h>
|
|
|
|
#include <cuda.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
namespace at::cuda {
|
|
|
|
inline Device getDeviceFromPtr(void* ptr) {
|
|
cudaPointerAttributes attr{};
|
|
|
|
AT_CUDA_CHECK(cudaPointerGetAttributes(&attr, ptr));
|
|
|
|
#if !defined(USE_ROCM)
|
|
TORCH_CHECK(attr.type != cudaMemoryTypeUnregistered,
|
|
"The specified pointer resides on host memory and is not registered with any CUDA device.");
|
|
#endif
|
|
|
|
return {c10::DeviceType::CUDA, static_cast<DeviceIndex>(attr.device)};
|
|
}
|
|
|
|
} // namespace at::cuda
|
|
|
|
#else
|
|
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
|
|
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
|