1 | //===- SCFOps.h - Structured Control Flow -----------------------*- C++ -*-===// |
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 structured control flow operations. |
10 | // |
11 | //===----------------------------------------------------------------------===// |
12 | |
13 | #ifndef MLIR_DIALECT_SCF_SCF_H |
14 | #define MLIR_DIALECT_SCF_SCF_H |
15 | |
16 | #include "mlir/Dialect/Arith/Utils/Utils.h" |
17 | #include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h" |
18 | #include "mlir/IR/Builders.h" |
19 | #include "mlir/IR/BuiltinTypes.h" |
20 | #include "mlir/IR/RegionKindInterface.h" |
21 | #include "mlir/Interfaces/ControlFlowInterfaces.h" |
22 | #include "mlir/Interfaces/DestinationStyleOpInterface.h" |
23 | #include "mlir/Interfaces/InferTypeOpInterface.h" |
24 | #include "mlir/Interfaces/LoopLikeInterface.h" |
25 | #include "mlir/Interfaces/ParallelCombiningOpInterface.h" |
26 | #include "mlir/Interfaces/SideEffectInterfaces.h" |
27 | #include "mlir/Interfaces/ViewLikeInterface.h" |
28 | |
29 | namespace mlir { |
30 | namespace scf { |
31 | void buildTerminatedBody(OpBuilder &builder, Location loc); |
32 | } // namespace scf |
33 | } // namespace mlir |
34 | |
35 | #include "mlir/Dialect/SCF/IR/SCFOpsDialect.h.inc" |
36 | |
37 | #define GET_OP_CLASSES |
38 | #include "mlir/Dialect/SCF/IR/SCFOps.h.inc" |
39 | |
40 | namespace mlir { |
41 | namespace scf { |
42 | |
43 | // Insert `loop.yield` at the end of the only region's only block if it |
44 | // does not have a terminator already. If a new `loop.yield` is inserted, |
45 | // the location is specified by `loc`. If the region is empty, insert a new |
46 | // block first. |
47 | void ensureLoopTerminator(Region ®ion, Builder &builder, Location loc); |
48 | |
49 | /// Returns the loop parent of an induction variable. If the provided value is |
50 | /// not an induction variable, then return nullptr. |
51 | ForOp getForInductionVarOwner(Value val); |
52 | |
53 | /// Returns the parallel loop parent of an induction variable. If the provided |
54 | /// value is not an induction variable, then return nullptr. |
55 | ParallelOp getParallelForInductionVarOwner(Value val); |
56 | |
57 | /// Returns the ForallOp parent of an thread index variable. |
58 | /// If the provided value is not a thread index variable, then return nullptr. |
59 | ForallOp getForallOpThreadIndexOwner(Value val); |
60 | |
61 | /// Return true if ops a and b (or their ancestors) are in mutually exclusive |
62 | /// regions/blocks of an IfOp. |
63 | // TODO: Consider moving this functionality to RegionBranchOpInterface. |
64 | bool insideMutuallyExclusiveBranches(Operation *a, Operation *b); |
65 | |
66 | /// Promotes the loop body of a scf::ForallOp to its containing block. |
67 | void promote(RewriterBase &rewriter, scf::ForallOp forallOp); |
68 | |
69 | /// An owning vector of values, handy to return from functions. |
70 | using ValueVector = SmallVector<Value>; |
71 | using LoopVector = SmallVector<scf::ForOp>; |
72 | struct LoopNest { |
73 | LoopVector loops; |
74 | ValueVector results; |
75 | }; |
76 | |
77 | /// Creates a perfect nest of "for" loops, i.e. all loops but the innermost |
78 | /// contain only another loop and a terminator. The lower, upper bounds and |
79 | /// steps are provided as `lbs`, `ubs` and `steps`, which are expected to be of |
80 | /// the same size. `iterArgs` points to the initial values of the loop iteration |
81 | /// arguments, which will be forwarded through the nest to the innermost loop. |
82 | /// The body of the loop is populated using `bodyBuilder`, which accepts an |
83 | /// ordered list of induction variables of all loops, followed by a list of |
84 | /// iteration arguments of the innermost loop, in the same order as provided to |
85 | /// `iterArgs`. This function is expected to return as many values as |
86 | /// `iterArgs`, of the same type and in the same order, that will be treated as |
87 | /// yielded from the loop body and forwarded back through the loop nest. If the |
88 | /// function is not provided, the loop nest is not expected to have iteration |
89 | /// arguments, the body of the innermost loop will be left empty, containing |
90 | /// only the zero-operand terminator. Returns the LoopNest containing the list |
91 | /// of perfectly nest scf::ForOp build during the call. |
92 | /// If bound arrays are empty, the body builder will be called |
93 | /// once to construct the IR outside of the loop with an empty list of induction |
94 | /// variables. |
95 | LoopNest buildLoopNest( |
96 | OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs, |
97 | ValueRange steps, ValueRange iterArgs, |
98 | function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)> |
99 | bodyBuilder = nullptr); |
100 | |
101 | /// A convenience version for building loop nests without iteration arguments |
102 | /// (like for reductions). Does not take the initial value of reductions or |
103 | /// expect the body building functions to return their current value. |
104 | /// The built nested scf::For are captured in `capturedLoops` when non-null. |
105 | LoopNest buildLoopNest(OpBuilder &builder, Location loc, ValueRange lbs, |
106 | ValueRange ubs, ValueRange steps, |
107 | function_ref<void(OpBuilder &, Location, ValueRange)> |
108 | bodyBuilder = nullptr); |
109 | |
110 | } // namespace scf |
111 | } // namespace mlir |
112 | #endif // MLIR_DIALECT_SCF_SCF_H |
113 | |