Index: mlir/examples/toy/Ch6/mlir/LowerToLLVM.cpp =================================================================== --- mlir/examples/toy/Ch6/mlir/LowerToLLVM.cpp +++ mlir/examples/toy/Ch6/mlir/LowerToLLVM.cpp @@ -79,7 +79,7 @@ if (i != e - 1) rewriter.create(loc, printfRef, rewriter.getIntegerType(32), newLineCst); - rewriter.create(loc); + rewriter.create(loc); rewriter.setInsertionPointToStart(loop.getBody()); } Index: mlir/examples/toy/Ch7/mlir/LowerToLLVM.cpp =================================================================== --- mlir/examples/toy/Ch7/mlir/LowerToLLVM.cpp +++ mlir/examples/toy/Ch7/mlir/LowerToLLVM.cpp @@ -79,7 +79,7 @@ if (i != e - 1) rewriter.create(loc, printfRef, rewriter.getIntegerType(32), newLineCst); - rewriter.create(loc); + rewriter.create(loc); rewriter.setInsertionPointToStart(loop.getBody()); } Index: mlir/include/mlir/Dialect/LoopOps/LoopOps.td =================================================================== --- mlir/include/mlir/Dialect/LoopOps/LoopOps.td +++ mlir/include/mlir/Dialect/LoopOps/LoopOps.td @@ -37,10 +37,10 @@ def ForOp : Loop_Op<"for", [DeclareOpInterfaceMethods, - SingleBlockImplicitTerminator<"TerminatorOp">]> { + SingleBlockImplicitTerminator<"YieldOp">]> { let summary = "for operation"; let description = [{ - The "loop.for" operation represents a loop nest taking 3 SSA value as + 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 @@ -51,17 +51,71 @@ lower bound but does not include the upper bound. The body region must contain exactly one block that terminates with - "loop.terminator". Calling ForOp::build will create such region and insert - the terminator, so will the parsing even in cases when it is absent from the - custom format. For example: + "loop.yield". Calling ForOp::build will create such 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 will have 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" result variables 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) -> (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) -> (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); + let results = (outs Variadic:$results); + let arguments = (ins Index:$lowerBound, Index:$upperBound, Index:$step, Variadic:$initArgs); let regions = (region SizedRegion<1>:$region); let skipDefaultBuilders = 1; @@ -76,19 +130,40 @@ OpBuilder getBodyBuilder() { return OpBuilder(getBody(), std::prev(getBody()->end())); } + iterator_range getRegionIterArgs() { + auto iterArgsBegin = getBody()->args_begin(); + std::advance(iterArgsBegin, 1); + return {iterArgsBegin, getBody()->args_end()}; + } + iterator_range getIterOperands() { + auto operands = getOperation()->getOperands(); + auto iterOpndBegin = operands.begin(); + std::advance(iterOpndBegin, getNumCtrlOperands()); + return {iterOpndBegin, operands.end()}; + } + 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 + constexpr unsigned getNumCtrlOperands() { return 3; } + // does the operation hold operands for loop-carried values + bool hasIterOperands() { return getOperation()->getNumOperands() > getNumCtrlOperands(); } + // get Number of loop-carried values + unsigned getNumIterOperands() { return getOperation()->getNumOperands() - getNumCtrlOperands(); } }]; } def IfOp : Loop_Op<"if", - [SingleBlockImplicitTerminator<"TerminatorOp">]> { + [SingleBlockImplicitTerminator<"YieldOp">]> { 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. The operation produces no results. For example: + is a boolean value. For example: ```mlir loop.if %b { @@ -98,15 +173,34 @@ } ``` - The 'else' block is optional, and may be omitted. For - example: + "loop.if" may also return results that are defined in its regions. The values + defined are determined by which execution path is taken. + For 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 explicitly. A + Also, if "loop.if" defines no values, the 'else' block is optional, and may be omitted. + For example: ```mlir loop.if %b { ... } ``` }]; + let results = (outs Variadic:$results); let arguments = (ins I1:$condition); let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion); @@ -131,7 +225,7 @@ } def ParallelOp : Loop_Op<"parallel", - [SameVariadicOperandSize, SingleBlockImplicitTerminator<"TerminatorOp">]> { + [SameVariadicOperandSize, SingleBlockImplicitTerminator<"YieldOp">]> { let summary = "parallel for operation"; let description = [{ The "loop.parallel" operation represents a loop nest taking 3 groups of SSA @@ -157,7 +251,7 @@ the same number of results as it has reduce operations. The body region must contain exactly one block that terminates with - "loop.terminator". Parsing ParallelOp will create such region and insert the + "loop.yield". Parsing ParallelOp will create such region and insert the terminator when it is absent from the custom format. For example: ```mlir @@ -251,25 +345,23 @@ let arguments = (ins AnyType:$result); } -def TerminatorOp : Loop_Op<"terminator", [Terminator]> { - let summary = "cf terminator operation"; +def YieldOp : Loop_Op<"yield", [Terminator]> { + let summary = "cf yield and termination operation"; let description = [{ - "loop.terminator" is a special terminator operation for blocks inside - loops. It terminates the region. This operation does _not_ have a custom - syntax. However, `std` control operations omit the terminator in their - custom syntax for brevity. - - ```mlir - loop.terminator - ``` + "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, then parent op must define the same + number and types of the operands. + If the parent op defines no values, then the "loop.yield" can be + left out and will be inserted implicitly. Otherwise, it has + to be explicitly present. }]; - // No custom parsing/printing form. - let parser = ?; + let arguments = (ins Variadic:$results); + let builders = [ + OpBuilder<"Builder *builder, OperationState &result"> + ]; let printer = ?; - - // Fully specified by traits. - let verifier = ?; } - #endif // LOOP_OPS Index: mlir/include/mlir/IR/OpImplementation.h =================================================================== --- mlir/include/mlir/IR/OpImplementation.h +++ mlir/include/mlir/IR/OpImplementation.h @@ -634,6 +634,11 @@ virtual ParseResult parseOptionalColonTypeList(SmallVectorImpl &result) = 0; + /// Parse a list of assignments of the form + /// (%x = %y : type, ...) + virtual ParseResult + parseAssignmentList(SmallVectorImpl &lhs, SmallVectorImpl &rhs, SmallVectorImpl &types) = 0; + /// Parse a keyword followed by a type. ParseResult parseKeywordType(const char *keyword, Type &result) { return failure(parseKeyword(keyword) || parseType(result)); Index: mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp =================================================================== --- mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp +++ mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp @@ -332,7 +332,7 @@ PatternMatchResult matchAndRewrite(AffineTerminatorOp op, PatternRewriter &rewriter) const override { - rewriter.replaceOpWithNewOp(op); + rewriter.replaceOpWithNewOp(op); return matchSuccess(); } }; Index: mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp =================================================================== --- mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp +++ mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp @@ -42,14 +42,13 @@ ConversionPatternRewriter &rewriter) const override; }; -/// Pattern to erase a loop::TerminatorOp. -class TerminatorOpConversion final - : public SPIRVOpLowering { +/// Pattern to erase a loop::YieldOp. +class TerminatorOpConversion final : public SPIRVOpLowering { public: - using SPIRVOpLowering::SPIRVOpLowering; + using SPIRVOpLowering::SPIRVOpLowering; PatternMatchResult - matchAndRewrite(loop::TerminatorOp terminatorOp, ArrayRef operands, + matchAndRewrite(loop::YieldOp terminatorOp, ArrayRef operands, ConversionPatternRewriter &rewriter) const override { rewriter.eraseOp(terminatorOp); return matchSuccess(); Index: mlir/lib/Dialect/LoopOps/LoopOps.cpp =================================================================== --- mlir/lib/Dialect/LoopOps/LoopOps.cpp +++ mlir/lib/Dialect/LoopOps/LoopOps.cpp @@ -76,18 +76,60 @@ // Check that the body defines as single block argument for the induction // variable. auto *body = op.getBody(); - if (body->getNumArguments() != 1 || !body->getArgument(0).getType().isIndex()) - return op.emitOpError("expected body to have a single index argument for " - "the induction variable"); + if (!body->getArgument(0).getType().isIndex()) + return op.emitOpError( + "expected body first argument to be an index argument for " + "the induction variable"); + + // If ForOp define values check that number and types of of defined + // values match ForOp initial iter operands and backedge basic block arguments + auto opNumResults = op.getNumResults(); + if (opNumResults != 0) { + if (op.getNumIterOperands() != opNumResults) + op.emitOpError( + "mismatch in number of op loop-carried values and defined values"); + if (op.getNumRegionIterArgs() != opNumResults) + op.emitOpError( + "mismatch in number of basic block args and defined values"); + auto iterOperands = op.getIterOperands(); + auto iterArgs = op.getRegionIterArgs(); + auto opResults = op.getResults(); + for (auto e : llvm::zip(iterOperands, iterArgs, opResults)) { + if (std::get<0>(e).getType() != std::get<2>(e).getType()) + op.emitOpError() + << "types mismatch between iter operands and defined values"; + if (std::get<1>(e).getType() != std::get<2>(e).getType()) + op.emitOpError() + << "types mismatch between iter region args and defined values"; + } + } return success(); } static void print(OpAsmPrinter &p, ForOp op) { + bool printBlockTerminators = false; p << op.getOperationName() << " " << op.getInductionVar() << " = " << op.lowerBound() << " to " << op.upperBound() << " step " << op.step(); + + if (op.hasIterOperands()) { + unsigned numIterOperands = op.getNumIterOperands(), i = 0; + p << " iter_args("; + auto regionArgs = op.getRegionIterArgs(); + auto operands = op.getIterOperands(); + for (auto e : llvm::zip(regionArgs, operands)) { + auto operand = std::get<1>(e); + p << std::get<0>(e) << " = " << std::get<1>(e) << " : "; + p << operand.getType(); + p << ((++i == numIterOperands) ? ")" : ", "); + } + } + if (!op.results().empty()) { + p << " -> (" << op.getResultTypes() << ")"; + printBlockTerminators = true; + } p.printRegion(op.region(), /*printEntryBlockArgs=*/false, - /*printBlockTerminators=*/false); + /*printBlockTerminators=*/printBlockTerminators); p.printOptionalAttrDict(op.getAttrs()); } @@ -108,9 +150,25 @@ parser.resolveOperand(step, indexType, result.operands)) return failure(); + // Parse the optional initial iteration arguments + SmallVector regionArgs, operands; + SmallVector argTypes; + regionArgs.push_back(inductionVariable); + + if (!parser.parseOptionalKeyword("iter_args")) { + parser.parseAssignmentList(regionArgs, operands, argTypes); + // resolve input operands + for (auto operand_type : llvm::zip(operands, argTypes)) + parser.resolveOperand(std::get<0>(operand_type),std::get<1>(operand_type), result.operands); + } + argTypes.insert(argTypes.begin(), indexType); + + // Parse optional results type list + if (parser.parseOptionalArrowTypeList(result.types)) + return failure(); // Parse the body region. Region *body = result.addRegion(); - if (parser.parseRegion(*body, inductionVariable, indexType)) + if (parser.parseRegion(*body, regionArgs, argTypes)) return failure(); ForOp::ensureTerminator(*body, builder, result.location); @@ -168,6 +226,9 @@ return op.emitOpError( "requires that child entry blocks have no arguments"); } + if (op.getNumResults() != 0 && op.elseRegion().empty()) + return op.emitOpError("must have an else block if defining values"); + return success(); } @@ -183,7 +244,9 @@ if (parser.parseOperand(cond) || parser.resolveOperand(cond, i1Type, result.operands)) return failure(); - + // Parse optional results type list + if (parser.parseOptionalArrowTypeList(result.types)) + return failure(); // Parse the 'then' region. if (parser.parseRegion(*thenRegion, /*arguments=*/{}, /*argTypes=*/{})) return failure(); @@ -199,15 +262,21 @@ // Parse the optional attribute list. if (parser.parseOptionalAttrDict(result.attributes)) return failure(); - return success(); } static void print(OpAsmPrinter &p, IfOp op) { + bool printBlockTerminators = false; + p << IfOp::getOperationName() << " " << op.condition(); + if (!op.results().empty()) { + p << " -> (" << op.getResultTypes() << ")"; + // print yield explicitly if the op defines values + printBlockTerminators = true; + } p.printRegion(op.thenRegion(), /*printEntryBlockArgs=*/false, - /*printBlockTerminators=*/false); + /*printBlockTerminators=*/printBlockTerminators); // Print the 'else' regions if it exists and has a block. auto &elseRegion = op.elseRegion(); @@ -215,7 +284,7 @@ p << " else"; p.printRegion(elseRegion, /*printEntryBlockArgs=*/false, - /*printBlockTerminators=*/false); + /*printBlockTerminators=*/printBlockTerminators); } p.printOptionalAttrDict(op.getAttrs()); @@ -434,6 +503,45 @@ << op.result().getType(); } +//===----------------------------------------------------------------------===// +// YieldOp +//===----------------------------------------------------------------------===// +void YieldOp::build(Builder *builder, OperationState &result) {} + +static LogicalResult verify(YieldOp op) { + auto parentOp = op.getParentOp(); + auto results = parentOp->getResults(); + auto operands = op.getOperands(); + + if (!isa(parentOp) && !isa(parentOp) && + !isa(parentOp)) + return op.emitError() << "yield terminate If, For or Parallel regions"; + + if (isa(parentOp) || isa(parentOp)) { + if (parentOp->getNumResults() != op.getNumOperands()) + return op.emitOpError() << "parent of yield must have same number of " + "results as the yield operands"; + for (auto e : llvm::zip(results, operands)) { + if (std::get<0>(e).getType() != std::get<1>(e).getType()) + op.emitOpError() << "types mismatch between yield op and its parent"; + } + } + return success(); +} + +static ParseResult parseYieldOp(OpAsmParser &parser, OperationState &result) { + SmallVector operands; + SmallVector types; + llvm::SMLoc loc = parser.getCurrentLocation(); + // parser variadic opreands list, their types, and resolve operands to SSA + // values + if (parser.parseOperandList(operands) || parser.parseColonTypeList(types) || + parser.resolveOperands(operands, types, loc, result.operands)) + return failure(); + + return success(); +} + //===----------------------------------------------------------------------===// // TableGen'd op method definitions //===----------------------------------------------------------------------===// Index: mlir/lib/Parser/Parser.cpp =================================================================== --- mlir/lib/Parser/Parser.cpp +++ mlir/lib/Parser/Parser.cpp @@ -4326,6 +4326,25 @@ return parser.parseTypeListNoParens(result); } + /// Parse an assignment list of the form + /// (%a = %b : type1, ... ) + /// The list must contain at least one entry + ParseResult + parseAssignmentList(SmallVectorImpl &lhs, SmallVectorImpl &rhs, SmallVectorImpl &types) { + auto parseElt = [&]() -> ParseResult { + OperandType regionArg, operand; + Type type; + if (parseRegionArgument(regionArg) || parseEqual() || parseOperand(operand) || parseColon() || parseType(type)) + return failure(); + lhs.push_back(regionArg); + types.push_back(type); + rhs.push_back(operand); + return success(); + }; + if (parseLParen()) + return failure(); + return parser.parseCommaSeparatedListUntil(Token::r_paren, parseElt); + } private: /// The source location of the operation name. SMLoc nameLoc; Index: mlir/test/Dialect/Linalg/parallel_loops.mlir =================================================================== --- mlir/test/Dialect/Linalg/parallel_loops.mlir +++ mlir/test/Dialect/Linalg/parallel_loops.mlir @@ -27,7 +27,7 @@ // CHECK: %[[SUM_ELEM:.*]] = load %[[SUM]][%[[I]], %[[J]]] // CHECK: %[[SUM:.*]] = addf %[[LHS_ELEM]], %[[RHS_ELEM]] : f32 // CHECK: store %[[SUM]], %{{.*}}[%[[I]], %[[J]]] -// CHECK: "loop.terminator"() : () -> () +// CHECK: "loop.yield"() : () -> () // ----- Index: mlir/test/Dialect/Loops/invalid.mlir =================================================================== --- mlir/test/Dialect/Loops/invalid.mlir +++ mlir/test/Dialect/Loops/invalid.mlir @@ -29,7 +29,7 @@ %c0 = constant 0 : index "loop.for"(%arg0, %arg0, %c0) ({ ^bb0(%arg1: index): - "loop.terminator"() : () -> () + "loop.yield"() : () -> () }) : (index, index, index) -> () return } @@ -39,8 +39,8 @@ func @loop_for_one_region(%arg0: index) { // expected-error@+1 {{incorrect number of regions: expected 1 but found 2}} "loop.for"(%arg0, %arg0, %arg0) ( - {"loop.terminator"() : () -> ()}, - {"loop.terminator"() : () -> ()} + {"loop.yield"() : () -> ()}, + {"loop.yield"() : () -> ()} ) : (index, index, index) -> () return } @@ -52,9 +52,9 @@ "loop.for"(%arg0, %arg0, %arg0) ( { ^bb1: - "loop.terminator"() : () -> () + "loop.yield"() : () -> () ^bb2: - "loop.terminator"() : () -> () + "loop.yield"() : () -> () } ) : (index, index, index) -> () return @@ -63,11 +63,11 @@ // ----- func @loop_for_single_index_argument(%arg0: index) { - // expected-error@+1 {{expected body to have a single index argument for the induction variable}} + // expected-error@+1 {{op expected body first argument to be an index argument for the induction variable}} "loop.for"(%arg0, %arg0, %arg0) ( { ^bb0(%i0 : f32): - "loop.terminator"() : () -> () + "loop.yield"() : () -> () } ) : (index, index, index) -> () return @@ -95,9 +95,9 @@ // expected-error@+1 {{expects region #0 to have 0 or 1 blocks}} "loop.if"(%arg0) ({ ^bb0: - "loop.terminator"() : () -> () + "loop.yield"() : () -> () ^bb1: - "loop.terminator"() : () -> () + "loop.yield"() : () -> () }, {}): (i1) -> () return } @@ -108,7 +108,7 @@ // expected-error@+1 {{requires that child entry blocks have no arguments}} "loop.if"(%arg0) ({ ^bb0(%0 : index): - "loop.terminator"() : () -> () + "loop.yield"() : () -> () }, {}): (i1) -> () return } @@ -130,7 +130,7 @@ // expected-error@+1 {{'loop.parallel' op expects arguments for the induction variable to be of index type}} "loop.parallel"(%arg0, %arg1, %arg2) ({ ^bb0(%i0: f32): - "loop.terminator"() : () -> () + "loop.yield"() : () -> () }): (index, index, index) -> () return } @@ -142,7 +142,7 @@ // expected-error@+1 {{'loop.parallel' op expects the same number of induction variables as bound and step values}} "loop.parallel"(%arg0, %arg1, %arg2) ({ ^bb0(%i0: index, %i1: index): - "loop.terminator"() : () -> () + "loop.yield"() : () -> () }): (index, index, index) -> () return } @@ -265,7 +265,7 @@ // expected-error@+1 {{the block inside reduce should be terminated with a 'loop.reduce.return' op}} loop.reduce(%arg1) { ^bb0(%lhs : f32, %rhs : f32): - "loop.terminator"(): () -> () + "loop.yield"(): () -> () } : f32 } : f32 return @@ -294,3 +294,63 @@ }): () -> () return } + +// ----- + +func @std_if_incorrect_yield(%arg0: i1, %arg1: f32) +{ + %x, %y = loop.if %arg0 -> (f32, f32) { + %0 = addf %arg1, %arg1 : f32 + // expected-error@+1 {{parent of yield must have same number of results as the yield operands}} + loop.yield %0 : f32 + } else { + %0 = subf %arg1, %arg1 : f32 + loop.yield %0 : f32 + } + return +} + +// ----- + +func @std_if_missing_else(%arg0: i1, %arg1: f32) +{ + // expected-error@+1 {{must have an else block if defining values}} + %x = loop.if %arg0 -> (f32) { + %0 = addf %arg1, %arg1 : f32 + loop.yield %0 : f32 + } + return +} + +// ----- + +func @std_for_operands_mismatch(%arg0 : index, %arg1 : index, %arg2 : index) { + %s0 = constant 0.0 : f32 + %t0 = constant 1 : i32 + // expected-error@+2 {{mismatch in number of basic block args and defined values}} + // expected-error@+1 {{mismatch in number of op loop-carried values and defined values}} + %result1:3 = loop.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0 : f32, %ti = %t0 : i32) -> (f32, i32, f32) { + %sn = addf %si, %si : f32 + %tn = addi %ti, %ti : i32 + loop.yield %sn, %tn, %sn : f32, i32, f32 + } + return +} + +// ----- + +func @std_for_operands_mismatch_2(%arg0 : index, %arg1 : index, %arg2 : index) { + %s0 = constant 0.0 : f32 + %t0 = constant 1 : i32 + %u0 = constant 1.0 : f32 + // expected-error@+2 {{mismatch in number of op loop-carried values and defined values}} + // expected-error@+1 {{mismatch in number of basic block args and defined values}} + %result1:2 = loop.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0 : f32, %ti = %t0 : i32, %ui = %u0 : f32) -> (f32, i32) { + %sn = addf %si, %si : f32 + %tn = addi %ti, %ti : i32 + %un = subf %ui, %ui : f32 + // expected-error@+1 {{parent of yield must have same number of results as the yield operands}} + loop.yield %sn, %tn, %un : f32, i32, f32 + } + return +} Index: mlir/test/Dialect/Loops/ops.mlir =================================================================== --- mlir/test/Dialect/Loops/ops.mlir +++ mlir/test/Dialect/Loops/ops.mlir @@ -90,6 +90,117 @@ // CHECK-NEXT: %[[RES:.*]] = addf %[[LHS]], %[[RHS]] : f32 // CHECK-NEXT: loop.reduce.return %[[RES]] : f32 // CHECK-NEXT: } : f32 -// CHECK-NEXT: "loop.terminator"() : () -> () +// CHECK-NEXT: "loop.yield"() : () -> () // CHECK-NEXT: } : f32 -// CHECK-NEXT: "loop.terminator"() : () -> () +// CHECK-NEXT: "loop.yield"() : () -> () + + +func @std_if_yield(%arg0: i1, %arg1: f32) +{ + %x, %y = loop.if %arg0 -> (f32, f32) { + %0 = addf %arg1, %arg1 : f32 + %1 = subf %arg1, %arg1 : f32 + loop.yield %0, %1 : f32, f32 + } else { + %0 = subf %arg1, %arg1 : f32 + %1 = addf %arg1, %arg1 : f32 + loop.yield %0, %1 : f32, f32 + } + return +} +// CHECK-LABEL: func @std_if_yield( +// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]: +// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]: +// CHECK-NEXT: %{{.*}}:2 = loop.if %[[ARG0]] -> (f32, f32) { +// CHECK-NEXT: %[[T1:.*]] = addf %[[ARG1]], %[[ARG1]] +// CHECK-NEXT: %[[T2:.*]] = subf %[[ARG1]], %[[ARG1]] +// CHECK-NEXT: "loop.yield"(%[[T1]], %[[T2]]) : (f32, f32) -> () +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[T3:.*]] = subf %[[ARG1]], %[[ARG1]] +// CHECK-NEXT: %[[T4:.*]] = addf %[[ARG1]], %[[ARG1]] +// CHECK-NEXT: "loop.yield"(%[[T3]], %[[T4]]) : (f32, f32) -> () +// CHECK-NEXT: } + +func @std_for_yield(%arg0 : index, %arg1 : index, %arg2 : index) { + %s0 = constant 0.0 : f32 + %result = loop.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0 : f32) -> (f32) { + %sn = addf %si, %si : f32 + loop.yield %sn : f32 + } + return +} +// CHECK-LABEL: func @std_for_yield( +// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]: +// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]: +// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]: +// CHECK-NEXT: %[[INIT:.*]] = constant +// CHECK-NEXT: %{{.*}} = loop.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]] +// CHECK-SAME: iter_args(%[[ITER:.*]] = %[[INIT]] : f32) -> (f32) { +// CHECK-NEXT: %[[NEXT:.*]] = addf %[[ITER]], %[[ITER]] : f32 +// CHECK-NEXT: "loop.yield"(%[[NEXT]]) : (f32) -> () +// CHECK-NEXT: } + + +func @std_for_yield_multi(%arg0 : index, %arg1 : index, %arg2 : index) { + %s0 = constant 0.0 : f32 + %t0 = constant 1 : i32 + %u0 = constant 1.0 : f32 + %result1:3 = loop.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0 : f32, %ti = %t0 : i32, %ui = %u0 : f32) -> (f32, i32, f32) { + %sn = addf %si, %si : f32 + %tn = addi %ti, %ti : i32 + %un = subf %ui, %ui : f32 + loop.yield %sn, %tn, %un : f32, i32, f32 + } + return +} +// CHECK-LABEL: func @std_for_yield_multi( +// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]: +// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]: +// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]: +// CHECK-NEXT: %[[INIT1:.*]] = constant +// CHECK-NEXT: %[[INIT2:.*]] = constant +// CHECK-NEXT: %[[INIT3:.*]] = constant +// CHECK-NEXT: %{{.*}}:3 = loop.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]] +// CHECK-SAME: iter_args(%[[ITER1:.*]] = %[[INIT1]] : f32, %[[ITER2:.*]] = %[[INIT2]] : i32, %[[ITER3:.*]] = %[[INIT3]] : f32) -> (f32, i32, f32) { +// CHECK-NEXT: %[[NEXT1:.*]] = addf %[[ITER1]], %[[ITER1]] : f32 +// CHECK-NEXT: %[[NEXT2:.*]] = addi %[[ITER2]], %[[ITER2]] : i32 +// CHECK-NEXT: %[[NEXT3:.*]] = subf %[[ITER3]], %[[ITER3]] : f32 +// CHECK-NEXT: "loop.yield"(%[[NEXT1]], %[[NEXT2]], %[[NEXT3]]) : (f32, i32, f32) -> () + + +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) -> (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 +} +// CHECK-LABEL: func @conditional_reduce( +// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]] +// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]] +// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]] +// CHECK-SAME: %[[ARG3:[A-Za-z0-9]+]] +// CHECK-NEXT: %[[INIT:.*]] = constant +// CHECK-NEXT: %[[ZERO:.*]] = constant +// CHECK-NEXT: %[[RESULT:.*]] = loop.for %[[IV:.*]] = %[[ARG1]] to %[[ARG2]] step %[[ARG3]] +// CHECK-SAME: iter_args(%[[ITER:.*]] = %[[INIT]] : f32) -> (f32) { +// CHECK-NEXT: %[[T:.*]] = load %[[ARG0]][%[[IV]]] +// CHECK-NEXT: %[[COND:.*]] = cmpf "ugt", %[[T]], %[[ZERO]] +// CHECK-NEXT: %[[IFRES:.*]] = loop.if %[[COND]] -> (f32) { +// CHECK-NEXT: %[[THENRES:.*]] = addf %[[ITER]], %[[T]] +// CHECK-NEXT: "loop.yield"(%[[THENRES]]) : (f32) -> () +// CHECK-NEXT: } else { +// CHECK-NEXT: "loop.yield"(%[[ITER]]) : (f32) -> () +// CHECK-NEXT: } +// CHECK-NEXT: "loop.yield"(%[[IFRES]]) : (f32) -> () +// CHECK-NEXT: } +// CHECK-NEXT: return %[[RESULT]]