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/service/interpreter/executor.h"
17
18 #include <cstring>
19
20 #include "tensorflow/compiler/xla/status_macros.h"
21
22 namespace stream_executor {
23 namespace interpreter {
24
AsExecutorStream(Stream * stream)25 host::HostStream *AsExecutorStream(Stream *stream) {
26 DCHECK(stream != nullptr);
27 return dynamic_cast<host::HostStream *>(stream->implementation());
28 }
29
XlaInterpreterExecutor(const PluginConfig & plugin_config)30 XlaInterpreterExecutor::XlaInterpreterExecutor(
31 const PluginConfig &plugin_config)
32 : plugin_config_(plugin_config) {}
33
~XlaInterpreterExecutor()34 XlaInterpreterExecutor::~XlaInterpreterExecutor() {}
35
Allocate(uint64 size)36 void *XlaInterpreterExecutor::Allocate(uint64 size) { return new char[size]; }
37
AllocateSubBuffer(DeviceMemoryBase * parent,uint64 offset_bytes,uint64)38 void *XlaInterpreterExecutor::AllocateSubBuffer(DeviceMemoryBase *parent,
39 uint64 offset_bytes,
40 uint64 /*size_bytes*/) {
41 return parent + offset_bytes;
42 }
43
Deallocate(DeviceMemoryBase * mem)44 void XlaInterpreterExecutor::Deallocate(DeviceMemoryBase *mem) {
45 if (!mem->is_sub_buffer()) {
46 delete[] static_cast<char *>(mem->opaque());
47 }
48 }
49
Memcpy(Stream * stream,void * host_dst,const DeviceMemoryBase & dev_src,uint64 size)50 bool XlaInterpreterExecutor::Memcpy(Stream *stream, void *host_dst,
51 const DeviceMemoryBase &dev_src,
52 uint64 size) {
53 AsExecutorStream(stream)->EnqueueTask([this, host_dst, dev_src, size]() {
54 port::Status ok = SynchronousMemcpy(host_dst, dev_src, size);
55 });
56 AsExecutorStream(stream)->BlockUntilDone();
57 return true;
58 }
59
Memcpy(Stream * stream,DeviceMemoryBase * dev_dst,const void * host_src,uint64 size)60 bool XlaInterpreterExecutor::Memcpy(Stream *stream, DeviceMemoryBase *dev_dst,
61 const void *host_src, uint64 size) {
62 AsExecutorStream(stream)->EnqueueTask([this, dev_dst, host_src, size]() {
63 port::Status ok = SynchronousMemcpy(dev_dst, host_src, size);
64 });
65 AsExecutorStream(stream)->BlockUntilDone();
66 return true;
67 }
68
SynchronousMemcpy(DeviceMemoryBase * dev_dst,const void * host_src,uint64 size)69 port::Status XlaInterpreterExecutor::SynchronousMemcpy(
70 DeviceMemoryBase *dev_dst, const void *host_src, uint64 size) {
71 memcpy(dev_dst->opaque(), host_src, size);
72 return port::Status::OK();
73 }
74
SynchronousMemcpy(void * host_dst,const DeviceMemoryBase & dev_src,uint64 size)75 port::Status XlaInterpreterExecutor::SynchronousMemcpy(
76 void *host_dst, const DeviceMemoryBase &dev_src, uint64 size) {
77 memcpy(host_dst, dev_src.opaque(), size);
78 return port::Status::OK();
79 }
80
HostCallback(Stream * stream,std::function<port::Status ()> callback)81 bool XlaInterpreterExecutor::HostCallback(
82 Stream *stream, std::function<port::Status()> callback) {
83 AsExecutorStream(stream)->EnqueueTask([callback]() {
84 port::Status s = callback();
85 if (!s.ok()) {
86 LOG(WARNING) << "Host callback failed: " << s;
87 }
88 });
89 return true;
90 }
91
CreateStreamDependency(Stream * dependent,Stream * other)92 bool XlaInterpreterExecutor::CreateStreamDependency(Stream *dependent,
93 Stream *other) {
94 AsExecutorStream(dependent)->EnqueueTask(
95 [other]() { SE_CHECK_OK(other->BlockHostUntilDone()); });
96 AsExecutorStream(dependent)->BlockUntilDone();
97 return true;
98 }
99
StartTimer(Stream * stream,Timer * timer)100 bool XlaInterpreterExecutor::StartTimer(Stream *stream, Timer *timer) {
101 dynamic_cast<host::HostTimer *>(timer->implementation())->Start(stream);
102 return true;
103 }
104
StopTimer(Stream * stream,Timer * timer)105 bool XlaInterpreterExecutor::StopTimer(Stream *stream, Timer *timer) {
106 dynamic_cast<host::HostTimer *>(timer->implementation())->Stop(stream);
107 return true;
108 }
109
BlockHostUntilDone(Stream * stream)110 port::Status XlaInterpreterExecutor::BlockHostUntilDone(Stream *stream) {
111 AsExecutorStream(stream)->BlockUntilDone();
112 return port::Status::OK();
113 }
114
PopulateDeviceDescription() const115 DeviceDescription *XlaInterpreterExecutor::PopulateDeviceDescription() const {
116 internal::DeviceDescriptionBuilder builder;
117
118 builder.set_device_address_bits(64);
119
120 builder.set_name("Interpreter");
121 builder.set_device_memory_size(static_cast<uint64>(4) * 1024 * 1024 * 1024);
122 builder.set_clock_rate_ghz(static_cast<float>(CLOCKS_PER_SEC) / 1e9);
123
124 return builder.Build().release();
125 }
126
127 } // namespace interpreter
128 } // namespace stream_executor
129