1 /* Copyright 2021 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/mlir/tensorflow/ir/tf_ops_layout_helper.h"
17
18 namespace mlir {
19 namespace TF {
20
ReversePermutation(ArrayRef<int64_t> permutation)21 SmallVector<int64_t, 4> ReversePermutation(ArrayRef<int64_t> permutation) {
22 SmallVector<int64_t, 4> reverse(permutation.size());
23 for (size_t i = 0; i < permutation.size(); ++i) {
24 reverse[permutation[i]] = i;
25 }
26 return reverse;
27 }
28
GetDataFormatPermutation(StringRef from,StringRef to)29 SmallVector<int64_t, 4> GetDataFormatPermutation(StringRef from, StringRef to) {
30 if (from == "NHWC" && to == "NCHW") {
31 return {0, 3, 1, 2};
32 } else if (from == "NCHW" && to == "NHWC") {
33 return {0, 2, 3, 1};
34 } else {
35 return {};
36 }
37 }
38
39 // Shuffle elements in the `attr` according to the permutation. Optional
40 // `inner_size` allows to shuffle array attributes created from rank 2 tensors
41 // on outer dimension only.
ShuffleArrayAttr(ArrayAttr attr,ArrayRef<int64_t> permutation,int inner_size)42 ArrayAttr ShuffleArrayAttr(ArrayAttr attr, ArrayRef<int64_t> permutation,
43 int inner_size) {
44 if (attr.empty()) return attr;
45
46 assert(attr.size() % inner_size == 0);
47 assert(attr.size() / inner_size == permutation.size());
48
49 SmallVector<Attribute, 8> values{attr.begin(), attr.end()};
50 SmallVector<Attribute, 8> shuffled(values.size());
51
52 for (size_t i = 0; i < permutation.size(); ++i) {
53 for (size_t j = 0; j < inner_size; ++j) {
54 shuffled[i * inner_size + j] = values[permutation[i] * inner_size + j];
55 }
56 }
57
58 return ArrayAttr::get(attr.getContext(), shuffled);
59 }
60
61 // Shuffle ranked tensor dimensions according to the permutation.
ShuffleRankedTensorType(Type type,ArrayRef<int64_t> permutation)62 Type ShuffleRankedTensorType(Type type, ArrayRef<int64_t> permutation) {
63 if (auto ranked_type = type.dyn_cast<RankedTensorType>()) {
64 ArrayRef<int64_t> shape = ranked_type.getShape();
65 assert(permutation.size() == shape.size());
66
67 SmallVector<int64_t, 4> new_shape(permutation.size());
68 for (size_t i = 0; i < permutation.size(); ++i)
69 new_shape[i] = shape[permutation[i]];
70
71 return RankedTensorType::get(new_shape, ranked_type.getElementType());
72 }
73
74 return type;
75 }
76
AreCancellablePermutations(DenseIntElementsAttr perm0,DenseIntElementsAttr perm1)77 bool AreCancellablePermutations(DenseIntElementsAttr perm0,
78 DenseIntElementsAttr perm1) {
79 if (perm0.getNumElements() == 0 || perm1.getNumElements() == 0) return false;
80 if (perm0.getNumElements() != perm1.getNumElements()) return false;
81
82 SmallVector<int64_t, 8> perm0_values;
83 for (const auto &value : perm0.getValues<APInt>())
84 perm0_values.push_back(value.getSExtValue());
85
86 SmallVector<int64_t, 8> perm1_values;
87 for (const auto &value : perm1.getValues<APInt>())
88 perm1_values.push_back(value.getSExtValue());
89
90 for (int i = 0; i < perm0_values.size(); ++i) {
91 if (perm0_values[perm1_values[i]] != i) return false;
92 }
93
94 return true;
95 }
96
97 } // namespace TF
98 } // namespace mlir
99