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_device_assignment(const DeviceAssignment & device_assignment)86 ExecutableBuildOptions& ExecutableBuildOptions::set_device_assignment(
87 const DeviceAssignment& device_assignment) {
88 device_assignment_ = device_assignment;
89 return *this;
90 }
91
ToString() const92 string ExecutableBuildOptions::ToString() const {
93 string result_layout = "nullopt";
94 if (result_layout_set_) {
95 result_layout = ShapeUtil::HumanStringWithLayout(result_layout_);
96 }
97 return absl::StrFormat(
98 "ExecutableBuildOptions{device_ordinal=%d, result_layout=%s, "
99 "num_replicas=%d}",
100 device_ordinal_, result_layout, num_replicas_);
101 }
102
CreateExecutionOptions(const ExecutableBuildOptions & build_options,const ProgramShape * program_shape)103 ExecutionOptions CreateExecutionOptions(
104 const ExecutableBuildOptions& build_options,
105 const ProgramShape* program_shape) {
106 ExecutionOptions execution_options = CreateDefaultExecutionOptions();
107 if (build_options.has_debug_options()) {
108 *execution_options.mutable_debug_options() = build_options.debug_options();
109 }
110 if (build_options.result_layout() != nullptr) {
111 *execution_options.mutable_shape_with_output_layout() =
112 build_options.result_layout()->ToProto();
113 } else {
114 Shape result_shape(program_shape->result());
115 LayoutUtil::SetToDefaultLayout(&result_shape);
116 *execution_options.mutable_shape_with_output_layout() =
117 result_shape.ToProto();
118 }
119 execution_options.set_num_replicas(build_options.num_replicas());
120 execution_options.set_num_partitions(build_options.num_partitions());
121 execution_options.set_use_spmd_partitioning(
122 build_options.use_spmd_partitioning());
123 execution_options.set_deduplicate_hlo(build_options.deduplicate_hlo());
124 if (build_options.has_device_assignment()) {
125 TF_CHECK_OK(build_options.device_assignment().Serialize(
126 execution_options.mutable_device_assignment()));
127 }
128 execution_options.set_alias_passthrough_params(
129 build_options.alias_passthrough_params());
130 return execution_options;
131 }
132
133 } // namespace xla
134