Implement lazy compilation for XLA
This CL implements lazy compilation for XLA with a very simple policy: compile a
cluster for a specific signature the second time we see the signature. This
policy is only for bootstrapping; I suspect we will have to substantially tune
this based on real world workloads.
This CL is organized as follows:
- Mark the compilation_successful output / input in the _XlaCompile/_XlaRun
kernels as HostMem. This was a bug in the CL that introduced
_XlaCompile/_XlaRun, but was unnoticed till now because we never read the
compilation_successful boolean.
* Also mark the key output / input as HostMem for symmetry even though
DT_STRING tensors are always placed on the host.
- Add a couple of new TF node matchers to make it easier to unit-test the
updated build_xla_ops pass.
- Update some tests that relied on JIT compilation to have a "warmup" phase so
that compilation is guaranteed.
- The main part of the CL is in kernels/xla_ops and jit/xla_compilation_cache.
PiperOrigin-RevId: 216786799
Loading
Please sign in to comment