xref: /llvm-project/mlir/unittests/Interfaces/ControlFlowInterfacesTest.cpp (revision dfaadf6b12b8ca5e45855dc417e56dc31ff2b2ae)
1 //===- ControlFlowInterfacesTest.cpp - Unit Tests for Control Flow Interf. ===//
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/Interfaces/ControlFlowInterfaces.h"
10 #include "mlir/IR/BuiltinOps.h"
11 #include "mlir/IR/Dialect.h"
12 #include "mlir/IR/DialectImplementation.h"
13 #include "mlir/IR/OpDefinition.h"
14 #include "mlir/IR/OpImplementation.h"
15 #include "mlir/Parser/Parser.h"
16 
17 #include <gtest/gtest.h>
18 
19 using namespace mlir;
20 
21 /// A dummy op that is also a terminator.
22 struct DummyOp : public Op<DummyOp, OpTrait::IsTerminator> {
23   using Op::Op;
24   static ArrayRef<StringRef> getAttributeNames() { return {}; }
25 
26   static StringRef getOperationName() { return "cftest.dummy_op"; }
27 };
28 
29 /// All regions of this op are mutually exclusive.
30 struct MutuallyExclusiveRegionsOp
31     : public Op<MutuallyExclusiveRegionsOp, RegionBranchOpInterface::Trait> {
32   using Op::Op;
33   static ArrayRef<StringRef> getAttributeNames() { return {}; }
34 
35   static StringRef getOperationName() {
36     return "cftest.mutually_exclusive_regions_op";
37   }
38 
39   // Regions have no successors.
40   void getSuccessorRegions(Optional<unsigned> index,
41                            ArrayRef<Attribute> operands,
42                            SmallVectorImpl<RegionSuccessor> &regions) {}
43 };
44 
45 /// Regions are executed sequentially.
46 struct SequentialRegionsOp
47     : public Op<SequentialRegionsOp, RegionBranchOpInterface::Trait> {
48   using Op::Op;
49   static ArrayRef<StringRef> getAttributeNames() { return {}; }
50 
51   static StringRef getOperationName() { return "cftest.sequential_regions_op"; }
52 
53   // Region 0 has Region 1 as a successor.
54   void getSuccessorRegions(Optional<unsigned> index,
55                            ArrayRef<Attribute> operands,
56                            SmallVectorImpl<RegionSuccessor> &regions) {
57     if (index == 0u) {
58       Operation *thisOp = this->getOperation();
59       regions.push_back(RegionSuccessor(&thisOp->getRegion(1)));
60     }
61   }
62 };
63 
64 /// A dialect putting all the above together.
65 struct CFTestDialect : Dialect {
66   explicit CFTestDialect(MLIRContext *ctx)
67       : Dialect(getDialectNamespace(), ctx, TypeID::get<CFTestDialect>()) {
68     addOperations<DummyOp, MutuallyExclusiveRegionsOp, SequentialRegionsOp>();
69   }
70   static StringRef getDialectNamespace() { return "cftest"; }
71 };
72 
73 TEST(RegionBranchOpInterface, MutuallyExclusiveOps) {
74   const char *ir = R"MLIR(
75 "cftest.mutually_exclusive_regions_op"() (
76       {"cftest.dummy_op"() : () -> ()},  // op1
77       {"cftest.dummy_op"() : () -> ()}   // op2
78   ) : () -> ()
79   )MLIR";
80 
81   DialectRegistry registry;
82   registry.insert<CFTestDialect>();
83   MLIRContext ctx(registry);
84 
85   OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
86   Operation *testOp = &module->getBody()->getOperations().front();
87   Operation *op1 = &testOp->getRegion(0).front().front();
88   Operation *op2 = &testOp->getRegion(1).front().front();
89 
90   EXPECT_TRUE(insideMutuallyExclusiveRegions(op1, op2));
91   EXPECT_TRUE(insideMutuallyExclusiveRegions(op2, op1));
92 }
93 
94 TEST(RegionBranchOpInterface, NotMutuallyExclusiveOps) {
95   const char *ir = R"MLIR(
96 "cftest.sequential_regions_op"() (
97       {"cftest.dummy_op"() : () -> ()},  // op1
98       {"cftest.dummy_op"() : () -> ()}   // op2
99   ) : () -> ()
100   )MLIR";
101 
102   DialectRegistry registry;
103   registry.insert<CFTestDialect>();
104   MLIRContext ctx(registry);
105 
106   OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
107   Operation *testOp = &module->getBody()->getOperations().front();
108   Operation *op1 = &testOp->getRegion(0).front().front();
109   Operation *op2 = &testOp->getRegion(1).front().front();
110 
111   EXPECT_FALSE(insideMutuallyExclusiveRegions(op1, op2));
112   EXPECT_FALSE(insideMutuallyExclusiveRegions(op2, op1));
113 }
114 
115 TEST(RegionBranchOpInterface, NestedMutuallyExclusiveOps) {
116   const char *ir = R"MLIR(
117 "cftest.mutually_exclusive_regions_op"() (
118       {
119         "cftest.sequential_regions_op"() (
120               {"cftest.dummy_op"() : () -> ()},  // op1
121               {"cftest.dummy_op"() : () -> ()}   // op3
122           ) : () -> ()
123         "cftest.dummy_op"() : () -> ()
124       },
125       {"cftest.dummy_op"() : () -> ()}           // op2
126   ) : () -> ()
127   )MLIR";
128 
129   DialectRegistry registry;
130   registry.insert<CFTestDialect>();
131   MLIRContext ctx(registry);
132 
133   OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
134   Operation *testOp = &module->getBody()->getOperations().front();
135   Operation *op1 =
136       &testOp->getRegion(0).front().front().getRegion(0).front().front();
137   Operation *op2 = &testOp->getRegion(1).front().front();
138   Operation *op3 =
139       &testOp->getRegion(0).front().front().getRegion(1).front().front();
140 
141   EXPECT_TRUE(insideMutuallyExclusiveRegions(op1, op2));
142   EXPECT_TRUE(insideMutuallyExclusiveRegions(op3, op2));
143   EXPECT_FALSE(insideMutuallyExclusiveRegions(op1, op3));
144 }
145