xref: /llvm-project/mlir/include/mlir/Interfaces/LoopLikeInterface.td (revision db791b278a414fb6df1acc1799adcf11d8fb9169)
1//===- LoopLikeInterface.td - LoopLike interface -----------*- tablegen -*-===//
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// Defines the interface for loop-like operations.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef MLIR_INTERFACES_LOOPLIKEINTERFACE
14#define MLIR_INTERFACES_LOOPLIKEINTERFACE
15
16include "mlir/IR/OpBase.td"
17
18//===----------------------------------------------------------------------===//
19// Interfaces
20//===----------------------------------------------------------------------===//
21
22def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
23  let description = [{
24    Contains helper functions to query properties and perform transformations
25    of a loop. Operations that implement this interface will be considered by
26    loop-invariant code motion.
27
28    Loop-carried variables can be exposed through this interface. There are
29    3 components to a loop-carried variable.
30    - The "region iter_arg" is the block argument of the entry block that
31      represents the loop-carried variable in each iteration.
32    - The "init value" is an operand of the loop op that serves as the initial
33      region iter_arg value for the first iteration (if any).
34    - The "yielded" value is the value that is forwarded from one iteration to
35      serve as the region iter_arg of the next iteration.
36
37    If one of the respective interface methods is implemented, so must the other
38    two. The interface verifier ensures that the number of types of the region
39    iter_args, init values and yielded values match.
40
41    Optionally, "loop results" can be exposed through this interface. These are
42    the values that are returned from the loop op when there are no more
43    iterations. The number and types of the loop results must match with the
44    region iter_args. Note: Loop results are optional because some loops
45    (e.g., `scf.while`) may produce results that do match 1-to-1 with the
46    region iter_args.
47  }];
48  let cppNamespace = "::mlir";
49
50  let methods = [
51    InterfaceMethod<[{
52        Returns true if the given value is defined outside of the loop.
53        A sensible implementation could be to check whether the value's defining
54        operation lies outside of the loops body region. If the loop uses
55        explicit capture of dependencies, an implementation could check whether
56        the value corresponds to a captured dependency.
57      }],
58      /*retTy=*/"bool",
59      /*methodName=*/"isDefinedOutsideOfLoop",
60      /*args=*/(ins "::mlir::Value ":$value),
61      /*methodBody=*/"",
62      /*defaultImplementation=*/[{
63        return !$_op->isAncestor(value.getParentRegion()->getParentOp());
64      }]
65    >,
66    InterfaceMethod<[{
67        Returns the regions that make up the body of the loop and should be
68        inspected for loop-invariant operations.
69      }],
70      /*retTy=*/"::llvm::SmallVector<::mlir::Region *>",
71      /*methodName=*/"getLoopRegions"
72    >,
73    InterfaceMethod<[{
74        Moves the given loop-invariant operation out of the loop.
75      }],
76      /*retTy=*/"void",
77      /*methodName=*/"moveOutOfLoop",
78      /*args=*/(ins "::mlir::Operation *":$op),
79      /*methodBody=*/"",
80      /*defaultImplementation=*/"op->moveBefore($_op);"
81    >,
82    InterfaceMethod<[{
83        Promotes the loop body to its containing block if the loop is known to
84        have a single iteration. Returns "success" if the promotion was
85        successful.
86      }],
87      /*retTy=*/"::llvm::LogicalResult",
88      /*methodName=*/"promoteIfSingleIteration",
89      /*args=*/(ins "::mlir::RewriterBase &":$rewriter),
90      /*methodBody=*/"",
91      /*defaultImplementation=*/[{
92        return ::mlir::failure();
93      }]
94    >,
95    InterfaceMethod<[{
96        Return all induction variables, if they exist. If the op has no notion of
97        induction variable, then return std::nullopt. If it does have
98        a notion but an instance doesn't have induction variables, then
99        return empty vector.
100      }],
101      /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::Value>>",
102      /*methodName=*/"getLoopInductionVars",
103      /*args=*/(ins),
104      /*methodBody=*/"",
105      /*defaultImplementation=*/[{
106        return ::std::nullopt;
107      }]
108    >,
109    InterfaceMethod<[{
110        Return all lower bounds, if they exist. If the op has no notion of
111        lower bounds, then return std::nullopt. If it does have
112        a notion but an instance doesn't have lower bounds, then
113        return empty vector.
114      }],
115      /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
116      /*methodName=*/"getLoopLowerBounds",
117      /*args=*/(ins),
118      /*methodBody=*/"",
119      /*defaultImplementation=*/[{
120        return ::std::nullopt;
121      }]
122    >,
123    InterfaceMethod<[{
124        Return all steps, if they exist. If the op has no notion of
125        steps, then return std::nullopt. If it does have
126        a notion but an instance doesn't have steps, then
127        return empty vector.
128      }],
129      /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
130      /*methodName=*/"getLoopSteps",
131      /*args=*/(ins),
132      /*methodBody=*/"",
133      /*defaultImplementation=*/[{
134        return ::std::nullopt;
135      }]
136    >,
137    InterfaceMethod<[{
138        Return all upper bounds, if they exist. If the op has no notion of
139        lower bounds, then return std::nullopt. If it does have
140        a notion but an instance doesn't have lower bounds, then
141        return empty vector.
142      }],
143      /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
144      /*methodName=*/"getLoopUpperBounds",
145      /*args=*/(ins),
146      /*methodBody=*/"",
147      /*defaultImplementation=*/[{
148        return ::std::nullopt;
149      }]
150    >,
151    InterfaceMethod<[{
152        Return the mutable "init" operands that are used as initialization
153        values for the region "iter_args" of this loop.
154      }],
155      /*retTy=*/"::llvm::MutableArrayRef<::mlir::OpOperand>",
156      /*methodName=*/"getInitsMutable",
157      /*args=*/(ins),
158      /*methodBody=*/"",
159      /*defaultImplementation=*/[{
160        return {};
161      }]
162    >,
163    InterfaceMethod<[{
164        Return the region "iter_args" (block arguments) that correspond to the
165        "init" operands. If the op has multiple regions, return the
166        corresponding block arguments of the entry region.
167      }],
168      /*retTy=*/"::mlir::Block::BlockArgListType",
169      /*methodName=*/"getRegionIterArgs",
170      /*args=*/(ins),
171      /*methodBody=*/"",
172      /*defaultImplementation=*/[{
173        return ::mlir::Block::BlockArgListType();
174      }]
175    >,
176    InterfaceMethod<[{
177        Return the mutable operand range of values that are yielded to the next
178        iteration by the loop terminator.
179
180        For loop operations that dont yield a value, this should return
181        std::nullopt.
182      }],
183      /*retTy=*/"std::optional<::llvm::MutableArrayRef<::mlir::OpOperand>>",
184      /*methodName=*/"getYieldedValuesMutable",
185      /*args=*/(ins),
186      /*methodBody=*/"",
187      /*defaultImplementation=*/[{
188        return std::nullopt;
189      }]
190    >,
191    InterfaceMethod<[{
192        Return the range of results that are return from this loop and
193        correspond to the "init" operands.
194
195        Note: This interface method is optional. If loop results are not
196        exposed via this interface, "std::nullopt" should be returned.
197        Otherwise, the number and types of results must match with the
198        region iter_args, inits and yielded values that are exposed via this
199        interface. If loop results are exposed but this loop op has no
200        loop-carried variables, an empty result range (and not "std::nullopt")
201        should be returned.
202      }],
203      /*retTy=*/"::std::optional<::mlir::ResultRange>",
204      /*methodName=*/"getLoopResults",
205      /*args=*/(ins),
206      /*methodBody=*/"",
207      /*defaultImplementation=*/[{
208        return ::std::nullopt;
209      }]
210    >,
211    InterfaceMethod<[{
212        Append the specified additional "init" operands: replace this loop with
213        a new loop that has the additional init operands. The loop body of
214        this loop is moved over to the new loop.
215
216        `newInitOperands` specifies the additional "init" operands.
217        `newYieldValuesFn` is a function that returns the yielded values (which
218        can be computed based on the additional region iter_args). If
219        `replaceInitOperandUsesInLoop` is set, all uses of the additional init
220        operands inside of this loop are replaced with the corresponding, newly
221        added region iter_args.
222
223        Note: Loops that do not support init/iter_args should return "failure".
224      }],
225      /*retTy=*/"::mlir::FailureOr<::mlir::LoopLikeOpInterface>",
226      /*methodName=*/"replaceWithAdditionalYields",
227      /*args=*/(ins "::mlir::RewriterBase &":$rewriter,
228                    "::mlir::ValueRange":$newInitOperands,
229                    "bool":$replaceInitOperandUsesInLoop,
230                    "const ::mlir::NewYieldValuesFn &":$newYieldValuesFn),
231      /*methodBody=*/"",
232      /*defaultImplementation=*/[{
233        return ::mlir::failure();
234      }]
235    >
236  ];
237
238  let extraClassDeclaration = [{
239    /// Returns if a block is inside a loop (within the current function). This
240    /// can either be because the block is nested inside a LoopLikeInterface, or
241    /// because the control flow graph is cyclic
242    static bool blockIsInLoop(Block *block);
243  }];
244
245  let extraSharedClassDeclaration = [{
246    /// If there is a single induction variable return it, otherwise return
247    /// std::nullopt.
248    ::std::optional<::mlir::Value> getSingleInductionVar() {
249      auto inductionVars = $_op.getLoopInductionVars();
250      if (inductionVars.has_value() && (*inductionVars).size() == 1)
251          return (*inductionVars)[0];
252        return std::nullopt;
253    }
254    /// Return the single lower bound value or attribute if it exists, otherwise
255    /// return std::nullopt.
256    ::std::optional<::mlir::OpFoldResult> getSingleLowerBound() {
257      auto lowerBounds = $_op.getLoopLowerBounds();
258      if (lowerBounds.has_value() && (*lowerBounds).size() == 1)
259          return (*lowerBounds)[0];
260      return std::nullopt;
261    }
262    /// Return the single step value or attribute if it exists, otherwise
263    /// return std::nullopt.
264    ::std::optional<::mlir::OpFoldResult> getSingleStep() {
265      auto steps = $_op.getLoopSteps();
266      if (steps.has_value() && (*steps).size() == 1)
267          return (*steps)[0];
268      return std::nullopt;
269    }
270    /// Return the single upper bound value or attribute if it exists, otherwise
271    /// return std::nullopt.
272    ::std::optional<::mlir::OpFoldResult> getSingleUpperBound() {
273      auto upperBounds = $_op.getLoopUpperBounds();
274      if (upperBounds.has_value() && (*upperBounds).size() == 1)
275          return (*upperBounds)[0];
276      return std::nullopt;
277    }
278
279    /// Append the specified additional "init" operands: replace this loop with
280    /// a new loop that has the additional init operands. The loop body of this
281    /// loop is moved over to the new loop.
282    ///
283    /// The newly added region iter_args are yielded from the loop.
284    ::mlir::FailureOr<::mlir::LoopLikeOpInterface>
285        replaceWithAdditionalIterOperands(::mlir::RewriterBase &rewriter,
286                                          ::mlir::ValueRange newInitOperands,
287                                          bool replaceInitOperandUsesInLoop) {
288      return $_op.replaceWithAdditionalYields(
289          rewriter, newInitOperands, replaceInitOperandUsesInLoop,
290          [](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
291            return SmallVector<Value>(newBBArgs);
292          });
293    }
294
295    /// Return the values that are yielded to the next iteration. If
296    /// the loop doesnt yield any values return `{}`.
297    ::mlir::ValueRange getYieldedValues() {
298      auto mutableValues = $_op.getYieldedValuesMutable();
299      if (!mutableValues || mutableValues->empty())
300        return {};
301      Operation *yieldOp = mutableValues->begin()->getOwner();
302      unsigned firstOperandIndex = mutableValues->begin()->getOperandNumber();
303      return OperandRange(
304          yieldOp->operand_begin() + firstOperandIndex,
305          yieldOp->operand_begin() + firstOperandIndex + mutableValues->size());
306    }
307
308    /// Return the "init" operands that are used as initialization values for
309    /// the region "iter_args" of this loop.
310    ::mlir::OperandRange getInits() {
311      auto initsMutable = $_op.getInitsMutable();
312      if (initsMutable.empty())
313        return ::mlir::OperandRange($_op->operand_end(), $_op->operand_end());
314      unsigned firstOperandIndex = initsMutable.begin()->getOperandNumber();
315      return OperandRange(
316          $_op->operand_begin() + firstOperandIndex,
317          $_op->operand_begin() + firstOperandIndex + initsMutable.size());
318    }
319
320    /// Return the region iter_arg that corresponds to the given init operand.
321    /// Return an "empty" block argument if the given operand is not an init
322    /// operand of this loop op.
323    BlockArgument getTiedLoopRegionIterArg(OpOperand *opOperand) {
324      auto initsMutable = $_op.getInitsMutable();
325      auto it = llvm::find(initsMutable, *opOperand);
326      if (it == initsMutable.end())
327        return {};
328      return $_op.getRegionIterArgs()[std::distance(initsMutable.begin(), it)];
329    }
330
331    /// Return the region iter_arg that corresponds to the given loop result.
332    /// Return an "empty" block argument if the given OpResult is not a loop
333    /// result or if this op does not expose any loop results.
334    BlockArgument getTiedLoopRegionIterArg(OpResult opResult) {
335      auto loopResults = $_op.getLoopResults();
336      if (!loopResults)
337        return {};
338      auto it = llvm::find(*loopResults, opResult);
339      if (it == loopResults->end())
340        return {};
341      return $_op.getRegionIterArgs()[std::distance(loopResults->begin(), it)];
342    }
343
344    /// Return the init operand that corresponds to the given region iter_arg.
345    /// Return "nullptr" if the given block argument is not a region iter_arg
346    /// of this loop op.
347    OpOperand *getTiedLoopInit(BlockArgument bbArg) {
348      auto iterArgs = $_op.getRegionIterArgs();
349      auto it = llvm::find(iterArgs, bbArg);
350      if (it == iterArgs.end())
351        return {};
352      return &$_op.getInitsMutable()[std::distance(iterArgs.begin(), it)];
353    }
354
355    /// Return the init operand that corresponds to the given loop result.
356    /// Return "nullptr" if the given OpResult is not a loop result or if this
357    /// op does not expose any loop results.
358    OpOperand *getTiedLoopInit(OpResult opResult) {
359      auto loopResults = $_op.getLoopResults();
360      if (!loopResults)
361        return nullptr;
362      auto it = llvm::find(*loopResults, opResult);
363      if (it == loopResults->end())
364        return nullptr;
365      return &$_op.getInitsMutable()[std::distance(loopResults->begin(), it)];
366    }
367
368    /// Return the yielded value that corresponds to the given region iter_arg.
369    /// Return "nullptr" if the given block argument is not a region iter_arg
370    /// of this loop op or if there is no yield corresponding to this `bbArg`.
371    OpOperand *getTiedLoopYieldedValue(BlockArgument bbArg) {
372      auto iterArgs = $_op.getRegionIterArgs();
373      auto it = llvm::find(iterArgs, bbArg);
374      if (it == iterArgs.end())
375        return {};
376      std::optional<llvm::MutableArrayRef<::mlir::OpOperand>> yieldValues =
377        $_op.getYieldedValuesMutable();
378      if (!yieldValues)
379        return {};
380      return &yieldValues.value()[std::distance(iterArgs.begin(), it)];
381    }
382
383    /// Return the loop result that corresponds to the given init operand.
384    /// Return an "empty" OpResult if the given operand is not an init operand
385    /// of this loop op or if this op does not expose any loop results.
386    OpResult getTiedLoopResult(OpOperand *opOperand) {
387      auto loopResults = $_op.getLoopResults();
388      if (!loopResults)
389        return {};
390      auto initsMutable = $_op.getInitsMutable();
391      auto it = llvm::find(initsMutable, *opOperand);
392      if (it == initsMutable.end())
393        return {};
394      return (*loopResults)[std::distance(initsMutable.begin(), it)];
395    }
396
397    /// Return the loop result that corresponds to the given region iter_arg.
398    /// Return an "empty" OpResult if the given block argument is not a region
399    /// iter_arg of this loop op or if this op does not expose any loop results.
400    OpResult getTiedLoopResult(BlockArgument bbArg) {
401      auto loopResults = $_op.getLoopResults();
402      if (!loopResults)
403        return {};
404      auto iterArgs = $_op.getRegionIterArgs();
405      auto it = llvm::find(iterArgs, bbArg);
406      if (it == iterArgs.end())
407        return {};
408      return (*loopResults)[std::distance(iterArgs.begin(), it)];
409    }
410  }];
411
412  let verifyWithRegions = 1;
413
414  let verify = [{
415    return detail::verifyLoopLikeOpInterface($_op);
416  }];
417}
418
419//===----------------------------------------------------------------------===//
420// Traits
421//===----------------------------------------------------------------------===//
422
423// Op contains a region with parallel execution semantics
424def HasParallelRegion : NativeOpTrait<"HasParallelRegion">;
425
426#endif // MLIR_INTERFACES_LOOPLIKEINTERFACE
427