| //===- Ops.td - Loop operation definitions ---------------*- tablegen -*-===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // Defines MLIR loop operations. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #ifndef LOOP_OPS |
| #define LOOP_OPS |
| |
| include "mlir/Interfaces/LoopLikeInterface.td" |
| include "mlir/Interfaces/SideEffects.td" |
| |
| def LoopOps_Dialect : Dialect { |
| let name = "loop"; |
| let cppNamespace = ""; |
| } |
| |
| // Base class for Loop dialect ops. |
| class Loop_Op<string mnemonic, list<OpTrait> traits = []> : |
| Op<LoopOps_Dialect, mnemonic, traits> { |
| // For every standard op, there needs to be a: |
| // * void print(OpAsmPrinter &p, ${C++ class of Op} op) |
| // * LogicalResult verify(${C++ class of Op} op) |
| // * ParseResult parse${C++ class of Op}(OpAsmParser &parser, |
| // OperationState &result) |
| // functions. |
| let printer = [{ return ::print(p, *this); }]; |
| let verifier = [{ return ::verify(*this); }]; |
| let parser = [{ return ::parse$cppClass(parser, result); }]; |
| } |
| |
| def ForOp : Loop_Op<"for", |
| [DeclareOpInterfaceMethods<LoopLikeOpInterface>, |
| SingleBlockImplicitTerminator<"YieldOp">, |
| RecursiveSideEffects]> { |
| let summary = "for operation"; |
| let description = [{ |
| The "loop.for" operation represents a loop taking 3 SSA value as operands |
| that represent the lower bound, upper bound and step respectively. The |
| operation defines an SSA value for its induction variable. It has one |
| region capturing the loop body. The induction variable is represented as an |
| argument of this region. This SSA value always has type index, which is the |
| size of the machine word. The step is a value of type index, required to be |
| positive. |
| The lower and upper bounds specify a half-open range: the range includes |
| the lower bound but does not include the upper bound. |
| |
| The body region must contain exactly one block that terminates with |
| "loop.yield". Calling ForOp::build will create such a region and insert |
| the terminator implicitly if none is defined, so will the parsing even in |
| cases when it is absent from the custom format. For example: |
| |
| ```mlir |
| loop.for %iv = %lb to %ub step %step { |
| ... // body |
| } |
| ``` |
| |
| `loop.for` can also operate on loop-carried variables and returns the final |
| values after loop termination. The initial values of the variables are |
| passed as additional SSA operands to the "loop.for" following the 3 loop |
| control SSA values mentioned above (lower bound, upper bound and step). The |
| operation region has equivalent arguments for each variable representing |
| the value of the variable at the current iteration. |
| |
| The region must terminate with a "loop.yield" that passes all the current |
| iteration variables to the next iteration, or to the "loop.for" result, if |
| at the last iteration. "loop.for" results hold the final values after the |
| last iteration. |
| |
| For example, to sum-reduce a memref: |
| |
| ```mlir |
| func @reduce(%buffer: memref<1024xf32>, %lb: index, |
| %ub: index, %step: index) -> (f32) { |
| // Initial sum set to 0. |
| %sum_0 = constant 0.0 : f32 |
| // iter_args binds initial values to the loop's region arguments. |
| %sum = loop.for %iv = %lb to %ub step %step |
| iter_args(%sum_iter = %sum_0) -> (f32) { |
| %t = load %buffer[%iv] : memref<1024xf32> |
| %sum_next = addf %sum_iter, %t : f32 |
| // Yield current iteration sum to next iteration %sum_iter or to %sum |
| // if final iteration. |
| loop.yield %sum_next : f32 |
| } |
| return %sum : f32 |
| } |
| ``` |
| |
| If the "loop.for" defines any values, a yield must be explicitly present. |
| The number and types of the "loop.for" results must match the initial |
| values in the "iter_args" binding and the yield operands. |
| |
| Another example with a nested "loop.if" (see "loop.if" for details) to |
| perform conditional reduction: |
| |
| ```mlir |
| func @conditional_reduce(%buffer: memref<1024xf32>, %lb: index, |
| %ub: index, %step: index) -> (f32) { |
| %sum_0 = constant 0.0 : f32 |
| %c0 = constant 0.0 : f32 |
| %sum = loop.for %iv = %lb to %ub step %step |
| iter_args(%sum_iter = %sum_0) -> (f32) { |
| %t = load %buffer[%iv] : memref<1024xf32> |
| %cond = cmpf "ugt", %t, %c0 : f32 |
| %sum_next = loop.if %cond -> (f32) { |
| %new_sum = addf %sum_iter, %t : f32 |
| loop.yield %new_sum : f32 |
| } else { |
| loop.yield %sum_iter : f32 |
| } |
| loop.yield %sum_next : f32 |
| } |
| return %sum : f32 |
| } |
| ``` |
| }]; |
| let arguments = (ins Index:$lowerBound, |
| Index:$upperBound, |
| Index:$step, |
| Variadic<AnyType>:$initArgs); |
| let results = (outs Variadic<AnyType>:$results); |
| let regions = (region SizedRegion<1>:$region); |
| |
| let skipDefaultBuilders = 1; |
| let builders = [ |
| OpBuilder<"Builder *builder, OperationState &result, " |
| "Value lowerBound, Value upperBound, Value step, " |
| "ValueRange iterArgs = llvm::None"> |
| ]; |
| |
| let extraClassDeclaration = [{ |
| Block *getBody() { return ®ion().front(); } |
| Value getInductionVar() { return getBody()->getArgument(0); } |
| OpBuilder getBodyBuilder() { |
| return OpBuilder(getBody(), std::prev(getBody()->end())); |
| } |
| Block::BlockArgListType getRegionIterArgs() { |
| return getBody()->getArguments().drop_front(); |
| } |
| Operation::operand_range getIterOperands() { |
| return getOperands().drop_front(getNumControlOperands()); |
| } |
| |
| void setLowerBound(Value bound) { getOperation()->setOperand(0, bound); } |
| void setUpperBound(Value bound) { getOperation()->setOperand(1, bound); } |
| void setStep(Value step) { getOperation()->setOperand(2, step); } |
| |
| /// Number of region arguments for loop-carried values |
| unsigned getNumRegionIterArgs() { |
| return getBody()->getNumArguments() - 1; |
| } |
| /// Number of operands controlling the loop: lb, ub, step |
| unsigned getNumControlOperands() { return 3; } |
| /// Does the operation hold operands for loop-carried values |
| bool hasIterOperands() { |
| return getOperation()->getNumOperands() > getNumControlOperands(); |
| } |
| /// Get Number of loop-carried values |
| unsigned getNumIterOperands() { |
| return getOperation()->getNumOperands() - getNumControlOperands(); |
| } |
| }]; |
| } |
| |
| def IfOp : Loop_Op<"if", |
| [SingleBlockImplicitTerminator<"YieldOp">, RecursiveSideEffects]> { |
| let summary = "if-then-else operation"; |
| let description = [{ |
| The `loop.if` operation represents an if-then-else construct for |
| conditionally executing two regions of code. The operand to an if operation |
| is a boolean value. For example: |
| |
| ```mlir |
| loop.if %b { |
| ... |
| } else { |
| ... |
| } |
| ``` |
| |
| `loop.if` may also return results that are defined in its regions. The |
| values defined are determined by which execution path is taken. |
| |
| Example: |
| |
| ```mlir |
| %x, %y = loop.if %b -> (f32, f32) { |
| %x_true = ... |
| %y_true = ... |
| loop.yield %x_true, %y_true : f32, f32 |
| } else { |
| %x_false = ... |
| %y_false = ... |
| loop.yield %x_false, %y_false : f32, f32 |
| } |
| ``` |
| |
| `loop.if` regions are always terminated with "loop.yield". If "loop.if" |
| defines no values, the "loop.yield" can be left out, and will be inserted |
| implicitly. Otherwise, it must be explicit. |
| Also, if "loop.if" defines one or more values, the 'else' block cannot be |
| omitted. |
| |
| Example: |
| |
| ```mlir |
| loop.if %b { |
| ... |
| } |
| ``` |
| }]; |
| let arguments = (ins I1:$condition); |
| let results = (outs Variadic<AnyType>:$results); |
| let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion); |
| |
| let skipDefaultBuilders = 1; |
| let builders = [ |
| OpBuilder<"Builder *builder, OperationState &result, " |
| "Value cond, bool withElseRegion">, |
| OpBuilder<"Builder *builder, OperationState &result, " |
| "TypeRange resultTypes, Value cond, bool withElseRegion"> |
| ]; |
| |
| let extraClassDeclaration = [{ |
| OpBuilder getThenBodyBuilder() { |
| assert(!thenRegion().empty() && "Unexpected empty 'then' region."); |
| Block &body = thenRegion().front(); |
| return OpBuilder(&body, |
| results().empty() ? std::prev(body.end()) : body.end()); |
| } |
| OpBuilder getElseBodyBuilder() { |
| assert(!elseRegion().empty() && "Unexpected empty 'else' region."); |
| Block &body = elseRegion().front(); |
| return OpBuilder(&body, |
| results().empty() ? std::prev(body.end()) : body.end()); |
| } |
| }]; |
| } |
| |
| def ParallelOp : Loop_Op<"parallel", |
| [AttrSizedOperandSegments, |
| DeclareOpInterfaceMethods<LoopLikeOpInterface>, |
| SingleBlockImplicitTerminator<"YieldOp">]> { |
| let summary = "parallel for operation"; |
| let description = [{ |
| The "loop.parallel" operation represents a loop nest taking 4 groups of SSA |
| values as operands that represent the lower bounds, upper bounds, steps and |
| initial values, respectively. The operation defines a variadic number of |
| SSA values for its induction variables. It has one region capturing the |
| loop body. The induction variables are represented as an argument of this |
| region. These SSA values always have type index, which is the size of the |
| machine word. The steps are values of type index, required to be positive. |
| The lower and upper bounds specify a half-open range: the range includes |
| the lower bound but does not include the upper bound. The initial values |
| have the same types as results of "loop.parallel". If there are no results, |
| the keyword `init` can be omitted. |
| |
| Semantically we require that the iteration space can be iterated in any |
| order, and the loop body can be executed in parallel. If there are data |
| races, the behavior is undefined. |
| |
| The parallel loop operation supports reduction of values produced by |
| individual iterations into a single result. This is modeled using the |
| loop.reduce operation (see loop.reduce for details). Each result of a |
| loop.parallel operation is associated with an initial value operand and |
| reduce operation that is an immediate child. Reductions are matched to |
| result and initial values in order of their appearance in the body. |
| Consequently, we require that the body region has the same number of |
| results and initial values as it has reduce operations. |
| |
| The body region must contain exactly one block that terminates with |
| "loop.yield" without operands. Parsing ParallelOp will create such a region |
| and insert the terminator when it is absent from the custom format. |
| |
| Example: |
| |
| ```mlir |
| loop.parallel (%iv) = (%lb) to (%ub) step (%step) -> f32 { |
| %zero = constant 0.0 : f32 |
| loop.reduce(%zero) : f32 { |
| ^bb0(%lhs : f32, %rhs: f32): |
| %res = addf %lhs, %rhs : f32 |
| loop.reduce.return %res : f32 |
| } |
| } |
| ``` |
| }]; |
| |
| let arguments = (ins Variadic<Index>:$lowerBound, |
| Variadic<Index>:$upperBound, |
| Variadic<Index>:$step, |
| Variadic<AnyType>:$initVals); |
| let results = (outs Variadic<AnyType>:$results); |
| let regions = (region SizedRegion<1>:$region); |
| |
| let skipDefaultBuilders = 1; |
| let builders = [ |
| OpBuilder<"Builder *builder, OperationState &result, " |
| "ValueRange lowerBounds, ValueRange upperBounds, " |
| "ValueRange steps, ValueRange initVals = {}">, |
| ]; |
| |
| let extraClassDeclaration = [{ |
| Block *getBody() { return ®ion().front(); } |
| unsigned getNumInductionVars() { |
| return getBody()->getNumArguments(); |
| } |
| Block::BlockArgListType getInductionVars() { |
| return getBody()->getArguments(); |
| } |
| unsigned getNumLoops() { return step().size(); } |
| unsigned getNumReductions() { return initVals().size(); } |
| }]; |
| } |
| |
| def ReduceOp : Loop_Op<"reduce", [HasParent<"ParallelOp">]> { |
| let summary = "reduce operation for parallel for"; |
| let description = [{ |
| "loop.reduce" is an operation occurring inside "loop.parallel" operations. |
| It consists of one block with two arguments which have the same type as the |
| operand of "loop.reduce". |
| |
| "loop.reduce" is used to model the value for reduction computations of a |
| "loop.parallel" operation. It has to appear as an immediate child of a |
| "loop.parallel" and is associated with a result value of its parent |
| operation. |
| |
| Association is in the order of appearance in the body where the first |
| result of a parallel loop operation corresponds to the first "loop.reduce" |
| in the operation's body region. The reduce operation takes a single |
| operand, which is the value to be used in the reduction. |
| |
| The reduce operation contains a region whose entry block expects two |
| arguments of the same type as the operand. As the iteration order of the |
| parallel loop and hence reduction order is unspecified, the result of |
| reduction may be non-deterministic unless the operation is associative and |
| commutative. |
| |
| The result of the reduce operation's body must have the same type as the |
| operands and associated result value of the parallel loop operation. |
| Example: |
| |
| ```mlir |
| %operand = constant 1.0 : f32 |
| loop.reduce(%operand) : f32 { |
| ^bb0(%lhs : f32, %rhs: f32): |
| %res = addf %lhs, %rhs : f32 |
| loop.reduce.return %res : f32 |
| } |
| ``` |
| }]; |
| |
| let skipDefaultBuilders = 1; |
| let builders = [ |
| OpBuilder<"Builder *builder, OperationState &result, " |
| "Value operand"> |
| ]; |
| |
| let arguments = (ins AnyType:$operand); |
| let regions = (region SizedRegion<1>:$reductionOperator); |
| } |
| |
| def ReduceReturnOp : |
| Loop_Op<"reduce.return", [HasParent<"ReduceOp">, NoSideEffect, |
| Terminator]> { |
| let summary = "terminator for reduce operation"; |
| let description = [{ |
| "loop.reduce.return" is a special terminator operation for the block inside |
| "loop.reduce". It terminates the region. It should have the same type as |
| the operand of "loop.reduce". Example for the custom format: |
| |
| ```mlir |
| loop.reduce.return %res : f32 |
| ``` |
| }]; |
| |
| let arguments = (ins AnyType:$result); |
| let assemblyFormat = "$result attr-dict `:` type($result)"; |
| } |
| |
| def YieldOp : Loop_Op<"yield", [NoSideEffect, Terminator]> { |
| let summary = "loop yield and termination operation"; |
| let description = [{ |
| "loop.yield" yields an SSA value from a loop dialect op region and |
| terminates the regions. The semantics of how the values are yielded is |
| defined by the parent operation. |
| If "loop.yield" has any operands, the operands must match the parent |
| operation's results. |
| If the parent operation defines no values, then the "loop.yield" may be |
| left out in the custom syntax and the builders will insert one implicitly. |
| Otherwise, it has to be present in the syntax to indicate which values are |
| yielded. |
| }]; |
| |
| let arguments = (ins Variadic<AnyType>:$results); |
| let builders = [ |
| OpBuilder<"Builder *builder, OperationState &result", |
| [{ /* nothing to do */ }]> |
| ]; |
| } |
| #endif // LOOP_OPS |