18cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 28cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 38cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlowerLicensed under the Apache License, Version 2.0 (the "License"); 48cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFloweryou may not use this file except in compliance with the License. 58cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlowerYou may obtain a copy of the License at 68cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 78cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower http://www.apache.org/licenses/LICENSE-2.0 88cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 98cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlowerUnless required by applicable law or agreed to in writing, software 108cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlowerdistributed under the License is distributed on an "AS IS" BASIS, 118cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlowerWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 128cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlowerSee the License for the specific language governing permissions and 138cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlowerlimitations under the License. 148cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower==============================================================================*/ 158cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 168cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_TRANSFER_MANAGER_H_ 178cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_TRANSFER_MANAGER_H_ 188cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 198cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#include <vector> 208cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 218cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#include "tensorflow/compiler/xla/service/generic_transfer_manager.h" 221856fe120d015fb17be303ef2da3873a10a6ffa6A. Unique TensorFlower#include "tensorflow/compiler/xla/service/gpu/infeed_manager.h" 238cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#include "tensorflow/compiler/xla/service/transfer_manager.h" 248cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#include "tensorflow/compiler/xla/statusor.h" 258cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#include "tensorflow/compiler/xla/xla_data.pb.h" 268cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#include "tensorflow/core/platform/macros.h" 278cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#include "tensorflow/core/platform/stream_executor_no_cuda.h" 288cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#include "tensorflow/core/platform/types.h" 298cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 308cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlowernamespace xla { 318cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 328cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower// An implementation of the XLA GenericTransferManager that 338cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower// handles GPU-specific infeed. 348cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlowerclass GpuTransferManager : public GenericTransferManager { 358cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower public: 368cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower GpuTransferManager(); 378cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower ~GpuTransferManager() override {} 388cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 398cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower Status TransferLiteralToInfeed(perftools::gputools::StreamExecutor* executor, 408cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower const Literal& literal) override; 41b52debb4e63cce1e0733d6d34975d4efb9934680Jacques Pienaar Status TransferBufferToInfeed(perftools::gputools::StreamExecutor* executor, 42b52debb4e63cce1e0733d6d34975d4efb9934680Jacques Pienaar int64 size, const void* source) override; 438cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 448cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower private: 45f8fc078a4ec357abbe5cb8c6bc4b4c2f546e31e6A. Unique TensorFlower // Initiates the infeed data transfers. InfeedBuffer->Done() must be 46f8fc078a4ec357abbe5cb8c6bc4b4c2f546e31e6A. Unique TensorFlower // called to clean up the memory allocated for InfeedBuffer. 47b52debb4e63cce1e0733d6d34975d4efb9934680Jacques Pienaar StatusOr<gpu::InfeedBuffer*> TransferBufferToInfeedInternal( 48b52debb4e63cce1e0733d6d34975d4efb9934680Jacques Pienaar perftools::gputools::StreamExecutor* executor, int64 size, 49b52debb4e63cce1e0733d6d34975d4efb9934680Jacques Pienaar const void* source); 50b52debb4e63cce1e0733d6d34975d4efb9934680Jacques Pienaar 51f8fc078a4ec357abbe5cb8c6bc4b4c2f546e31e6A. Unique TensorFlower // Enqueues infeed data buffers with the infeed manager after their 52f8fc078a4ec357abbe5cb8c6bc4b4c2f546e31e6A. Unique TensorFlower // transfer completes. 53f8fc078a4ec357abbe5cb8c6bc4b4c2f546e31e6A. Unique TensorFlower Status EnqueueBuffersToInfeed(perftools::gputools::StreamExecutor* executor, 54f8fc078a4ec357abbe5cb8c6bc4b4c2f546e31e6A. Unique TensorFlower std::vector<gpu::InfeedBuffer*> buffers); 55f8fc078a4ec357abbe5cb8c6bc4b4c2f546e31e6A. Unique TensorFlower 568cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower TF_DISALLOW_COPY_AND_ASSIGN(GpuTransferManager); 578cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower}; 588cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 598cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower} // namespace xla 608cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower 618cb5e9867482a8e05f756fad35634e1674fe7f16A. Unique TensorFlower#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_TRANSFER_MANAGER_H_ 62