| 145 | } |
| 146 | |
| 147 | void* allocate(size_t num_bytes) const override { |
| 148 | void* ret = allocator_->AllocateRaw(32 /* alignment */, num_bytes); |
| 149 | if (ret == nullptr) { |
| 150 | if (context_) { |
| 151 | context_->SetStatus(errors::ResourceExhausted( |
| 152 | strings::StrCat("Ran out of GPU memory when allocating ", num_bytes, |
| 153 | " bytes for ", operation_))); |
| 154 | } else { |
| 155 | LOG(FATAL) |
| 156 | << "EigenAllocator for GPU ran out of memory when allocating " |
| 157 | << num_bytes << ". See error logs for more detailed info."; |
| 158 | } |
| 159 | } |
| 160 | if (LogMemory::IsEnabled() && ret != nullptr) { |
| 161 | LogMemory::RecordRawAllocation(operation_, step_id_, num_bytes, ret, |
| 162 | allocator_); |
| 163 | } |
| 164 | return ret; |
| 165 | } |
| 166 | void deallocate(void* buffer) const override { |
| 167 | if (LogMemory::IsEnabled() && buffer != nullptr) { |
| 168 | LogMemory::RecordRawDeallocation(operation_, step_id_, buffer, allocator_, |
no test coverage detected