xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/compile_options.proto (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1syntax = "proto3";
2
3package xla;
4
5import "tensorflow/compiler/xla/xla.proto";
6import "tensorflow/compiler/xla/xla_data.proto";
7
8// A serialization of xla::ExecutableBuildOptions.
9// Next id: 13.
10message ExecutableBuildOptionsProto {
11  // If set, this is the device to build the computation for. Valid
12  // device_ordinal values are: 0 to # of devices - 1. These values are
13  // identical to the device ordinal values used by StreamExecutor. The built
14  // executable will be executable on any device equivalent to the specified
15  // device as determined by Backend::devices_equivalent(). A value of -1
16  // indicates this option has not been set.
17  int64 device_ordinal = 1;
18
19  // If set, this specifies the layout of the result of the computation. If not
20  // set, the service will chose the layout of the result. A Shape is used to
21  // store the layout to accommodate tuple result shapes. A value of nullptr
22  // indicates the option has not been set.
23  xla.ShapeProto result_layout = 2;
24
25  // Expose access to the XLA debug options which will be passed to the
26  // compilation process.
27  xla.DebugOptions debug_options = 3;
28
29  // The number of replicas of this computation that are to be executed.
30  // Defaults to 1.
31  int64 num_replicas = 4;
32
33  // The number of partitions in this computation. Defaults to 1.
34  int64 num_partitions = 5;
35
36  // Indicates whether to use SPMD (true) or MPMD (false) partitioning when
37  // num_partitions > 1 and XLA is requested to partition the input program.
38  bool use_spmd_partitioning = 6;
39
40  // Whether to automatically generate XLA shardings for SPMD partitioner.
41  bool use_auto_spmd_partitioning = 7;
42
43  // Whether HLOs should be deduplicated.
44  bool deduplicate_hlo = 8;
45
46  // If set, this specifies a static device assignment for the computation.
47  // Otherwise, the computation will be compiled generically and can be run with
48  // any device assignment compatible with the computation's replica and
49  // partition counts.
50  xla.DeviceAssignmentProto device_assignment = 9;
51
52  // Whether input and output buffers are aliased if the associated parameter is
53  // passed-through XLA modules without being changed.
54  bool alias_passthrough_params = 10;
55
56  // By default, XLA builds an executable by invoking standard compilation, i.e.
57  // running Compiler::Compile, or both Compiler::RunHloPasses and
58  // Compiler::RunBackend. When run_backend_only is set to true, XLA builds an
59  // executable by invoking only RunBackend and skip invoking RunHloPasses,
60  // which can be used to compile post-optimizations HLO modules.
61  bool run_backend_only = 11;
62
63  // Allows sharding propagation to propagate to the outputs. This changes the
64  // output shape of the computation (which is undesirable), but it can be used
65  // to allow to run partial compilation to determine what would be the output
66  // sharding of a computation if XLA would be allowed to propagate the sharding
67  // which can be used by higher level framework as a way to query intermediate
68  // sharding of operations when multiple computation would be chained and
69  // merged together.
70  bool allow_spmd_sharding_propagation_to_output = 12;
71}
72
73// TODO(b/240299401): Add attributes to the proto to mirror CompileOptions.
74message CompileOptionsProto {}
75