From: Justin Lebar <jlebar@google.com>
Date: Wed, 9 May 2018 01:16:47 +0000 (-0700)
Subject: [XLA] Make XlaAllocator obey retry_on_failure arg.
X-Git-Tag: upstream/v1.9.0_rc1~116^2^2~217
X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=d8cc88a19d8a8c61023c34395cce55593a498cbf;p=platform%2Fupstream%2Ftensorflow.git

[XLA] Make XlaAllocator obey retry_on_failure arg.

Previously we ignored it.

PiperOrigin-RevId: 195908178
---

diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc
index 0223f97..e12e88f 100644
--- a/tensorflow/compiler/jit/xla_launch_util.cc
+++ b/tensorflow/compiler/jit/xla_launch_util.cc
@@ -62,7 +62,10 @@ XlaAllocator::~XlaAllocator() {}
 
 xla::StatusOr<se::DeviceMemoryBase> XlaAllocator::Allocate(
     int device_ordinal, uint64 size, bool retry_on_failure) {
-  void* data = wrapped_->AllocateRaw(Allocator::kAllocatorAlignment, size);
+  AllocationAttributes attrs;
+  attrs.no_retry_on_failure = !retry_on_failure;
+  void* data =
+      wrapped_->AllocateRaw(Allocator::kAllocatorAlignment, size, attrs);
   if (data == nullptr) {
     return errors::ResourceExhausted("Out of memory while trying to allocate ",
                                      size, " bytes.");