experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add TpuExecutorInterface::GetCoreLocationExternal() to TPU API.

Browse Source 
PiperOrigin-RevId: 327082333
Change-Id: Ifd130fd49d635a137f407b28fe89243126cef800
This commit is contained in:
Skye Wanderman-Milne 2020-08-17 13:05:45 -07:00 committed by TensorFlower Gardener
parent 6684bae7c6
commit f8f221e367
6 changed files with 17 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/core/tpu/tpu_library_init_fns.inc
View File

@ -96,6 +96,7 @@ tensorflow::Status SetExecutorStructFn(void* library_handle) {
TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateStream);
TFTPU_SET_FN(executor_fn, TpuExecutor_CreateStreamDependency);
TFTPU_SET_FN(executor_fn, TpuExecutor_GetStatus);
TFTPU_SET_FN(executor_fn, TpuExecutor_GetCoreLocation);
TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateEvent);
TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateEvent);
TFTPU_SET_FN(executor_fn, TpuExecutor_PollForEventStatus);

1
tensorflow/stream_executor/tpu/BUILD
View File

@ -308,6 +308,7 @@ cc_library(
visibility = ["//visibility:public"],
deps = [
":tpu_platform_interface",
":tpu_topology_external",
"//tensorflow/stream_executor:stream_executor_headers",
"//tensorflow/stream_executor/lib",
],

5
tensorflow/stream_executor/tpu/tpu_executor.cc
View File

@ -80,6 +80,11 @@ Status TpuExecutor::GetStatus(Stream* stream) {
return status.status();
}
tpu::TpuCoreLocationExternal TpuExecutor::GetCoreLocationExternal() const {
return tpu::TpuCoreLocationExternal(
tpu::ExecutorApiFn()->TpuExecutor_GetCoreLocationFn(executor_));
}
bool TpuExecutor::AllocateStream(Stream* stream) {
return tpu::ExecutorApiFn()->TpuExecutor_AllocateStreamFn(
executor_, stream_map().at(stream->implementation()));

2
tensorflow/stream_executor/tpu/tpu_executor.h
View File

@ -100,6 +100,8 @@ class TpuExecutor : public tensorflow::tpu::TpuExecutorInterface {
absl::optional<stream_executor::AllocatorStats> GetAllocatorStats() override;
tpu::TpuCoreLocationExternal GetCoreLocationExternal() const override;
Status GetStatus(Stream* stream) override;
std::unique_ptr<::stream_executor::internal::StreamInterface>

3
tensorflow/stream_executor/tpu/tpu_executor_c_api.h
View File

@ -65,6 +65,8 @@ bool TpuExecutor_CreateStreamDependency(SE_StreamExecutor* executor,
void TpuExecutor_GetStatus(SE_StreamExecutor* executor, SE_Stream* stream,
SE_Status* status);
SE_TpuTopology_Core* TpuExecutor_GetCoreLocation(SE_StreamExecutor* executor);
void TpuExecutor_AllocateEvent(SE_StreamExecutor* executor, SE_Event* event,
SE_Status* status);
void TpuExecutor_DeallocateEvent(SE_StreamExecutor* executor, SE_Event* event,
@ -304,6 +306,7 @@ struct TfTpu_ExecutorApiFn {
TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_DeallocateStream);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_CreateStreamDependency);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_GetStatus);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_GetCoreLocation);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_AllocateEvent);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_DeallocateEvent);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_PollForEventStatus);

5
tensorflow/stream_executor/tpu/tpu_executor_interface.h
View File

@ -22,6 +22,7 @@ limitations under the License.
#include "tensorflow/stream_executor/lib/statusor.h"
#include "tensorflow/stream_executor/stream_executor_internal.h"
#include "tensorflow/stream_executor/tpu/tpu_platform_interface.h"
#include "tensorflow/stream_executor/tpu/tpu_topology.h"
namespace tpu {
class TpuCore;
@ -53,6 +54,10 @@ class TpuExecutorInterface
}
virtual TpuPlatformInterface& platform() { LOG(FATAL) << "Unimplemented."; }
virtual TpuCoreLocationExternal GetCoreLocationExternal() const {
LOG(FATAL) << "Unimplemented.";
}
};
} // namespace tpu