Commit cb8f301d authored by Sanjoy Das's avatar Sanjoy Das Committed by TensorFlower Gardener
Browse files

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
parent 3b5d3466
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment