| 211 | platform_info_(PlatformInfoFromContext(ctx)) {} |
| 212 | |
| 213 | static Status BuildCompilationCache(OpKernelContext* ctx, |
| 214 | const XlaPlatformInfo& platform_info, |
| 215 | XlaCompilationCache** cache) { |
| 216 | if (platform_info.xla_device_metadata()) { |
| 217 | *cache = new XlaCompilationCache( |
| 218 | platform_info.xla_device_metadata()->client(), |
| 219 | platform_info.xla_device_metadata()->jit_device_type()); |
| 220 | return Status::OK(); |
| 221 | } |
| 222 | |
| 223 | auto platform = |
| 224 | se::MultiPlatformManager::PlatformWithId(platform_info.platform_id()); |
| 225 | if (!platform.ok()) { |
| 226 | return platform.status(); |
| 227 | } |
| 228 | |
| 229 | xla::StatusOr<xla::Compiler*> compiler_for_platform = |
| 230 | xla::Compiler::GetForPlatform(platform.ValueOrDie()); |
| 231 | if (!compiler_for_platform.ok()) { |
| 232 | // In some rare cases (usually in unit tests with very small clusters) we |
| 233 | // may end up transforming an XLA cluster with at least one GPU operation |
| 234 | // (which would normally force the cluster to be compiled using XLA:GPU) |
| 235 | // into an XLA cluster with no GPU operations (i.e. containing only CPU |
| 236 | // operations). Such a cluster can fail compilation (in way that |
| 237 | // MarkForCompilation could not have detected) if the CPU JIT is not linked |
| 238 | // in. |
| 239 | // |
| 240 | // So bail out of _XlaCompile in this case, and let the executor handle the |
| 241 | // situation for us. |
| 242 | const Status& status = compiler_for_platform.status(); |
| 243 | if (status.code() == error::NOT_FOUND) { |
| 244 | return errors::Unimplemented("Could not find compiler for platform ", |
| 245 | platform.ValueOrDie()->Name(), ": ", |
| 246 | status.ToString()); |
| 247 | } |
| 248 | } |
| 249 | |
| 250 | xla::LocalClientOptions client_options; |
| 251 | client_options.set_platform(platform.ValueOrDie()); |
| 252 | client_options.set_intra_op_parallelism_threads( |
| 253 | ctx->device()->tensorflow_cpu_worker_threads()->num_threads); |
| 254 | auto client = xla::ClientLibrary::GetOrCreateLocalClient(client_options); |
| 255 | if (!client.ok()) { |
| 256 | return client.status(); |
| 257 | } |
| 258 | const XlaOpRegistry::DeviceRegistration* registration; |
| 259 | if (!XlaOpRegistry::GetCompilationDevice(platform_info.device_type().type(), |
| 260 | ®istration)) { |
| 261 | return errors::InvalidArgument("No JIT device registered for ", |
| 262 | platform_info.device_type().type()); |
| 263 | } |
| 264 | *cache = new XlaCompilationCache( |
| 265 | client.ValueOrDie(), DeviceType(registration->compilation_device_name)); |
| 266 | return Status::OK(); |
| 267 | } |
| 268 | |
| 269 | static Status CompileToLocalExecutable( |
| 270 | OpKernelContext* ctx, const NameAttrList& function, |
no test coverage detected