Commit 7d5fb9c6 authored by Davis King's avatar Davis King

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.
parent 8ece5428
...@@ -136,6 +136,8 @@ namespace dlib ...@@ -136,6 +136,8 @@ namespace dlib
try try
{ {
CHECK_CUDA(cudaGetDevice(&the_device_id));
void* data; void* data;
CHECK_CUDA(cudaMallocHost(&data, new_size*sizeof(float))); CHECK_CUDA(cudaMallocHost(&data, new_size*sizeof(float)));
// Note that we don't throw exceptions since the free calls are invariably // Note that we don't throw exceptions since the free calls are invariably
......
...@@ -40,7 +40,7 @@ namespace dlib ...@@ -40,7 +40,7 @@ namespace dlib
public: public:
gpu_data( 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 ...@@ -52,6 +52,7 @@ namespace dlib
gpu_data(gpu_data&& item) : gpu_data() { swap(item); } gpu_data(gpu_data&& item) : gpu_data() { swap(item); }
gpu_data& operator=(gpu_data&& item) { swap(item); return *this; } gpu_data& operator=(gpu_data&& item) { swap(item); return *this; }
int device_id() const { return the_device_id; }
#ifdef DLIB_USE_CUDA #ifdef DLIB_USE_CUDA
void async_copy_to_device() const; void async_copy_to_device() const;
...@@ -153,6 +154,7 @@ namespace dlib ...@@ -153,6 +154,7 @@ namespace dlib
std::swap(data_host, item.data_host); std::swap(data_host, item.data_host);
std::swap(data_device, item.data_device); std::swap(data_device, item.data_device);
std::swap(cuda_stream, item.cuda_stream); std::swap(cuda_stream, item.cuda_stream);
std::swap(the_device_id, item.the_device_id);
} }
private: private:
...@@ -177,6 +179,7 @@ namespace dlib ...@@ -177,6 +179,7 @@ namespace dlib
std::shared_ptr<float> data_host; std::shared_ptr<float> data_host;
std::shared_ptr<float> data_device; std::shared_ptr<float> data_device;
std::shared_ptr<void> cuda_stream; std::shared_ptr<void> cuda_stream;
int the_device_id;
}; };
inline void serialize(const gpu_data& item, std::ostream& out) inline void serialize(const gpu_data& item, std::ostream& out)
......
...@@ -45,6 +45,7 @@ namespace dlib ...@@ -45,6 +45,7 @@ namespace dlib
- #device() == nullptr - #device() == nullptr
- #host_ready() == true - #host_ready() == true
- #device_ready() == true - #device_ready() == true
- #device_id() == 0
!*/ !*/
// This object is not copyable, however, it is movable. // This object is not copyable, however, it is movable.
...@@ -53,6 +54,14 @@ namespace dlib ...@@ -53,6 +54,14 @@ namespace dlib
gpu_data(gpu_data&& item); gpu_data(gpu_data&& item);
gpu_data& operator=(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( void async_copy_to_device(
); );
......
...@@ -52,6 +52,8 @@ namespace dlib ...@@ -52,6 +52,8 @@ namespace dlib
virtual float* device() = 0; virtual float* device() = 0;
virtual float* device_write_only() = 0; virtual float* device_write_only() = 0;
int device_id() const { return data().device_id(); }
tensor& operator= (float val) tensor& operator= (float val)
{ {
#ifdef DLIB_USE_CUDA #ifdef DLIB_USE_CUDA
......
...@@ -187,6 +187,15 @@ namespace dlib ...@@ -187,6 +187,15 @@ namespace dlib
every memory location in the returned memory block. 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= ( tensor& operator= (
float val float val
); );
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment