1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/compiler/xla/client/executable_build_options.h"
17
18 #include "absl/strings/str_format.h"
19 #include "tensorflow/compiler/xla/debug_options_flags.h"
20 #include "tensorflow/compiler/xla/execution_options_util.h"
21 #include "tensorflow/compiler/xla/shape_util.h"
22
23 namespace xla {
24
set_device_allocator(se::DeviceMemoryAllocator * allocator)25 ExecutableBuildOptions& ExecutableBuildOptions::set_device_allocator(
26 se::DeviceMemoryAllocator* allocator) {
27 device_allocator_ = allocator;
28 return *this;
29 }
30
device_allocator() const31 se::DeviceMemoryAllocator* ExecutableBuildOptions::device_allocator() const {
32 return device_allocator_;
33 }
34
set_device_ordinal(int device_ordinal)35 ExecutableBuildOptions& ExecutableBuildOptions::set_device_ordinal(
36 int device_ordinal) {
37 CHECK_GE(device_ordinal, 0);
38 device_ordinal_ = device_ordinal;
39 return *this;
40 }
41
device_ordinal() const42 int ExecutableBuildOptions::device_ordinal() const { return device_ordinal_; }
43
mutable_debug_options()44 DebugOptions* ExecutableBuildOptions::mutable_debug_options() {
45 if (!has_debug_options()) {
46 debug_options_ = GetDebugOptionsFromFlags();
47 }
48 return &debug_options_.value();
49 }
50
set_result_layout(const Shape & shape_with_layout)51 ExecutableBuildOptions& ExecutableBuildOptions::set_result_layout(
52 const Shape& shape_with_layout) {
53 result_layout_set_ = true;
54 result_layout_ = shape_with_layout;
55 return *this;
56 }
57
result_layout() const58 const Shape* ExecutableBuildOptions::result_layout() const {
59 return result_layout_set_ ? &result_layout_ : nullptr;
60 }
61
set_num_replicas(int num_replicas)62 ExecutableBuildOptions& ExecutableBuildOptions::set_num_replicas(
63 int num_replicas) {
64 num_replicas_ = num_replicas;
65 return *this;
66 }
67
set_num_partitions(int num_partitions)68 ExecutableBuildOptions& ExecutableBuildOptions::set_num_partitions(
69 int num_partitions) {
70 num_partitions_ = num_partitions;
71 return *this;
72 }
73
set_use_spmd_partitioning(bool use_spmd_partitioning)74 ExecutableBuildOptions& ExecutableBuildOptions::set_use_spmd_partitioning(
75 bool use_spmd_partitioning) {
76 use_spmd_partitioning_ = use_spmd_partitioning;
77 return *this;
78 }
79
set_deduplicate_hlo(bool deduplicate_hlo)80 ExecutableBuildOptions& ExecutableBuildOptions::set_deduplicate_hlo(
81 bool deduplicate_hlo) {
82 deduplicate_hlo_ = deduplicate_hlo;
83 return *this;
84 }
85
set_broadcast_replicated_params(bool broadcast_replicated_params)86 ExecutableBuildOptions& ExecutableBuildOptions::set_broadcast_replicated_params(
87 bool broadcast_replicated_params) {
88 broadcast_replicated_params_ = broadcast_replicated_params;
89 return *this;
90 }
91
set_device_assignment(const DeviceAssignment & device_assignment)92 ExecutableBuildOptions& ExecutableBuildOptions::set_device_assignment(
93 const DeviceAssignment& device_assignment) {
94 device_assignment_ = device_assignment;
95 return *this;
96 }
97
ToString() const98 string ExecutableBuildOptions::ToString() const {
99 string result_layout = "nullopt";
100 if (result_layout_set_) {
101 result_layout = ShapeUtil::HumanStringWithLayout(result_layout_);
102 }
103 return absl::StrFormat(
104 "ExecutableBuildOptions{device_ordinal=%d, result_layout=%s, "
105 "num_replicas=%d}",
106 device_ordinal_, result_layout, num_replicas_);
107 }
108
CreateExecutionOptions(const ExecutableBuildOptions & build_options,const ProgramShape * program_shape)109 ExecutionOptions CreateExecutionOptions(
110 const ExecutableBuildOptions& build_options,
111 const ProgramShape* program_shape) {
112 ExecutionOptions execution_options = CreateDefaultExecutionOptions();
113 if (build_options.has_debug_options()) {
114 *execution_options.mutable_debug_options() = build_options.debug_options();
115 }
116 if (build_options.result_layout() != nullptr) {
117 *execution_options.mutable_shape_with_output_layout() =
118 build_options.result_layout()->ToProto();
119 } else {
120 Shape result_shape(program_shape->result());
121 LayoutUtil::SetToDefaultLayout(&result_shape);
122 *execution_options.mutable_shape_with_output_layout() =
123 result_shape.ToProto();
124 }
125 execution_options.set_num_replicas(build_options.num_replicas());
126 execution_options.set_num_partitions(build_options.num_partitions());
127 execution_options.set_use_spmd_partitioning(
128 build_options.use_spmd_partitioning());
129 execution_options.set_deduplicate_hlo(build_options.deduplicate_hlo());
130 if (build_options.has_device_assignment()) {
131 TF_CHECK_OK(build_options.device_assignment().Serialize(
132 execution_options.mutable_device_assignment()));
133 }
134 execution_options.set_alias_passthrough_params(
135 build_options.alias_passthrough_params());
136 return execution_options;
137 }
138
139 } // namespace xla
140