1 /* Copyright 2019 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 #include "tensorflow/compiler/tf2xla/lib/data_format.h"
17
18 #include "tensorflow/compiler/xla/util.h"
19 #include "tensorflow/core/lib/core/errors.h"
20
21 namespace tensorflow {
22 namespace {
23
Contract(xla::XlaOp input,int64 dim)24 xla::StatusOr<xla::XlaOp> Contract(xla::XlaOp input, int64 dim) {
25 xla::XlaBuilder* builder = input.builder();
26 TF_ASSIGN_OR_RETURN(xla::Shape input_shape, builder->GetShape(input));
27
28 if (input_shape.dimensions().back() != 4) {
29 return errors::InvalidArgument("Expected last dimension to be 4; got ",
30 input_shape.dimensions().back());
31 }
32
33 // Transpose the input so C is directly followed by VECT_C.
34 std::vector<int64> permutation;
35 for (int64 i = 0; i != input_shape.rank() - 1; ++i) {
36 permutation.push_back(i);
37 if (i == dim) {
38 permutation.push_back(input_shape.rank() - 1);
39 }
40 }
41
42 // Now merge the adjacent dimensions with a reshape.
43 std::vector<int64> contracted_shape(input_shape.dimensions().begin(),
44 input_shape.dimensions().end() - 1);
45 contracted_shape[dim] *= 4;
46
47 return xla::Reshape(xla::Transpose(input, permutation), contracted_shape);
48 }
49
Expand(xla::XlaOp input,int64 dim)50 xla::StatusOr<xla::XlaOp> Expand(xla::XlaOp input, int64 dim) {
51 xla::XlaBuilder* builder = input.builder();
52 TF_ASSIGN_OR_RETURN(xla::Shape input_shape, builder->GetShape(input));
53
54 if (input_shape.dimensions(dim) % 4 != 0) {
55 return errors::InvalidArgument(
56 "Expected vectorized dimension to be evenly divisible by 4; got ",
57 input_shape.dimensions(dim));
58 }
59
60 // Split the `dim` into two dimensions with a reshape. The size of the new
61 // dimension is always 4.
62 std::vector<int64> expanded_shape =
63 xla::SpanToVector(input_shape.dimensions());
64 expanded_shape[dim] /= 4;
65 expanded_shape.insert(expanded_shape.begin() + dim + 1, 4);
66
67 // Move the newly created dimension to the end with a transpose.
68 std::vector<int64> permutation;
69 for (int64 i = 0, end = expanded_shape.size(); i != end; ++i) {
70 permutation.push_back(i);
71 if (i == dim) {
72 ++i;
73 }
74 }
75 permutation.push_back(dim + 1);
76
77 return xla::Transpose(xla::Reshape(input, expanded_shape), permutation);
78 }
79
80 } // namespace
81
NCHW_VECT_CToNCHW(xla::XlaOp input)82 xla::StatusOr<xla::XlaOp> NCHW_VECT_CToNCHW(xla::XlaOp input) {
83 return Contract(input, 1);
84 }
85
NCHWToNCHW_VECT_C(xla::XlaOp input)86 xla::StatusOr<xla::XlaOp> NCHWToNCHW_VECT_C(xla::XlaOp input) {
87 return Expand(input, 1);
88 }
89
90 } // namespace tensorflow
91