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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
|
/* 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.
==============================================================================*/
#include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h"
#include <functional>
#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
namespace xla {
namespace cpu {
namespace runtime {
XfeedManager* GetXfeedManager() {
static XfeedManager* manager = new XfeedManager;
return manager;
}
extern const char* const kEigenMatMulF16SymbolName =
"__xla_cpu_runtime_EigenMatMulF16";
extern const char* const kEigenMatMulF32SymbolName =
"__xla_cpu_runtime_EigenMatMulF32";
extern const char* const kEigenMatMulF64SymbolName =
"__xla_cpu_runtime_EigenMatMulF64";
extern const char* const kMKLConvF32SymbolName = "__xla_cpu_runtime_MKLConvF32";
extern const char* const kMKLMatMulF32SymbolName =
"__xla_cpu_runtime_MKLMatMulF32";
extern const char* const kMKLMatMulF64SymbolName =
"__xla_cpu_runtime_MKLMatMulF64";
extern const char* const kMKLSingleThreadedMatMulF32SymbolName =
"__xla_cpu_runtime_MKLSingleThreadedMatMulF32";
extern const char* const kMKLSingleThreadedMatMulF64SymbolName =
"__xla_cpu_runtime_MKLSingleThreadedMatMulF64";
extern const char* const kEigenConvF16SymbolName =
"__xla_cpu_runtime_EigenConvF16";
extern const char* const kEigenConvF32SymbolName =
"__xla_cpu_runtime_EigenConvF32";
extern const char* const kEigenFftSymbolName = "__xla_cpu_runtime_EigenFft";
extern const char* const kEigenSingleThreadedMatMulF16SymbolName =
"__xla_cpu_runtime_EigenSingleThreadedMatMulF16";
extern const char* const kEigenSingleThreadedMatMulF32SymbolName =
"__xla_cpu_runtime_EigenSingleThreadedMatMulF32";
extern const char* const kEigenSingleThreadedMatMulF64SymbolName =
"__xla_cpu_runtime_EigenSingleThreadedMatMulF64";
extern const char* const kEigenSingleThreadedConvF16SymbolName =
"__xla_cpu_runtime_EigenSingleThreadedConvF16";
extern const char* const kEigenSingleThreadedConvF32SymbolName =
"__xla_cpu_runtime_EigenSingleThreadedConvF32";
extern const char* const kAcquireInfeedBufferForDequeueSymbolName =
"__xla_cpu_runtime_AcquireInfeedBufferForDequeue";
extern const char* const kReleaseInfeedBufferAfterDequeueSymbolName =
"__xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue";
extern const char* const kAcquireOutfeedBufferForPopulationSymbolName =
"__xla_cpu_runtime_AcquireOutfeedBufferForPopulation";
extern const char* const kReleaseOutfeedBufferAfterPopulationSymbolName =
"__xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation";
extern const char* const kParallelForkJoinSymbolName =
"__xla_cpu_runtime_ParallelForkJoin";
extern const char* const kXlaCpuRuntimeSymbolNamePrefix = "__xla_cpu_runtime_";
} // namespace runtime
} // namespace cpu
} // namespace xla
namespace {
tensorflow::string ShapeString(const void* shape_ptr, xla::int32 shape_length) {
xla::StatusOr<xla::Shape> shape =
xla::llvm_ir::DecodeSelfDescribingShapeConstant(shape_ptr, shape_length);
if (shape.ok()) {
return xla::ShapeUtil::HumanStringWithLayout(shape.ValueOrDie());
}
return "<invalid shape>";
}
} // namespace
void* __xla_cpu_runtime_AcquireInfeedBufferForDequeue(xla::int32 buffer_length,
const void* shape,
xla::int32 shape_length) {
if (VLOG_IS_ON(2)) {
LOG(INFO) << "AcquireInfeedBufferForDequeue: "
<< ShapeString(shape, shape_length);
}
xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager();
// Wait until there's a buffer to dequeue.
xla::cpu::runtime::XfeedBuffer* buffer =
xfeed->infeed()->BlockingDequeueBuffer();
CHECK_EQ(buffer->length(), buffer_length)
<< "XLA program infeed request buffer size " << buffer_length
<< " did not match the runtime's infed buffer length " << buffer->length()
<< "; program reports desired shape: "
<< ShapeString(shape, shape_length);
return buffer->data();
}
void __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue(
xla::int32 buffer_length, void* buffer_ptr, const void* shape_ptr,
xla::int32 shape_length) {
if (VLOG_IS_ON(2)) {
LOG(INFO) << "ReleaseInfeedBufferAfterDeque: "
<< ShapeString(shape_ptr, shape_length);
}
xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager();
xla::StatusOr<xla::Shape> shape =
xla::llvm_ir::DecodeSelfDescribingShapeConstant(shape_ptr, shape_length);
xfeed->infeed()->ReleaseCurrentBuffer(buffer_length, buffer_ptr,
std::move(shape));
}
void* __xla_cpu_runtime_AcquireOutfeedBufferForPopulation(
xla::int32 buffer_length, const void* shape_ptr, xla::int32 shape_length) {
if (VLOG_IS_ON(2)) {
LOG(INFO) << "AcquireOutfeedBufferForPopulation: "
<< ShapeString(shape_ptr, shape_length);
}
xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager();
// Wait until there's a buffer to dequeue.
xla::cpu::runtime::XfeedBuffer* buffer =
xfeed->outfeed()->BlockingDequeueBuffer();
CHECK_EQ(buffer->length(), buffer_length)
<< "XLA program outfeed request buffer size " << buffer_length
<< " did not match the runtime's outfeed buffer length "
<< buffer->length() << "; program reports outfed shape: "
<< ShapeString(shape_ptr, shape_length);
return buffer->data();
}
void __xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation(
xla::int32 buffer_length, void* buffer_ptr, const void* shape_ptr,
xla::int32 shape_length) {
if (VLOG_IS_ON(2)) {
LOG(INFO) << "ReleaseOutfeedBufferAfterPopulation: "
<< ShapeString(shape_ptr, shape_length);
}
xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager();
xla::StatusOr<xla::Shape> shape =
xla::llvm_ir::DecodeSelfDescribingShapeConstant(shape_ptr, shape_length);
xfeed->outfeed()->ReleaseCurrentBuffer(buffer_length, buffer_ptr,
std::move(shape));
}
|
