Commit 1fb724f2 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower Committed by TensorFlower Gardener
Browse files

Normally tf2xla (autoclustering, jit_scope and rewrite) rely on graph optimization

passes to outline subgraphs. The XLA device itself only sees Compute() calls for
_XlaLaunch ops. All other ops are registered with a dummy op factory that just
prints an error.

This patch adds an alternative, selected at registration time, that disables
default graph optimization and instead registers a non-dummy op implementation.

This op implementation compiles the op "on demand"; it generates a fake graph containing
_Arg and _Retval nodes and calls into the XlaCompiler code as usual.

This allows the device to be used as a "normal" TensorFlow device, as well as from
Eager mode, at the expense of performance.

Later additions will add the ability to create traces to amortize kernel launch overhead,
and the ability to combine op-by-op/tracing and autoclustering with jit_scope annotations.

PiperOrigin-RevId: 189463593
parent 6a0b4e17
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment