diff --git a/include/mlir/IR/MLFunction.h b/include/mlir/IR/MLFunction.h
index 248828c..ed12afc 100644
--- a/include/mlir/IR/MLFunction.h
+++ b/include/mlir/IR/MLFunction.h
@@ -80,6 +80,10 @@
 
   ~MLFunction();
 
+  // Return the 'return' statement of this MLFunction.
+  const OperationStmt *getReturnStmt() const;
+  OperationStmt *getReturnStmt();
+
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(const Function *func) {
     return func->getKind() == Function::Kind::MLFunc;
diff --git a/lib/IR/Function.cpp b/lib/IR/Function.cpp
index 4719b35..1c89de8 100644
--- a/lib/IR/Function.cpp
+++ b/lib/IR/Function.cpp
@@ -152,3 +152,11 @@
   this->~MLFunction();
   free(this);
 }
+
+const OperationStmt *MLFunction::getReturnStmt() const {
+  return cast<OperationStmt>(&back());
+}
+
+OperationStmt *MLFunction::getReturnStmt() {
+  return cast<OperationStmt>(&back());
+}