1 | //===-- MyExtension.cpp - Transform dialect tutorial ----------------------===// |
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 | // This file defines Transform dialect extension operations used in the |
10 | // Chapter 4 of the Transform dialect tutorial. |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #include "MyExtension.h" |
15 | #include "mlir/Dialect/Transform/IR/TransformDialect.h" |
16 | #include "llvm/Support/Debug.h" |
17 | |
18 | #define DEBUG_TYPE_MATCHER "transform-matcher" |
19 | #define DBGS_MATCHER() (llvm::dbgs() << "[" DEBUG_TYPE_MATCHER "] ") |
20 | #define DEBUG_MATCHER(x) DEBUG_WITH_TYPE(DEBUG_TYPE_MATCHER, x) |
21 | |
22 | #define GET_OP_CLASSES |
23 | #include "MyExtension.cpp.inc" |
24 | |
25 | //===---------------------------------------------------------------------===// |
26 | // MyExtension |
27 | //===---------------------------------------------------------------------===// |
28 | |
29 | // Define a new transform dialect extension. This uses the CRTP idiom to |
30 | // identify extensions. |
31 | class MyExtension |
32 | : public ::mlir::transform::TransformDialectExtension<MyExtension> { |
33 | public: |
34 | // The TypeID of this extension. |
35 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MyExtension) |
36 | |
37 | // The extension must derive the base constructor. |
38 | using Base::Base; |
39 | |
40 | // This function initializes the extension, similarly to `initialize` in |
41 | // dialect definitions. List individual operations and dependent dialects |
42 | // here. |
43 | void init(); |
44 | }; |
45 | |
46 | void MyExtension::init() { |
47 | // Register the additional match operations with the dialect similarly to |
48 | // other transform operations. List all operations generated from ODS. This |
49 | // call will perform additional checks that the operations implement the |
50 | // transform and memory effect interfaces required by the dialect interpreter |
51 | // and assert if they do not. |
52 | registerTransformOps< |
53 | #define GET_OP_LIST |
54 | #include "MyExtension.cpp.inc" |
55 | >(); |
56 | } |
57 | |
58 | //===---------------------------------------------------------------------===// |
59 | // HasOperandSatisfyingOp |
60 | //===---------------------------------------------------------------------===// |
61 | |
62 | /// Returns `true` if both types implement one of the interfaces provided as |
63 | /// template parameters. |
64 | template <typename... Tys> |
65 | static bool implementSameInterface(mlir::Type t1, mlir::Type t2) { |
66 | return ((llvm::isa<Tys>(t1) && llvm::isa<Tys>(t2)) || ... || false); |
67 | } |
68 | |
69 | /// Returns `true` if both types implement one of the transform dialect |
70 | /// interfaces. |
71 | static bool implementSameTransformInterface(mlir::Type t1, mlir::Type t2) { |
72 | return implementSameInterface< |
73 | mlir::transform::TransformHandleTypeInterface, |
74 | mlir::transform::TransformParamTypeInterface, |
75 | mlir::transform::TransformValueHandleTypeInterface>(t1, t2); |
76 | } |
77 | |
78 | // Matcher ops implement `apply` similarly to other transform ops. They are not |
79 | // expected to modify payload, but use the tri-state result to signal failure or |
80 | // success to match, as well as potential irrecoverable errors. |
81 | mlir::DiagnosedSilenceableFailure |
82 | mlir::transform::HasOperandSatisfyingOp::apply( |
83 | mlir::transform::TransformRewriter &rewriter, |
84 | mlir::transform::TransformResults &results, |
85 | mlir::transform::TransformState &state) { |
86 | // For simplicity, only handle a single payload op. Actual implementations |
87 | // can use `SingleOpMatcher` trait to simplify implementation and document |
88 | // this expectation. |
89 | auto payloadOps = state.getPayloadOps(getOp()); |
90 | if (!llvm::hasSingleElement(payloadOps)) |
91 | return emitSilenceableError() << "expected single payload" ; |
92 | |
93 | // Iterate over all operands of the payload op to see if they can be matched |
94 | // using the body of this op. |
95 | Operation *payload = *payloadOps.begin(); |
96 | for (OpOperand &operand : payload->getOpOperands()) { |
97 | // Create a scope for transform values defined in the body. This corresponds |
98 | // to the syntactic scope of the region attached to this op. Any values |
99 | // associated with payloads from now on will be automatically dissociated |
100 | // when this object is destroyed, i.e. at the end of the iteration. |
101 | // Associate the block argument handle with the operand. |
102 | auto matchScope = state.make_region_scope(getBody()); |
103 | if (failed(state.mapBlockArgument(getBody().getArgument(0), |
104 | {operand.get()}))) { |
105 | return DiagnosedSilenceableFailure::definiteFailure(); |
106 | } |
107 | |
108 | // Iterate over all nested matchers with the current mapping and see if they |
109 | // succeed. |
110 | bool matchSucceeded = true; |
111 | for (Operation &matcher : getBody().front().without_terminator()) { |
112 | // Matcher ops are applied similarly to any other transform op. |
113 | DiagnosedSilenceableFailure diag = |
114 | state.applyTransform(cast<TransformOpInterface>(matcher)); |
115 | |
116 | // Definite failures are immediately propagated as they are irrecoverable. |
117 | if (diag.isDefiniteFailure()) |
118 | return diag; |
119 | |
120 | // On success, keep checking the remaining conditions. |
121 | if (diag.succeeded()) |
122 | continue; |
123 | |
124 | // Report failure-to-match for debugging purposes and stop matching this |
125 | // operand. |
126 | assert(diag.isSilenceableFailure()); |
127 | DEBUG_MATCHER(DBGS_MATCHER() |
128 | << "failed to match operand #" << operand.getOperandNumber() |
129 | << ": " << diag.getMessage()); |
130 | (void)diag.silence(); |
131 | matchSucceeded = false; |
132 | break; |
133 | } |
134 | // If failed to match this operand, try other operands. |
135 | if (!matchSucceeded) |
136 | continue; |
137 | |
138 | // If we reached this point, the matching succeeded for the current operand. |
139 | // Remap the values associated with terminator operands to be associated |
140 | // with op results, and also map the parameter result to the operand's |
141 | // position. Note that it is safe to do here despite the end of the scope |
142 | // as `results` are integrated into `state` by the interpreter after `apply` |
143 | // returns rather than immediately. |
144 | SmallVector<SmallVector<MappedValue>> yieldedMappings; |
145 | transform::detail::prepareValueMappings( |
146 | yieldedMappings, getBody().front().getTerminator()->getOperands(), |
147 | state); |
148 | results.setParams(cast<OpResult>(getPosition()), |
149 | {rewriter.getI32IntegerAttr(operand.getOperandNumber())}); |
150 | for (auto &&[result, mapping] : llvm::zip(getResults(), yieldedMappings)) |
151 | results.setMappedValues(result, mapping); |
152 | return DiagnosedSilenceableFailure::success(); |
153 | } |
154 | |
155 | // If we reached this point, none of the operands succeeded the match. |
156 | return emitSilenceableError() |
157 | << "none of the operands satisfied the conditions" ; |
158 | } |
159 | |
160 | // By convention, operations implementing MatchOpInterface must not modify |
161 | // payload IR and must therefore specify that they only read operand handles and |
162 | // payload as their effects. |
163 | void mlir::transform::HasOperandSatisfyingOp::getEffects( |
164 | llvm::SmallVectorImpl<mlir::MemoryEffects::EffectInstance> &effects) { |
165 | onlyReadsPayload(effects); |
166 | onlyReadsHandle(getOpMutable(), effects); |
167 | producesHandle(getOperation()->getOpResults(), effects); |
168 | } |
169 | |
170 | // Verify well-formedness of the operation and emit diagnostics if it is |
171 | // ill-formed. |
172 | llvm::LogicalResult mlir::transform::HasOperandSatisfyingOp::verify() { |
173 | mlir::Block &bodyBlock = getBody().front(); |
174 | if (bodyBlock.getNumArguments() != 1 || |
175 | !isa<TransformValueHandleTypeInterface>( |
176 | bodyBlock.getArgument(0).getType())) { |
177 | return emitOpError() |
178 | << "expects the body to have one value handle argument" ; |
179 | } |
180 | if (bodyBlock.getTerminator()->getNumOperands() != getNumResults() - 1) { |
181 | return emitOpError() << "expects the body to yield " |
182 | << (getNumResults() - 1) << " values, got " |
183 | << bodyBlock.getTerminator()->getNumOperands(); |
184 | } |
185 | for (auto &&[i, operand, result] : |
186 | llvm::enumerate(bodyBlock.getTerminator()->getOperands().getTypes(), |
187 | getResults().getTypes())) { |
188 | if (implementSameTransformInterface(operand, result)) |
189 | continue; |
190 | return emitOpError() << "expects terminator operand #" << i |
191 | << " and result #" << (i + 1) |
192 | << " to implement the same transform interface" ; |
193 | } |
194 | |
195 | for (Operation &op : bodyBlock.without_terminator()) { |
196 | if (!isa<TransformOpInterface>(op) || !isa<MatchOpInterface>(op)) { |
197 | InFlightDiagnostic diag = emitOpError() |
198 | << "expects body to contain match ops" ; |
199 | diag.attachNote(op.getLoc()) << "non-match operation" ; |
200 | return diag; |
201 | } |
202 | } |
203 | |
204 | return success(); |
205 | } |
206 | |
207 | void registerMyExtension(::mlir::DialectRegistry ®istry) { |
208 | registry.addExtensions<MyExtension>(); |
209 | } |
210 | |