Reduce overhead of constructing an `OpKernelContext` on non-GPU devices.
The current `OpKernelContext` constructor unconditionally makes a virtual call to get an allocator that is only used when the kernel is running on a GPU device. Since we can tell if this is necessary after the call to `params_->ensure_eigen_gpu_device()`, we can avoid the call, plus the virtual call to `ReinitializeGpuDevice()`, when we know that it is not a GPU device. PiperOrigin-RevId: 203523500
Loading
Please sign in to comment