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/Patterns.h"
13 #include "mlir/Dialect/SCF/SCF.h"
14 #include "mlir/Dialect/SCF/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 DiagnosedSilencableFailure
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         DiagnosedSilencableFailure diag = emitSilencableError()
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 DiagnosedSilencableFailure::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 DiagnosedSilencableFailure
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       DiagnosedSilencableFailure diag = emitSilencableError()
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 DiagnosedSilencableFailure::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 DiagnosedSilencableFailure::success();
124 }
125 
126 //===----------------------------------------------------------------------===//
127 // LoopPeelOp
128 //===----------------------------------------------------------------------===//
129 
130 FailureOr<scf::ForOp> transform::LoopPeelOp::applyToOne(scf::ForOp loop) {
131   scf::ForOp result;
132   IRRewriter rewriter(loop->getContext());
133   LogicalResult status =
134       scf::peelAndCanonicalizeForLoop(rewriter, loop, result);
135   if (failed(status)) {
136     if (getFailIfAlreadyDivisible())
137       return reportUnknownTransformError(loop);
138     return loop;
139   }
140   return result;
141 }
142 
143 //===----------------------------------------------------------------------===//
144 // LoopPipelineOp
145 //===----------------------------------------------------------------------===//
146 
147 /// Callback for PipeliningOption. Populates `schedule` with the mapping from an
148 /// operation to its logical time position given the iteration interval and the
149 /// read latency. The latter is only relevant for vector transfers.
150 static void
151 loopScheduling(scf::ForOp forOp,
152                std::vector<std::pair<Operation *, unsigned>> &schedule,
153                unsigned iterationInterval, unsigned readLatency) {
154   auto getLatency = [&](Operation *op) -> unsigned {
155     if (isa<vector::TransferReadOp>(op))
156       return readLatency;
157     return 1;
158   };
159 
160   DenseMap<Operation *, unsigned> opCycles;
161   std::map<unsigned, std::vector<Operation *>> wrappedSchedule;
162   for (Operation &op : forOp.getBody()->getOperations()) {
163     if (isa<scf::YieldOp>(op))
164       continue;
165     unsigned earlyCycle = 0;
166     for (Value operand : op.getOperands()) {
167       Operation *def = operand.getDefiningOp();
168       if (!def)
169         continue;
170       earlyCycle = std::max(earlyCycle, opCycles[def] + getLatency(def));
171     }
172     opCycles[&op] = earlyCycle;
173     wrappedSchedule[earlyCycle % iterationInterval].push_back(&op);
174   }
175   for (const auto &it : wrappedSchedule) {
176     for (Operation *op : it.second) {
177       unsigned cycle = opCycles[op];
178       schedule.push_back(std::make_pair(op, cycle / iterationInterval));
179     }
180   }
181 }
182 
183 FailureOr<scf::ForOp> transform::LoopPipelineOp::applyToOne(scf::ForOp loop) {
184   scf::PipeliningOption options;
185   options.getScheduleFn =
186       [this](scf::ForOp forOp,
187              std::vector<std::pair<Operation *, unsigned>> &schedule) mutable {
188         loopScheduling(forOp, schedule, getIterationInterval(),
189                        getReadLatency());
190       };
191 
192   scf::ForLoopPipeliningPattern pattern(options, loop->getContext());
193   SimpleRewriter rewriter(getContext());
194   rewriter.setInsertionPoint(loop);
195   FailureOr<scf::ForOp> patternResult =
196       pattern.returningMatchAndRewrite(loop, rewriter);
197   if (failed(patternResult))
198     return reportUnknownTransformError(loop);
199   return patternResult;
200 }
201 
202 //===----------------------------------------------------------------------===//
203 // LoopUnrollOp
204 //===----------------------------------------------------------------------===//
205 
206 LogicalResult transform::LoopUnrollOp::applyToOne(scf::ForOp loop) {
207   if (failed(loopUnrollByFactor(loop, getFactor())))
208     return reportUnknownTransformError(loop);
209   return success();
210 }
211 
212 //===----------------------------------------------------------------------===//
213 // Transform op registration
214 //===----------------------------------------------------------------------===//
215 
216 namespace {
217 class SCFTransformDialectExtension
218     : public transform::TransformDialectExtension<
219           SCFTransformDialectExtension> {
220 public:
221   SCFTransformDialectExtension() {
222     declareDependentDialect<AffineDialect>();
223     declareDependentDialect<func::FuncDialect>();
224     registerTransformOps<
225 #define GET_OP_LIST
226 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
227         >();
228   }
229 };
230 } // namespace
231 
232 #define GET_OP_CLASSES
233 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
234 
235 void mlir::scf::registerTransformDialectExtension(DialectRegistry &registry) {
236   registry.addExtensions<SCFTransformDialectExtension>();
237 }
238