1 | /* Copyright 2015 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 | // See docs in ../ops/parse_ops.cc. |
17 | |
18 | #include <algorithm> |
19 | |
20 | #include "tensorflow/core/framework/op_kernel.h" |
21 | #include "tensorflow/core/framework/tensor.h" |
22 | #include "tensorflow/core/framework/tensor_shape.h" |
23 | #include "tensorflow/core/framework/types.h" |
24 | #include "tensorflow/core/lib/core/errors.h" |
25 | #include "tensorflow/core/platform/bfloat16.h" |
26 | #include "tensorflow/core/platform/byte_order.h" |
27 | |
28 | namespace tensorflow { |
29 | |
30 | template <typename T> |
31 | class DecodeRawOp : public OpKernel { |
32 | public: |
33 | explicit DecodeRawOp(OpKernelConstruction* context) : OpKernel(context) { |
34 | OP_REQUIRES_OK(context, context->GetAttr("out_type" , &out_type_)); |
35 | |
36 | const bool host_is_little_endian = port::kLittleEndian; |
37 | bool data_is_little_endian; |
38 | OP_REQUIRES_OK(context, |
39 | context->GetAttr("little_endian" , &data_is_little_endian)); |
40 | convert_data_endianness_ = host_is_little_endian != data_is_little_endian; |
41 | } |
42 | |
43 | void Compute(OpKernelContext* context) override { |
44 | const auto& input = context->input(0); |
45 | int64_t str_size = -1; |
46 | auto flat_in = input.flat<tstring>(); |
47 | for (int64_t i = 0; i < flat_in.size(); ++i) { |
48 | const tstring& in_str = flat_in(i); |
49 | if (str_size == -1) { |
50 | str_size = in_str.size(); |
51 | } else { |
52 | OP_REQUIRES(context, str_size == in_str.size(), |
53 | errors::InvalidArgument( |
54 | "DecodeRaw requires input strings to all be the same " |
55 | "size, but element " , |
56 | i, " has size " , str_size, " != " , in_str.size())); |
57 | } |
58 | } |
59 | TensorShape out_shape = input.shape(); |
60 | if (str_size == -1 || str_size == 0) { // Empty input |
61 | out_shape.AddDim(0); |
62 | Tensor* output_tensor = nullptr; |
63 | OP_REQUIRES_OK(context, context->allocate_output("output" , out_shape, |
64 | &output_tensor)); |
65 | return; |
66 | } |
67 | OP_REQUIRES( |
68 | context, str_size % sizeof(T) == 0, |
69 | errors::InvalidArgument("Input to DecodeRaw has length " , str_size, |
70 | " that is not a multiple of " , sizeof(T), |
71 | ", the size of " , DataTypeString(out_type_))); |
72 | const int64_t added_dim = str_size / sizeof(T); |
73 | out_shape.AddDim(added_dim); |
74 | Tensor* output_tensor = nullptr; |
75 | OP_REQUIRES_OK( |
76 | context, context->allocate_output("output" , out_shape, &output_tensor)); |
77 | auto out = output_tensor->flat_inner_dims<T>(); |
78 | DCHECK_EQ(flat_in.size(), out.dimensions()[0]); |
79 | T* out_data = out.data(); |
80 | |
81 | // If the data is already in the host's byte order, or if the width of the |
82 | // output type is a single byte, we can copy the memory directly. |
83 | if (!convert_data_endianness_ || sizeof(T) == 1) { |
84 | for (int64_t i = 0; i < flat_in.size(); ++i) { |
85 | const T* in_data = reinterpret_cast<const T*>(flat_in(i).data()); |
86 | memcpy(out_data, in_data, str_size); |
87 | out_data += added_dim; |
88 | } |
89 | } else { |
90 | // Otherwise, the data is not in the host's byte order, and rather than a |
91 | // direct copy, we need to reverse the byte ordering of each element. |
92 | int64_t element_size; |
93 | if (out_type_ == DT_COMPLEX64 || out_type_ == DT_COMPLEX128) { |
94 | // For Complex data type, real and imaginary parts need to be swapped |
95 | // separately |
96 | element_size = sizeof(T) / 2; |
97 | } else { |
98 | element_size = sizeof(T); |
99 | } |
100 | for (int64_t i = 0; i < flat_in.size(); ++i) { |
101 | const char* in_data_bytes = |
102 | reinterpret_cast<const char*>(flat_in(i).data()); |
103 | char* out_data_bytes = reinterpret_cast<char*>(out_data); |
104 | const char* p = in_data_bytes; |
105 | char* q = out_data_bytes; |
106 | for (; p < in_data_bytes + str_size; |
107 | p += element_size, q += element_size) { |
108 | std::reverse_copy(p, p + element_size, q); |
109 | } |
110 | out_data += added_dim; |
111 | } |
112 | } |
113 | } |
114 | |
115 | private: |
116 | // True if the endianness of the data and the endianness of the host are |
117 | // different, and the data needs conversion. |
118 | bool convert_data_endianness_; |
119 | |
120 | // True if the input data is in little endian format. |
121 | bool data_is_little_endian_; |
122 | DataType out_type_; |
123 | }; |
124 | |
125 | #define REGISTER(type) \ |
126 | REGISTER_KERNEL_BUILDER( \ |
127 | Name("DecodeRaw").Device(DEVICE_CPU).TypeConstraint<type>("out_type"), \ |
128 | DecodeRawOp<type>) |
129 | |
130 | REGISTER(Eigen::half); |
131 | REGISTER(float); |
132 | REGISTER(double); |
133 | REGISTER(int32); |
134 | REGISTER(uint16); |
135 | REGISTER(uint8); |
136 | REGISTER(int16); |
137 | REGISTER(int8); |
138 | REGISTER(int64_t); |
139 | REGISTER(bool); |
140 | REGISTER(complex64); |
141 | REGISTER(complex128); |
142 | REGISTER(bfloat16); |
143 | |
144 | #undef REGISTER |
145 | |
146 | } // namespace tensorflow |
147 | |