| 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_, |
| 169 | true); |
| 170 | } |
| 171 | AsyncFreeData* afData = |
| 172 | new AsyncFreeData(allocator_, buffer, operation_, step_id_); |
| 173 | #if GOOGLE_CUDA |
| 174 | cudaError_t err = cudaStreamAddCallback(*stream_, asyncFree, afData, 0); |
| 175 | CHECK_EQ(err, cudaSuccess); |
| 176 | #elif TENSORFLOW_USE_ROCM |
| 177 | hipError_t err = hipStreamAddCallback(*stream_, asyncFree, afData, 0); |
| 178 | CHECK_EQ(err, hipSuccess); |
| 179 | #endif |
| 180 | } |
| 181 | |
| 182 | // Return a pointer to a per stream scratchpad of 1024 bytes residing |
| 183 | // in global memory. |
no test coverage detected