Commit 7dea8383 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower Committed by TensorFlower Gardener
Browse files

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
parent 86a9610e
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment