Allow tensorflow devices to report their load. This may be used to improve batch scheduling.

PiperOrigin-RevId: 171716813

3 files changed, 10 insertions, 0 deletions

diff --git a/tensorflow/stream_executor/stream_executor_internal.h b/tensorflow/stream_executor/stream_executor_internal.h
index 802ef755eb..12593e31d4 100644
--- a/tensorflow/stream_executor/stream_executor_internal.h
+++ b/tensorflow/stream_executor/stream_executor_internal.h
@@ -225,6 +225,8 @@ class StreamExecutorInterface {
   virtual port::Status SetDeviceSharedMemoryConfig(
       SharedMemoryConfig config) = 0;
 
+  virtual int64 GetDeviceLoad() { return -1; }
+
   virtual bool DeviceMemoryUsage(int64 *free, int64 *total) const {
     return false;
   }
diff --git a/tensorflow/stream_executor/stream_executor_pimpl.cc b/tensorflow/stream_executor/stream_executor_pimpl.cc
index 9bbfe7f04a..9dc1749327 100644
--- a/tensorflow/stream_executor/stream_executor_pimpl.cc
+++ b/tensorflow/stream_executor/stream_executor_pimpl.cc
@@ -268,6 +268,10 @@ const DeviceDescription &StreamExecutor::GetDeviceDescription() const {
   return *device_description_;
 }
 
+int64 StreamExecutor::GetDeviceLoad() const {
+  return implementation_->GetDeviceLoad();
+}
+
 int StreamExecutor::PlatformDeviceCount() const {
   return implementation_->PlatformDeviceCount();
 }
diff --git a/tensorflow/stream_executor/stream_executor_pimpl.h b/tensorflow/stream_executor/stream_executor_pimpl.h
index f354317a6e..9c225e5fae 100644
--- a/tensorflow/stream_executor/stream_executor_pimpl.h
+++ b/tensorflow/stream_executor/stream_executor_pimpl.h
@@ -312,6 +312,10 @@ class StreamExecutor {
   // The value is cached on first use.
   const DeviceDescription &GetDeviceDescription() const;
 
+  // If implemented, returns device specific measurement of load
+  // (e.g. pending requests).
+  int64 GetDeviceLoad() const;
+
   // Returns the underlying device memory usage information, if it is available.
   // If it is not available (false is returned), free/total may not be
   // initialized.