| 590 | } |
| 591 | |
| 592 | XlaDeviceOpRegistrations* RegisterXlaDeviceKernels(const char* device, |
| 593 | const char* jit_device) { |
| 594 | // Any op assigned to the device that isn't rewritten by the graph rewriter |
| 595 | // gets executed by an XlaCompileOnDemandOp, which compiles it and executes |
| 596 | // it just-in-time. |
| 597 | OpKernel* (*factory)(OpKernelConstruction*) = |
| 598 | [](OpKernelConstruction* context) -> OpKernel* { |
| 599 | return new XlaCompileOnDemandOp(context); |
| 600 | }; |
| 601 | XlaOpRegistry::RegisterCompilationKernels(); |
| 602 | XlaDeviceOpRegistrations* registrations = new XlaDeviceOpRegistrations; |
| 603 | for (const KernelDef* jit_def : XlaOpRegistry::DeviceKernels( |
| 604 | jit_device, |
| 605 | /*include_compilation_only_kernels=*/false)) { |
| 606 | KernelDef* def = new KernelDef(*jit_def); |
| 607 | def->set_device_type(device); |
| 608 | registrations->op_kernel_registrars.emplace_back( |
| 609 | new kernel_factory::OpKernelRegistrar(def, "XlaCompileOnDemandOp", |
| 610 | factory)); |
| 611 | } |
| 612 | return registrations; |
| 613 | } |
| 614 | |
| 615 | } // namespace tensorflow |
no test coverage detected