| 211 | } |
| 212 | |
| 213 | StatusOr<std::shared_ptr<Device>> DeviceForDLContext( |
| 214 | const PyLocalClient& client, const DLContext& context) { |
| 215 | se::Platform::Id platform_id; |
| 216 | switch (context.device_type) { |
| 217 | case kDLCPU: |
| 218 | platform_id = se::host::kHostPlatformId; |
| 219 | break; |
| 220 | case kDLGPU: |
| 221 | platform_id = se::cuda::kCudaPlatformId; |
| 222 | break; |
| 223 | default: |
| 224 | return InvalidArgument("Unknown/unsupported DLPack device type %d", |
| 225 | context.device_type); |
| 226 | } |
| 227 | auto it = absl::c_find_if( |
| 228 | client.local_devices(), [&](const std::shared_ptr<Device>& device) { |
| 229 | return device->local_device_state()->executor()->platform()->id() == |
| 230 | platform_id && |
| 231 | device->local_device_state()->device_ordinal() == |
| 232 | context.device_id; |
| 233 | }); |
| 234 | if (it == client.local_devices().end()) { |
| 235 | return InvalidArgument( |
| 236 | "No matching device found for DLPack device_type %d device_id %d", |
| 237 | context.device_type, context.device_id); |
| 238 | } |
| 239 | return *it; |
| 240 | } |
| 241 | |
| 242 | } // namespace |
| 243 |
no test coverage detected