xref: /llvm-project/mlir/unittests/IR/OperationSupportTest.cpp (revision 195730a6505ac352ae9347010ba701ab5267aad7) 
1 //===- OperationSupportTest.cpp - Operation support unit tests ------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/IR/OperationSupport.h"
10 #include "mlir/IR/Builders.h"
11 #include "mlir/IR/BuiltinTypes.h"
12 #include "llvm/ADT/BitVector.h"
13 #include "llvm/Support/FormatVariadic.h"
14 #include "gtest/gtest.h"
15 
16 using namespace mlir;
17 using namespace mlir::detail;
18 
19 static Operation *createOp(MLIRContext *context,
20                            ArrayRef<Value> operands = llvm::None,
21                            ArrayRef<Type> resultTypes = llvm::None,
22                            unsigned int numRegions = 0) {
23   context->allowUnregisteredDialects();
24   return Operation::create(UnknownLoc::get(context),
25                            OperationName("foo.bar", context), resultTypes,
26                            operands, llvm::None, llvm::None, numRegions);
27 }
28 
29 namespace {
30 TEST(OperandStorageTest, NonResizable) {
31   MLIRContext context;
32   Builder builder(&context);
33 
34   Operation *useOp =
35       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
36   Value operand = useOp->getResult(0);
37 
38   // Create a non-resizable operation with one operand.
39   Operation *user = createOp(&context, operand);
40 
41   // The same number of operands is okay.
42   user->setOperands(operand);
43   EXPECT_EQ(user->getNumOperands(), 1u);
44 
45   // Removing is okay.
46   user->setOperands(llvm::None);
47   EXPECT_EQ(user->getNumOperands(), 0u);
48 
49   // Destroy the operations.
50   user->destroy();
51   useOp->destroy();
52 }
53 
54 TEST(OperandStorageTest, Resizable) {
55   MLIRContext context;
56   Builder builder(&context);
57 
58   Operation *useOp =
59       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
60   Value operand = useOp->getResult(0);
61 
62   // Create a resizable operation with one operand.
63   Operation *user = createOp(&context, operand);
64 
65   // The same number of operands is okay.
66   user->setOperands(operand);
67   EXPECT_EQ(user->getNumOperands(), 1u);
68 
69   // Removing is okay.
70   user->setOperands(llvm::None);
71   EXPECT_EQ(user->getNumOperands(), 0u);
72 
73   // Adding more operands is okay.
74   user->setOperands({operand, operand, operand});
75   EXPECT_EQ(user->getNumOperands(), 3u);
76 
77   // Destroy the operations.
78   user->destroy();
79   useOp->destroy();
80 }
81 
82 TEST(OperandStorageTest, RangeReplace) {
83   MLIRContext context;
84   Builder builder(&context);
85 
86   Operation *useOp =
87       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
88   Value operand = useOp->getResult(0);
89 
90   // Create a resizable operation with one operand.
91   Operation *user = createOp(&context, operand);
92 
93   // Check setting with the same number of operands.
94   user->setOperands(/*start=*/0, /*length=*/1, operand);
95   EXPECT_EQ(user->getNumOperands(), 1u);
96 
97   // Check setting with more operands.
98   user->setOperands(/*start=*/0, /*length=*/1, {operand, operand, operand});
99   EXPECT_EQ(user->getNumOperands(), 3u);
100 
101   // Check setting with less operands.
102   user->setOperands(/*start=*/1, /*length=*/2, {operand});
103   EXPECT_EQ(user->getNumOperands(), 2u);
104 
105   // Check inserting without replacing operands.
106   user->setOperands(/*start=*/2, /*length=*/0, {operand});
107   EXPECT_EQ(user->getNumOperands(), 3u);
108 
109   // Check erasing operands.
110   user->setOperands(/*start=*/0, /*length=*/3, {});
111   EXPECT_EQ(user->getNumOperands(), 0u);
112 
113   // Destroy the operations.
114   user->destroy();
115   useOp->destroy();
116 }
117 
118 TEST(OperandStorageTest, MutableRange) {
119   MLIRContext context;
120   Builder builder(&context);
121 
122   Operation *useOp =
123       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
124   Value operand = useOp->getResult(0);
125 
126   // Create a resizable operation with one operand.
127   Operation *user = createOp(&context, operand);
128 
129   // Check setting with the same number of operands.
130   MutableOperandRange mutableOperands(user);
131   mutableOperands.assign(operand);
132   EXPECT_EQ(mutableOperands.size(), 1u);
133   EXPECT_EQ(user->getNumOperands(), 1u);
134 
135   // Check setting with more operands.
136   mutableOperands.assign({operand, operand, operand});
137   EXPECT_EQ(mutableOperands.size(), 3u);
138   EXPECT_EQ(user->getNumOperands(), 3u);
139 
140   // Check with inserting a new operand.
141   mutableOperands.append({operand, operand});
142   EXPECT_EQ(mutableOperands.size(), 5u);
143   EXPECT_EQ(user->getNumOperands(), 5u);
144 
145   // Check erasing operands.
146   mutableOperands.clear();
147   EXPECT_EQ(mutableOperands.size(), 0u);
148   EXPECT_EQ(user->getNumOperands(), 0u);
149 
150   // Destroy the operations.
151   user->destroy();
152   useOp->destroy();
153 }
154 
155 TEST(OperandStorageTest, RangeErase) {
156   MLIRContext context;
157   Builder builder(&context);
158 
159   Type type = builder.getNoneType();
160   Operation *useOp = createOp(&context, /*operands=*/llvm::None, {type, type});
161   Value operand1 = useOp->getResult(0);
162   Value operand2 = useOp->getResult(1);
163 
164   // Create an operation with operands to erase.
165   Operation *user =
166       createOp(&context, {operand2, operand1, operand2, operand1});
167   llvm::BitVector eraseIndices(user->getNumOperands());
168 
169   // Check erasing no operands.
170   user->eraseOperands(eraseIndices);
171   EXPECT_EQ(user->getNumOperands(), 4u);
172 
173   // Check erasing disjoint operands.
174   eraseIndices.set(0);
175   eraseIndices.set(3);
176   user->eraseOperands(eraseIndices);
177   EXPECT_EQ(user->getNumOperands(), 2u);
178   EXPECT_EQ(user->getOperand(0), operand1);
179   EXPECT_EQ(user->getOperand(1), operand2);
180 
181   // Destroy the operations.
182   user->destroy();
183   useOp->destroy();
184 }
185 
186 TEST(OperationOrderTest, OrderIsAlwaysValid) {
187   MLIRContext context;
188   Builder builder(&context);
189 
190   Operation *containerOp =
191       createOp(&context, /*operands=*/llvm::None, /*resultTypes=*/llvm::None,
192                /*numRegions=*/1);
193   Region &region = containerOp->getRegion(0);
194   Block *block = new Block();
195   region.push_back(block);
196 
197   // Insert two operations, then iteratively add more operations in the middle
198   // of them. Eventually we will insert more than kOrderStride operations and
199   // the block order will need to be recomputed.
200   Operation *frontOp = createOp(&context);
201   Operation *backOp = createOp(&context);
202   block->push_back(frontOp);
203   block->push_back(backOp);
204 
205   // Chosen to be larger than Operation::kOrderStride.
206   int kNumOpsToInsert = 10;
207   for (int i = 0; i < kNumOpsToInsert; ++i) {
208     Operation *op = createOp(&context);
209     block->getOperations().insert(backOp->getIterator(), op);
210     ASSERT_TRUE(op->isBeforeInBlock(backOp));
211     // Note verifyOpOrder() returns false if the order is valid.
212     ASSERT_FALSE(block->verifyOpOrder());
213   }
214 
215   containerOp->destroy();
216 }
217 
218 TEST(OperationFormatPrintTest, CanUseVariadicFormat) {
219   MLIRContext context;
220   Builder builder(&context);
221 
222   Operation *op = createOp(&context);
223 
224   std::string str = formatv("{0}", *op).str();
225   ASSERT_STREQ(str.c_str(), "\"foo.bar\"() : () -> ()");
226 
227   op->destroy();
228 }
229 
230 TEST(NamedAttrListTest, TestAppendAssign) {
231   MLIRContext ctx;
232   NamedAttrList attrs;
233   Builder b(&ctx);
234 
235   attrs.append("foo", b.getStringAttr("bar"));
236   attrs.append("baz", b.getStringAttr("boo"));
237 
238   {
239     auto it = attrs.begin();
240     EXPECT_EQ(it->first, b.getStringAttr("foo"));
241     EXPECT_EQ(it->second, b.getStringAttr("bar"));
242     ++it;
243     EXPECT_EQ(it->first, b.getStringAttr("baz"));
244     EXPECT_EQ(it->second, b.getStringAttr("boo"));
245   }
246 
247   attrs.append("foo", b.getStringAttr("zoo"));
248   {
249     auto dup = attrs.findDuplicate();
250     ASSERT_TRUE(dup.hasValue());
251   }
252 
253   SmallVector<NamedAttribute> newAttrs = {
254       b.getNamedAttr("foo", b.getStringAttr("f")),
255       b.getNamedAttr("zoo", b.getStringAttr("z")),
256   };
257   attrs.assign(newAttrs);
258 
259   auto dup = attrs.findDuplicate();
260   ASSERT_FALSE(dup.hasValue());
261 
262   {
263     auto it = attrs.begin();
264     EXPECT_EQ(it->first, b.getStringAttr("foo"));
265     EXPECT_EQ(it->second, b.getStringAttr("f"));
266     ++it;
267     EXPECT_EQ(it->first, b.getStringAttr("zoo"));
268     EXPECT_EQ(it->second, b.getStringAttr("z"));
269   }
270 
271   attrs.assign({});
272   ASSERT_TRUE(attrs.empty());
273 }
274 } // end namespace
275