1 //===- ControlFlowInterfaces.cpp - ControlFlow Interfaces -----------------===//
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/StandardTypes.h"
11 #include "llvm/ADT/SmallPtrSet.h"
12 
13 using namespace mlir;
14 
15 //===----------------------------------------------------------------------===//
16 // ControlFlowInterfaces
17 //===----------------------------------------------------------------------===//
18 
19 #include "mlir/Interfaces/ControlFlowInterfaces.cpp.inc"
20 
21 //===----------------------------------------------------------------------===//
22 // BranchOpInterface
23 //===----------------------------------------------------------------------===//
24 
25 /// Returns the `BlockArgument` corresponding to operand `operandIndex` in some
26 /// successor if 'operandIndex' is within the range of 'operands', or None if
27 /// `operandIndex` isn't a successor operand index.
28 Optional<BlockArgument>
29 detail::getBranchSuccessorArgument(Optional<OperandRange> operands,
30                                    unsigned operandIndex, Block *successor) {
31   // Check that the operands are valid.
32   if (!operands || operands->empty())
33     return llvm::None;
34 
35   // Check to ensure that this operand is within the range.
36   unsigned operandsStart = operands->getBeginOperandIndex();
37   if (operandIndex < operandsStart ||
38       operandIndex >= (operandsStart + operands->size()))
39     return llvm::None;
40 
41   // Index the successor.
42   unsigned argIndex = operandIndex - operandsStart;
43   return successor->getArgument(argIndex);
44 }
45 
46 /// Verify that the given operands match those of the given successor block.
47 LogicalResult
48 detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
49                                       Optional<OperandRange> operands) {
50   if (!operands)
51     return success();
52 
53   // Check the count.
54   unsigned operandCount = operands->size();
55   Block *destBB = op->getSuccessor(succNo);
56   if (operandCount != destBB->getNumArguments())
57     return op->emitError() << "branch has " << operandCount
58                            << " operands for successor #" << succNo
59                            << ", but target block has "
60                            << destBB->getNumArguments();
61 
62   // Check the types.
63   auto operandIt = operands->begin();
64   for (unsigned i = 0; i != operandCount; ++i, ++operandIt) {
65     if ((*operandIt).getType() != destBB->getArgument(i).getType())
66       return op->emitError() << "type mismatch for bb argument #" << i
67                              << " of successor #" << succNo;
68   }
69   return success();
70 }
71 
72 //===----------------------------------------------------------------------===//
73 // RegionBranchOpInterface
74 //===----------------------------------------------------------------------===//
75 
76 /// Verify that types match along all region control flow edges originating from
77 /// `sourceNo` (region # if source is a region, llvm::None if source is parent
78 /// op). `getInputsTypesForRegion` is a function that returns the types of the
79 /// inputs that flow from `sourceIndex' to the given region, or llvm::None if
80 /// the exact type match verification is not necessary (e.g., if the Op verifies
81 /// the match itself).
82 static LogicalResult
83 verifyTypesAlongAllEdges(Operation *op, Optional<unsigned> sourceNo,
84                          function_ref<Optional<TypeRange>(Optional<unsigned>)>
85                              getInputsTypesForRegion) {
86   auto regionInterface = cast<RegionBranchOpInterface>(op);
87 
88   SmallVector<RegionSuccessor, 2> successors;
89   unsigned numInputs;
90   if (sourceNo) {
91     Region &srcRegion = op->getRegion(sourceNo.getValue());
92     numInputs = srcRegion.getNumArguments();
93   } else {
94     numInputs = op->getNumOperands();
95   }
96   SmallVector<Attribute, 2> operands(numInputs, nullptr);
97   regionInterface.getSuccessorRegions(sourceNo, operands, successors);
98 
99   for (RegionSuccessor &succ : successors) {
100     Optional<unsigned> succRegionNo;
101     if (!succ.isParent())
102       succRegionNo = succ.getSuccessor()->getRegionNumber();
103 
104     auto printEdgeName = [&](InFlightDiagnostic &diag) -> InFlightDiagnostic & {
105       diag << "from ";
106       if (sourceNo)
107         diag << "Region #" << sourceNo.getValue();
108       else
109         diag << "parent operands";
110 
111       diag << " to ";
112       if (succRegionNo)
113         diag << "Region #" << succRegionNo.getValue();
114       else
115         diag << "parent results";
116       return diag;
117     };
118 
119     Optional<TypeRange> sourceTypes = getInputsTypesForRegion(succRegionNo);
120     if (!sourceTypes.hasValue())
121       continue;
122 
123     TypeRange succInputsTypes = succ.getSuccessorInputs().getTypes();
124     if (sourceTypes->size() != succInputsTypes.size()) {
125       InFlightDiagnostic diag = op->emitOpError(" region control flow edge ");
126       return printEdgeName(diag) << ": source has " << sourceTypes->size()
127                                  << " operands, but target successor needs "
128                                  << succInputsTypes.size();
129     }
130 
131     for (auto typesIdx :
132          llvm::enumerate(llvm::zip(*sourceTypes, succInputsTypes))) {
133       Type sourceType = std::get<0>(typesIdx.value());
134       Type inputType = std::get<1>(typesIdx.value());
135       if (sourceType != inputType) {
136         InFlightDiagnostic diag = op->emitOpError(" along control flow edge ");
137         return printEdgeName(diag)
138                << ": source type #" << typesIdx.index() << " " << sourceType
139                << " should match input type #" << typesIdx.index() << " "
140                << inputType;
141       }
142     }
143   }
144   return success();
145 }
146 
147 /// Verify that types match along control flow edges described the given op.
148 LogicalResult detail::verifyTypesAlongControlFlowEdges(Operation *op) {
149   auto regionInterface = cast<RegionBranchOpInterface>(op);
150 
151   auto inputTypesFromParent = [&](Optional<unsigned> regionNo) -> TypeRange {
152     if (regionNo.hasValue()) {
153       return regionInterface.getSuccessorEntryOperands(regionNo.getValue())
154           .getTypes();
155     }
156 
157     // If the successor of a parent op is the parent itself
158     // RegionBranchOpInterface does not have an API to query what the entry
159     // operands will be in that case. Vend out the result types of the op in
160     // that case so that type checking succeeds for this case.
161     return op->getResultTypes();
162   };
163 
164   // Verify types along control flow edges originating from the parent.
165   if (failed(verifyTypesAlongAllEdges(op, llvm::None, inputTypesFromParent)))
166     return failure();
167 
168   // RegionBranchOpInterface should not be implemented by Ops that do not have
169   // attached regions.
170   assert(op->getNumRegions() != 0);
171 
172   // Verify types along control flow edges originating from each region.
173   for (unsigned regionNo : llvm::seq(0U, op->getNumRegions())) {
174     Region &region = op->getRegion(regionNo);
175 
176     // Since the interface cannot distinguish between different ReturnLike
177     // ops within the region branching to different successors, all ReturnLike
178     // ops in this region should have the same operand types. We will then use
179     // one of them as the representative for type matching.
180 
181     Operation *regionReturn = nullptr;
182     for (Block &block : region) {
183       Operation *terminator = block.getTerminator();
184       if (!terminator->hasTrait<OpTrait::ReturnLike>())
185         continue;
186 
187       if (!regionReturn) {
188         regionReturn = terminator;
189         continue;
190       }
191 
192       // Found more than one ReturnLike terminator. Make sure the operand types
193       // match with the first one.
194       if (regionReturn->getOperandTypes() != terminator->getOperandTypes())
195         return op->emitOpError("Region #")
196                << regionNo
197                << " operands mismatch between return-like terminators";
198     }
199 
200     auto inputTypesFromRegion =
201         [&](Optional<unsigned> regionNo) -> Optional<TypeRange> {
202       // If there is no return-like terminator, the op itself should verify
203       // type consistency.
204       if (!regionReturn)
205         return llvm::None;
206 
207       // All successors get the same set of operands.
208       return TypeRange(regionReturn->getOperands().getTypes());
209     };
210 
211     if (failed(verifyTypesAlongAllEdges(op, regionNo, inputTypesFromRegion)))
212       return failure();
213   }
214 
215   return success();
216 }
217