tf_1.8_xla_doc
Public Member Functions | Private Member Functions | List of all members
tensorflow::XlaCompiler Class Reference

#include <xla_compiler.h>

Public Member Functions

Status CompileGraph (const CompileOptions &options, string const &name, std::unique_ptr< Graph > graph, const std::vector< Argument > &args, CompilationResult *result)
 Compiles a tensorflow::Graph into an xla::Computation. More...
 

Private Member Functions

Status BuildArguments (const Graph &graph, const std::vector< XlaCompiler::Argument > &args, bool use_tuple_arg, xla::ComputationBuilder *builder, XlaContext *context, std::vector< int > *arg_cores, std::vector< XlaExpression > *arg_expressions, std::vector< int > *input_mapping, std::vector< xla::Shape > *input_shapes, bool is_entry_computation)
 Build XLA computations for each of the arguments to the UserComputation. More...
 

Detailed Description

The XlaCompiler class is responsible for compilation of a self-contained subgraph of a TensorFlow computation using the XLA linear algebra runtime. It does a symbolic execution of the graph starting from specific input shapes, using a JIT device to convert operators into XLA computations.

XlaCompiler is typically invoked from an _XlaLaunch operator once the shapes of all input parameters to the computation are known. This is because the symbolic execution requires known shapes for all operations.

XlaCompiler compiles Tensorflow graphs that received inputs via _Arg nodes, and return outputs via _Retval nodes.

The XlaCompiler requires one Argument struct for each _Arg index, that describes each argument. Arguments can be compile-time constants (kind kConstant), run-time parameters (kind kParameter), or resources (kind kResource).

Only kParameter and initialized kResource arguments become runtime parameters to the generated XLA computation. The XLA computation will have run-time parameters in the following order: +------------------—+--------------------------------------—+ | kParameter values | Initial values of kResource arguments | +------------------—+--------------------------------------—+ Within each block, the arguments are arranged by the _Arg index from which they were derived.

The run-time outputs of the XLA computation are arranged in the following order: +---------------—+--------------------------------------—+ | _Retval values | Updated values of kResource arguments | +---------------—+--------------------------------------—+ _Retval values are ordered by _Retval index, whereas kResource values are ordered by the original _Arg position of the variable.

In both inputs and outputs, kResource values are placed the end. When emitting While loop bodies, we must ensure that the loop body has identical input and output signatures. By moving variable values to the end of the argument list and using the return_updated_values_for_all_variables option, we can ensure that the input and output values of resources appear at the same positions.

Resources are passed as parameters or returned as resource updates in "packed" form. kStack resources are packed as (array, size of stack) XLA tuples. kTensorArray resources without gradients are packed as the array that backs the TensorArray. If gradients are present (tensor_array_gradients), the packed representation is a (array, gradient0, gradient1, ...) tuple, where gradient_k is the value of the k-th gradient in the tensor_array_gradients ordered set.

Member Function Documentation

◆ BuildArguments()

Status tensorflow::XlaCompiler::BuildArguments ( const Graph &  graph,
const std::vector< XlaCompiler::Argument > &  args,
bool  use_tuple_arg,
xla::ComputationBuilder *  builder,
XlaContext *  context,
std::vector< int > *  arg_cores,
std::vector< XlaExpression > *  arg_expressions,
std::vector< int > *  input_mapping,
std::vector< xla::Shape > *  input_shapes,
bool  is_entry_computation 
)
private

Build XLA computations for each of the arguments to the UserComputation.

  1. Traverse args (arguments to the UserComputation) to set index of non-constant arguments in input_mapping (in order of kParameter to kResource) and set up arg_expressions.
  2. Traverse input_mapping to compute the shapes of non-constant arguments and record in arg_shapes. Based on the option use_tuple_arg, it decides how to store arg_shapes to input_shapes.
  3. Use _Arg nodes in graph to resolve core assignments. The result is written in function argument arg_cores.
  4. Build parameter handles for non-constant arguments.
  5. Fill in the handles in non-constant arguments (update the arg_expressions set up in 1.).

Unclear parts:

  1. Usage of XlaResource in 1.
  2. Order of kParameter & kResource in 1.
  3. Will input_shapes be used in the future? in 2.
  4. What's sharding for? The detail of retrieving sharding info. in 3.
Here is the call graph for this function:
Here is the caller graph for this function:

◆ CompileGraph()

Status tensorflow::XlaCompiler::CompileGraph ( const CompileOptions &  options,
string const &  name,
std::unique_ptr< Graph >  graph,
const std::vector< Argument > &  args,
CompilationResult *  result 
)

Compiles a tensorflow::Graph into an xla::Computation.

  1. Check status of initialization (set during construction).
  2. Functionalize control flow in Graph in order to be compiled into XLA code.
  3. Contruct ComputationBuilder & XlaContext objects.
  4. Build argument (Build XLA computations for each of arguments to the UserComputation.)
  5. Todo:
    Rest of the code
Here is the call graph for this function:
The documentation for this class was generated from the following files:
Generated by   doxygen 1.8.14