Clean up cpu_allocator() uses to use ProcessState where possible.
1. Rename cpu_allocator() to cpu_allocator_base(). 2. Add new definition of cpu_allocator() that forwards the request to ProcessState::GetCPUAllocator() when possible. The legacy cpu_allocator() definition is an artifact of a simpler era in TensorFlow design. It continues to have a use in compilation and execution contexts where either ProcessState is unavailable or a very simple need is to be met. However, its use can cause performance or correctness problems in contexts where NUMA affinity is desired or pre-registered memory for RDMA is in use. This change makes the former semantics of cpu_allocator() available via the new function cpu_allocator_base(). The new cpu_allocator() function takes an optional numa_node parameter and calls the ProcessState::GetCPUAllocator() when available (which it should always be in a real model execution, but may not be in e.g. some test contexts). Some existing uses of cpu_allocator are changed to provide a NUMA-node argument. Also fixes creation of GPUCompatibleCPUDeviceFactory for optional NUMA-aware CPU devices. Also ensures that ProcessState generates unique CPU allocators per NUMA node and does not fall back to cpu_allocator_base() when optional NUMA affinity is turned on. PiperOrigin-RevId: 225858100
Loading
Please sign in to comment