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, bool resizableOperands,
18                            ArrayRef<Value> operands = llvm::None,
19                            ArrayRef<Type> resultTypes = llvm::None) {
20   context->allowUnregisteredDialects();
21   return Operation::create(
22       UnknownLoc::get(context), OperationName("foo.bar", context), resultTypes,
23       operands, llvm::None, llvm::None, 0, resizableOperands);
24 }
25 
26 namespace {
27 TEST(OperandStorageTest, NonResizable) {
28   MLIRContext context;
29   Builder builder(&context);
30 
31   Operation *useOp =
32       createOp(&context, /*resizableOperands=*/false, /*operands=*/llvm::None,
33                builder.getIntegerType(16));
34   Value operand = useOp->getResult(0);
35 
36   // Create a non-resizable operation with one operand.
37   Operation *user = createOp(&context, /*resizableOperands=*/false, operand,
38                              builder.getIntegerType(16));
39 
40   // Sanity check the storage.
41   EXPECT_EQ(user->hasResizableOperandsList(), false);
42 
43   // The same number of operands is okay.
44   user->setOperands(operand);
45   EXPECT_EQ(user->getNumOperands(), 1u);
46 
47   // Removing is okay.
48   user->setOperands(llvm::None);
49   EXPECT_EQ(user->getNumOperands(), 0u);
50 
51   // Destroy the operations.
52   user->destroy();
53   useOp->destroy();
54 }
55 
56 TEST(OperandStorageDeathTest, AddToNonResizable) {
57   MLIRContext context;
58   Builder builder(&context);
59 
60   Operation *useOp =
61       createOp(&context, /*resizableOperands=*/false, /*operands=*/llvm::None,
62                builder.getIntegerType(16));
63   Value operand = useOp->getResult(0);
64 
65   // Create a non-resizable operation with one operand.
66   Operation *user = createOp(&context, /*resizableOperands=*/false, operand,
67                              builder.getIntegerType(16));
68 
69   // Sanity check the storage.
70   EXPECT_EQ(user->hasResizableOperandsList(), false);
71 
72   // Adding operands to a non resizable operation should result in a failure.
73   ASSERT_DEATH(user->setOperands({operand, operand}), "");
74 }
75 
76 TEST(OperandStorageTest, Resizable) {
77   MLIRContext context;
78   Builder builder(&context);
79 
80   Operation *useOp =
81       createOp(&context, /*resizableOperands=*/false, /*operands=*/llvm::None,
82                builder.getIntegerType(16));
83   Value operand = useOp->getResult(0);
84 
85   // Create a resizable operation with one operand.
86   Operation *user = createOp(&context, /*resizableOperands=*/true, operand,
87                              builder.getIntegerType(16));
88 
89   // Sanity check the storage.
90   EXPECT_EQ(user->hasResizableOperandsList(), true);
91 
92   // The same number of operands is okay.
93   user->setOperands(operand);
94   EXPECT_EQ(user->getNumOperands(), 1u);
95 
96   // Removing is okay.
97   user->setOperands(llvm::None);
98   EXPECT_EQ(user->getNumOperands(), 0u);
99 
100   // Adding more operands is okay.
101   user->setOperands({operand, operand, operand});
102   EXPECT_EQ(user->getNumOperands(), 3u);
103 
104   // Destroy the operations.
105   user->destroy();
106   useOp->destroy();
107 }
108 
109 } // end namespace
110