1 //===- SCFTransformOps.cpp - Implementation of SCF transformation ops -----===//
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/Dialect/SCF/TransformOps/SCFTransformOps.h"
10 #include "mlir/Dialect/Affine/IR/AffineOps.h"
11 #include "mlir/Dialect/Func/IR/FuncOps.h"
12 #include "mlir/Dialect/SCF/IR/SCF.h"
13 #include "mlir/Dialect/SCF/Transforms/Patterns.h"
14 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
15 #include "mlir/Dialect/SCF/Utils/Utils.h"
16 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
17 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
18 #include "mlir/Dialect/Vector/IR/VectorOps.h"
19 
20 using namespace mlir;
21 
22 namespace {
23 /// A simple pattern rewriter that implements no special logic.
24 class SimpleRewriter : public PatternRewriter {
25 public:
26   SimpleRewriter(MLIRContext *context) : PatternRewriter(context) {}
27 };
28 } // namespace
29 
30 //===----------------------------------------------------------------------===//
31 // GetParentForOp
32 //===----------------------------------------------------------------------===//
33 
34 DiagnosedSilenceableFailure
35 transform::GetParentForOp::apply(transform::TransformResults &results,
36                                  transform::TransformState &state) {
37   SetVector<Operation *> parents;
38   for (Operation *target : state.getPayloadOps(getTarget())) {
39     scf::ForOp loop;
40     Operation *current = target;
41     for (unsigned i = 0, e = getNumLoops(); i < e; ++i) {
42       loop = current->getParentOfType<scf::ForOp>();
43       if (!loop) {
44         DiagnosedSilenceableFailure diag = emitSilenceableError()
45                                            << "could not find an '"
46                                            << scf::ForOp::getOperationName()
47                                            << "' parent";
48         diag.attachNote(target->getLoc()) << "target op";
49         return diag;
50       }
51       current = loop;
52     }
53     parents.insert(loop);
54   }
55   results.set(getResult().cast<OpResult>(), parents.getArrayRef());
56   return DiagnosedSilenceableFailure::success();
57 }
58 
59 //===----------------------------------------------------------------------===//
60 // LoopOutlineOp
61 //===----------------------------------------------------------------------===//
62 
63 /// Wraps the given operation `op` into an `scf.execute_region` operation. Uses
64 /// the provided rewriter for all operations to remain compatible with the
65 /// rewriting infra, as opposed to just splicing the op in place.
66 static scf::ExecuteRegionOp wrapInExecuteRegion(RewriterBase &b,
67                                                 Operation *op) {
68   if (op->getNumRegions() != 1)
69     return nullptr;
70   OpBuilder::InsertionGuard g(b);
71   b.setInsertionPoint(op);
72   scf::ExecuteRegionOp executeRegionOp =
73       b.create<scf::ExecuteRegionOp>(op->getLoc(), op->getResultTypes());
74   {
75     OpBuilder::InsertionGuard g(b);
76     b.setInsertionPointToStart(&executeRegionOp.getRegion().emplaceBlock());
77     Operation *clonedOp = b.cloneWithoutRegions(*op);
78     Region &clonedRegion = clonedOp->getRegions().front();
79     assert(clonedRegion.empty() && "expected empty region");
80     b.inlineRegionBefore(op->getRegions().front(), clonedRegion,
81                          clonedRegion.end());
82     b.create<scf::YieldOp>(op->getLoc(), clonedOp->getResults());
83   }
84   b.replaceOp(op, executeRegionOp.getResults());
85   return executeRegionOp;
86 }
87 
88 DiagnosedSilenceableFailure
89 transform::LoopOutlineOp::apply(transform::TransformResults &results,
90                                 transform::TransformState &state) {
91   SmallVector<Operation *> transformed;
92   DenseMap<Operation *, SymbolTable> symbolTables;
93   for (Operation *target : state.getPayloadOps(getTarget())) {
94     Location location = target->getLoc();
95     Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(target);
96     SimpleRewriter rewriter(getContext());
97     scf::ExecuteRegionOp exec = wrapInExecuteRegion(rewriter, target);
98     if (!exec) {
99       DiagnosedSilenceableFailure diag = emitSilenceableError()
100                                          << "failed to outline";
101       diag.attachNote(target->getLoc()) << "target op";
102       return diag;
103     }
104     func::CallOp call;
105     FailureOr<func::FuncOp> outlined = outlineSingleBlockRegion(
106         rewriter, location, exec.getRegion(), getFuncName(), &call);
107 
108     if (failed(outlined)) {
109       (void)reportUnknownTransformError(target);
110       return DiagnosedSilenceableFailure::definiteFailure();
111     }
112 
113     if (symbolTableOp) {
114       SymbolTable &symbolTable =
115           symbolTables.try_emplace(symbolTableOp, symbolTableOp)
116               .first->getSecond();
117       symbolTable.insert(*outlined);
118       call.setCalleeAttr(FlatSymbolRefAttr::get(*outlined));
119     }
120     transformed.push_back(*outlined);
121   }
122   results.set(getTransformed().cast<OpResult>(), transformed);
123   return DiagnosedSilenceableFailure::success();
124 }
125 
126 //===----------------------------------------------------------------------===//
127 // LoopPeelOp
128 //===----------------------------------------------------------------------===//
129 
130 FailureOr<scf::ForOp> transform::LoopPeelOp::applyToOne(scf::ForOp loop,
131                                                         TransformState &state) {
132   scf::ForOp result;
133   IRRewriter rewriter(loop->getContext());
134   LogicalResult status =
135       scf::peelAndCanonicalizeForLoop(rewriter, loop, result);
136   if (failed(status)) {
137     if (getFailIfAlreadyDivisible())
138       return reportUnknownTransformError(loop);
139     return loop;
140   }
141   return result;
142 }
143 
144 //===----------------------------------------------------------------------===//
145 // LoopPipelineOp
146 //===----------------------------------------------------------------------===//
147 
148 /// Callback for PipeliningOption. Populates `schedule` with the mapping from an
149 /// operation to its logical time position given the iteration interval and the
150 /// read latency. The latter is only relevant for vector transfers.
151 static void
152 loopScheduling(scf::ForOp forOp,
153                std::vector<std::pair<Operation *, unsigned>> &schedule,
154                unsigned iterationInterval, unsigned readLatency) {
155   auto getLatency = [&](Operation *op) -> unsigned {
156     if (isa<vector::TransferReadOp>(op))
157       return readLatency;
158     return 1;
159   };
160 
161   DenseMap<Operation *, unsigned> opCycles;
162   std::map<unsigned, std::vector<Operation *>> wrappedSchedule;
163   for (Operation &op : forOp.getBody()->getOperations()) {
164     if (isa<scf::YieldOp>(op))
165       continue;
166     unsigned earlyCycle = 0;
167     for (Value operand : op.getOperands()) {
168       Operation *def = operand.getDefiningOp();
169       if (!def)
170         continue;
171       earlyCycle = std::max(earlyCycle, opCycles[def] + getLatency(def));
172     }
173     opCycles[&op] = earlyCycle;
174     wrappedSchedule[earlyCycle % iterationInterval].push_back(&op);
175   }
176   for (const auto &it : wrappedSchedule) {
177     for (Operation *op : it.second) {
178       unsigned cycle = opCycles[op];
179       schedule.emplace_back(op, cycle / iterationInterval);
180     }
181   }
182 }
183 
184 FailureOr<scf::ForOp>
185 transform::LoopPipelineOp::applyToOne(scf::ForOp loop, TransformState &state) {
186   scf::PipeliningOption options;
187   options.getScheduleFn =
188       [this](scf::ForOp forOp,
189              std::vector<std::pair<Operation *, unsigned>> &schedule) mutable {
190         loopScheduling(forOp, schedule, getIterationInterval(),
191                        getReadLatency());
192       };
193 
194   scf::ForLoopPipeliningPattern pattern(options, loop->getContext());
195   SimpleRewriter rewriter(getContext());
196   rewriter.setInsertionPoint(loop);
197   FailureOr<scf::ForOp> patternResult =
198       pattern.returningMatchAndRewrite(loop, rewriter);
199   if (failed(patternResult))
200     return reportUnknownTransformError(loop);
201   return patternResult;
202 }
203 
204 //===----------------------------------------------------------------------===//
205 // LoopUnrollOp
206 //===----------------------------------------------------------------------===//
207 
208 LogicalResult transform::LoopUnrollOp::applyToOne(scf::ForOp loop,
209                                                   TransformState &state) {
210   if (failed(loopUnrollByFactor(loop, getFactor())))
211     return reportUnknownTransformError(loop);
212   return success();
213 }
214 
215 //===----------------------------------------------------------------------===//
216 // Transform op registration
217 //===----------------------------------------------------------------------===//
218 
219 namespace {
220 class SCFTransformDialectExtension
221     : public transform::TransformDialectExtension<
222           SCFTransformDialectExtension> {
223 public:
224   SCFTransformDialectExtension() {
225     declareDependentDialect<AffineDialect>();
226     declareDependentDialect<func::FuncDialect>();
227     registerTransformOps<
228 #define GET_OP_LIST
229 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
230         >();
231   }
232 };
233 } // namespace
234 
235 #define GET_OP_CLASSES
236 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
237 
238 void mlir::scf::registerTransformDialectExtension(DialectRegistry &registry) {
239   registry.addExtensions<SCFTransformDialectExtension>();
240 }
241