| 317 | : memory_limit_(memory_limit), total_byte_size_(0), context_(context) {} |
| 318 | int64 GetMemoryLimitInBytes() override { return memory_limit_; } |
| 319 | se::port::StatusOr<se::DeviceMemory<uint8>> AllocateBytes( |
| 320 | int64 byte_size) override { |
| 321 | Tensor temporary_memory; |
| 322 | if (byte_size > memory_limit_) { |
| 323 | return se::port::StatusOr<se::DeviceMemory<uint8>>(); |
| 324 | } |
| 325 | AllocationAttributes allocation_attr; |
| 326 | allocation_attr.no_retry_on_failure = true; |
| 327 | Status allocation_status(context_->allocate_temp( |
| 328 | DT_UINT8, TensorShape({byte_size}), &temporary_memory, |
| 329 | AllocatorAttributes(), allocation_attr)); |
| 330 | if (!allocation_status.ok()) { |
| 331 | return se::port::StatusOr<se::DeviceMemory<uint8>>(); |
| 332 | } |
| 333 | // Hold the reference of the allocated tensors until the end of the |
| 334 | // allocator. |
| 335 | allocated_tensors_.push_back(temporary_memory); |
| 336 | total_byte_size_ += byte_size; |
| 337 | return se::port::StatusOr<se::DeviceMemory<uint8>>( |
| 338 | AsDeviceMemory(temporary_memory.flat<uint8>().data(), |
| 339 | temporary_memory.flat<uint8>().size())); |
| 340 | } |
| 341 | int64 TotalByteSize() { return total_byte_size_; } |
| 342 | |
| 343 | private: |
no test coverage detected