Replace enum BusAdjacency with protobuf DeviceLocality for describing the

topological neighborhood of a device.
Change: 136663586

1 files changed, 7 insertions, 8 deletions

diff --git a/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc b/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc
index a7e3773ac8..878256d5d7 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc
@@ -26,10 +26,10 @@ namespace tensorflow {
 class GPUDevice : public BaseGPUDevice {
  public:
   GPUDevice(const SessionOptions& options, const string& name,
-            Bytes memory_limit, BusAdjacency bus_adjacency, int gpu_id,
+            Bytes memory_limit, const DeviceLocality& locality, int gpu_id,
             const string& physical_device_desc, Allocator* gpu_allocator,
             Allocator* cpu_allocator)
-      : BaseGPUDevice(options, name, memory_limit, bus_adjacency, gpu_id,
+      : BaseGPUDevice(options, name, memory_limit, locality, gpu_id,
                       physical_device_desc, gpu_allocator, cpu_allocator,
                       false /* sync every op */, 1 /* max_streams */) {}
 
@@ -51,11 +51,11 @@ class GPUDeviceFactory : public BaseGPUDeviceFactory {
  private:
   BaseGPUDevice* CreateGPUDevice(const SessionOptions& options,
                                  const string& name, Bytes memory_limit,
-                                 BusAdjacency bus_adjacency, int gpu_id,
+                                 const DeviceLocality& locality, int gpu_id,
                                  const string& physical_device_desc,
                                  Allocator* gpu_allocator,
                                  Allocator* cpu_allocator) override {
-    return new GPUDevice(options, name, memory_limit, bus_adjacency, gpu_id,
+    return new GPUDevice(options, name, memory_limit, locality, gpu_id,
                          physical_device_desc, gpu_allocator, cpu_allocator);
   }
 };
@@ -69,10 +69,9 @@ REGISTER_LOCAL_DEVICE_FACTORY("GPU", GPUDeviceFactory);
 class GPUCompatibleCPUDevice : public ThreadPoolDevice {
  public:
   GPUCompatibleCPUDevice(const SessionOptions& options, const string& name,
-                         Bytes memory_limit, BusAdjacency bus_adjacency,
+                         Bytes memory_limit, const DeviceLocality& locality,
                          Allocator* allocator)
-      : ThreadPoolDevice(options, name, memory_limit, bus_adjacency,
-                         allocator) {}
+      : ThreadPoolDevice(options, name, memory_limit, locality, allocator) {}
   ~GPUCompatibleCPUDevice() override {}
 
   Allocator* GetAllocator(AllocatorAttributes attr) override {
@@ -99,7 +98,7 @@ class GPUCompatibleCPUDeviceFactory : public DeviceFactory {
     for (int i = 0; i < n; i++) {
       string name = strings::StrCat(name_prefix, "/cpu:", i);
       devices->push_back(new GPUCompatibleCPUDevice(
-          options, name, Bytes(256 << 20), BUS_ANY, cpu_allocator()));
+          options, name, Bytes(256 << 20), DeviceLocality(), cpu_allocator()));
     }
 
     return Status::OK();