1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
|
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_EXECUTABLE_H_
#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_EXECUTABLE_H_
#include <memory>
#include <string>
#include "tensorflow/compiler/xla/service/buffer_assignment.h"
#include "tensorflow/compiler/xla/service/device_memory_allocator.h"
#include "tensorflow/compiler/xla/service/executable.h"
#include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h"
#include "tensorflow/compiler/xla/service/gpu/stream_assignment.h"
#include "tensorflow/compiler/xla/service/gpu/thunk.h"
#include "tensorflow/compiler/xla/service/gpu/thunk_schedule.h"
#include "tensorflow/compiler/xla/service/hlo_execution_profile.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_module_config.h"
#include "tensorflow/compiler/xla/service/shaped_buffer.h"
#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/stream_executor_no_cuda.h"
namespace xla {
namespace gpu {
// GPU-targeting implementation of the XLA Executable interface.
//
// Launches the given CUDA kernel via the StreamExecutor.
//
// This is an immutable data type after initialization, and thus thread safe.
class GpuExecutable : public Executable {
public:
GpuExecutable(tensorflow::StringPiece ptx,
std::unique_ptr<ThunkSchedule> thunk_schedule,
std::unique_ptr<HloModule> hlo_module,
std::unique_ptr<HloModuleConfig> module_config,
std::unique_ptr<BufferAssignment> assignment);
// This should be called after set_ir_module_string.
const string& ir_module_string() const { return ir_module_string_; }
// This should be called before ExecuteOnStream.
void set_ir_module_string(const string& ir_module_string) {
ir_module_string_ = ir_module_string;
}
// Returns the compiled PTX for the computation.
tensorflow::StringPiece ptx() const { return ptx_; }
StatusOr<perftools::gputools::DeviceMemoryBase> ExecuteOnStream(
const ServiceExecutableRunOptions* run_options,
tensorflow::gtl::ArraySlice<perftools::gputools::DeviceMemoryBase>
arguments,
HloExecutionProfile* hlo_execution_profile) override;
StatusOr<std::unique_ptr<ShapedBuffer>> ExecuteOnStream(
const ServiceExecutableRunOptions* run_options,
tensorflow::gtl::ArraySlice<const ShapedBuffer*> arguments,
HloExecutionProfile* hlo_execution_profile) override;
StatusOr<perftools::gputools::DeviceMemoryBase> ExecuteAsyncOnStream(
const ServiceExecutableRunOptions* run_options,
tensorflow::gtl::ArraySlice<perftools::gputools::DeviceMemoryBase>
arguments) override;
private:
Status ExecuteThunks(const ServiceExecutableRunOptions* run_options,
const BufferAllocations& buffer_allocations,
HloExecutionProfile* hlo_execution_profile);
// Returns the points-to set of the root instruction of the entry
// computation. Uses points-to analysis from buffer assignment.
const PointsToSet& GetRootPointsToSet() const;
// The LLVM IR, in string format, of the unoptimized module generated for this
// GpuExecutable. We save a string instead of an llvm::Module* because leaving
// llvm::Module* in a singleton can cause the heap checker to emit false
// positives.
//
// This string should be modified only before ExecuteOnStream.
string ir_module_string_;
// The reference to the compiled PTX for the computation.
const tensorflow::StringPiece ptx_;
// The thunks to be invoked by this GpuExecutable. They are generated by the
// IrEmitter.
const std::unique_ptr<ThunkSchedule> thunk_schedule_;
// Owns the buffer data at runtime. It provides information to allocate
// memory for every output/temp buffers.
const std::unique_ptr<BufferAssignment> assignment_;
TF_DISALLOW_COPY_AND_ASSIGN(GpuExecutable);
};
} // namespace gpu
} // namespace xla
#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_EXECUTABLE_H_
|
