[XLA] Exposes the IsValueAllowedInAlternateMemory() for better reuse.
PiperOrigin-RevId: 327304648 Change-Id: I849662031df6bdadd9afb5d6d530df5c193f853b
This commit is contained in:
parent
62bf64d77e
commit
22fa4416cf
@ -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) &&
|
||||
|
@ -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
|
||||
|
Loading…
x
Reference in New Issue
Block a user