xref: /llvm-project/mlir/unittests/IR/OperationSupportTest.cpp (revision a776ecb6c2b83454f4cd22a0e302aef558381dab)
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 "gtest/gtest.h"
14 
15 using namespace mlir;
16 using namespace mlir::detail;
17 
18 static Operation *createOp(MLIRContext *context,
19                            ArrayRef<Value> operands = llvm::None,
20                            ArrayRef<Type> resultTypes = llvm::None,
21                            unsigned int numRegions = 0) {
22   context->allowUnregisteredDialects();
23   return Operation::create(UnknownLoc::get(context),
24                            OperationName("foo.bar", context), resultTypes,
25                            operands, llvm::None, llvm::None, numRegions);
26 }
27 
28 namespace {
29 TEST(OperandStorageTest, NonResizable) {
30   MLIRContext context;
31   Builder builder(&context);
32 
33   Operation *useOp =
34       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
35   Value operand = useOp->getResult(0);
36 
37   // Create a non-resizable operation with one operand.
38   Operation *user = createOp(&context, operand);
39 
40   // The same number of operands is okay.
41   user->setOperands(operand);
42   EXPECT_EQ(user->getNumOperands(), 1u);
43 
44   // Removing is okay.
45   user->setOperands(llvm::None);
46   EXPECT_EQ(user->getNumOperands(), 0u);
47 
48   // Destroy the operations.
49   user->destroy();
50   useOp->destroy();
51 }
52 
53 TEST(OperandStorageTest, Resizable) {
54   MLIRContext context;
55   Builder builder(&context);
56 
57   Operation *useOp =
58       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
59   Value operand = useOp->getResult(0);
60 
61   // Create a resizable operation with one operand.
62   Operation *user = createOp(&context, operand);
63 
64   // The same number of operands is okay.
65   user->setOperands(operand);
66   EXPECT_EQ(user->getNumOperands(), 1u);
67 
68   // Removing is okay.
69   user->setOperands(llvm::None);
70   EXPECT_EQ(user->getNumOperands(), 0u);
71 
72   // Adding more operands is okay.
73   user->setOperands({operand, operand, operand});
74   EXPECT_EQ(user->getNumOperands(), 3u);
75 
76   // Destroy the operations.
77   user->destroy();
78   useOp->destroy();
79 }
80 
81 TEST(OperandStorageTest, RangeReplace) {
82   MLIRContext context;
83   Builder builder(&context);
84 
85   Operation *useOp =
86       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
87   Value operand = useOp->getResult(0);
88 
89   // Create a resizable operation with one operand.
90   Operation *user = createOp(&context, operand);
91 
92   // Check setting with the same number of operands.
93   user->setOperands(/*start=*/0, /*length=*/1, operand);
94   EXPECT_EQ(user->getNumOperands(), 1u);
95 
96   // Check setting with more operands.
97   user->setOperands(/*start=*/0, /*length=*/1, {operand, operand, operand});
98   EXPECT_EQ(user->getNumOperands(), 3u);
99 
100   // Check setting with less operands.
101   user->setOperands(/*start=*/1, /*length=*/2, {operand});
102   EXPECT_EQ(user->getNumOperands(), 2u);
103 
104   // Check inserting without replacing operands.
105   user->setOperands(/*start=*/2, /*length=*/0, {operand});
106   EXPECT_EQ(user->getNumOperands(), 3u);
107 
108   // Check erasing operands.
109   user->setOperands(/*start=*/0, /*length=*/3, {});
110   EXPECT_EQ(user->getNumOperands(), 0u);
111 
112   // Destroy the operations.
113   user->destroy();
114   useOp->destroy();
115 }
116 
117 TEST(OperandStorageTest, MutableRange) {
118   MLIRContext context;
119   Builder builder(&context);
120 
121   Operation *useOp =
122       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
123   Value operand = useOp->getResult(0);
124 
125   // Create a resizable operation with one operand.
126   Operation *user = createOp(&context, operand);
127 
128   // Check setting with the same number of operands.
129   MutableOperandRange mutableOperands(user);
130   mutableOperands.assign(operand);
131   EXPECT_EQ(mutableOperands.size(), 1u);
132   EXPECT_EQ(user->getNumOperands(), 1u);
133 
134   // Check setting with more operands.
135   mutableOperands.assign({operand, operand, operand});
136   EXPECT_EQ(mutableOperands.size(), 3u);
137   EXPECT_EQ(user->getNumOperands(), 3u);
138 
139   // Check with inserting a new operand.
140   mutableOperands.append({operand, operand});
141   EXPECT_EQ(mutableOperands.size(), 5u);
142   EXPECT_EQ(user->getNumOperands(), 5u);
143 
144   // Check erasing operands.
145   mutableOperands.clear();
146   EXPECT_EQ(mutableOperands.size(), 0u);
147   EXPECT_EQ(user->getNumOperands(), 0u);
148 
149   // Destroy the operations.
150   user->destroy();
151   useOp->destroy();
152 }
153 
154 TEST(OperandStorageTest, RangeErase) {
155   MLIRContext context;
156   Builder builder(&context);
157 
158   Type type = builder.getNoneType();
159   Operation *useOp = createOp(&context, /*operands=*/llvm::None, {type, type});
160   Value operand1 = useOp->getResult(0);
161   Value operand2 = useOp->getResult(1);
162 
163   // Create an operation with operands to erase.
164   Operation *user =
165       createOp(&context, {operand2, operand1, operand2, operand1});
166   llvm::BitVector eraseIndices(user->getNumOperands());
167 
168   // Check erasing no operands.
169   user->eraseOperands(eraseIndices);
170   EXPECT_EQ(user->getNumOperands(), 4u);
171 
172   // Check erasing disjoint operands.
173   eraseIndices.set(0);
174   eraseIndices.set(3);
175   user->eraseOperands(eraseIndices);
176   EXPECT_EQ(user->getNumOperands(), 2u);
177   EXPECT_EQ(user->getOperand(0), operand1);
178   EXPECT_EQ(user->getOperand(1), operand2);
179 
180   // Destroy the operations.
181   user->destroy();
182   useOp->destroy();
183 }
184 
185 TEST(OperationOrderTest, OrderIsAlwaysValid) {
186   MLIRContext context;
187   Builder builder(&context);
188 
189   Operation *containerOp =
190       createOp(&context, /*operands=*/llvm::None, /*resultTypes=*/llvm::None,
191                /*numRegions=*/1);
192   Region &region = containerOp->getRegion(0);
193   Block *block = new Block();
194   region.push_back(block);
195 
196   // Insert two operations, then iteratively add more operations in the middle
197   // of them. Eventually we will insert more than kOrderStride operations and
198   // the block order will need to be recomputed.
199   Operation *frontOp = createOp(&context);
200   Operation *backOp = createOp(&context);
201   block->push_back(frontOp);
202   block->push_back(backOp);
203 
204   // Chosen to be larger than Operation::kOrderStride.
205   int kNumOpsToInsert = 10;
206   for (int i = 0; i < kNumOpsToInsert; ++i) {
207     Operation *op = createOp(&context);
208     block->getOperations().insert(backOp->getIterator(), op);
209     ASSERT_TRUE(op->isBeforeInBlock(backOp));
210     // Note verifyOpOrder() returns false if the order is valid.
211     ASSERT_FALSE(block->verifyOpOrder());
212   }
213 
214   containerOp->destroy();
215 }
216 
217 } // end namespace
218