| //===- Ops.cpp - Loop MLIR Operations -------------------------------------===// |
| // |
| // Part of the MLIR 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 |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Dialect/LoopOps/LoopOps.h" |
| #include "mlir/Dialect/StandardOps/Ops.h" |
| #include "mlir/IR/AffineExpr.h" |
| #include "mlir/IR/AffineMap.h" |
| #include "mlir/IR/Builders.h" |
| #include "mlir/IR/Function.h" |
| #include "mlir/IR/Matchers.h" |
| #include "mlir/IR/Module.h" |
| #include "mlir/IR/OpImplementation.h" |
| #include "mlir/IR/PatternMatch.h" |
| #include "mlir/IR/StandardTypes.h" |
| #include "mlir/IR/Value.h" |
| #include "mlir/Support/MathExtras.h" |
| #include "mlir/Support/STLExtras.h" |
| #include "mlir/Transforms/SideEffectsInterface.h" |
| |
| using namespace mlir; |
| using namespace mlir::loop; |
| |
| //===----------------------------------------------------------------------===// |
| // LoopOpsDialect Interfaces |
| //===----------------------------------------------------------------------===// |
| namespace { |
| |
| struct LoopSideEffectsInterface : public SideEffectsDialectInterface { |
| using SideEffectsDialectInterface::SideEffectsDialectInterface; |
| |
| SideEffecting isSideEffecting(Operation *op) const override { |
| if (isa<IfOp>(op) || isa<ForOp>(op)) { |
| return Recursive; |
| } |
| return SideEffectsDialectInterface::isSideEffecting(op); |
| }; |
| }; |
| |
| } // namespace |
| |
| //===----------------------------------------------------------------------===// |
| // LoopOpsDialect |
| //===----------------------------------------------------------------------===// |
| |
| LoopOpsDialect::LoopOpsDialect(MLIRContext *context) |
| : Dialect(getDialectNamespace(), context) { |
| addOperations< |
| #define GET_OP_LIST |
| #include "mlir/Dialect/LoopOps/LoopOps.cpp.inc" |
| >(); |
| addInterfaces<LoopSideEffectsInterface>(); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // ForOp |
| //===----------------------------------------------------------------------===// |
| |
| void ForOp::build(Builder *builder, OperationState &result, Value lb, Value ub, |
| Value step) { |
| result.addOperands({lb, ub, step}); |
| Region *bodyRegion = result.addRegion(); |
| ForOp::ensureTerminator(*bodyRegion, *builder, result.location); |
| bodyRegion->front().addArgument(builder->getIndexType()); |
| } |
| |
| LogicalResult verify(ForOp op) { |
| if (auto cst = dyn_cast_or_null<ConstantIndexOp>(op.step().getDefiningOp())) |
| if (cst.getValue() <= 0) |
| return op.emitOpError("constant step operand must be positive"); |
| |
| // 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"); |
| return success(); |
| } |
| |
| static void print(OpAsmPrinter &p, ForOp op) { |
| p << op.getOperationName() << " " << op.getInductionVar() << " = " |
| << op.lowerBound() << " to " << op.upperBound() << " step " << op.step(); |
| p.printRegion(op.region(), |
| /*printEntryBlockArgs=*/false, |
| /*printBlockTerminators=*/false); |
| p.printOptionalAttrDict(op.getAttrs()); |
| } |
| |
| static ParseResult parseForOp(OpAsmParser &parser, OperationState &result) { |
| auto &builder = parser.getBuilder(); |
| OpAsmParser::OperandType inductionVariable, lb, ub, step; |
| // Parse the induction variable followed by '='. |
| if (parser.parseRegionArgument(inductionVariable) || parser.parseEqual()) |
| return failure(); |
| |
| // Parse loop bounds. |
| Type indexType = builder.getIndexType(); |
| if (parser.parseOperand(lb) || |
| parser.resolveOperand(lb, indexType, result.operands) || |
| parser.parseKeyword("to") || parser.parseOperand(ub) || |
| parser.resolveOperand(ub, indexType, result.operands) || |
| parser.parseKeyword("step") || parser.parseOperand(step) || |
| parser.resolveOperand(step, indexType, result.operands)) |
| return failure(); |
| |
| // Parse the body region. |
| Region *body = result.addRegion(); |
| if (parser.parseRegion(*body, inductionVariable, indexType)) |
| return failure(); |
| |
| ForOp::ensureTerminator(*body, builder, result.location); |
| |
| // Parse the optional attribute list. |
| if (parser.parseOptionalAttrDict(result.attributes)) |
| return failure(); |
| |
| return success(); |
| } |
| |
| Region &ForOp::getLoopBody() { return region(); } |
| |
| bool ForOp::isDefinedOutsideOfLoop(Value value) { |
| return !region().isAncestor(value.getParentRegion()); |
| } |
| |
| LogicalResult ForOp::moveOutOfLoop(ArrayRef<Operation *> ops) { |
| for (auto op : ops) |
| op->moveBefore(this->getOperation()); |
| return success(); |
| } |
| |
| ForOp mlir::loop::getForInductionVarOwner(Value val) { |
| auto ivArg = val.dyn_cast<BlockArgument>(); |
| if (!ivArg) |
| return ForOp(); |
| assert(ivArg.getOwner() && "unlinked block argument"); |
| auto *containingInst = ivArg.getOwner()->getParentOp(); |
| return dyn_cast_or_null<ForOp>(containingInst); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // IfOp |
| //===----------------------------------------------------------------------===// |
| |
| void IfOp::build(Builder *builder, OperationState &result, Value cond, |
| bool withElseRegion) { |
| result.addOperands(cond); |
| Region *thenRegion = result.addRegion(); |
| Region *elseRegion = result.addRegion(); |
| IfOp::ensureTerminator(*thenRegion, *builder, result.location); |
| if (withElseRegion) |
| IfOp::ensureTerminator(*elseRegion, *builder, result.location); |
| } |
| |
| static LogicalResult verify(IfOp op) { |
| // Verify that the entry of each child region does not have arguments. |
| for (auto ®ion : op.getOperation()->getRegions()) { |
| if (region.empty()) |
| continue; |
| |
| for (auto &b : region) |
| if (b.getNumArguments() != 0) |
| return op.emitOpError( |
| "requires that child entry blocks have no arguments"); |
| } |
| return success(); |
| } |
| |
| static ParseResult parseIfOp(OpAsmParser &parser, OperationState &result) { |
| // Create the regions for 'then'. |
| result.regions.reserve(2); |
| Region *thenRegion = result.addRegion(); |
| Region *elseRegion = result.addRegion(); |
| |
| auto &builder = parser.getBuilder(); |
| OpAsmParser::OperandType cond; |
| Type i1Type = builder.getIntegerType(1); |
| if (parser.parseOperand(cond) || |
| parser.resolveOperand(cond, i1Type, result.operands)) |
| return failure(); |
| |
| // Parse the 'then' region. |
| if (parser.parseRegion(*thenRegion, /*arguments=*/{}, /*argTypes=*/{})) |
| return failure(); |
| IfOp::ensureTerminator(*thenRegion, parser.getBuilder(), result.location); |
| |
| // If we find an 'else' keyword then parse the 'else' region. |
| if (!parser.parseOptionalKeyword("else")) { |
| if (parser.parseRegion(*elseRegion, /*arguments=*/{}, /*argTypes=*/{})) |
| return failure(); |
| IfOp::ensureTerminator(*elseRegion, parser.getBuilder(), result.location); |
| } |
| |
| // Parse the optional attribute list. |
| if (parser.parseOptionalAttrDict(result.attributes)) |
| return failure(); |
| |
| return success(); |
| } |
| |
| static void print(OpAsmPrinter &p, IfOp op) { |
| p << IfOp::getOperationName() << " " << op.condition(); |
| p.printRegion(op.thenRegion(), |
| /*printEntryBlockArgs=*/false, |
| /*printBlockTerminators=*/false); |
| |
| // Print the 'else' regions if it exists and has a block. |
| auto &elseRegion = op.elseRegion(); |
| if (!elseRegion.empty()) { |
| p << " else"; |
| p.printRegion(elseRegion, |
| /*printEntryBlockArgs=*/false, |
| /*printBlockTerminators=*/false); |
| } |
| |
| p.printOptionalAttrDict(op.getAttrs()); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // ParallelOp |
| //===----------------------------------------------------------------------===// |
| |
| static LogicalResult verify(ParallelOp op) { |
| // Check that there is at least one value in lowerBound, upperBound and step. |
| // It is sufficient to test only step, because it is ensured already that the |
| // number of elements in lowerBound, upperBound and step are the same. |
| Operation::operand_range stepValues = op.step(); |
| if (stepValues.empty()) |
| return op.emitOpError( |
| "needs at least one tuple element for lowerBound, upperBound and step"); |
| |
| // Check whether all constant step values are positive. |
| for (Value stepValue : stepValues) |
| if (auto cst = dyn_cast_or_null<ConstantIndexOp>(stepValue.getDefiningOp())) |
| if (cst.getValue() <= 0) |
| return op.emitOpError("constant step operand must be positive"); |
| |
| // Check that the body defines the same number of block arguments as the |
| // number of tuple elements in step. |
| Block *body = &op.body().front(); |
| if (body->getNumArguments() != stepValues.size()) |
| return op.emitOpError( |
| "expects the same number of induction variables as bound and step " |
| "values"); |
| for (auto arg : body->getArguments()) |
| if (!arg.getType().isIndex()) |
| return op.emitOpError( |
| "expects arguments for the induction variable to be of index type"); |
| |
| // Check that the number of results is the same as the number of ReduceOps. |
| SmallVector<ReduceOp, 4> reductions(body->getOps<ReduceOp>()); |
| if (op.results().size() != reductions.size()) |
| return op.emitOpError( |
| "expects number of results to be the same as number of reductions"); |
| |
| // Check that the types of the results and reductions are the same. |
| for (auto resultAndReduce : llvm::zip(op.results(), reductions)) { |
| auto resultType = std::get<0>(resultAndReduce).getType(); |
| auto reduceOp = std::get<1>(resultAndReduce); |
| auto reduceType = reduceOp.operand().getType(); |
| if (resultType != reduceType) |
| return reduceOp.emitOpError() |
| << "expects type of reduce to be the same as result type: " |
| << resultType; |
| } |
| return success(); |
| } |
| |
| static ParseResult parseParallelOp(OpAsmParser &parser, |
| OperationState &result) { |
| auto &builder = parser.getBuilder(); |
| // Parse an opening `(` followed by induction variables followed by `)` |
| SmallVector<OpAsmParser::OperandType, 4> ivs; |
| if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1, |
| OpAsmParser::Delimiter::Paren)) |
| return failure(); |
| |
| // Parse loop bounds. |
| SmallVector<OpAsmParser::OperandType, 4> lower; |
| if (parser.parseEqual() || |
| parser.parseOperandList(lower, ivs.size(), |
| OpAsmParser::Delimiter::Paren) || |
| parser.resolveOperands(lower, builder.getIndexType(), result.operands)) |
| return failure(); |
| |
| SmallVector<OpAsmParser::OperandType, 4> upper; |
| if (parser.parseKeyword("to") || |
| parser.parseOperandList(upper, ivs.size(), |
| OpAsmParser::Delimiter::Paren) || |
| parser.resolveOperands(upper, builder.getIndexType(), result.operands)) |
| return failure(); |
| |
| // Parse step value. |
| SmallVector<OpAsmParser::OperandType, 4> steps; |
| if (parser.parseKeyword("step") || |
| parser.parseOperandList(steps, ivs.size(), |
| OpAsmParser::Delimiter::Paren) || |
| parser.resolveOperands(steps, builder.getIndexType(), result.operands)) |
| return failure(); |
| |
| // Now parse the body. |
| Region *body = result.addRegion(); |
| SmallVector<Type, 4> types(ivs.size(), builder.getIndexType()); |
| if (parser.parseRegion(*body, ivs, types)) |
| return failure(); |
| |
| // Parse attributes and optional results (in case there is a reduce). |
| if (parser.parseOptionalAttrDict(result.attributes) || |
| parser.parseOptionalColonTypeList(result.types)) |
| return failure(); |
| |
| // Add a terminator if none was parsed. |
| ForOp::ensureTerminator(*body, builder, result.location); |
| |
| return success(); |
| } |
| |
| static void print(OpAsmPrinter &p, ParallelOp op) { |
| p << op.getOperationName() << " ("; |
| p.printOperands(op.body().front().getArguments()); |
| p << ") = (" << op.lowerBound() << ") to (" << op.upperBound() << ") step (" |
| << op.step() << ")"; |
| p.printRegion(op.body(), /*printEntryBlockArgs=*/false); |
| p.printOptionalAttrDict(op.getAttrs()); |
| if (!op.results().empty()) |
| p << " : " << op.getResultTypes(); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // ReduceOp |
| //===----------------------------------------------------------------------===// |
| |
| static LogicalResult verify(ReduceOp op) { |
| // The region of a ReduceOp has two arguments of the same type as its operand. |
| auto type = op.operand().getType(); |
| Block &block = op.reductionOperator().front(); |
| if (block.empty()) |
| return op.emitOpError("the block inside reduce should not be empty"); |
| if (block.getNumArguments() != 2 || |
| llvm::any_of(block.getArguments(), [&](const BlockArgument &arg) { |
| return arg.getType() != type; |
| })) |
| return op.emitOpError() |
| << "expects two arguments to reduce block of type " << type; |
| |
| // Check that the block is terminated by a ReduceReturnOp. |
| if (!isa<ReduceReturnOp>(block.getTerminator())) |
| return op.emitOpError("the block inside reduce should be terminated with a " |
| "'loop.reduce.return' op"); |
| |
| return success(); |
| } |
| |
| static ParseResult parseReduceOp(OpAsmParser &parser, OperationState &result) { |
| // Parse an opening `(` followed by the reduced value followed by `)` |
| OpAsmParser::OperandType operand; |
| if (parser.parseLParen() || parser.parseOperand(operand) || |
| parser.parseRParen()) |
| return failure(); |
| |
| // Now parse the body. |
| Region *body = result.addRegion(); |
| if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{})) |
| return failure(); |
| |
| // And the type of the operand (and also what reduce computes on). |
| Type resultType; |
| if (parser.parseColonType(resultType) || |
| parser.resolveOperand(operand, resultType, result.operands)) |
| return failure(); |
| |
| return success(); |
| } |
| |
| static void print(OpAsmPrinter &p, ReduceOp op) { |
| p << op.getOperationName() << "(" << op.operand() << ") "; |
| p.printRegion(op.reductionOperator()); |
| p << " : " << op.operand().getType(); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // ReduceReturnOp |
| //===----------------------------------------------------------------------===// |
| |
| static LogicalResult verify(ReduceReturnOp op) { |
| // The type of the return value should be the same type as the type of the |
| // operand of the enclosing ReduceOp. |
| auto reduceOp = cast<ReduceOp>(op.getParentOp()); |
| Type reduceType = reduceOp.operand().getType(); |
| if (reduceType != op.result().getType()) |
| return op.emitOpError() << "needs to have type " << reduceType |
| << " (the type of the enclosing ReduceOp)"; |
| return success(); |
| } |
| |
| static ParseResult parseReduceReturnOp(OpAsmParser &parser, |
| OperationState &result) { |
| OpAsmParser::OperandType operand; |
| Type resultType; |
| if (parser.parseOperand(operand) || parser.parseColonType(resultType) || |
| parser.resolveOperand(operand, resultType, result.operands)) |
| return failure(); |
| |
| return success(); |
| } |
| |
| static void print(OpAsmPrinter &p, ReduceReturnOp op) { |
| p << op.getOperationName() << " " << op.result() << " : " |
| << op.result().getType(); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // TableGen'd op method definitions |
| //===----------------------------------------------------------------------===// |
| |
| #define GET_OP_CLASSES |
| #include "mlir/Dialect/LoopOps/LoopOps.cpp.inc" |