Added device_id() methods to gpu_data and tensor objects. These functions

allow you to find out which device owns the memory inside these objects.
This commit is contained in:
Davis King 2016-04-27 20:18:57 -04:00
parent 8ece542852
commit 7d5fb9c60c
5 changed files with 26 additions and 1 deletions

View File

@ -136,6 +136,8 @@ namespace dlib
try
{
CHECK_CUDA(cudaGetDevice(&the_device_id));
void* data;
CHECK_CUDA(cudaMallocHost(&data, new_size*sizeof(float)));
// Note that we don't throw exceptions since the free calls are invariably

View File

@ -40,7 +40,7 @@ namespace dlib
public:
gpu_data(
) : data_size(0), host_current(true), device_current(true),have_active_transfer(false),device_in_use(false)
) : data_size(0), host_current(true), device_current(true),have_active_transfer(false),device_in_use(false), the_device_id(0)
{
}
@ -52,6 +52,7 @@ namespace dlib
gpu_data(gpu_data&& item) : gpu_data() { swap(item); }
gpu_data& operator=(gpu_data&& item) { swap(item); return *this; }
int device_id() const { return the_device_id; }
#ifdef DLIB_USE_CUDA
void async_copy_to_device() const;
@ -153,6 +154,7 @@ namespace dlib
std::swap(data_host, item.data_host);
std::swap(data_device, item.data_device);
std::swap(cuda_stream, item.cuda_stream);
std::swap(the_device_id, item.the_device_id);
}
private:
@ -177,6 +179,7 @@ namespace dlib
std::shared_ptr<float> data_host;
std::shared_ptr<float> data_device;
std::shared_ptr<void> cuda_stream;
int the_device_id;
};
inline void serialize(const gpu_data& item, std::ostream& out)

View File

@ -45,6 +45,7 @@ namespace dlib
- #device() == nullptr
- #host_ready() == true
- #device_ready() == true
- #device_id() == 0
!*/
// This object is not copyable, however, it is movable.
@ -53,6 +54,14 @@ namespace dlib
gpu_data(gpu_data&& item);
gpu_data& operator=(gpu_data&& item);
int device_id(
) const;
/*!
ensures
- returns the ID of the CUDA device that allocated this memory. I.e. the
number returned by cudaGetDevice() when the memory was allocated.
- If CUDA is not being used then this function always returns 0.
!*/
void async_copy_to_device(
);

View File

@ -52,6 +52,8 @@ namespace dlib
virtual float* device() = 0;
virtual float* device_write_only() = 0;
int device_id() const { return data().device_id(); }
tensor& operator= (float val)
{
#ifdef DLIB_USE_CUDA

View File

@ -187,6 +187,15 @@ namespace dlib
every memory location in the returned memory block.
!*/
int device_id(
) const;
/*!
ensures
- returns the ID of the CUDA device that allocated this memory. I.e. the
number returned by cudaGetDevice() when the memory was allocated.
- If CUDA is not being used then this function always returns 0.
!*/
tensor& operator= (
float val
);