1 /* Copyright 2017 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/executable_run_options.h"
17 
18 #include <atomic>
19 
20 namespace xla {
21 
RunId()22 RunId::RunId() {
23   static std::atomic<int64> counter{0};
24   data_ = counter.fetch_add(1);
25 }
26 
operator ==(const RunId & a,const RunId & b)27 bool operator==(const RunId& a, const RunId& b) { return a.data_ == b.data_; }
28 
ToString() const29 std::string RunId::ToString() const {
30   return "RunId: " + std::to_string(data_);
31 }
32 
ToInt() const33 int64 RunId::ToInt() const { return data_; }
34 
set_device_ordinal(int device_ordinal)35 ExecutableRunOptions& ExecutableRunOptions::set_device_ordinal(
36     int device_ordinal) {
37   device_ordinal_ = device_ordinal;
38   return *this;
39 }
40 
device_ordinal() const41 int ExecutableRunOptions::device_ordinal() const { return device_ordinal_; }
42 
set_allocator(stream_executor::DeviceMemoryAllocator * allocator)43 ExecutableRunOptions& ExecutableRunOptions::set_allocator(
44     stream_executor::DeviceMemoryAllocator* allocator) {
45   allocator_ = allocator;
46   return *this;
47 }
48 
allocator() const49 stream_executor::DeviceMemoryAllocator* ExecutableRunOptions::allocator()
50     const {
51   return allocator_;
52 }
53 
set_stream(stream_executor::Stream * stream)54 ExecutableRunOptions& ExecutableRunOptions::set_stream(
55     stream_executor::Stream* stream) {
56   stream_ = stream;
57   return *this;
58 }
59 
stream() const60 stream_executor::Stream* ExecutableRunOptions::stream() const {
61   return stream_;
62 }
63 
set_host_to_device_stream(stream_executor::Stream * stream)64 ExecutableRunOptions& ExecutableRunOptions::set_host_to_device_stream(
65     stream_executor::Stream* stream) {
66   host_to_device_stream_ = stream;
67   return *this;
68 }
69 
host_to_device_stream() const70 stream_executor::Stream* ExecutableRunOptions::host_to_device_stream() const {
71   return host_to_device_stream_;
72 }
73 
set_intra_op_thread_pool(const Eigen::ThreadPoolDevice * intra_op_thread_pool)74 ExecutableRunOptions& ExecutableRunOptions::set_intra_op_thread_pool(
75     const Eigen::ThreadPoolDevice* intra_op_thread_pool) {
76   intra_op_thread_pool_ = intra_op_thread_pool;
77   return *this;
78 }
79 
intra_op_thread_pool() const80 const Eigen::ThreadPoolDevice* ExecutableRunOptions::intra_op_thread_pool()
81     const {
82   return intra_op_thread_pool_;
83 }
84 
set_execution_profile(ExecutionProfile * profile)85 ExecutableRunOptions& ExecutableRunOptions::set_execution_profile(
86     ExecutionProfile* profile) {
87   execution_profile_ = profile;
88   return *this;
89 }
90 
execution_profile() const91 ExecutionProfile* ExecutableRunOptions::execution_profile() const {
92   return execution_profile_;
93 }
94 
set_device_assignment(const DeviceAssignment * device_assignment)95 ExecutableRunOptions& ExecutableRunOptions::set_device_assignment(
96     const DeviceAssignment* device_assignment) {
97   device_assignment_ = device_assignment;
98   return *this;
99 }
100 
device_assignment() const101 const DeviceAssignment* ExecutableRunOptions::device_assignment() const {
102   return device_assignment_;
103 }
104 
set_gpu_executable_run_options(const gpu::GpuExecutableRunOptions * gpu_executable_run_options)105 ExecutableRunOptions& ExecutableRunOptions::set_gpu_executable_run_options(
106     const gpu::GpuExecutableRunOptions* gpu_executable_run_options) {
107   gpu_executable_run_options_ = gpu_executable_run_options;
108   return *this;
109 }
110 
111 const gpu::GpuExecutableRunOptions*
gpu_executable_run_options() const112 ExecutableRunOptions::gpu_executable_run_options() const {
113   return gpu_executable_run_options_;
114 }
115 
set_rng_seed(int rng_seed)116 ExecutableRunOptions& ExecutableRunOptions::set_rng_seed(int rng_seed) {
117   rng_seed_ = rng_seed;
118   return *this;
119 }
120 
rng_seed() const121 int ExecutableRunOptions::rng_seed() const { return rng_seed_; }
122 
set_run_id(RunId id)123 ExecutableRunOptions& ExecutableRunOptions::set_run_id(RunId id) {
124   run_id_ = id;
125   return *this;
126 }
127 
run_id() const128 RunId ExecutableRunOptions::run_id() const { return run_id_; }
129 
130 }  // namespace xla
131