| 153 | } |
| 154 | |
| 155 | void ScopedShapedBuffer::Deallocate() { |
| 156 | // allocator_ will be null if we were moved-from. |
| 157 | if (allocator_ == nullptr) { |
| 158 | return; |
| 159 | } |
| 160 | // Deallocate all non-null buffers. A buffer may appear in more than one spot |
| 161 | // in the shape (eg, a tuple with a repeated element) so keep track of what |
| 162 | // has been deallocated. |
| 163 | absl::flat_hash_set<void*> deallocated_ptrs; |
| 164 | for (auto& pair : buffers_) { |
| 165 | se::DeviceMemoryBase& memory_base = pair.second; |
| 166 | if (!memory_base.is_null() && |
| 167 | deallocated_ptrs.insert(memory_base.opaque()).second) { |
| 168 | TF_CHECK_OK(allocator_->Deallocate(device_ordinal(), memory_base)); |
| 169 | } |
| 170 | } |
| 171 | } |
| 172 | |
| 173 | ScopedShapedBuffer ScopedShapedBuffer::TakeSubTree(ShapeIndexView index) { |
| 174 | const xla::Shape& sub_on_host_shape = |
no test coverage detected