• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 #ifndef TENSORFLOW_COMPILER_XLA_CLIENT_EXECUTABLE_BUILD_OPTIONS_H_
17 #define TENSORFLOW_COMPILER_XLA_CLIENT_EXECUTABLE_BUILD_OPTIONS_H_
18 
19 #include "absl/strings/string_view.h"
20 #include "absl/types/optional.h"
21 #include "tensorflow/compiler/xla/service/device_memory_allocator.h"
22 #include "tensorflow/compiler/xla/shape.h"
23 #include "tensorflow/compiler/xla/util.h"
24 #include "tensorflow/compiler/xla/xla.pb.h"
25 #include "tensorflow/compiler/xla/xla_data.pb.h"
26 
27 namespace xla {
28 
29 // Class containing options for building an LocalExecutable with
30 // LocalClient::Compile.
31 class ExecutableBuildOptions {
32  public:
33   // If set, this is the device to build the computation for. Valid
34   // device_ordinal values are: 0 to # of devices - 1. These values are
35   // identical to the device ordinal values used by StreamExecutor. The built
36   // executable will be executable on any device equivalent to the specified
37   // device as determined by Backend::devices_equivalent(). A value of -1
38   // indicates this option has not been set.
39   ExecutableBuildOptions& set_device_ordinal(int device_ordinal);
40   int device_ordinal() const;
41 
42   // If set, this specifies the layout of the result of the computation. If not
43   // set, the service will chose the layout of the result. A Shape is used to
44   // store the layout to accommodate tuple result shapes. A value of nullptr
45   // indicates the option has not been set.
46   ExecutableBuildOptions& set_result_layout(const Shape& shape_with_layout);
47   const Shape* result_layout() const;
48 
49   // Expose access to the XLA debug options which will be passed to the
50   // compilation process.
has_debug_options()51   bool has_debug_options() const { return debug_options_.has_value(); }
debug_options()52   const DebugOptions& debug_options() const { return *debug_options_; }
53   DebugOptions* mutable_debug_options();
54 
55   // If set, this specifies an allocator that can be used to allocate temporary
56   // space on the device during compilation.  For example, the compiler might
57   // want to run various algorithms on the device and pick the fastest one -- it
58   // might allocate buffers for use by these algorithms using this allocator.
59   //
60   // This does not need to be the same as the DeviceMemoryAllocator passed when
61   // running the executable.
62   ExecutableBuildOptions& set_device_allocator(
63       DeviceMemoryAllocator* allocator);
64   DeviceMemoryAllocator* device_allocator() const;
65 
66   // Returns a string representation of the build options, suitable for
67   // debugging.
68   string ToString() const;
69 
70   // The number of replicas of this computation that are to be executed.
71   // Defaults to 1.
num_replicas()72   int num_replicas() const { return num_replicas_; }
73   ExecutableBuildOptions& set_num_replicas(int num_replicas);
74 
75  private:
76   int device_ordinal_ = -1;
77   Shape result_layout_;
78   bool result_layout_set_ = false;
79   absl::optional<DebugOptions> debug_options_;
80   DeviceMemoryAllocator* device_allocator_ = nullptr;
81   int num_replicas_ = 1;
82 };
83 
84 }  // namespace xla
85 
86 #endif  // TENSORFLOW_COMPILER_XLA_CLIENT_EXECUTABLE_BUILD_OPTIONS_H_
87