| 175 | } // namespace |
| 176 | |
| 177 | ScopedActivateContext::ScopedActivateContext(GpuContext* context) { |
| 178 | if (FLAGS_gpuexec_rocm_sync_around_driver_calls) { |
| 179 | SynchronizeOrDie(); |
| 180 | } |
| 181 | |
| 182 | auto* tls = &tls_data.get(); |
| 183 | if (tls->depth == 0) { |
| 184 | tls->current_device_ordinal = CurrentDeviceOrDie(); |
| 185 | } |
| 186 | |
| 187 | if (kVerifyGpuContext) { |
| 188 | CHECK_EQ(CurrentDeviceOrDie(), tls->current_device_ordinal); |
| 189 | } |
| 190 | |
| 191 | tls->depth++; |
| 192 | |
| 193 | to_restore_ = context; |
| 194 | |
| 195 | if (context->device_ordinal() == tls->current_device_ordinal) { |
| 196 | DCHECK_EQ(CurrentDeviceOrDie(), context->device_ordinal()); |
| 197 | return; |
| 198 | } |
| 199 | |
| 200 | VLOG(3) << "ScopedActivateContext switching device from " |
| 201 | << tls->current_device_ordinal << " to " << context->device_ordinal(); |
| 202 | |
| 203 | // Set the device and update thread local. |
| 204 | CHECK_EQ(hipSuccess, |
| 205 | tensorflow::wrap::hipSetDevice(context->device_ordinal())); |
| 206 | tls->current_device_ordinal = context->device_ordinal(); |
| 207 | } |
| 208 | |
| 209 | ScopedActivateContext::~ScopedActivateContext() { |
| 210 | if (FLAGS_gpuexec_rocm_sync_around_driver_calls) { |
nothing calls this directly
no test coverage detected