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