package xla
Get desktop application:
View/edit binary Protocol Buffers messages
rpc Connect (, )
protocol.proto:126Connects a node to the distributed coordinator node. Blocks until all tasks have connected. The service receives the number of nodes to expect as an option passed to its constructor.
message
protocol.proto:56int32 protocol_version = 1
int32 timeout_milliseconds = 2
int32 node_id = 3
We assume that each node knows its globally-unique node ID, provided by whatever mechanism launches the tasks. Node IDs should form a dense range of integers [0, num_nodes).
uint64 client_id = 4
A unique ID number for the client.
message
protocol.proto:69uint64 session_id = 1
rpc EnumerateDevices (, )
protocol.proto:131Blocking enumeration of devices, used by the GPU backend only. In parallel, all clients call EnumerateDevices() with their local device topology, and receive back a global topology in response.
message
protocol.proto:73uint64 session_id = 1
optional local_topology = 3
message
protocol.proto:78optional global_topology = 1
rpc Heartbeat (, )
protocol.proto:137Health-checking RPC. Workers send heartbeats to the coordinator at regular intervals. If the worker does not hear from the coordinator or the coordinator does not hear from the tasks, the tasks abort.
message
protocol.proto:110uint64 session_id = 1
int32 node_id = 2
message
protocol.proto:114(message has no fields)
rpc KeyValueGet (, )
protocol.proto:150Looks up a key in the key-value service. Blocks until the key is present or until `timeout` expires.
message
protocol.proto:82uint64 session_id = 1
bytes key = 2
int32 timeout_milliseconds = 3
message
protocol.proto:88bool found = 1
bytes value = 2
rpc KeyValueSet (, )
protocol.proto:153Updates the value associated with a key.
message
protocol.proto:93uint64 session_id = 1
bytes key = 2
bytes value = 3
message
protocol.proto:99(message has no fields)
rpc Shutdown (, )
protocol.proto:142Shutdown RPC. Workers send this RPC when they are ready to shut down; the RPC blocks until all tasks have indicated they are ready to shut down, or a timeout is reached.
message
protocol.proto:116uint64 session_id = 1
int32 node_id = 2
message
protocol.proto:120(message has no fields)
rpc WaitAtBarrier (, )
protocol.proto:156Blocks until all nodes are at the barrier or the barrier times out.
message
protocol.proto:101uint64 session_id = 1
bytes barrier_id = 2
int32 node_id = 3
int32 timeout_milliseconds = 4
message
protocol.proto:108(message has no fields)
service XlaService
xla_service.proto:49/////////////////////// Global data requests
rpc Compile (, )
xla_service.proto:132Compiles the provided computation into executable. Returns the handle of the executable.
message
xla.proto:594optional computation = 1
The graph to be compiled.
optional execution_options = 2
Options that affect how XLA compiles code to service this request.
repeated input_shape_with_layout = 3
The layouts of the input arguments. If not set, the default layout will be used. Although the real arguments are not needed in compilation, the layouts of the arguments can affect the compilation.
message
xla.proto:607optional handle = 1
The handle to the executable.
Computes the value of a constant expression. The request contains the computation graph for the constant expression.
message
xla.proto:652optional computation = 1
optional output_layout = 2
message
xla.proto:657optional literal = 1
A LiteralProto is returned directly for this request.
Creates a channel handle that can be used to transfer data between two computations via a pair of Send and Recv instructions.
message
xla.proto:580ChannelHandle.ChannelType channel_type = 1
message
xla.proto:584optional channel = 1
rpc DeconstructTuple (, )
xla_service.proto:62Deconstructs a tuple. Returns a newly created GlobalDataHandle for each element in the tuple.
message
xla.proto:662optional tuple_handle = 2
message
xla.proto:666repeated element_handles = 1
rpc Execute (, )
xla_service.proto:137Invokes the provided executable with the provided global data passed as immutable arguments. The request contains the handle to the executable. Returns global data output and execution timing.
message
xla.proto:612optional handle = 1
repeated arguments = 2
The shape and layout of the arguments must be the same as the those of the executable's parameters.
Invokes the provided list of computations in parallel with the provided global data for each computation. Returns a list of global data output and execution timing.
message
xla.proto:630repeated requests = 1
message
xla.proto:639repeated responses = 1
Requests the statistics of the given computation.
message
xla.proto:571optional computation = 1
optional debug_options = 2
message
xla.proto:576optional stats = 1
Requests one or more device handles from the target. The returned device handles can be used to specify the device on which to execute computations or transfer data.
message
xla.proto:515int64 device_count = 1
message
xla.proto:519repeated device_handles = 1
rpc GetShape (, )
xla_service.proto:74Requests the shape of the referenced global data.
message
xla.proto:700optional data = 1
message
xla.proto:704optional shape = 1
rpc LoadData (, )
xla_service.proto:83Loads a variable number of values with a given element type from ColumnIO.
message
xla.proto:670string columnio_tablet_path = 1
Describes the path of the ColumnIO tablet to load.
string columnio_field = 2
Describes the field to load within the ColumnIO tablet.
optional element_shape = 3
Individual element shape, excluding rows.
int64 offset = 4
Warning: ColumnIO does not support random-access, so use offset with caution in performance-critical scenarios.
int64 limit = 5
Maximum number of elements (with shape element_shape) to load.
bool zip = 6
If more than one item is requested (via limit > 1), then this request attribute zips together the produced vectors.
message
xla.proto:692optional data = 1
optional data_shape = 2
int64 available_rows = 3
int64 rows_loaded = 4
int64 nanoseconds = 5
rpc ResetDevice (, )
xla_service.proto:108Resets the device, clearing all existing state on the device.
message
xla.proto:565optional device_handle = 1
message
xla.proto:569(message has no fields)
Transferred literal from the Outfeed buffer of the device.
message
xla.proto:552optional shape_with_layout = 1
This optional field directs the service to return the literal in this layout. A shape is used to hold the layout to accommodate tuples.
int64 replica_id = 2
optional device_handle = 3
message
xla.proto:561optional literal = 1
rpc TransferToClient (, )
xla_service.proto:87Transfers the given global data to the client in the form of a Literal.
message
xla.proto:523optional data = 1
optional shape_with_layout = 2
This optional field directs the service to return the literal in this layout. A shape is used to hold the layout to accommodate tuples.
message
xla.proto:531optional literal = 1
rpc TransferToInfeed (, )
xla_service.proto:98Transfers the given literal to the Infeed buffer of the device.
message
xla.proto:544optional literal = 1
int64 replica_id = 2
optional device_handle = 3
message
xla.proto:550(message has no fields)
rpc TransferToServer (, )
xla_service.proto:93Transfers the given literal to the server to be stored in a global allocation, which is returned.
message
xla.proto:535optional literal = 1
optional device_handle = 2
message
xla.proto:540optional data = 1
rpc Unpack (, )
xla_service.proto:70Unpack requests that a global data handle, with a tuple shape, has global data handles created for each of its constituent members. This is the equivalent of the "destructuring assignment" present in various programming languages.
message
xla.proto:708optional data = 1
message
xla.proto:712repeated tied_data = 1
rpc Unregister (, )
xla_service.proto:57Unregisters a global allocation. If the handle given is not currently allocated, a NOT_FOUND status is returned.
message
xla.proto:588repeated data = 1
message
xla.proto:592(message has no fields)
Waits until the given execution (aysnchronously launched) is complete, and returns the global data output.
message
xla.proto:643optional execution = 1
message
xla.proto:647optional output = 1
optional profile = 2
message
hlo.proto:544Serialization of BufferAllocation.
Used in:
int64 index = 1
int64 size = 2
bool is_thread_local = 3
bool is_tuple = 11
bool is_entry_computation_parameter = 5
bool is_constant = 12
int64 parameter_number = 6
repeated int64 parameter_shape_index = 10
bool maybe_live_out = 7
int64 color = 8
repeated assigned = 9
message
hlo.proto:547Assigned represents a single LogicalBuffer that is assigned to this BufferAllocation.
Used in:
int64 logical_buffer_id = 1
int64 offset = 2
int64 size = 3
message
hlo.proto:610Serialization of BufferAssignment.
Used in: ,
repeated logical_buffers = 1
repeated buffer_aliases = 2
repeated buffer_allocations = 3
repeated heap_simulator_traces = 4
message
hlo.proto:613Alias represents a source LogicalBuffer, and the buffer location that aliases it.
Used in:
int64 source_buffer_id = 1
optional location = 2
message
xla_data.proto:404Handle given to a user to represent a channel between two computations via a Send and Recv instruction pair. Channels are unbuffered, so Send Send instructions will be blocked until the data is transferred.
Used in:
int64 handle = 1
ChannelHandle.ChannelType type = 2
Used in: ,
CHANNEL_TYPE_INVALID = 0
Invalid primitive type to serve as default.
DEVICE_TO_DEVICE = 1
A channel for sending data between devices.
DEVICE_TO_HOST = 2
A channel for sending data from the device to the host. Can only be used with a Send operation.
HOST_TO_DEVICE = 3
A channel for sending data from the host to the device. Can only be used with a Recv operation.
message
xla_data.proto:667Used in:
bool lower = 1
If true, uses the lower triangle of `a`. If false, uses the upper triangle of `a`.
The compilation event that triggered the use of the profile.
Used in: ,
COMPILATION_EVENT_UNKNOWN_EVENT = 0
COMPILATION_EVENT_FIRST_COMPILATION = 1
COMPILATION_EVENT_RECOMPILATION = 2
message
compile_options.proto:73repeated argument_layouts = 1
Refer CompileOptions for documentation of fields.
bool parameter_is_tupled_arguments = 2
optional executable_build_options = 3
bool compile_portable_executable = 4
int64 profile_version = 5
bytes serialized_multi_slice_config = 6
message
xla_data.proto:252Statistics of a computation.
Used in:
double flop_count = 1
The number of floating point operations in the computation.
double transcendental_count = 2
The number of transcendental operations (e.g., exp) in the computation.
message
xla_data.proto:564Used in:
int64 input_batch_dimension = 7
The number of the dimension that represents batch in the input.
int64 input_feature_dimension = 8
The number of the dimension that represents features in the input.
repeated int64 input_spatial_dimensions = 11
The dimension numbers for the spatial dimensions that the window moves through in the input.
int64 kernel_input_feature_dimension = 3
The number of the dimension that represents input features in the convolutional kernel (rhs).
int64 kernel_output_feature_dimension = 4
The number of the dimension that represents output features in the convolutional kernel (rhs).
repeated int64 kernel_spatial_dimensions = 6
The dimension numbers for the spatial dimensions that the window moves through in the kernel (rhs). window.strides(0) is the stride in the kernel_spatial_dimensions(0) dimension.
int64 output_batch_dimension = 9
The number of the dimension that represents batch in the output.
int64 output_feature_dimension = 10
The number of the dimension that represents features in the output.
repeated int64 output_spatial_dimensions = 12
The dimension numbers for the spatial dimensions that the window moves through in the output.
message
hlo.proto:451Used in:
int64 parameter = 1
repeated int64 index = 2
The version of the API used by the custom call function. The signatures for each version are given below. TODO(b/189822916): Remove this enum when all clients are migrated to the status-returning API.
Used in:
API_VERSION_UNSPECIFIED = 0
API_VERSION_ORIGINAL = 1
The first version of the API, with the following signatures: CPU: void do_custom_call(void* out, const void** in); GPU: void do_custom_call(CUstream stream, void** buffers, const char* opaque, size_t opaque_len);
API_VERSION_STATUS_RETURNING = 2
When the ability to return success/failure status was added: CPU: void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status); GPU: void do_custom_call(CUstream stream, void** buffers, const char* opaque, size_t opaque_len, XlaCustomCallStatus* status);
API_VERSION_STATUS_RETURNING_UNIFIED = 3
Fixes the API signatures on the CPU side of the version STATUS_RETURNING by adding the opaque string so that the custom call API is consistent across CPUs and GPUs. For GPUs, the behaviors invoked by API_VERSION_STATUS_RETURNING and API_VERSION_STATUS_RETURNING_UNIFIED are the same. CPU: void do_custom_call(void* out, const void** in, const char* opaque, size_t opaque_len, XlaCustomCallStatus* status); GPU: void do_custom_call(CUstream stream, void** buffers, const char* opaque, size_t opaque_len, XlaCustomCallStatus* status);
message
xla_data.proto:795Specifies a pair of output/operand buffers for kCustomCall that alias each other.
Used in:
repeated int64 output_shape_index = 1
int64 operand_index = 2
repeated int64 operand_shape_index = 3
Used in:
SCHEDULE_NONE = 0
SCHEDULE_LATEST = 1
SCHEDULE_EARLIEST = 2
message
xla.proto:25Debugging options for XLA. These options may change at any time - there are no guarantees about backward or forward compatibility for these fields.
Used in: , , , ,
bool xla_hlo_graph_addresses = 2
Show addresses of HLO ops in graph dump.
bool xla_hlo_profile = 9
Instrument the computation to collect per-HLO cycle counts.
repeated string xla_disable_hlo_passes = 30
List of HLO passes to disable/enable. These names must exactly match the pass names as specified by the HloPassInterface::name() method. At least one of xla_disable_hlo_passes and xla_enable_hlo_passes_only must be empty.
repeated string xla_enable_hlo_passes_only = 124
bool xla_disable_all_hlo_passes = 104
Disables all HLO passes. Notes that some passes are necessary for correctness and the invariants that must be satisfied by "fully optimized" HLO are different for different devices and may change over time. The only "guarantee", such as it is, is that if you compile XLA and dump the optimized HLO for some graph, you should be able to run it again on the same device with the same build of XLA.
int32 xla_backend_optimization_level = 31
Numerical optimization level for the XLA compiler backend; the specific interpretation of this value is left to the backends.
bool xla_embed_ir_in_executable = 33
Embed the compiler IR as a string in the executable.
bool xla_eliminate_hlo_implicit_broadcast = 35
Eliminate implicit broadcasts when lowering user computations to HLO instructions; use explicit broadcast instead.
bool xla_cpu_multi_thread_eigen = 60
When generating calls to Eigen in the CPU backend, use multi-threaded Eigen mode.
string xla_gpu_cuda_data_dir = 61
Path to directory with cuda/ptx tools and libraries.
bool xla_gpu_ftz = 62
Enable flush-to-zero semantics in the GPU backend.
bool xla_llvm_enable_alias_scope_metadata = 70
If true, in LLVM-based backends, emit !alias.scope metadata in generated IR.
bool xla_llvm_enable_noalias_metadata = 71
If true, in LLVM-based backends, emit !noalias metadata in the generated IR.
bool xla_llvm_enable_invariant_load_metadata = 72
If true, in LLVM-based backends, emit !invariant.load metadata in the generated IR.
bool xla_llvm_disable_expensive_passes = 73
If true, a set of expensive LLVM optimization passes will not be run.
bool xla_test_all_output_layouts = 90
This is used by ClientLibraryTestBase::ComputeAndCompare*. If true, the computation will run n! times with all permunations of layouts for the output shape in rank n. For example, with a 3D shape, all permutations of the set {0, 1, 2} are tried.
bool xla_test_all_input_layouts = 91
This is used by ClientLibraryTestBase::ComputeAndCompare*. If true, the computation will run for all permunations of layouts of all input arguments. For example, with 2 input arguments in 2D and 4D shapes, the computation will run 2! * 4! times.
bool xla_hlo_graph_sharding_color = 92
Assign colors based on sharding information when generating the Graphviz HLO graph.
bool xla_cpu_use_mkl_dnn = 97
Generate calls to MKL-DNN in the CPU backend.
bool xla_cpu_use_xla_runtime = 177
Enable XLA Runtime in the CPU backend.
int32 xla_gpu_max_kernel_unroll_factor = 98
Maximum kernel unroll factor for the GPU backend.
bool xla_cpu_enable_fast_math = 99
When true, "unsafe" mathematical optimizations are enabled. These transformations include but are not limited to: - Reducing the precision of operations (e.g. using an approximate sin function, or transforming x/y into x * (1/y)). - Assuming that operations never produce or consume NaN or +/- Inf (this behavior can be adjusted using xla_cpu_fast_math_allow_{nans|infs}). - Assuming that +0 and -0 are indistinguishable.
bool xla_cpu_fast_math_honor_nans = 120
When xla_cpu_enable_fast_math is true then this controls whether we allow operations to produce NaNs. Ignored when xla_cpu_enable_fast_math is false.
bool xla_cpu_fast_math_honor_infs = 121
When xla_cpu_enable_fast_math is true then this controls whether we allow operations to produce infinites. Ignored when xla_cpu_enable_fast_math is false.
bool xla_cpu_fast_math_honor_division = 126
When xla_cpu_enable_fast_math is true then this controls whether we forbid to use the reciprocal of an argument instead of division. Ignored when xla_cpu_enable_fast_math is false.
bool xla_cpu_fast_math_honor_functions = 129
When xla_cpu_enable_fast_math is true then this controls whether we forbid to approximate calculations for functions. Ignored when xla_cpu_enable_fast_math is false.
bool xla_cpu_enable_fast_min_max = 140
When false we lower the Minimum and Maximum hlos in the CPU backend such that Min(NotNaN, NaN) = Min(NaN, NotNaN) = NaN. In other words, if flag this is false we always propagate NaNs through Min and Max. Note, this does not correspond to the exact same behavior as the gpu flag below!
bool xla_gpu_enable_fast_min_max = 100
When true we lower the Minimum and Maximum hlos in the GPU backend such that Min(NotNaN, NaN) = Min(NaN, NotNaN) = NotNaN. In other words, if flag this is true we don't propagate NaNs through Min and Max. Note, this does not correspond to the exact same behavior as the cpu flag above!
bool xla_allow_excess_precision = 122
Allows xla to increase the output precision of floating point operations.
bool xla_gpu_crash_on_verification_failures = 101
Crashes the program when any kind of verification fails, instead of just logging the failures. One example is cross checking of convolution results among different algorithms.
int32 xla_gpu_autotune_level = 123
0: Disable gemm and convolution autotuning. 1: Enable autotuning, but disable correctness checking. 2: Also set output buffers to random numbers during autotuning. 3: Also reset output buffers to random numbers after autotuning each algorithm. 4+: Also check for correct outputs and for out-of-bounds reads/writes. Default: 4.
int32 xla_force_host_platform_device_count = 102
Force the host platform to pretend that there are these many host "devices". All these devices are backed by the same threadpool. Defaults to 1. Setting this to anything other than 1 can increase overhead from context switching but we let the user override this behavior to help run tests on the host that run models in parallel across multiple devices.
bool xla_gpu_disable_gpuasm_optimizations = 103
If set to true XLA:GPU invokes `ptxas` with -O0 (default is -O3).
DebugOptions.ShapeChecks xla_gpu_shape_checks = 170
bool xla_cpu_enable_mlir_lowering = 171
Enable MLIR-based lowering in XLA:CPU instead of LLVM emitters.
bool xla_gpu_enable_mlir_lowering = 173
If true, use MLIR instead of IR emitter to generate device code for supported lmhlo.fusion ops. See xla::gpu::RewriteFusionOps() for details.
bool xla_hlo_evaluator_use_fast_path = 106
Enable fast math with eigen in the HLO evaluator.
bool xla_allow_scalar_index_dynamic_ops = 107
Temporary option to allow support for both the R1 and the scalar index versions of DynamicSlice and DynamicUpdateSlice. Only used for testing.
DebugOptions.StepMarkerLocation xla_step_marker_location = 108
Option to emit a target-specific marker to indicate the start of a training step. The location of the marker (if any) is determined by the option value.
string xla_dump_to = 109
Directory to dump into.
string xla_dump_hlo_module_re = 110
If specified, will only dump modules which match this regexp.
string xla_dump_hlo_pass_re = 111
If this flag is specified, will also dump HLO before and after passes that match this regular expression. Set to .* to dump before/after all passes.
bool xla_dump_hlo_as_text = 112
Specifies the format that HLO is dumped in. Multiple of these may be specified.
bool xla_dump_hlo_as_proto = 113
bool xla_dump_hlo_as_dot = 114
bool xla_dump_hlo_as_url = 115
bool xla_dump_hlo_as_html = 116
Dump HLO graphs as an HTML (DOT -> SVG inlined in HTML)
bool xla_dump_fusion_visualization = 149
Dump the visualization of the fusion progress.
bool xla_dump_hlo_snapshots = 118
If true, every time an HLO module is run, we will dump an HloSnapshot (essentially, a serialized module plus its inputs) to the --xla_dump_to directory.
bool xla_dump_include_timestamp = 131
Include a timestamp in the dumped filenames.
int32 xla_dump_max_hlo_modules = 132
Max number of hlo module dumps in a directory. Set to < 0 for unbounded.
bool xla_dump_module_metadata = 144
Dump HloModuleMetadata as a text proto for each HLO module.
bool xla_dump_compress_protos = 151
GZip-compress protos dumped via --xla_dump_hlo_as_proto.
bool xla_dump_hlo_as_long_text = 164
Dump HLO in long text format. Ignored unless xla_dump_hlo_as_text is true.
bool xla_gpu_force_conv_nchw = 125
Overrides for XLA GPU's convolution layout heuristic.
bool xla_gpu_force_conv_nhwc = 146
repeated string xla_gpu_ptx_file = 127
Paths to files with ptx code.
bool xla_gpu_dump_llvmir = 155
Whether to dump llvm ir when compiling to ptx.
string xla_gpu_algorithm_denylist_path = 128
Denylist for cuDNN convolutions.
bool xla_tpu_detect_nan = 135
Debug options that trigger execution errors when NaN or Inf are detected.
bool xla_tpu_detect_inf = 136
bool xla_cpu_enable_xprof_traceme = 137
True if TraceMe annotations are enabled for XLA:CPU.
bool xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found = 138
It is usually preferable to not fallback to the driver; it can consume more memory, or have bugs.
string xla_gpu_asm_extra_flags = 141
Extra parameters to pass the GPU assembler.
int32 xla_multiheap_size_constraint_per_heap = 142
Per-heap size constraint. New heaps will be created if per-heap max size is reached.
bool xla_detailed_logging_and_dumping = 143
Enable detailed logging into vlog and xla dumping. If this is disabled, no compilation summary will be printed in the end of computation and no hlo modules will be dumped.
int32 xla_gpu_force_compilation_parallelism = 147
Overrides normal multi-threaded compilation settting to use this many threads. Setting to 0 (the default value) means no enforcement.
bool xla_gpu_deterministic_ops = 148
Guarantees run-to-run determinism. At present, the HLO ops Scatter and SelectAndScatter do not have deterministic XLA:GPU implementations. Compilation errors out if these ops are encountered.
repeated string xla_gpu_llvm_ir_file = 150
Paths to files with LLVM code.
bool xla_gpu_enable_async_all_reduce = 152
Convert synchronous all-reduces ops into asynchronous.
int64 xla_gpu_all_reduce_combine_threshold_bytes = 157
Size threshold (in bytes) for the GPU all-reduce combiner.
bool xla_gpu_all_reduce_contiguous = 158
Combine GPU all-reduces into a single operation over a contiguous buffer.
int32 xla_gpu_all_reduce_blueconnect_num_devices_per_host = 159
Number of devices per host for first stage of BlueConnect decomposition pass. The pass will attempt to decompose all-reduces ops into a ReduceScatter-AllReduce-AllGather sequence, with the initial ReduceScatter being performed over all of the devices in the same host. Set to < 1 to disable all-reduce decomposition.
bool xla_gpu_enable_cudnn_frontend = 160
Whether to use the cuDNN frontend API for convolutions when possible.
bool xla_dump_disable_metadata = 153
Disable dumping metadata in HLO dumps.
string xla_dump_hlo_pipeline_re = 154
If this flag is specified, will only dump HLO before and after passes in the pass pipeline that matches this regular expression. Default empty value enables dumping in all pipelines.
bool xla_gpu_strict_conv_algorithm_picker = 156
If true, abort immediately when conv algorithm picker fails, rather than logging a warning and proceeding with fallback.
bool xla_gpu_enable_xla_runtime_executable = 169
If true, use XLA runtime for XLA:GPU backend.
int64 xla_gpu_nccl_termination_timeout_seconds = 163
Timeout in seconds before terminating jobs that are stuck in a NCCL Rendezvous. Negative value disables the timeout and will not terminate.
bool xla_gpu_enable_shared_constants = 165
Enables shared constants for XLA/GPU. This allows large constants to be shared among multiple GPU executables.
bool xla_gpu_enable_cublaslt = 166
Whether to use cuBLASLt for GEMMs on GPUs.
int64 xla_gpu_redzone_scratch_max_megabytes = 167
Size threshold (in megabytes) for the GPU redzone scratch allocator.
bool xla_gpu_simplify_all_fp_conversions = 168
Allows all floating-point conversions to be simplified, including those that affect the numerics. The `BFloat16Normalization` pass inserts many `f32 -> bf16 -> f32` conversion pairs. These are not removed by the `AlgebraicSimplifier`, as that will only simplify conversions that are no-ops, e.g. `bf16 -> f32 -> bf16`. Removing these improves accuracy.
bool xla_gpu_normalize_layouts = 172
An experimental option to force all layouts present in the after-optimizations HLO to be descending, e.g. ShapeUtil::MakeShapeWithDescendingLayout is an identity on all instructions.
bool xla_cpu_use_acl = 174
Generate calls to Arm Compute Library in the CPU backend.
bool xla_cpu_strict_dot_conv_math = 175
By default, XLA:CPU will run fp16 dot/conv as fp32, as this is generally (much) faster on our hardware. Set this flag to disable this behavior.
map<string, string> xla_backend_extra_options = 500
Extra options to pass to the compilation backend (e.g. LLVM); specific interpretation of these values is left to the backend.
Used in:
IGNORE = 0
Do not insert any shape checks for dynamically shaped operations; output buffers might contain garbage data if shapes don't match.
RUNTIME = 1
Check shapes at runtime, will insert an extra synchronization if shapes cannot be proven correct at compile time.
COMPILE_TIME = 2
Will refuse to compile any program where shape correctness can not be established at compile time.
Used in: ,
STEP_MARK_AT_ENTRY = 0
Generate a step marker at the program entry. This handles the case where each step is done by one or multiple program execution(s). Only the first program will be tagged for generating a step marker at the program entry. This is the default.
STEP_MARK_AT_TOP_LEVEL_WHILE_LOOP = 1
Generate a step marker at each iteration of the top level while loop, which is assumed to be a training loop.
STEP_MARK_AT_SECOND_LEVEL_WHILE_LOOP = 3
Generate a step marker at each iteration of the second level while loops, which is assumed to be a training or eval loop.
STEP_MARK_NONE = 2
No step marker generated.
message
xla_data.proto:427DeviceAssignmentProto is a serialized form of DeviceAssignment class, which represents the device ids assigned to a set of replicated computations. See xla::DeviceAssignment class comment for more details.
Used in: , , , , ,
int32 replica_count = 1
int32 computation_count = 2
repeated computation_devices = 3
message
xla_data.proto:433Each logical computation runs on replica_count physical devices. ComputationDevice represents the device ids assinged to the replicas.
Used in:
repeated int32 replica_device_ids = 1
message
xla_data.proto:393Handle given to a user that represents a replicated virtual device. Each replicated device represents N physical devices for execution where N is the number of replicas.
Used in: , , , , ,
int64 handle = 1
int64 device_count = 2
The number of model-parallel virtual devices that communicate via XLA Send/Recv instructions.
message
protocol.proto:35Describes a device local to a host.
Used in:
int32 local_device_ordinal = 1
string name = 2
string vendor = 3
int32 global_device_id = 4
The following fields are present in the GlobalTopologyProto message returned by EnumerateDevices() but not in the LocalTopologyProto messages passed to EnumerateDevices(). In other words, the coordinator node determines the global device IDs during EnumerateDevices().
Globally unique ID number.
A DimLevelType indicates the encoding method for a dimension in an array. The semantics of this field are identical to those of the MLIR SparseTensor dialect. This should be kept in sync with the SparseTensor DimLevelType enum: https://github.com/llvm/llvm-project/blob/5674a3c88088e668b684326c2194a6282e8270ff/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td#L86
Used in:
DIM_DENSE = 0
The corresponding dimension is Dense, every entry is stored.
DIM_COMPRESSED = 1
The corresponding dimension is Compressed, only nonzeros are stored.
DIM_SINGLETON = 2
The corresponding dimension contains a single coordinate, no sibling elements for each parent.
message
xla_data.proto:615Used in: ,
repeated int64 lhs_contracting_dimensions = 1
The dimension numbers that represent the 'lhs' contracting dimensions.
repeated int64 rhs_contracting_dimensions = 2
The dimension numbers that represent the 'rhs' contracting dimensions.
repeated int64 lhs_batch_dimensions = 3
The dimension numbers that represent the 'lhs' batch dimensions.
repeated int64 rhs_batch_dimensions = 4
The dimension numbers that represent the 'rhs' batch dimensions.
message
hlo.proto:416Used in:
repeated entries = 1
message
hlo.proto:440A list of bindings which indicates that the `target_dim_num` in the subshape `target_param_index` of parameter `target_param_num` is a dynamic dimension and its real dynamic size is represented by `dynamic_param_index` in parameter `dynamic_param_num`. As an example, imagine we have a program: ENTRY main { a = f32[] parameter(0) b = f32[10] parameter(1) ROOT root = (f32[], f32[10]) tuple(%a, %b) } Let's say 'b' (param index 1) is a dynamic shape whose input has an upperbound of 10 and real size is determined at runtime.'a' represents the real size of b's first dimension. In this case, the fields are set in the following way: dynamic_param_num = 1 dynamic_param_index = {} target_param_num = 0 target_param_index = {} target_param_dim = 0
Used in:
int64 dynamic_param_num = 1
repeated int64 dynamic_param_index = 2
int64 target_param_num = 3
repeated int64 target_param_index = 4
int64 target_param_dim_num = 5
message
hlo.proto:715Encodes attributes for an entry function.
Used in:
repeated buffers = 1
string result_xla_shape = 2
xla::Shape in string format.
message
hlo.proto:722Encodes attributes for a single buffer parameter.
Used in:
int64 lmhlo_params = 1
Represents an lmhlo.params function argument attribute.
bool lmhlo_params_present = 6
TODO(hanbinyoon): Deprecate when optional fields are available in proto3 (Protocol Buffers v3.15.0).
optional lmhlo_param_shape_index = 2
Represents an lmhlo.param_shape_index function argument attribute.
string lmhlo_constant_name = 3
Represents an lmhlo.constant_name function argument attribute.
bool lmhlo_must_alias = 4
Represents an lmhlo.must_alias function argument attribute.
optional lmhlo_output_index = 5
Represents an lmhlo.params function argument attribute.
message
hlo.proto:717Acts as the underlying container for an xla::ShapeIndex.
Used in:
repeated int64 indices = 1
message
compile_options.proto:10A serialization of xla::ExecutableBuildOptions. Next id: 13.
Used in:
int64 device_ordinal = 1
If set, this is the device to build the computation for. Valid device_ordinal values are: 0 to # of devices - 1. These values are identical to the device ordinal values used by StreamExecutor. The built executable will be executable on any device equivalent to the specified device as determined by Backend::devices_equivalent(). A value of -1 indicates this option has not been set.
optional result_layout = 2
If set, this specifies the layout of the result of the computation. If not set, the service will chose the layout of the result. A Shape is used to store the layout to accommodate tuple result shapes. A value of nullptr indicates the option has not been set.
optional debug_options = 3
Expose access to the XLA debug options which will be passed to the compilation process.
int64 num_replicas = 4
The number of replicas of this computation that are to be executed. Defaults to 1.
int64 num_partitions = 5
The number of partitions in this computation. Defaults to 1.
bool use_spmd_partitioning = 6
Indicates whether to use SPMD (true) or MPMD (false) partitioning when num_partitions > 1 and XLA is requested to partition the input program.
bool use_auto_spmd_partitioning = 7
Whether to automatically generate XLA shardings for SPMD partitioner.
bool deduplicate_hlo = 8
Whether HLOs should be deduplicated.
optional device_assignment = 9
If set, this specifies a static device assignment for the computation. Otherwise, the computation will be compiled generically and can be run with any device assignment compatible with the computation's replica and partition counts.
bool alias_passthrough_params = 10
Whether input and output buffers are aliased if the associated parameter is passed-through XLA modules without being changed.
bool run_backend_only = 11
By default, XLA builds an executable by invoking standard compilation, i.e. running Compiler::Compile, or both Compiler::RunHloPasses and Compiler::RunBackend. When run_backend_only is set to true, XLA builds an executable by invoking only RunBackend and skip invoking RunHloPasses, which can be used to compile post-optimizations HLO modules.
bool allow_spmd_sharding_propagation_to_output = 12
Allows sharding propagation to propagate to the outputs. This changes the output shape of the computation (which is undesirable), but it can be used to allow to run partial compilation to determine what would be the output sharding of a computation if XLA would be allowed to propagate the sharding which can be used by higher level framework as a way to query intermediate sharding of operations when multiple computation would be chained and merged together.
message
xla.proto:622TODO(b/118493728): Remove this and ExecuteGraphParallelRequest and replace the uses with calls to Compile and Execute.
Used in:
optional computation = 1
repeated arguments = 2
optional execution_options = 3
Options that affect how XLA compiles and runs code to service this request.
message
xla.proto:634Used as response type in: XlaService.Execute
Used as field type in:
optional output = 1
optional profile = 2
message
xla_data.proto:379Handle given to a user that represents an execution that the user launched asynchronously on the device.
Used in: , ,
int64 handle = 1
message
xla.proto:444These settings control how XLA compiles and/or runs code. Not all settings will have an effect on every platform. When adding new fields, keep in mind that boolean fields default to false.
Used in: ,
optional shape_with_output_layout = 2
This optional field's layout is used as a hint when storing the output of this computation. Subsequent transfers of this output array to the client may be faster when using this layout. We use a Shape here to accommodate computations that return a tuple.
uint64 seed = 3
Used to seed random-number generators used in this computation. If this is 0, we generate a seed ourselves. TODO(b/32083678): Changing the seed unnecessarily forces a recompilation.
optional debug_options = 4
repeated device_handles = 5
This optional field specifies a particular set of devices to run the computation on. The computation will be partitioned across these devices. If not provided, the default device will be chosen.
int32 num_replicas = 6
Number of replicas of the computation to run. If zero, uses the default number of replicas for the XLA service.
optional device_assignment = 7
This optional field specifies the device assignment if known at compile time.
bool alias_passthrough_params = 8
Alias input and output buffers for parameters that are passed-through XLA modules without being changed.
int32 num_partitions = 9
Number of partitions of the computation to run (model parallelism). If zero, uses the default number of partitions for the XLA service.
int32 launch_id = 10
Used to identify a set of programs that should be launch together.
bool use_spmd_partitioning = 11
Indicates whether to use SPMD (true) or MPMD (false) partitioning when num_partitions > 1 and XLA is requested to partition the input program.
bool use_auto_spmd_partitioning = 15
Whether to automatically generate XLA shardings for SPMD partitioner.
repeated int64 auto_spmd_partitioning_mesh_shape = 16
Device mesh shape used to create the sharding search space when use_auto_spmd_partitioning=true.
repeated int64 auto_spmd_partitioning_mesh_ids = 17
Device mesh ids compatible with the above mesh_shape used when use_auto_spmd_partitioning=true.
bool deduplicate_hlo = 12
If set, deduplicate hlo into function calls to reduce binary size. Only works on TPU.
bool allow_spmd_sharding_propagation_to_output = 14
Allows sharding propagation to propagate to the outputs. This changes the output shape of the computation (which is undesirable), but it can be used to allow to run partial compilation to determine what would be the output sharding of a computation if XLA would be allowed to propagate the sharding which can be used by higher level framework as a way to query intermediate sharding of operations when multiple computation would be chained and merged together.
message
xla_data.proto:346Profile data from the execution of a computation.
Used in: ,
bool compilation_cache_hit = 1
Whether the executable was read from the compilation cache.
int64 compile_time_ms = 2
The time in milliseconds spent to compile the computation. This only set if the executable was not read from the compilation cache (compilation_cache_hit == false).
int64 compute_cycle_count = 3
The number of cycles spent for the computation. This does not include the time taken for the data transfers between the host and the device. This is a target-dependent field and only used for debugging purposes.
int64 compute_time_ns = 4
The time in nanoseconds spent for the computation, without data transfer.
int64 compute_and_transfer_time_ns = 5
The time in nanoseconds spent for the entire computation, including the result data transfer time. Current implementation does not spend any cycles for the input data transfer since the memory is initialized with the proper values before the execution.
int64 executable_size_in_bytes = 6
The size of the binary code in the executable.
bool profile_cache_hit = 7
Whether this profile was drawn from a cache of profiles instead of from execution on the hardware.
Used in:
FFT = 0
Forward FFT; complex in, complex out.
IFFT = 1
Inverse FFT; complex in, complex out.
RFFT = 2
Forward real FFT; real in, fft_length / 2 + 1 complex out
IRFFT = 3
Inverse real FFT; fft_length / 2 + 1 complex in,
message
xla_data.proto:675Generic map of attributes used to pass hints / configuration options from the Python frontend to the XLA backend.
Used in:
map<string, string> map = 1
message
xla_data.proto:523Describes the dimension numbers for a gather operation. See https://www.tensorflow.org/performance/xla/operation_semantics#gather for more details.
Used in:
repeated int64 offset_dims = 1
"Window indices" is a term for a set of indices that index into the interior of a dynamic-slice from the input tensor, the starting indices for which were computed from output_gather_dims (see the operation semantic for how this is defined) and the start_indices tensor. The window indices for a specific output index Out is computed as: i = 0 for (k : [0, input_tensor_shape.rank)) window_indices[k] = if k in collapsed_slice_dims then 0 else Out[offset_dims[i++]]
repeated int64 collapsed_slice_dims = 2
repeated int64 start_index_map = 3
This is interpreted as a map from i to start_index_map[i]. It transforms the gather index looked up from the start_indices tensor into the starting index in the input space.
int64 index_vector_dim = 4
The dimension in the start_indices input that contains the starting indices.
message
xla_data.proto:386Handle given to a user that represents a globally accessible allocation. Contrast this against a ComputationDataHandle, which is not globally accessible, since it only exists within a specific computation.
Used in: , , , , , , , , , , , ,
int64 handle = 1
message
protocol.proto:52Used in: ,
repeated nodes = 1
message
hlo.proto:567A trace of a HeapSimulator run.
Used in:
repeated events = 1
bool whole_module_simulation = 2
int64 buffer_allocation_index = 3
message
hlo.proto:570The trace includes a list of events, where each event describes one action performed by the heap simulator.
Used in:
Event.Kind kind = 1
int64 buffer_id = 2
The id of the LogicalBuffer that the event applies to.
string computation_name = 3
The HloInstruction that the simulation was processing that caused this event to occur, identified by its computation and instruction name. E.g. buffers defined by instruction A are allocated when processing A.
string instruction_name = 4
int64 share_with_canonical_id = 5
The id of the canonical LogicalBuffer that the buffer shares with. Only set for SHARE_WITH events.
Used in:
ALLOC = 0
A memory region was allocated for the buffer.
FREE = 1
A memory region was freed for the buffer.
SHARE_WITH = 2
A buffer was shared with another (canonical) buffer. This is similar to ALLOC, except that instead of allocating a new region of memory, the memory region of the canonical buffer is directly re-used. Multiple buffers may share with the same canonical buffer. The lifetime of the canonical buffer is extended to the union of all lifetimes.
message
hlo.proto:338Serialization of HloComputation.
Used in:
string name = 1
repeated instructions = 2
The array of instructions is always in a valid dependency order, where operands appear before their users.
optional program_shape = 4
int64 id = 5
The id of this computation.
int64 root_id = 6
The id of the root of the computation.
bool is_fusion_computation = 7
Whether this is a fusion computation. Fusion computations should use this to determine whether they are a fusion in CreateFromProto since the parent fusion_instruction_ may get removed and be nullptr.
string execution_thread = 8
The name of execution thread this computation belongs to.
optional printer_data = 1
repeated int64 profile_counters = 2
message
hlo.proto:388Used in:
repeated entries = 1
message
hlo.proto:402The following proto describes a pair of aliased an input (described by parameter number and a ShapeIndex of the parameter) and an output (described by a ShapeIndex of the root instruction). For example: entry = { output_shape_index={1}, parameter_number=0, parameter_shape_index={1, 2}, } This entry indicates that the first paremter's {1, 2} element is aliased with the {1} element of the root instruction.
Used in:
repeated int64 output_shape_index = 1
ShapeIndex of the root hlo.
int64 parameter_number = 2
Number of the parameter in entry computation.
repeated int64 parameter_shape_index = 3
ShapeIndex of the parameter instruction.
Kind kind = 4
The kind of alias to be setup.
message
hlo.proto:94Serialization of HloInstruction. Next ID: 80
Used in: ,
string name = 1
string opcode = 2
optional shape = 3
optional metadata = 7
optional literal = 8
Literal, only present for kConstant.
int64 parameter_number = 9
Parameter number is only present for kParameter.
string fusion_kind = 11
Fusion state, only present for kFusion.
int64 tuple_index = 13
Index for kGetTupleElement.
repeated int64 dimensions = 14
Dimensions present for some operations that require reshaping or broadcasting, including Reshape, Reduce, ReduceWindow, and Reverse.
optional window = 15
Describes the window in a windowed operation such as convolution.
optional convolution_dimension_numbers = 16
Describes the dimension numbers used for a convolution.
int64 feature_group_count = 50
The number of feature groups. Used for a convolution. Must be a divisor of the input feature dimension and output feature dimension. If not specified, it will use a default value of 1.
int64 batch_group_count = 58
repeated slice_dimensions = 17
int32 exponent_bits = 18
The bit sizes for a reduce-precision operation.
int32 mantissa_bits = 19
repeated int64 dynamic_slice_sizes = 20
Describes the [start, start + size) range size for a dynamic slice ('start' is specified dynamically in the second operand of the operation).
optional padding_config = 21
The padding configuration that describes the edge padding and interior padding of this pad instruction. Only set for pad instructions.
bytes outfeed_config = 22
Outfeed configuration information, only present for kOutfeed.
RandomDistribution distribution = 23
The distribution requested for random number generation. Only present for kRng.
float epsilon = 24
A small float number added to the variance to avoid divide-by-zero error. Only present for kBatchNormTraining, kBatchNormInference, and kBatchNormGrad.
int64 feature_index = 25
An integer value representing the index of the feature dimension. Only present for kBatchNormTraining, kBatchNormInference, and kBatchNormGrad.
int64 channel_id = 26
Represents a unique identifier for each Send/Recv instruction pair or optionally for collective instructions (AllReduce, CollectivePermute, AllToAll). Non-positive channel_id is equivalent to no channel id.
bytes infeed_config = 27
The string representation of the infeed configuration.
string custom_call_target = 28
Name of a external target (eg, global symbol) to call, only present for kCustomCall.
optional outfeed_shape = 29
Shape of outfeed request.
optional dot_dimension_numbers = 30
Describes the dimension numbers used for a dot operation
FftType fft_type = 31
FFT type (FFT, IFFT, etc).
repeated int64 fft_length = 32
FFT length.
string comparison_direction = 63
Comparison direction only used for kCompare.
optional gather_dimension_numbers = 33
Gather dimension numbers.
repeated int64 gather_slice_sizes = 34
int64 id = 35
The id of this instruction.
repeated int64 operand_ids = 36
repeated int64 control_predecessor_ids = 37
repeated int64 called_computation_ids = 38
optional sharding = 40
bytes backend_config = 43
Backend configuration for the instruction. Has backend-specific meaning.
repeated replica_groups = 49
Cross replica op fields.
int64 all_reduce_id = 45
Deprecated, but keeping it for backward compatibility. Use channel_id. Non-positive all_reduce_id is equivalent to no all_reduce_id.
bool use_global_device_ids = 71
If true, interprets ids in ReplicaGroup as global device ids, which is a linearized id of `replica_id * partition_count + partition_id`.
bool is_host_transfer = 47
Whether this Send/Recv instruction transfers data to/from the host. Only present for Send and Recv instructions and their SendDone and RecvDone partners.
bool is_stable = 60
Whether this Sort instruction should be stable.
optional scatter_dimension_numbers = 48
optional precision_config = 51
Precision configuration for the instruction. Has backend-specific meaning.
repeated source_target_pairs = 52
Collective permute field.
optional domain_entry_sharding = 54
Sharding for kDomain instructions.
optional domain_exit_sharding = 55
bool constrain_layout = 56
For custom call this indicates that the layouts are constrained. If constrain_layout is true then the 'shape' field must contain a layout, and 'operand_shapes_with_layout' must contain a shape with layout for each operand.
repeated operand_shapes_with_layout = 57
optional triangular_solve_options = 59
Options for TriangularSolve
optional cholesky_options = 62
Options for Cholesky
optional parameter_replication = 61
Describes how parameters behave with regards to replicas.
bool custom_call_has_side_effect = 65
Whether the kCustomCall instruction has side-effects, only present for kCustomCall.
repeated custom_call_output_operand_aliasing = 74
A list of CustomCallOutputOperandAliasing pairs that specifies aliasing buffers between output and operands for kCustomCall.
CustomCallSchedule custom_call_schedule = 76
Specifies the desired schedule for the custom-call. The field is only present for custom-call.
int64 delta = 66
The delta value for kRngGetAndUpdateState.
bool indices_are_sorted = 67
Specifies if the gather/scatter indices are guaranteed to be sorted by the caller.
optional frontend_attributes = 68
Frontend attributes to pass to the XLA backend.
bool unique_indices = 69
Specifies if all elements updated are guaranteed to be unique by the caller.
RandomAlgorithm rng_algorithm = 70
RNG algorithm used by kRngBitGenerator.
string comparison_type = 72
The comparison type used for kCompare.
bool is_cross_program_prefetch = 73
Specifies if this is a cross-program-prefetch, used by kCopyStart.
PaddingType padding_type = 75
If a convolution is dynamic, a dynamic padding type will be specified.
CustomCallApiVersion custom_call_api_version = 77
The API version used by the custom call function. This field is only present for custom-call. TODO(b/189822916): Remove this field when all clients are migrated to the status-returning API.
int64 async_group_id = 78
Represents a unique identifier for an async group which consists of an async start, async done, and zero or more async update operations. Negative async_group_id is equivalent to no async group id.
string async_execution_thread = 79
Represents a unique execution thread name for one or more async groups. Each HLO module may contain a main thread and one or more parallel threads. Empty async_execution_thread is equivalent to main thread.
message
hlo.proto:150Describes the [begin, end) index range and stride for slices.
Used in:
int64 start = 1
int64 limit = 2
int64 stride = 3
message
hlo.proto:604An abstraction representing a set of HLO module built to run concurrently across different devices.
string name = 1
repeated hlo_modules = 2
message
hlo.proto:653Metadata for an HLO module. Dumped after HLO passes and before LLO lowering with filename module_####.metadata.textproto, where #### is canonical_module_id.
int64 canonical_module_id = 1
Uniquely identifies an HloModuleMetadata. Equal to the first unique_id of the module (a module may go through multiple unique_ids). If a module is partitioned into multiple modules, those modules will each have a new HloModuleMetadata with a different canonical_module_id.
string module_group_name = 2
Name of the module group that the module is part of.
int64 original_module_id = 3
The canonical module id of the module that this one is partitioned from, if applicable.
repeated int64 partitioned_module_ids = 4
The canonical module ids of the modules that this one is partitioned into, if applicable.
repeated pass_metadata = 5
Metadata for the HLO passes that are run on the module.
message
hlo.proto:457Serialization of HloModule.
Used in: , , , , , , , ,
string name = 1
string entry_computation_name = 2
int64 entry_computation_id = 6
repeated computations = 3
The array of computations is always in a valid dependency order, where callees appear before their callers.
optional host_program_shape = 4
The host program shape (with layout) of the entry computation.
int64 id = 5
The id of this module.
optional schedule = 7
The schedule for this module.
optional input_output_alias = 8
Describes alias information between inputs and outputs.
optional dynamic_parameter_binding = 9
repeated cross_program_prefetches = 10
bool is_dynamic = 11
True if the module contains dynamic computation.
optional spmd_output_sharding = 12
repeated spmd_parameters_shardings = 14
bool use_auto_spmd_partitioning = 16
Uses AutoSharding pass or not.
repeated profile_info = 13
Profile information for the HLO module.
optional device_assignment = 15
DeviceAssignment object information.
message
hlo.proto:502Information about the optimization profile that this module contains.
Used in:
ProfileType profile_type = 1
The optimization profiles that this module contains.
double relative_speedup = 2
Speedup of tuned config compared to default config.
ProfileSource profile_source = 3
The source of the optimization profile that this module contains.
CompilationEvent compilation_event = 4
The compilation event that triggered the use of the profile.
The type of optimization profile in use for module-level optimizations.
Used in:
INVALID = 0
FLAG = 1
FUSION = 2
LAYOUT = 3
DOT = 4
message
hlo.proto:680Metadata for one run of an HLO pass on a module. Provides more information when processing debug dumps of HloProtos about the order of HLO passes and various other stats like duration. `pass_id` may also be used to identify a particular run of a pass in debug info that propagates through stages of compilation.
Used in:
int64 pass_id = 1
For a given module, pass_id uniquely identifies a run of an HLO pass on that module. Note that a pass_id may not always refer to the same pass because the order of passes during compilation may change. For finding metadata for a particular pass, pass_name and pipeline_name would be more reliable, although note that they may not be unique.
string pass_name = 2
string pipeline_name = 3
repeated string dump_filenames = 4
Filenames of the dumps of the module after this pass ran. Module may be dumped in multiple formats, and the order of formats in this field will stay consistent across passes.
bool module_changed = 5
Return value of pass.Run(). True if this pass changed the module, or, in the case where the module was run through this pass as part of a module group, true if this pass changed any module in the same module group.
int64 module_id = 6
The unique_id of the module that this pass is run on. May be different from the canonical_module_id of the HloModuleMetadata that this HloPassMetadata is inside.
repeated int64 module_group_module_ids = 7
If the module went through this pass as part of a module group, this is set as the ids of all the modules in the module group. Empty otherwise.
int64 start_timestamp_usec = 8
Timestamp before and after the pass is run. Note they may be equal.
int64 end_timestamp_usec = 9
Describes how to pretty-print a profile counter array gathered for a specific HloModule.
Used in:
repeated computation_infos = 1
HloComputationInfos for every HloComputation in the HloModule.
int64 profile_counters_size = 2
The size of the profile counters array we will pretty-print.
map<string, int64> extra_metrics = 3
Maps extra metric name to the index into the profile counters array.
string entry_computation = 4
Name of the entry computation.
Pretty-printer information about an HloComputation.
Used in:
string name = 1
int64 profile_index = 2
The index into the profile counters array for the HloComputation corresponding to this HloComputationInfo.
repeated instruction_infos = 3
HloInstructionInfos for every HloInstruction in the HloComputation for corresponding to this HloComputattionInfo.
Pretty-printer information about an HloInstruction.
Used in:
string long_name = 1
string short_name = 2
string category = 3
float flop_count = 4
Metrics computed by HloCostAnalysis.
float transcendental_count = 5
int64 bytes_accessed = 9
float optimal_seconds = 7
int64 profile_index = 8
The index into the profile counters array for the HloInstruction corresponding to this HloInstructionInfo.
message
hlo.proto:625Grouping message that contains all of the information above.
Used in: , ,
optional hlo_module = 1
optional buffer_assignment = 3
message
hlo.proto:369Serialization of an HLO schedule. An HLO schedule contains a total order of instructions for each non-fusion computation in the module.
Used in:
map<int64, > sequences = 1
Map from computation id to sequence.
message
hlo.proto:370Used in:
repeated int64 instruction_ids = 1
message
hlo.proto:636Encapsulates HloProto together with the arguments, result, and execution_platform. This message is used for purposes such as analysis/replay/file-storage.
Used in: ,
optional hlo = 1
The hlo graph.
repeated arguments = 2
The arguments passed to the graph.
optional result = 3
The result of the graph.
string execution_platform = 4
The name of the platform used to run the graph.
enum Kind
hlo.proto:378Used in:
UNDEFINED_ALIAS = 0
Define a UNDEFINED_ALIAS equal to zero to get around the default-0 proto3 behavior and missing has_*() APIs.
MAY_ALIAS = 1
The buffers may or may not alias at runtime.
MUST_ALIAS = 2
The buffers must alias at runtime.
message
xla_data.proto:148A layout describes how the array is placed in (1D) memory space. This includes the minor-to-major ordering of dimensions within a shape. Clients must specify the layouts of input Literals to the computation. Layouts specified in interior operations which take Shapes (for example, Convert) are ignored. See the XLA documentation for more information on shapes and layouts. LINT.IfChange
Used in: , ,
repeated DimLevelType dim_level_types = 9
The dimension level type list for this array, specifying the way in which each array dimension is represented in memory. If this list is empty, the array is assumed to be dense.
repeated int64 minor_to_major = 1
Sequence of dimension numbers, from minor (fastest varying index) to major (slowest varying index). This field is required.
repeated tiles = 6
A sequence of tiles, starting from the tile that's applied first to the Shape. TODO(b/119839262): implement tiling in each backend or add Unimplemented error.
int64 element_size_in_bits = 7
Bit size of each element. If the size is bigger than what the element type requires, the value is stored in the least significant bits and the additional most significant bits are filled with 0's. TODO(b/119839262): implement in each backend or add Unimplemented error.
int64 memory_space = 8
Memory space where this array resides. The integer field is interpreted in a backend-specific manner.
optional physical_shape = 10
The physical, on-device shape used to represent the shape this layout belongs to. Only used for sparse arrays. The layout(s) contained within the physical shape should not also contain a physical shape.
message
xla_data.proto:445Literals are used when the server and client need to exchange materialized data / results. Literals are also used to describe constants used in computations. Transfers to/from the client are encoded in literal form, and the structure of the repeated fields is implied by the shape.
Used in: , , , , , , , ,
optional shape = 1
repeated bool preds = 2
bytes s8s = 15
bytes u8s = 3
repeated int32 s32s = 4
repeated int64 s64s = 5
repeated uint32 u32s = 6
repeated uint64 u64s = 7
repeated float f32s = 8
repeated double f64s = 9
repeated float c64s = 12
Stored as interleaved real, imag floats.
repeated double c128s = 18
Stored as interleaved real, imag doubles.
repeated tuple_literals = 10
bytes f16s = 11
The F16s, BF16s, U16s and S16s are encoded in little endian byte order
bytes bf16s = 13
bytes u16s = 16
bytes s16s = 17
repeated int64 sparse_indices = 14
Next = 19
message
protocol.proto:47Used in: ,
int32 node_id = 1
repeated devices = 2
message
hlo.proto:521Serialization of LogicalBuffer.
Used in:
int64 id = 1
int64 size = 2
optional defined_at = 3
The location where the buffer is defined.
int64 color = 4
message
hlo.proto:524Location represents an instruction and its shape index, which uniquely identifies a point where a buffer is needed.
Used in: ,
string computation_name = 1
NOTE: module_name isn't necessary, since all LogicalBuffers are associated with a single HloModule. TODO(b/239098765): Remove instruction_name and computation_name.
string instruction_name = 2
int64 instruction_id = 4
repeated int64 shape_index = 3
message
xla_data.proto:286Symbolization metadata for HLO Instructions. This metadata is used for debugging XLA code generation, as well as performance profiling of XLA-generated executables.
Used in: ,
string op_type = 1
The framework op name that generated this XLA op. Frameworks that build on top of XLA should mirror the names of their ops back to users by specifying the op_type. In this way, even if the framework's "ops" are implemented as multiple XLA HLO Ops, they can be grouped appropriately. (e.g. if a SoftMax layer is emitted into XLA as multiple ops, then each op should have the op_type be "SoftMax".)
string op_name = 2
The user-specified name of the op. This name is often unique within a computation. Note: some frameworks add auto-generated names if the user does not provide one.
string source_file = 3
Indicate a file and line that this op is associated to in a user's program. e.g. it could be the file and line of user code that generated the op.
int32 source_line = 4
repeated ProfileType profile_type = 5
Deprecated, use [ProfileInfo][profile_type] instead.
int64 creation_pass_id = 6
HloPassMetadata.pass_id of the pass that created this HLO instruction object. Should never be copied between HLO instructions. Zero if unset and -1 if the instruction was created before HLO passes began.
int64 logical_creation_pass_id = 7
HloPassMetadata.pass_id of the pass that created the logical functionality that this HLO instruction represents. Should be copied between HLO instructions that correspond across compilation passes. Zero if unset and -1 if the instruction was created before HLO passes began.
int64 size_of_generated_code_in_bytes = 8
The footprint of the generated code for the instruction.
int64 size_of_memory_working_set_in_bytes = 9
The size of the working set, i.e., the amount of memory, used by the instruction in a compiler-managed fast device memory.
optional profile_info = 10
Profile information for the Op.
message
xla_data.proto:327Information about the optimization profile that this operation contains.
Used in:
repeated ProfileType profile_type = 1
The type of optimization profiles that this operation contains.
double relative_speedup = 2
Speedup of tuned config compared to default config. TODO(b/203817882) Set the relative_speedup.
ProfileSource profile_source = 3
The source of the optimization profiles that this operation contains.
CompilationEvent compilation_event = 4
The compilation event that triggered the use of the profiles.
message
xla_data.proto:680LINT.IfChange
Used in: , , ,
OpSharding.Type type = 1
optional tile_shape = 2
The shape of the sharded tile.
repeated int64 tile_assignment_dimensions = 3
The shape of the tile assignment tensor - this must be the same rank as tile_shape and the product of its dimensions must equal tile_assignment_devices.size().
repeated int64 tile_assignment_devices = 4
Flattened list of device IDs. The order of flattening is the same as used by IndexUtil::MultiToLinearIndex(tile_assignment_shape).
repeated tuple_shardings = 5
If type == TUPLE, the sub-shardings, one per leaf node in the tuple shape, in pre-order. The tuple shape could be nested; here we store just a flattened list of all leaves in the tuple shape. Note that the tuple shape is not stored here; shardings do not store the shapes to which they are applied, this is inferred from the instruction this sharding gets attached to.
bool replicate_on_last_tile_dim = 6
Only used for OTHER type. If true, data is sharded according to other dimensions of tile_assignment(), but replicated across devices along the last dimension. (Experimental)
repeated metadata = 7
This field is used to track the source of this sharding, usually derived from instructions. Multple metadata may be populated if sharding is combined with other shardings. Metadata are to not be populated when type == TUPLE and instead metadata should be set on individual tuple elements.
repeated OpSharding.Type last_tile_dims = 8
This field is used to represented the sharding type of each subgroup. For example, sharding={devices=[2,2,2,2]0,1,2,...,15 last_tile_dims={ replicate, manual, unreduced}} means that each of the last 3 dimensions in [2,2,2,2] represents a subgrouping in replicate, manual, unreduced sharding type respectively.
Used in:
REPLICATED = 0
This sharding is replicated across all devices (implies maximal, all other fields are unused).
MAXIMAL = 1
This sharding is maximal - one device runs the entire operation.
TUPLE = 2
This sharding is a tuple - only the tuple_shardings field is valid.
OTHER = 3
None of the above; tile_shape and tile_assignment are both used.
MANUAL = 4
This op is manually sharded: the shapes are already partitioned and the partitioner should not change this op.
message
xla_data.proto:94Describes the padding configuration for Pad operation. The padding amount on both edges as well as between the elements are specified for each dimension.
Used in:
repeated dimensions = 1
The padding configuration for all dimensions.
message
xla_data.proto:96Describes the padding configuration for a dimension.
Used in:
int64 edge_padding_low = 1
Padding amount on the low-end (next to the index 0). May be negative.
int64 edge_padding_high = 2
Padding amount on the high-end (next to the highest index). May be negative.
int64 interior_padding = 3
Padding amount between the elements. May not be negative.
Used in:
PADDING_INVALID = 0
PADDING_VALID = 1
Only valid portion of the base are covered.
PADDING_SAME = 2
Extra is added to produce same output size as the input.
message
xla_data.proto:766Describes whether all data-parallelism replicas will receive the same parameter data at each buffer.
Used in:
repeated bool replicated_at_leaf_buffers = 1
A list of boolean values for the flattened leaf buffers. Each value indicates whether the corresponding leaf buffer is replicated. If this field is empty, it means no buffer is replicated. Otherwise, the number of elements in this field must match the number of leaf buffers in the HLO instruction's shape.
message
xla_data.proto:749Used to indicate the precision configuration. It has backend specific meaning.
Used in: ,
repeated PrecisionConfig.Precision operand_precision = 1
Used in:
DEFAULT = 0
HIGH = 1
HIGHEST = 2
PACKED_NIBBLE = 3
Each U8/S8 value in a tensor actually represents 2 nibble values.
Primitive types are the individual values that can be held in rectangular multidimensional arrays. A description of the rectangular multidimensional array dimensions / primitive type is given by Shape, below. LINT.IfChange
Used in: ,
PRIMITIVE_TYPE_INVALID = 0
Invalid primitive type to serve as default.
PRED = 1
Predicates are two-state booleans.
S8 = 2
Signed integral values of fixed width.
S16 = 3
S32 = 4
S64 = 5
U8 = 6
Unsigned integral values of fixed width.
U16 = 7
U32 = 8
U64 = 9
F16 = 10
Floating-point values of fixed width. Note: if f16s are not natively supported on the device, they will be converted to f16 from f32 at arbirary points in the computation.
F32 = 11
BF16 = 16
Truncated 16 bit floating-point format. This is similar to IEEE's 16 bit floating-point format, but uses 1 bit for the sign, 8 bits for the exponent and 7 bits for the mantissa.
F64 = 12
C64 = 15
Complex values of fixed width.
Paired F32 (real, imag), as in std::complex<float>.
C128 = 18
Paired F64 (real, imag), as in std::complex<double>.
TUPLE = 13
A tuple is a polymorphic sequence; e.g. a shape that holds different sub-shapes. They are used for things like returning multiple values from a computation; e.g. a computation that returns weights and biases may have a signature that results in a tuple like (f32[784x2000], f32[2000]) If a shape proto has the tuple element type, it may not have any entries in the dimensions field.
OPAQUE_TYPE = 14
An opaque type used for passing context-specific data to a custom operation. Shapes of this primitive type will have empty dimensions and tuple_shapes fields. (OPAQUE would be a better name for this identifier, but that conflicts with a macro defined in windows.h.)
TOKEN = 17
A token type threaded between side-effecting operations. Shapes of this primitive type will have empty dimensions and tuple_shapes fields.
The source of the optimization profile.
Used in: ,
PROFILE_SOURCE_UNKNOWN_SOURCE = 0
PROFILE_SOURCE_EMBEDDED = 1
PROFILE_SOURCE_REMOTE = 2
The type optimization profiles in use for Op-level optimizations.
Used in: ,
INVALID = 0
WINDOW = 1
FLAG = 2
INTEGER = 3
message
xla_data.proto:245Shape of the parameters and output of a computation (like a traditional function signature).
Used in: , , ,
repeated parameters = 1
optional result = 2
repeated string parameter_names = 3
Used in:
RNG_DEFAULT = 0
Backend dependent default algorithm.
RNG_THREE_FRY = 1
RNG_PHILOX = 2
Next: 2
Used in:
RNG_INVALID = 0
RNG_UNIFORM = 1
Creates a uniform-distribution-generated random number on the semi-open interval [parameter[0], parameter[1]).
RNG_NORMAL = 2
Creates a normal-distribution-generated random number with mean parameter[0] and standard deviation parameter[1].
message
xla_data.proto:735Describes the replica groups in a cross replica op (e.g., all-reduce and all-to-all).
Used in:
repeated int64 replica_ids = 1
The ids of the replicas that belongs to the same group. The ordering of the ids matters in some ops (e.g., all-to-all).
message
run_hlo_module.proto:7Used in:
repeated arguments = 2
Arguments used by the iteration.
optional result = 3
Ressult of the iteration on the target platform.
optional reference_result = 4
Result of the iteration on the reference platform.
message
run_hlo_module.proto:18repeated iterations = 1
Iterations of run hlo module.
message
xla_data.proto:554Describes the dimension numbers for a scatter operation. All the fields are similar to the corresponding fields in GatherDimensionNumbers. Differences are noted below.
Used in:
repeated int64 update_window_dims = 1
The set of dimensions in the updates shape that are window dimensions.
repeated int64 inserted_window_dims = 2
The set of window dimensions that must be inserted into the updates shape.
repeated int64 scatter_dims_to_operand_dims = 3
int64 index_vector_dim = 4
message
xla_data.proto:207A shape describes the number of dimensions in the array, the size of each dimension, and the primitive component type. Tuples are a special case in that they have rank zero and have tuple_shapes defined. See the XLA documentation for more information on shapes and layouts. LINT.IfChange
Used in: , , , , , , , , , , , , , , , , ,
PrimitiveType element_type = 2
The element type for this shape.
repeated int64 dimensions = 3
The size (number of elements) for each dimension, or an upper bound on the size if the dimension is dynamic. In XLA, dimensions are numbered from 0 to N-1 for an N-dimensional array. The first element of 'dimensions' is the size of dimension 0, the second element is the size of dimension 1, and so forth. Empty list indicates a scalar. If the respective element in 'is_dimension_dynamic' is true then the value in this field represents an upper bound on the size of the dimension.
repeated tuple_shapes = 4
For tuples only, the shapes of constituent shapes in the tuple sequence.
optional layout = 5
The layout used to back this shape.
repeated bool is_dynamic_dimension = 6
For arrays, this indicates whether or not each dimension is dynamically-sized. The number of elements in this repeated field should be zero (indicating that no dimensions are dynamic) or equal to the number of elements in the 'dimensions' field.
message
xla_data.proto:742Describes the source target pair in the collective permute op.
Used in:
int64 source = 1
int64 target = 2
message
xla_data.proto:130Describes a tile used in tiling-based layout. Refer to g3doc/third_party/tensorflow/compiler/xla/g3doc/tiled_layout.md for details about tiling-based layout.
Used in:
repeated int64 dimensions = 1
Number of elements in each dimension of the tile. It's ordered from the most major dimension of the tile to the most minor dimension of the tile. The dimensions correspond to a suffix of the dimensions of the shape being tiled.
message
xla_data.proto:647Used in:
bool left_side = 1
If true, solves ax = b. If false, solves xa = b.
bool lower = 2
If true, 'a' is lower triangular. If false, 'a' is upper triangular.
bool unit_diagonal = 3
If true, the diagonal elements of 'a' are assumed to be 1 and not accessed.
TriangularSolveOptions.Transpose transpose_a = 4
Should we transpose or use the adjoint of 'a'?
Used in:
TRANSPOSE_INVALID = 0
NO_TRANSPOSE = 1
Don't transpose 'a'.
TRANSPOSE = 2
Transpose 'a'.
ADJOINT = 3
Complex conjugate and transpose 'a'.
message
xla_data.proto:784A backend-config for kWhile loops that stores the loop's trip count, if it is known. This is useful for backends that can implement a `for i in 0..N` loop more efficiently than a `while` loop. For example, on GPUs, we can implement a `for i in 0..N` loop by enqueueing the kernels for the loop body N times, whereas implementing a `while` loop requires a host-device sync on each iteration.
optional known_trip_count = 1
This indirection lets us distinguish between known-trip-count == 0 and unknown-trip-count.
message
xla_data.proto:785Used in:
int64 n = 1
message
xla_data.proto:515Describes the windowing in an operation such as convolution. The window is moved across a base area and for each position of the window a computation is performed. The field below describes the window and the movement of the window across a base area.
Used in:
repeated dimensions = 1
message
xla_data.proto:468Used in:
int64 size = 1
The size of the window in this dimension. For a rectangle, this would be the width or height.
int64 stride = 2
The stride at which the window moves across the base area in this dimension. In other words, this is the spacing between different positions of the window in this dimension.
int64 padding_low = 3
If positive, means the amount of padding to add to the base area at the low end of this dimension; if negative, its negative means the number of elements removed from the low end of this dimension. For example, in the horizontal dimension of a rectangle, this would be the number of padding values to pad on the left, given that indices increase when going right. The actual padding value depends upon the context. Convolution pads with zeros. ReduceWindow and SelectAndScatter pads with the reduce function's init value.
int64 padding_high = 4
As padding_low, but on the high end of this dimension. For example, in the horizontal dimension of a rectangle, this would be the number of values to pad on the right, given that indices increase when going right.
int64 window_dilation = 5
Dilation factor of the sliding window in this dimension. A dilation factor of 1 means no dilation. window_dilation - 1 no-op entries ("holes") are implicitly placed between each kernel element. This value may not be less than 1. See documentation for convolution.
int64 base_dilation = 6
Dilation factor of the base area in this dimension. A dilation factor of 1 means no dilation. base_dilation - 1 no-op entries ("holes") are implicitly placed between each base area element. This value may not be less than 1. See documentation for convolution.
bool window_reversal = 7
Window reversal means that this dimension was logically reversed before the operation.
message
hlo.proto:749Encodes the underlying Xla runtime executable compiled from the XLA module.
Used in: ,
optional hlo_module_proto = 1
optional entry_func_attrs = 2
XLA-specific attributes of the executable's entry function.
bytes obj_file = 3
Serialized object file compiled from the XLA module.
string mlir_module = 4
Serialized MLIR module corresponding to compiled object file.