Implement return statement as RetOp operation. Add verification of the return statement placement and operands. Add parser and parsing error tests for return statements with non-zero number of operands. Add a few missing tests for ForStmt parsing errors.
Prior to this CL, return statement had no explicit representation in MLIR. Now, it is represented as ReturnOp standard operation and is pretty printed according to the return statement syntax. This way statement walkers can process ML function return operands without making special case for them.
PiperOrigin-RevId: 208092424
diff --git a/lib/IR/Verifier.cpp b/lib/IR/Verifier.cpp
index dde196c..272395c 100644
--- a/lib/IR/Verifier.cpp
+++ b/lib/IR/Verifier.cpp
@@ -207,7 +207,7 @@
for (unsigned i = 0, e = results.size(); i != e; ++i)
if (inst.getOperand(i)->getType() != results[i])
return failure("type of return operand " + Twine(i) +
- " doesn't match result function result type",
+ " doesn't match function result type",
inst);
return false;
@@ -313,7 +313,11 @@
llvm::PrettyStackTraceFormat fmt("MLIR Verifier: mlfunc @%s",
fn.getName().c_str());
- // TODO: check basic structural properties.
+ // TODO: check basic structural properties
+ // TODO: check that operation is not a return statement unless it's
+ // the last one in the function.
+ if (verifyReturn())
+ return true;
return verifyDominance();
}
@@ -321,6 +325,9 @@
/// Walk all of the code in this MLFunc and verify that the operands of any
/// operations are properly dominated by their definitions.
bool verifyDominance();
+
+ /// Verify that function has a return statement that matches its signature.
+ bool verifyReturn();
};
} // end anonymous namespace
@@ -390,6 +397,37 @@
return walkBlock(fn);
}
+bool MLFuncVerifier::verifyReturn() {
+ // TODO: fold return verification in the pass that verifies all statements.
+ const char missingReturnMsg[] = "ML function must end with return statement";
+ if (fn.getStatements().empty())
+ return failure(missingReturnMsg, fn);
+
+ const auto &stmt = fn.getStatements().back();
+ if (const auto *op = dyn_cast<OperationStmt>(&stmt)) {
+ if (!op->isReturn())
+ return failure(missingReturnMsg, fn);
+
+ // The operand number and types must match the function signature.
+ // TODO: move this verification in ReturnOp::verify() if printing
+ // of the error messages below can be made to work there.
+ const auto &results = fn.getType()->getResults();
+ if (op->getNumOperands() != results.size())
+ return failure("return has " + Twine(op->getNumOperands()) +
+ " operands, but enclosing function returns " +
+ Twine(results.size()),
+ *op);
+
+ for (unsigned i = 0, e = results.size(); i != e; ++i)
+ if (op->getOperand(i)->getType() != results[i])
+ return failure("type of return operand " + Twine(i) +
+ " doesn't match function result type",
+ *op);
+ return false;
+ }
+ return failure(missingReturnMsg, fn);
+}
+
//===----------------------------------------------------------------------===//
// Entrypoints
//===----------------------------------------------------------------------===//