[XLA] Exposes the IsValueAllowedInAlternateMemory() for better reuse.

PiperOrigin-RevId: 327304648
Change-Id: I849662031df6bdadd9afb5d6d530df5c193f853b
This commit is contained in:
A. Unique TensorFlower 2020-08-18 14:14:51 -07:00 committed by TensorFlower Gardener
parent 62bf64d77e
commit 22fa4416cf
2 changed files with 5 additions and 5 deletions

View File

@ -17,9 +17,8 @@ limitations under the License.
namespace xla {
namespace {
bool IsValueAllowedInAlternateMemory(const HloValue* value) {
bool MemorySpaceAssignmentUtils::IsValueAllowedInAlternateMemory(
const HloValue* value) {
// If the buffer is a tuple, don't use this algorithm for now. The buffers
// that are pointed to by the tuple will still use this algorithm. Because
// tuples are cheap to place in the alternate memory (they are just pointers)
@ -93,8 +92,6 @@ bool IsValueAllowedInAlternateMemory(const HloValue* value) {
return true;
}
} // namespace
bool MemorySpaceAssignmentUtils::IsIntervalAllowedInAlternateMemory(
const GlobalDecreasingSizeBestFitHeap<HloValue>::BufferInterval& interval) {
return IsValueAllowedInAlternateMemory(interval.buffer) &&

View File

@ -28,6 +28,9 @@ class MemorySpaceAssignmentUtils {
static bool IsIntervalAllowedInAlternateMemory(
const GlobalDecreasingSizeBestFitHeap<HloValue>::BufferInterval&
interval);
// Returns true if the HloValue is allowed to be placed in alternate memory.
static bool IsValueAllowedInAlternateMemory(const HloValue* value);
};
} // namespace xla