Move the ReturnOp type checks to ReturnOp::verify.

This was left as a TODO in the code.  Move the type verification from
MLFuncVerifier::verifyReturn to ReturnOp::verify.  Since the return operation
can only appear as the last statement of an MLFunction, i.e. where the
surrounding block is the function itself, it is easy to access the function
descriptor (ReturnOp::verify already relies on this).  From the function
descriptor, one can easily access the type information.  Note that this
slightly modifies the error message due to the use of emitOpError instead of a
plain emitError.

Drop the obsolete TODO comment in MLFunction::verify about checking that
"return" only appears as the last operation of an MLFunction since
ReturnOp::verify explicitly checks for that.
PiperOrigin-RevId: 218347843

diff --git a/mlir/lib/Analysis/Verifier.cpp b/mlir/lib/Analysis/Verifier.cpp
index f4223863488b606eb0475b044d60ca7d6f48df34..9e0033dbd0ff3b3069464fc7da822aa600a43e89 100644 (file)
--- a/mlir/lib/Analysis/Verifier.cpp
+++ b/mlir/lib/Analysis/Verifier.cpp
@@ -371,8 +371,6 @@ struct MLFuncVerifier : public Verifier, public StmtWalker<MLFuncVerifier> {
     if (hadError)
       return true;
 
-    // TODO: check that operation is not a return statement unless it's
-    // the last one in the function.
     // TODO: check that loop bounds and if conditions are properly formed.
     if (verifyReturn())
       return true;
@@ -463,21 +461,6 @@ bool MLFuncVerifier::verifyReturn() {
     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);

diff --git a/mlir/lib/IR/BuiltinOps.cpp b/mlir/lib/IR/BuiltinOps.cpp
index c8b570af2e324ac732a82ff1d54ab705a2dc0258..94bcd16fcbd2056194548a34cf29338fe6238fcf 100644 (file)
--- a/mlir/lib/IR/BuiltinOps.cpp
+++ b/mlir/lib/IR/BuiltinOps.cpp
@@ -322,10 +322,24 @@ bool ReturnOp::verify() const {
     if (!block || !isa<MLFunction>(block) || &block->back() != stmt)
       return emitOpError("must be the last statement in the ML function");
 
+    // The operand number and types must match the function signature.
+    MLFunction *function = cast<MLFunction>(block);
+    const auto &results = function->getType()->getResults();
+    if (stmt->getNumOperands() != results.size())
+      return emitOpError("has " + Twine(stmt->getNumOperands()) +
+                         " operands, but enclosing function returns " +
+                         Twine(results.size()));
+
+    for (unsigned i = 0, e = results.size(); i != e; ++i)
+      if (stmt->getOperand(i)->getType() != results[i]) {
+        emitError("type of return operand " + Twine(i) +
+                  " doesn't match function result type");
+        return true;
+      }
+
     // Return success. Checking that operand types match those in the function
     // signature is performed in the ML function verifier.
     return false;
   }
   return emitOpError("cannot occur in a CFG function");
 }
-

diff --git a/mlir/test/IR/invalid.mlir b/mlir/test/IR/invalid.mlir
index 11aa1b1b3527c713bbe70dae3b5237edc35974d3..f1621a0bc4c30363e15e5686d8eacad65f9cac21 100644 (file)
--- a/mlir/test/IR/invalid.mlir
+++ b/mlir/test/IR/invalid.mlir
@@ -325,7 +325,7 @@ bb42:
 // -----
 
 mlfunc @mlfunc_resulterror() -> i32 {
-  return // expected-error {{return has 0 operands, but enclosing function returns 1}}
+  return // expected-error {{'return' op has 0 operands, but enclosing function returns 1}}
 }
 
 // -----