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
Loading
Please sign in to comment