| 61 | namespace { |
| 62 | |
| 63 | XlaPlatformInfo PlatformInfoFromContext(OpKernelConstruction* ctx) { |
| 64 | DeviceType device_type = ctx->device_type(); |
| 65 | se::Platform::Id platform_id = nullptr; |
| 66 | const XlaDevice::Metadata* xla_device_metadata = nullptr; |
| 67 | std::shared_ptr<se::DeviceMemoryAllocator> custom_allocator; |
| 68 | |
| 69 | if (ctx->device_type() == DeviceType(DEVICE_CPU)) { |
| 70 | platform_id = se::host::kHostPlatformId; |
| 71 | } else if (ctx->device_type() == DeviceType(DEVICE_GPU)) { |
| 72 | platform_id = ctx->device() |
| 73 | ->tensorflow_gpu_device_info() |
| 74 | ->stream->parent() |
| 75 | ->platform() |
| 76 | ->id(); |
| 77 | } else if (XlaDevice::GetMetadata(ctx, &xla_device_metadata).ok()) { |
| 78 | // If we are on an XlaDevice, use the underlying XLA platform's allocator |
| 79 | // directly. We could use the StreamExecutor's allocator which may |
| 80 | // theoretically be more correct, but XLA returns a nice OOM message in a |
| 81 | // Status and StreamExecutor does not. |
| 82 | // |
| 83 | // Importantly we can't use ctx->device()->GetAllocator() as the allocator |
| 84 | // (which xla_allocator above uses) as on an XlaDevice, this is a dummy |
| 85 | // allocator that returns XlaTensor objects. The XlaCompiler needs a real |
| 86 | // allocator to allocate real buffers. |
| 87 | platform_id = xla_device_metadata->platform()->id(); |
| 88 | custom_allocator = |
| 89 | xla_device_metadata->client()->backend().shared_memory_allocator(); |
| 90 | } |
| 91 | |
| 92 | return XlaPlatformInfo(device_type, platform_id, xla_device_metadata, |
| 93 | custom_allocator); |
| 94 | } |
| 95 | |
| 96 | // A closure describing how to run a compiled version of a TensorFlow function. |
| 97 | // |
no test coverage detected