From 34be80aa6edda60e240e4ea46f28b1b54ababf72 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou@gmail.com>
Date: Thu, 15 May 2025 09:19:39 +0800
Subject: [PATCH] [mlir-runner] Check entry function does not expect arguments
 (#136825)

This PR fixes a crash if entry function has inputs. Fixes #136143.
---
 mlir/lib/ExecutionEngine/JitRunner.cpp        | 10 +++-
 .../verify-entry-point-result.mlir            |  7 ---
 mlir/test/mlir-runner/verify-entry-point.mlir | 48 +++++++++++++++++++
 3 files changed, 56 insertions(+), 9 deletions(-)
 delete mode 100644 mlir/test/mlir-runner/verify-entry-point-result.mlir
 create mode 100644 mlir/test/mlir-runner/verify-entry-point.mlir

diff --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp
index cf462ddf6f17..2107df37d199 100644
--- a/mlir/lib/ExecutionEngine/JitRunner.cpp
+++ b/mlir/lib/ExecutionEngine/JitRunner.cpp
@@ -222,9 +222,14 @@ static Error compileAndExecuteVoidFunction(
     CompileAndExecuteConfig config, std::unique_ptr<llvm::TargetMachine> tm) {
   auto mainFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(
       SymbolTable::lookupSymbolIn(module, entryPoint));
-  if (!mainFunction || mainFunction.empty())
+  if (!mainFunction || mainFunction.isExternal())
     return makeStringError("entry point not found");
 
+  if (cast<LLVM::LLVMFunctionType>(mainFunction.getFunctionType())
+          .getNumParams() != 0)
+    return makeStringError(
+        "JIT can't invoke a main function expecting arguments");
+
   auto resultType = dyn_cast<LLVM::LLVMVoidType>(
       mainFunction.getFunctionType().getReturnType());
   if (!resultType)
@@ -274,7 +279,8 @@ Error compileAndExecuteSingleReturnFunction(
 
   if (cast<LLVM::LLVMFunctionType>(mainFunction.getFunctionType())
           .getNumParams() != 0)
-    return makeStringError("function inputs not supported");
+    return makeStringError(
+        "JIT can't invoke a main function expecting arguments");
 
   if (Error error = checkCompatibleReturnType<Type>(mainFunction))
     return error;
diff --git a/mlir/test/mlir-runner/verify-entry-point-result.mlir b/mlir/test/mlir-runner/verify-entry-point-result.mlir
deleted file mode 100644
index ad46e0b5fe1b..000000000000
--- a/mlir/test/mlir-runner/verify-entry-point-result.mlir
+++ /dev/null
@@ -1,7 +0,0 @@
-// RUN: not mlir-runner %s -e entry -entry-point-result=void 2>&1 | FileCheck %s
-
-// CHECK: Error: expected void function
-llvm.func @entry() -> (i32) {
-  %0 = llvm.mlir.constant(0 : index) : i32
-  llvm.return %0 : i32
-}
diff --git a/mlir/test/mlir-runner/verify-entry-point.mlir b/mlir/test/mlir-runner/verify-entry-point.mlir
new file mode 100644
index 000000000000..c7165bd46302
--- /dev/null
+++ b/mlir/test/mlir-runner/verify-entry-point.mlir
@@ -0,0 +1,48 @@
+// RUN: not mlir-runner %s -e entry_point_void -entry-point-result=void 2>&1 | FileCheck %s --check-prefix=CHECK-ENTRY-POINT-VOID
+// RUN: not mlir-runner %s -e entry_inputs_void -entry-point-result=void 2>&1 | FileCheck %s --check-prefix=CHECK-ENTRY-INPUTS-VOID
+// RUN: not mlir-runner %s -e entry_result_void -entry-point-result=void 2>&1 | FileCheck %s --check-prefix=CHECK-ENTRY-RESULT-VOID
+// RUN: not mlir-runner %s -e entry_point_i32 -entry-point-result=i32 2>&1 | FileCheck %s --check-prefix=CHECK-ENTRY-POINT-I32
+// RUN: not mlir-runner %s -e entry_inputs_i32 -entry-point-result=i32 2>&1 | FileCheck %s --check-prefix=CHECK-ENTRY-INPUTS-I32
+// RUN: not mlir-runner %s -e entry_result_i32 -entry-point-result=i32 2>&1 | FileCheck %s --check-prefix=CHECK-ENTRY-RESULT-I32
+// RUN: not mlir-runner %s -e entry_result_i64 -entry-point-result=i64 2>&1 | FileCheck %s --check-prefix=CHECK-ENTRY-RESULT-I64
+// RUN: not mlir-runner %s -e entry_result_f32 -entry-point-result=f32 2>&1 | FileCheck %s --check-prefix=CHECK-ENTRY-RESULT-F32
+
+// CHECK-ENTRY-POINT-VOID: Error: entry point not found
+llvm.func @entry_point_void() -> ()
+
+// CHECK-ENTRY-INPUTS-VOID: Error: JIT can't invoke a main function expecting arguments
+llvm.func @entry_inputs_void(%arg0: i32) {
+  llvm.return
+}
+
+// CHECK-ENTRY-RESULT-VOID: Error: expected void function
+llvm.func @entry_result_void() -> (i32) {
+  %0 = llvm.mlir.constant(0 : index) : i32
+  llvm.return %0 : i32
+}
+
+// CHECK-ENTRY-POINT-I32: Error: entry point not found
+llvm.func @entry_point_i32() -> (i32)
+
+// CHECK-ENTRY-INPUTS-I32: Error: JIT can't invoke a main function expecting arguments
+llvm.func @entry_inputs_i32(%arg0: i32) {
+  llvm.return
+}
+
+// CHECK-ENTRY-RESULT-I32: Error: only single i32 function result supported
+llvm.func @entry_result_i32() -> (i64) {
+  %0 = llvm.mlir.constant(0 : index) : i64
+  llvm.return %0 : i64
+}
+
+// CHECK-ENTRY-RESULT-I64: Error: only single i64 function result supported
+llvm.func @entry_result_i64() -> (i32) {
+  %0 = llvm.mlir.constant(0 : index) : i32
+  llvm.return %0 : i32
+}
+
+// CHECK-ENTRY-RESULT-F32: Error: only single f32 function result supported
+llvm.func @entry_result_f32() -> (i32) {
+  %0 = llvm.mlir.constant(0 : index) : i32
+  llvm.return %0 : i32
+}