Minor improvements to the return operation implementation.
PiperOrigin-RevId: 208166863
diff --git a/include/mlir/IR/StandardOps.h b/include/mlir/IR/StandardOps.h
index 4c71b7e..99cbc59 100644
--- a/include/mlir/IR/StandardOps.h
+++ b/include/mlir/IR/StandardOps.h
@@ -299,7 +299,6 @@
explicit StoreOp(const Operation *state) : OpBase(state) {}
};
-/// TODO: change comment.
/// The "return" operation represents a return statement of an ML function.
/// The operation takes variable number of operands and produces no results.
/// The operand number and types must match the signature of the ML function
diff --git a/lib/IR/StandardOps.cpp b/lib/IR/StandardOps.cpp
index c3f815f..e78290e 100644
--- a/lib/IR/StandardOps.cpp
+++ b/lib/IR/StandardOps.cpp
@@ -369,7 +369,6 @@
bool ReturnOp::parse(OpAsmParser *parser, OperationState *result) {
SmallVector<OpAsmParser::OperandType, 2> opInfo;
SmallVector<Type *, 2> types;
- SmallVector<SSAValue *, 2> operands;
return parser->parseOperandList(opInfo, -1, OpAsmParser::Delimiter::None) ||
(!opInfo.empty() && parser->parseColonTypeList(types)) ||
@@ -391,17 +390,15 @@
const char *ReturnOp::verify() const {
// ReturnOp must be part of an ML function.
if (auto *stmt = dyn_cast<OperationStmt>(getOperation())) {
- StmtBlock *block = stmt->getBlock();
-
- if (!block || !isa<MLFunction>(block) ||
- &cast<MLFunction>(block)->back() != stmt)
+ MLFunction *func = dyn_cast_or_null<MLFunction>(stmt->getBlock());
+ if (!func || &func->back() != stmt)
return "must be the last statement in the ML function";
// Return success. Checking that operand types match those in the function
// signature is performed in the ML function verifier.
return nullptr;
}
- return "cannot occur in a CFG function.";
+ return "cannot occur in a CFG function";
}
//===----------------------------------------------------------------------===//
@@ -470,6 +467,6 @@
/// Install the standard operations in the specified operation set.
void mlir::registerStandardOperations(OperationSet &opSet) {
opSet.addOperations<AddFOp, AffineApplyOp, AllocOp, ConstantOp, DimOp, LoadOp,
- StoreOp, ReturnOp>(
+ ReturnOp, StoreOp>(
/*prefix=*/"");
}