Get cuda_common.h from master.

This commit is contained in:
Sergii Dymchenko 2020-04-09 16:56:52 -07:00
parent 84773c61c6
commit 0e4080f1d6

View file

@ -92,8 +92,6 @@ class CudaKernel : public OpKernel {
provider_->AddDeferredReleaseCPUPtr(p);
}
const cudaDeviceProp& GetDeviceProp() const { return provider_->GetDeviceProp(); };
// To support cudaMemcpyAsync, the cpu memory should be allocated in pinned memory
// and it can only be released after the copy has finished
template <typename T>
@ -224,6 +222,38 @@ inline bool CalculateFdmStrides(gsl::span<fast_divmod> p, const std::vector<int6
return true;
}
struct DeviceProp {
static const std::vector<cudaDeviceProp>& GetCachedDeviceProps() {
std::call_once(s_cachedDevicePropsInitFlag, [=] {
int numDevices;
// must wait GPU idle, otherwise cudaGetDeviceProperties might fail
CUDA_CALL_THROW(cudaDeviceSynchronize());
CUDA_CALL_THROW(cudaGetDeviceCount(&numDevices));
s_cachedDeviceProps.resize(numDevices);
for (int i = 0; i < numDevices; i++)
CUDA_CALL_THROW(cudaGetDeviceProperties(&s_cachedDeviceProps[i], i));
});
return s_cachedDeviceProps;
}
static size_t GetCurrentDeviceId() {
int deviceId;
cudaGetDevice(&deviceId);
return (size_t)deviceId;
}
// get device properties of current device
static const cudaDeviceProp& GetDeviceProps() {
const auto& cachedDevicesProps = GetCachedDeviceProps();
return cachedDevicesProps[GetCurrentDeviceId()];
}
private:
static std::vector<cudaDeviceProp> s_cachedDeviceProps;
static std::once_flag s_cachedDevicePropsInitFlag;
};
class CublasMathModeSetter {
public:
CublasMathModeSetter(cublasHandle_t handle, cublasMath_t mode) : handle_(handle) {