[OpenCL] Stats tracking

tensorflow/core/common_runtime/sycl/sycl_allocator.cc

@@ -19,6 +19,14 @@ limitations under the License.
 
 namespace tensorflow {
 
+SYCLAllocator::SYCLAllocator(Eigen::QueueInterface *queue)
+    : sycl_device_(new Eigen::SyclDevice(queue)) {
+  cl::sycl::queue& sycl_queue = sycl_device_->sycl_queue();
+  const cl::sycl::device& device = sycl_queue.get_device();
+  stats_.bytes_limit =
+      device.get_info<cl::sycl::info::device::max_mem_alloc_size>();
+}
+
 SYCLAllocator::~SYCLAllocator() {
   if(sycl_device_) {
     delete sycl_device_;
@@ -30,18 +38,44 @@ string SYCLAllocator::Name() { return "device:SYCL"; }
 void *SYCLAllocator::AllocateRaw(size_t alignment, size_t num_bytes) {
   assert(sycl_device_);
   if (num_bytes == 0) {
-    return sycl_device_->allocate(1);
+    // Cannot allocate no bytes in SYCL, so instead allocate a single byte
+    num_bytes = 1;
   }
   auto p = sycl_device_->allocate(num_bytes);
+  const auto& allocated_buffer = sycl_device_->get_sycl_buffer(p);
+  const std::size_t bytes_allocated = allocated_buffer.get_range().size();
+
+  mutex_lock lock(mu_);
+  ++stats_.num_allocs;
+  stats_.bytes_in_use += bytes_allocated;
+  stats_.max_bytes_in_use =
+      std::max<int64>(stats_.max_bytes_in_use, stats_.bytes_in_use);
+  stats_.max_alloc_size =
+      std::max<int64>(stats_.max_alloc_size, bytes_allocated);
+
   return p;
 }
 
 void SYCLAllocator::DeallocateRaw(void *ptr) {
+  const auto& buffer_to_delete = sycl_device_->get_sycl_buffer(ptr);
+  const std::size_t dealloc_size = buffer_to_delete.get_range().size();
+  mutex_lock lock(mu_);
+  stats_.bytes_in_use -= dealloc_size;
   if (sycl_device_) {
     sycl_device_->deallocate(ptr);
   }
 }
 
+void SYCLAllocator::GetStats(AllocatorStats* stats) {
+  mutex_lock lock(mu_);
+  *stats = stats_;
+}
+
+size_t SYCLAllocator::RequestedSize(void* ptr) {
+  const auto& buffer = sycl_device_->get_sycl_buffer(ptr);
+  return buffer.get_size();
+}
+
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_USE_SYCL

tensorflow/core/common_runtime/sycl/sycl_allocator.h

@@ -21,14 +21,15 @@ limitations under the License.
 #define TENSORFLOW_COMMON_RUNTIME_SYCL_SYCL_ALLOCATOR_H_
 
 #include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/types.h"
 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 
 namespace tensorflow {
 
 class SYCLAllocator : public Allocator {
  public:
-  SYCLAllocator(Eigen::QueueInterface *queue) : sycl_device_(new Eigen::SyclDevice(queue)) {}
+  SYCLAllocator(Eigen::QueueInterface *queue);
   virtual ~SYCLAllocator() override;
   string Name() override;
   void *AllocateRaw(size_t alignment, size_t num_bytes) override;
@@ -37,10 +38,20 @@ class SYCLAllocator : public Allocator {
   virtual bool ShouldAllocateEmptyTensors() override final { return true; }
   void Synchronize() { sycl_device_->synchronize(); }
   bool Ok() { return sycl_device_->ok(); }
+  void GetStats(AllocatorStats* stats) override;
+  // The SYCL buffers keep track of their size, so we already have tracking.
+  bool TracksAllocationSizes() override { return true; }
+  // Get the size of the corresponding SYCL buffer.
+  // Implementing this also provides an implementation of
+  // AllocatedSize(void* ptr) by default.
+  size_t RequestedSize(void* ptr) override;
   Eigen::SyclDevice* getSyclDevice() { return sycl_device_; }
  private:
   Eigen::SyclDevice *sycl_device_;  // owned
 
+  mutable mutex mu_;
+  AllocatorStats stats_ GUARDED_BY(mu_);
+
   TF_DISALLOW_COPY_AND_ASSIGN(SYCLAllocator);
 };