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_CORE_TPU_KERNELS_INFEED_OPS_H_ 17 #define TENSORFLOW_CORE_TPU_KERNELS_INFEED_OPS_H_ 18 19 #include "tensorflow/core/framework/op_kernel.h" 20 #include "tensorflow/core/framework/tensor_shape.h" 21 #include "tensorflow/core/framework/types.pb.h" 22 #include "tensorflow/core/platform/status.h" 23 #include "tensorflow/core/tpu/kernels/transfer_ops.h" 24 25 namespace tensorflow { 26 27 // TODO(b/65200690): Rework this when there is a callback based infeed API to 28 // StreamExecutor. 29 30 // The InfeedEnqueue op is used to deliver data to the device infeed queue. 31 class TpuInfeedEnqueueOp : public TpuTransferAsyncOpKernel { 32 public: 33 explicit TpuInfeedEnqueueOp( 34 OpKernelConstruction* ctx, 35 std::unique_ptr<TpuTransferOpInterface> transfer_op); 36 Status DoWork(OpKernelContext* ctx, int device_ordinal) override; 37 38 private: 39 TensorShape shape_; 40 DataType dtype_; 41 xla::Shape xla_shape_; 42 43 // TpuInfeedEnqueueOp is neither copyable nor movable. 44 TpuInfeedEnqueueOp(const TpuInfeedEnqueueOp&) = delete; 45 TpuInfeedEnqueueOp& operator=(const TpuInfeedEnqueueOp&) = delete; 46 }; 47 48 // The InfeedEnqueueTuple op is used on the host to deliver multiple tensors to 49 // the device infeed queue as an XLA tuple. 50 class TpuInfeedEnqueueTupleOp : public TpuTransferAsyncOpKernel { 51 public: 52 explicit TpuInfeedEnqueueTupleOp( 53 OpKernelConstruction* ctx, 54 std::unique_ptr<TpuTransferOpInterface> transfer_op); 55 Status DoWork(OpKernelContext* ctx, int device_ordinal) override; 56 57 private: 58 std::vector<TensorShape> shapes_; 59 DataTypeVector dtypes_; 60 xla::Shape tuple_shape_; 61 62 // TpuInfeedEnqueueTupleOp is neither copyable nor movable. 63 TpuInfeedEnqueueTupleOp(const TpuInfeedEnqueueTupleOp&) = delete; 64 TpuInfeedEnqueueTupleOp& operator=(const TpuInfeedEnqueueTupleOp&) = delete; 65 }; 66 67 // The InfeedEnqueuePrelinearizedBufferOp op is used to transfer prelinearized 68 // buffers to the device infeed queue. 69 class InfeedEnqueuePrelinearizedBufferOp : public TpuTransferAsyncOpKernel { 70 public: 71 explicit InfeedEnqueuePrelinearizedBufferOp( 72 OpKernelConstruction* ctx, 73 std::unique_ptr<TpuTransferOpInterface> transfer_op); 74 75 Status DoWork(OpKernelContext* ctx, int device_ordinal) override; 76 77 private: 78 // InfeedEnqueuePrelinearizedBufferOp is neither copyable nor movable. 79 InfeedEnqueuePrelinearizedBufferOp( 80 const InfeedEnqueuePrelinearizedBufferOp&) = delete; 81 InfeedEnqueuePrelinearizedBufferOp& operator=( 82 const InfeedEnqueuePrelinearizedBufferOp&) = delete; 83 }; 84 } // namespace tensorflow 85 86 #endif // TENSORFLOW_CORE_TPU_KERNELS_INFEED_OPS_H_ 87