Deprecate `DeviceBase::GetStepAllocator()` and replace with calls to `GetAllocator()`.

The `GetStepAllocator()` API relied on the existence of a "step resource manager",
which is no longer a concept in the runtime (it was replaced by "step containers").
Since the additional flexibility does not appear to be used in the codebase, the
`GetScopedAllocator()` seems to provide a similar extension point (based on step IDs),
and the `OpKernelContext::get_allocator()` method is called frequently, this change
simplifies the implementation somewhat.

The `GetStepAllocator()` method is retained as a non-virtual stub that forwards to
`GetAllocator()`, because at least one third-party library (libxsmm) calls this
interface directly.

PiperOrigin-RevId: 197922154

1 files changed, 4 insertions, 7 deletions

diff --git a/tensorflow/core/framework/device_base.h b/tensorflow/core/framework/device_base.h
index 223b74857d..ec26d92a61 100644
--- a/tensorflow/core/framework/device_base.h
+++ b/tensorflow/core/framework/device_base.h
@@ -169,13 +169,10 @@ class DeviceBase {
     return nullptr;
   }
 
-  // Return the Allocator implementation to use based on the allocator
-  // attributes requested and the supplied resource manager. By
-  // default this ignores the resource manager and calls the base
-  // implementation but devices can override if they want to consult
-  // the resource manager when choosing the allocator.
-  virtual Allocator* GetStepAllocator(AllocatorAttributes attr,
-                                      ResourceMgr* /*step_resource_manager*/) {
+  // DEPRECATED: Use `this->GetAllocator()` or `this->GetScopedAllocator()`.
+  // This method is provided for backwards compatibility, and will be removed
+  // in a future release.
+  Allocator* GetStepAllocator(AllocatorAttributes attr, ResourceMgr*) {
     return GetAllocator(attr);
   }