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/StandardTypes.h"
12 #include "gtest/gtest.h"
13 
14 using namespace mlir;
15 using namespace mlir::detail;
16 
17 static Operation *createOp(MLIRContext *context,
18                            ArrayRef<Value> operands = llvm::None,
19                            ArrayRef<Type> resultTypes = llvm::None) {
20   context->allowUnregisteredDialects();
21   return Operation::create(UnknownLoc::get(context),
22                            OperationName("foo.bar", context), resultTypes,
23                            operands, llvm::None, llvm::None, 0);
24 }
25 
26 namespace {
27 TEST(OperandStorageTest, NonResizable) {
28   MLIRContext context;
29   Builder builder(&context);
30 
31   Operation *useOp =
32       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
33   Value operand = useOp->getResult(0);
34 
35   // Create a non-resizable operation with one operand.
36   Operation *user = createOp(&context, operand);
37 
38   // The same number of operands is okay.
39   user->setOperands(operand);
40   EXPECT_EQ(user->getNumOperands(), 1u);
41 
42   // Removing is okay.
43   user->setOperands(llvm::None);
44   EXPECT_EQ(user->getNumOperands(), 0u);
45 
46   // Destroy the operations.
47   user->destroy();
48   useOp->destroy();
49 }
50 
51 TEST(OperandStorageTest, Resizable) {
52   MLIRContext context;
53   Builder builder(&context);
54 
55   Operation *useOp =
56       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
57   Value operand = useOp->getResult(0);
58 
59   // Create a resizable operation with one operand.
60   Operation *user = createOp(&context, operand);
61 
62   // The same number of operands is okay.
63   user->setOperands(operand);
64   EXPECT_EQ(user->getNumOperands(), 1u);
65 
66   // Removing is okay.
67   user->setOperands(llvm::None);
68   EXPECT_EQ(user->getNumOperands(), 0u);
69 
70   // Adding more operands is okay.
71   user->setOperands({operand, operand, operand});
72   EXPECT_EQ(user->getNumOperands(), 3u);
73 
74   // Destroy the operations.
75   user->destroy();
76   useOp->destroy();
77 }
78 
79 TEST(OperandStorageTest, RangeReplace) {
80   MLIRContext context;
81   Builder builder(&context);
82 
83   Operation *useOp =
84       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
85   Value operand = useOp->getResult(0);
86 
87   // Create a resizable operation with one operand.
88   Operation *user = createOp(&context, operand);
89 
90   // Check setting with the same number of operands.
91   user->setOperands(/*start=*/0, /*length=*/1, operand);
92   EXPECT_EQ(user->getNumOperands(), 1u);
93 
94   // Check setting with more operands.
95   user->setOperands(/*start=*/0, /*length=*/1, {operand, operand, operand});
96   EXPECT_EQ(user->getNumOperands(), 3u);
97 
98   // Check setting with less operands.
99   user->setOperands(/*start=*/1, /*length=*/2, {operand});
100   EXPECT_EQ(user->getNumOperands(), 2u);
101 
102   // Check inserting without replacing operands.
103   user->setOperands(/*start=*/2, /*length=*/0, {operand});
104   EXPECT_EQ(user->getNumOperands(), 3u);
105 
106   // Check erasing operands.
107   user->setOperands(/*start=*/0, /*length=*/3, {});
108   EXPECT_EQ(user->getNumOperands(), 0u);
109 
110   // Destroy the operations.
111   user->destroy();
112   useOp->destroy();
113 }
114 
115 TEST(OperandStorageTest, MutableRange) {
116   MLIRContext context;
117   Builder builder(&context);
118 
119   Operation *useOp =
120       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
121   Value operand = useOp->getResult(0);
122 
123   // Create a resizable operation with one operand.
124   Operation *user = createOp(&context, operand);
125 
126   // Check setting with the same number of operands.
127   MutableOperandRange mutableOperands(user);
128   mutableOperands.assign(operand);
129   EXPECT_EQ(mutableOperands.size(), 1u);
130   EXPECT_EQ(user->getNumOperands(), 1u);
131 
132   // Check setting with more operands.
133   mutableOperands.assign({operand, operand, operand});
134   EXPECT_EQ(mutableOperands.size(), 3u);
135   EXPECT_EQ(user->getNumOperands(), 3u);
136 
137   // Check with inserting a new operand.
138   mutableOperands.append({operand, operand});
139   EXPECT_EQ(mutableOperands.size(), 5u);
140   EXPECT_EQ(user->getNumOperands(), 5u);
141 
142   // Check erasing operands.
143   mutableOperands.clear();
144   EXPECT_EQ(mutableOperands.size(), 0u);
145   EXPECT_EQ(user->getNumOperands(), 0u);
146 
147   // Destroy the operations.
148   user->destroy();
149   useOp->destroy();
150 }
151 
152 } // end namespace
153