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