[MLIR][LLVM] Avoid importing broken calls and invokes

[Dinistro]

This commit adds a check to catch calls/invokes that do not satisfy the type constraints of their callee. This is not verified in LLVM IR but is considered UB. Importing this into MLIR will lead to verification errors, thus we should avoid this early on.

[llvmbot]

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Christian Ulmann (Dinistro)

Changes

This commit adds a check to catch calls/invokes that do not satisfy the return type of their callee. This is not verified in LLVM IR but is considered UB. Importing this into MLIR will lead to verification errors, thus we should avoid this early on.

---

Full diff: https://github.com/llvm/llvm-project/pull/125041.diff

2 Files Affected:

• (modified) mlir/lib/Target/LLVMIR/ModuleImport.cpp (+27-15)
• (modified) mlir/test/Target/LLVMIR/Import/import-failure.ll (+31-1)

diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index 1d1a985c46fb5b..d2e5390e397933 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1620,7 +1620,7 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
     return success();
   }
   if (inst->getOpcode() == llvm::Instruction::Call) {
-    auto callInst = cast<llvm::CallInst>(inst);
+    auto *callInst = cast<llvm::CallInst>(inst);
     llvm::Value *calledOperand = callInst->getCalledOperand();
 
     FailureOr<SmallVector<Value>> operands =
@@ -1629,10 +1629,10 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
       return failure();
 
     auto callOp = [&]() -> FailureOr<Operation *> {
-      if (auto asmI = dyn_cast<llvm::InlineAsm>(calledOperand)) {
-        Type resultTy = convertType(callInst->getType());
-        if (!resultTy)
-          return failure();
+      Type resultTy = convertType(callInst->getType());
+      if (!resultTy)
+        return failure();
+      if (auto *asmI = dyn_cast<llvm::InlineAsm>(calledOperand)) {
         return builder
             .create<InlineAsmOp>(
                 loc, resultTy, *operands,
@@ -1642,17 +1642,21 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
                 /*is_align_stack=*/false, /*asm_dialect=*/nullptr,
                 /*operand_attrs=*/nullptr)
             .getOperation();
-      } else {
-        LLVMFunctionType funcTy = convertFunctionType(callInst);
-        if (!funcTy)
-          return failure();
-
-        FlatSymbolRefAttr callee = convertCalleeName(callInst);
-        auto callOp = builder.create<CallOp>(loc, funcTy, callee, *operands);
-        if (failed(convertCallAttributes(callInst, callOp)))
-          return failure();
-        return callOp.getOperation();
       }
+      LLVMFunctionType funcTy = convertFunctionType(callInst);
+      if (!funcTy)
+        return failure();
+
+      if (funcTy.getReturnType() != resultTy)
+        return emitError(loc)
+               << "incompatible call and function return types: " << resultTy
+               << " vs. " << funcTy.getReturnType();
+
+      FlatSymbolRefAttr callee = convertCalleeName(callInst);
+      auto callOp = builder.create<CallOp>(loc, funcTy, callee, *operands);
+      if (failed(convertCallAttributes(callInst, callOp)))
+        return failure();
+      return callOp.getOperation();
     }();
 
     if (failed(callOp))
@@ -1720,6 +1724,14 @@ LogicalResult ModuleImport::convertInstruction(llvm::Instruction *inst) {
     if (!funcTy)
       return failure();
 
+    Type resultTy = convertType(invokeInst->getType());
+    if (!resultTy)
+      return failure();
+    if (funcTy.getReturnType() != resultTy)
+      return emitError(loc)
+             << "incompatible invoke and function return types: " << resultTy
+             << " vs. " << funcTy.getReturnType();
+
     FlatSymbolRefAttr calleeName = convertCalleeName(invokeInst);
 
     // Create the invoke operation. Normal destination block arguments will be
diff --git a/mlir/test/Target/LLVMIR/Import/import-failure.ll b/mlir/test/Target/LLVMIR/Import/import-failure.ll
index b616cb81e0a8a5..910f60fd81e8a0 100644
--- a/mlir/test/Target/LLVMIR/Import/import-failure.ll
+++ b/mlir/test/Target/LLVMIR/Import/import-failure.ll
@@ -1,4 +1,4 @@
-; RUN: not mlir-translate -import-llvm -emit-expensive-warnings -split-input-file %s 2>&1 | FileCheck %s
+; RUN: not mlir-translate -import-llvm -emit-expensive-warnings -split-input-file %s 2>&1 -o /dev/null | FileCheck %s
 
 ; CHECK:      <unknown>
 ; CHECK-SAME: error: unhandled instruction: indirectbr ptr %dst, [label %bb1, label %bb2]
@@ -353,3 +353,33 @@ declare void @llvm.experimental.noalias.scope.decl(metadata)
 ; CHECK:      import-failure.ll
 ; CHECK-SAME: warning: unhandled data layout token: ni:42
 target datalayout = "e-ni:42-i64:64"
+
+; // -----
+
+; CHECK:      <unknown>
+; CHECK-SAME: incompatible call and function return types: 'i64' vs. '!llvm.ptr'
+define i64 @incompatible_call_and_callee_types() {
+L.entry:
+  %0 = call i64 @callee()
+  ret i64 %0
+}
+
+declare ptr @callee()
+
+; // -----
+
+; CHECK:      <unknown>
+; CHECK-SAME: incompatible invoke and function return types: '!llvm.void' vs. 'i32'
+define void @f() personality ptr @__gxx_personality_v0 {
+entry:
+  invoke void @g() to label %bb1 unwind label %bb2
+bb1:
+  ret void
+bb2:
+  %0 = landingpad i32 cleanup
+  unreachable
+}
+
+declare i32 @g()
+
+declare i32 @__gxx_personality_v0(...)

[Dinistro]

Note that I've refined this to check not just the return types, but the actual full function types.

[gysit]

Nice!

LGTM

[xlauko]

LGTM