Function GetAllocator

tensorflow/compiler/jit/kernels/xla_ops.cc:183–199

Return allocator from platform info if non-null, or populate and return a pointer to the allocator adapter with allocator from context. This is necessary because for XLA devices the underlying TF allocator returns dummy tensors.

// This is necessary because for XLA devices the underlying TF allocator returns
// dummy tensors.
static std::shared_ptr<se::DeviceMemoryAllocator> GetAllocator(
    OpKernelContext* ctx, const XlaPlatformInfo& platform_info) {
  if (platform_info.custom_allocator()) {
    return platform_info.custom_allocator();
  }
  auto *alloc = ctx->device()->GetAllocator({});
  auto dc = ctx->op_device_context();
  if (!dc) {
    // Stream is not set for the host platform.
    se::Platform* platform =
        se::MultiPlatformManager::PlatformWithId(platform_info.platform_id())
            .ValueOrDie();
    return std::make_shared<se::TfAllocatorAdapter>(alloc, platform);
  } else {
    return std::make_shared<se::TfAllocatorAdapter>(alloc, dc->stream());
  }
}

}  // namespace