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/mlir/lite/utils/validators.h"
17
18 #include <algorithm>
19
20 #include "mlir/Dialect/Traits.h" // from @llvm-project
21 #include "mlir/IR/Builders.h" // from @llvm-project
22 #include "mlir/IR/BuiltinAttributeInterfaces.h" // from @llvm-project
23
24 namespace mlir {
25 namespace TFL {
26
27 // Returns true if the given `op`
28 // * has an attribute with the given `name`,
29 // * and the attribute is an integer list of the form [1, X, Y, 1],
30 // and writes X, Y as 32-bit integer attribute to `x`, `y`.
TFIntListIs1XY1(Operation * op,StringRef name,IntegerAttr * x,IntegerAttr * y)31 bool TFIntListIs1XY1(Operation *op, StringRef name, IntegerAttr *x,
32 IntegerAttr *y) {
33 auto attr = op->getAttrOfType<ArrayAttr>(name);
34 if (!attr) return false;
35
36 auto elements = attr.getValue();
37 if (elements.size() != 4 ||
38 std::any_of(elements.begin(), elements.end(),
39 [](Attribute e) { return !e.isa<IntegerAttr>(); }))
40 return false;
41
42 if (elements.front().cast<IntegerAttr>().getInt() != 1 ||
43 elements.back().cast<IntegerAttr>().getInt() != 1)
44 return false;
45
46 Builder b(op->getContext());
47 *x = b.getI32IntegerAttr(elements[1].cast<IntegerAttr>().getInt());
48 *y = b.getI32IntegerAttr(elements[2].cast<IntegerAttr>().getInt());
49
50 return true;
51 }
52
53 // Returns true if the attribute is an integer list of the form [1, X, Y, 1],
TFIntListIs1XY1(const Attribute attr)54 bool TFIntListIs1XY1(const Attribute attr) {
55 const auto &elements = attr.cast<ArrayAttr>().getValue();
56 if (elements.size() != 4 ||
57 std::any_of(elements.begin(), elements.end(),
58 [](Attribute e) { return !e.isa<IntegerAttr>(); }))
59 return false;
60
61 if (elements.front().cast<IntegerAttr>().getValue() != 1 ||
62 elements.back().cast<IntegerAttr>().getValue() != 1)
63 return false;
64 return true;
65 }
66
67 // Returns true if the given `op`
68 // * has an attribute with the given `name`,
69 // * and the attribute is an integer list of the form [1, X, Y, Z, 1],
70 // and writes X, Y as 32-bit integer attribute to `x`, `y`, z.
TFIntListIs1XYZ1(Operation * op,StringRef name,IntegerAttr * x,IntegerAttr * y,IntegerAttr * z)71 bool TFIntListIs1XYZ1(Operation *op, StringRef name, IntegerAttr *x,
72 IntegerAttr *y, IntegerAttr *z) {
73 auto attr = op->getAttrOfType<ArrayAttr>(name);
74 if (!attr) return false;
75
76 auto elements = attr.getValue();
77 if (elements.size() != 5 ||
78 std::any_of(elements.begin(), elements.end(),
79 [](Attribute e) { return !e.isa<IntegerAttr>(); }))
80 return false;
81
82 if (elements.front().cast<IntegerAttr>().getInt() != 1 ||
83 elements.back().cast<IntegerAttr>().getInt() != 1)
84 return false;
85
86 Builder b(op->getContext());
87 *x = b.getI32IntegerAttr(elements[1].cast<IntegerAttr>().getInt());
88 *y = b.getI32IntegerAttr(elements[2].cast<IntegerAttr>().getInt());
89 *z = b.getI32IntegerAttr(elements[3].cast<IntegerAttr>().getInt());
90
91 return true;
92 }
93
94 // Returns true if every element of the attribute is 1. All elements of `attr`
95 // must be `IntegerAttr`.
TFIntListIsAllOnes(const Attribute attr)96 bool TFIntListIsAllOnes(const Attribute attr) {
97 const auto &elements = attr.cast<ArrayAttr>().getValue();
98
99 return !std::any_of(elements.begin(), elements.end(), [](Attribute e) {
100 return e.cast<IntegerAttr>().getValue() != 1;
101 });
102 }
103
IsBroadcastableElementsAttrs(mlir::TypedAttr a,mlir::TypedAttr b)104 bool IsBroadcastableElementsAttrs(mlir::TypedAttr a, mlir::TypedAttr b) {
105 // This would return false if we had unranked tensors (where they should
106 // probably be considered as broadcastable), but given we are working with
107 // attributes here that shouldn't be an issue,
108 return OpTrait::util::getBroadcastedType(a.getType(), b.getType()) != Type();
109 }
110
IsDimensionsDegenerateExceptLastOne(ArrayRef<int64_t> elements_shape)111 bool IsDimensionsDegenerateExceptLastOne(ArrayRef<int64_t> elements_shape) {
112 if (elements_shape.empty()) return true;
113
114 for (auto dim : elements_shape.drop_back(1)) {
115 if (dim != 1) return false;
116 }
117 return true;
118 }
119
IsDimensionsDegenerateExceptLastOne(TypedAttr val)120 bool IsDimensionsDegenerateExceptLastOne(TypedAttr val) {
121 if (auto ranked_type = val.getType().dyn_cast<RankedTensorType>()) {
122 return IsDimensionsDegenerateExceptLastOne(ranked_type.getShape());
123 }
124 return false;
125 }
126
127 } // namespace TFL
128 } // namespace mlir
129