1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
28namespace tensorflow {
29
30template <typename T>
31class 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
130REGISTER(Eigen::half);
131REGISTER(float);
132REGISTER(double);
133REGISTER(int32);
134REGISTER(uint16);
135REGISTER(uint8);
136REGISTER(int16);
137REGISTER(int8);
138REGISTER(int64_t);
139REGISTER(bool);
140REGISTER(complex64);
141REGISTER(complex128);
142REGISTER(bfloat16);
143
144#undef REGISTER
145
146} // namespace tensorflow
147