11e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 21e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 31e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsLicensed under the Apache License, Version 2.0 (the "License"); 41e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinsyou may not use this file except in compliance with the License. 51e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsYou may obtain a copy of the License at 61e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 71e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins http://www.apache.org/licenses/LICENSE-2.0 81e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 91e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsUnless required by applicable law or agreed to in writing, software 101e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinsdistributed under the License is distributed on an "AS IS" BASIS, 111e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 121e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsSee the License for the specific language governing permissions and 131e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinslimitations under the License. 141e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins==============================================================================*/ 151e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 161e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GENERIC_TRANSFER_MANAGER_H_ 171e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#define TENSORFLOW_COMPILER_XLA_SERVICE_GENERIC_TRANSFER_MANAGER_H_ 181e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 191e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include <vector> 201e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 211e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include "tensorflow/compiler/xla/service/transfer_manager.h" 221e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include "tensorflow/compiler/xla/statusor.h" 231e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include "tensorflow/compiler/xla/xla_data.pb.h" 241e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include "tensorflow/core/platform/macros.h" 251e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include "tensorflow/core/platform/stream_executor_no_cuda.h" 261e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include "tensorflow/core/platform/types.h" 271e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 281e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinsnamespace xla { 291e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 301e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins// A generic implementation of the XLA TransferManager interface 311e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins// that is the base class for both CPU and GPU. For GPU, it transfers 321e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins// data between host and device (GPU). For CPU, since the "device" 331e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins// here is the host itself, there's not much for this transfer manager 341e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins// to do except memcpy the result. There is a CpuTransferManager that 351e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins// inherits from GenericTransferManager and handles CPU-specific 361e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins// infeed. 371e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinsclass GenericTransferManager : public TransferManager { 381e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins public: 3906deeea373c93ea36547648481c5daf4dc56126fMark Heffernan GenericTransferManager(perftools::gputools::Platform::Id platform_id, 4006deeea373c93ea36547648481c5daf4dc56126fMark Heffernan size_t pointer_size); 411e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins ~GenericTransferManager() override {} 421e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 431e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins perftools::gputools::Platform::Id PlatformId() const override; 441e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 4522d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan StatusOr<std::unique_ptr<Literal>> TransferLiteralFromDevice( 4622d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan perftools::gputools::StreamExecutor* executor, 4722d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan const ShapedBuffer& device_buffer) override; 4822d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan 4922d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan Status TransferLiteralToDevice(perftools::gputools::StreamExecutor* executor, 5022d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan const Literal& literal, 5122d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan const ShapedBuffer& device_buffer) override; 5222d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan 531e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins Status TransferLiteralToInfeed(perftools::gputools::StreamExecutor* executor, 541e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins const Literal& literal) override; 55efc8f98d45df835bac2373e19f1da57e3a1ea2d0Jacques Pienaar Status TransferLiteralFromOutfeed( 56efc8f98d45df835bac2373e19f1da57e3a1ea2d0Jacques Pienaar perftools::gputools::StreamExecutor* executor, const Shape& literal_shape, 57efc8f98d45df835bac2373e19f1da57e3a1ea2d0Jacques Pienaar Literal* literal) override; 58efc8f98d45df835bac2373e19f1da57e3a1ea2d0Jacques Pienaar 5999e1b19ceba32b8354dddc2841b81864c9ba96bbJacques Pienaar Status ResetDevices( 6099e1b19ceba32b8354dddc2841b81864c9ba96bbJacques Pienaar tensorflow::gtl::ArraySlice<perftools::gputools::StreamExecutor*> 6199e1b19ceba32b8354dddc2841b81864c9ba96bbJacques Pienaar executors) override; 621e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 6322d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan int64 GetByteSizeRequirement(const Shape& shape) const override; 6422d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan 6522d948d2739ecaadfb4091302f2050ba9cf0d0c1Mark Heffernan protected: 66fc2526a8c1cf0bc2a93c8cc819ff7209eb4628c9A. Unique TensorFlower Status TransferBufferToInfeed(perftools::gputools::StreamExecutor* executor, 67fc2526a8c1cf0bc2a93c8cc819ff7209eb4628c9A. Unique TensorFlower int64 size, const void* source) override; 68fc2526a8c1cf0bc2a93c8cc819ff7209eb4628c9A. Unique TensorFlower 69fc2526a8c1cf0bc2a93c8cc819ff7209eb4628c9A. Unique TensorFlower Status WriteSingleTupleIndexTable( 7006deeea373c93ea36547648481c5daf4dc56126fMark Heffernan perftools::gputools::StreamExecutor* executor, 7106deeea373c93ea36547648481c5daf4dc56126fMark Heffernan tensorflow::gtl::ArraySlice<perftools::gputools::DeviceMemoryBase> 7206deeea373c93ea36547648481c5daf4dc56126fMark Heffernan elements, 7306deeea373c93ea36547648481c5daf4dc56126fMark Heffernan const Shape& shape, 7406deeea373c93ea36547648481c5daf4dc56126fMark Heffernan perftools::gputools::DeviceMemoryBase* region) override; 7506deeea373c93ea36547648481c5daf4dc56126fMark Heffernan 761e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins private: 771e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins // The platform this transfer manager targets. 787de939bb74c5edbc2f45e77a5d4696e70bb59e5bKay Zhu const perftools::gputools::Platform::Id platform_id_; 791e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 8006deeea373c93ea36547648481c5daf4dc56126fMark Heffernan // The size in bytes of pointers on this platform. 8106deeea373c93ea36547648481c5daf4dc56126fMark Heffernan const size_t pointer_size_; 8206deeea373c93ea36547648481c5daf4dc56126fMark Heffernan 831e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins TF_DISALLOW_COPY_AND_ASSIGN(GenericTransferManager); 841e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins}; 851e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 861e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins} // namespace xla 871e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 881e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GENERIC_TRANSFER_MANAGER_H_ 89