1 /* Copyright 2020 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 #ifndef TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_PLATFORM_H_ 17 #define TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_PLATFORM_H_ 18 19 #include <memory> 20 21 #include "absl/container/flat_hash_map.h" 22 #include "tensorflow/core/platform/mutex.h" 23 #include "tensorflow/core/platform/types.h" 24 #include "tensorflow/stream_executor/executor_cache.h" 25 #include "tensorflow/stream_executor/platform.h" 26 #include "tensorflow/stream_executor/stream_executor_internal.h" 27 #include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h" 28 #include "tensorflow/stream_executor/tpu/tpu_platform_interface.h" 29 30 namespace tensorflow { 31 namespace tpu { 32 33 class TpuPlatform : public ::tensorflow::tpu::TpuPlatformInterface { 34 public: 35 using StreamMap = 36 absl::flat_hash_map<stream_executor::internal::StreamInterface*, 37 SE_Stream*>; 38 using EventMap = 39 absl::flat_hash_map<stream_executor::internal::EventInterface*, 40 SE_Event*>; 41 42 static const ::stream_executor::Platform::Id kId; 43 44 using Status = ::stream_executor::port::Status; 45 template <typename T> 46 using StatusOr = ::stream_executor::port::StatusOr<T>; 47 48 TpuPlatform(); 49 50 ~TpuPlatform() override; 51 52 static TpuPlatform* GetRegisteredPlatform(); 53 54 Id id() const override; 55 56 const std::string& Name() const override; 57 58 int VisibleDeviceCount() const override; 59 60 int64 TpuMemoryLimit() override; 61 62 bool ShouldRegisterTpuDeviceToDeviceCopy() override; 63 64 const tensorflow::tpu::TpuTopologyPtr GetTopologyPtr() override; 65 66 const tensorflow::tpu::TpuHostLocationExternal GetTpuHostLocation() 67 const override; 68 69 TpuRuntimeVersion version() const override; 70 71 bool Initialized() const override; 72 73 Status Initialize( 74 const std::map<std::string, std::string>& platform_options) override; 75 Reset(bool only_tear_down,absl::string_view reason)76 Status Reset(bool only_tear_down, absl::string_view reason) override { 77 LOG(FATAL) << "Not yet implemented"; 78 } 79 80 StatusOr<std::unique_ptr<::stream_executor::DeviceDescription>> DescriptionForDevice(int ordinal)81 DescriptionForDevice(int ordinal) const override { 82 LOG(FATAL) << "Not yet implemented"; 83 } 84 ExecutorForDevice(int ordinal)85 StatusOr<::stream_executor::StreamExecutor*> ExecutorForDevice( 86 int ordinal) override { 87 stream_executor::StreamExecutorConfig config; 88 config.ordinal = ordinal; 89 return GetExecutor(config); 90 } 91 92 StatusOr<::stream_executor::StreamExecutor*> ExecutorForDeviceWithPluginConfig(int ordinal,const::stream_executor::PluginConfig & plugin_config)93 ExecutorForDeviceWithPluginConfig( 94 int ordinal, 95 const ::stream_executor::PluginConfig& plugin_config) override { 96 stream_executor::StreamExecutorConfig config; 97 config.ordinal = ordinal; 98 config.plugin_config = plugin_config; 99 return GetExecutor(config); 100 } 101 102 StatusOr<::stream_executor::StreamExecutor*> GetExecutor( 103 const ::stream_executor::StreamExecutorConfig& config) override; 104 105 StatusOr<std::unique_ptr<::stream_executor::StreamExecutor>> 106 GetUncachedExecutor( 107 const ::stream_executor::StreamExecutorConfig& config) override; 108 RegisterTraceListener(std::unique_ptr<stream_executor::TraceListener> listener)109 void RegisterTraceListener( 110 std::unique_ptr<stream_executor::TraceListener> listener) override { 111 LOG(FATAL) << "Not yet implemented"; 112 } 113 UnregisterTraceListener(stream_executor::TraceListener * listener)114 void UnregisterTraceListener( 115 stream_executor::TraceListener* listener) override { 116 LOG(FATAL) << "Not yet implemented"; 117 } 118 stream_map()119 StreamMap* stream_map() { return &stream_map_; } 120 121 void InsertEvent(stream_executor::internal::EventInterface* key, 122 SE_Event* val); 123 SE_Event* LookupEvent(stream_executor::internal::EventInterface* key); LookupStream(stream_executor::internal::StreamInterface * key)124 SE_Stream* LookupStream(stream_executor::internal::StreamInterface* key) { 125 mutex().lock(); 126 auto stream = stream_map_.at(key); 127 mutex().unlock(); 128 return stream; 129 } 130 void EraseEvent(stream_executor::internal::EventInterface* key); 131 se_platform()132 SE_Platform* se_platform() const { return platform_; } 133 134 // Returns the number of TPUs per host. 135 static Status TpusPerHost(int* tpus); 136 137 // Returns the memory capacity of the TPUs on this host. 138 static Status TpuMemoryLimit(int64* memory_limit); 139 mutex()140 tensorflow::mutex& mutex() { return event_map_mu_; } 141 142 private: 143 mutable SE_Platform* platform_; 144 std::string name_; 145 stream_executor::ExecutorCache executor_cache_; 146 StreamMap stream_map_; 147 EventMap event_map_; 148 tensorflow::mutex event_map_mu_; 149 }; 150 151 bool RegisterTpuPlatform(); 152 153 } // namespace tpu 154 } // namespace tensorflow 155 156 #endif // TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_PLATFORM_H_ 157