[mlir-runner] Check entry function does not expect arguments (#136825)

This PR fixes a crash if entry function has inputs. Fixes #136143.
This commit is contained in:
Longsheng Mou 2025-05-15 09:19:39 +08:00 committed by GitHub
parent 18b885f66b
commit 34be80aa6e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 56 additions and 9 deletions

View File

@ -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;

View File

@ -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
}

View File

@ -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
}