1 /* Copyright 2017 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 // Utilities for working with XLA Literals. 17 18 #ifndef TENSORFLOW_COMPILER_TF2XLA_LITERAL_UTIL_H_ 19 #define TENSORFLOW_COMPILER_TF2XLA_LITERAL_UTIL_H_ 20 21 #include "absl/types/span.h" 22 #include "tensorflow/compiler/xla/literal.h" 23 #include "tensorflow/compiler/xla/xla_data.pb.h" 24 #include "tensorflow/core/framework/tensor.h" 25 #include "tensorflow/core/lib/core/status.h" 26 27 namespace tensorflow { 28 29 // Returns a BorrowingLiteral that utilizes the same underlying buffer owned by 30 // 'host_tensor'. 31 Status HostTensorToBorrowingLiteral(const Tensor& host_tensor, 32 xla::BorrowingLiteral* literal); 33 // Similar as above, except the literal shape is explicitly provided and used 34 // instead of obtaining it from the 'host_tensor'. The provided literal shape 35 // 'xla_shape' must be compatible with the shape of 'host_tensor'. 36 Status HostTensorToBorrowingLiteral(const xla::Shape& xla_shape, 37 const Tensor& host_tensor, 38 xla::BorrowingLiteral* literal); 39 40 // Returns a Literal with the contents of 'host_tensor', backed by its own 41 // storage (i.e., not reusing 'host_tensor's buffers.) 42 StatusOr<xla::Literal> HostTensorToLiteral(const Tensor& host_tensor); 43 44 // Returns a MutableBorrowingLiteral that utilizes the same underlying buffer 45 // owned by 'host_tensor', but is mutable via the xla::Literal methods. 46 Status HostTensorToMutableBorrowingLiteral( 47 Tensor* host_tensor, xla::MutableBorrowingLiteral* literal); 48 // Similar as above, except the literal shape is explicitly provided and used 49 // instead of obtaining it from the 'host_tensor'. The provided literal shape 50 // 'xla_shape' must be compatible with the shape of 'host_tensor'. 51 Status HostTensorToMutableBorrowingLiteral( 52 const xla::Shape& xla_shape, Tensor* host_tensor, 53 xla::MutableBorrowingLiteral* literal); 54 55 // Returns a BorrowingLiteral tuple that utilizes the same underlying buffers 56 // owned by 'host_tensors'. 57 Status HostTensorsToBorrowingLiteralTuple(absl::Span<const Tensor> host_tensors, 58 xla::BorrowingLiteral* literal); 59 60 // Copies 'literal' to freshly allocated 'host_tensor', which is allocated of 61 // type <target_type>. 62 // Fails if the literal's primitive type != 63 // DataTypeToPrimitiveType(target_type). Note that <target_type> is not 64 // derivable from the type of <literal>, because multiple tensorflow types map 65 // to the same XLA type (e.g. INT32 and QINT32 both map to INT32 in 66 // XLA). 67 Status LiteralToHostTensor(const xla::LiteralSlice& literal, 68 DataType target_type, Tensor* host_tensor); 69 70 // Copies the contents of 'literal' to a previously allocated tensor 71 // 'host_tensor'. The tensor and the literal must have the same number of 72 // elements and the same type. 73 Status CopyLiteralToHostTensor(const xla::LiteralSlice& literal, 74 Tensor* host_tensor); 75 76 } // namespace tensorflow 77 78 #endif // TENSORFLOW_COMPILER_TF2XLA_LITERAL_UTIL_H_ 79